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

TL;DR

WM-DAgger synthesizes OOD recovery data using World Models, achieving a 93.3% success rate in soft bag pushing with five demonstrations.

cs.RO 🔴 Advanced 2026-04-13 5 views
Anlan Yu Zaishu Chen Peili Song Zhiqing Hong Haotian Wang Desheng Zhang Tian He Yi Ding Daqing Zhang
imitation learning data aggregation world models robotic manipulation unsupervised learning

Key Findings

Methodology

WM-DAgger synthesizes OOD recovery data using World Models, eliminating the need for human intervention. It employs a Corrective Action Synthesis Module and a Consistency-Guided Filtering Module to ensure physical consistency and task orientation of the generated data.

Key Results

  • In the soft bag pushing task, WM-DAgger achieved a 93.3% success rate with only five demonstrations, significantly outperforming standard Behavioral Cloning and DMD methods.
  • Validated the method's effectiveness across multiple real-world tasks, significantly improving success rates.
  • Ablation studies confirmed the critical roles of the Corrective Action Synthesis Module and Consistency-Guided Filtering Module.

Significance

This research significantly enhances the scalability of imitation learning by reducing reliance on human intervention. By synthesizing high-quality OOD recovery data, it addresses the compounding error problem in traditional DAgger frameworks, holding significant implications for the field of robotic learning.

Technical Contribution

Introduces an unsupervised data synthesis method based on World Models, significantly enhancing the robustness of imitation learning. The introduction of the Corrective Action Synthesis Module and Consistency-Guided Filtering Module ensures the physical consistency of generated data.

Novelty

First to utilize World Models to synthesize OOD recovery data, eliminating the need for manual labeling. Compared to existing methods, it significantly improves the physical consistency and task orientation of the data.

Limitations

  • In complex physical interaction scenarios, the generated data may still exhibit some physical inconsistencies.
  • High quality of training data for World Models is required.

Future Work

Future work could explore applying this framework in more complex physical interaction scenarios and integrating with other learning paradigms to enhance model generalization.

AI Executive Summary

Imitation learning is crucial for training robotic policies but is limited by compounding errors. Existing DAgger frameworks rely on human intervention, limiting scalability. WM-DAgger addresses this by using World Models to synthesize OOD recovery data. The method was validated across multiple real-world tasks, significantly improving success rates, notably achieving a 93.3% success rate in soft bag pushing with only five demonstrations.

WM-DAgger introduces a Corrective Action Synthesis Module and a Consistency-Guided Filtering Module to ensure the physical consistency and task orientation of the generated data. By reducing reliance on human intervention, the method significantly enhances the scalability of imitation learning, holding significant implications for the field of robotic learning.

Despite significant achievements, the method still faces challenges in complex physical interaction scenarios. Future work will explore applying this framework in more complex scenarios and integrating with other learning paradigms to enhance model generalization.

Deep Analysis

Background

Imitation learning is a vital paradigm for transferring human expertise to robotic systems. However, its effectiveness is often limited by distribution shifts and compounding execution errors in real-world settings. The DAgger framework addresses this by iteratively incorporating expert feedback to recover from OOD states, but its reliance on human intervention limits scalability.

Core Problem

The core problem in imitation learning is compounding errors: minor policy inaccuracies can drive robots into unseen out-of-distribution states, leading to larger errors and eventual task failure.

Innovation

WM-DAgger synthesizes OOD recovery data using World Models, eliminating the need for human intervention. It introduces a Corrective Action Synthesis Module and a Consistency-Guided Filtering Module to ensure the physical consistency and task orientation of the generated data.

Methodology

  • �� Use World Models to synthesize OOD recovery data, reducing the need for human intervention.
  • �� Corrective Action Synthesis Module generates task-oriented recovery actions to prevent misleading supervision.
  • �� Consistency-Guided Filtering Module compares synthesized frames with real frames to discard physically implausible trajectories.

Experiments

WM-DAgger was validated across multiple real-world tasks, including soft bag pushing, pick-and-place, ballot insertion, and towel folding. Experiments used standard Behavioral Cloning and DMD as baselines, evaluating success rates and model robustness.

Results

WM-DAgger achieved a 93.3% success rate in the soft bag pushing task, significantly outperforming standard Behavioral Cloning and DMD methods. Ablation studies confirmed the critical roles of the Corrective Action Synthesis Module and Consistency-Guided Filtering Module.

Applications

This method can be directly applied to robotic manipulation tasks requiring efficient data aggregation, such as automated sorting in logistics and robotic learning in complex physical interaction scenarios.

Limitations & Outlook

Generated data may still exhibit some physical inconsistencies in complex physical interaction scenarios. High quality of training data for World Models is required. Future work will explore applying this framework in more complex scenarios.

Plain Language Accessible to non-experts

Imagine a robot learning to cook in a kitchen. Traditional methods require a chef to constantly guide it, telling it what to do at each step. WM-DAgger is like a smart assistant that can figure out how to handle unexpected situations, like a pot falling or running out of ingredients, by watching the chef a few times. It simulates the kitchen environment to generate possible solutions without needing the chef to be there all the time.

ELI14 Explained like you're 14

Imagine you're playing a robot game where the robot needs to learn tasks like pushing a soft bag into a basket. Usually, you have to keep telling it what to do, but that's a hassle. WM-DAgger is like a super-smart game assistant that only needs to watch you play a few times to learn how to handle unexpected situations, like the bag falling. It simulates the game environment to figure things out on its own. Isn't that cool?

Glossary

Imitation Learning

A method of training robotic policies by mimicking human expert behavior.

Used to transfer human expertise to robotic systems.

Data Aggregation

The process of collecting and synthesizing data to improve model training.

Used in the DAgger framework to generate recovery data.

World Models

Internal models that simulate environmental dynamics to help agents predict the outcomes of their actions.

Used to synthesize OOD recovery data, reducing human intervention.

Corrective Action Synthesis Module

Generates task-oriented recovery actions to prevent misleading supervision.

Used in WM-DAgger to generate physically consistent recovery data.

Consistency-Guided Filtering Module

Compares synthesized frames with real frames to discard physically implausible trajectories.

Ensures the physical consistency of generated data.

Open Questions Unanswered questions from this research

  • 1 How can WM-DAgger be applied in more complex physical interaction scenarios?
  • 2 How to improve the generalization ability of generated data across diverse tasks?

Applications

Immediate Applications

Logistics Automation

Implement automated sorting in logistics, reducing reliance on manual operations and increasing efficiency.

Long-term Vision

Robotic Learning in Complex Physical Interactions

Achieve efficient robotic learning in complex physical interaction scenarios, advancing robotics technology.

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.

cs.RO