DOCSEARCHbyAlgolia
Submit Docs

Stable Baselines3 (SB3)

repository·master·Indexed 12 days ago

https://github.com/dlr-rm/stable-baselines3

A set of reliable, high-quality implementations of reinforcement learning algorithms in PyTorch. SB3 provides a common interface for state-of-the-art RL methods, supporting algorithms such as PPO, A2C, DQN, and SAC across various Gymnasium action spaces including Box, Discrete, MultiDiscrete, and MultiBinary.

Tokens
50K
Snippets
113
Records
183
Agent score
95%

What's inside Stable Baselines3

  1. Overview of Stable Baselines3 (SB3)

    master

    Stable Baselines3 (SB3) is a collection of reliable reinforcement learning (RL) algorithm implementations built on PyTorch. It is designed to provide a unified structure for all algorithms, following PEP 8 standards, with high code coverage, type hints, and TensorBoard support.

    Key ecosystem components include:

    • RL Baselines3 Zoo: A training framework for SB3 that provides pre-trained agents, scripts for training, evaluating, hyperparameter tuning, plotting results, and recording videos.
    • SB3-Contrib: A repository for experimental RL features and the latest algorithms.
    • SBX (Stable-Baselines Jax): A version of SB3 that utilizes Jax for performance.

  2. What is Stable Baselines Jax (SBX)?

    master

    Stable Baselines Jax (SBX) is a proof-of-concept implementation of Stable-Baselines3 using JAX. While it offers a more minimal feature set than SB3, it can provide significant performance improvements, potentially up to 20x faster.

    SBX follows the SB3 API and is compatible with the RL Zoo framework.

  3. Use SB3-Contrib algorithms

    master

    SB3-Contrib provides implementations of experimental RL algorithms that follow the standard Stable-Baselines3 API and folder structure. This makes it easy to switch from standard SB3 algorithms to contrib algorithms by changing the import and class name.

    Available RL Algorithms

    • Augmented Random Search (ARS)
    • Quantile Regression DQN (QR-DQN)
    • Maskable PPO: PPO with invalid action masking
    • RecurrentPPO: PPO with a recurrent policy (PPO LSTM)
    • Truncated Quantile Critics (TQC)
    • Trust Region Policy Optimization (TRPO)
    • CrossQ: Batch Normalization in Deep Reinforcement Learning

    Available Gym Wrappers

    • Time Feature Wrapper

  4. Explore projects using Stable Baselines3

    master

    Stable Baselines3 is used across a wide variety of specialized research and application projects. Notable examples include:

    • Autonomous Driving: DriverGym (l5kit) and highway-env.
    • Robotics: RL Reach, Furuta Pendulum Robot (using gSDE), and tactile-gym.
    • Drones/UAVs: gym-pybullet-drones and UAV_Navigation_DRL_AirSim.
    • Specialized Control: gym-electric-motor (electric drives), Rocket League Gym (game environments), and mobile-env (wireless networks).
    • Exploration & Research: RLeXplore (intrinsic reward exploration) and Pink Noise Exploration (noise for off-policy algorithms).

  5. Available algorithms in the imitation library

    master

    The imitation library implements several imitation learning algorithms that leverage Stable-Baselines3:

    • Behavioral Cloning
    • DAgger (with synthetic examples)
    • Adversarial Inverse Reinforcement Learning (AIRL)
    • Generative Adversarial Imitation Learning (GAIL)
    • Deep RL from Human Preferences (DRLHP)

  6. Explore SB3 ecosystem and related repositories

    master

    Stable Baselines3 is part of a larger ecosystem for Reinforcement Learning:

    • SB3-Contrib: Contains experimental RL features (e.g., Recurrent PPO, Maskable PPO, QR-DQN) to keep the core SB3 stable.
    • RL Baselines3 Zoo: A training framework for training, evaluating, tuning hyperparameters, and benchmarking agents.
    • Stable-Baselines Jax (SBX): A high-performance proof-of-concept implementation using Jax, capable of up to 20x speedups.
    • Integrations: Supports Weights & Biases for tracking and Hugging Face for model sharing.

  7. Soft Actor Critic (SAC) Overview

    master

    Soft Actor Critic (SAC) is an Off-Policy Maximum Entropy Deep Reinforcement Learning algorithm with a Stochastic Actor. It is the successor to Soft Q-Learning (SQL) and incorporates the double Q-learning trick from TD3.

    A key feature of SAC is that it maximizes a trade-off between expected return and entropy (a measure of policy randomness).

    Implementation Details:

    • Uses an entropy coefficient (equivalent to the inverse of reward scale in the original paper) to avoid high errors during Q-function updates.
    • When automatically adjusting the temperature (alpha/entropy coefficient), the implementation optimizes the logarithm of the entropy coefficient for better stability.
    • MlpPolicy uses ReLU activation instead of tanh to match the original paper.

  8. Use EventCallback to trigger child callbacks

    master

    An EventCallback is a specialized BaseCallback used to trigger a secondary (child) callback when a specific event occurs. This is useful for conditional logic, such as stopping training only when a certain performance threshold is met.

    For example, EvalCallback is an EventCallback that triggers its child callback when a new best model is identified. A child callback like StopTrainingOnRewardThreshold can then be used to abort training based on that event.

    class EventCallback(BaseCallback):
    """
    Base class for triggering callback on event.

    :param callback: Callback that will be called when an event is triggered.
    :param verbose: Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages
    """
    def __init__(self, callback: BaseCallback, verbose: int = 0):
        super().__init__(verbose=verbose)
        self.callback = callback
        # Give access to the parent
        self.callback.parent = self

    def _on_event(self) -> bool:
        return self.callback()

  9. Handle incompatible Discrete and MultiDiscrete spaces

    master

    Stable Baselines3 has specific limitations regarding certain gymnasium.spaces configurations. Use the following wrappers to make your environment compatible:

    1. Discrete spaces with start != 0

    SB3 does not support Discrete spaces where the starting value is not 0. Use ShiftWrapper to shift the action space to start at 0 and adjust the step function accordingly.

    2. MultiDiscrete spaces with multi-dimensional arrays

    SB3 does not support MultiDiscrete spaces where the nvec attribute is a multi-dimensional array. Use ReshapeWrapper to flatten the nvec and reshape the action back to its original dimensions during the step function.

    import numpy as np
