PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcement Learning from Rich Feedback with Distributional DAgger

Rishabh Agrawal, Jacob Fein-Ashley, Paria Rashidinejad

arXiv:2606.05152v2cs.LGcs.AIcs.CL

TL;DR

RLVR often reduces rich supervision to a binary final-answer signal, while existing self-distillation can fail to improve reward monotonically or credit early decisions. The paper introduces DistIL, a distributional DAgger method using forward cross-entropy and full sequence-level gradients. DistIL provides theoretical guarantees and improves performance across scientific reasoning, coding, and hard mathematics.

  • Problem

    RLVR commonly uses a single terminal correctness bit, motivating methods that can exploit richer feedback while reliably converting a better feedback-conditioned teacher into student policy improvement.

  • Method

    DistIL applies distributional DAgger with forward cross-entropy on student-visited states and future-aware sequence-level credit assignment.

  • Results

    DistIL outperforms RLVR and self-distillation baselines across scientific reasoning, coding, and hard mathematical reasoning, with gains up to 9.6 points on scientific reasoning Avg@16 and 3.8 points on AIME25 Avg@16.

  • Takeaways & Limitations

    Forward cross-entropy provides monotonic policy-improvement and regret guarantees while supporting black-box rich-feedback teachers and improved Pass@N.

  • Takeaways & Limitations

    The formulation assumes that feedback induces a teacher whose expected reward is at least as high as the current student.

Abstract

from arXiv · show

Reasoning models have advanced rapidly, but the dominant reinforcement learning from verifiable rewards (RLVR) recipe remains surprisingly narrow: sample many responses and reward each with a single bit indicating whether the final answer is correct. Yet many settings provide rich feedback, including execution traces, tool outputs, expert corrections, and model self-evaluations. We study how to use such feedback through a distributional variant of the classic imitation learning algorithm DAgger, where the learner has local access to an expert distribution on states visited by the current policy. This yields a simple forward cross-entropy objective that admits a blackbox expert and whose sequence-level gradient {conduct rich credit assignment by propagating} future expert-student disagreement back to earlier decisions. We show that prior RL with self-distillation objectives based on reverse KL or Jensen-Shannon fail to guarantee monotonic policy improvement: even when the expert has higher reward, their updates may increase probability on worse actions. In contrast, we show that forward cross-entropy admits monotonic policy improvement and enjoys guarantees on regret. We further show that our objective optimizes a lower bound on teacher-weighted likelihood of success, leading to improved Pass@N. Empirically, our approach, DistIL, improves over RLVR and RL with self-distillation baselines across a variety of domains: scientific reasoning, coding, and solving hard mathematical problems.

1 Introduction

The paper argues that RLVR’s binary terminal feedback underuses richer supervision and proposes DistIL, a distributional DAgger method with theoretical guarantees and broad empirical gains.

  • RLVR typically provides only a binary terminal correctness signal, despite many settings offering richer feedback such as critiques, traces, and expert responses.
  • Existing f-divergence self-distillation can both violate monotonic policy improvement and lose delayed credit for early decisions.
  • DistIL uses forward cross-entropy on student-visited states and propagates later teacher–student disagreement back to earlier decisions.
  • DistIL guarantees monotonic improvement, sublinear regret, and a teacher-weighted success-likelihood bound connected to Pass@N.
  • Across scientific reasoning, coding, and challenging mathematics, DistIL outperforms strong RL and self-distillation baselines across feedback regimes.

2 Background and Problem Formulation

The paper formulates rich-feedback learning as contextual-MDP reinforcement learning in which feedback induces a privileged teacher that is assumed to be at least as good as the student.

  • The environment is a finite-horizon contextual MDP where prompts are sampled, policies generate actions, and transitions produce subsequent states.
  • RL seeks a policy maximizing expected cumulative reward from per-step rewards.
  • RLVR uses sparse outcome supervision, commonly assigning zero intermediate rewards and a terminal binary correctness reward.
  • Rich-feedback distillation trains a student policy to match a feedback-conditioned teacher along student-generated trajectories.
  • Prior methods use divergence measures including reverse KL and Jensen–Shannon divergence for teacher–student matching.
  • The central assumption is that feedback induces a teacher with expected reward at least as high as the current student.

