alterlab-stable-baselines3

Stable Baselines3

概述

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

何时使用此技能

当用户希望执行以下操作时，请使用此技能：

• 在 Gymnasium 环境中训练和评估单 Agent RL 策略（PPO、A2C、SAC、TD3、DDPG、DQN、HER）。
• 将研究模拟器封装为自定义 gymnasium.Env，并使用 check_env 对其进行验证。
• 通过回调添加评估、检查点、提前停止或自定义日志记录。
• 使用向量化环境加快训练，或以可复现的方式保存、加载和报告已训练的 Agent。

不适用的场景

场景	改用
每秒数百万步、原生多 Agent 环境，或 PufferEnv/Atari 级别的吞吐量	alterlab-pufferlib
监督式深度学习训练循环（LightningModule、多 GPU 策略）	alterlab-pytorch-lightning
在没有序贯决策过程的情况下进行参数黑盒优化或多目标优化	alterlab-pymoo
无需学习策略的队列或资源竞争模拟	alterlab-simpy

核心功能

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

关键限制：

• 图像观测必须为 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
from stable_baselines3.common.vec_env import SubprocVecEnv

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

使用 SubprocVecEnv 时，请将调用位置放在 if __name__ == "__main__": 中（macOS/Windows 上使用 spawn 时必须如此）。

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

API 差异：

• reset() 仅返回观测（信息可通过 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
)

视频录制（包装后的环境必须使用 render_mode="rgb_array" 创建，并且必须安装 moviepy，例如运行 uv pip install "gymnasium[other]"）：

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

# Built-in equivalent (SB3 >= 2.7; replaces the deprecated get_linear_fn/constant_fn/get_schedule_fn)
from stable_baselines3.common.utils import LinearSchedule
model = PPO("MlpPolicy", env, learning_rate=LinearSchedule(1e-3, 0.0, 1.0))

N 步回报（离策略，SB3 >= 2.7）： SAC、TD3、DDPG 和 DQN 接受 n_steps=，可根据 N 步回报进行自举（由 NStepReplayBuffer 支持）。

多输入策略（Dict 观测）：

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

安装

需要 Python 3.10+ 和 PyTorch >= 2.8（SB3 2.9 将最低版本从 2.3 提高到了 2.8）。[extra] 选项会安装 TensorBoard、OpenCV、ale-py（Atari）、pygame-ce、tqdm/rich（进度条），以及自 2.9 起的 pandas 和 matplotlib；它们已不再是核心依赖（results_plotter/load_results 需要这些依赖）。导入 ale_py 后会注册 "PongNoFrameskip-v4" 或 "ALE/Pong-v5" 等 Atari ID（import ale_py; gym.register_envs(ale_py)）。

# Basic installation
uv pip install stable-baselines3

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

属于 AlterLab Academic Skills 套件的一部分。