Source-linked AI summary

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

Yuhang Zhou, Kai Zheng, Haoling Li, Dengyun Peng, Can Xu, Jingjing Chen

arXiv:2607.05804v1cs.AIcs.CL

TL;DR

Long-horizon agent OPD must allocate supervision across turns efficiently because shallow-token feedback can undertrain later decisions and full rollouts can waste computation. TurnOPD adapts rollout depth and progressively balances KL loss across turns, advancing the accuracy–time frontier with up to 2.29× faster training.

  • Problem

    Long-horizon agent OPD lacks effective supervision allocation across turns, where shallow-token feedback may under-supervise crucial later decisions and full rollouts may waste computation.

  • Method

    TurnOPD combines adaptive rollout-depth budgeting with progressive turn-normalized loss budgeting to select rollout lengths and increase weighting of deeper turns.

  • Results

    TurnOPD advances the accuracy–time frontier across ALFWorld, WebShop, and Multi-Hop Search, with up to 2.29× faster training.

  • Takeaways & Limitations

    Long-horizon agent OPD benefits from treating supervision as turn-conditioned decisions rather than flat token positions.

  • Takeaways & Limitations

    Equalizing turn weights from the beginning may amplify weaker or lower-support supervision estimates too early.

Abstract

from arXiv · show

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.

1 Introduction

TurnOPD diagnoses turn-dependent inefficiencies in long-horizon on-policy distillation: fixed rollout depths waste computation, while trajectory-level KL allocation under-trains deeper decisions. It addresses both mismatches with adaptive rollout-depth and progressive turn-normalized loss budgeting, improving the accuracy–time frontier across agent benchmarks.

  • Motivation: Long-horizon agent rollouts involve multiple turns, tool calls, environmental shifts, and state changes, making token-level feedback insufficient for crucial later decisions.Early decisions affect subsequent states, while later turns contain key but infrequent decisions.
  • Diagnosis: Only 3.6–4.5% of KL loss reaches the deepest third of ALFWorld turns, compared with 11–13% on Multi-Hop Search.Shallow turns dominate because they are overrepresented and carry higher KL.
  • Diagnosis: Fixed rollout depth wastes compute on low-signal tail turns, while trajectory-level normalization concentrates KL on easy shallow turns and starves deeper informative ones.These external and internal mismatches arise because correction signal and survivor count vary by turn, and uniform token weights favor shallow decisions.
  • Method: TurnOPD combines adaptive rollout-depth budgeting via periodic probes with progressive turn-normalized loss budgeting that increases deep-turn weighting.The rollout controller uses survivor-weighted KL and coverage thresholds; the loss controller transitions from trajectory-level to turn-balanced aggregation.
  • Results: 83.0 to 86.3 Same-Step Avg@4 and 4.42h to 1.93h 100-step wall time are achieved by TurnOPD on ALFWorld-1.7B.The introduction reports improvements from 83.0 to 86.3 and a reduction in 100-step wall time from 4.42h to 1.93h.
  • Results: Up to 2.29× faster training is achieved while maintaining comparable or better performance, advancing the accuracy–time frontier on tested benchmarks.TurnOPD improves minimal-time Avg@4 across all evaluated tasks and models.

2 Related Work

Prior work established on-policy distillation recipes and later framed OPD as KL-constrained policy optimization. Long-horizon agent research has progressed from controlled interaction environments toward realistic web, GUI, and practical software-engineering tasks.

  • On-policy distillation: MiniLLM introduced reverse-KL training on student-generated samples, while GKD combined on- and off-policy data and explored multiple divergence objectives.
  • On-policy distillation: Later OPD work interpreted the teacher–student log-ratio as a token-level reward within KL-constrained policy optimization.
  • Long-horizon agent tasks: ALFWorld connects text planning with embodied environments, whereas WebShop evaluates grounded e-commerce navigation.
  • Long-horizon agent tasks: Mind2Web, WebArena, OSWorld, and AndroidWorld broaden agent evaluation to realistic websites and desktop or mobile interfaces.

3 Preliminaries: On-Policy Distillation

