stable-baselines3

Stable Baselines3

概览

Stable Baselines3（SB3）是一个基于 PyTorch 的库，提供可靠的强化学习算法实现。本技能提供全面的指导，帮助你使用 SB3 的统一 API 训练 RL Agent、创建自定义环境、实现回调并优化训练工作流。

当前上游版本： SB3 2.8.0（2026 年 4 月）。文档：stable-baselines3.readthedocs.io。

安装

已针对 stable-baselines3 2.8.0 测试。需要 Python 3.10+（2.8.0 起不再支持 3.9）和 PyTorch >= 2.3。

# Basic installation
uv pip install "stable-baselines3>=2.8"

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

在 zsh 中，请为方括号加引号：uv pip install 'stable-baselines3[extra]>=2.8'。

用于 MuJoCo 连续控制基准测试：

uv pip install "gymnasium[mujoco]"

检查版本：

import stable_baselines3
print(stable_baselines3.__version__)

相关项目

• SB3-Contrib：实验性算法（MaskablePPO、CrossQ、QR-DQN、RecurrentPPO）——独立的 sb3-contrib 软件包
• RL Baselines3 Zoo：预训练 Agent、超参数和训练脚本
• SBX：SB3 + JAX 实现，适用于偏好使用 JAX 而非 PyTorch 的用户

核心功能

1. 训练 RL Agent

基本训练模式：

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)

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

重要说明：

• total_timesteps 是下限；由于批次采集，实际训练步数可能会超过此值
• 使用 model.load() 作为静态方法，而不是在现有实例上调用
• 为节省空间，模型不会保存回放缓冲区

算法选择： 使用 references/algorithms.md 了解详细的算法特性和选择指南。快速参考：

• PPO/A2C：通用，支持所有动作空间类型，适合多进程
• SAC/TD3：连续控制、离策略、样本效率高
• DQN：离散动作、离策略
• HER：目标条件任务

请参阅 scripts/train_rl_agent.py，获取包含最佳实践的完整训练模板。

2. 自定义环境

要求： 自定义环境必须继承 gymnasium.Env 并实现：

• __init__()：定义 action_space 和 observation_space
• reset(seed, options)：返回初始观测和 info 字典
• step(action)：返回观测、奖励、terminated、truncated 和 info
• render()：可视化（可选）
• close()：清理资源

主要限制：

• 图像观测必须为 np.uint8，取值范围为 [0, 255]
• 尽可能使用通道优先格式（通道、高度、宽度）
• SB3 会自动将图像除以 255 进行归一化
• 如果图像已预先归一化，请在 policy_kwargs 中设置 normalize_images=False
• SB3 不支持 start!=0 的 Discrete 或 MultiDiscrete 空间

验证：

from stable_baselines3.common.env_checker import check_env

check_env(env, warn=True)

请参阅 scripts/custom_env_template.py 获取完整的自定义环境模板，并参阅 references/custom_environments.md 获取全面的指导。

3. 向量化环境

用途： 向量化环境会并行运行多个环境实例，从而加快训练，并支持某些封装器（帧堆叠、归一化）。

类型：

• DummyVecEnv：在当前进程中顺序执行（适用于轻量级环境）
• SubprocVecEnv：跨进程并行执行（适用于计算密集型环境）

快速设置：

from stable_baselines3.common.env_util import make_vec_env

# Create 4 parallel environments
env = make_vec_env("CartPole-v1", n_envs=4, vec_env_cls=SubprocVecEnv)

model = PPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=25000)

离策略优化： 将多个环境与离策略算法（SAC、TD3、DQN）配合使用时，将 gradient_steps=-1，以便每个环境步执行一次梯度更新，从而平衡实际运行时间和样本效率。

API 差异：

• reset() 仅返回观测（info 可通过 vec_env.reset_infos 获取）
• step() 返回 4 元组：(obs, rewards, dones, infos)，而不是 5 元组
• 环境会在每个回合结束后自动重置
• 可通过 infos[env_idx]["terminal_observation"] 获取终止观测

请参阅 references/vectorized_envs.md，了解有关封装器和高级用法的详细信息。

4. 用于监控和控制的回调

用途： 回调可用于监控指标、保存检查点、实现提前停止，以及自定义训练逻辑，而无需修改核心算法。

常用回调：

• EvalCallback：定期评估并保存最佳模型
• CheckpointCallback：按间隔保存模型检查点
• StopTrainingOnRewardThreshold：达到目标奖励时停止训练
• ProgressBarCallback：显示训练进度和耗时

自定义回调结构：

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

可用属性：

• self.model：RL 算法实例
• self.num_timesteps：环境步数总计
• self.training_env：训练环境

串联回调：

from stable_baselines3.common.callbacks import CallbackList

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

