💬 NLP

Reinforcement Learning from Rich Feedback with Distributional DAgger

This paper introduces DistIL, a distributional DAgger-based approach that leverages rich feedback through a forward cross-entropy objective to achieve monotonic policy improvement and superior performance in reasoning, coding, and mathematical tasks compared to standard RLVR and self-distillation methods.

Original authors: Rishabh Agrawal, Jacob Fein-Ashley, Paria Rashidinejad

Published 2026-06-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Rishabh Agrawal, Jacob Fein-Ashley, Paria Rashidinejad

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a brilliant but inexperienced student how to solve complex puzzles, like writing code, solving math problems, or explaining scientific concepts.

The Old Way (RLVR): The "Pass/Fail" Exam
Currently, the most popular way to train these AI models is like giving them a final exam with only one piece of feedback: "Pass" or "Fail."

The student tries to solve a problem.
If the final answer is right, they get a gold star. If it's wrong, they get a red X.
The Problem: The teacher doesn't tell the student why they failed. Did they make a mistake in the first step? The middle? The end? Because the feedback is just a single "Pass/Fail" at the very end, the student has to guess which specific words or steps caused the error. It's like trying to learn to drive by only being told "You crashed" after the car hits a wall, without knowing if you turned the wheel too early, braked too late, or forgot to check your mirrors.

The New Idea: Rich Feedback
In the real world, we often have much better feedback.

In coding: We get error logs showing exactly which line broke.
In math: We might have a step-by-step solution or a hint.
In science: We might have a critique of the reasoning process.
The paper argues that we should use this "rich" feedback instead of just the final Pass/Fail.

The Problem with Current "Rich" Methods
Researchers have tried to use this rich feedback by having the AI model act as both the student and the teacher. The model generates a solution, gets feedback, and then tries to mimic the "corrected" version of itself.
However, the authors found two major flaws in how this is currently done:

The "Wrong Direction" Trap: Sometimes, even if the teacher is smarter than the student, the math used to teach the student can accidentally push them away from the right answer. It's like a coach telling a player, "You're doing it wrong," but the specific instruction causes the player to trip over their own feet.
The "Blind Spot" Problem: Current methods only look at the immediate mistake. If the student makes a tiny error in step 1 that leads to a huge disaster in step 10, current methods often ignore step 1 because the error didn't show up until step 10. They fail to give credit (or blame) to the early decisions that caused the later problem.

The Solution: DistIL (Distributional Imitation Learning)
The authors propose a new algorithm called DistIL. Think of it as a new, smarter coaching method based on a classic technique called "DAgger."

Here is how DistIL works, using a simple analogy:

