用途:Stable Baselines3(SB3)是基于 PyTorch 的强化学习算法库,该技能面向训练 RL 智能体、创建自定义环境、实现回调以及优化训练流程等任务,统一 API 覆盖 PPO、SAC、DQN、TD3、DDPG、A2C 等算法。
主要能力:
- 训练智能体。基本模式为用 gym.make 创建环境,初始化如 PPO("MlpPolicy", env, verbose=1),调用 model.learn(total_timesteps=10000),再 model.save 保存。total_timesteps 是下界,实际可能因批次采集超出;加载用 PPO.load 静态方法,replay buffer 不随模型保存。
- 自定义环境。需继承 gymnasium.Env 并实现 init、reset、step、render(可选)、close;图像观测须为 [0,255] 的 np.uint8,优先通道在前,SB3 自动除以 255 归一化;不支持 start!=0 的 Discrete/MultiDiscrete;训练前应用 check_env 校验。
- 向量化环境。DummyVecEnv 单进程顺序执行,SubprocVecEnv 多进程并行;可用 make_vec_env("CartPole-v1", n_envs=4, vec_env_cls=SubprocVecEnv)。离策略算法多环境时设 gradient_steps=-1。向量化后 reset 只返回观测,step 返回 4 元组,回合结束自动重置。
- 回调。常用 EvalCallback、CheckpointCallback、StopTrainingOnRewardThreshold、ProgressBarCallback;自定义回调继承 BaseCallback 并实现 _on_training_start、_on_step、_on_rollout_end,_on_step 返回 False 可停止训练;可用 CallbackList 串联。
- 模型持久化与检查。model.save/model.load 保存加载模型,VecNormalize.save 与 load 保存归一化统计;get_parameters/set_parameters 及 policy.state_dict() 访问参数。
- 评估与录制。evaluate_policy 返回平均与标准差奖励;VecVideoRecorder 可按触发条件录制视频。
- 高级特性。学习率调度函数进度从 1 到 0;字典观测用 MultiInputPolicy;HER 用 replay_buffer_class=HerReplayBuffer;支持 tensorboard_log。
使用方式:基础安装 pip install stable-baselines3,含 Tensorboard 等附加依赖用 pip install stable-baselines3[extra]。
适用场景:连续控制、离散动作、目标条件任务及需要并行加速或自定义环境的研究与工程训练。
限制:不保存 replay buffer;不支持 start!=0 的离散空间;图像观测需符合指定格式;内存不足时可减小 buffer_size 或减少并行环境数;训练慢可用 SubprocVecEnv;训练不稳定可换算法、调超参或检查奖励缩放。