3 Limitations of Existing On-Policy Self-Distillation Methods

The paper identifies two failures in existing on-policy self-distillation: divergence objectives need not improve reward monotonically, and local tokenwise gradients can miss delayed consequences of early actions.

  • 3.1 Does divergence minimization guarantee monotonic policy improvement?: The update direction depends on how teacher–student mismatch aligns with reward, not merely on the teacher’s average advantage.
  • 3.1 Does divergence minimization guarantee monotonic policy improvement?: A three-action bandit demonstrates that correcting excess student probability on mediocre actions can increase probability on a zero-reward action.
  • 3.1 Does divergence minimization guarantee monotonic policy improvement?: Reverse-KL natural policy-gradient updates can decrease expected reward despite a constant teacher advantage.
  • 3.1 Does divergence minimization guarantee monotonic policy improvement?: F-divergence self-distillation does not guarantee monotonic policy improvement even when the teacher is better than the student.
  • 3.2 Local credit assignment leads to suboptimal policies: In a two-step MDP, local credit assignment converges to reward 1/3, whereas the full sequence-level gradient reaches reward 2.
  • 3.2 Local credit assignment leads to suboptimal policies: Local tokenwise approximations omit the future-credit term, so early decisions may receive zero gradient even when their downstream consequences matter.

4 RL from Rich Feedback with Distributional Imitation Learning

DistIL casts rich-feedback self-distillation as distributional DAgger: it matches a local teacher distribution on student-visited states while preserving full sequence-level credit assignment.

  • DistIL replaces DAgger’s single expert action with a local teacher distribution, using either teacher probabilities or teacher samples.
  • The objective does not require teacher probabilities and supports unbiased sample-based estimates, enabling black-box teachers and human experts.
  • DistIL’s gradient combines teacher-weighted local imitation with future-aware credit assignment.
  • The local term directly increases token log-probabilities in proportion to teacher probabilities rather than weighting them by student–teacher mismatch.
  • The future term propagates later disagreement through earlier action log-probabilities, crediting decisions that made later prefixes likely.
  • The practical algorithm normalizes cumulative future loss for long generations and uses PPO-style trust-region updates.

5 Theoretical Properties of DistIL

DistIL’s forward cross-entropy aligns student updates with a better teacher, yielding monotonic improvement and sublinear regret under stated assumptions. It also optimizes a teacher-weighted lower bound on success likelihood, supporting its relevance to Pass@N.

  • Monotonic policy improvement: DistIL’s forward cross-entropy objective yields monotonic policy improvement under fixed-state and local realizability assumptions.The natural-gradient update moves the student toward the teacher, so sufficiently small movement toward a higher-reward teacher improves expected reward.
  • Monotonic policy improvement: Reverse-KL self-distillation can distort the reward-improving direction and increase probability on worse actions, even when the teacher is better.Its log-ratio weighting may suppress good actions that the student overrepresents relative to the teacher.
  • Regret guarantee: DistIL’s online procedure achieves a decaying suboptimality gap relative to the teacher under finite concentrability assumptions.The regret analysis accounts for changing state distributions after each student update and uses KL contraction in a recurrence over intermediate errors.
  • Regret guarantee: Teacher stochasticity determines the regret rate: nearly deterministic teachers yield O(n^-1/2), whereas highly stochastic teachers yield O(n^-1/4).The bound includes teacher-policy variance and a recoverability parameter, alongside student–teacher coverage coefficients.
  • Teacher-weighted likelihood of success: For binary rewards, DistIL maximizes a lower bound on teacher-weighted log-likelihood of success, emphasizing prompts where the teacher is likely correct.This connection is specific to teacher-weighted forward cross-entropy and provides an explanation for strong Pass@N performance.

6 Experimental results

