stable-baselines3

Stable Baselines3 (SB3) is a set of reliable, well-tested implementations of reinforcement learning algorithms in PyTorch, built as the successor to the original TensorFlow-based Stable Baselines. It ships eight core algorithms — A2C, DDPG, DQN, HER, PPO, SAC, and TD3, plus the HerReplayBuffer for goal-conditioned tasks — behind a single, sklearn-like interface (model = PPO("MlpPolicy", env).learn(total_timesteps)), so switching algorithms rarely requires touching training code.

Each algorithm is validated against published benchmarks and tracked on the OpenRL Benchmark platform, and the library is typed, linted, and covered by an extensive pytest suite that runs across four Python versions in CI. It supports custom environments and policies, Dict observation spaces, TensorBoard logging, and a callback system for evaluation, checkpointing, and early stopping — with a companion ecosystem (RL Baselines3 Zoo for tuned hyperparameters and training scripts, SB3-Contrib for experimental algorithms like Recurrent PPO and QR-DQN, and SBX for a JAX-accelerated variant) covering what the stable core intentionally leaves out.

What You Get

• Eight production-grade RL algorithms (A2C, DDPG, DQN, HER, PPO, SAC, TD3) sharing one training/prediction API
• A vectorized environment layer (DummyVecEnv, SubprocVecEnv, VecNormalize) for parallel rollouts and observation/reward normalization
• MLP and CNN policy networks with support for Box, Discrete, MultiDiscrete, MultiBinary, and Dict observation/action spaces
• A callback framework for TensorBoard logging, periodic evaluation, checkpointing, and custom training hooks
• Zip-based model save/load utilities for reproducible checkpointing and deployment
• Type hints and a py.typed marker throughout, backed by mypy in CI

Common Use Cases

• Training baseline RL agents on classic control, robotics, or Atari benchmarks before trying custom algorithms
• Prototyping a custom Gymnasium environment and validating it against known-good algorithm implementations
• Running reproducible RL experiments for research papers, citing the JMLR-published benchmark results
• Teaching or learning reinforcement learning with a simple, documented, sklearn-style API

Under The Hood

Architecture SB3 is organized as a layered hierarchy rooted in stable_baselines3/common/base_class.py’s abstract BaseAlgorithm, which is specialized into OnPolicyAlgorithm (common/on_policy_algorithm.py, backing A2C/PPO) and OffPolicyAlgorithm (common/off_policy_algorithm.py, backing DQN/SAC/TD3/DDPG); each concrete algorithm lives in its own top-level package (a2c/, ppo/, sac/, etc.) and plugs into shared infrastructure — common/policies.py for network heads, common/buffers.py for rollout/replay storage, common/callbacks.py for a composable callback chain, common/vec_env/ for environment vectorization and normalization, and common/save_util.py for zip-based serialization. This separation means adding a new algorithm mostly means implementing a train() loop against existing buffer/policy/callback contracts rather than rebuilding infrastructure, and changing the base class’s public methods (learn, predict, save, load) would ripple through every algorithm at once.

Tech Stack The library targets Python 3.10+ and depends on PyTorch (>=2.8,<3.0) as its sole deep-learning backend, Gymnasium (>=0.29.1,<2.0) for the environment API, NumPy for array operations, and cloudpickle for serialization; an extra install group adds TensorBoard, OpenCV, ale-py for Atari, and pandas/matplotlib for results analysis. Builds use plain setuptools (setup.py), with ruff and black for linting/formatting, mypy for static typing, and Sphinx (hosted on Read the Docs) for documentation; a Makefile drives lint, type, and pytest targets, and a Dockerfile is provided for containerized use.

Code Quality The tests/ directory contains over two dozen focused test modules covering buffers, callbacks, CNN and custom policies, determinism, dict/vectorized environments, distributions, GAE, HER, logging, monitoring, prediction, save/load round-tripping, and stochastic policies, run via pytest with coverage tracked in CI. The GitHub Actions CI matrix exercises Python 3.10 through 3.13 and multiple Gymnasium versions, running ruff lint, mypy type checks, and the full test suite on every push and pull request — combined with comprehensive docstrings and type hints, this gives the codebase a notably high bar for a research-adjacent library.

API Design The public API deliberately mirrors scikit-learn’s fit/predict conventions: constructing an algorithm with a policy name and environment, then calling .learn() and .predict(), is enough to train and run any of the eight algorithms with no per-algorithm boilerplate. Consistent constructor signatures across algorithms make swapping one for another close to a one-line change, and the documented custom-policy and custom-environment extension points keep the learning curve for going beyond the defaults manageable — the tradeoff is that the underlying implementation, being algorithm-agnostic infrastructure, does not attempt algorithmic novelty of its own.