ATOD: Annealed Turn-Aware On-Policy Distillation for Multi-Turn Agentic Tasks

Recent advances in large language models (LLMs) have enabled impressive static text generation, but their application to multi-turn interactive tasks remains limited. Early works focused on supervised learning and knowledge distillation, such as teacher-student frameworks, to improve training efficiency. Reinforcement learning (RL), especially policy optimization algorithms like PPO and GRPO, was introduced to directly maximize environment rewards, enabling models to surpass imitation limits. OPD emerged as a promising hybrid, leveraging teacher token distributions for dense supervision, improving sample efficiency. However, these methods face challenges in balancing rapid initial learning with long-term exploration, especially in long-horizon tasks with sparse rewards. Recent research emphasizes dynamic, fine-grained control mechanisms, but comprehensive solutions remain elusive.

The core challenge lies in effectively combining imitation and exploration in multi-turn agent training. Pure distillation accelerates early learning but saturates near teacher performance, limiting potential. Pure RL can surpass this ceiling but struggles with sparse, delayed rewards, leading to inefficient exploration, especially for small models. Additionally, identifying critical decision points within long trajectories is difficult; uniform supervision often wastes resources on routine turns, while key turns may be underweighted. These issues hinder the development of small, efficient models capable of complex multi-turn reasoning, restricting their deployment in real-world scenarios.

The paper proposes a hybrid approach with two main innovations. First, a linear annealed schedule smoothly transitions the training focus from teacher-guided distillation to reward-driven exploration, balancing early rapid learning with later performance gains. Second, T-DUR introduces turn-level disagreement and uncertainty metrics, which are combined via a soft OR function to dynamically weight teacher supervision. This mechanism emphasizes turns with high divergence or low confidence, ensuring the model concentrates on critical decision points. Together, these innovations enable more efficient training, better utilization of supervision signals, and improved performance ceiling, especially for small models in complex tasks.

• �� Design a linear schedule: define κ(s) and ρ(s) to gradually shift emphasis from OPD to RL over training steps. • Compute turn-level disagreement (dk): token-wise log probability differences between teacher and student, normalized within each trajectory. • Calculate turn-level uncertainty (hk): negative log probability of student tokens, normalized similarly. • Fuse signals: combine dk and hk using a soft OR to produce turn weights wk, emphasizing turns with high divergence or uncertainty. • Formulate advantage function: At = κ(s)AOPDt + ρ(s)AGRPOt, where AOPDt is token-level distillation advantage, and AGRPOt is RL advantage. • Optimize policy: use PPO-style clipped surrogate with advantage At, updating the model iteratively. • Implement annealing: linearly adjust κ(s) and ρ(s) based on training progress p(s), ensuring smooth transition from imitation to exploration.

The evaluation involves three datasets: ALFWorld (embodied instruction), WebShop (web navigation), and Search-QA (multi-hop QA). Models include Qwen3-0.6B, 1.7B, and 4B, trained with a 150-step window. Baselines include vanilla, GRPO, SDAR, OPD, SOD, and TCOD. Success rate and trajectory length are primary metrics. Ablation studies test the impact of scheduling, T-DUR, and reweighting strategies. Results show ATOD consistently outperforms baselines across all sizes and tasks, with significant gains especially for smaller models. Sensitivity analyses confirm the importance of the annealing schedule and turn-level reweighting.

ATOD achieves success rates of 82.81% (ALFWorld, 0.6B), surpassing OPD by 6.25 points and teacher models by 2.16 points on average. For larger models, success rates reach over 85%, with notable improvements over pure distillation and RL. Ablation results indicate removing scheduling or T-DUR reduces success by 4-6 points, confirming their critical role. The method effectively balances imitation and exploration, leading to faster convergence and higher performance ceilings. The success in diverse tasks demonstrates robustness and scalability.

This approach can be directly applied to training efficient multi-turn dialogue agents, autonomous systems, and interactive AI assistants. It requires environment feedback and a teacher model, making it suitable for resource-limited settings. Long-term, the framework can extend to multi-modal tasks, integrating vision and speech, to develop more versatile intelligent agents. Its ability to improve small models opens pathways for deploying AI in embedded devices, customer service bots, and educational tools, democratizing access to advanced AI capabilities.

The method relies on manually tuned scheduling parameters, which may not generalize across tasks. In highly complex or noisy environments, turn-level signals might be insufficient to identify critical turns. Computational costs increase with model size and interaction length, limiting scalability. Future work should focus on adaptive parameter tuning, more sophisticated uncertainty estimation, and reducing training overhead to broaden applicability.