☰
自适应监督机制:解决on-policy蒸馏中reward shaping的动态调参难题
2026/10/9 4:21:48 网站建设 项目流程

从 6 月到 9 月,我大部分精力都花在 on-policy distillation 这条线上。原本只是想解决 reward shaping 系数难调的问题,结果一路推到了自适应监督这个方向。这篇算是续篇,把 6 到 9 月的进展、踩过的坑、以及一些还没完全解决的开放问题,做一个比较系统的梳理。

1. 为什么最后还是绕不开 on-policy 蒸馏

先说一个我自己反复碰壁之后才想明白的问题:做策略蒸馏的时候,为什么 on-policy 比 off-policy 难这么多,甚至很多人一开始会刻意绕开它。

1.1 Off-policy 蒸馏看起来很美好,但埋了数据分布的雷

我最早做的版本是标准的 off-policy 蒸馏。teacher 网络先跑一批数据,把状态、动作、价值估计存下来,然后 student 网络在这些静态数据上去拟合 teacher 的输出。前几个 epoch 效果很好,loss 蹭蹭往下掉,student 的模仿误差也在收敛。但等 student 真正部署到环境里,问题就暴露了:它看到的 state 分布和训练时完全不同。

这其实就是经典的 covariate shift 问题。teacher 在训练时见过的状态空间,和 student 自己探索到的状态空间,两者之间的交集可能很小。不是 student 学得不好,是它根本没在"自己的轨迹"上被训练过。常见做法是加 DAgger 这类交互式采样,但 DAgger 的代价是每轮都要 teacher 重新标注,训练成本立刻翻倍。

1.2 On-policy 蒸馏的数据分布一致性

on-policy 蒸馏的出发点非常直接:让 student 在环境里自己采样,拿这些 trajectory 去请求 teacher 的监督信号,再用这个监督信号去更新 student。也就是说,student 学的状态和它实际要部署的状态来自同一个分布。

这个思路和 DAgger 一脉相承,但是蒸馏场景下有一些独特的问题。最典型的是:teacher 和 student 的能力差距过大时,student 的探索轨迹质量很低,teacher 给出来的"正确示范"可能和 student 能执行的动作差得很远。于是 student 会不停尝试超出自己能力范围的动作,采样效率和学习效果都不理想。

这是我这个季度最早想解决的问题。也是从这里开始,我意识到 reward shaping 也许能帮上忙。

1.3 为什么 reward shaping 成为最初的抓手

在 on-policy 设定下,student 的策略更新需要有一个 reward 信号。标准做法是用 KL 散度作为约束项,让 student 不要离 teacher 太远。但 KL 散度只约束"分布接近",不保证 student 的行为在环境里有好的表现。行为表现还要靠环境 reward 来引导,于是就成了一个多目标优化:既要像 teacher,又要在环境里拿高分。

reward shaping 就是给这个多目标优化配权重的方法。具体一点说,我当时采用的是类似 potential-based reward shaping 的套路,给 student 的每个状态一个额外奖励,这个奖励通常由 teacher 的价值函数或置信度来提供。

2. Reward shaping 的局限:静态系数是最大的不稳定源

Reward shaping 的好处这里不多说,重点是它的问题。我先把整个框架展开,再讲我踩到的具体坑。

2.1 静态权重的矛盾:早期太弱,晚期太强

我在 6 月的版本里用的是固定系数,形式大概是:

total_reward = env_reward + alpha * shaping_signal

其中shaping_signal来自 teacher 的 value prediction 和 student 自身 value prediction 的差值,或者直接取 teacher 的 confidence。

alpha的值我试了 0.1、0.3、0.5、1.0 几个档位。结果是:alpha 太小,student 前期模仿速度太慢,基本靠自己瞎探索;alpha 太大,student 完全变成 teacher 的影子,teacher 的某些系统性偏差会被完整复制。而且这不是一个可以提前调好的参数。

举个例子。训练早期,student 的动作分布还很散,teacher 给出的 shaping 信号信息量很大,此时 alpha 应该大一些,帮 student 快速收敛到 teacher 的策略附近。但到了训练后期,student 已经接近 teacher 的水平,shaping 信号里可提供的增量信息变少,甚至会成为阻碍——student 自己探索得到的新经验,可能被过大的 shaping 信号压制住,导致学习停滞。

2.2 从两个失败实验里获得的直觉

7 月上旬我做了一组对比实验。一个环境是连续控制类的运动任务,另一个是带稀疏奖励的导航任务。前者 shaping 权重 0.5,后者 shaping 权重 0.5,参数一模一样。

