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

2026-07-05

AI Summary

This work addresses the distribution mismatch commonly faced by large language model agents in supervised fine-tuning, where training relies on complete teacher demonstrations while testing depends on student-generated contexts. The authors formulate online policy data construction as a budget allocation problem and propose replacing lengthy or costly filtered teacher trajectories with a small number of unfiltered, short-step teacher continuations, strategically injected into student-induced critical contexts. By systematically exploring the design space of rollout policies, switching time distributions, continuation lengths, and filtering rules—and incorporating a dual-cost model accounting for both teacher inference and supervision signal retention—the method demonstrates strong empirical performance on HotpotQA, ALFWorld, and Terminal-Bench-Dev. Notably, it matches or exceeds existing critical-context filtering baselines on the first two benchmarks at lower computational cost, indicating that a few well-placed teacher steps can substantially enhance training efficiency.

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.

Problem

• on-policy data augmentation
• supervised fine-tuning
• distribution mismatch
• teacher-student alignment
• cost-efficient supervision

Innovation

• on-policy data augmentation
• budget allocation
• supervised fine-tuning
• teacher-student mismatch
• few-step continuation