The framework models agents as multi-turn systems whose actions alter future observations and contexts. On-policy distillation trains a student on its own rollouts using token-level supervision from a frozen teacher, with reverse KL as the objective.

  • Multi-turn agent interaction: At turn t, the student conditions on the full interaction history h_t before sampling response r_t from π_θ.The history includes the prompt, observations, and prior responses or rewards.
  • Multi-turn agent interaction: The environment maps executable response content to the next observation, producing a rollout that ends at task completion, a terminal state, or the maximum horizon.A rollout is represented as τ = (x, o_1, r_1, o_2, r_2, . . . , o_T, r_T).
  • Multi-turn agent interaction: Because early student actions change later observations, they also change the future contexts on which the teacher is queried.This turn dependence motivates treating interaction structure explicitly in agent distillation.
  • Multi-turn OPD: In multi-turn OPD, the student samples an on-policy rollout while a frozen teacher supplies token-level supervision on student-visited prefixes.Each prefix includes previous observations, previous student responses, and the current partial response.
  • Multi-turn OPD: The OPD objective uses reverse KL over response tokens selected by a response-token mask.The reverse-KL formulation is introduced for token positions in the concatenated model responses.

4 Diagnosis: Signal Structure in Agent OPD

The diagnosis shows that vanilla agent OPD’s reverse-KL supervision is turn-dependent, front-loaded, and increasingly unreliable as a measure of deep-turn policy disagreement. Context-induced forced tokens compress late-turn KL, while shallow token counts dominate the loss budget.

  • Turn-resolved signal: Reverse-KL signal is non-uniform and non-stationary across turns and training, with ALFWorld and Multi-Hop Search exhibiting task-dependent teacher-entropy profiles.ALFWorld often has lower teacher entropy at deeper turns, whereas Multi-Hop Search can have higher later-turn entropy despite front-loaded reverse-KL curves.
  • Outcome separation: Deep-turn KL lacks reliable outcome-predictive power, and ALFWorld’s separation is mostly negative, decreasing with depth as successful trajectories often have higher KL than failed ones.This contradicts the interpretation of KL as local correction demand, under which failed rollouts would be expected to require larger teacher corrections.
  • Contamination compression: Context-induced forced mass can compress observable KL even when free-component student–teacher disagreement remains nonzero, making raw KL an ambiguous proxy for policy mismatch.Longer student-generated contexts can increase surface-determined continuations, suppressing measured disagreement in later turns.
  • Loss-budget concentration: In ALFWorld-4B, turn 0 uses about a quarter of the KL budget, the first three turns nearly half, and the reliable deep third only 3.6–4.5%.Multi-Hop Search likewise assigns 38–40% of the budget to its first three turns, while its reliable deep third receives only 11–13%.
  • Loss-budget concentration: Equal turn weighting would remove shallow token-count bias but could amplify weak estimates: ALFWorld-4B’s reliable deep third has only 31–42% of shallow-prefix raw KL, while Multi-Hop Search has 16–17% survivor support.The diagnosis therefore identifies reliability as a constraint on aggressive turn-level reweighting.

5 Method: TurnOPD

TurnOPD addresses two inefficiencies in long-horizon OPD with complementary controllers: adaptive rollout-depth budgeting and progressive shift from trajectory-level toward turn-balanced KL supervision. Together, they regulate interaction collection and supervision allocation over collected tokens.

  • Overview: TurnOPD uses two budget controllers: one adaptively budgets rollout depth, while the other progressively shifts KL allocation toward turn-balanced decision refinement.The controllers target complementary resources: how much interaction is collected and how supervision is distributed.
  • Adaptive rollout-depth budgeting: The rollout-depth controller estimates an unobservable optimal horizon by combining efficiency-side effective supervision with a coverage-based lower bound.The target horizon balances supervision value per unit cost against task-completion coverage.
  • Adaptive rollout-depth budgeting: Effective supervision is estimated from survivor-weighted positive reverse-KL mass across turns, whose centroid remains shallow or deepens when meaningful correction tails persist.Relative weighting makes the statistic robust to changes in absolute KL magnitude.
  • Adaptive rollout-depth budgeting: The coverage lower bound uses the success-conditioned completion-depth p-quantile to ensure rollouts include at least 80% of successful trajectories.Periodic probe rollouts estimate this completion-depth distribution.
  • Adaptive rollout-depth budgeting: Periodic full-length probe rollouts update uncensored depth statistics, while routine truncated rollouts use the current cap and still contribute to OPD updates.Excluding truncated rollouts from depth estimation avoids bias from unobserved later turns.
  • Progressive turn-normalized loss budgeting: A linear blend starts with trajectory-level token-mass weighting and progressively shifts toward uniform turn-wise weighting, improving supervision of deeper decisions later in training.This directly addresses shallow-turn loss concentration under standard trajectory normalization.

