Dream to Manipulate: Compositional World Models Empowering Robot Imitation Learning with Imagination
This work addresses the challenge of data-efficient and generalizable robot imitation learning for real-world applications, though it builds incrementally on existing world model and digital twin concepts.
The paper tackles the problem of enabling robots to learn manipulation tasks from limited demonstrations by proposing DreMa, a compositional world model that acts as a learnable digital twin, integrating Gaussian Splatting and physics simulators to generate novel data for imitation learning. The result includes significant improvements in accuracy and robustness, with a real robot achieving one-shot policy learning from a single example per task variation.
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/.