Dream to Manipulate: Compositional World Models Empowering Robot Imitation Learning with Imagination

TL;DR

This paper introduces DreMa, a novel digital twin-based world model that enables robots to imagine and predict environmental changes, significantly improving imitation learning efficiency and enabling one-shot physical task learning on real robots.

Contribution

DreMa bridges digital twins and world models by integrating explicit representations with physics simulators, enhancing realism and generalization in robot imitation learning.

Findings

DreMa improves accuracy and robustness in environment modeling.

It reduces data requirements for policy learning.

Enables one-shot learning on a real robot.

Abstract

A world model provides an agent with a representation of its environment, enabling it to predict the causal consequences of its actions. Current world models typically cannot directly and explicitly imitate the actual environment in front of a robot, often resulting in unrealistic behaviors and hallucinations that make them unsuitable for real-world robotics applications. To overcome those challenges, we propose to rethink robot world models as learnable digital twins. We introduce DreMa, a new approach for constructing digital twins automatically using learned explicit representations of the real world and its dynamics, bridging the gap between traditional digital twins and world models. DreMa replicates the observed world and its structure by integrating Gaussian Splatting and physics simulators, allowing robots to imagine novel configurations of objects and to predict the future consequences of robot actions thanks to its compositionality. We leverage this capability to generate new data for imitation learning by applying equivariant transformations to a small set of demonstrations. Our evaluations across various settings demonstrate significant improvements in accuracy and robustness by incrementing actions and object distributions, reducing the data needed to learn a policy and improving the generalization of the agents. As a highlight, we show that a real Franka Emika Panda robot, powered by DreMa's imagination, can successfully learn novel physical tasks from just a single example per task variation (one-shot policy learning). Our project page can be found in: https://dreamtomanipulate.github.io/.