Reinforcement Learning from Rich Feedback with Distributional DAgger

Recent advances in large-scale pretraining have propelled large language models (LLMs) to the forefront of AI research, enabling impressive performance in natural language understanding, reasoning, and generation tasks. Early methods such as REINFORCE, Actor-Critic, and proximal policy optimization (PPO) laid the foundation for reinforcement learning (RL) in these models. However, the reliance on sparse, delayed rewards—often only indicating whether the final answer is correct—limits the ability to learn intermediate reasoning steps effectively. To address this, reinforcement learning from verifiable rewards (RLVR) was developed, exemplified by algorithms like GRPO and SDPO, which leverage automatic verification signals such as code correctness or mathematical solution validity. While successful in specific domains, these approaches suffer from issues like poor credit assignment, instability, and inability to utilize richer feedback signals.

In parallel, knowledge distillation and self-distillation techniques emerged, aiming to leverage intermediate signals like execution traces, expert critiques, and ground-truth annotations. These methods attempt to turn sparse supervision into dense, token-level guidance, improving learning efficiency. Nonetheless, existing self-distillation objectives, especially those based on f-divergences like reverse KL or Jensen-Shannon, lack theoretical guarantees of monotonic policy improvement. Empirical results show that these objectives can sometimes degrade performance, especially when the divergence measures misalign with reward signals. Additionally, local token-wise gradient approximations used in prior work ignore the delayed effects of early decisions, further limiting their effectiveness.

This context motivates the development of a new framework that can harness rich feedback effectively while guaranteeing policy improvement. The authors revisit classical imitation learning algorithms, particularly DAgger, and propose a distributional variant—DistIL—that models expert feedback as a local distribution over visited states. By optimizing a forward cross-entropy loss and propagating future disagreements through sequence-level gradients, DistIL ensures that each update moves the policy closer to the optimal, with theoretical guarantees of monotonic improvement and regret bounds. This approach bridges the gap between imitation learning and reinforcement learning, enabling stable, data-efficient training in complex reasoning tasks.

The core challenge addressed in this paper is the effective utilization of rich, structured feedback signals in reinforcement learning for complex reasoning tasks. Traditional RL methods rely on sparse, delayed rewards, which hinder credit assignment for intermediate steps, leading to slow convergence and suboptimal policies. Existing self-distillation approaches attempt to leverage richer signals but often rely on divergence measures like reverse KL or Jensen-Shannon, which do not guarantee policy improvement and can even cause performance degradation. Furthermore, local token-wise gradient approximations ignore the influence of early decisions on future states and rewards, resulting in policies that may converge to suboptimal solutions. The fundamental problem is designing an algorithm that can incorporate rich feedback, propagate future disagreements backward, and ensure monotonic policy improvement, all while being applicable to black-box experts and sample-based estimation scenarios.

The key innovations introduced in this work include: 1) a theoretical analysis revealing the limitations of existing f-divergence-based self-distillation objectives, demonstrating their failure to guarantee monotonic policy improvement; 2) the formulation of DistIL, a distributional imitation learning algorithm that employs a forward cross-entropy loss combined with a future-aware gradient propagation mechanism, enabling sequence-level credit assignment; 3) the development of a full-gradient optimization strategy that propagates future disagreements back to earlier decisions, ensuring each policy update improves the expected reward monotonically. This approach supports black-box expert integration and sample-based estimation, broadening its practical applicability. Theoretically, the authors prove that DistIL guarantees monotonic improvement and sublinear regret, providing a solid foundation for stable learning in rich feedback environments. Empirically, the method outperforms existing baselines across multiple challenging tasks, validating its effectiveness and robustness.

• �� Model the expert feedback as a local distribution over visited states, allowing black-box and sample-based supervision.
• �� Define a distributional imitation learning objective based on the forward cross-entropy between student and expert state distributions, supporting both explicit probabilities and sample estimates.
• �� Incorporate a future-aware gradient component that propagates sequence-level disagreements (measured via cross-entropy) back to earlier decisions, enabling sequence-level credit assignment.
• �� Optimize the combined objective using full gradients, avoiding local token-wise approximations that ignore delayed effects.
• �� Theoretically analyze the properties of the proposed objective, proving monotonic policy improvement and sublinear regret bounds.
• �� Implement the algorithm with a PPO-style trust-region update, updating the student policy and an exponential moving average of the teacher, conditioned on rich feedback signals.
• �� Conduct experiments on scientific reasoning, coding, and mathematical datasets, comparing against RLVR, SDPO, and other baselines, with ablation studies on gradient components and feedback types.

• �� Use the Qwen3-8B model fine-tuned on scientific reasoning datasets (e.g., MATH, ARC) to evaluate validation Best@16 and Maj@16 metrics, tracking training dynamics and early performance gains.
• �� Test on code generation benchmarks like CodeX and OpenAI Codex, measuring code accuracy and logical reasoning capabilities.
• �� Evaluate on mathematical reasoning datasets such as HardMath and MATH, assessing problem-solving accuracy on high-difficulty tasks.
• �� Compare DistIL with baselines including RLVR, SDPO, OPSD, and GRPO, using consistent hyperparameters and multiple random seeds for statistical robustness.
• �� Perform ablation studies to isolate the impact of future gradient propagation and expert distribution modeling, analyzing convergence speed, stability, and performance improvements.
• �� Analyze the covariance between reward and divergence measures to understand the conditions under which the method guarantees improvement.

• �� DistIL achieves an average of 12% higher validation Best@16 scores across scientific reasoning tasks, with early gains appearing within the first 20 training steps and maintained throughout training. It exhibits less oscillation and more stable convergence compared to SDPO and RLVR.
• �� In code generation, DistIL improves accuracy by 8%, particularly excelling in tasks requiring complex reasoning and long sequence generation, demonstrating better sample efficiency.
• �� On high-difficulty math problems, DistIL increases problem-solving accuracy by 10%, validating its effectiveness in high-stakes reasoning scenarios.
• �� Ablation results confirm that the future-aware gradient component significantly enhances convergence stability and prevents policy degradation, outperforming local gradient approximations.
• �� Theoretical analysis shows that the proposed objective guarantees monotonic improvement under mild assumptions, aligning well with empirical observations.

• �� The approach is directly applicable to training large language models for scientific reasoning, automated coding, and complex mathematical problem-solving, where rich feedback signals are available.
• �� It can be integrated into AI assistants, automated theorem proving, and intelligent tutoring systems, enhancing their learning efficiency and robustness.
• �� Long-term, the framework could enable autonomous scientific discovery systems that learn from detailed experimental logs, expert critiques, and multi-modal feedback, accelerating innovation in research and development.

• �� The reliance on high-quality, dense feedback signals limits applicability in environments with sparse or noisy annotations.
• �� The computational cost of full-gradient propagation and sequence-level credit assignment can be high, especially for very large models and long sequences.
• �� The current evaluation focuses on controlled benchmarks; robustness in real-world, noisy, or less-structured feedback scenarios remains to be validated. Future work should address scalability and feedback quality issues.