结果很有意思:导航任务上 student 学得不错,因为稀疏奖励本身给不了太多指导,teacher 的 dense shaping 正好补上了这个空缺。但运动任务里环境奖励本身就比较密集,shaping 权重的叠加让 student 的策略出现振荡,reward 曲线忽高忽低,学完之后的策略不如纯环境奖励训练出来的版本。

这两个实验让我确定了一件事:shaping 权重不应该是一个全局常量,它要和"student 当前需要 teacher 的程度"挂钩。换句话说,监督强度应该是自适应的。

2.3 潜在的一个误区:adaptive 不等于简单 annealing

很多人想到"自适应"的第一反应是:按训练步数让 alpha 逐渐衰减。我一开始也这么做了,用了一个简单的 warmup + decay 调度。效果确实比固定系数好一些,但问题没有根本解决。

原因是,student 的学习进度不完全是训练步数的函数。它可能因为某个超参数变动突飞猛进,也可能因为一个不稳定的 batch 突然退化。按步数衰减的 alpha 完全感知不到这些变化,本质上还是开环控制。

真正需要的是一个闭环的自适应机制:根据 student 当前对 teacher 的依赖程度,动态调整监督信号的强度。这个想法听起来不难,但实际操作中需要回答一个关键问题——怎么度量 student 对 teacher 的依赖程度。

3. 6 到 9 月之间的关键改动:从固定监督到自适应监督

这三个月里,我最大的改动可以概括为:把 reward shaping 从人工设计的静态项,逐渐改造成一个由 student 学习状态驱动的动态监督信号。

3.1 第一版自适应:基于 Student 不确定度

第一个可行方案是用 student 自身的分布熵来调制 shaping 权重。核心逻辑是:当 student 对当前状态的动作选择很不确定时,说明它需要 teacher 的指导;当它已经形成明确偏好时,监督权重应该降下来。

具体实现里,我用的是 student policy 输出分布的信息熵,经过一个映射函数转成权重。熵越高权重越大,熵越低权重越小。这个方案有几个优点:实现简单、不需要额外的模型、训练稳定。

但也暴露了一个问题:熵这个量太"粗"。student 可能对某个状态非常自信,但自信的方向是错的。这种情况下,熵很低,权重被压下来,student 就用自己的错误策略去采样,错误被进一步强化。这个现象在高维状态空间里尤其明显。

3.2 第二版改进:引入 Q 值的置信度度量

为了克服熵的局限,我把度量方式从"分布熵"换成了"teacher 和 student 的价值一致性"。具体做法:在采样过程中同时维护 teacher 和 student 对每个状态的价值预测,计算两者的差值。如果差值大,说明 student 对当前状态的理解和 teacher 差距还很大,监督权重提高;如果差值小,说明 student 已经接近 teacher 的水平,监督权重降低。

这个设计比熵的更合理,因为它直接度量了"学生和老师之间的知识差距"这个本质变量,而不是度量学生自身的统计特性。我跑了 ablation:在同样的训练预算下,基于价值差距的自适应方案,最终策略的 reward 和模仿精度都高于基于熵的方案。

3.3 第三版:训练过程中的实际实现要点

实现层面有很细节的问题需要处理。因为 value network 本身就是动态更新的,直接用原始差值会导致权重震荡非常剧烈。我在中间加了一个 EMA 平滑,窗口大概几百步,让权重变化更稳定。

另一个问题是 rew shaping 信号和 teacher 的监督信号在梯度上会互相干扰。我采用了一种类似 stop-gradient 的做法,shaping 信号在计算时对 teacher 的梯度做阻断,只用它来影响 student 策略的 reward term,不让反向传播去耦合两个网络的更新。这一步看似微小,但对稳定性影响极大。

4. 自适应监督机制的核心部件:它其实是一个带状态的门控系统

到了 8 月中旬,我开始把整个机制抽象成一个更清晰的框架。我意识到,我构建的自适应监督本质上是一个门控系统,它决定 teacher 的监督信号在多大程度上注入 student 的优化过程。

4.1 门控信号的三个来源

在这套框架里,监督强度由三个信号共同决定:

第一个是上述的 teacher-student 价值差距。这是最主要的信号,代表知识差距。第二个是 student 的轨迹回报方差。当前阶段 student 的回报方差大,说明行为还不够稳定,需要更多的 teacher 指导来收敛。第三个是相对熵变化率,即 student 策略相对于 teacher 策略的 KL 散度变化趋势。如果 KL 正在快速下降,可以适当降低监督强度,让 student 有更多自由探索的空间。

三个信号通过一个小型 MLP 融合成一个系数。这个 MLP 的输入是三个特征,输出是一个 0 到 1 之间的权重。为了让 MLP 的输出行为可控,我没有直接让这个权重乘以 shaping 项,而是先经过一个 sigmoid 后的线性插值:

