PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Anlan Yu, Zaishu Chen, Peili Song, Zhiqing Hong, Haotian Wang, Desheng Zhang, Tian He, Yi Ding, Daqing Zhang

arXiv:2604.11351v1cs.RO

TL;DR

Compounding errors in out-of-distribution states and DAgger’s dependence on human intervention limit scalable imitation learning, while single-frame synthesis does not capture continuous recovery dynamics. WM-DAgger uses world models with corrective-action and consistency-filtering mechanisms to synthesize recovery data, reporting 93.3% success in soft bag pushing with five demonstrations.

  • Problem

    DAgger depends on human intervention, limiting scalability, while single-frame recovery synthesis cannot model continuous dynamics and struggles with complex physical interactions.

    Source: §I · p. 1

  • Method

    WM-DAgger uses world models to synthesize OOD recovery data, with corrective actions and consistency-guided filtering to address misleading supervision and hallucinated trajectories.

    Source: §I · p. 2

  • Results

    93.3% success in soft bag pushing is reported with only five demonstrations.

    Source: §I · p. 2

  • Takeaways & Limitations

    Real-world evaluations on four diverse tasks support WM-DAgger as a framework for learning recovery behaviors without extra human demonstrations.

    Source: §V · p. 8

Abstract

from arXiv · show

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.

I. INTRODUCTION

Imitation learning can compound small policy errors into failures in out-of-distribution states, while DAgger’s reliance on human intervention limits scalability. WM-DAgger uses world models to synthesize recovery data and reports strong real-world results with few demonstrations.

I. INTRODUCTION

  • Small policy inaccuracies can lead to unseen states and further errors, eventually causing task failure.These out-of-distribution states are not covered by human demonstrations.Source: §I · p. 1
  • DAgger uses human intervention to collect recovery data, but manual operation limits scalability; single-frame synthesis also misses continuous recovery dynamics.Diffusion-based approaches are also described as struggling with complex physical interactions and deformable objects.Source: §I · p. 1
  • WM-DAgger uses world models to synthesize continuous OOD recovery data, reducing the human effort required for DAgger-style collection.World models use historical frames and intended actions to synthesize future frames.Source: §I · pp. 1–2
  • Corrective actions follow the expert trajectory’s direction, while consistency filtering validates synthesized data against real expert frames to address misleading labels and hallucinations.The filtering module compares each synthesized trajectory’s terminal frame with the corresponding expert frame.Source: §I · p. 2
  • 93.3% success in soft bag pushing is reported with only five demonstrations.Real-world evaluations across multiple tasks are reported.Source: §I · p. 2

II. RELATED WORK

Prior approaches address distribution shifts through expert feedback or generated data, while world models offer a way to represent environment dynamics. WM-DAgger trains an eye-in-hand world model, synthesizes and filters recovery trajectories, and aggregates validated data with expert demonstrations.

A. Dataset Aggregation for Imitation Learning

  • DAgger and its variants address OOD recovery through expert feedback, but remain dependent on costly human supervision; DMD synthesizes OOD data without physical awareness.HG-DAgger and CR-DAgger use controlled gating and compliant intervention, respectively.Source: A. Dataset Aggregation for Imitation… · p. 2

B. World Models in Robot Learning

  • World models predict action consequences, and prior robot-learning systems use imagined rollouts or physical-robot learning to support policy development.The cited examples include Dreamer, Cosmos, DayDreamer, and World4RL.Source: B. World Models in Robot Learning · p. 2

III. METHODOLOGY

  • WM-DAgger trains an eye-in-hand action-conditioned world model on explorative play data and expert task data, then synthesizes, filters, and aggregates OOD recovery trajectories.The corrective-action module deviates from expert trajectories and steers back; the filtering module discards trajectories with physically implausible hallucinations.Source: §III · pp. 2–3

A. Eye-in-Hand Action-Conditioned World Model Design

  • Figure 3 presents the EAC-WM architecture, an eye-in-hand world model used to synthesize visual dynamics in the recovery-data pipeline.The pipeline trains an eye-in-hand, action-conditioned world model before generating recovery data.Source: §III · p. 2, Figure 3 · p. 3

1) Overall Architecture:

EAC-WM predicts eye-in-hand visual dynamics from robot actions and historical observations. Action2Image converts sparse actions into dense, pixel-aligned geometric conditions for the world model.

