🤖 machine learning

Revisiting DAgger in the Era of LLM-Agents

This paper revisits Dataset Aggregation (DAgger) for long-horizon LLM agents to effectively mitigate covariate shift by combining on-policy interaction with dense teacher supervision, demonstrating significant performance gains in software engineering tasks where smaller models trained with this method outperform larger baseline systems.

Original authors: Changhao Li, Rushi Qiang, Jiawei Huang, Chenxiao Gao, Chao Zhang, Niao He, Bo Dai

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Changhao Li, Rushi Qiang, Jiawei Huang, Chenxiao Gao, Chao Zhang, Niao He, Bo Dai

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a brilliant but inexperienced apprentice (the Student) how to fix complex software bugs in a massive, messy codebase. You, the Master (the Teacher), are an expert who knows exactly how to solve these problems.

The goal is to get the apprentice to solve these problems on their own, eventually without your help. The paper argues that the two most common ways we currently try to teach apprentices are flawed, and it proposes a new, smarter method called DAgger (revisited for the age of Large Language Models).

Here is the breakdown of the problem and the solution using simple analogies.

The Problem: Two Flawed Teaching Styles

The paper identifies two existing ways to train these AI agents, both of which have a major weakness:

1. The "Shadowing" Method (Supervised Fine-Tuning / SFT)

How it works: The apprentice watches the Master solve a problem from start to finish and tries to copy every move perfectly.
The Flaw (Covariate Shift): This is like teaching a driver by only showing them videos of perfect driving on empty highways. But in the real world, the apprentice might make a tiny mistake early on (like turning the wheel slightly too early). Because they were only trained on "perfect" scenarios, they don't know how to recover from that mistake. They panic, the car goes off the road, and the whole trip fails.
In AI terms: The model learns to mimic the teacher, but if it makes a small error, it enters a "state" (a situation) it has never seen before, causing it to spiral into failure.

2. The "Trial and Error" Method (Reinforcement Learning / RL)

How it works: The apprentice is thrown into the codebase alone. They try to solve the problem, and if they succeed at the very end, they get a gold star. If they fail, they get nothing.
The Flaw (Sparse Feedback): This is like teaching someone to cook by only telling them "Good job" or "Bad job" after they have finished the entire meal. If they burned the first ingredient, they don't know which step caused the disaster. They have to guess, and it takes a massive amount of time and wasted food (computing power) to learn.
In AI terms: The model gets feedback only at the very end of a long task, making it hard to learn from specific mistakes along the way.

The Solution: The "Safety Net" Method (DAgger)

The paper proposes a new training method that combines the best of both worlds. Imagine a driving instructor who sits in the passenger seat but has a safety net.

How the New Method Works:

The Mix: The apprentice drives the car (solves the problem) for a few steps.
The Intervention: If the apprentice starts to drift or make a mistake, the Master (Teacher) gently takes the wheel for a moment to correct the path and get the car back on the right road.
The Lesson: Crucially, the apprentice doesn't just watch the Master fix it; they are taught what the Master would have done at that exact moment of confusion.
Fading Support: Over time, the Master intervenes less and less. The apprentice gets to drive more of the way, but whenever they hit a bump, the Master is there to show them the correct move for that specific bump.

Why this is better:

Realism: The apprentice learns how to handle the messy, real-world situations where they actually make mistakes (unlike the "Shadowing" method).
Rich Feedback: The apprentice gets a specific lesson for every single step, not just a grade at the end (unlike the "Trial and Error" method).
Efficiency: Because the Master helps them recover from early mistakes, the apprentice doesn't waste time spiraling into dead ends.

The Results: Small Models, Big Wins

The researchers tested this method on Software Engineering Agents (AI that fixes code). They used two sizes of student models: a smaller one (4 Billion parameters) and a medium one (8 Billion parameters).

The 4B Model: Using this new method, the small 4B model performed better than many published 8B models. It was like a small car with a perfect safety net outperforming a larger car with a broken one.
The 8B Model: The medium model got even better, nearly catching up to massive 32B models (which are much larger and more expensive).

What changed in the AI's behavior?

Less Panic: The models stopped getting stuck in loops or giving up when they made a mistake.
Better Recovery: When they did make a mistake, they knew how to fix it and get back on track.
Stability: The training process was much smoother and didn't crash as often as the other methods.

Summary

The paper argues that to teach AI agents to handle long, complex tasks (like fixing software bugs), we shouldn't just show them perfect examples or let them fail blindly. Instead, we should let them try, let them make mistakes, but immediately show them the correct way to handle those specific mistakes. This "safety net" approach (DAgger) allows smaller, cheaper AI models to perform as well as much larger, more expensive ones.

Technical Summary: Revisiting DAgger in the Era of LLM-Agents

Problem Statement

Long-horizon Large Language Model (LLM) agents operate in multi-turn interactive environments (e.g., software engineering, web browsing) where a single early error can alter the state distribution, causing the agent to diverge from the trajectory of successful expert demonstrations. This phenomenon, known as covariate shift, poses a significant challenge for post-training algorithms. Existing methods suffer from structural limitations in this setting:

