Stable Baselines3

Stable Baselines3

Overview

Stable Baselines3 (SB3) is a PyTorch-based library providing reliable implementations of reinforcement learning algorithms. This skill provides comprehensive guidance for training RL agents, creating custom environments, implementing callbacks, and optimizing training workflows using SB3's unified API.

Current upstream: SB3 2.9.0 (June 15, 2026). Docs: stable-baselines3.readthedocs.io.

Installation

Tested against stable-baselines3 2.9.0. Requires Python 3.10+ (3.9 dropped in 2.8.0) and PyTorch >= 2.8.

# Basic installation
uv pip install "stable-baselines3==2.9.0"

# With extra dependencies (TensorBoard, ale-py for Atari, etc.)
uv pip install "stable-baselines3[extra]==2.9.0"

The 2.9.0 release supports Gymnasium >=0.29.1,<2.0; this review exercised Gymnasium 1.3.0, PyTorch 2.14.1 and Python 3.13 on CPU. pandas/matplotlib are now optional extras. Installation requires network unless packages are cached; local RL training needs no credentials or service endpoints.

On zsh, quote brackets: uv pip install 'stable-baselines3[extra]==2.9.0'.

For MuJoCo continuous-control benchmarks:

uv pip install "gymnasium[mujoco]"

Check your version:

import stable_baselines3
print(stable_baselines3.__version__)

Related Projects

• SB3-Contrib: experimental algorithms (MaskablePPO, CrossQ, QR-DQN, RecurrentPPO) — separate sb3-contrib package
• RL Baselines3 Zoo: pre-trained agents, hyperparameters, training scripts
• SBX: SB3 + JAX implementations for users who prefer JAX over PyTorch

Core Capabilities

1. Training RL Agents

Basic Training Pattern:

import gymnasium as gym
from stable_baselines3 import PPO

# Create environment
env = gym.make("CartPole-v1")

# Initialize agent (device="cpu" is often faster for MlpPolicy on small envs)
model = PPO("MlpPolicy", env, verbose=1, device="cpu", seed=0)

# Train the agent
model.learn(total_timesteps=10000)

# Save the model
model.save("ppo_cartpole")

# Load the model (without prior instantiation)
model = PPO.load("ppo_cartpole", env=env, device="cpu")
env.close()

Important Notes:

• total_timesteps is a lower bound; actual training may exceed this due to batch collection
• Call the class method PPO.load(...) and keep the returned new model
• The replay buffer is NOT saved with the model to save space

Algorithm Selection: Use references/algorithms.md for detailed algorithm characteristics and selection guidance. Quick reference:

• PPO/A2C: General-purpose, supports Box, Discrete, flat MultiDiscrete and MultiBinary actions, good for multiprocessing
• SAC/TD3: Continuous control, off-policy, sample-efficient
• DQN: Discrete actions, off-policy
• HER: Replay-buffer strategy for goal-conditioned off-policy tasks

See scripts/train_rl_agent.py for a complete training template with best practices.

2. Custom Environments

Requirements: Custom environments must inherit from gymnasium.Env and implement:

• __init__(): Define action_space and observation_space
• reset(seed, options): Return initial observation and info dict
• step(action): Return observation, reward, terminated, truncated, info
• render(): Visualization (optional)
• close(): Cleanup resources

Key Constraints:

• Default CNN image preprocessing expects np.uint8 in range [0, 255]
• Use channel-first format when possible (channels, height, width)
• SB3 normalizes images automatically by dividing by 255
• For pre-normalized float images, use channel-first layout and policy_kwargs={"normalize_images": False}
• SB3 does NOT support Discrete or MultiDiscrete spaces with start!=0

Validation:

from stable_baselines3.common.env_checker import check_env

check_env(env, warn=True)

See the template and environment guide. The template now observes both agent and random goal coordinates, shape (4,); old shape (2,) checkpoints require retraining. gym.make("CustomEnv-v0") adds the 100-step time limit; direct CustomEnv() does not. check_env checks API consistency, not Markov sufficiency, reward correctness or learnability.

3. Vectorized Environments

Purpose: Vectorized environments run multiple environment instances in parallel, accelerating training and enabling certain wrappers (frame-stacking, normalization).

Types:

• DummyVecEnv: Sequential execution on current process (for lightweight environments)
• SubprocVecEnv: Parallel execution across processes (for compute-heavy environments)

Quick Setup:

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# DummyVecEnv batches 4 lightweight environments sequentially.
env = make_vec_env("CartPole-v1", n_envs=4, seed=0)
try:
    model = PPO("MlpPolicy", env, verbose=1, device="cpu")
    model.learn(total_timesteps=25000)
finally:
    env.close()

Off-Policy Optimization: With step-based train_freq, gradient_steps=-1 matches gradient updates to collected transitions (train_freq * n_envs) after warmup. This changes compute and reuse of data; benchmark it rather than assuming it is always faster. SubprocVecEnv creation belongs under a main guard in a Python file.

API Differences:

• reset() returns only observations (info available in vec_env.reset_infos)
• step() returns 4-tuple: (obs, rewards, dones, infos) not 5-tuple
• Environments auto-reset after episodes
• Terminal observations available via infos[env_idx]["terminal_observation"]

See references/vectorized_envs.md for detailed information on wrappers and advanced usage.

4. Callbacks for Monitoring and Control

Purpose: Callbacks enable monitoring metrics, saving checkpoints, implementing early stopping, and custom training logic without modifying core algorithms.

Common Callbacks:

• EvalCallback: Evaluate periodically and save best model
• CheckpointCallback: Save model checkpoints at intervals
• StopTrainingOnRewardThreshold: Stop when target reward reached
• ProgressBarCallback: Display training progress with timing

Custom Callback Structure:

from stable_baselines3.common.callbacks import BaseCallback

class CustomCallback(BaseCallback):
    def _on_training_start(self):
        # Called before first rollout
        pass

    def _on_step(self):
        # Called after each environment step
        # Return False to stop training
        return True

    def _on_rollout_end(self):
        # Called at end of rollout
        pass

Available Attributes:

• self.model: The RL algorithm instance
• self.num_timesteps: Total environment steps
• self.training_env: The training environment

Chaining Callbacks:

from stable_baselines3.common.callbacks import CallbackList

callback = CallbackList([eval_callback, checkpoint_callback, custom_callback])
model.learn(total_timesteps=10000, callback=callback)

See references/callbacks.md for comprehensive callback documentation.

5. Model Persistence and Inspection

Saving and Loading:

from stable_baselines3.common.vec_env import VecNormalize

# After training with VecNormalize, save a matching pair:
model.save("model_name")
model.get_vec_normalize_env().save("vec_normalize.pkl")

# Build the same underlying environment and wrappers before loading:
vec_env = make_vec_env("Pendulum-v1", n_envs=1, seed=20000)
vec_env = VecNormalize.load("vec_normalize.pkl", vec_env)
vec_env.training = False
vec_env.norm_reward = False
model = PPO.load("model_name", env=vec_env, device="cpu")

This is a continuation fragment for a PPO/Pendulum run with normalization. Load only trusted model/statistics files. For off-policy training continuation, save_replay_buffer() / load_replay_buffer() are separate from save() / load(). Resume with a live environment and learn(..., reset_num_timesteps=False).

Parameter Access:

# Get parameters
params = model.get_parameters()

# Set parameters
model.set_parameters(params)

# Access PyTorch state dict
state_dict = model.policy.state_dict()