1) Overall Architecture:

  • EAC-WM uses an Action2Image module to condition the world model with fine-grained geometry for predicting eye-in-hand visual dynamics.The architecture is built on GE-Sim with a Cosmos-Predict2.5 (2B) backbone.Source: 1) Overall Architecture · p. 3
  • The model predicts subsequent visual states from a context window of historical observations and a sequence of future actions.The robot control input comprises Cartesian translation, quaternion orientation, and continuous gripper state.Source: 1) Overall Architecture · p. 3

2) Action2Image Conditioning:

  • Action2Image projects sparse actions into dense, pixel-aligned geometric conditions to bridge the gap between low-dimensional actions and visual inputs.It encodes relative translation and per-pixel directional shifts, then adds gripper state as a feature map.Source: 2) Action2Image Conditioning · p. 3

3) Video-Action Tokenization:

EAC-WM represents visual and action inputs in a shared latent space and is trained to predict future video-action sequences. Its data strategy combines foundation-model pretraining with adaptation to manipulation dynamics.

3) Video-Action Tokenization:

  • A pretrained VAE maps historical and target frames to visual tokens, while dense geometric conditions are down-sampled into action tokens for shared latent-space processing.This unified tokenization facilitates cross-modal interaction within the world foundation model.Source: 3) Video-Action Tokenization · pp. 3–4

4) World Model Training and Data Strategy:

  • EAC-WM starts from Cosmos-Predict2.5 (2B) pretrained on internet-scale video-only data, then adapts to manipulation dynamics using Play Data and Task Data.Play Data consists of unscripted human exploration; Task Data comprises expert demonstrations.Source: 4) World Model Training and Data… · p. 4
  • Video-action post-training uses historical context tokens as unnoised anchors and future tokens interpolated with Gaussian noise as targets for a Video Diffusion Transformer.The transformer predicts the velocity field, and training minimizes its mean squared error with adaptive weighting.Source: 4) World Model Training and Data… · p. 4

B. World Model for Data Aggregation

WM-DAgger uses a trained world model to synthesize OOD recovery data, deriving corrective actions before filtering generated data for physical and visual consistency.

B. World Model for Data Aggregation

1) Corrective Action Synthesis Module:

  • For each expert demonstration, a random pivot and deviation direction define an OOD perturbation, while directions less than 120° from the next expert action are rejected.The synthesized deviation displacement is set to the average action displacement in Task Data to match expert action speed.Source: 1) Corrective Action Synthesis Module · p. 4
  • Each synthesized trajectory has a deviation phase away from the expert pose and a reverse-direction recovery phase returning toward the expert manifold.The world model predicts visual states from historical frames and the synthesized action sequence.Source: 1) Corrective Action Synthesis Module · p. 4

B. World Model for Data Aggregation

  • Only the recovery phase is retained for policy training, so the policy learns to recover from OOD states rather than imitate the deviation.Figure 5 also depicts the filtering module and retained frames.Source: Figure 5 · p. 4

2) Consistency-Guided Filtering Module:

Consistency-guided filtering screens synthesized trajectories for hallucinations before they join expert demonstrations in policy training. Evaluation spans four manipulation tasks, with cross-task baseline comparisons and soft-bag-pushing ablations described.

2) Consistency-Guided Filtering Module:

  • The filter uses the terminal synthesized frame, which is most temporally distant from conditioning frames, as a proxy for rollout quality by comparing it with the expert frame at the same viewing location.A pretrained DINOv2 encoder computes cosine similarity, and trajectories below the average similarity score are discarded.Source: 2) Consistency-Guided Filtering Module · p. 5
  • Figure 5 examples show discarded frames with image morphing or incorrect object position, alongside retained frames without significant hallucination.These examples include soft-bag-pushing and pick-and-place demonstrations.Source: 2) Consistency-Guided Filtering Module · p. 5

C. Policy Training

  • The augmented training set combines expert demonstrations with synthesized trajectories, and training on it is intended to impart expert behavior and OOD recovery ability.The policy predicts action chunks from observations and minimizes mean squared error against training-set actions.Source: C. Policy Training · p. 5

IV. EVALUATION

  • Evaluation covers soft bag pushing, pick-and-place, ballot insertion, and towel folding across 6-DoF control, rigid and deformable objects, novel instances, and contact-rich precision tasks.The setup includes an eye-in-hand camera and 6-DoF pose tracking; the paper describes success-rate comparisons across tasks and soft-bag-pushing ablations.Source: §IV · p. 5, Figure 6 · p. 5

A. Experiment Setups

