TurnOPD: Making On-Policy Distillation Turn-Aware for Efficient Long-Horizon Agent Training

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

Why this paper was selected

Turn-aware on-policy distillation for long-horizon agents; directly improves student training efficiency

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

TurnOPD addresses two inefficiencies in on-policy distillation (OPD) for long-horizon agent training: wasted compute on low-signal tail turns and loss budget concentration on shallow turns. It introduces two turn-level budget controllers—adaptive rollout-depth budgeting and progressive turn-normalized loss weighting—achieving up to 2.29× faster training with equal or better accuracy.

Problem

Applying OPD to long-horizon agent tasks suffers from two mismatches: (1) external mismatch—fixed full-horizon rollouts waste compute on tail turns where the survivor-weighted KL signal is sparse and noisy due to context-forced compression of policy disagreement; (2) internal mismatch—trajectory-level token normalization concentrates the KL loss on shallow turns (e.g., the deepest third of turns receives only 3.6–4.5% of total KL loss on ALFWorld), starving critical deep decision turns of gradient signal.

Method

TurnOPD introduces two budget controllers atop standard reverse-KL OPD:

1. Adaptive rollout-depth controller: Periodically probes turn-level statistics (survivor-weighted effective KL \(H_{\text{eff}}\) and a coverage lower-bound \(H_{\text{cov}}\)) to determine rollout horizon \(H = \max(H_{\text{eff}}, H_{\text{cov}})\), cutting collection early when the signal-to-noise ratio of deep turns falls below threshold.

2. Progressive turn-normalized loss controller: Linearly interpolates between trajectory-level and turn-balanced KL aggregation over training: $\(\mathcal{L}_{\text{mix}} = (1-\alpha)\mathcal{L}_{\text{token}} + \alpha \mathcal{L}_{\text{round}},\)$ where \(\alpha\) increases from 0 to 1 across training progress, shifting weight toward deeper turns only after shallow turns are already aligned.

The diagnosis motivates these choices via a contamination-compression theorem (Proposition 1): for context-forced mass \(\lambda(c)\), observable KL is bounded as \(D_{\text{KL}}(\pi_S \| \pi_T) \leq (1-\lambda(c))\,\Delta^{\text{free}}(c)\), explaining why raw KL at deep turns underestimates true policy disagreement on student-generated trajectories.

Key Contributions

• Turn-level empirical diagnosis of KL signal distribution and loss-budget allocation in vanilla OPD across ALFWorld and Multi-Hop Search.
• Formal contamination-compression mechanism (Proposition 1) explaining why observable KL shrinks at depth and why the success/failure KL gap can invert on failed rollouts.
• TurnOPD: adaptive rollout-depth budgeting + progressive turn-normalized loss budgeting.
• Demonstration of an accuracy–time Pareto improvement over vanilla OPD on three multi-turn agent benchmarks.

Results

• ALFWorld-1.7B: Same-Step Avg@4 improves from 83.0 → 86.3; wall time for 100 steps cut from 4.42h → 1.93h (~2.29× speedup).
• WebShop: Both accuracy and wall time improve (1.57h → 1.24h).
• Multi-Hop Search: Improves minimal-time Avg@4 across all tested models.
• Ablations show adaptive depth is the dominant efficiency lever (nearly halves wall time on ALFWorld-1.7B alone), while linear KL blending provides complementary accuracy gains by reallocating supervision to deep turns.
• Deep third of turns receives 3.6–4.5% of KL loss under vanilla OPD on ALFWorld; TurnOPD corrects this allocation.

Limitations

• Evaluated on three benchmarks (ALFWorld, WebShop, Multi-Hop Search); generalization to GUI/computer-use tasks or code agents is not demonstrated.
• The adaptive depth controller requires periodic probe rollouts, introducing an implementation overhead not fully quantified.
• The coverage floor hyperparameter must be tuned; the paper notes the controller can be swept between conservative and aggressive horizons but no automatic schedule is provided.
• The contamination-compression bound (Eq. 6) is an upper bound; it does not directly quantify how much genuine policy disagreement is masked in practice.
• Paper text is truncated; full ablation tables and all hyperparameter sensitivity results may not be reflected here.

Relevance to Agentic AI / LLM Agents

TurnOPD directly addresses the training efficiency bottleneck in multi-turn LLM agents—an increasingly central problem as agent rollouts grow in length and cost. The contamination-compression analysis provides a principled theoretical account of why standard KL-based distillation degrades over long interaction traces, which is relevant to any on-policy training method (RLVR, OPD, GRPO) applied to agents. The turn-level budgeting framework is task-agnostic and complements prior step-reweighting work (SOD, TCOD), offering a practical recipe for scaling agent distillation without proportional compute increases. This work is particularly relevant for researchers building smaller student agents from stronger teachers in settings like tool-use, embodied planning, and multi-hop reasoning.