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.05,target_kl 取 0.02。这个参数组合在 mujoco 的 Hopper 和 Walker2d 上都跑出了比固定权重更好的结果。主要收益出现在后期稳定性上,震荡尖峰消失了,最终性能比六月份的版本高了大约 12%。但这时候我意识到这只是一个工程补丁,它仍然依赖人工设定的 target_kl,而且比例控制器的反应速度跟不上策略变化的速度。在某些快速变化的环境中,比如 Ant 或者 Humanoid,KL 散度会在几千步内从 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.5)
- PPO + 固定 reward shaping + 固定 KL 蒸馏(六月份方案)
这三个基线加上九月份的自适应监督方案,在 Hopper、Walker2d、Ant、Humanoid 四个 mujoco 环境上做了完整的对比实验。每个实验跑 4 个随机种子,取均值和中位数分别做统计。
3.2 关键超参数记录与调参过程
自适应监督方案里最关键的三个超参数是 initial_coeff、min_coeff 和 max_coeff。我一开始把 initial_coeff 设为 1.0,min_coeff 设为 0.1,max_coeff 设为 2.0。但实验发现 max_coeff 设置为 2.0 有点危险,在学生策略初期性能非常差的时候,performance_ratio 接近于 0,adaptive_kl_coeff 会直接冲高到 max_coeff,导致 KL 惩罚项过强,学生策略被死死压在教师策略的分布附近,探索能力被严重抑制。
后来我把 max_coeff 降到 1.2,min_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