Stable-Baselines3 — Reliable Reinforcement Learning in PyTorch

Introduction

Stable-Baselines3 (SB3) is a PyTorch-based library providing reliable, well-tested implementations of popular reinforcement learning algorithms. It is the successor to the original Stable Baselines (TensorFlow) and focuses on code clarity, reproducibility, and ease of use. SB3 is widely adopted in both academic research and applied RL projects.

What Stable-Baselines3 Does

• Implements core RL algorithms including PPO, A2C, SAC, TD3, DQN, and HER
• Provides a consistent API across all algorithms for training, evaluation, and deployment
• Supports custom environments via the Gymnasium interface
• Includes callback system for logging, early stopping, and checkpoint saving
• Offers vectorized environments for parallel data collection during training

Architecture Overview

SB3 follows a modular design with base classes for on-policy and off-policy algorithms. Each algorithm inherits from these bases and implements its specific update logic. Policies (actor-critic networks) are separate modules that can be customized or replaced. The training loop handles rollout collection, buffer management, and gradient updates, while callbacks hook into the loop for monitoring and control.

Self-Hosting & Configuration

• Install via pip install stable-baselines3 with optional extras for additional environments
• Requires Python 3.8+ with PyTorch and Gymnasium as core dependencies
• Configure hyperparameters (learning rate, batch size, network architecture) via constructor arguments
• Save and load trained models with model.save() and model.load() for deployment
• Use the rl-zoo3 companion project for hyperparameter optimization with Optuna

Key Features

• Thoroughly tested implementations with unit tests and integration tests for every algorithm
• Consistent API: switching between PPO, SAC, or DQN requires changing one line of code
• TensorBoard integration for real-time training visualization
• Support for Dict and MultiDiscrete observation and action spaces
• SB3-Contrib extends the library with RecurrentPPO, TQC, CrossQ, and other experimental algorithms

Comparison with Similar Tools

• CleanRL — single-file RL implementations for maximum transparency; SB3 provides production-ready modular code
• RLlib (Ray) — distributed RL library for large-scale training; SB3 is simpler and focused on single-machine use
• Tianshou — PyTorch RL library with broader algorithm coverage; SB3 prioritizes reliability and testing rigor
• Spinning Up — educational RL implementations by OpenAI; SB3 is designed for practical research and deployment

FAQ

Q: Which algorithm should I start with? A: PPO is a good default for most tasks. Use SAC or TD3 for continuous control, and DQN for discrete action spaces.

Q: Can I use custom neural network architectures? A: Yes. Pass a custom policy class or use the policy_kwargs parameter to modify layer sizes, activation functions, or feature extractors.

Q: How do I tune hyperparameters? A: Use the rl-zoo3 companion project, which provides Optuna-based hyperparameter optimization with predefined search spaces for each algorithm.

Q: Does SB3 support multi-agent RL? A: SB3 is designed for single-agent settings. For multi-agent RL, consider PettingZoo with compatible wrappers or other frameworks.

Sources

• https://github.com/DLR-RM/stable-baselines3
• https://stable-baselines3.readthedocs.io/