On-policy distillation preserves student gains beyond teacher performance
Beyond the Teacher: Understanding and Preserving Distillation Overshoot in Agentic On-Policy Distillation
Machine LearningComputation and Language
Summary
When training an AI agent by learning from a stronger teacher, sometimes the student temporarily does better than the teacher, only to get worse later. The authors found out this happens because the training both fixes student mistakes and sometimes undoes good student actions the teacher doesn't like. They created a new method, SPOD, that protects the student's good decisions while still learning from the teacher. This helps the student keep and even improve on its early advantages and beat the teacher consistently.
What this means in practice
- •For reinforcement learning engineers: Improve training of AI agents in multi-step tasks by preserving early student improvements beyond the teacher’s behavior using SPOD.
- •For robotics development teams: Enhance robotic control policies that learn from expert demonstrations by maintaining beneficial actions during iterative policy refinement.
Authors
Jike Zhong, Ming Li, Yuxiang Lai, Xi Xiao, Sheng Zhang, Ritwick Chaudhry, Tianle Chen, Jingyu Xie, Zheda Mai, Shao-Yuan Lo, Konstantinos Psounis, Yuanyuan Lei
Abstract
On-policy distillation (OPD) trains agents on their own trajectories using supervision from a stronger teacher. We uncover a counterintuitive phenomenon in multi-turn agentic tasks: a smaller, initially weaker student can substantially outperform its teacher during training, only to regress as distillation continues. We term this phenomenon distillation overshoot. Through decision- and task-level analyses, we trace overshoot to two competing effects of OPD: repairing weak student decisions and eroding successful student choices that the teacher disfavors. Early in training, repair outpaces erosion, allowing the student to combine the teacher's corrections with its own strengths and surpass the teacher. As training continues, erosion overtakes repair, causing these gains to vanish. Motivated by this finding, we propose Success-Preserving On-Policy Distillation (SPOD), which selectively masks conflicting teacher supervision on successful student choices and introduces a soft probability floor to counteract their erosion, while retaining teacher guidance elsewhere. SPOD requires only binary episode success and model probabilities already available during OPD. Across two model families and three agentic benchmarks, SPOD consistently outperforms ordinary OPD and evaluated variants, improving final success by up to 15.2 percentage points and surpassing the teacher by up to 20.0 points. Moreover, SPOD's final checkpoints outperform even ordinary OPD's best intermediate checkpoints in all six settings, demonstrating that beyond-teacher gains can be not only preserved but further extended. Code: https://github.com/Jike338/SPOD