SuperAPI AI 工具集
合作ENAI 行业资讯收藏
Skill 广场stable-baselines3

stable-baselines3

智能体与模型

查看项目

更新于 10-15 20:21

技能介绍

用途

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

主要能力

训练强化学习智能体

基本训练模式：

import gymnasium as gym
from stable_baselines3 import PPO

# 创建环境
env = gym.make("CartPole-v1")

# 初始化智能体
model = PPO("MlpPolicy", env, verbose=1)

# 训练智能体
model.learn(total_timesteps=10000)

# 保存模型
model.save("ppo_cartpole")

# 加载模型（无需事先实例化）
model = PPO.load("ppo_cartpole", env=env)

重要说明：

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

算法选择快速参考：

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

自定义环境

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

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

关键约束：

  • 图像观测必须为 np.uint8，范围 [0, 255]
  • 尽可能使用通道优先格式（channels, height, width）
  • 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)

向量化环境

向量化环境并行运行多个环境实例，加速训练并启用某些包装器（帧堆叠、归一化）。

类型：

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

快速设置：

from stable_baselines3.common.env_util import make_vec_env

# 创建 4 个并行环境
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() 仅返回观测（信息可通过 vec_env.reset_infos 获取）
  • step() 返回 4 元组：(obs, rewards, dones, infos)，而非 5 元组
  • 环境在回合结束后自动重置
  • 终端观测可通过 infos[env_idx]["terminal_observation"] 获取

回调监控与控制

回调可在不修改核心算法的情况下监控指标、保存检查点、实现早停和自定义训练逻辑。

常用回调：

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

自定义回调结构：

from stable_baselines3.common.callbacks import BaseCallback

class CustomCallback(BaseCallback):
    def _on_training_start(self):
        # 在首次 rollout 前调用
        pass

    def _on_step(self):
        # 在每个环境步后调用
        # 返回 False 以停止训练
        return True

    def _on_rollout_end(self):
        # 在 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)

模型持久化与检查

保存与加载：

# 保存模型
model.save("model_name")

# 保存归一化统计（若使用 VecNormalize）
vec_env.save("vec_normalize.pkl")

# 加载模型
model = PPO.load("model_name", env=env)

# 加载归一化统计
vec_env = VecNormalize.load("vec_normalize.pkl", vec_env)

参数访问：

# 获取参数
params = model.get_parameters()

# 设置参数
model.set_parameters(params)

# 访问 PyTorch 状态字典
state_dict = model.policy.state_dict()

评估与录制

评估：

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

# 用视频录制器包装环境
env = VecVideoRecorder(
    env,
    "videos/",
    record_video_trigger=lambda x: x % 2000 == 0,
    video_length=200
)

高级功能

学习率调度：

def linear_schedule(initial_value):
    def func(progress_remaining):
        # progress_remaining 从 1 到 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. 选择算法：参考算法选择指南
  3. 创建/适配环境：如需自定义环境，使用自定义环境模板
  4. 验证环境：训练前始终运行 check_env()
  5. 设置训练：使用训练脚本模板作为起点
  6. 添加监控：实现评估和检查点回调
  7. 优化性能：考虑使用向量化环境加速
  8. 评估与迭代：使用评估脚本进行评估

安装：

# 基本安装
uv pip install stable-baselines3

# 包含额外依赖（Tensorboard 等）
uv pip install stable-baselines3[extra]

适用场景

  • 训练强化学习智能体（PPO、SAC、DQN、TD3、DDPG、A2C 等）
  • 创建自定义 Gym 环境
  • 实现回调以监控和控制训练
  • 使用向量化环境加速训练
  • 模型持久化、评估与视频录制
  • 目标条件任务（HER）
  • 多输入策略（字典观测）

限制

  • 图像观测必须为 np.uint8，范围 [0, 255]
  • SB3 不支持 start!=0 的 Discrete 或 MultiDiscrete 空间
  • 回放缓冲区不会随模型保存
  • total_timesteps 是下界，实际训练可能超过该值
  • 常见问题：内存错误（减少 buffer_size 或并行环境数）、训练缓慢（考虑 SubprocVecEnv）、训练不稳定（尝试不同算法、调参或检查奖励缩放）、导入错误（确保已安装 stable_baselines3）