6 Experiments

TurnOPD is evaluated across three long-horizon agent tasks and multiple student–teacher pairs, consistently improving the accuracy–time tradeoff over vanilla OPD and TCOD-F2B. It achieves strong accuracy under both wall-clock and fixed-step evaluations while reducing training time.

  • Experimental Setup: Experiments span ALFWorld, Multi-Hop Search, and WebShop, using multiple student–teacher pairs to test TurnOPD across environments and models.The primary methods are vanilla OPD, TCOD-F2B, and TurnOPD, evaluated under Least-Time and Same-Step regimes.
  • Accuracy: Under Least-Time, TurnOPD achieves the highest overall avg@4 across task–model combinations and matches or surpasses the best baseline in most Same-Step cases.On ALFWorld with a Qwen3-1.7B student, TurnOPD reaches 86.29 avg@4 versus 83.00 for vanilla OPD and 80.06 for TCOD-F2B.
  • Accuracy: With the Qwen3-4B student, TurnOPD exceeds the ALFWorld teacher reference in overall avg@4.This result is reported alongside TurnOPD’s best overall Least-Time score on ALFWorld.
  • Training Efficiency: TurnOPD reduces 100-step wall-clock time on ALFWorld-1.7B from 4.42 hours for vanilla OPD to 1.93 hours while improving accuracy.Comparable reductions occur on Multi-Hop Search, 2.94h versus 4.45h, and WebShop, 1.24h versus 1.57h.
  • Controller Validation: The main results confirm that TurnOPD improves both accuracy and compute-efficiency frontiers, motivating further validation of its adaptive controller against diagnostic metrics.Figure 5 likewise reports improved accuracy–time frontiers on each task type.
  • Controller Validation: TurnOPD’s rollout-depth controller targets the maximum of the survivor-weighted raw-KL centroid Heff and success-conditioned completion quantile Hcov at each step.During warm-up, Hcov is set to zero before sufficient successful rollouts appear.

7 Ablation Studies and Analysis

The ablations show that TurnOPD’s two controllers have complementary roles: adaptive rollout depth improves efficiency, while progressive turn-normalized loss improves optimization and intermediate progress. KL normalization and rollout-depth hyperparameters further reveal tradeoffs between deep-turn supervision, accuracy, and training time.

  • Controller decomposition: Adaptive depth alone cuts 100-step wall time from 4.42 h to 1.96 h, while accuracy drops from 83.0 to 82.8.The result isolates rollout-depth budgeting as TurnOPD’s primary efficiency lever.
  • Controller decomposition: TurnOPD’s combined controllers deliver faster progress and better intermediate results, despite alternatives sometimes matching or slightly exceeding its final performance.Adaptive depth provides the efficiency lever, progressive loss budgeting provides the optimization lever, and their combination gives the best accuracy–time tradeoff.
  • KL normalization: Trajectory-level KL assigns only 3.2/0.7/1.2% of early/mid/late budget to the deepest third, whereas hard turn-level KL abruptly assigns around one-third.The hard strategy may overweight unreliable deep-turn estimates early, motivating the smoother linear blend.
  • KL normalization: The linear blend increases α from 0.17 to 0.83 across early/mid/late training, smoothly increasing deep-turn budget instead of making an abrupt allocation shift.Here α = 0 corresponds to trajectory-level KL and α = 1 to hard turn-level KL.
  • Coverage-floor sensitivity: With the full-population CDF, p = 0.6 reaches 85.1 accuracy in 1.66 hours, while p = 0.8 reaches 85.8 and requires 1.83 hours.Increasing p targets deeper rollout horizons, slightly improving final success rates while increasing training time.
  • Coverage-floor sensitivity: The full-population CDF consistently increases ema-H and wall time because failed trajectories can stall at maximum depth, overestimating the required rollout horizon.It remains effective but requires lower p values to match the success-conditioned CDF’s efficiency and performance.

