Imitation Learning for Multi-turn LM Agents via On-policy Expert Corrections
Proposes On-policy Expert Corrections (OEC) to mitigate covariate shift in multi-turn LM training, improving software engineering task success rates by 13-14%.
Key Findings
Methodology
Inspired by DAgger, OEC involves partially rolling out trajectories with a student model and switching mid-way to an expert model, combining on-policy data with expert demonstrations. The process includes random switch points, trajectory rewriting, and reward-based filtering. This hybrid approach addresses covariate shift by maintaining trajectory diversity and relevance, leveraging rejection sampling based on environment feedback, such as unit test success, to select high-quality data for fine-tuning. The method balances exploration and expert guidance, enabling more robust multi-turn interactions.
Key Results
- In software engineering benchmarks, 7B and 32B models trained with OEC trajectories achieved 14% and 13% higher resolution rates respectively over traditional imitation learning, demonstrating improved sample efficiency and robustness. The experiments showed that pure autonomous trajectories often degrade performance, while OEC-guided trajectories significantly outperform behavior cloning and full autonomous rollouts, especially in complex multi-turn tasks.
- Embedding analyses revealed that covariate shift accumulates over trajectory steps, but OEC trajectories effectively mitigate this drift, leading to more stable representations across states. Ablation studies confirmed that trajectory filtering and expert switching are critical for maximizing gains.
- The results highlight that combining expert demonstrations with on-policy data creates a synergistic effect, enabling models to generalize better in multi-turn environments, especially when dealing with complex software bugs or reasoning tasks.
Significance
This work offers a practical solution to a long-standing challenge in training multi-turn language models—covariate shift—by integrating expert corrections dynamically during training. It bridges the gap between pure imitation learning and reinforcement learning, providing a scalable, data-efficient approach that enhances model robustness and performance in real-world applications like automated software debugging. The methodology paves the way for more reliable, adaptable AI agents capable of complex multi-step reasoning and interaction, with broad implications for AI deployment in industry and research.
Technical Contribution
The core innovation lies in adapting DAgger’s iterative data aggregation to large-scale LM training through partial trajectory switching, combined with rejection sampling and trajectory rewriting. This hybrid approach allows for on-policy data collection that retains expert guidance, overcoming the limitations of traditional off-policy imitation learning. The method also introduces a filtering mechanism to eliminate poor-quality or repetitive trajectories, ensuring high data fidelity. Extensive experiments validate that OEC outperforms existing methods across multiple parameters and task complexities, establishing a new benchmark for multi-turn LM training.
Novelty
This is the first application of a partial, mid-trajectory expert switching mechanism in large-scale LM training, inspired by DAgger but tailored for autoregressive models and multi-turn interactions. Unlike pure behavior cloning or reinforcement learning, OEC combines the strengths of both, effectively addressing covariate shift. Its novelty lies in the dynamic, on-policy data augmentation during training, which significantly improves performance and robustness in software engineering tasks, setting a new paradigm for multi-turn agent training.
Limitations
- The approach relies heavily on the quality of the expert model, which can be computationally expensive to train and maintain. In extremely complex or multi-modal tasks, the expert guidance may not be sufficient to prevent drift.
- Trajectory filtering, while reducing poor data, may introduce bias by excluding certain behaviors, potentially limiting generalization.
- Scaling to larger models or real-time applications requires further optimization to manage computational costs and data management complexity.
Future Work
Future research will explore extending OEC to multi-modal and multi-task settings, integrating reinforcement learning rewards for adaptive trajectory sampling. Developing more efficient expert models and filtering strategies could reduce costs and improve scalability. Additionally, investigating automated curriculum strategies for expert switching points and broader application domains such as robotics or dialogue systems will be key directions.
AI Executive Summary
Over recent years, multi-turn language models (LMs) have advanced rapidly, especially in domains like software engineering, reasoning, and tool use. However, training these models effectively remains challenging due to the phenomenon of covariate shift—where models encounter states during deployment that differ from training data, leading to performance degradation. Traditional imitation learning methods, such as behavioral cloning, often fail to address this issue because they rely on off-policy data, which does not adapt to the model’s evolving behavior.
This paper introduces a novel approach called On-policy Expert Corrections (OEC), inspired by the classic DAgger algorithm, to mitigate covariate shift in multi-turn LM training. The core idea is to generate trajectories where the model begins with the student policy but switches mid-way to an expert policy at randomly chosen points. This partial on-policy rollout allows the collection of data that is more representative of the model’s current state distribution while still leveraging expert knowledge. The trajectories are then rewritten to ensure consistency and filtered based on environment feedback, such as unit test success, to retain only high-quality data.
Extensive experiments in the software engineering domain, specifically on the SWE-bench benchmark, demonstrate that models trained with OEC trajectories outperform traditional imitation learning by 14% (7B models) and 13% (32B models). These results highlight the effectiveness of combining expert demonstrations with on-policy data, leading to more robust and generalizable models capable of handling complex multi-turn interactions. The approach addresses key limitations of existing methods, offering a scalable, data-efficient pathway for training advanced AI agents.
The significance of this work lies in its potential to transform multi-turn AI agent training, making it more resilient to distributional shifts and better suited for real-world applications. By bridging imitation learning and reinforcement learning principles, the method opens new avenues for developing autonomous systems that can learn from limited supervision while adapting dynamically to evolving environments. Future directions include expanding to multi-modal tasks, optimizing expert model training, and integrating reinforcement signals for further improvements, promising a broad impact across AI research and industry.
Deep Dive
Abstract
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.