PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Niklas Lauffer, Xiang Deng, Srivatsa Kundurthy, Brad Kenstler, Jeff Da

arXiv:2512.14895v1cs.LGcs.AI

TL;DR

Multi-turn imitation learning can encounter covariate shift as student behavior diverges from expert demonstrations. The paper introduces on-policy expert corrections, which switch from student to expert partway through a rollout. In SWE-agent experiments, OEC trajectories improve sample efficiency over baseline approaches, while the paper notes limits to their theoretical guarantees and on-policy character.

  • Problem

    In multi-turn imitation learning, student errors can compound and produce deployment states unlike those in expert demonstrations, creating covariate shift.

    Source: §1 · p. 1

  • Method

    On-policy expert corrections start rollouts with a student model and switch to an expert partway through, retaining expert data while incorporating on-policy states.

    Source: §1 · p. 2

  • Results

    OEC trajectories improve sample efficiency over behavioral-cloning and on-policy trajectory baselines; in the 32B setting, 961 OEC plus 961 BC trajectories resolve 23 (4.6%) more SWE-bench verified instances than 2,000 BC trajectories.

    Source: §1 · p. 2, §4.2 · p. 7

  • Takeaways & Limitations

    The experiments highlight the importance of assessing trajectory quality beyond verifiable rewards, since a small proportion of low-quality positive trajectories can destabilize learning.

    Source: §5 · p. 9

    OECs lack DAgger's no-regret learning guarantees, and their on-policy portions become increasingly off-policy during fine-tuning.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Multi-turn imitation learning can suffer covariate shift as student errors compound, taking agents into states unlike expert demonstrations. The paper introduces on-policy expert corrections (OECs) and reports improved sample efficiency over behavioral cloning and on-policy trajectory baselines.

1 INTRODUCTION

  • Covariate shift grows across multi-turn trajectories as student inaccuracies compound and lead agents into states unlike those in expert demonstrations.For LM agents, divergence includes both environment state and the history of prior tokens, tool calls, and actions supplied as model input.Source: §1 · p. 1
  • On-policy expert corrections (OECs) start a rollout with the student and switch to an expert partway through, combining on-policy data with expert trajectories.The trajectory can continue to completion, allowing environment-verifier rewards to be used for rejection sampling.Source: §1 · p. 2
  • OEC trajectories improve sample efficiency over behavioral cloning and on-policy trajectory baselines in SWE-agent training.The reported experiments fine-tune Qwen2.5-Coder models using unit-test rejection sampling and supervised fine-tuning.Source: §1 · p. 2
  • The study evaluates OECs against behavioral cloning and on-policy trajectories in SWE tasks, where agents interact with buggy codebases to submit patches.The experiments use rejection sampling based on unit tests with supervised fine-tuning.Source: §1 · p. 2
  • The reported contributions include Qwen2.5-Instruct 7B and 32B models reaching state-of-the-art performance in their respective classes on SWE-bench verified.The ablations also report that unit-test rejection sampling alone does not ensure high-quality training trajectories.Source: §1 · p. 3

2 RELATED WORKS

Related work distinguishes off-policy imitation learning from on-policy reinforcement-learning approaches and situates OECs between them. It also traces covariate shift in imitation learning and describes the multi-turn software-engineering agent setting.

Training LM Agents

  • Off-policy agent training commonly uses supervised fine-tuning on expert demonstrations, including teacher-student distillation (Hinton et al. (2015)); on-policy approaches typically use reinforcement learning.Rejection-sampling fine-tuning also trains on successful on-policy trajectories.Source: §2 · p. 3

Imitation Learning and Covariate Shift

  • Covariate shift arises when test-time states diverge from training data and can compound in multi-turn imitation learning; DAgger addresses it through policy rollouts with expert corrective actions (Ross et al., 2011).The paper draws on DAgger to motivate partially on-policy trajectories for multi-turn LM agents.Source: §2 · p. 3

LM Agents for SWE Tasks

  • SWE-bench evaluates agents on real GitHub issues requiring codebase navigation, issue diagnosis, and a working patch.More general agent scaffolds permit arbitrary command sequences and longer trajectories, increasing the setting's covariate-shift challenge.Source: §2 · p. 3

3 METHODOLOGY

The methodology generates partially on-policy expert-correction trajectories, then filters and prepares them for supervised fine-tuning. It addresses the data-quality risks of repetitive trajectories while retaining environment-based rejection sampling.

