Skip to content

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

🕒 Published (v1): 2026-06-26 07:56 UTC · Source: Arxiv · link

Why this paper was selected

Annealed on-policy distillation for multi-turn agents; practical recipe for training small long-horizon agents

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

ATOD is a hybrid post-training algorithm for small language-model agents that combines on-policy distillation (OPD) and reinforcement learning (GRPO) via an annealed schedule and turn-level distillation reweighting. It addresses the complementary failure modes of pure OPD (teacher-ceiling plateau) and pure RL (cold-start inefficiency under sparse rewards). Across ALFWorld, WebShop, and Search-QA with Qwen3 students (0.6B–4B), ATOD outperforms all baselines and even exceeds its teacher models.

Problem

Training small LM agents for long-horizon interactive tasks faces two opposing failure modes: (1) on-policy distillation (OPD) bootstraps rapidly via dense teacher supervision but plateaus near the teacher ceiling, suppressing reward-improving deviations; (2) RL (GRPO) can improve beyond imitation but suffers slow cold-start learning under sparse, delayed trajectory-level rewards. A second issue is that uniform supervision across all turns in a long trajectory wastes signal on routine turns and dilutes supervision on high-utility decision points.

Method

ATOD combines OPD and GRPO into a single token-level hybrid advantage:

\[A_t = \kappa(s)\, A_t^{\text{OPD}} + \rho(s)\, A_t^{\text{GRPO}}\]

with \(A_t^{\text{OPD}} = (\log \pi_T(a_t|s_t) - \log \pi_\theta(a_t|s_t))\, w_{k(t)}\), where \(k(t)\) maps token \(t\) to its containing turn and \(w_{k(t)}\) is the T-DUR weight.

Annealed OPD–RL schedule: A progress variable \(p(s) = \min(s/T, 1)\) linearly increases \(\rho(s)\) and decreases \(\kappa(s)\) over a configurable annealing window, so teacher guidance dominates early and environment reward dominates late. A floor \(\kappa_{\min} > 0\) retains a weak teacher anchor to prevent reward hacking.

Turn-level Disagreement-Uncertainty Reweighting (T-DUR): For each turn \(k\), two statistics are computed from sampled tokens: disagreement \(d_k\) (mean teacher–student log-prob gap) and uncertainty \(h_k\) (mean negative log-prob under the student). Both are per-trajectory min-max normalized to \(\tilde{d}_k, \tilde{h}_k \in [0,1]\). The turn weight is fused via soft-OR:

\[w_k = 1 - (1-\tilde{d}_k)(1-\tilde{h}_k)\]

This is large when either signal is high, upweighting turns where the student is uncertain or confidently disagrees with the teacher. T-DUR is applied only to the OPD advantage term, leaving GRPO rewards unchanged. The actor is updated with the standard clipped GRPO surrogate using \(A_t\) as advantage; no separate KL penalty is added.

Key Contributions

  • Annealed OPD–RL objective that shifts from teacher-guided bootstrap to reward-driven exploration over training, encoded as a single hybrid token-level advantage.
  • T-DUR: turn-level (not token-level) soft reweighting of distillation based on disagreement and uncertainty, computed without full-vocabulary teacher sampling.
  • Empirical demonstration that a student can surpass its teacher model on multiple benchmarks using the hybrid training regime.

Results

  • Average success rate across 3 benchmarks: ATOD improves over the second-best baseline by +3.89 pp (0.6B), +5.07 pp (1.7B), +2.13 pp (4B).
  • vs. pure OPD: +3.03 pp average across all student sizes.
  • vs. pure GRPO: +23.62 pp average across all student sizes.
  • vs. teacher: ATOD surpasses its teacher model by +2.16 pp average; teacher-exceeding on ALFWorld and WebShop for all sizes.
  • ALFWorld (Qwen3-1.7B): ATOD 80.47% vs. OPD 72.66%, GRPO 31.25%, teacher 76.56%.
  • WebShop (Qwen3-1.7B): ATOD 89.06% vs. OPD 77.34%, teacher 82.03%.
  • 0.6B on ALFWorld: Vanilla 0.78% → ATOD 82.81% (>100× relative gain); GRPO only reaches 30.47%.
  • Ablation (ALFWorld): Removing annealing causes the largest drop (e.g., 0.6B: 82.8% → 75.8%); removing T-DUR or downgrading to token-level reweighting also degrades all sizes.
  • Training efficiency: ATOD produces shorter average trajectories than GRPO throughout training, indicating more efficient interaction policies.

Limitations

  • Evaluation confined to three benchmarks (ALFWorld, WebShop, Search-QA); no code-execution or GUI-based agent tasks.
  • Teacher must support token-level log-probability inference on student-sampled trajectories; requires a capable, accessible teacher at training time.
  • The annealing schedule introduces hyperparameters (\(\kappa_{\text{init}}, \kappa_{\min}, \rho_{\text{init}}, \rho_{\max}, T\)) that may need tuning per task or student size.
  • Search-QA gains over OPD are small (<1 pp), suggesting the method's advantage is stronger in embodied/web interaction than in search-augmented QA.
  • No analysis of wall-clock or compute cost relative to baselines beyond training-step count.

Relevance to Agentic AI / LLM Agents

ATOD directly targets the core bottleneck in small-agent deployment: how to efficiently transfer multi-turn agentic behavior from a large teacher while still enabling reward-driven improvement that exceeds pure imitation. The T-DUR mechanism advances turn-level credit assignment in long trajectories, a fundamental challenge in multi-step agent training where not all decision points are equally consequential. The result that a 0.6B student can exceed a 4B teacher model on ALFWorld is a practically significant data point for resource-constrained agent deployment. The paper also provides a clean formulation of the OPD/RL complementarity that applies broadly to any agentic post-training pipeline using GRPO-style reward optimization.