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

2026-07-06

AI Summary

This work addresses the inefficiency of conventional On-Policy Distillation (OPD) in long-horizon tasks, where full-trajectory replay wastes computation on information-sparse tail episodes, and token-level KL divergence overemphasizes shallow decisions while neglecting deeper reasoning steps. To overcome these limitations, the authors propose TurnOPD, which introduces, for the first time, a turn-aware adaptive replay mechanism that dynamically allocates replay depth based on probe-driven episode statistics. Additionally, they design a progressive turn-normalized KL loss to better balance computational resources with supervisory signals. Evaluated on ALFWorld, WebShop, and Multi-Hop Search benchmarks, TurnOPD significantly outperforms OPD under identical wall-clock time constraints, achieving higher validation accuracy and advancing the efficiency–performance trade-off frontier in long-horizon agent distillation.

Abstract

On-policy distillation (OPD) trains a student policy by matching a stronger teacher on the student's own trajectories, offering a promising framework for language agent training. However, its application to long-horizon agentic tasks remains insufficiently explored. We identify two key inefficiencies in vanilla agent OPD: (1) full-horizon rollouts often waste wall-clock resources on tail turns that provide weak and noisy KL supervision, and (2) trajectory-level KL objectives concentrate most of the loss on shallow tokens, leaving deeper decision turns under-trained once initial behaviors are aligned. To address these challenges, we propose TurnOPD, a turn-level budgeting strategy for efficient on-policy distillation of long-horizon agents. TurnOPD consists of two budget controllers: adaptive rollout-depth budgeting, which uses probe-based turn statistics to determine rollout length, and progressive turn-normalized loss budgeting, which gradually shifts KL weighting from token-level to turn-balanced supervision. Experiments on ALFWorld, WebShop, and Multi-Hop Search with task-specialized teacher models show that TurnOPD achieves superior validation accuracy under equal wall-clock training budgets and advances the accuracy--time frontier beyond vanilla OPD.

Problem

• on-policy distillation
• long-horizon agent
• turn-level efficiency
• KL supervision
• trajectory rollout

Innovation

• on-policy distillation
• turn-aware budgeting
• long-horizon agents
• adaptive rollout
• turn-normalized loss