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

TL;DR

This paper introduces on-policy expert corrections (OECs), a novel data generation method inspired by DAgger, to improve multi-turn language model training by addressing covariate shift, showing significant performance gains in software engineering tasks.

Contribution

The paper proposes OECs, a new on-policy data collection technique for multi-turn LM training that mitigates covariate shift and enhances performance over traditional imitation learning.

Findings

OEC trajectories outperform traditional imitation learning by 14-13% in software engineering tasks.

Combining expert demonstrations with on-policy data improves multi-turn LM agent training.

Experiments validate the effectiveness of OECs in multi-turn language modeling scenarios.

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.