Supervised Fine-Tuning (SFT): Trains on expert-induced states but fails during deployment because the student generates its own states, leading to a mismatch between training and test distributions.
Reinforcement Learning with Verifiable Rewards (RLVR): Learns from on-policy rollouts to avoid distribution mismatch but relies on sparse, outcome-level rewards, leading to difficult credit assignment and high computational costs.
On-Policy Distillation (OPD): Attempts to combine on-policy state coverage with teacher supervision but faces a "cold-start" bottleneck where weak students generate unproductive trajectories early in training, forcing the teacher to supervise failed prefixes. Additionally, OPD often requires access to teacher logits, limiting its use with black-box models.

The central question addressed is: Is there a method that can simultaneously exploit dense feedback, maintain on-policy state coverage, and provide early access to successful trajectories?

Methodology

The authors revisit Dataset Aggregation (DAgger), a classical imitation learning algorithm, and adapt it for multi-turn LM agents. The core principle is to train the student on states it is actually likely to visit, rather than only on expert-induced states.

Teacher-Interleaved Rollout

The proposed method employs a stochastic mixture of student and teacher policies during trajectory collection. Two specific rollout strategies are introduced:

DAgger-style (Turn-level Mixture): At each turn $t$ , an action is executed by the teacher policy $\pi_e$ with probability $\beta_i$ or by the student policy $\pi_\theta$ with probability $1-\beta_i$ . The mixing parameter $\beta_i$ is annealed from 1.0 to a lower bound (e.g., 0.6) over training iterations.
AggreVaTe-style (Trajectory-level Mixture): The student generates an initial prefix of length $\kappa$ , after which the teacher takes over to complete the trajectory. The prefix length distribution is scheduled to grow over time.

Key Mechanism: Regardless of whether the student or teacher executes the action at a given state, the teacher is queried for the correct action at every visited state. The student is then trained via supervised learning (cross-entropy loss) to mimic these teacher actions on the states it actually encountered.

Unified Framework

The paper situates this approach within a unified post-training framework defined by three components:

Context Distribution ( $p_s$ ): The distribution of states visited (mixture of student and teacher in DAgger vs. pure teacher in SFT or pure student in RL).
Label Distribution ( $p_a$ ): The source of action labels (Teacher in DAgger vs. Student in RL/OPD).
Scoring Function ( $w(s,a)$ ): The weight assigned to samples (Uniform in DAgger vs. Advantage-based in RL).

By using a mixture context distribution with teacher labels and uniform weights, the method bridges the gap between SFT (dense supervision, covariate shift) and RL/OPD (on-policy coverage, sparse/cold-start issues).

Key Contributions

Algorithmic Adaptation: Successfully adapts DAgger for long-horizon LLM agents, enabling the model to learn from realistic deployment states while receiving dense, token-level supervision from a teacher.
Cold-Start Mitigation: Unlike OPD, the teacher-interleaved approach ensures early trajectories are guided by the teacher, preventing the student from getting stuck in unproductive states during the initial training phase.
Black-Box Compatibility: The method only requires the teacher to provide action labels (tokens), not logits, making it compatible with black-box LLMs (e.g., GPT, Gemini).
Empirical Validation: Demonstrates consistent improvements over SFT, RL (GRPO), and On-Policy Distillation on software engineering tasks.

Experimental Results

The method was evaluated on SWE-Gym (training) and SWE-Bench Verified (out-of-domain evaluation) using Qwen3-4B and Qwen3-8B student models with a Qwen3-30B teacher.

Performance Gains:
- 4B Model: Achieved 27.3% on SWE-Bench Verified, outperforming the strongest post-training baseline (OPD) by +3.9 points and surpassing representative published 8B SWE-agent systems.
- 8B Model: Achieved 29.8% on SWE-Bench Verified, exceeding OPD by +3.6 points and surpassing the 32B-scale SWE-Gym-32B system by +9.2 points.
Training Stability: The method exhibited more stable scaling trajectories compared to SFT and OPD, particularly in the early stages (cold-start).
Covariate Shift Reduction: Analysis of reverse KL divergence showed that DAgger-style training significantly reduced the divergence between the student's deployment distribution and the teacher's distribution compared to SFT, which suffered from a rebound in divergence as training progressed.
Behavioral Improvements: Qualitative analysis indicated that DAgger-trained agents made fewer syntax/runtime errors and exhibited better recovery behaviors, shifting failure modes from "passive" errors (e.g., no submission) to harder limitations like context overflow.

Significance and Claims

The paper claims that explicit state-distribution correction via teacher-interleaved DAgger is a powerful ingredient for post-training long-horizon LM agents. The results suggest that:

Smaller models (4B/8B) trained with this method can approach or exceed the performance of significantly larger systems (32B).
The approach effectively mitigates the train-test distribution mismatch inherent in SFT while avoiding the sparse reward and cold-start issues of RL and OPD.
This methodology offers a practical, sample-efficient recipe for deploying robust tool-using agents in complex, multi-turn environments.

The authors emphasize that while the method improves reliability and sample efficiency, the remaining bottlenecks for these agents are increasingly related to long-context capacity (e.g., context overflow) rather than the training algorithm itself.