8 Conclusion

TurnOPD addresses two allocation mismatches in vanilla on-policy distillation for long-horizon language agents by adapting rollout depth and progressively shifting KL normalization toward turn-balanced supervision. It is evaluated across ALFWorld, WebShop, and Multi-Hop Search.

  • Diagnosis: Vanilla OPD wastes compute on low-yield tail turns and concentrates trajectory-level KL loss on shallow tokens.These are the two allocation mismatches identified for long-horizon language agent tasks.
  • TurnOPD: TurnOPD uses turn-level budgeting to adapt rollout depth and progressively shift KL normalization toward turn-balanced supervision.The method targets both rollout allocation and loss allocation mismatches.
  • Evaluation: TurnOPD is evaluated across ALFWorld, WebShop, and Multi-Hop Search.The supplied passage names these three long-horizon language agent benchmarks.

A Environments, Tool Configurations, And Task Examples … A.2 Task-Specialized Teacher Training

The paper evaluates agents across ALFWorld, Multi-Hop Search, and WebShop using tool-interactive turns, task-specific environment protocols, and specialized GRPO teachers. These setups span household planning, retrieval-based question answering, and online shopping, with task-dependent action spaces and reward definitions.

  • A Environments, Tool Configurations, And Task Examples: Experiments cover ALFWorld, Multi-Hop Search, and WebShop as the paper’s three environments and task types.Multi-Hop Search includes PopQA, NQ, 2WikiMultiHopQA, and HotpotQA.
  • A Environments, Tool Configurations, And Task Examples: Each model turn combines chain-of-thought with a tool call, while client-side validation blocks malformed or inadmissible actions before environment execution.Experiments run on 32 NVIDIA H20 GPUs with 80 GB of memory each.
  • A.1.1 ALFWorld (Embodied Household Planning): ALFWorld uses the alfworld_step tool for natural-language household commands across six task types, with binary reward on the valid_unseen split.The TextWorld-backed environment returns observations and refreshed admissible commands, rejecting invalid actions client-side.
  • A.1.2 Multi-Hop Search: Multi-Hop Search uses a Search-R1-style retrieval tool over a fixed Wikipedia corpus, supports up to 50 turns, and scores extracted answers by exact match.The interface separates reasoning, search, returned information, and final answer channels.
  • A.1.2 Multi-Hop Search: The Multi-Hop Search test mixture combines direct open-domain retrieval from PopQA and NQ with chained retrieval dependencies from 2WikiMultiHopQA and HotpotQA.This mixture tests both direct retrieval and long-horizon dependency cases.
  • A.1.3 WebShop (Web Navigation / Online Shopping): WebShop agents use search and click actions in an episode progressing from product search through inspection, attribute selection, and Buy Now termination.The dense terminal reward lies in [0, 1] and provides partial credit for matching product type, attributes, options, and price.
  • A.2 Task-Specialized Teacher Training: Task-specialized teachers are trained with GRPO on all three environments and evaluated using avg@4 validation accuracy.ALFWorld and WebShop use Qwen3-8B-GRPO teachers, whereas Multi-Hop Search uses a Qwen3.5-9B-GRPO teacher.

B Proof Of The Contamination-Compression Bound · C Existence of a Latent Rollout-Depth Optimum

