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

Authors: Leonardo Barcellona, Andrii Zadaianchuk, Davide Allegro, Samuele Papa, Stefano Ghidoni, Efstratios Gavves
Primary area: applications to robotics, autonomy, planning
Keywords: World model;  Imagination; Imitation Learning; Gaussian Splatting; Compositional; Physics-informed; Object-centric;

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/.

Decision: accept (poster)
Reviewer scores: 5, 8, 8

Reviews and discussion:
<thread>
creation date:
1729451851072

thread_type:
official_review

<thread_object>
creator: reviewer_mne8

<content>
summary:
The paper presents a novel paradigm for constructing world models that serve as explicit representations of real-world environments and their dynamics. By integrating advances in real-time photorealism, such as Gaussian Splatting, with physics simulators, the authors propose a system capable of generating new data for imitation learning. Additionally, the paper demonstrates the application of this model in real-world scenarios, showing how data collected from the world model can be used to train robots via imitation learning, with promising results when transferring learned behaviors to real-world tasks.

**Strengths:**

1. The paper introduces an innovative approach by leveraging world models to generate robotic data for imitation learning, which is a contribution to the field.
2. The experiments are detailed, covering both simulation environments and real-world robot demonstrations, providing a robust evaluation of the approach.
3. A creative method for augmenting data used in imitation learning is introduced, which could lead to improved learning efficiency.

**Weaknesses:**

1. The absence of publicly available source code limits the reproducibility of the results. It is suggested to release the code during the rebuttal stage.
2. Some figures in the paper need improvement, as the text in several instances is too small to read clearly.
3. The predictions demonstrated in the paper are limited to simple tasks and physics environments, and future work should focus on extending these predictions to more challenging tasks and complex physical simulations.

In conclusion, the paper presents a compelling framework that blends world modeling with imitation learning, but there are areas for improvement, particularly in terms of figure clarity, task complexity, and providing source code for reproducibility.

strengths:
1. The paper introduces an innovative approach by leveraging world models to generate robotic data for imitation learning, which is a contribution to the field.
2. The experiments are detailed, covering both simulation environments and real-world robot demonstrations, providing a robust evaluation of the approach.
3. A creative method for augmenting data used in imitation learning is introduced, which could lead to improved learning efficiency.

weaknesses:
1. The absence of publicly available source code limits the reproducibility of the results. It is suggested to release the code during the rebuttal stage.
2. Some figures in the paper need improvement, as the text in several instances is too small to read clearly.
3. The predictions demonstrated in the paper are limited to simple tasks and physics environments, and future work should focus on extending these predictions to more challenging tasks and complex physical simulations.

questions:
Could you please show some performance results in more complex physical environments and challenging tasks? Even if they were unsuccessful, it would be helpful to see such results, even though they are not included in the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698436755

thread_type:
official_review

<thread_object>
creator: reviewer_kqvg

<content>
summary:
The paper augments a small set of demonstrations with imagination to improve few-shot imitation learning in both RL bench and real-world robotic tasks. The imagination comes from learning compositional objects models through Gaussian Spatting and replaying demonstrations with varied objects poses in a simulator.

strengths:
The paper presents a novel strategy of data augmentation by first acquiring object models and then leveraging simulations to ensure correct dynamics of imagined demonstrations instead of learning worlds models of both objects and dynamics simultaneously. The strategy is shown to have meaningful improvement on few-shot imitation performance in sufficient sim and real tasks. The approach is also thoroughly invested with ablations showing the significance of the imaging with roto-translation of original demos. The paper is well written with clear motivations and goals and sufficient results to support the claim.

weaknesses:
1. The paper should compare to other data augmentation approaches such as MimicGen or Digital Twin, which is significant extensive line of work worthy of more elaboration and discussion. Similar approach such as [1] should be discussed or cited. 
2. The method relies on open-vocabulary tracking models to segment objects, which limit the approach to non-articulated objects. It is unclear how such segmentation models can capture individual robot link or object parts connected with articulated joints accurately. Also, it is unclear how to extend the simulator to incorporate articulated objects interaction after the parts have been learned. 
3. Existing approach of verify imagined demonstration is rudimentary. 

[1] MIRA: Mental Imagery for Robotic Affordances

questions:
1. What’s the relationship with digital twin line of work?
2. How to learn 3d models for parts of articulated objects and how to imagine demonstrations that manipulate articulated objects?
3. Can you replay imagined trajectories in simulator to verify correctness? Will that help improve imitation success?
4. How does error build up in the pipeline of segmenting object masks -> learning objects models through Gaussian Spatting -> imagine demonstrations -> learning policy? Perhaps some ablations or quantitative metrics to measure error will be useful! 
5. Can you add one more baseline method of data augmentation to compare to?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730706267503

