# WM-DAgger: Enabling Efficient Data Aggregation for Imitation Learning with World Models

arXiv: https://arxiv.org/abs/2604.11351 · HTML: https://mlsanity.com/paper/2604.11351/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 10 of its 25 references that are in the map.
- First posted on arXiv 13 Apr 2026.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)
- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 10 of 10)

- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [GR00T N1: An Open Foundation Model for Generalist Humanoid Robots](https://mlsanity.com/paper/2503.14734/) (2025)
- [Cosmos World Foundation Model Platform for Physical AI](https://mlsanity.com/paper/2501.03575/) (2025)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [Understanding World or Predicting Future? A Comprehensive Survey of World Models](https://mlsanity.com/paper/2411.14499/) (2024)
- [A Comprehensive Survey on World Models for Embodied AI](https://mlsanity.com/paper/2510.16732/) (2025)
- [Diffusion Meets DAgger: Supercharging Eye-in-hand Imitation Learning](https://mlsanity.com/paper/2402.17768/) (2024)
- [World4RL: Diffusion World Models for Policy Refinement with Reinforcement Learning for Robotic Manipulation](https://mlsanity.com/paper/2509.19080/) (2025)
- [ManiGaussian++: General Robotic Bimanual Manipulation with Hierarchical Gaussian World Model](https://mlsanity.com/paper/2506.19842/) (2025)
- [SimWorld: A Unified Benchmark for Simulator-Conditioned Scene Generation via World Model](https://mlsanity.com/paper/2503.13952/) (2025)

## Built on by (most cited first, 7 of 7)

- [BWM: A Low-Cost High-Fidelity World Simulator for Robot Learning](https://mlsanity.com/paper/2607.29302/) (2026)
- [Do Robotic World Models Really Follow Actions? Diagnosing and Aligning Action-Conditioned Generation for Policy Learning](https://mlsanity.com/paper/2608.24885/) (2026)
- [TACO: TActile World Model as a Self-COrrector for Scalable Robot Policy Post-Training](https://mlsanity.com/paper/2607.02840/) (2026)
- [XPACE: Joint World and Action Modeling from Heterogeneous Experience](https://mlsanity.com/paper/2609.17372/) (2026)
- [Towards a Data Flywheel for Embodied Intelligence in Logistics](https://mlsanity.com/paper/2606.05960/) (2026)
- [EgoRecovery: Acquiring Failure Recovery Ability Through Human Recovery Demonstration](https://mlsanity.com/paper/2607.19745/) (2026)
- [Policy-Calibrated DAgger: Offline Calibrated Noise Injection for Imitation Learning](https://mlsanity.com/paper/2609.30462/) (2026)

## Abstract (arXiv)

Imitation learning is a powerful paradigm for training robotic policies, yet its performance is limited by compounding errors: minor policy inaccuracies could drive robots into unseen out-of-distribution (OOD) states in the training set, where the policy could generate even bigger errors, leading to eventual failures. While the Data Aggregation (DAgger) framework tries to address this issue, its reliance on continuous human involvement severely limits scalability. In this paper, we propose WM-DAgger, an efficient data aggregation framework that leverages World Models to synthesize OOD recovery data without requiring human involvement. Specifically, we focus on manipulation tasks with an eye-in-hand robotic arm and only few-shot demonstrations. To avoid synthesizing misleading data and overcome the hallucination issues inherent to World Models, our framework introduces two key mechanisms: (1) a Corrective Action Synthesis Module that generates task-oriented recovery actions to prevent misleading supervision, and (2) a Consistency-Guided Filtering Module that discards physically implausible trajectories by anchoring terminal synthesized frames to corresponding real frames in expert demonstrations. We extensively validate WM-DAgger on multiple real-world robotic tasks. Results that our method significantly improves success rates, achieving a 93.3% success rate in soft bag pushing with only five demonstrations. The source code is publicly available at https://github.com/czs12354-xxdbd/WM-Dagger.