Section B proves the contamination-compression bound from joint convexity of KL divergence and a specific parameter substitution. Section C characterizes the useful rollout horizon as a two-sided efficiency–coverage target bounded by signal exhaustion and successful-completion coverage.

  • B Proof Of The Contamination-Compression Bound: Joint convexity of KL divergence yields the mixture inequality in Equation 17.The proof applies the log-sum inequality tokenwise and sums over the vocabulary.
  • B Proof Of The Contamination-Compression Bound: Substituting α = λ, P1 = Q1 = pF, and P2 = pfree into the mixture inequality proves Equation 6.
  • C Existence of a Latent Rollout-Depth Optimum: The useful rollout horizon balances coverage of successful trajectories’ decision turns against collection cost after teacher-correction signal decays.This defines a two-sided efficiency–coverage sandwich rather than a horizon determined by a single force.
  • C Existence of a Latent Rollout-Depth Optimum: Adding turn H + 1 improves efficiency only when its marginal rate is at least the current average.The condition follows by cross-multiplying the ratio comparison for positive denominators and increments.
  • C Existence of a Latent Rollout-Depth Optimum: Beyond a signal-exhaustion frontier τc, zero marginal teacher-correction value or below-average marginal rates decrease the efficiency ratio.
  • C Existence of a Latent Rollout-Depth Optimum: The coverage floor Hcov ensures at least a fraction p of successful trajectories have completed by the rollout horizon.Stopping before Hcov misses completion decisions from more than a (1 −p) fraction of successful trajectories.
  • C Existence of a Latent Rollout-Depth Optimum: For finite empirical support, full-coverage completion depth equals the largest observed successful completion length, max supp(Lsucc).The definition is the first point that has already covered all successful completions, not the last point before full coverage.
  • C Existence of a Latent Rollout-Depth Optimum: The latent target horizon must be deep enough to cover successful completions but not extend beyond signal exhaustion and the deepest successful completion.

D Complete TurnOPD Hyperparameters

This section specifies TurnOPD’s implementation configuration and training algorithm. The procedure uses bounded rollout-depth control, periodic probe batches, teacher reverse-KL losses, and a progressive blend of trajectory- and turn-normalized weights.

  • Configuration: TurnOPD is configured through implementation keys mapped to formal-model symbols, with ema_alpha denoting the EMA weight on Hctrl rather than the loss-blend coefficient α.Table 7 gives the full main-method configuration for the ALFWorld-1.7B reference run.
  • Training algorithm: The algorithm takes πθ, πT, task distribution D, rollout bounds Hmin and Hmax, training horizon K, and controller hyperparameters as inputs, and returns an updated student policy πθ.The controller hyperparameters include min and (s, e) in the algorithm specification.
  • Depth control: The rollout-depth controller initializes at Hmax, clips the next depth between Hmin and Hmax, and activates probing during warmup or every rprobe iterations.The clipped estimate is ˆHk ← clip(round(¯Hk−1) + 1, Hmin, Hmax).
  • Training algorithm: Each iteration collects on-policy trajectories up to Hroll turns, queries πT on supervised tokens, and computes top-K reverse-KL token losses ℓi.The rollout batch is denoted Bk and is collected from πθ on D.
  • Depth control: Probe batches compute per-turn raw-KL means and survivor counts, then estimate successful-trajectory coverage depth Hcov when the successful probe batch reaches the minimum threshold.The estimated depth satisfies Fsucc(H) ≥ p.
  • Loss budgeting: The loss budget increases αk from 0 to 1 across normalized training progress using endpoints s and e, blending trajectory-normalized and turn-normalized weights before forming Lk.The blend is wblend = (1 − αk)wtraj + αkwturn.

E OPD Data Scale And Train/Test Construction

OPD uses online-generated trajectories from task-instance prompts rather than stored teacher trajectories, with task-specific dataset scales and split-construction procedures. ALFWorld and WebShop enforce non-overlapping environment or goal partitions, while Multi-Hop Search uses externally prepared questions and a separate held-out test set.

  • Online data generation: OPD samples are task instances, while trajectories are generated online during training against the live environment.A sample is a goal or question, not a fixed teacher trajectory.
  • ALFWorld: ALFWorld synthesizes reproducible, non-overlapping train and validation prompts by seeded sampling across six task types and object–receptacle vocabularies.Training instances use train games, whereas evaluation instances are routed to a separate environment partition.
  • Multi-Hop Search: Multi-Hop Search trains on ≈30k externally prepared multi-hop and open-domain question–answer instances and evaluates on a separate held-out test set of ≈400 questions.The evaluation questions come from PopQA, Natural Questions (NQ), 2WikiMultiHopQA, and HotpotQA, with source metadata preserved for reward routing and breakdowns.
  • WebShop: WebShop partitions data by strictly disjoint goal-id ranges, assigning evaluation ids [0, Neval) and training ids [Neval, Neval + Ntrain).The reference configuration uses ≈3000 training instances, while the supplied passage truncates the evaluation count.
Loading 2607.05804v1…