跑通了几套内置环境之后,最兴奋也最容易受挫的一件事,就是把自己的任务写成一个能直接丢给 SB3 训练的自定义环境。我印象特别深,第一次干这事时,光一个reset的返回值就折腾了一晚上:新版 Gym 接口要求返回两个值,网上教程一半是老接口一半是新接口,对着抄都能抄出 bug。这篇文章是强化学习实战系列的第 5 篇,主题很聚焦:用 Stable-Baselines3(下面统一叫 SB3)训练自定义环境,并且把“输入状态向量”这条链路彻底讲透——从环境接口怎么写、观测空间怎么声明,到训练脚本怎么搭、模型训练完怎么验证。适合的读者是:已经把 CartPole 这类内置环境跑通、正准备把自己手头任务迁移到强化学习上的同学。
这篇文章不会讲太多晦涩的理论,重点是把能直接复用的代码、配置和经验教训给到你。我会用一个非常典型的二维导航任务做例子:智能体是一个点,目标点随机出现,状态向量是智能体坐标、目标坐标和当前距离。这种任务麻雀虽小五脏俱全,你把它理解透之后,换成你自己的业务状态输入只是改几个字段的事。
- 自定义环境真正要过的三关:接口、观测、奖励
1.1 Gym 接口:reset 和 step 并不是“照抄就行”
很多同学第一次写自定义环境,习惯性地去找所谓“标准模板”。但实际上 Gymnasium 的接口在近两年有过一次比较大的调整,网上搜到的老代码很可能让你白白踩坑。当前 SB3 主版本对接的是 Gymnasium 的新接口,要求环境实现五个方法:__init__、reset、step、render、close。其中close和render在训练时可以不实现成完整功能,但接口必须存在,否则 SB3 内部某些工具链会报错。
核心是这两个方法的签名:
def reset(self, seed=None, options=None): # 返回 (obs, info),注意是新接口! return obs, info def step(self, action): # 返回五元组 (obs, reward, terminated, truncated, info) return obs, reward, terminated, truncated, info老版本 Gym 的step返回的是(obs, reward, done, info),只有四个值;新版把done拆成了terminated和truncated两个布尔值。这个拆分不是没事找事,它直接影响了强化学习算法对“回合结束”这件事的理解。terminated表示任务真正达成或彻底失败,比如到达目标;truncated表示这次对话被外部条件截断,比如超出了最大步数。算法在更新价值函数时,terminated会被当作真正的终端状态处理,而truncated通常还需要继续做一步“自举”,二者混用会带来很隐蔽的学习错误,后面我会专门讲这个坑。
1.2 状态向量怎么声明:Box 空间里的 shape 和 dtype
标题里特别强调“输入状态向量”,对应的就是observation_space里的Box空间。所谓状态向量,就是一串数字拼起来的数组,比如[x, y, target_x, target_y, distance]这样一个长度为 5 的向量。
self.observation_space = spaces.Box( low=-10.0, high=10.0, shape=(5,), dtype=np.float32 )这里有几个细节我要多说一句。第一,low和high是边界,很多教程习惯写-np.inf和np.inf,能用但不推荐。如果你的状态有明确的范围,尽量写实际边界,后面接VecNormalize或做归一化时会更稳。第二,dtype一定要用np.float32,不要图省事省略。PyTorch 默认的浮点类型就是float32,SB3 在把观测转换给策略网络时也会做类型转换,但如果你在环境内部返回的是float64,某些版本会留下很难排查的隐藏警告,甚至导致数值表现异常。
动作空间如果是连续动作,也用Box来声明。我这里任务的动作是二维增量,所以动作空间是shape=(2,)的Box。如果你做的是离散动作任务,那就用Discrete,对应的算法策略也要配合调整。
1.3 奖励函数是唯一没有标准答案的地方
环境接口和空间声明都是“有标准答案”的,唯独奖励函数的设计没有标准答案,这也是自定义环境最灵魂的部分。我的习惯是第一版先用稠密奖励:每一步都根据“离目标近了还是远了”给出反馈。
dist = float(np.linalg.norm(self._agent - self._target)) reward = -dist * 0.1 terminated = dist < 0.1这里-dist * 0.1的含义是:离目标越远惩罚越大,每一步都诱导智能体往目标方向走。系数 0.1 是为了把单步奖励控制在零点几的量级,避免梯度波动太剧烈。奖励的尺度问题非常容易被忽略,后面我会专门讲它带来的训练灾难。稠密奖励不一定适合所有任务,如果你的任务里“中间状态好不好”很难定义,那就只能上稀疏奖励,但稀疏奖励的训练难度会明显上升,最好配合额外手段。
- 从零搭一个二维导航环境:代码逐段拆解
2.1 环境整体代码:先跑通再优化
先给出一个完整可跑的环境代码。这个环境对应的是我前面说的任务:智能体在二维平面里移动,目标随机刷新,状态向量包含智能体坐标、目标坐标和当前距离,动作是每一步的位移增量。
import numpy as np import gymnasium as gym from gymnasium import spaces class TargetReachEnv(gym.Env): def __init__(self): super().__init__() # 状态向量: [x, y, target_x, target_y, distance] self.observation_space = spaces.Box( low=-10.0, high=10.0, shape=(5,), dtype=np.float32 ) # 动作: 两个方向上的位移增量,范围 [-1, 1] self.action_space = spaces.Box( low=-1.0, high=1.0, shape=(2,), dtype=np.float32 ) self.max_steps = 100 self.step_count = 0 self._agent = None self._target = None def reset(self, seed=None, options=None): super().reset(seed=seed) self.step_count = 0 self._agent = np.zeros(2, dtype=np.float32) self._target = self.np_random.uniform( low=-5.0, high=5.0, size=2 ).astype(np.float32) obs = self._get_obs() return obs, {} def step(self, action): action = np.clip(action, -1.0, 1.0) self._agent += action * 0.2 self.step_count += 1 dist = float(np.linalg.norm(self._agent - self._target)) reward = -dist * 0.1 terminated = dist < 0.1 truncated = self.step_count >= self.max_steps obs = self._get_obs() info = {"dist": dist, "steps": self.step_count} return obs, reward, terminated, truncated, info def _get_obs(self): dist = float(np.linalg.norm(self._agent - self._target)) return np.array( [ self._agent[0], self._agent[1], self._target[0], self._target[1], dist, ], dtype=np.float32, ) def render(self): # 简单的控制台输出,训练时可以什么都不做 print( f"agent=({self._agent[0]:.2f}, {self._agent[1]:.2f}), " f"target=({self._target[0]:.2f}, {self._target[1]:.2f}), " f"dist={np.linalg.norm(self._agent - self._target):.2f}" ) def close(self): pass这个代码我故意写得朴素,一是方便你读懂,二是方便你做最小改造。你只需要改observation_space、_get_obs返回的向量内容和step里的状态转移逻辑,就能把它改成你自己的任务。注意我在_get_obs里每次都重新算了距离,这实际上让状态里包含了一个冗余特征。冗余特征不是坏事,距离能显著降低策略网络的学习难度,属于性价比很高的“特征工程”。
2.2 reset() 为什么必须返回两个值,以及 seed 的隐藏作用
我见过太多人在reset里忘记调super().reset(seed=seed),这行看起来多余,其实非常重要。Gymnasium 的Env基类内部维护了一个np_random随机数生成器,你只有调用了父类的reset(seed=seed),self.np_random才会被正确初始化。如果你自己用np.random生成随机目标,那么实验的可复现性就打了折扣——特别是做对比实验时,环境随机性不一致会直接影响结论。
另外,reset返回的info是一个字典,第一回合的额外信息也建议放在里面,比如初始距离。注意不要漏掉这个返回值,新接口下 SB3 的VecEnv包装器会严格检查返回值数量,漏掉就会直接抛异常或者产生诡异的“数据对不上”错误。
2.3 step() 五元组里最容易理解错的两个布尔值
新版接口里terminated和truncated是最容易写错的地方。以我的环境为例:智能体到达目标附近,属于“任务本身结束了”,该置terminated = True;而步数耗尽,属于“这次尝试被外部时间限制切断了”,该置truncated = True。两者可以同时为True吗?可以。比如智能体在第 100 步恰好到达目标,那terminated和truncated都是True,这是完全合法的。
很多老教程写的done = dist < 0.1 or step_count >= max_steps,然后返回(obs, reward, done, info),这在旧版没问题。但如果你用新版接口却把两种结束混在一起,SB3 在部分算法上会把这个标志当作terminated处理,导致回合超时被误判成“任务成功/失败”,价值函数的学习就会出现偏差。我的建议很简单:从写第一版代码起就严格区分这两个布尔值,养成习惯。
2.4 check_env 是第一个守门员
SB3 自带一个环境检查器,在你花一小时调试训练脚本之前,先让它花一秒钟告诉你环境有没有低级错误。
from stable_baselines3.common.env_checker import check_env env = TargetReachEnv() check_env(env)这个检查器会验证观测空间和动作空间的声明是否和实际返回值一致、dtype 是否正确、reset和step的返回值类型是否符合规范。如果检查器报了 warning,不要跳过,最好全部处理掉。我自己写环境的规定动作就是:写完第一版先跑check_env,然后再碰训练脚本。
- 训练脚本:从环境实例到 PPO 跑起来
3.1 为什么必须用 DummyVecEnv 包一层
SB3 的训练入口不接受裸环境,必须传入VecEnv。第一次接触时很容易问:我明明已经写好了环境,为什么还要再包一层?答案在于 SB3 的采样循环是建立在“向量化环境”概念上的。所谓向量化,就是把多个环境并排跑,一次采样同时得到 N 条轨迹,这样数据利用率高、训练也稳定。
最省事的是DummyVecEnv,它在单进程里用循环模拟多个环境并行。首次跑通时我强烈建议用它,调试方便,出了问题能直接看到 traceback。
from stable_baselines3.common.vec_env import DummyVecEnv env = TargetReachEnv() vec_env = DummyVecEnv([lambda: TargetReachEnv()])注意这里传入的是“创建环境的工厂函数”,也就是lambda表达式,而不是环境的实例。原因在于SubprocVecEnv这种多进程模式需要把工厂函数发给子进程去创建环境,如果你传的是实例,子进程就没法复制了。
3.2 首版 PPO 参数:先用默认,再谈调参
SB3 内置的 PPO 默认参数本身就是经过很多基准任务验证过的,第一版训练时我建议只改必要的几个,其余保持默认。下面是我在这个任务上的首版配置:
from stable_baselines3 import PPO model = PPO( "MlpPolicy", vec_env, verbose=1, learning_rate=3e-4, n_steps=2048, batch_size=64, n_epochs=10, gamma=0.99, gae_lambda=0.95, clip_range=0.2, )简单解释一下我为什么只动这几个。MlpPolicy是“多层感知机策略”,适用于状态向量输入的任务,跟标题里的“输入状态向量”正好对应。learning_rate=3e-4是无数实验验证过的稳健值,不追求极限性能就不要动。n_steps=2048是每次收集的步数,代表策略更新前要攒多少数据。n_epochs=10是每批数据被重复学习的轮数,PPO 的核心约束决定了这个值不能太大,否则策略更新容易过激。gamma和gae_lambda是奖励折扣和泛化估计参数,默认值对这个量级的任务足够。
如果你用的是SubprocVecEnv,可以同时开 4 或 8 个环境,n_steps就按“总步数 = 单环境步数 × 环境数”来理解,实际上n_steps已经是所有环境加起来的合计步数了,不用你手动乘。
3.3 训练日志里四个关键指标怎么看
verbose=1会让训练过程打印实时的滚动日志,很多同学盯着一堆数字不知道看什么。我平时重点看四个:
rollout/ep_rew_mean:最近一轮采样的平均回合奖励,是智能体学没学到的第一直观指标。rollout/ep_len_mean:平均回合长度,如果它一直顶着max_steps不降,说明智能体还没找到缩短路径的办法。time/fps:每秒环境交互帧数,主要用于判断训练速度是否正常。如果 fps 突然掉到个位数,检查是不是环境里不小心加了昂贵的渲染或者打印。train/approx_kl:策略更新的 KL 散度,数值长期过大说明更新太激进,需要降低学习率或n_epochs。
训练调用本身很简单:
model.learn(total_timesteps=200_000) model.save("target_reach_ppo_v1")200000 步在这个玩具任务上大概几分钟能跑完。从ep_rew_mean的走势看,如果从负几十稳步爬升到接近 0,说明智能体确实在学习靠近目标。
- 训练中常见的静默失败与排错实录
这里我不打算列“问题大全”,只想把我实际踩过、且具有很强隐蔽性的四个问题拆开讲。它们有个共同特点:训练不报错,甚至日志还挺好看,但最后模型就是不对。
4.1 观测值 shape 正确但 dtype 不对
check_env会检查观测的 dtype,但如果你用了自定义的info逻辑或者某些写法绕过了检查,float64的观测就可能一路流进网络。PyTorch 层面通常会自动转成float32,但“能跑”不代表“没问题”。状态向量的数值如果因为 dtype 转换出现精度波动,在训练早期可能不会立刻暴露,但会让你的实验结果变得不可复现。我自己的教训是:在_get_obs的return里强制astype(np.float32),这样无论你前面怎么折腾numpy,最终返回值一定是合规的。这是成本最低、收益最稳定的防呆设计。
4.2 terminated/truncated 混用导致“白训练”
这是一个特别阴间的错误:如果你把超时和成功混在一起都设置成terminated=True,PPO 会认为“每回合都是真正终结的”。这会导致算法在回合超时后把最后的观测状态当成一个真实的终点状态,去学习它的价值。但实际上这个状态根本不是终点,智能体现在的位置还有救。这个错误的后果不会立刻体现在 loss 上,而是让价值函数产生系统性偏差,最终表现就是训练很久,ep_rew_mean反复横跳但不稳定上升。
我的修复方式是用两个独立的布尔变量,然后在info里额外记录结束原因,方便事后统计:
terminated = dist < 0.1 truncated = self.step_count >= self.max_steps info = {"terminated_reason": "success" if terminated else "timeout"}4.3 每回合步数过短:模型根本来不及学
自定义环境最容易被忽略的参数就是max_steps。如果设成 10,智能体从初始位置到目标可能要 20 步,那它每次都在拿到有效反馈前就被掐断了,训练信号基本是无效噪声。PPO 是在收集一整批轨迹后统一做优势估计的,回合样本长度太短会让 GAE 的优势估计失去意义。
我的经验是:第一版把max_steps设得宽裕一点,宁可牺牲一点训练效率,也要保证智能体有足够时间完成任务。比如这个任务,智能体每步最多移动 0.2,初始位置到随机目标的平均距离大约 2 到 3,那 100 步的上限基本够用。如果设计一个新任务,你可以先手动运行环境,统计一下专家策略或随机策略的平均回合长度,然后把max_steps设为其 1.5 到 2 倍。
4.4 reward 绝对值过大:梯度直接爆炸
奖励尺度问题我用一张表来说明,因为它的表现形态太多了:
| 奖励设置 | 现象 | 根本原因 |
|---|---|---|
reward = -dist * 10 | loss 巨大,策略崩坏 | 梯度范数过大,优势估计被少数极端样本主导 |
reward = -dist * 0.001 | 学得很慢,甚至不动 | 奖励信号被 value 网络当作近似零,学习信号太弱 |
reward = -dist * 0.1 | 稳定收敛 | 单步奖励量级在 0.1~1,属于常见算法的舒适区 |
| 到达目标给 +1000,其余 0 | 前期完全学不动 | 稀疏奖励,探索阶段几乎拿不到正反馈 |
这里不是说你不能给大数值奖励,而是如果给了大数值,配套的归一化措施也要跟上。SB3 提供VecNormalize,能对观测和奖励做运行均值归一化,但用它的代价是会改变奖励的原始尺度,自己心里要有数。第一版我还是建议把奖励设计得温和一些,先让整个训练链路跑通,再去玩高阶技巧。
- 验证和复用:如何确认智能体真的学会了
5.1 用 evaluate_policy 和手动 demo 双重验证
训练结束后,不要只盯着日志里的曲线说“看起来不错”。SB3 提供了evaluate_policy函数,可以快速给出一个统计指标:
from stable_baselines3.common.evaluation import evaluate_policy mean_reward, std_reward = evaluate_policy( model, vec_env, n_eval_episodes=20, deterministic=True ) print(f"mean_reward={mean_reward:.2f} +/- {std_reward:.2f}")如果mean_reward显著高于随机策略的水平,说明模型学到了东西。但这还不够,我还会写一个手动 demo,用model.predict一步步跑环境,观测每一步的动作和状态变化。这能帮你确认智能体是不是真的在“走向目标”,而不是找到了某个奖励漏洞。历史上我就见过模型学会原地转圈刷奖励的鬼畜行为,只有手动看轨迹才能发现。
obs, _ = env.reset() total_reward = 0.0 for _ in range(200): action, _ = model.predict(obs, deterministic=True) obs, reward, terminated, truncated, info = env.step(action) total_reward += reward if terminated or truncated: break print(f"demo total_reward: {total_reward:.2f}")5.2 保存、加载与继续训练的正确姿势
model.save()保存的只是模型权重和超参数,不包括环境本身。加载时你需要自己重新构建环境,再把它设置给模型。这里有一个非常常见的问题:加载后直接调learn()会报错,因为模型内部没有关联环境。
model = PPO.load("target_reach_ppo_v1") vec_env = DummyVecEnv([lambda: TargetReachEnv()]) model.set_env(vec_env) model.learn(total_timesteps=100_000) # 继续训练如果你用了VecNormalize,光保存模型是不够的,还要单独保存归一化统计量:
vec_env.save("vec_normalize_stats.pkl")训练环境与推理环境的差异是另一个问题。如果你的自定义环境在训练时用了随机目标,推理时想要固定目标,我的建议是在环境里增加一个set_target之类的方法,而不是靠修改全局变量,否则多进程环境下会失效。
5.3 首版训练后的调参思路总结
如果你的第一版结果不理想,按照我踩坑的顺序,先排查环境本身:
- 用
check_env复查环境。 - 手动 demo 观察智能体行为,确认不是奖励漏洞。
- 检查
max_steps是否给得太紧。 - 检查奖励尺度是否落在 0.1~1 的舒适区。
这些都排除了,再考虑调 PPO 参数。我试出来的优先级一般是:total_timesteps不够 → 奖励尺度和状态特征有问题 → 学习率需要衰减或调小 →n_steps不够大导致样本相关性太强。不要一上来就堆参数,那只会让你分不清是环境错还是算法错。
5.4 加载后测试新随机场景的一个小技巧
最后分享一个我个人的习惯性操作。每次改完环境或调完参数,我会在固定几个 seed 上重新训练一小段时间,而不是每次都随机初始化。这样能快速对比改动前后的效果差异。做法也很简单:
import gymnasium as gym from stable_baselines3.common.monitor import Monitor env = Monitor(TargetReachEnv(seed=2024), filename="logs/run1")Monitor包装器会把每个回合的奖励、长度、时间都记录到 CSV 文件,方便你训练完之后打开logs/run1.csv做逐回合分析。用固定 seed 跑,再配合表格里的曲线数据,判断“这次到底改好了还是改差了”就变得非常客观。
实际训练中我还发现一个小习惯很管用:每次调参只改一个变量。以前我总喜欢同时调学习率、批量大小和奖励系数,结果模型崩了都不知道是哪一刀切坏的。现在我把每个实验的配置和 CSV 日志放在独立目录,记录改动点,虽然看着繁琐,但长期下来省的时间不是一点半点。希望这篇实战记录能帮你在自定义环境这条路上少走几个来回。