有关回调的完整文档，请参阅 references/callbacks.md。

5. 模型持久化和检查

保存和加载：

# Save model
model.save("model_name")

# Save normalization statistics (if using VecNormalize)
vec_env.save("vec_normalize.pkl")

# Load model
model = PPO.load("model_name", env=env)

# Load normalization statistics
vec_env = VecNormalize.load("vec_normalize.pkl", vec_env)

参数访问：

# Get parameters
params = model.get_parameters()

# Set parameters
model.set_parameters(params)

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

6. 评估和录制

评估：

from stable_baselines3.common.evaluation import evaluate_policy

mean_reward, std_reward = evaluate_policy(
    model,
    env,
    n_eval_episodes=10,
    deterministic=True
)

视频录制：

from stable_baselines3.common.vec_env import VecVideoRecorder

# Wrap environment with video recorder
env = VecVideoRecorder(
    env,
    "videos/",
    record_video_trigger=lambda x: x % 2000 == 0,
    video_length=200
)

完整的评估和录制模板，请参阅 scripts/evaluate_agent.py。

7. 高级功能

学习率调度：

def linear_schedule(initial_value):
    def func(progress_remaining):
        # progress_remaining goes from 1 to 0
        return progress_remaining * initial_value
    return func

model = PPO("MlpPolicy", env, learning_rate=linear_schedule(0.001))

多输入策略（字典观测）：

model = PPO("MultiInputPolicy", env, verbose=1)

当观测值是字典时使用（例如，将图像与传感器数据结合）。

后见经验回放：

from stable_baselines3 import SAC, HerReplayBuffer

model = SAC(
    "MultiInputPolicy",
    env,
    replay_buffer_class=HerReplayBuffer,
    replay_buffer_kwargs=dict(
        n_sampled_goal=4,
        goal_selection_strategy="future",
    ),
)

TensorBoard 集成：

model = PPO("MlpPolicy", env, tensorboard_log="./tensorboard/")
model.learn(total_timesteps=10000)

工作流指南

启动新的 RL 项目：

1. 定义问题：确定观测空间、动作空间和奖励结构
2. 选择算法：使用 references/algorithms.md 获取选择建议
3. 创建或调整环境：如有需要，使用 scripts/custom_env_template.py
4. 验证环境：训练前务必运行 check_env()
5. 设置训练：使用 scripts/train_rl_agent.py 作为起始模板
6. 添加监控：实现用于评估和检查点保存的回调
7. 优化性能：考虑使用向量化环境以提升速度
8. 评估并迭代：使用 scripts/evaluate_agent.py 进行评估

常见问题：

• 内存错误：减小离策略算法的 buffer_size，或减少并行环境数量
• 训练缓慢：考虑使用 SubprocVecEnv 来运行并行环境
• 训练不稳定：尝试不同的算法、调整超参数，或检查奖励缩放
• 导入错误：确保已安装 stable_baselines3：uv pip install 'stable-baselines3[extra]>=2.8'

资源

scripts/

• train_rl_agent.py：包含最佳实践的完整训练脚本模板
• evaluate_agent.py：Agent 评估和视频录制模板
• custom_env_template.py：自定义 Gym 环境模板

references/

• algorithms.md：详细的算法比较和选择指南
• custom_environments.md：全面的自定义环境创建指南
• callbacks.md：完整的回调系统参考
• vectorized_envs.md：向量化环境的用法和包装器

Agent 操作流程

1. 检查环境。 确认 Python 环境和库版本（python -c "import pkg; print(pkg.__version__)"），并检查数据的形状、类型和缺失值。
2. 明确输入。 根据数据或用户提供的信息确认格式、标识符和参数。任何会改变结果的值都应询问，不要猜测。
3. 先运行小规模版本。 先在样本或单个折上运行，并检查运行时间和内存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 使用留出数据、固定随机种子和适当的指标；检查是否存在数据泄漏；报告不确定性（置信区间、不同随机种子下的标准差）。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果发生这种情况	请这样做
内存不足或执行速度很慢	对数据进行子采样，使用分块或惰性计算，或缩小模型规模，并告知用户做了哪些调整。
这些说明中的函数、标志或端点在已安装版本中不存在	查阅已安装版本自己的文档（help()、--help、官方文档），进行适配并告知用户。切勿编造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明所作假设。

诚信准则

• 切勿伪造结果、参数、标识符、引文或统计数据。如果某项内容无法运行或验证，请如实说明。
• 切勿报告本次会话中未计算过的指标；应展示生成每个数值的代码路径。
• 将此处与特定版本相关的细节视为可能已过时：请对照已安装版本的官方文档进行确认。
• 对于需要花费资金、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，须事先征求同意。