3.1 ON-POLICY EXPERT CORRECTIONS

  • OECs begin with a student rollout and switch to an expert at a randomly chosen turn, preserving the trajectory history while rewriting it into the expert scaffold's format.The expert continues the rollout, so the expert portion follows a student-generated history.Source: §3.1 · p. 4, §3 · p. 4
  • OEC trajectories do not retain DAgger's theoretical no-regret benefits, but they support environment-feedback rejection sampling and efficient multi-turn fine-tuning with shared context.The underlying state distribution remains not completely on-policy.Source: §3.1 · p. 4

3.2 REJECTION SAMPLING FINE-TUNING AND ON-POLICY MASKING

  • Fine-tuning treats each agent turn as a sample and masks environment-step losses; trajectories failing unit tests are removed, and student-generated portions are masked from training.A trajectory submitted before the expert switch is removed because its entire training portion would be masked.Source: §3.2 · pp. 4–5

3.3 FILTERING REPETITIVE TRAJECTORIES

  • Even unit-test-passing on-policy trajectories can degrade performance when repetitive actions dominate their gradients; the filter removes three identical consecutive actions or 20-plus consecutive file-reading commands.The paper describes repeated file reading as a recurring loop that can reinforce negative behavior.Source: §3.3 · p. 5

4 EXPERIMENTS

The experiments compare trajectory-generation strategies for SWE agents and examine the effects of data filtering, on-policy masking, covariate shift, and failure patterns. OEC–BC mixtures outperform pure BC in the reported 32B comparison, while fully on-policy training can reduce performance.

4.1 EXPERIMENT SETUP

  • OEC trajectories switch from student to expert at a random index drawn from U(0, 30), and are compared with expert-only behavioral cloning and fully student-generated trajectories.All methods use rejection-sampling fine-tuning and repetition filtering, with performance tested across dataset sizes.Source: §4.1 · p. 6

4.2 RESULTS

  • In the 7B setting, OEC training produced larger performance gains than BC, while fully on-policy trajectories yielded little, if any, improvement.Each model also used 1,000 original Claude-3.7-Sonnet trajectories during fine-tuning.Source: Figure 2 · p. 6, §4.2 · p. 7
  • 23 (4.6%) additional SWE-bench verified instances were resolved with 961 OEC and 961 BC trajectories than with 2,000 BC trajectories; fully on-policy fine-tuning severely decreased performance.The fully on-policy decrease occurred despite unit-test rejection sampling and repetition filtering.Source: Figure 2 · p. 6, §4.2 · p. 7
  • A mixture of OEC and BC trajectories performed best, with 961 OEC and 2,000 BC trajectories producing the highest-performing model across the experiments.The figure varies one trajectory type while holding the other at 961 samples; the lines meet at an even 961–961 split.Source: §4.2 · p. 7, Figure 3 · p. 8
  • Removing repetition filtering corresponded to absolute resolution-rate differences of 6% in the 7B setting and 3.6% in the 32B setting.On-policy masking made no difference in 7B but was crucial in 32B, where its removal essentially negated the expert portion’s benefit.Source: §4.2 · pp. 7–8, Table 2 · p. 8

4.3 COVARIATE SHIFT BETWEEN THE STUDENT AND EXPERT

  • The covariate-shift analysis embeds rollout histories and compares student and expert embedding distributions turn by turn; measured divergence rises after a brief dip until around turn 10.The Fréchet Distance increases sharply through turn 30 and continues rising afterward, despite behavioral-cloning training.Source: §4.3 · pp. 8–9, Figure 4 · p. 8

4.4 QUALITATIVE ANALYSIS

  • The qualitative analysis uses GPT-5 to classify one failure mode for each unresolved instance from its truncated trajectory and groups failures by whether a patch was submitted.The categories cover submitted-patch errors and non-submission issues such as tool-calling errors, context horizons, and loops.Source: §4.4 · p. 9, Table 3 · p. 9

5 CONCLUSION

The paper introduces OECs to address covariate shift in multi-turn LM imitation learning and reports that combining on-policy data with expert demonstrations matters. Its SWE-agent findings also highlight data-quality concerns and limitations of the approach.

5 CONCLUSION

  • OECs combine on-policy rollouts, expert data, and rejection sampling to address covariate shift in multi-turn LM imitation learning.The experiments highlight limitations of relying only on on-policy training or only on expert demonstrations.Source: §5 · p. 9
  • A small proportion of low-quality positive trajectories can greatly destabilize learning, so data quality needs evaluation beyond verifiable rewards.Source: §5 · p. 9
  • The experiments focus on SWE agents, and OEC trajectories require a source of expert trajectories.The authors identify testing other multi-turn LM domains and using human or enhanced-compute experts as possible future directions.Source: §5 · pp. 9–10
  • OECs are not known to share DAgger’s no-regret learning guarantees, and their on-policy portions can become increasingly off-policy during fine-tuning.The authors suggest multiple intermediate generation rounds or online trajectory generation as possible mitigations.Source: §5 · p. 10

