A Few Teacher Steps Go a Long Way: Cost-Efficient On-Policy Data Augmentation for Agent Post-Training
cs.LG, cs.AI
Submitted: 2026-07-06
Updated: 2026-08-29
Comments: 8/28 update: added funding acknowledgements
License: http://creativecommons.org/licenses/by/4.0/
The gist: For LLM agents, supervised fine-tuning is not only about teacher labels' quality, but also about which interaction contexts those labels condition on.
Terminology
Abstract
For LLM agents, supervised fine-tuning is not only about teacher labels' quality, but also about which interaction contexts those labels condition on. Pure behavioral cloning uses full teacher demonstrations, creating a mismatch between teacher-induced contexts seen in training and student-induced contexts encountered at test time. Recent work addresses this mismatch by querying a teacher at contexts reached by the student, often with increasingly elaborate filtering of the teacher's continuations. We instead frame on-policy data construction as a budget-allocation problem: under matched supervision resources, should teacher output be spent on more start-to-finish demos, longer continuations, outcome filtering, or broader coverage of learner-induced contexts? We formalize this design space through the rollout policy, switch-time distribution, continuation horizon, filtering rules, and two complementary costs: teacher inference generated before filtering and teacher supervision retained for SFT. Across HotpotQA, ALFWorld, and Terminal-Bench-Dev, bounded unfiltered teacher continuations at learner-induced contexts improve over pure behavioral cloning at matched budgets. On HotpotQA and ALFWorld, where we run the full comparison, few-step continuations match or exceed success-filtered and critical-context-filtered alternatives. Our findings suggest that a few teacher steps, placed at learner-induced contexts, can be a more cost-efficient supervision allocation than longer or more heavily curated teacher completions.
Sources
- Training LLM-Based Agents with Synthetic Self-Reflected Trajectories and Partial Masking
- ATLaS: Agent Tuning via Learning Critical Steps
- MiniLLM: On-Policy Distillation of Large Language Models
- Exploring Expert Failures Improves LLM Agent Tuning
- Imitation Learning for Multi-turn LM Agents via On-policy Expert Corrections
- Revisiting DAgger in the Era of LLM-Agents
- NL2Bash: A Corpus and Semantic Parser for Natural Language Interface to the Linux Operating System
- Student-Centered Distillation Narrows the Agentic Gap Between Small and Large LLMs
- Supervised Fine Tuning on Curated Data is Reinforcement Learning (and can be improved)
- Reinforcement and Imitation Learning via Interactive No-Regret Learning
- ALFWorld: Aligning Text and Embodied Environments for Interactive Learning
- OVD: On-policy Verbal Distillation
- ReAct: Synergizing Reasoning and Acting in Language Models
- Black-Box On-Policy Distillation of Large Language Models
- On-Policy Context Distillation for Language Models
- Agent-R: Training Language Model Agents to Reflect via Iterative Self-Training
- Decoupling KL and Trajectories: A Unified Perspective for SFT, DAgger, Offline RL, and OPD in LLM Distillation
- Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks