Research map / Multi-agent systems and games / LLM agents

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

ICML 2026 · Multi-agent systems and games › LLM agents · approach: Agentic benchmarks and orchestration

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes on-policy expert corrections (OECs), which roll out a student model and switch to an expert mid-trajectory to create partially on-policy SWE-agent training data, and shows gains over behavioral cloning on SWE-bench Verified.

Problem

Imitation learning for multi-turn LM agents suffers from covariate shift: the student encounters states absent from expert-generated training data, and errors compound over long trajectories.

Approach

Generate training trajectories by starting rollouts with the student and switching to the expert at a random turn, so the expert corrects on-policy states while rollouts run to completion and can be filtered by unit-test rejection sampling; the student-generated prefix is masked out of the loss.

Key results

• OEC trajectories give a relative 14% (7B) and 13% (32B) improvement over traditional imitation learning on SWE-bench Verified.
• Student-32B rises from 31.2% with no training to 40.0% with OEC, versus 36.0% with 3,000 BC trajectories.
• Fine-tuning on purely on-policy trajectories severely degrades 32B performance, and covariate shift (Frechet distance) grows through a trajectory.

Surprising findings

• Fine-tuning on purely on-policy student trajectories, even after unit-test filtering, can significantly degrade performance in the 32B setting.
• Unit-test rejection sampling alone does not guarantee high-quality trajectories; repetitive loops in positive samples harm training.

Limitations

• No theoretical no-regret guarantees as in DAgger, since the state distribution remains partially off-policy.
• Requires a stronger expert model as a source of trajectories.
• On-policy portions become increasingly off-policy as fine-tuning proceeds, limiting benefit.

Experimental setup

Datasets: SWE-smith SWE-bench Verified SWE-bench Lite

Benchmarks: SWE-bench Verified SWE-bench Lite

Metrics: resolution rate covariate shift Frechet distance LLM-as-judge failure categories

Baselines: Behavioral cloning On-policy trajectory generation SWE-smith-LM-32B SWE-smith-LM-7B SWE-Gym-32B Skywork-SWE-32B SkyRL-Agent-7B-v0

Builds on: DAgger behavioral cloning rejection sampling fine-tuning SWE-smith

Closest papers by meaning

• Mitigating Covariate Shift in Behavioral Cloning via Robust Stationary Distribution Correction (NeurIPS 2024 · similarity 0.88)
• Behavior Cloning is Not All You Need: The Optimality of On-Policy Distillation for Noisy Expert Feedback (NeurIPS 2026 · similarity 0.87)
• Learning When to Stop: Selective Imitation Learning Under Arbitrary Dynamics Shift (NeurIPS 2026 · similarity 0.87)
• When Does Online Imitation Learning Help in LLM Post-Training? The Role of (Non-)Realizability Beyond Horizon (NeurIPS 2026 · similarity 0.87)
• Learning on One Mode: Addressing Multi-modality in Offline Reinforcement Learning (ICLR 2025 · similarity 0.87)
• Skill-SD: Skill-Conditioned Self-Distillation for Multi-turn LLM Agents (NeurIPS 2026 · similarity 0.87)