thread_type:
official_review

<thread_object>
creator: reviewer_zmyw

<content>
summary:
This paper proposes DreMa, which integrates object-centric Gaussian splatting with a rigid-body physics simulator, to “imagine” new demonstrations to train imitation learning models. These imagined demonstrations are obtained in simulation by applying robot transformations and rotations to objects extracted from Gaussian splatting.

strengths:
- The paper tackles the data efficient regime, and their proposed approach uses a single demonstration

- The paper validates the proposed approach with real-world experiments, and demonstrates a system that can perform manipulation tasks

weaknesses:
My main issue with the paper is that the work is positioned within the “world model” literature, while the proposed method seems to fall under real2sim and data augmentation strategies. The “world model” is used to generate an augmented set of demonstrations in simulation in order to train an imitation learning model offline, and it is not used during online control. The set of equivariant transformations used to generate the augmented demo set is hand-designed, and likely task-specific. 

The simulated results would be more convincing if they were expanded to include more than 3 RLBench tasks, but the method is limited to non-articulated objects. Additionally, the real-world tasks seem to only include blocks or boxes. How does DreMa perform with more complex shapes? How does DreMa perform when the scene includes a mix of objects with different shapes? 

- It is unclear what base imitation learning model is used by the proposed DreMa method. Is this just PerAct? Are observations to the policy directly captured from cameras, or by rendering Gaussian splats?

- There are no metrics on runtime performance of the proposed method, while the introduction mentions the “real-time” performance of Gaussian splatting. 

- Table 1: Comparisons to PerAct (Shridhar et al., 2023) are done using only 5 episodes per task, while PerAct uses 25 episodes per task. 

- The related work section would benefit from a broader discussion of particle-based simulation approaches, such as: https://arxiv.org/abs/1810.01566, https://arxiv.org/abs/2312.05359, https://arxiv.org/abs/2405.01044, https://arxiv.org/abs/2303.05512. A comparison to ManiGaussian (https://arxiv.org/abs/2403.08321) would also be helpful.

questions:
Figure 2: It would be helpful to include how “open-vocabulary tracking models” fit into the pipeline. 

L107: “Two recent works demonstrated predicting future states can be applied to robotics” - This statement can be made more precise, since predicting future states / modeling forward dynamics for control is not a new idea in robotics.

L128: Why is (Cheng et al., 2023) used as a reference for the zero-shot capabilities of foundation models?

L157: There is some overloading of the term “manipulation”, which in the context of the paper seems to refer to a “manipulable” or controllable world model, rather than a world model explicitly designed for robot manipulation tasks. 

L174: Set $\mathcal{X}$ notation is used for a sequence.

L267-268: Other objects in the world that the robot arm may interact with could also be articulated? ie. cabinets 

L409: Is the PerAct baseline in the multi-task setting only trained on the subset of RLBench tasks you’re evaluating on? How many demonstrations were used? Was PerAct trained with data augmentation?

L412-413: Was this model selection using a validation set done over the entire course of training? How does this compare to just using the final model after training for 600k steps?

L418: Why “112, 61, and 81” demonstrations for the three tasks? How were these number of “imagined” demonstrations chosen? 

L485: How were OOD test conditions chosen? Would they be within the set of equivariant transformations evaluated in Table 2? 

[minor editing]

Figure 2: part 3, typo in Gaussian

L105: “gaussian” -> “Gaussian”

L233: $x_{n,k}$ should be $y$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733313527231

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response and summary of rebuttal

comment:
Dear Reviewers, Dear AC,

Thank you for all your help to bring this paper forward. We believe it introduces a new way to ground world representations with the physical world, and trying to fit into a single and specific category is not the point. We think this is evident from the substantial improvements (up to 50% absolute improvements in some tasks, 33.3% on average) with few-shot imitation learning with real robots, not only on the original 3 tasks, but also 2 more that we run during the rebuttal.

Reviewers kqVG and Mne8 acknowledge our message and positioning (Mne8: “Congratulations, I think your work is valuable.”).

As far as we can tell, the only remaining difference in perspective
is with Reviewer zmyw on the positioning: is the proposed model a world model (as we claim), a data augmentation model or a real2sim model. We summarize why we think our proposed model should be better thought of as a world model, and why a ‘data augmentation’ or ‘real2sim’ positioning does not perfectly fit.

**Why call it a world model approach?**

- A world model is a learned representation, and our work includes learning components.

-  Future state predictions, observations and actions are key characteristics present in our approach.

