A Few Teacher Steps Go a Long Way: Cost-Efficient On-Policy Data Augmentation for Agent Post-Training
2026-07-06 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors study how to best use teacher examples when training language model agents. Instead of always copying full teacher demonstrations, they explore smarter ways to use limited teacher feedback, like giving shorter teacher responses at points where the student model actually finds itself. They find that providing a few steps of teacher guidance during these learner-generated situations can work better and be more efficient than longer or heavily filtered teacher examples. Their experiments on tasks like HotpotQA and ALFWorld show this approach improves performance under the same supervision budget.
LLM agentssupervised fine-tuningbehavioral cloningon-policy datateacher continuationsHotpotQAALFWorldrollout policysupervision budget
Authors
Junze Ye, Jiayi Cheng, Miao Lu, Michal Mankowski, Jose Blanchet, Mohsen Bayati
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.