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

Abstract: 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.

• The paper proposes WM-DAgger, a framework that uses world models to synthesize off-distribution recovery data for imitation learning.
• It integrates an eye-in-hand action-conditioned world model, corrective action synthesis, and consistency-guided filtering to ensure physical and visual plausibility.
• Experiments demonstrate robust policy improvements in 6-DoF robotic manipulation tasks with as few as five demonstrations, outperforming state-of-the-art baselines.

WM-DAgger: Data Aggregation for Imitation Learning via Physically Consistent World Model Synthesis

Introduction

The paper "WM-DAgger: Enabling Efficient Data Aggregation for Imitation Learning with World Models" (2604.11351) introduces an imitation learning framework, WM-DAgger, designed to overcome compounding error in behavioral cloning (BC) by using world models (WMs) to synthesize out-of-distribution (OOD) recovery data. Unlike traditional DAgger, which requires continual human supervision for data aggregation, WM-DAgger eliminates human-in-the-loop requirements through autonomous recovery trajectory generation and systematic hallucination suppression. The proposed approach primarily targets robotic manipulation within the eye-in-hand paradigm under a limited human demonstration budget.

Figure 1: WM-DAgger mitigates compounding errors in standard BC by generating large-scale recovery supervision from a world model.

WM-DAgger Framework

WM-DAgger leverages world models to synthesize recovery behaviors for OOD states encountered during policy rollouts without human intervention. The framework is instantiated through a pipeline that combines expert demonstrations with generated recovery trajectories, improving the robustness of the learned policy.

The key modular components in WM-DAgger are:

• Eye-in-Hand Action-Conditioned World Model (EAC-WM): An architecture conditioned on both visual feedback and action-space information, built atop Cosmos-Predict2.5 and enhanced with dense, pixel-aligned geometric encoding of robot actions for higher-fidelity future frame synthesis.
• Corrective Action Synthesis Module: Given expert demonstrations, this module generates recovery action trajectories from selected "deviation" points, enforcing geometric consistency by constraining synthesized actions to be task-oriented and aligned with the demonstrated manifold.
• Consistency-Guided Filtering Module: To address the hallucination and cumulative simulation error inherent in long-horizon predictions, this module utilizes DINOv2 embeddings for feature-space consistency checks. Generated rollouts whose terminal frames deviate visually/physically from real expert states are discarded.

  Figure 2: The overall pipeline of WM-DAgger, highlighting synthesis, filtering, and aggregation of recovery data.

EAC-WM Architecture and Conditioning

The EAC-WM is the central mechanism for generating action-conditioned visual trajectories. Its core innovation is the Action2Image module, which maps control actions into pixel-wise displacement and orientation fields, thus providing explicit camera-motion context required for accurate prediction in eye-in-hand setups. Both play data (random exploratory trajectories) and task data (expert demonstrations) are collected for post-training the world model. Synthesis and tokenization leverage VAEs for latent representation and Rectified Flow for robust space-time interpolation.

Figure 3: EAC-WM Architecture, featuring geometric action conditioning for each pixel and VAE-based tokenization.

Corrective Action and Hallucination Filtering

The Corrective Action Synthesis Module systematically generates deviation and recovery action sequences around the expert manifold, ensuring that recovery data reflects plausible, task-centric behaviors rather than arbitrary perturbations. The Conistency-Guided Filtering Module operates at visual feature level: by comparing DINOv2 embeddings between generated terminal frames and their corresponding real expert frames, the model automatically discards trajectories with evidence of object morphing, spatial displacement, or unphysical artifacts.

Figure 4: Corrective Action Synthesis Module—generation of structured OOD recovery action sequences for robust policy correction.

Figure 5: Consistency-Guided Filtering—visualization of real, filtered (hallucinated), and retained (plausible) terminal frames.

Policy Training with Aggregated Data

The robot policy is trained using a temporally chunked action-prediction model, such as Gr00t N1.5, on a training set comprising both expert and validated synthetic recovery trajectories. This approach directly internalizes corrective behaviors into the policy, significantly improving resilience to state drift and execution error over baseline BC.

Experimental Evaluation

Comprehensive experiments are conducted across four manipulation tasks—soft bag pushing, pick-and-place (seen/unseen objects), ballot insertion, and towel folding—in 6-DoF action space environments. The results indicate that WM-DAgger achieves dramatic improvements over both standard BC and the SOTA DMD baseline, especially in settings with a small demonstration budget (e.g., 93.3% success in soft bag pushing with only five demonstrations). The improvement is robust to data scaling, saturating after moderate amounts of synthesized data. Ablations confirm catastrophic degradation if either corrective action constraints or hallucination filtering is removed.

Figure 6: Robotic experimental setup and the four real-world manipulation tasks evaluated.

Figure 7: Visualization of EAC-WM-generated frames in six directions, illustrating fine-grained physical consistency and accurate camera-motion.

Figure 8: Comparison of EAC-WM and DMD for soft bag pushing. EAC-WM better captures structure and physics of interaction, especially under complex deformable object dynamics.

Implications and Future Directions

WM-DAgger establishes that physically consistent and action-conditioned data synthesis via large world models can obviate costly human feedback loops in imitation learning, enabling efficient and robust policy training in sample-limited regimes. The system generalizes recovery across OOD states, supports both rigid and deformable object manipulation, and scales to complex tasks such as towel folding under 6-DoF control. The findings imply that model-based augmentation, when coupled with principled constraint and filtering mechanisms, offers a more scalable path for embodied control than previously dominant diffusion-based augmentation.

A primary limitation is the current inapplicability to multi-finger dexterous manipulation due to compounded DoF and articulation complexity. Future work is likely to incorporate morphological priors, higher-fidelity kinematic representations, and cross-task transfer mechanisms into world models for dexterous hands and more intricate articulated objects.

Conclusion

WM-DAgger provides an effective methodology for overcoming compounding errors in imitation learning through integrated world model-based data aggregation. Its combination of geometric action conditioning, structured recovery synthesis, and semantic hallucination filtering results in significant performance gains on a broad suite of manipulation tasks. The framework validates the practical utility of high-capacity world models as scalable, autonomous supervisors for embodied policy learning, setting a solid foundation for extending these techniques to the dexterous and high-complexity domains of future AI-driven robotics.