stable-baselines3

Stable Baselines3

概述

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

核心功能

1. 训练 RL Agent

基础训练模式：

import gymnasium as gym
from stable_baselines3 import PPO

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

# Initialize agent
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()：清理资源

主要限制：

• 图像观测必须是范围为 [0, 255] 的 np.uint8
• 尽可能使用通道优先格式（通道数、高度、宽度）
• 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]

资源

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：向量化环境的用法和包装器

安装

# Basic installation
uv pip install stable-baselines3

# With extra dependencies (Tensorboard, etc.)
uv pip install stable-baselines3[extra]