UOPD: Uncertainty-Aware Intervention for On-Policy Distillation of Multi-Turn Agents
cs.LG, cs.AI
Submitted: 2026-09-27
Updated: 2026-09-27
Code: https://github.com/onepounchman/UOPD
License: http://creativecommons.org/licenses/by/4.0/
The gist: On-policy distillation (OPD) trains a student on its own rollouts using dense supervision from a teacher.
Terminology
Abstract
On-policy distillation (OPD) trains a student on its own rollouts using dense supervision from a teacher. In multi-turn environments, a mistake at a critical decision step can redirect the subsequent rollout toward poor outcomes. We use low teacher confidence on student actions to select high-uncertainty steps for correction. In a controlled ALFWorld study, a single teacher correction at a low-confidence step improves subsequent student behavior and task success, motivating selective intervention during distillation. We propose UOPD, an uncertainty-aware intervention method for on-policy distillation. At low-uncertainty turns, UOPD executes student actions and applies the standard OPD loss. At high-uncertainty turns, it samples and executes teacher actions and trains the student to imitate them through supervised fine-tuning, which minimizes forward Kullback-Leibler divergence in expectation. UOPD utilizes adaptive uncertainty thresholds to target a scheduled intervention rate. Empirically, we evaluate UOPD across a broad range of agentic tasks, including ALFWorld, WebShop, and Search, demonstrating its superior performance over OPD methods and their variants. UOPD improves WebShop score by up to 15.8% relative to standard OPD.
Sources
- Reinforcement Learning from Rich Feedback with Distributional DAgger
- Look Ahead Before You Distill: Future Trajectory Validation of Teacher Guidance for Agentic On-Policy Distillation
- Rubric-based On-policy Distillation
- Group-in-Group Policy Optimization for LLM Agent Training
- Distilling the Knowledge in a Neural Network
- Asymmetric On-Policy Distillation: Bridging Exploitation and Imitation at the Token Level
- Trajectory-Refined Distillation
- Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning
- Entropy-Aware On-Policy Distillation of Language Models
- Respecting Self-Uncertainty in On-Policy Self-Distillation for Efficient LLM Reasoning
- On-Policy Distillation with Curriculum Turn-level Guidance for Multi-turn Agents
- What and When to Distill: Selective Hindsight Distillation for Multi-Turn Agents
- Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe
- ADWIN: Adaptive Windows for Horizon-Aware On-Policy Distillation
- Multi-Turn On-Policy Distillation with Prefix Replay
- Your Teacher Can't Help You Here: Combating Supervision Fidelity Decay in On-Policy Distillation
- Self-Distilled Agentic Reinforcement Learning
- KL for a KL: On-Policy Distillation with Control Variate Baseline
- Trinity-RFT: A General-Purpose and Unified Framework for Reinforcement Fine-Tuning of Large Language Models
- SPOT: Sparse Probing and Outcome Calibration for On-Policy Distillation
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks