Stable Baselines 3 PyTorch - Weights & Biases Documentation

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()

​Log your SB3 experiments

​WandbCallback arguments

​Basic example

Log your SB3 experiments

`WandbCallback` arguments

Basic example