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

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.

• The paper introduces a hybrid paradigm that combines on-policy distillation with reinforcement learning to overcome teacher performance ceilings and sample inefficiency.
• The paper leverages an annealed coefficient schedule and turn-level disagreement-uncertainty reweighting to dynamically allocate supervision in long-horizon tasks.
• The paper demonstrates that ATOD outperforms pure RL, direct OPD, and recent variants, achieving higher success rates and more efficient trajectories on multi-turn benchmarks.

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

Introduction and Motivation

ATOD introduces a hybrid policy optimization paradigm that targets the efficient training of small-scale LLM (LM) agents for multi-turn, long-horizon tasks. The method is constructed around the observation that on-policy distillation (OPD) and reinforcement learning (RL) — specifically, Group Relative Policy Optimization (GRPO) — have distinct advantages and deficiencies for agentic policy improvement. OPD enables rapid behavioral bootstrapping via dense, token-level teacher guidance but saturates at the teacher performance ceiling. RL, while providing direct reward-driven improvement beyond imitation, is stymied by high sample inefficiency and credit assignment issues in the initial phase, particularly for small students. ATOD synergistically leverages both paradigms through an annealed hybridization and a turn-level weighting mechanism.

Methodology

Hybrid Advantage with Annealing

ATOD integrates OPD and GRPO using a hybrid token-level advantage function, with schedule-controlled mixing coefficients. Early in training, OPD's coefficient $\kappa(s)$ is set high and decays according to training progress; RL's coefficient $\rho(s)$ starts low and increases, ensuring teacher-guided optimization at the outset and reward maximization late in training. A minimal non-zero $\kappa_{\min}$ is maintained to regularize policy drift and prevent reward hacking. This approach ensures initial sample efficiency without entrenching the policy below the teacher’s reward ceiling.

Figure 1: (a) The training schedule anneals from dense teacher supervision to reward-based exploration. (b) Turn-level T-DUR selectively amplifies supervision on key trajectory turns via disagreement and uncertainty proxies.

Turn-level Disagreement-Uncertainty Reweighting (T-DUR)

To further optimize dense teacher supervision in long agentic trajectories, ATOD introduces the Turn-level Disagreement-Uncertainty Reweighting (T-DUR) mechanism. T-DUR calculates, for each turn $k$, (1) a disagreement proxy $d_k$ — mean absolute logit divergence between teacher and student on sampled tokens; and (2) an uncertainty proxy $h_k$ — average student negative log-probability (entropy). Both are normalized per trajectory, and fused via a Soft-OR t-conorm to produce the turn weight $w_k = 1-(1-\tilde d_k)(1-\tilde h_k)$. Only the OPD component is reweighted, ensuring that teacher guidance is concentrated on semantically informative turns where supervision is most valuable, such as pivotal decisions and high-uncertainty states, while discounting routine steps with low learning potential.

Empirical Results

Main Comparative Results

ATOD was validated on three multi-turn benchmarks: ALFWorld (embodied interaction), WebShop (web shopping), and Search-QA (open-domain QA). Qwen3 models of varying sizes (0.6B, 1.7B, 4B) were employed as students, with strong Qwen3-4B or Qwen3-30B teachers.

Figure 2: Aggregate performance of the Qwen3-1.7B model across ALFWorld, Search-QA, and WebShop tasks.

ATOD consistently outperforms all baselines, including pure RL (GRPO), direct OPD, and recent agentic OPD variants (SOD, TCOD), and surpasses the teacher policy across most settings. For instance, the 0.6B student achieves an 82.8% ALFWorld success rate versus 30.5% for GRPO and 76.6% for OPD; for 1.7B and 4B students, ATOD also achieves the highest average success rates. Notably, ATOD exceeds the teacher baseline in both ALFWorld and WebShop for all student sizes, demonstrating efficient teacher imitation coupled with robust reward maximization.

Training Dynamics and Analysis

Figure 3: Training trajectories reveal that ATOD exhibits rapid early improvements, avoids the saturation of OPD, and transcends the slow reward accretion of GRPO.

Figure 4: On ALFWorld, ATOD outperforms in training reward, converges to a higher validation success rate, and yields more efficient policies (fewer average trajectory turns) relative to OPD and GRPO.

Early-stage dynamics confirm that ATOD matches the fast convergence of OPD, but after the OPD signal wanes, the RL signal continues to refine policy, pushing performance beyond the imitation ceiling. Trajectory length metrics further show that ATOD learns more efficient policies than pure GRPO, avoiding trajectory bloat typical of inefficient exploration.

Figure 5: Internal diagnostics show OPD signal decays while RL remains robust; mean turn weight evidences active reallocation of teacher supervision; teacher-student policy gap shrinks progressively.

Ablation Studies

Figure 6: Ablations in ALFWorld indicate that removing the turn-level T-DUR or annealing schedule impairs performance, confirming both mechanisms' necessity.

Results indicate both T-DUR and annealing are indispensable. Using token-level instead of turn-level T-DUR, or removing it entirely, degrades performance and stability. Omitting hybrid coefficient annealing (using fixed mixing) substantially hurts performance, particularly in small models, supporting the claim that static weighting cannot balance dense early supervision and late reward maximization.

Diagnostic and Interpretive Insights

T-DUR diagnostics suggest that the model gradually reallocates teacher supervision resources toward high-utility turns (e.g., error-prone states, exploration frontiers), as confirmed by turn weight and disagreement/uncertainty trajectories (see also Figure 7, not reproduced here). By dynamically adapting which trajectory segments receive teacher guidance, ATOD avoids distillation inefficiency and error compounding typical in long-horizon agent tasks.

Implications and Future Directions

ATOD offers a paradigm for training compact, resource-efficient agents that do not just imitate larger teachers but successfully maximize environment-defined returns beyond teacher-level competence. Practically, this method enables deployment of small footprint language agents in latency- and privacy-constrained applications without sacrificing multi-turn decision quality. Theoretically, hybrid schedules and turn-aware weighting suggest new research directions in fine-grained credit assignment, efficient knowledge transfer in sequential settings, and curriculum-based RL for LMs.

Potential future extensions include: (1) explicit modeling of state-action uncertainty for dynamic schedule adaptation; (2) meta-learned or task-adaptive annealing policies; (3) application to multimodal agent settings requiring cross-modal reasoning and tool use.

Conclusion

ATOD effectively resolves sample inefficiency in pure RL and performance plateauing in pure OPD by combining annealed hybrid optimization with turn-level supervision reallocation. The method achieves state-of-the-art small agent performance on multi-turn language tasks, surpassing strong teacher baselines and underscoring the efficacy of scheduled, turn-aware distillation for training autonomous language agents (2606.27814).