final_supervision = gate * teacher_signal + (1 - gate) * student_signal

这样 gate 表示的是"在 teacher 和学生信号之间插值"的比例,而不是简单的开关。语文上更接近"soft switch",而不是 gate。

4.2 为什么这一类改造能稳定提升性能

从实验结果看,这套机制的收益主要不是来自单点大幅提升,而是来自训练曲线的稳定性。

使用固定 alpha 时,训练曲线在中后段会出现反复震荡。使用自适应监督后,这种震荡明显减少。原因很容易理解:fixed alpha 的情况下,shaping 权重不会因为 student 水平的提升而自动减弱,造成策略反复被拉向 teacher;而自适应版本相当于在"student 已经足够强"的时候自动放开约束,让它顺着环境梯度走。

4.3 关于 off-policy 数据的处理:蒸馏数据缓冲池的设计

这个部分聊聊数据层面的改造。

在 on-policy 蒸馏中,数据是不断动态产生的。如果每次都用最新的一批数据来更新,样本效率会比较低。我在 8 月引入了一个次级缓冲池,把最近若干轮的 trajectory 存起来,每轮更新时以一定比例混合新旧数据。

但这里有个隐蔽的风险:过旧的 teacher 监督信号可能和当前 student 的状态分布不匹配。我加入了一个采样权重——根据当前 student 的分布与缓冲数据生成时的分布之间的差异来动态调整采样概率。这个权重和前面说的门控系统是联动的,保证过时数据不会被过度使用。

5. 实验设计与评估:除了最终 reward,还需要看的四个指标

实验部分我没有只用最终 reward 作为评估标准。原因很直接:奖励值受到环境随机性影响很大,一次实验涨一点点不一定代表真的有效。我在这个季度定了一套评估协议,把指标拆成四个。

5.1 四个核心评估指标

指标测什么观测方式
最终环境回报策略在环境中获得的累计回报每 N 轮评估一次
与 teacher 的动作一致性蒸馏质量student 与 teacher 在相同状态下输出动作的余弦相似度
固定状态集上的 KL 散度策略分布距离在预采样状态集上计算 KL
训练曲线方差稳定性多次运行同一实验的回报方差

之前很多项目只看第一个指标,结果训练过程中出现退化完全没察觉。加上后面三个指标后,很多之前看不懂的现象都暴露了原因。

5.2 实验结果的横向比较

我用三个版本做了对比:固定 alpha 0.5、熵自适应、价值差距自适应。三个版本使用相同的数据采样量、相同的优化器、相同的训练步数。

固定 alpha 版本最终回报是四个实验中最差的,训练曲线中段有非常明显的波峰然后下降,典型的"先学后忘"现象。熵自适应版本稍好,但确实存在我在前面说的"自信但错误"问题——它在某些状态上 KL 散度居高不下,因为熵低导致监督信号提前消失,错误被固定住了。价值差距自适应版本表现最稳定,最终回报最高,且多次实验的方差最小。

5.3 一个让我警觉的边界 case

9 月有一个实验让我比较警觉。在最高难度的环境配置下,自适应监督的表现反而不如固定 alpha。排查后发现:这个环境里 reward 本身非常稀疏,student 几乎无法从环境信号中学习,此时自适应机制会把监督权重保持在接近 1 的高位。理论上没问题,但因为 teacher 在这个高难度场景下本身的性能也不够好,student 把 teacher 的坏习惯也彻底学会了。

这说明自适应监督并不是一个彻底解决稀疏奖励的方案。它的前提是 teacher 在目标状态空间上有一定的 quality 保证。如果 teacher 本身 coverage 不够,自适应监督只会更高效地传播 teacher 的不足。

6. 训练稳定性相关的工程实践:那些论文不会写的细节

最后这部分把我这个季度实际工程中遇到并解决的问题集中分享一下,每个都值得单独记录。

6.1 Teacher 和 Student 使用相同的优化器不一定最优

初期我让 teacher 和 student 共用同一套优化器配置。后来发现 teacher 网络通常更深、表达容量更大,它需要的 learning rate 往往比 student 更小。给两边各自设置了独立的 learning rate、gradient clip 和 warmup 之后,训练稳定性上了一个明显的台阶。

我的参考配置:

teacher_opt = Adam(teacher.parameters(), lr=1e-4, clip_norm=0.5) student_opt = Adam(student.parameters(), lr=3e-4, clip_norm=1.0)

6.2 门控网络的重置策略

门控网络的训练本身也是一个问题。它很容易在训练早期就收敛到某个固定的输出值,之后就基本不再变化,相当于自适应机制"名存实亡"。我试过几种手段,最有效的是周期性重置门控网络的最后一层,同时保留前面几层的特征表示。这样既不会让门控完全失稳,又能保持一定的自适应能力。

