WM-DAgger: Enabling Efficient Data Aggregation for Imitation Learning with World Models
Authors: Anlan Yu, Zaishu Chen, Peili Song, Zhiqing Hong, Haotian Wang, Desheng Zhang, Tian He, Yi Ding, +1 more
Organizations: Peking University, Beijing, China · JD Logistics, China · Nankai University, Tianjin, China · The Hong Kong University of Science and Technology (Guangzhou), China · Rutgers University, Piscataway, USA · the University of Texas at Dallas, Richardson, TX, USA · Institut Polytechnique de Paris, France
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.
Figures & tables
Fig. 1 : WM-DAgger mitigates the compounding errors of standard Behavioral Cloning (BC) by generating massive recovery supervision with a world model. (e.g., visual transitions a→b→c and d→e→f ).
Fig. 2 : Overall pipeline of our WM-DAgger framework.
Fig. 3 : EAC-WM Architecture.
Fig. 4 : Corrective Action Synthesis Module.
Fig. 5 : Consistency-Guided Filtering Module and Visualization of Real, Filtered, and Retained Frames. Frames with morphing and incorrect object position are filtered.
Fig. 6 : Experimental setup and manipulation tasks. (a) Hardware configuration for data collection, featuring a handheld gripper integrated with a GoPro camera for visual observations and an HTC Vive Tracker for high-precision 6-DoF trajectory tracking. (b–e) The four evaluation tasks: (b) Soft Bag Pushing, (c) Pick-and-Place, (d) Ballot Insertion, and (e) Towel Folding.
Fig. 7 : Visualization results of EAC-WM-generated frames. Each row depicts a specific task, starting from the real expert frame It (first column). Subsequent columns ( I^t+k ) show synthesized observations conditioned on action sequences in six cardinal directions (right, left, up, down, forward, and backward).
Fig. 8 : Visual evaluation of EAC-WM versus DMD in generating soft bag pushing data. EAC-WM exhibits superior structural and physical consistency. Notably, at t=16, EAC-WM accurately simulates the dynamics of the bag dropping.
Method
1-shot
3-shot
5-shot
10-shot
20-shot
Standard BC
6.7
20.0
26.7
30.0
30.0
DMD [ 6 ]
13.3
33.3
40.0
53.3
56.7
Ours
73.3
86.7
93.3
93.3
96.7
TABLE I : Success Rates(%) for Task 1: Soft Bag Pushing across varying numbers of expert demonstrations.
Synthesized Samples
300
900
1500
3000
Success Rate (%)
46.7
63.3
96.7
96.7
TABLE II: Policy Performance vs. Data Scaling.
Method / Variant
Success Rate (%)
WM-DAgger (Ours Full)
96.7
Data Ablation
w/o Play Data
83.3
Module Ablation
w/o Filter
66.7
w/o Dir.
0.0
TABLE III : Ablation Studies for 20-shot Soft Bag Pushing.
Seen
Unseen
Method
O1
O2
O3
O4
O5
Standard BC
13.3
13.3
6.7
0.0
10.0
DMD
33.3
36.7
26.7
6.8
16.7
Ours
83.3
90.0
80.0
63.3
76.7
TABLE IV : Success Rates(%) for Task 2: Pick-and-Place.
Collecting human demonstrations via teleoperation is a common approach for teaching robots task-specific skills. However, when only a limited number of demonstrations are available, policies are prone to entering out-of-distribution (OOD) states due to compounding errors or environmental stochasticity. Existing interactive imitation learning or human-in-the-loop methods try to address this issue by following the Human-Gated DAgger (HG-DAgger) paradigm, an approach that augments demonstrations through selective human intervention during policy execution. Nevertheless, these approaches struggle to balance dexterity and generality: they either provide fine-grained corrections but are limited to specific kinematic structures, or achieve generality at the cost of precise control. To overcome this limitation, we propose the Human-Robot Copilot framework that can leverage a scaling factor for dexterous teleoperation while maintaining compatibility with a wide range of industrial and research manipulators. Experimental results demonstrate that our framework achieves higher performance with the same number of demonstration trajectories. Moreover, since corrective interventions are required only intermittently, the overall data collection process is more efficient and less time-consuming.
Policies trained with imitation learning can accumulate errors over time, causing the robot to drift outside the training distribution. Existing methods mitigate this covariate shift by collecting additional data where the policy fails or is likely to fail. The first places the robot in unsafe conditions and the second requires choosing an appropriate noise distribution to collect new expert demonstrations under that noise. We propose Policy-Calibrated DAgger, a method that makes use of the properties of recent generative policies to estimate the policy's noise offline by using its own predicted action distribution. We measure a diffusion policy's spread of predicted actions at observations along the expert trajectory and measure its closed-loop error relative to a recorded trajectory. To address issues with measuring error in a multimodal action space, we guide the policy towards the trajectory during closed-loop control through partial denoising, and use properties of a diffusion model to unnormalize the measured error as if we did not guide it. We experiment in a scenario where a robot is tasked to reach an engine lever in a cluttered and narrow environment and show results in a 3D photorealistic simulator and a 2D planar reacher environment. We show that our method surpasses policies trained with dataset aggregation without noising and matches the performance of the best noise level in hindsight, without requiring a sweep over noise levels.
In-context imitation learning (ICIL) enables robots to learn new tasks from a small number of demonstrations by conditioning a pre-trained policy on task-specific examples, without retraining at test time. Despite this promise, training generalizable and scalable in-context imitation policies remains an open challenge. We present SynthICL, a scalable framework that trains ICIL policies entirely from RGB-only synthetic data. Specifically, we build a data generation pipeline to produce high-fidelity ICIL data and train a flow-matching transformer policy on the resulting dataset. SynthICL avoids the need for depth sensing, precise camera calibration, and real-world training data in prior approaches, offering a simpler and more scalable alternative. We further incorporate subgoal prediction by training the model to predict the next subgoal images, enabling more precise and visually grounded control. Evaluated on 16 unseen real-world manipulation tasks, SynthICL achieves an average success rate of 79% with only one demonstration provided at test time and outperforms prior methods. Project page: https://synth-icl.github.io