6 REPRODUCIBILITY STATEMENT

The reproducibility statement describes planned releases of code, models, and OEC trajectories, and points to the paper sections and appendix containing experiment details.

6 REPRODUCIBILITY STATEMENT

  • The supplementary materials include code for generating OEC trajectories, computing covariate shift, and conducting the LLM-as-judge analysis.Source: §6 · p. 10
  • The authors state that OEC-SWE-32B, OEC-SWE-7B, and the collected OEC trajectories will be open-sourced upon publication.Source: §6 · p. 10
  • The trajectory problem instances come from SWE-smith, with their distribution described in Section 4.1.Source: §6 · p. 10
  • Section 3 and Algorithm 1 detail OEC generation and model training, while Appendix C lists supervised fine-tuning hyperparameters.Source: §6 · p. 10

B AN EXAMPLE ON-POLICY EXPERT CORRECTION TRAJECTORY

In the example OEC trajectory, Student-32B spends its first 29 turns failing to localize a bug; Claude-3.7-Sonnet then completes localization, patches the issue, and resolves it. The transcript concerns a YamlLint ignore-pattern issue and shows the resulting changes passing tests.

B AN EXAMPLE ON-POLICY EXPERT CORRECTION TRAJECTORY

C HYPERPARAMETERS

Table 4 identifies the hyperparameters used in the 7B and 32B supervised fine-tuning experiments.

C HYPERPARAMETERS

  • Table 4 presents hyperparameters for the 7B and 32B supervised fine-tuning experiments.Source: Table 4 · p. 27

D MODEL FAILURE ANALYSIS

The judge prompt asks an expert software engineer to analyze why a failed SWE trajectory ended and assign one predefined error category. It supplies tool context, the problem, and recent actions and observations as input.

D MODEL FAILURE ANALYSIS

  • The prompt frames the task as expert software-engineering failure analysis and documents file commands—view, create, str_replace, insert, and undo_edit—with a required absolute path and command-specific arguments for file creation and replacement.Source: §D.1 · p. 28
  • The judge receives the problem statement, recent actions and observations, trajectory summary, and exit-status guidance before analyzing a failed run.The prompt requests an explanation of the issue and why the trajectory failed, followed by a category.Source: §D.1 · p. 29
  • The category set covers file identification, missed edge cases, problem misunderstanding, wrong solutions, tool errors, loops, repeated reading, context overflow, syntax errors, and other failures.The prompt says to choose exactly one category and not invent new ones.Source: §D.1 · p. 30

D.2 SAMPLE JUDGEMENTS

The sample judgments illustrate two failed SWE trajectories: one stalled in repeated file reading without producing a patch, while another pursued an incorrect fix without confirming that edits affected runtime behavior.

D.2 SAMPLE JUDGEMENTS

  • Repeatedly reopening backend files without editing left the duration-expression bug unfixed and produced no patch.The agent made a view-range misstep and received truncated outputs; the judgment classifies the failure as endless_file_reading.Source: §D.2 · p. 30
  • For the xarray issue, speculative edits to where did not preserve attributes or dtype, and no tests confirmed the modified code path was running.The judgment notes that the code already passed keep_attrs to apply_ufunc and suggests the fix may have belonged elsewhere.Source: §D.2 · p. 30
  • The xarray trajectory is classified as a wrong solution: it changed the wrong place and did not ensure its edits were executed.The proposed edits included unrelated dtype casting and brittle attribute logic.Source: §D.2 · p. 30

E STUDENT-EXPERT COVARIATE SHIFT KL-DIVERGENCE & BASELINE

The appendix captions describe divergence between student and expert across trajectory turns, including an expert-expert baseline. It also notes that LLMs assisted with related-work searches, LaTeX formatting, and plotting.

E STUDENT-EXPERT COVARIATE SHIFT KL-DIVERGENCE & BASELINE

  • KL-measured student-expert distributional divergence increases over trajectory turns, with Figure 5 relating the measure to Fréchet Distance.Figure 6 adds a dashed expert-expert partition baseline; the caption does not state its numerical values.Source: Figure 5 · p. 31, Figure 6 · p. 31

F USE OF LLMS

  • LLMs were used to find related work, format LaTeX syntax such as tables, and generate result plots.Source: F USE OF LLMS · p. 32

Related papers

Loading 2512.14895v1…