跳转至

Reinforcement Learning from Rich Feedback with Distributional DAgger

会议: NeurIPS 2026
arXiv: 2606.05152
OpenReview: https://openreview.net/forum?id=6ux7PyTXx9
代码: https://github.com/rishabh-1086/distIL
领域: 强化学习
关键词: 强化学习, 模仿学习, DAgger, 自蒸馏, 信用分配

一句话总结

针对传统 RLVR 标量奖励信用分配缺失以及已有自蒸馏目标(逆向 KL、JS 散度)缺乏单调提升保证的缺陷,本文提出分布级 DAgger 变体 DistIL,利用前向交叉熵与全局序列级未来信用分配,实现了理论单调策略提升、次线性遗憾界与显著的 Pass@N 收益。

研究背景与动机

在大语言模型推理能力的强化学习后训练中,可验证奖励强化学习(RLVR,如 GRPO 及其变体)已成为主流范式。然而,RLVR 严重依赖环境反馈的二值终末奖励:整个推理轨迹仅获得一个最终正确与否的标量反馈,导致序列中所有 token 均摊相同的更新权重。这种极度稀疏的反馈不仅无法识别哪些关键推理步骤导致了成败,而且在组内采样全部正确或全部错误时会出现优势信号归零的梯度塌陷,在缺乏可自动验证答案的非确定性任务中更难以拓展。

许多现实推理任务天然包含比单个正确性标量丰富得多的反馈信号,如代码执行轨迹与报错信息、单元测试用例、人类或强模型生成的自然语言批判与提示等。为充分利用此类丰富反馈,近期兴起了基于策略内自蒸馏(on-policy self-distillation)的方法(如采用逆向 KL 的 SDPO 以及采用 JS 散度的 OPSD),通过在当前模型输入中拼入特权反馈构造特权教师策略,再将其单步蒸馏回无特权学生策略。然而,这类方法隐含着严峻的理论与算法瓶颈:一方面,基于 \(f\)-散度(如逆向 KL、JS 散度)的单步蒸馏在优化理论上并不保证策略收益的单调提升,当教师分布与学生分布存在局部频次偏置时,自然策略梯度更新甚至可能降低更优动作的概率;另一方面,现有方法普遍采用逐 token 的局部梯度近似,割裂了早期动作对后续状态分布的影响,使得早期关键决策无法获得来自下游反馈的延迟信用分配,易收敛至严格次优策略。

受到经典模仿学习算法 DAgger 的启发,本文将自蒸馏重构为在学生访问状态分布上的分布级模仿学习问题。核心 idea:提出基于前向交叉熵的分布级 DAgger 算法 DistIL,利用学生自主探索状态上的特权教师分布作为监督,并引入包含未来长程信用分配的完整序列级策略梯度,证明其具备单调策略提升保证、次线性遗憾界以及显式优化成功似然下界的优良特性。

方法详解

整体框架

DistIL(Distributional Imitation Learning)建立在有限步长上下文马尔可夫决策过程(Contextual MDP)的框架下。给定提示词输入 \(x \sim \rho\),学生策略 \(\pi_\theta\) 自回归生成轨迹 \(y = (y_1, \dots, y_H)\) 并与环境交互收集丰富反馈 \(f\)(例如代码报错日志、参考解答或验证反馈)。模型将该丰富反馈注入当前模型参数(或历史 EMA 策略参数),构造出具有特权信息的教师分布 \(\pi_T(\cdot \mid x, f, y_{1:t-1})\)。不同于传统 DAgger 仅提供单一步骤的确定性动作标签,DistIL 在学生策略访问的前缀状态 \(s_t = (x, y_{1:t-1})\) 上对特权教师策略分布进行直接匹配。