DistIL is evaluated across sparse-reward science, execution-feedback coding, and ground-truth-supervised mathematics, outperforming strong RL and self-distillation baselines. Ablations indicate that future-aware sequence-level credit assignment and an intermediate number of teacher tokens are important for performance and stability.

  • Scientific reasoning: DistIL achieves the best avg@16 performance across most scientific-reasoning columns, with physics gains reaching 8.1/9.6 points over the next-best method on Qwen3-8B/OLMo3.Its lead appears within ∼20 steps and is largely sustained, while SDPO becomes unstable on several domains.
  • Coding: DistIL reaches Accuracy/Avg@16 = 0.656 and Score/Avg@16 = 0.482 on LiveCodeBench, exceeding SDPO and GRPO.The advantage persists across Best@k and Maj@k score and accuracy metrics for all evaluated k.
  • Hard mathematical reasoning: On hard mathematics, DistIL leads most benchmark columns, including AIME25 Avg@16 gains of 3.8 points on Qwen3-4B and 1.4 points on Qwen3-8B.GRPO matches the base model because zero pass rate produces zero advantage; DistIL remains competitive without degrading the stronger base model.
  • Credit assignment: Full sequence-level credit assignment consistently outperforms the local CE baseline, whose training has higher variance and worse final performance.The comparison attributes the difference to propagating credit across the entire response rather than only token-wise locally.
  • Teacher-token ablation: Top-100 teacher-token distillation achieves the best final performance and most stable trajectory, outperforming Top-1 and Full distillation.Full distillation can add noise from low-probability tokens, while Top-1 provides insufficient signal; the experiments therefore use K=100.

7 Related Work

The related work positions DistIL within RLVR, on-policy distillation, and DAgger. Its distinction is a cross-entropy objective with stronger optimization properties and supervision on states visited by the learner.

  • Credit Assignment in RLVR: RLVR obtains automatically verifiable supervision but typically provides sparse outcome-level rewards, making token-level credit assignment difficult for long-horizon reasoning.Examples include unit-test verification for code and final-answer checks for mathematical or scientific problems.
  • On-Policy Distillation and On-Policy Self-Distillation: On-policy distillation reduces distribution mismatch by training on the student’s own rollouts, but typically depends on a stronger external teacher.The paper situates its approach among token-level divergence-based distillation methods.
  • On-Policy Distillation and On-Policy Self-Distillation: DistIL revisits self-distillation with a cross-entropy formulation linked to DAgger’s expert supervision on learner-induced state distributions.The formulation is presented as having well-behaved optimization, more stable training dynamics, and improved gradient estimation.

8 Conclusion

The paper frames rich-feedback reinforcement learning as distillation from a privileged feedback-conditioned teacher and introduces DistIL to address failures in existing self-distillation objectives and local credit assignment. It proves policy-improvement, regret, and success-likelihood guarantees, then reports strong results across feedback regimes.

  • Conclusion: DistIL addresses two identified failures: self-distillation objectives may not guarantee monotonic reward improvement, and local token-wise gradients may miss delayed consequences of early decisions.The paper connects these issues to its distributional DAgger formulation and future-aware credit assignment.
  • Conclusion: DistIL uses forward cross-entropy and future-aware credit assignment, with proofs of monotonic policy improvement, a regret guarantee, and a teacher-weighted lower bound on success probability.The lower bound is presented as explaining improved Pass@N for every N.
  • Conclusion: Empirically, DistIL performs strongly across scientific reasoning, coding, and hard mathematical reasoning with sparse correctness, execution traces, and ground-truth solutions.These experiments compare the method against RLVR and self-distillation baselines.

Appendices

The appendix proof expands the policy-value change under a natural-gradient update and expresses the first-order effect as a reward–gradient covariance. This links the update’s policy improvement behavior to the alignment between rewards and the objective-induced gradient signal.

  • A.1 Proof of Proposition 1: The proof defines teacher-to-student probability ratios and then introduces the score vector to rewrite policy derivatives.These definitions support the subsequent first-order expansion of the policy map.
  • A.1 Proof of Proposition 1: A natural-gradient update is expanded to first order, using ∇θπθ(a) = πθ(a)sa to obtain the induced change in action probabilities.The derivation then substitutes this change into the policy-value expression.
  • A.1 Proof of Proposition 1: The resulting policy-value change is J(πθ′) − J(πθ) = −η Covy∼πθ(r(y), g(y)) + O(η2).Thus, the first-order effect is represented by a covariance between reward and the objective-dependent signal g.

