A Few Teacher Steps Go a Long Way: Cost-Efficient On-Policy Data Augmentation for Agent Post-Training¶

🕒 Published (v1): 2026-07-06 01:02 UTC · Source: Arxiv · Venue: ICML 2026 · link

Why this paper was selected

ICML 2026; on-policy data augmentation via teacher steps; practical post-training recipe for agents

Ask a follow-up

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

TL;DR¶

This paper frames on-policy SFT data construction for LLM agents as a budget-allocation problem, separating teacher inference cost \(C_i\) from retained training cost \(C_{tr}\). The key finding is that short, unfiltered teacher continuations (OP-SHORT, capped at \(K\) turns) placed at student-induced contexts outperform both pure behavioral cloning and heavily filtered alternatives at matched budgets. A few teacher steps at learner-visited states are consistently the most cost-efficient form of additional supervision.

Problem¶

Pure behavioral cloning (pure BC) trains on teacher-induced contexts, but at test time the student encounters its own self-induced contexts — a compounding distribution mismatch in multi-turn agents. Prior DAgger-style remedies query the teacher at student-visited states but leave open the central question: how to allocate a fixed supervision budget across more end-to-end demos, longer continuations, outcome filtering, or broader coverage of learner states. Existing work conflates these choices and uses incommensurable units (e.g., "one training sample"), making fair comparison impossible.

Method¶

The authors define a unified data-construction recipe \(G = (\rho, \pi_e, n, \Psi, T_{\text{switch}}, K, \Phi)\), where \(\rho\) is the prefix-context rollout policy, \(\pi_e\) the teacher, \(n\) proposals per task, \(\Psi\) a prefix-context filter, \(T_{\text{switch}}\) a switch-time sampler, \(K\) the teacher continuation horizon (in assistant turns), and \(\Phi\) a post-continuation filter. This parametrizes the induced context distribution as:

\[d^{\rho;w}_\mu(B) = \frac{\mathbb{E}_{\tau \sim P^\rho_\mu}\left[\sum_{t=1}^{T(\tau)} w_t(\tau) \cdot \mathbf{1}\{H_{t-1}(\tau) \in B\}\right]}{\mathbb{E}_{\tau \sim P^\rho_\mu}\left[\sum_{t=1}^{T(\tau)} w_t(\tau)\right]}\]

Two averaging conventions are formalized: pooled-turn (\(w_t = 1\), longer trajectories contribute more mass) and trajectory-uniform (\(w_t = 1/T(\tau)\), each episode gets equal mass). The main variant OP-SHORT uses \(\rho = \pi_1\) (a BC-trained student checkpoint), trajectory-uniform switch sampling, no filters (\(\Psi = \Phi = 1\)), and caps continuations at \(K\) teacher turns. The BC+OP recipe trains the final model on \(D_{\text{BC}} \cup D_{\text{OP}}\) reinitializing from the base checkpoint. Comparisons are made at matched \(C_i\) and \(C_{tr}\) against: PURE-BC (scaled teacher demos), OP-SUCCESS (successful-continuation filter), and OP-CRITICAL (prefix filter restricting to failed tasks). Experiments use Qwen2.5-3B-Instruct as the student (\(\pi_0\)) and Qwen3-235B-A22B-Instruct as the teacher on HotpotQA, ALFWorld, and Terminal-Bench-Dev.

Key Contributions¶

Reframes SFT dataset construction as a budget-allocation problem with two orthogonal cost axes, \(C_i\) (teacher inference before filtering) and \(C_{tr}\) (accepted training tokens per epoch), enabling fair cross-protocol comparison.
Provides a formal population-level framework making explicit the context-generation policy, switch-time distribution, continuation horizon, filtering rules, and resulting context distribution for variable-length agentic rollouts.
Empirically demonstrates across three benchmarks that bounded unfiltered teacher continuations at learner-induced contexts are consistently more cost-effective than longer or more heavily curated alternatives at matched budgets.
Identifies a practical heuristic: use pure BC to bootstrap a student, then spend additional budget on short on-policy teacher continuations before committing to full rollouts, success filtering, or RL.

Results¶

HotpotQA: OP-SHORT improves over PURE-BC at matched \(C_i\) and \(C_{tr}\); unfiltered few-step continuations match or exceed OP-SUCCESS and OP-CRITICAL depending on cost axis.
ALFWorld: OP-SHORT at \(K=3\) exceeds pure BC at all four tested training-instance budgets, with both in-distribution and out-of-distribution evaluation.
Terminal-Bench-Dev: OP-SHORT with a small fraction of the SFT corpus and no RL stage matches the OpenThoughts-Agent baseline that uses the full corpus followed by reinforcement learning.
Additional teacher turns do not monotonically improve performance; best \(K\) is domain-dependent (larger on ALFWorld and Terminal-Bench-Dev than on HotpotQA).

Limitations¶

Single-round augmentation only; no iterative DAgger-style updates studied.
Optimal \(K\) is domain-dependent with no principled selection rule provided.
Student and teacher are both from the Qwen family; cross-architecture generalization is untested.
Full experimental figures are not shown in the provided text excerpt; some numeric comparisons are described qualitatively.
\(|Q_{\text{train}}|\) is held fixed as a hyperparameter, so interactions between task diversity and on-policy augmentation are not studied.
Accepted at an ICML 2026 workshop (not the main track), suggesting findings are preliminary.

Relevance to Agentic AI / LLM Agents¶

This paper directly addresses one of the most practically important questions in LLM agent training: how to efficiently spend a teacher-model budget when the teacher is large and expensive (e.g., a frontier API call). The formal decomposition into \(C_i\) vs. \(C_{tr}\) gives agent practitioners a concrete vocabulary for comparing data recipes that previously appeared incommensurable. The result that short, unfiltered on-policy corrections outperform elaborate filtering pipelines challenges the common assumption that higher curation always pays off, and the Terminal-Bench-Dev finding — matching RL-trained baselines with just a fraction of the SFT data — is directly relevant to teams deciding when to invest in RL post-training versus cheaper SFT augmentation.