在优化目标上,DistIL 摒弃了 SDPO 等方法使用的逆向 KL 散度 \(\mathrm{KL}(\pi_S \| \pi_T)\),改用学生访问状态分布上的前向交叉熵(Forward Cross-Entropy)目标: $\(\mathcal{L}_{\mathrm{DistIL}}(\theta) = \mathbb{E}_{x \sim \rho, y \sim \pi_\theta(\cdot \mid x)} \left[ \sum_{t=1}^{H} H^{\times}\big(\pi_T(\cdot \mid x, y_{1:t-1}), \pi_\theta(\cdot \mid x, y_{1:t-1})\big) \right]\)$ 该目标天然支持概率已知模型以及黑盒教师(仅需从 \(\pi_T\) 采样即可构造无偏估计),更重要的是推导出的完整序列级策略梯度自然分解为两项:一是当前步骤基于教师概率加权的局部模仿项,二是跨越未来长程交互、将下游交叉熵不一致性回传至上游决策的未来信用分配项。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入提示词与环境状态<br/>x ~ ρ"] --> B["学生策略自主采样轨迹<br/>y ~ π_θ(·|x)"]
    B --> C["获取丰富反馈并构造特权教师<br/>f 与 π_T = sg(π_ϕ(·|s_t, f))"]
    C --> D["前向交叉熵分布级目标构建<br/>H×(π_T(·|s_t), π_θ(·|s_t))"]
    D --> E["未来感知长程信用分配<br/>累积未来散度 C_fut 赋权 ∇log π_θ"]
    D --> F["教师概率加权局部模仿更新<br/>-Σ π_T(a|s_t) ∇log π_θ(a|s_t)"]
    E --> G["序列级策略更新与教师EMA动量迭代<br/>PPO剪裁更新 θ 并更新 ϕ ← βϕ + (1-β)θ"]
    F --> G

关键设计

1. 教师概率加权局部模仿:解决逆向散度的非单调退化

以往自蒸馏方法(如 SDPO、OPSD)使用逆向 KL 或 JS 散度,其梯度中 student log-probabilities 被局部对数比率 \(\log \frac{\pi_T(y \mid s_t)}{\pi_\theta(y \mid s_t)}\) 加权。论文在 Proposition 1 与 Proposition 2 中严谨证明:当使用 \(f\)-散度时,即使教师期望奖励显著高于学生(\(\Delta > 0\)),自然策略梯度更新第一阶变动量仍由奖励与似然比差值的协方差 \(\mathbb{C}\mathrm{ov}_{y \sim \pi_\theta}(r(y), g(y))\) 决定;一旦学生对某一中等奖励动作过度采样,逆向 KL 更新会将其视为误差予以惩罚,归一化后反而导致零奖励极差动作的选取概率上升,导致策略整体期望收益单调下降。

