Stable Baselines 3 PyTorch

Integrate W&B with Stable Baselines3 to track reinforcement learning experiments and log training performance.

Stable Baselines 3 (SB3) is a set of reliable implementations of reinforcement learning algorithms in PyTorch. The W&B SB3 integration:

• Records metrics such as losses and episodic returns.
• Uploads videos of agents playing the games.
• Saves the trained model.
• Logs the model’s hyperparameters.
• Logs the model’s gradient histograms.

Review an example SB3 training run.

Log your SB3 experiments

To log SB3 training to W&B, pass WandbCallback to your model’s learn method:

from wandb.integration.sb3 import WandbCallback

model.learn(..., callback=WandbCallback())

WandbCallback arguments

The following table describes the arguments you can pass to WandbCallback:

Argument	Usage
verbose	The verbosity of SB3 output.
model_save_path	Path to the folder where the model is saved. The default is None, so the model isn’t logged.
model_save_freq	Frequency to save the model.
gradient_save_freq	Frequency to log gradients. The default is 0, so gradients aren’t logged.

Basic example

The W&B SB3 integration uses the logs output from TensorBoard to log your metrics.

import gym
from stable_baselines3 import PPO
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv, VecVideoRecorder
import wandb
from wandb.integration.sb3 import WandbCallback


config = {
    "policy_type": "MlpPolicy",
    "total_timesteps": 25000,
    "env_name": "CartPole-v1",
}
run = wandb.init(
    project="sb3",
    config=config,
    sync_tensorboard=True,  # auto-upload sb3's tensorboard metrics
    monitor_gym=True,  # auto-upload the videos of agents playing the game
    save_code=True,  # optional
)


def make_env():
    env = gym.make(config["env_name"])
    env = Monitor(env)  # record stats such as returns
    return env


env = DummyVecEnv([make_env])
env = VecVideoRecorder(
    env,
    f"videos/{run.id}",
    record_video_trigger=lambda x: x % 2000 == 0,
    video_length=200,
)
model = PPO(config["policy_type"], env, verbose=1, tensorboard_log=f"runs/{run.id}")
model.learn(
    total_timesteps=config["total_timesteps"],
    callback=WandbCallback(
        gradient_save_freq=100,
        model_save_path=f"models/{run.id}",
        verbose=2,
    ),
)
run.finish()