- We demonstrated that the agent could learn directly within the learned its “own hallucination" [5].
- Our approach offers benefits beyond planning, including offline training in imitation learning, by leveraging world model predictions from unseen states.

- Prior works such as [6] and [7] support using world model definitions in our context.

 - As we discussed in Sections 2.1 and 2.2, along with Appendix E, there is a strong connection between world models in robotics and real2sim approaches. 

- Existing world models also include hand-designed components, such as reward functions and inductive biases in neural architectures. 

- DreMa autonomously generates imagined data for new tasks without manual adjustments, consistently improving baseline performance (+9.1%) and outperforming the traditional augmented baselines (+5%).

- The need to collect new observations when transitioning between environments is also common in world models, including DreamerV2 [1], which requires re-optimizing its network when switching between Atari games. Similarly, our approach gathers observations and re-optimizes when encountering new tasks and  environments.

**Why not call it only a data augmentation approach?**

- Traditional augmentation modifies existing data (e.g., patches or color changes). In contrast, our method generates entirely new data based on DreMa’s state predictions, creating novel trajectories using learned representations.

- The DreMa process is conceptually distinct from standard augmentation, as it involves imagining new possibilities rather than altering existing ones.

**Why not call it only a real2sim approach?**

- Real2Sim aims to create high-fidelity simulations, often requiring manual intervention [2] or not using the reconstructed environment actively [3]. Our approach differs by leveraging state predictions from DreMa to facilitate task learning, focusing on autonomous generation of novel demonstrations rather than recreating precise environmental replicas.

- While world models may seem real2sim-like (e.g., DreamerV2 reconstructing aspects of Atari games), they autonomously learn from observational data to enable agent learning. Applications like DayDreamer [4] blur this distinction further when used with real robots.

In the end, we think that an added value comes from comparing different sub-fields, finding what is common between them rather than what is not, and bridging communities.

Thank you for your constructive insights, and we hope this clarifies our perspective.

The Authors

[1] Hafner, Danijar, et al. "Mastering Atari with Discrete World Models." International Conference on Learning Representations 2021.

[2] Marcel Torne, Anthony Simeonov, Zechu Li, April Chan, Tao Chen, Abhishek Gupta, and Pulkit Agrawal. Reconciling reality through simulation: A real-to-sim-to-real approach for robust manipulation. arXiv preprint arXiv:2403.03949, 2024

[3] Allan Zhou, Moo Jin Kim, Lirui Wang, Pete Florence, and Chelsea Finn. NeRF in the palm of your hand: Corrective augmentation for robotics via novel-view synthesis. CVPR 2023

[4] Wu, Philipp, et al. "Daydreamer: World models for physical robot learning." Conference on robot learning. PMLR, 2023.

[5] Ha, David, and Jürgen Schmidhuber. "World Models." arXiv preprint arXiv:1803.10122, 2018.

[6] Yang, Mengjiao, et al. "Learning Interactive Real-World Simulators." The Twelfth International Conference on Learning Representations, 2024.

[7] Abou-Chakra, Jad, et al. "Physically Embodied Gaussian Splatting: A Realtime Correctable World Model for Robotics." arXiv preprint arXiv:2406.10788, 2024.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a novel method for constructing photo-realistic 3D simulations (world models) of real-world scenes using object-centric Gaussian splatting. The resulting simulation is then used to generate data for augmenting few-shot imitation learning.

The reviewers agree that the method is novel and impactful. However, key concerns raised include: (1) the positioning of the paper as a world-model approach versus a real2sim data augmentation method, (2) insufficient comparisons with relevant baselines, and (3) limited evaluation on simple pick-and-place tasks. Concerns regarding the paper's clarity were addressed during the discussion phase.

Regarding point (1), I find that the current presentation effectively communicates the method's intent and does not risk confusing the audience. Nonetheless, I concur with reviewer zmywwe's suggestion to downplay the emphasis on the world-model aspect to avoid potential misalignment with the paper's contributions.

On point (2), the evaluation against baselines and design choice ablations appear insufficient for the following reasons:
- The choice of PerAct as the downstream policy learning algorithm does not fully capitalize on the photo-realistic renderings generated by Gaussian splatting, raising questions about whether this choice aligns with the method's strengths.
- Additionally, the same data augmentation approach could potentially be applied to segmented voxels or point clouds, suggesting that the benefits of Gaussian splatting for this task remain underexplored in comparison.

Point (3) is a major limitation, the authors should further clarify and elaborate how they envision the proposed approach may generalize to articulated objects. 

Overall, I believe addressing these gaps in evaluation and contextualization would strengthen this paper further.