DistIL 的前向交叉熵局部梯度为 \(\nabla_\theta H^{\times}(\pi_T(\cdot \mid s_t), \pi_\theta(\cdot \mid s_t)) = -\sum_y \pi_T(y \mid s_t) \nabla_\theta \log \pi_\theta(y \mid s_t)\)。这一设计使更新方向直接沿教师概率分布加权,将学生策略一阶向高收益教师拉近。在切空间局部可实现性假设下(Proposition 4),DistIL 的自然策略梯度步直接满足 \(J(\pi_{\theta'}) = J(\pi_\theta) + \eta \Delta + O(\eta^2)\),只要特权教师平均质量优于学生(\(\Delta > 0\)),就严格保证策略单调提升。

2. 未来感知长程信用分配:打破局部近似的次优收敛陷阱

SDPO 和 OPSD 在计算梯度时,人为截断了学生策略对状态转移分布 \(\mathbb{P}_\theta(s_t)\) 的微分,退化为纯单步局部梯度。论文通过 Proposition 3 的两步决策反例证实:当奖励延迟出现且只有特定分支才产生教师-学生分歧时,局部梯度在初始对称状态下对第一步决策的梯度恒为 0,只能停留在期望奖励为 \(1/3\) 的对称陷阱;而保留完整状态分布梯度的序列级优化则能够驱动第一步转移概率偏移,收敛至最优期望奖励 \(2/5\)。

DistIL 通过对目标函数进行关于轨迹概率分布的完整求导,自然导出了未来长程信用分配项: $\(\mathbb{E}_{x \sim \rho, y \sim \pi_\theta(\cdot \mid x)}\left[\sum_{t=1}^{H}\nabla_{\theta}\log\pi_{\theta}(y_{t}\mid s_{t})\left(\sum_{\tau>t}H^{\times}\Big(\pi_{T}(\cdot\mid s_{\tau}),\;\pi_{\theta}(\cdot\mid s_{\tau})\Big)\right)\right]\)$ 该设计让每个中间 token 的动作似然更新直接与未来轨迹的交叉熵惩罚挂钩,若当前 token 导致后续状态与教师发生剧烈偏离,该分支即刻受到强负反馈抑制。为防止在长思维链推理中因累积步数过多产生长度惩罚偏差,算法在具体实现时对未来损失做了归一化因子 \(\frac{1}{|y_j| - t - 1}\) 处理。

3. 成功似然下界优化:打通 Pass@N 收益机理

在具有二值校验结果的推理任务中,测试集衡量指标往往聚焦于采样能力(如 Pass@N、Best@N 与 Maj@N)。DistIL 在理论上建立了与最大似然强化学习的直接联系(Proposition 5):对于任意输入 \(x\) 与二值奖励 \(r(x, y) \in \{0, 1\}\),最小化前向交叉熵本质上等价于最大化成功概率对数似然的加权下界: $\(-H^\times\big(\pi_T(\cdot \mid x), \pi_\theta(\cdot \mid x)\big) \leq p_{\pi_T}(x) \log p_{\pi_\theta}(x)\)$ 其中 \(p_\pi(x) := \mathbb{E}_{y \sim \pi}[r(x, y)]\)。这一数学性质是逆向 KL 所不具备的。它表明 DistIL 会对特权教师成功率更高的难样本赋予更大权重,从而在采样空间中显式抬升成功解的生成似然,在数学机制上解释了为何 DistIL 能够在各类 Best@N 与 Maj@N 上全面超越基线。

损失函数 / 训练策略

在训练实践中,DistIL 采用成组采样与 PPO-style 信任域裁剪算法结合。对于采样得到的 \(G\) 条轨迹,计算每步的局部模仿梯度与未来感知信用项,形成批次联合梯度: $\(\nabla_{\theta}\mathcal{L}_{\mathrm{DistIL}} = \frac{1}{G}\sum_{j=1}^{G}\sum_{t=1}^{|y_{j}|} \left( C^{\mathrm{fut}}_{j,t} \cdot \nabla_{\theta}\log\pi_{\theta}(y_{j,t}\mid s_{t}) - \sum_{a \in \mathcal{K}} \mathsf{sg}\big(\pi_{\phi}(a\mid s_{t},f_{j})\big) \cdot \nabla_{\theta}\log\pi_{\theta}(a\mid s_{t}) \right)\)$ 其中教师词表动作集合 \(\mathcal{K}\) 在实验中截断为 Top-100 token,兼顾计算开销与避免尾部低概率噪声干扰。教师模型参数 \(\phi\) 采用学生参数 \(\theta\) 的指数移动平均(EMA)进行动量更新:\(\phi \leftarrow \beta \phi + (1-\beta) \theta\)。

实验关键数据

主实验

论文在三大基准任务上展开评估: 1. 科学推理(SciKnowEval L3):在化学、物理、生物、材料四个本科级别学科上评估,模型基础为 Qwen3-8B 和 Olmo3-7B-Instruct,对比 GRPO(on-policy / off-policy)、SDPO 与 DistIL 在 1h 与 5h 墙钟训练下的 Avg@16。 2. 代码生成(LiveCodeBench v6):包含丰富执行报错日志,评估 Exact Pass Score 及 Partial Accuracy。 3. 高难度数学竞赛推理:精选 738 道 Qwen3-4B pass@512=0 的极难题目,以标准题解作为特权反馈,评测 AIME24、AIME25、HMMT25、AMC23 及 Minerva。

下表摘录 SciKnowEval L3 科学推理主对比(Avg@16):

模型与方法 化学 (1h / 5h) 物理 (1h / 5h) 生物 (1h / 5h) 材料 (1h / 5h) 平均 (1h / 5h)
Qwen3-8B (Zero-shot) 41.2 / 59.2 30.8 / 58.9 47.5 / - - / - - / -
+ GRPO (off-policy) 65.9 / 74.5 63.8 / 72.7 35.1 / 59.9 74.3 / 77.1 59.8 / 71.1
+ GRPO (on-policy) 63.3 / 63.4 63.6 / 63.6 49.8 / 49.8 73.9 / 74.1 62.7 / 62.7
+ SDPO 73.0 / 80.2 68.0 / 72.4 52.9 / 63.6 72.2 / 75.9 66.5 / 73.0
+ DistIL (本文) 75.8 / 80.8 72.7 / 80.8 53.3 / 66.6 74.9 / 76.2 69.2 / 76.1
Olmo3-7B-Instruct 22.8 / 37.7 16.2 / 36.7 28.4 / - - / - - / -
+ GRPO (off-policy) 39.7 / 56.7 55.3 / 63.3 35.6 / 55.8 70.9 / 75.0 50.4 / 62.7
+ GRPO (on-policy) 51.4 / 57.5 62.7 / 62.7 49.8 / 49.8 73.3 / 73.5 59.3 / 60.9
+ SDPO 70.2 / 79.2 59.8 / 64.9 49.5 / 52.9 71.8 / 78.1 62.8 / 68.8
+ DistIL (本文) 72.1 / 81.0 67.4 / 74.5 47.8 / 55.3 73.5 / 76.9 65.2 / 71.9

下表摘录极难数学推理竞赛基准主对比(Qwen3-4B 与 Qwen3-8B,指标为 Avg@16 / Pass@16):

骨干模型 优化方法 AIME24 (Avg/Pass) AIME25 (Avg/Pass) HMMT25 (Avg/Pass) AMC23 (Avg/Pass) Minerva (Avg/Pass)
Qwen3-4B Base Model 61.4 / 82.0 50.3 / 69.5 30.3 / 48.8 93.8 / 98.8 43.2 / 48.6
SFT 55.8 / 80.9 43.1 / 67.4 29.7 / 46.0 91.7 / 97.0 41.9 / 50.4
GRPO 61.4 / 82.0 50.3 / 69.5 30.3 / 48.8 93.8 / 98.8 43.2 / 48.6
SDPO 60.9 / 85.4 49.6 / 74.1 32.8 / 50.9 93.9 / 99.8 43.4 / 52.0
OPSD 63.2 / 85.4 51.5 / 76.1 33.0 / 54.2 94.5 / 99.6 43.4 / 51.8
DistIL (本文) 65.3 / 87.5 55.3 / 77.6 33.0 / 56.6 94.8 / 100.0 44.2 / 52.9
Qwen3-8B Base Model 75.9 / 86.8 66.9 / 79.6 44.7 / 65.6 95.6 / 100.0 48.9 / 56.3
SFT 74.2 / 85.3 65.3 / 80.7 42.2 / 64.4 96.3 / 100.0 48.2 / 56.2
GRPO 75.9 / 86.8 66.9 / 79.6 44.7 / 65.6 95.6 / 100.0 48.9 / 56.3
SDPO 76.9 / 89.5 68.3 / 84.0 45.6 / 71.2 96.5 / 100.0 49.2 / 57.7
OPSD 76.5 / 91.3 69.7 / 83.3 45.5 / 68.2 96.2 / 100.0 49.2 / 57.8
DistIL (本文) 76.4 / 90.7 71.1 / 85.0 46.4 / 71.4 96.6 / 100.0 49.5 / 58.4

消融实验

论文对未来信用分配项的有效性以及教师分布词表截断深度 \(K\) 进行了细致消融:

消融配置 评测基准与设置 核心表现与指标 结论说明
DistIL (Full) 科学推理 (Material) 训练曲线平稳上升,最终 Avg@16 达到 76.2 全局序列级未来信用分配有效抑制局部过拟合
CE baseline (仅局部信用) 科学推理 (Material) 训练方差剧烈增大,性能显著落后 DistIL 仅保留局部交叉熵项时,无法捕捉多步推理的延迟信用
Top-1 蒸馏 (\(K=1\)) 科学推理 (Physics) 整体准确率最低,收敛速度缓慢 单一离散 token 无法传递教师连续分布的不确定性与软边界
Top-100 蒸馏 (\(K=100\)) 科学推理 (Physics) 达到最高最终精度与最平稳训练轨迹 既覆盖教师核心概率质量,又滤除了尾部随机噪声
Full 词表蒸馏 (\(K=\mathrm{Full}\)) 科学推理 (Physics) 性能轻微落后于 Top-100 全词表长尾低概率词引入了不必要的梯度噪声

关键发现

  • 训练平稳性与单调性:在科学推理多任务中,SDPO 在训练进行至 100 步左右时,生物学域的 Best@16 发生了明显崩溃衰减,并在物理与化学域呈现大幅振荡;相反,DistIL 的验证指标在 20 步以内迅速确立领先,并在整个训练周期内呈现平稳单调增长,印证了 Proposition 4 的单调提升保证。
  • 冷启动与极端困难样本突围:在 base 模型 pass@512=0 的极难数学题中,由于无任何采样命中答案,GRPO 的组内优势严格为零,无法推进训练;直接 SFT 则因长程解题记忆机制不足而导致泛化性能全面崩塌(低于 base 模型);而 DistIL 借助特权参考解答,将 Avg@16 在 AIME25 上大幅提升 5.0 个百分点。
  • 跨环境执行信号的利用效率:在 LiveCodeBench 上,DistIL 的 Accuracy/Avg@16 达到 0.656、Score/Avg@16 达到 0.482,不仅碾压无法利用报错信息的 GRPO(0.600 / 0.405),也显著优于受制于逆向散度偏置的 SDPO(0.643 / 0.467)。

亮点与洞察

  • 理论澄清了自蒸馏中散度选择的本质缺陷:通过构造反例严格揭示了逆向 KL 与通用 \(f\)-散度在策略优化中的不可靠性,指出了散度偏置与动作协方差可能诱发奖励逆向退化,为 LLM 后训练散度设计树立了理论标杆。
  • 全梯度展开破除局部视野:明确指出了以往工作将状态访问分布视为常数冻结的隐性弊端,借助序列级完整导数将后续状态的不一致性显式反传给早期 token,无须额外训练过程奖励模型(PRM)即可内生实现精细化步骤信用分配。
  • 黑盒特权教师的无缝兼容:前向交叉熵结构使得目标函数对教师策略的概率形式无硬性要求,仅依赖从教师处获取的样本即可完成无偏梯度估计,这为引入非开源大模型、编译器执行器或人类专家即时互动等更广泛的黑盒特权反馈开辟了可行通路。

局限与展望

  • 教师策略的质量假设依赖:理论分析均建立在特权教师期望奖励优于或等于学生的前提假定(\(\Delta \geq 0\))之上。若外部环境反馈质量极低,或提示诱导的特权模型产生严重幻觉与误导性步骤,该前向模仿更新可能会强化错误归纳。
  • 未来信用计算的显存与计算开销:序列级全梯度推导要求计算每一步后继序列的累积交叉熵差值,尽管采用了归一化与并行张量优化,相比纯局部前向传递仍带来了一定的额外计算开销。
  • 未来方向:探索与动态强化探索机制(如主动不确定性探索)的结合,进一步研究如何在完全开放域(无地面真实标签且执行器不可用)的环境下,由模型自适应构建兼具高奖励与高信息量的特权教师分布。

相关工作与启发

  • vs GRPO / RLVR (Shao et al., 2024): GRPO 将全序列粗糙地赋予同一标量优势,且面临全负/全正批次的梯度归零问题;DistIL 借由特权反馈引入密集 token 级监督,并以未来交叉熵反传实现内生信用分配,在全负样本冷启动阶段表现出降维打击的优势。
  • vs SDPO (Hübotter et al., 2026) & OPSD (Zhao et al., 2026): SDPO 基于逆向 KL 散度且仅使用局部梯度截断,OPSD 使用 JS 散度;DistIL 证明了两者在策略提升上缺乏单调性保证且容易收敛至次优解,改用具备单调性与次线性遗憾界的前向交叉熵目标,同时通过全梯度实现了未来感知信用传递。
  • vs 过程奖励模型 PRM (Lightman et al., 2023): PRM 需要高昂的人工逐步标注或单独训练步级判别模型;DistIL 将丰富环境反馈直接嵌入条件提示,免去了辅助判定模型的训练与对齐负担。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从理论上证明了现有自蒸馏散度的非单调退化与局部梯度缺陷,创新性地提出分布级 DAgger 范式。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖科学推理、代码执行与数学难题三大经典场景,模型跨越 Qwen3 与 Olmo3,指标全面包含 Best@N、Maj@N 与 Avg@N。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导与算法动机紧密咬合,反例构造直击要害,行文逻辑清晰严谨。
  • 价值: ⭐⭐⭐⭐⭐ 为大模型利用复杂外部环境交互信号与执行轨迹后训练提供了兼具理论保证与实用价值的全新基础算法范式。