ATOD: Annealed Turn-aware On-policy Distillation for Multi-turn Autonomous Agents

2026-06-26

AI Summary

This work addresses the challenge that small language model agents struggle to simultaneously achieve rapid imitation learning and reward-driven continual optimization in long-horizon, multi-turn interactive tasks. To this end, the authors propose a hybrid online distillation algorithm that integrates early-stage teacher-guided policy distillation with late-stage progressive reinforcement learning. They further introduce an annealing-based OPD-RL scheduling mechanism and a turn-level divergence-uncertainty reweighting (T-DUR) strategy to enhance learning efficiency from high-value trajectories. Evaluated on ALFWorld, WebShop, and Search-QA, the method achieves an average success rate improvement of 3.03 points over OPD, 23.62 points over GRPO, and even surpasses the teacher model by 2.16 points.

Abstract

Training small language-model agents for long-horizon interactive tasks requires both fast imitation and reward-driven improvement. On-policy distillation (OPD) provides dense teacher guidance and typically improves rapidly in the early stage, but its gains saturate once the student approaches the teacher, limiting the final performance ceiling. Reinforcement learning (RL) directly optimizes environment rewards and encourages exploratory improvement toward a higher reward-defined ceiling, but sparse and delayed feedback makes early-stage learning much less efficient than OPD. In this paper, we propose ATOD (Annealed Turn-aware On-policy Distillation), a hybrid online distillation algorithm that explicitly exploits this complementarity. (1) ATOD uses an annealed OPD-RL schedule: OPD dominates early training to approach teacher-level behavior, while RL is gradually strengthened to drive reward-based exploration. (2) ATOD introduces Turn-level Disagreement-Uncertainty Reweighting (T-DUR), which softly amplifies high-utility turns and improves dense supervision in long trajectories. Experiments on ALFWorld, WebShop, and Search-QA show that ATOD consistently outperforms competing post-training baselines: across the three student sizes, ATOD improves average success rate by 3.03 points over OPD and 23.62 points over GRPO, while surpassing the corresponding teacher models by 2.16 points.

Problem

• on-policy distillation
• reinforcement learning
• multi-turn agents
• long-horizon tasks
• imitation learning

Innovation

• On-policy Distillation
• Reinforcement Learning
• Annealed Scheduling
• Turn-level Reweighting
• Multi-turn Agents