6.3 关于 BatchNorm 在蒸馏中的使用

如果 student 和 teacher 都用了 BatchNorm,那在蒸馏时会有一个额外的麻烦:两者的统计量更新节奏不同,导致 teacher 的监督信号在训练中不稳定。8 月底我把两个网络的 BatchNorm 都换成了 LayerNorm,问题直接消失。蒸馏场景下,LayerNorm 比 BatchNorm 稳定得多,因为它的归一化不依赖 batch 统计量,彻底绕开了"两个网络在同一 batch 上的统计量不一致"这个坑。

6.4 训练耗时的优化:Gradient Accumulation 与异步推理

蒸馏训练有个效率瓶颈,就是每一步都需要 teacher 前向传播产生监督信号,计算量比普通 RL 多了一倍。我用了几种手段优化:

第一是梯度累积。student 能力较弱的时候,一个 batch 里面的数据差异性很大,用较大的有效 batch size 可以让更新更平滑。我在实际配置里是每 4 个小 batch 累积一次梯度更新。

第二是将 teacher 的前向推理和 student 的环境采样放到异步线程里。因为 teacher 前向推理不参与梯度回传,完全可以在独立线程中以更低的精度(FP16)做推理,再把结果转成 tensor 传回训练线程。用这样的方式,我几乎把 teacher 推理的开销完全隐藏在环境采样时间之内。

第三是共享特征层。如果 teacher 和 student 的输入空间相同,可以让两者共享底层的特征提取模块,从源头减少一次前向计算。但要注意,这个方法只在输入分布比较稳定的场景下有效,如果环境动态分布变化很大,共享层反而会让两个网络互相干扰。

7. 9 月的最后一次调整:引入了基于轨迹的 Early Stop

9 月底,我加了一个比较小的改动,但效果出奇地好:基于 student 轨迹质量的 early stop。

逻辑很朴素。当 student 在连续多个评估周期内,动作一致性和环境回报都超过了某个阈值,就说明它已经达到"足够接近 teacher"的状态。此时继续训练不仅收益极低,还可能导致灾难性遗忘。于是我在训练循环里加了一个早停判断,当连续 3 个评估周期的指标都没有明显提升,且绝对指标超过目标值时,结束训练。

加了这层之后,平均训练时间缩短了大约四分之一,策略稳定性反而提升了。这是我这季度投入产出比最高的一次改动。

8. 还没解决的问题和下一步的方向

最后写点当前的开放问题,希望能和做这个方向的朋友交流。

目前最让我困惑的问题是:当 teacher 本身水平不够高的时候,自适应监督机制应该如何正确响应。从实验结果看,它会倾向于高比重地模仿 teacher,但这在 teacher 性能较差时是一个缺点。一种可能的方向是引入 student 自身的 exploration 信号作为"teacher 质量"的估计器——当 student 能从环境中独立获得正向回报时,即使 teacher 给出的监督信号和环境的矛盾很大,也要降低 teacher 的权重。这个思路听起来合理但还缺乏有效的实现。我打算下一阶段重点试一下 offline evaluation of teacher quality,用一个小型 reward model 来预测 teacher 在特定状态下的可信度,而不是简单地用价值差距来间接度量。

另一个还没完全解决的是多 teacher 融合的问题。蒸馏场景下一步很自然的扩展是多个 teacher 各有所长、各自覆盖不同的状态子空间。现在的门控机制只考虑了单个 teacher 的场景,如果推广到多个 teacher,需要处理不同 teacher 信号之间的冲突和优先级问题。我现在想到的方案是用每个 teacher 在对应状态区间上的历史 accuracy 来作为融合权重,但对如何高效地做状态区间划分还没有特别好的思路。

再有一点关于 evaluation 的想法。蒸馏领域常用的评估方式还是看最终策略的 reward 和环境交互表现。但我越来越觉得,评估维度应该把"策略的可迁移性"也算进去——也就是把蒸馏得到的 student 放到一个与训练环境分布略有不同的评测环境里,测试它的泛化能力。这个指标比静态 benchmark 更有意义,也更贴近实际部署场景。目前这只是个想法,还没有跑系统的实验。

这几个月做下来,最大的体会是:自适应监督本质上不是一种具体的算法,而是一种看待"知识传递"的视角——把 teacher 当作一个随环境变化的监督源,而不是一个静态的数据生产者。这个视角替换之后,很多设计上的取舍都变得清晰了。希望大家在复现和改造的时候,也多往这个方向思考。如果你们在这个问题上有什么进展,欢迎交流。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询