1. 项目整体定位为什么我一直坚持 On-Policy Distillation 这条线先交代一下背景。过去三个月我一直在推进一个和策略蒸馏强相关的研究课题早期版本的核心思路是 reward shaping 引导下的 on-policy 学习后来逐步演进到自适应监督信号的框架。这个项目不是从零开始拍脑袋设计的而是基于前几轮实验踩坑后的自然迭代整体节奏从 6 月延续到 9 月目前已经完成第四轮关键实验。标题里带个“续”字是因为之前的阶段性总结已经整理过前五个月的进展这次主要聚焦最近三个月的技术演变、问题暴露和方案修正。这个项目解决的核心问题其实很朴素当我们在做 on-policy distillation 时教师策略和学生策略之间的分布差异会被时间放大尤其是学生策略快速改善后教师给出的监督信号可能已经从“引导”变成了“约束”。如果一直用固定权重去平衡 reward 和 distillation loss训练后期很容易出现策略退化或者探索崩塌。所以整个 6 到 9 月的核心工作就是在尝试把 reward shaping 的静态设计改造成一种能跟随策略状态动态调整的自适应监督机制。适合读这篇文章的人主要是正在做强化学习策略蒸馏、模仿学习、或者是 reward shaping 相关实验的研究者和工程同学。如果你只是单纯想了解概念也能从里面拿到不少可以直接复用的实验设计思路。我自己跑实验的时候最烦那种只讲理论不给参数的文章所以这篇会把每个阶段的实验配置、关键曲线变化、还有我踩过的坑都摊开来讲。2. 从 reward shaping 到自适应监督整体思路演进的三个阶段2.1 六月份的核心矛盾固定 reward shaping 权重在蒸馏中后期开始失效六月份的时候我的方案还比较简单粗暴。教师策略用 PPO 训好之后冻结学生策略通过 on-policy 交互收集数据同时用教师策略的 action logits 做 KL 蒸馏。为了保证学生初期不跑偏我加了一个 reward shaping 项具体形式是给学生每个 step 的 reward 额外加上一项和教师策略动作分布相似度成正比的 bonus。这样做初期效果非常明显学生策略在 1e6 步以内就能追上教师策略 80% 的水平比纯靠环境 reward 快了一倍多。但问题出在训练中后期。当学生策略的分布逐渐接近教师之后这个 shaping bonus 反而成了阻碍。因为教师策略也不是完美的它本身存在一些探索性的随机动作学生一旦和学生分布对齐得过于紧密就会把教师的次优动作也学进去。更麻烦的是reward shaping 的权重是固定的后期 KL 散度已经很小了但 bonus 项还在持续抬高学生策略对教师分布的拟合压力导致学生策略在局部最优附近来回震荡loss 曲线出现明显的周期性尖峰。我一开始怀疑是 KL 系数太大调小了之后训练确实平稳了一些但收敛速度又掉回去了。后来我意识到问题不只是一个系数的大小而是整个监督信号的结构设计缺少对训练阶段的感知。固定 shaping 的逻辑天然假设教师策略在所有阶段都是同等可靠的但这个假设在后期明显不成立。2.2 七八月份的关键转机把监督信号从“静态加权”变成“动态调节”七月初我开始尝试一个相对温和的改进把奖励 shaping 的 bonus 系数做成 KL 散度的函数。具体来说当学生和教师分布差异大的时候bonus 权重自动调高加速初期对齐当两者分布接近时权重自动衰减避免后期过度拟合。这个改动听起来不复杂但实现的时候需要考虑的细节很多比如如何平滑地计算 KL 变化率、如何防止权重在训练中后期频繁震荡、以及在何种频率下更新这个系数。我做了一个比较简单的自适应调节器每 5000 步计算一次学生策略和教师策略在当前状态分布下的平均 KL 散度然后用一个比例控制器调节 shaping bonus 的权重# 自适应调节器核心逻辑简化版 if step % 5000 0: current_kl compute_mean_kl(student_policy, teacher_policy, replay_buffer) # 目标 KL 设定为 0.02当前 KL 高于目标则增加 shaping 权重 error current_kl - target_kl shaping_weight min(max(shaping_weight alpha * error, 0.0), 1.0)alpha 我取的是 0.05target_kl 取 0.02。这个参数组合在 mujoco 的 Hopper 和 Walker2d 上都跑出了比固定权重更好的结果。主要收益出现在后期稳定性上震荡尖峰消失了最终性能比六月份的版本高了大约 12%。但这时候我意识到这只是一个工程补丁它仍然依赖人工设定的 target_kl而且比例控制器的反应速度跟不上策略变化的速度。在某些快速变化的环境中比如 Ant 或者 HumanoidKL 散度会在几千步内从 0.1 掉到 0.01控制器还没来得及衰减 shaping 权重学生策略就已经在向教师拟合的过程中失去了对环境的适应性。2.3 九月份的形态自适应监督信号彻底替代固定 reward shaping八月底到九月初我把思路再往前推了一步。之前的自适应调节器本质上还是在调一个外部奖励项的权重但整个目标函数的结构没有变化。九月份的版本做了更彻底的改动不再单独设置 shaping bonus而是直接把蒸馏损失本身变成一个自适应项。具体的做法是将 on-policy distillation loss 的系数变成学生策略当前性能的函数用滑动窗口内的环境 reward 均值作为性能指标。当学生策略性能低于教师策略性能的某个阈值时蒸馏系数保持较高当学生策略开始超越教师时蒸馏系数自动降低让学生有更多自由度去探索自己的策略空间。这个设计的核心逻辑是监督信号的价值是随着学生能力提升而递减的自适应监督的本质是让这个信号在学生最需要它的时候出现在它开始产生副作用的时候退场。适配后的目标函数如下# 九月份版本核心目标函数 total_loss policy_loss - entropy_bonus adaptive_kl_coeff * kl_divergence(student, teacher) # adaptive_kl_coeff 的计算方式 performance_ratio student_mean_reward / teacher_mean_reward adaptive_kl_coeff max(min(initial_coeff * (1.0 - performance_ratio), max_coeff), min_coeff)这个版本在多个环境中都跑出了让人满意的结果尤其是 Humanoid 这种高维连续控制任务训练稳定性和最终性能都有显著提升。这时候整个方案已经基本从“reward shaping 引导下的蒸馏”转变成了“自适应监督机制下的策略学习”reward shaping 这个词依然存在但它在系统中的角色已经变成了自适应监督的一个特例。3. 核心实验设置与实现细节配置、环境、超参与代码级描述3.1 实验环境和对比基线说明整个实验基于 PyTorch 实现PPO 作为基础 RL 算法策略网络是两层 MLP每层 256 个单元激活函数用 tanh。蒸馏教师策略是一个已经收敛的 PPO 模型固定参数后不做更新。学生策略从随机初始化开始训练训练总步数在 5e6 到 1e7 之间根据环境复杂度调整。对比基线设了三个纯 PPO无蒸馏作为性能下界参考PPO 固定 KL 蒸馏KL 系数固定为 0.5PPO 固定 reward shaping 固定 KL 蒸馏六月份方案这三个基线加上九月份的自适应监督方案在 Hopper、Walker2d、Ant、Humanoid 四个 mujoco 环境上做了完整的对比实验。每个实验跑 4 个随机种子取均值和中位数分别做统计。3.2 关键超参数记录与调参过程自适应监督方案里最关键的三个超参数是 initial_coeff、min_coeff 和 max_coeff。我一开始把 initial_coeff 设为 1.0min_coeff 设为 0.1max_coeff 设为 2.0。但实验发现 max_coeff 设置为 2.0 有点危险在学生策略初期性能非常差的时候performance_ratio 接近于 0adaptive_kl_coeff 会直接冲高到 max_coeff导致 KL 惩罚项过强学生策略被死死压在教师策略的分布附近探索能力被严重抑制。后来我把 max_coeff 降到 1.2min_coeff 降到 0.05情况好了很多。这个经验也说明自适应监督虽然能解决后期问题但前期如果约束太强同样会伤害探索。初始阶段学生策略需要足够的自由度去尝试不同的动作而不是被教师策略完全框住。所以在实际使用中max_coeff 不要太激进初期让学生策略在蒸馏信号和环境 reward 之间有一个相对平衡的取舍反而是更好的选择。另外滑动窗口的 size 我设定为 20000 个时间步也就是每收集 20000 步的数据计算一次均值。窗口太短会导致 performance_ratio 噪声很大窗口太长又会降低自适应的响应速度20000 步是一个比较合适的折中点。3.3 蒸馏损失的具体实现与数值稳定性处理on-policy distillation 的损失函数用的是 KL 散度但这里有一个容易踩坑的细节KL 散度本身对分布中的零概率非常敏感。如果教师策略的某个动作概率极小而学生策略在这个动作上分配了非零概率KL 散度就会产生一个很大的值导致梯度爆炸。我采取的防护措施是给两个策略的动作 logits 都加上一个小的数值稳定项比如 1e-8然后在计算 log_softmax 之前对 logits 做一次标准化缩放。具体实现如下def kl_divergence_loss(student_logits, teacher_logits): student_log_probs F.log_softmax(student_logits 1e-8, dim-1) teacher_probs F.softmax(teacher_logits 1e-8, dim-1) kl (teacher_probs * (teacher_probs.log() - student_log_probs)).sum(dim-1) return kl.mean()另外一个细节是学生策略和教师策略的 action space 维度和范围必须完全一致特别是连续控制任务中两个策略输出的动作分布需要映射到同一个 scale 上。如果教师策略用了 tanh 输出层学生策略也一定要用相同的输出层结构否则 KL 散度计算出来的数值会非常奇怪。提示在连续动作空间中直接用高斯分布参数的 KL 散度公式会涉及协方差矩阵的 log 行列式计算此时要特别注意数值稳定性。我用的方法是先算出两个分布的均值和 log_std再借助稳定版公式计算避免直接用 torch.distributions.kl_divergence因为那个接口在某些情况下对极端输入的容错性不够好。4. 实验数据结果与关键曲线解读自适应监督到底带来了什么提升4.1 Hopper 和 Walker2d 上的性能对比Hopper 上的实验最直观。固定 KL 蒸馏方案在 3e6 步时能拿到大约 3200 的平均 reward但曲线波动很大尤其是 2.5e6 到 3.5e6 步之间出现了一个明显的性能塌陷reward 一度跌回 2500 左右。六月份的固定 shaping 方案把塌陷出现的时间延后了但并没有消除最终收敛值也只有 3500 左右。九月份的自适应监督方案则平稳得多全程没有明显的性能塌陷最终收敛值稳定在 4100 附近。Walker2d 的结果类似但差异没有 Hopper 那么极端。固定方案收敛值大约在 4800自适应方案在 5600 左右主要优势集中在训练后期也就是 5e6 步之后自适应方案仍然能保持缓慢上升的趋势而固定方案已经明显进入平台期。我特意查看了训练后期的 KL 散度曲线发现固定方案的 KL 散度在后期并不会降到零而是维持在一个较高的水平上震荡这说明学生策略一直在试图拟合教师策略中不稳定的动作输出从而干扰了环境 reward 的优化。自适应方案的 KL 散度则稳步下降最终稳定在一个较低水平。4.2 Ant 和 Humanoid 上的现象与新问题Ant 环境里出现了一个有意思的现象。自适应方案前期1e6 步以内比固定方案慢一些大约落后 15% 的 reward但 2e6 步之后开始反超最终成绩领先固定方案约 20%。这个现象其实很好理解自适应方案初期为了探索给了学生策略更多自由度丢弃了一部分教师策略的成熟经验所以前期慢一些但正因为没有死板地复刻教师学生策略后期探索出了教师策略没有发现的一些高效步态所以最终性能反而更高。Humanoid 是唯一一个让我折腾了比较久的环境。因为它的动作维度是 17 维状态维度也很高KL 散度计算时协方差矩阵的数值特性极其复杂。我在 Humanoid 上跑固定方案的时候发现训练经常在 4e6 步左右崩掉loss 变成 NaN原因是学生策略的 log_std 在训练中过度下降方差趋近于零导致 KL 散度计算中出现分母为零的情况。自适应方案里我把这个风险处理掉了具体做法是在策略网络中限制 log_std 的下限不允许它小于 -5同时也给 KL 散度公式加上物理上合理的 epsilon。_log_std 下限限制其实非常重要不只是为了数值稳定它还关乎探索能力。如果一个策略的方差降得过低它的分布只集中在均值附近一个极小的区域内后续几乎无法探索新动作。这在蒸馏场景中尤其危险因为学生策略方差一旦过小就会对教师策略的错误动作细节过度拟合丧失自己的泛化空间。4.3 关键曲线形态变化与可视化思路我把四种方案在四个环境上的训练曲线整理成对比图x 轴是步数y 轴是滑动平均后的 reward。这里有一个值得强调的细节只画平均曲线是不够的一定要把不同种子的 std 带画出来。因为我发现固定方案在某些种子上表现很好但换一个种子就崩了如果只看平均曲线会严重高估方案的稳定性。自适应方案除了 Humanoid 环境外在四个种子上的表现都比较集中std 带很窄这也是它最让我满意的一点。如果你也想复现这个对比我建议至少跑 4 个随机种子并且在绘图时用 interquartile range 而不是标准差。因为强化学习训练曲线经常出现长尾分布个别种子跑飞会把标准差撑得很大从而掩盖了大多数种子的真实表现。中位数曲线配合 IQR 带是更加稳健的展示方式。5. 训练中的常见问题与排查实录数值爆炸、策略停滞和过拟合5.1 KL 散度数值爆炸的排查与防护前面提到 KL 散度容易数值爆炸这里展开讲一下我踩过的具体坑。有次在 Ant 环境上跑训练到 3.8e6 步时loss 突然从 0.2 跳到 1e4再往后几步直接 NaN。我用梯度裁剪也挽救不回来。后来排查发现是某种极端情况下教师策略输出动作概率分布出现了一个概率值极其接近 1 的动作学生策略在该动作的 log_prob 则因为数值误差变成了一个非常大的负值导致 KL 散度里出现了正负无穷相乘的隐患。我的解决方式分为三层第一层是所有概率计算前加上 1e-8 的 epsilon第二层是将 loss 计算限制在 batch 内的有效样本上过滤掉那些 KL 散度超过 10 的离群样本第三层是给 KL 散度项挂一个梯度裁剪单独设置 max_grad_norm 为 0.5避免单条样本对整个网络造成过大冲击。这三层措施叠加以后从 7 月到现在没有再次出现数值爆炸的情况。5.2 学生策略停滞在局部最优的表现与识别另一个常见问题是训练中后期学生策略陷入局部最优看起来性能已经收敛但和教师策略的差距其实仍然很显著。这时候如果你只看 reward 曲线很难发现问题因为曲线已经平了。我的识别技巧是同时监控学生策略和教师策略在相同状态下的动作分布重叠度也就是计算一些小批量状态上两个策略的 KL 散度。如果这个值一直保持在一个高位说明学生策略并没有真正学到教师策略的核心行为只是用另一种方式拿到了相似的 reward。出现这种情况时光靠调蒸馏系数是没用的我通常会加一层辅助监督也就是让学生策略不定期地采样教师策略曾经访问过的状态然后在这些状态上重新计算蒸馏损失。这个做法有点像经验回放但回放的不是 transition而是状态集。运行了几轮之后学生策略对教师策略的分布覆盖度明显提高局部最优被有效打破。5.3 自适应系数震荡问题的处理方法自适应系数本身也有可能出现震荡主要表现为 adaptive_kl_coeff 在训练过程中频繁上下跳动导致学生策略的训练目标一直不稳定。我用的是滑动平均平滑策略每次更新时不直接使用当前计算的 performance_ratio而是使用它的指数移动平均版本系数移动平均的 decay 取 0.99。这样处理后adaptive_kl_coeff 的曲线变得平滑很多而且还没有明显牺牲自适应的响应速度。值得提醒的是如果你发现 adaptive_kl_coeff 在某一段训练中持续震荡大概率不是超参数的问题而是学生策略本身的表现波动过大也就是它的 reward 曲线本身就很不稳定。这时候要先解决策略的稳定性再谈自适应系数的平滑顺序反了就事倍功半。6. 自适应监督方案的工程化要点如何把它嵌入现有训练框架6.1 模块化设计与调试接口整个自适应监督逻辑可以封装成一个独立的调度器和 PPO 的 rollout 收集、策略更新解耦。我这边实现的类接口大概长这样class AdaptiveSupervisionScheduler: def __init__(self, initial_coeff, min_coeff, max_coeff, alpha): self.initial_coeff initial_coeff self.min_coeff min_coeff self.max_coeff max_coeff self.alpha alpha self.ema_performance_ratio None def update(self, student_reward, teacher_reward): performance_ratio student_reward / (teacher_reward 1e-8) performance_ratio np.clip(performance_ratio, 0.0, 1.5) if self.ema_performance_ratio is None: self.ema_performance_ratio performance_ratio else: self.ema_performance_ratio 0.99 * self.ema_performance_ratio 0.01 * performance_ratio