Microduck 训练防崩工程：NaN 怎么从奖励函数一路杀到 PPO 崩溃
做过强化学习训练的人应该都遇到过这种场景：训练跑了两个小时，loss 突然变成 NaN，策略彻底废掉，日志里找不到任何有意义的错误信息。在 microduck_rl 仓库的 mdp.py 文件里，Pollen Robotics 的工程师写了五个猴子补丁来解决这个问题。这不是过度工程，而是被真实的训练崩溃逼出来的。
NaN 的传播链：从奖励到崩溃
MuJoCo Warp 在 4096 个并行环境里跑 50Hz 仿真，任何一个环境出现物理状态异常，对应的奖励项就会返回 NaN。问题在于，NaN 不会自己消失，它会沿着 PPO 的计算图一路传播。
具体路径是这样的：奖励函数返回 NaN → 奖励累加到回合总回报 → TD 误差变成 NaN → 优势函数变成 NaN → 策略 loss 变成 NaN → 梯度变成 NaN → 优化器更新后策略参数全 NaN → 下一个 mini-batch 采样时 torch.normal 崩溃。
这条链路里最危险的一环是 torch.normal。PPO 在采样动作时需要从策略分布中抽取样本，如果标准差变成 NaN 或负数，PyTorch 直接抛异常，训练终止。而标准差变成 NaN 的原因，通常就是上一步梯度里的 NaN 流入了策略网络的对数标准差参数。
补丁一：奖励计算的 NaN 兜底
mdp.py 里的第一个补丁直接包装了 RewardManager.compute。原函数计算完奖励后，补丁调用 torch.nan_to_num 把所有 NaN 奖励替换成 0。
为什么要这么做？因为 mjlab 在环境重置之前计算奖励。这意味着如果一个环境的物理状态已经崩了（比如鸭子穿进了地面、关节角度飞到了 NaN），这一步的奖励一定是 NaN。如果不拦截，这个 NaN 就进入 PPO buffer 污染整个 batch。
补丁还做了一件事：把 _episode_sums 里的 NaN 也清零。_episode_sums 是 compute 内部更新的回合统计量，如果不清理，WandB 上会看到 NaN 的奖励曲线，调试时根本分不清是”真的崩了”还是”历史统计被污染了”。
补丁二：优势函数的 NaN 防御
第二个补丁针对 PPO 的 compute_returns。即便奖励层的 NaN 被堵住了，还有一种情况能产生 NaN 优势函数：课程学习的突变。
mdp.py 的注释里记录了一个具体案例。当某个奖励权重突然乘以 2.5 倍时，所有 TD 误差同时偏移一个量级。这时候优势函数的标准差会趋近于零，而 (A – mean) / (std + 1e-8) 会变成一个极大的数值。这个数值灌进梯度，优化器会把策略标准差推到负数——然后 NaN 就出现了。
补丁的做法很直接：在 compute_returns 之后，对 advantages 和 returns 都做 nan_to_num，把 NaN 和正负无穷都替换成 0。这不是优雅的解决方案，但它有效——在 4096 个环境的规模下，任何一个环境的数值异常都不应该拖垮整个训练。
补丁三：标准差钳制（已退役）
第三个补丁曾经存在但已被移除。它原本钳制策略分布的标准差下限，防止优化器把它推到负数或 NaN。在 mjlab 1.3.0 迁移到 rsl_rl 5.0.1 后，策略类被重构，原来的 ActorCritic 类不存在了，分布逻辑搬到了 rsl_rl.modules.distribution.GaussianDistribution。
注释里写得很坦诚：这是一个防御性的创可贴，mjlab 1.3.0 下没有它也能正常跑。但如果标准差爆炸的问题在新版本下复发，需要在新的分布类上重新实现。这种”补丁移除但教训保留”的做法值得注意——注释本身就是知识传递。
补丁四：ONNX 导出的被动关节过滤
第四个补丁解决了一个导出时的崩溃。新版 Microduck 模型有被动关节（下颌连杆通过等式约束闭合，没有执行器），但导出器的元数据函数遍历所有 16 个关节名，试图在 14 维的控制 ID 映射里查找——KeyError，导出失败。
补丁的逻辑是过滤掉 passive_ 前缀的关节，只导出 14 个舵机关节的元数据。这保证了 ONNX 策略的动作空间和运行时一致——不管模型里有多少被动关节，策略永远输出 14 维动作。
补丁五：热启动的课程重置
最后一个补丁解决了一个非常隐蔽的问题。mjlab 的 runner 在加载 checkpoint 时会恢复 common_step_counter，这对于断点续训是正确的。但如果你在做热启动——加载另一个任务的 checkpoint 作为新任务的起点——恢复的步数计数器会让所有基于步数的课程直接跳到最终阶段。
AGENTS.md 里记录了 2026-09-09 的一个案例：从步行的 checkpoint 热启动到新任务，恢复的步数是 90000，于是所有保护性奖励权重直接拉满，摔倒检测也关了——策略在第一次迭代就面对最终难度的环境，根本学不到东西。
补丁通过环境变量 MICRODUCK_WARM_START=1 控制：加载后把步数计数器和迭代计数器都重置为 0，但保留权重、归一化器和优化器状态。课程从头开始，策略基础保留。这个区分——resume 保持计数器、warm start 重置计数器——是所有课程学习项目都会遇到的陷阱。
粗糙地形的生成 Bug：一个真实案例
mdp.py 里还记录了一个导致训练崩溃的具体 bug。在粗糙地形上，斜坡金字塔的平台远高于 z=0。如果生成函数把鸭子的躯干 z 设成绝对值 0.05-0.09 米，鸭子就会被塞进地形内部。这导致接触力爆炸——一帧之内产生巨大的碰撞力——然后回报、价值函数、标准差全部炸掉。
2026-09-10 的一次训练（run id wy8gcaus）就是这么死的：稳定训练到 700 轮，趴姿生成开启，value loss 从 0.6 飙到 42，entropy 从 5 飙到 33。修复方法是加一个 _env_origin_z 辅助函数，在每个生成函数写绝对 z 值时加上地形高度偏移。平地训练永远不会遇到这个问题，因为地形高度就是 0。
奖励设计：预防胜于补丁
五个补丁是在数值层面兜底，但 Pollen 的工程师更强调在奖励设计层面预防。AGENTS.md 里列了几条用血泪换来的规则。
符号约定是第一条：mdp.py 有两种惩罚函数风格，一种返回正值配负权重，一种自反返回负值配正权重。如果搞混了，一个负权重配在自反惩罚上就变成双重否定——策略会去”追求”惩罚条件。比如屁股着地弹跳、故意摔倒坐着。检查方法很简单：WandB 上每个惩罚项的回合奖励必须 ≤ 0。
另一个规则是”不要给坏状态发正奖励”。如果摔倒后还能拿到站立奖励的大部分分数，策略就会选择趴着不动——因为趴着比站着省力。正确做法是用基于势函数的整形奖励：奖励的是进步量（Δcos(tilt)），站起来有奖励，保持站立没有额外奖励，趴着没有任何收益——无法被刷。
训练稳定性是工程问题不是算法问题
4096 个并行环境、50Hz 仿真、1-2 小时训出一个可用步态——在这个规模下，任何一个环境的数值异常都可能传播到整个 batch。Pollen 的做法是在两个层面防御：奖励计算和 PPO 内部做 NaN 兜底，奖励设计层面消除产生 NaN 的激励。两层加起来，训练稳定性才有保障。
这些补丁不是写完就不管的代码。第三个补丁在版本迁移时被移除，第五个补丁通过环境变量开关区分 resume 和 warm start。每个补丁的注释都记录了具体的失败案例、run id、发生日期。这种”补丁即文档”的做法，比单纯写一段防御代码更有价值——它让后来者不仅能看到怎么修的，还能看到为什么修、什么时候会再出问题。
—
来源：
– pollen-robotics/microduck_rl 仓库 src/mjlab_microduck/tasks/mdp.py：五个补丁的完整实现与注释，包含 NaN 传播链描述、课程突变案例、粗糙地形 bug 记录。
– pollen-robotics/microduck_rl 仓库 AGENTS.md：奖励设计规则、符号约定、势函数整形原则、课程对齐策略、warm start 与 resume 区分。
– BAM 项目：github.com/Rhoban/bam，Duclusaud et al., ICRA 2025。
– pollen-robotics/microduck_rl README：任务列表、训练配置、ONNX 导出流程。
