Imitation Learning for Multi-turn LM Agents via On-policy Expert Corrections

A popular paradigm for training LM agents relies on imitation learning, fine-tuning on expert trajectories. However, we show that the off-policy nature of imitation learning for multi-turn LM agents suffers from the fundamental limitation known as covariate shift: as the student policy's behavior diverges from the expert's, it encounters states not present in the training data, reducing the effectiveness of fine-tuning. Taking inspiration from the classic DAgger algorithm, we propose a novel data generation methodology for addressing covariate shift for multi-turn LLM training. We introduce on-policy expert corrections (OECs), partially on-policy data generated by starting rollouts with a student model and then switching to an expert model part way through the trajectory. We explore the effectiveness of our data generation technique in the domain of software engineering (SWE) tasks, a multi-turn setting where LLM agents must interact with a development environment to fix software bugs. Our experiments compare OEC data against various other on-policy and imitation learning approaches on SWE agent problems and train models using a common rejection sampling (i.e., using environment reward) combined with supervised fine-tuning technique. Experiments find that OEC trajectories show a relative 14% and 13% improvement over traditional imitation learning in the 7b and 32b setting, respectively, on SWE-bench verified. Our results demonstrate the need for combining expert demonstrations with on-policy data for effective multi-turn LM agent training.

Language models are increasingly used as agents that take many steps, use tools, and interact with external environments, but the standard imitation learning methods used to train these agents can fail when a student model makes different choices than the expert and reaches unfamiliar situations where errors compound. We introduce on-policy expert corrections, a data generation method where the student begins a task and an expert model takes over partway through to correct and complete the trajectory. This gives the student training examples that better reflect the situations it actually reaches while still preserving the benefits of expert demonstrations. We evaluate this method on software engineering tasks, where agents must inspect code, run commands, and write bug fixes over many steps, and find that training on on-policy expert corrections improves over standard imitation learning, with relative gains of 14% for 7B models and 13% for 32B models on SWE-bench verified. These results suggest that training long-horizon language model agents should combine expert demonstrations with data from the student’s own behavior.