Experiments use an eye-in-hand robotic setup and EAC-WM with a Gr00t N1.5 policy. Each task uses collected play and expert data, with 1,500 recovery episodes generated by default.

1) Hardware Platform:

  • The setup pairs a handheld two-finger gripper and eye-in-hand GoPro with a UR7e arm, Robotiq gripper, and HTC Vive pose capture.Models are trained on four NVIDIA L20 GPUs and inferred on one.Source: 1) Hardware Platform · p. 5

2) World Model and Policy Architectures:

3) Baselines:

  • EAC-WM frames are conditioned on actions in six cardinal directions, while its comparison with DMD shows greater structural and physical consistency.The visual comparison highlights EAC-WM simulating the bag dropping at t=16.Source: Figure 8 · p. 6, Figure 7 · p. 6

4) Data Preparation:

  • Each task uses five minutes of play data and 20 human demonstrations, with 1,500 recovery-action episodes generated by default.Source: 4) Data Preparation · p. 6

B. Visualization of Generated Data

EAC-WM generates action-conditioned observations with visual and physical consistency, including deformable-object dynamics. In soft bag pushing, WM-DAgger reports higher success rates than BC and DMD, including 73.3% with one demonstration.

B. Visualization of Generated Data

  • EAC-WM generates action-conditioned observations that capture soft-bag deformation during pushing and towel wrinkling during folding.The visualizations begin with a real expert frame and show observations conditioned on actions in six spatial directions.Source: B. Visualization of Generated Data · p. 6
  • EAC-WM preserves structural and physical integrity better than DMD in soft bag pushing, including the bag-dropping dynamics at t=16.Both models receive the same action sequences, with predictions compared at t=1, 6, 11, and 16.Source: B. Visualization of Generated Data · p. 6

1) Task Description:

  • The soft-bag task requires pushing a randomly initialized deformable bag across the workspace into a target basket, with 30 online validation runs.Validation covers ten bag locations with three trials per location.Source: 1) Task Description · p. 6

2) Overall Results:

3) Impact of Scaling Generated Data:

The section examines how synthetic-data scale, play-data use, and two WM-DAgger modules affect policy success in soft bag pushing.

3) Impact of Scaling Generated Data:

4) Impact of Play Data:

  • Without play data, success falls from 96.7% to 83.3%, indicating that diverse environmental interactions help the world model learn physical dynamics.Source: 4) Impact of Play Data · p. 7

5) Module Ablation Studies:

  • 0.0% success rate follows removal of the directional constraint, as unconstrained recovery actions can drive policy supervision away from the task objective.The reported failure involves erratic movements in arbitrary directions rather than toward the soft bag.Source: 5) Module Ablation Studies · p. 7, Table III · p. 7
  • 46.7% success rate results when filtering is bypassed, as visually and physically inconsistent generated frames can bias the training set.Source: 5) Module Ablation Studies · p. 7, Table III · p. 7

D. Task 2: Pick-and-Place

The pick-and-place task tests grasping and depositing rigid objects into a target box, including objects excluded from training and data aggregation. WM-DAgger reports moderate success on unseen objects where BC success rates are low.

1) Task Description:

  • The task requires reaching, grasping, and depositing rigid objects; evaluation covers three seen and two unseen objects, with 30 trials per object.Source: 1) Task Description · p. 7

2) Online Validation Results:

3) Generalizability across Unseen Objects:

2) Online Validation Results:

WM-DAgger improves success in ballot insertion and towel folding, while the conclusion reports gains across four diverse real-world tasks. The towel experiment tests performance in a 6-DoF action space with severe object-shape variation.

2) Online Validation Results:

1) Task Description:

  • The towel-folding task requires grasping and folding a deformable towel through complex 6-DoF spatial trajectories.It evaluates scalability to more complex action spaces and severe object-shape variation.Source: 1) Task Description · p. 8

2) Online Validation Results:

V. CONCLUSION AND FUTURE WORK

  • WM-DAgger combines corrective-action synthesis and consistency-guided filtering to produce physically consistent recovery data, with reported success-rate improvements across four diverse real-world tasks.The conclusion says this supports learning robust recovery behaviors without expensive extra human demonstrations.Source: §V · p. 8
  • Scaling beyond two-finger grippers to dexterous multi-finger hands remains challenging because high degrees of freedom complicate consistent frame synthesis.The authors propose adding morphological priors and kinematic topologies to world-model conditioning.Source: §V · p. 8

Related papers

Loading 2604.11351v1…