The "Future-Proof" Coach: Instead of just looking at the mistake happening right now, DistIL looks at the whole journey. If the student makes a small choice at the beginning that leads to a bad outcome later, DistIL traces that bad outcome all the way back to the beginning and says, "Hey, that first choice was the problem!" This is called future-aware credit assignment.
The "Direct Imitation" Rule: Instead of using complex math that sometimes gets the direction wrong, DistIL uses a simple "forward cross-entropy" rule. Imagine the teacher says, "Do exactly what I do." DistIL simply tries to match the teacher's probability of choosing the right path. The authors prove mathematically that this method always moves the student in the right direction (monotonic improvement) and never accidentally makes them worse.
Black-Box Friendly: This method is flexible. It can learn from a human expert, a computer program, or another AI model, even if that "teacher" is a black box (you can't see inside its brain, you just see its answers).

The Results
The team tested DistIL on three difficult tasks:

Scientific Reasoning: (Biology, Chemistry, Physics). DistIL learned faster and more stably than previous methods, which tended to get confused and unstable after a while.
Coding: When the AI had to write code that actually ran, DistIL used the error logs to improve much better than methods that only looked at the final "Pass/Fail."
Hard Math: On extremely difficult math problems where the AI usually fails completely, DistIL was able to learn from the correct solutions and significantly improve its success rate, while other methods failed to improve at all.

In Summary
The paper introduces DistIL, a new way to teach AI models using detailed feedback (like error logs or step-by-step hints) rather than just a final grade. It fixes the flaws of current methods by ensuring the AI always learns in the right direction and by connecting early mistakes to their later consequences. The result is an AI that learns faster, more stably, and solves harder problems.

Technical Summary: Reinforcement Learning from Rich Feedback with Distributional DAgger

Problem Formulation

The paper addresses the limitations of the dominant paradigm in reasoning model training: Reinforcement Learning from Verifiable Rewards (RLVR). In RLVR, models generate multiple candidate solutions, and each is assigned a single binary reward (correct/incorrect) based on the final answer. This approach suffers from two critical flaws:

Poor Credit Assignment: The terminal reward is broadcast uniformly across all tokens in a response, failing to distinguish which specific reasoning steps contributed to success or failure.
Inapplicability to Unverifiable Domains: RLVR relies on automatic verification of final answers, making it unsuitable for tasks where ground-truth answers are unavailable or difficult to check.

While many settings offer "rich feedback" (e.g., execution traces, error logs, expert corrections, or self-evaluations), current methods attempting to leverage this via on-policy self-distillation (e.g., SDPO, OPSD) face theoretical and practical hurdles. Existing approaches typically minimize $f$ -divergences (like reverse KL or Jensen-Shannon) between a student policy and a feedback-conditioned teacher. The authors argue these methods lack guarantees for monotonic policy improvement and fail to perform sequence-level credit assignment, often leading to suboptimal convergence.

Methodology: DistIL

The authors propose DistIL (Distributional Imitation Learning), a distributional variant of the classic DAgger algorithm adapted for reinforcement learning with rich feedback.

Core Objective

DistIL optimizes a forward cross-entropy loss between the student policy ( $\pi_\theta$ ) and a feedback-conditioned teacher policy ( $\pi_T$ ) on states visited by the student:
$\mathcal{L}_{\text{DistIL}}(\theta) := \mathbb{E}_{x \sim \rho, y \sim \pi_\theta(\cdot|x)} \left[ \sum_{t=1}^H H_\times(\pi_T(\cdot | s_t), \pi_\theta(\cdot | s_t)) \right]$
where $H_\times(p, q) = -\mathbb{E}_{y \sim p}[\log q(y)]$ is the forward cross-entropy. Crucially, the teacher $\pi_T$ is conditioned on rich feedback $f$ (e.g., execution logs or ground truth) but is held fixed during the student update.

Gradient and Credit Assignment

Unlike prior self-distillation methods that approximate gradients using only local token-wise discrepancies (effectively freezing the state distribution), DistIL utilizes full gradients. The gradient decomposes into two terms:

Local Imitation: A teacher-weighted term that updates the student to match the teacher's distribution at the current state.
Future-Aware Credit Assignment: A term that propagates the cross-entropy disagreement observed at future timesteps back to earlier decisions. This ensures that early token choices influencing the trajectory to a state with high teacher-student disagreement receive appropriate credit or penalty.

The algorithm supports black-box teachers (external models or human experts) as it does not require access to teacher log-probabilities, only samples or the ability to estimate the objective.

Key Contributions

1. Theoretical Limitations of Existing Methods

The paper provides formal proofs demonstrating that existing on-policy self-distillation objectives based on $f$ -divergences (including reverse KL and Jensen-Shannon) suffer from two fundamental issues:

No Monotonic Improvement Guarantee: Even if the teacher has a strictly higher expected reward than the student, minimizing an $f$ -divergence does not guarantee that a natural gradient step will increase the student's expected reward. The update direction may inadvertently increase the probability of suboptimal actions due to the weighting by log-ratios.
Suboptimal Convergence via Local Approximation: By ignoring the "future-credit" term (the impact of current actions on future state distributions), local gradient approximations can converge to strictly suboptimal policies. The authors prove that in simple two-step MDPs, local updates can fail to assign credit to early decisions that are crucial for reaching high-reward states.

2. Theoretical Guarantees of DistIL

DistIL is shown to possess properties that align optimization with reward maximization:

Monotonic Policy Improvement: Under standard local realizability assumptions, the forward cross-entropy objective guarantees that moving toward a better teacher increases the student's expected reward.
Regret Bounds: The algorithm admits a sublinear regret bound, with convergence rates dependent on the concentrability coefficients and the variance of the teacher.
Pass@N Optimization: The objective is proven to maximize a lower bound on the teacher-weighted likelihood of success. This provides a principled explanation for why DistIL improves Pass@N metrics (the probability of generating at least one correct answer in $N$ samples) across all $N$ .

3. Empirical Performance

DistIL was evaluated across three domains:

Scientific Reasoning: On SciKnowEval (Chemistry, Physics, Biology, Materials), DistIL outperformed SDPO, GRPO (on-policy and off-policy), and achieved higher Best@16 and Maj@16 scores. Notably, DistIL showed stable training, whereas SDPO exhibited significant oscillations and performance collapse in certain domains.
Coding: On LiveCodeBench, utilizing execution traces as feedback, DistIL outperformed SDPO and GRPO in both exact correctness and partial accuracy.
Hard Mathematical Reasoning: On challenging problems (AIME, HMMT) where the base model has a zero pass rate, RLVR methods (GRPO) failed to improve due to zero advantage signals. DistIL successfully leveraged ground-truth solutions to improve performance, outperforming self-distillation baselines (SDPO, OPSD) and SFT.

Significance and Claims

The paper claims that DistIL bridges the gap between imitation learning and reinforcement learning by providing a theoretically grounded framework for learning from rich feedback. Its significance lies in:

Correcting the Objective: Moving away from $f$ -divergence minimization (which can be reward-misaligned) to forward cross-entropy, which guarantees monotonic improvement when the teacher is better.
Enabling Sequence-Level Learning: By utilizing full gradients, DistIL solves the credit assignment problem inherent in long-horizon reasoning, ensuring early decisions are evaluated based on their downstream consequences.
Practical Robustness: The method works with black-box teachers and diverse feedback types (from binary correctness to execution logs), making it applicable to a wider range of reasoning tasks than standard RLVR.

The authors conclude that the combination of a reward-aligned objective and future-aware credit assignment is essential for effectively distilling rich feedback into reasoning models, leading to consistent empirical gains over state-of-the-art baselines.