ATOD: Annealed Turn-aware On-policy Distillation for Multi-turn Autonomous Agents
06:00 · June 29, 2026 · arXiv cs.AI RSS
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.
Summary
ATOD is a hybrid training method that combines on-policy distillation with reinforcement learning to improve small language-model agents on long-horizon interactive tasks. The approach addresses the complementary limitations of its two components: on-policy distillation supplies dense token-level guidance from a teacher and accelerates early learning, yet performance plateaus once the student nears the teacher’s level; reinforcement learning optimizes directly for environment rewards and can exceed imitation, but sparse terminal signals slow progress in the initial stages.
The algorithm therefore applies an annealed schedule in which the distillation loss dominates early training and the reinforcement-learning term, implemented via group-relative policy optimization, is gradually strengthened. In parallel, ATOD introduces turn-level disagreement-uncertainty reweighting, which estimates the utility of each decision step from the student’s entropy and the divergence between teacher and student distributions. Higher weights are assigned to turns that are both uncertain and informative, concentrating supervision on decisive actions rather than routine ones within extended trajectories.
Evaluations on ALFWorld, WebShop, and Search-QA across three student sizes show consistent gains. On average, ATOD raises success rate by 3.03 points relative to pure on-policy distillation and by 23.62 points relative to group-relative policy optimization alone, while also exceeding the performance of the corresponding teacher models by 2.16 points. The results indicate that the combination of scheduled objectives and turn-aware weighting improves both sample efficiency and final reward ceiling for resource-constrained agents.
Why it matters
This research is highly relevant for Dutch AI researchers and developers focusing on efficient AI and autonomous agents. By enabling small language models to achieve teacher-level performance through a novel distillation and RL approach, it supports the development of cost-effective, high-performing AI solutions suitable for widespread SME adoption in the Netherlands.