import gymnasium as gym

class ShiftWrapper(gym.Wrapper):
    """Allow to use Discrete() action spaces with start!=0"""
    def __init__(self, env: gym.Env) -> None:
        super().__init__(env)
        assert isinstance(env.action_space, gym.spaces.Discrete)
        self.action_space = gym.spaces.Discrete(env.action_space.n, start=0)
    
    def step(self, action: int):
        return self.env.step(action + env.action_space.start)

class ReshapeWrapper(gym.Wrapper):
    """Allow to use MultiDiscrete() action spaces with len(nvec.shape) > 1:""
    def __init__(self, env: gym.Env) -> None:
        super().__init__(env)
        assert isinstance(env.action_space, gym.spaces.MultiDiscrete)
        self.original_shape = env.action_space.nvec.shape
        self.action_space = gym.spaces.MultiDiscrete(env.action_space.nvec.flatten())
    
    def step(self, action: np.ndarray):
        return self.env.step(action.reshape(self.original_shape))

  10. How Hindsight Experience Replay (HER) works in Stable Baselines3

    master

    Hindsight Experience Replay (HER) is an algorithm designed for off-policy methods (such as DQN, SAC, TD3, and DDPG) that works by relabeling transitions. It creates "virtual" transitions by changing the desired goal to a goal that was actually achieved during a rollout, even if the original target goal was not met.

    Important Implementation Note: Starting from Stable Baselines3 v1.1.0, HER is no longer a standalone algorithm. Instead, it is implemented as the HerReplayBuffer class. To use HER, you must pass this buffer class to an off-policy algorithm when using a MultiInputPolicy (which provides support for Dict observation spaces).

    Environment Requirements: To use HER, your environment must follow the legacy gym_robotics.GoalEnv interface. Specifically, the gym.Env must have:

    1. A vectorized implementation of compute_reward().
    2. A dictionary observation space containing three keys: observation, achieved_goal, and desired_goal.
    from stable_baselines3 import HerReplayBuffer, SAC

# HER is used by passing HerReplayBuffer to the algorithm
model = SAC(
    "MultiInputPolicy",
    env,
    replay_buffer_class=HerReplayBuffer,
    replay_buffer_kwargs=dict(
        n_sampled_goal=4,
        goal_selection_strategy="future",
    ),
    verbose=1,
)

  11. Handle gSDE noise during PPO inference

    master

    When using PPO models trained with use_sde=True (Generalized State-Dependent Exploration), the automatic noise resetting that occurs during training does not happen during model.predict(). This can result in deterministic behavior even if deterministic=False is passed.

    Recommendations:

    1. For continuous control tasks, use deterministic=True during inference.
    2. If you require stochastic behavior, you must manually reset the noise by calling model.policy.reset_noise(env.num_envs) at intervals matching your training sde_sample_freq.

  12. How probability distributions are handled in policies

    master

    Policies in SB3 handle different probability distributions to support various action spaces. All distributions are located in common/distributions.py and follow a unified interface.

    Common distribution types include:

    • Categorical: Used for discrete action spaces.
    • DiagGaussian: Used for continuous action spaces.
    • SquashedGaussian: Used for continuous action spaces.
    • StateDependentDistribution: Used for state-dependent exploration.