A.2 Proof of Proposition 2: Reverse-KL distillation can decrease reward

The proof constructs examples showing reverse-KL distillation can increase probability on a suboptimal action and that local tokenwise updates can miss delayed feedback. In contrast, the full sequence-level objective propagates future disagreement to earlier decisions and reaches a better policy.

  • Reverse-KL counterexample: Reverse-KL natural-policy-gradient updates can increase probability on a suboptimal action even when the teacher is better.The first-order reward change depends on a covariance term whose sign need not be favorable.
  • Two-step construction: The two-step construction uses a teacher that is uniform initially but recommends R after first action a.This creates delayed feedback that is only visible through the future state reached after the first decision.
  • Local versus sequence-level gradients: The local tokenwise approximation treats prefix visitation as fixed and therefore omits the gradient term caused by changing earlier decisions.That omitted term is the mechanism by which future teacher–student disagreement is assigned to earlier actions.
  • Local versus sequence-level gradients: At initialization, the local update gives no signal to the first-step parameter, whereas the full sequence-level gradient gives a strict update signal.The local approximation leaves the first-step probability unchanged, while the full objective responds to downstream mismatch.
  • Result: The full sequence-level optimization converges to a policy with J(πseq) = 2, whereas the local construction remains strictly suboptimal.The comparison establishes that dropping future-credit terms can produce a strictly worse policy.

B.4.3 Proof of Theorem 1: DistIL regret guarantee

The regret proof bounds NPG-DistIL’s averaged cross-entropy error and converts it into policy suboptimality using concentrability, Hellinger, teacher-variance, and recoverability arguments. The resulting analysis establishes a regret guarantee under finite coverage assumptions.

  • Proof conversion: A KL-to-Hellinger inequality converts the cross-entropy control into an averaged Hellinger error.This intermediate conversion enables the subsequent policy-suboptimality argument.
  • Problem-dependent terms: The final policy suboptimality bound depends on teacher variance and the signed teacher recoverability parameter.Teacher variance measures stochastic imitation difficulty, while recoverability captures the value loss from deviating before returning to the teacher.
  • Assumptions: NPG-DistIL’s regret analysis assumes finite ratio-based and initial KL concentrability coefficients.These coefficients quantify teacher–student coverage and support overlap needed by the bound.
  • Cross-entropy bound: The proof first establishes an averaged cross-entropy suboptimality bound for the iterates generated by NPG-DistIL.The bound is stated in Theorem 2 for the statewise forward cross-entropy loss.

C.1 Hyperparameters for Science and Coding Benchmark

The Science and Coding benchmark configuration follows the SDPO implementation and uses its reported baseline hyperparameters, with additional evaluation settings specified for LCBv6.

  • Science and Coding: Science and Coding experiments follow SDPO’s implementation and reported baseline hyperparameters.The hyperparameters for SDPO, DistIL, and GRPO variants are provided in Tables 3–5.
  • Evaluation settings: Figure 5 evaluates LCBv6 at τ=0.2 and checkpoint-step 80 using Score and Accuracy at Best@k and Maj@k.The evaluation covers k ∈ {2, 4, 8, 16}.
  • Scope boundary: For mathematics benchmarks, the paper instead follows OPSD’s implementation and reports separate hyperparameter tables.OPSD uses Forward-KL divergence, whereas SDPO uses reverse-KL divergence in that setting.

D.1 More results for Coding.

At temperature 0.6, DistIL outperforms SDPO and GRPO on the reported coding metrics.

  • D.1 More results for Coding.: At temperature = 0.6, DistIL achieves Acc/Mean@16 = 0.660 and Score/Mean@16 = 0.482.These results are reported for the coding experiment.
  • D.1 More results for Coding.: DistIL’s Acc/Mean@16 of 0.660 exceeds SDPO’s 0.643 and GRPO’s 0.639.The comparison is reported for the temperature = 0.6 coding setting.
  • D.1 More results for Coding.: DistIL’s Score/Mean@16 of 0.482 exceeds SDPO’s 0.467 and GRPO’s 0.411.The reported gap over GRPO is attributed to GRPO’s inability to exploit execution feedback.
Loading 2606.05152v2…