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
Venue: ICLR 2025 Poster (2025)
Keywords: World model;  Imagination; Imitation Learning; Gaussian Splatting; Compositional; Physics-informed; Object-centric;
Primary area: applications to robotics, autonomy, planning

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

Published as a conference paper at ICLR 2025
DREAM TO MANIPULATE: COMPOSITIONAL WORLD MODELS EMPOWERING ROBOT IMITATION LEARNING WITH IMAGINATION
Leonardo Barcellona 1.2+ Andrii Zadaianchuk
Davide Allegro
Samuele Papa*
Stefano "hidoni' Efstratios Gavves
University of Padua, Italy Politecnico di Torino, Italy SUniversity Amsterdam, The Netherlands "Archimedes/A Laatddess//Attnt RC, Greece
ABSTRACT
A world model provides an agent with representation of its environment, en- abling it to predict the causal consequences of its actions. Current world models typically cannot directly and explicitly imitate the actual environment in front of robot, often resulting in unrealistic behaviors and hallucinations that make them unsuitable for real-world robotics upplications. To overcome those chal- lenges, 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, bridg- ing the gap between traditional digital twins and world models. DREMA repli- cates 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 composition- ality. We leverage this capability to generate new data for imitation learning by applying equivariant transformations to small set of demonstrations. Our eval- uations across various settings demonstrate significant improvements in accuracy and robustness by incrementing actions and object distributions, reducing the data needed to learn 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 single exam- ple per task variation (one-shot policy learning). Our project page can be found in: https: drr.mt oman www.ttaaiiiaaaaamaamaammmmaaaaae git.ub.iiii io/.
INTRODUCTION
World models are learnable representations of the real world that an agent can use to predict the consequences of its actions (Ha & Schmidhuber, 2018). They predict future states given the current state and an action, empowering the agents to learn new skills from the inferred evolution of the model (Hafner et al., 2019; 2020; Gumbsch et al., 2024; Hansen et al., 2024; Yang et al., 2023). Thus, they are fundamental building block for robots to imagine, a capability that we humans exploit to interact with new environments (Hawkins et al., 2017; Ha & Schmidhuber, 2018).
Applying world models to real robots implies rendering realistic observations (Yang et al., 2023; Zhou et al., 2024) while simulating the dynamics (Li et al., 2024). Whereas most of the world models learn a statistical representation of the dynamics (Hafner et al., 2019; Gumbsch et al., 2024: Hansen et al., 2024; Yang et al., 2023), is unrealistic to expect any training distribution to encompass all possible configurations of the world. This limits the current literature, as agents need the ability to imagine beyond direct experience. We identify three requirements needed by an effective world model for robotic agent. First, it must be inherently compositional, allowing the generation of novel and valid combinations of learned concepts (Zhou et al., 2024). Second, must be object-centric, as the world is primarily composed of objects of varying number, type, and
"Correspondence leonardo mamaarooharreellona@pooli @oolitoo,,,, (a.zadaianchuk, e.gavves) @uva.nl

---

Published as a conference paper at ICLR 2025
Scenario
Object-centric Reconstruction
Physics Simulation
Imagination
Gaussians Update
Detraisttation
Figure Overview of imagination with DREMA, which builds a compositional manipulation world model from environment images using object-centric Gaussian Splatting to generate novel demonstrations by transforming real ones.
configuration (Spelke, 1990; Kipf et al., 2019). Third, manipulating the model should be simple to ensure that actions and consequences can be controllably modelled (Li et al., 2024).
In this work, we propose a compositional manipulation world model, named DREMA, that allows robotic agents to exploit physically reliable predictions to generalize to novel situations. Thus, it an interactable phosice-constrained representation that bridges robot imagination with real-world demonstrations following recent advances in learnable digital twins and explicit world models, which both enable agents to simulate and reason about their environment (discussed in App. E). We shape DREMA by reconstructing the scene with object-centric Gaussian Splatting (Kerbl et al., 2023) and by embedding the obtained compositional representation into physics simulator, e.g., PyBullet (Coumans & Bai, 2016). Gaussian Splatting reconstructs 3D scenes set of independent Gaussians that can be optimized and rendered, allowing for extremely high resolution and real-time graphics rendering. We infer compositional properties of rendered Gaussian Splats using foundational models (Ravi et al., 2024; Ren et al., 2024), allowing the agent to manipulate the object-centric model to create new scenes.
Through DREMA, it is possible to execute robot policies directly in imagination while rendering realistic observations. This capability benefits imitation learning, whose objective is to solve tasks given example demonstrations (James et al., 2022). To ensure reliable generalization of the learned policy, demonstrations should contain sufficient variations, which is impractical and time-consuming with real applications (Mandlekar et al., 2023). We show that imposing equivariant transformations, namely roto-translation of the demonstrated actions and the objects, DREMA produces new high-quality data inside the world model and reduces the required demonstrations. We exploit imagination to generate new demonstrations that augment the original ones, while also using the proposed world model to verify that the new actions result in successful task execution and rendering of realistic observations.
Our contributions are: (1) We the first to propose a compositional manipulation world model, DREMA, that empowers robots with the ability to imagine in a grounded manner; (2) We propose to exploit the compositional structure of the world model to generate novel, valid, and realistic training demonstrations for imitation learning; (3) We empirically demonstrate in simulations and with real robots that DREMA improves agent generalizability even with minimal number of demonstrations.
2
RELATED WORK
2.1 WORLD MODELS
World models were originally conceived to improve policy learning in reinforcement learning by predicting the consequences of actions on the environment (Ha & Schmidhuber, 2018 8; Hafner al., 2019; 2020; 2023). State-of-the--te -o-the--rttttttte approaches often leverage generative models to build these world models (Jia et al., 2023; Hu et al., 2023; Yang et al., 2023; Bruce et al., 2024). While initially applied to game simulations (Ha & Schmidhuber, 2018; Hafner et al., 2019) and navigation tasks (Hu et al., 2023; Jia et al., 2023), world models now being extended to robotics (Yang et al., 2023; Wu
2

---

Published as a conference paper at ICLR 2025
et al., 2023; Zhou al., 2024). For example, UniSim (Yang et al., 2023) uses diffusion to generate possible future frames conditioned on robot actions. DayDreamer (Wu et al., 2023) demonstrates that skills such locomotion and manipulation can be learned using Dreamer (Hafner et al., 2019), while RoboDreamer (Zhou et al., 2024) uses compositionality to generate novel actions. The main limitation of these approaches is their inability to comprehensively capture the motion dynamics and physics specific to the robot's environment, particularly when faced with scenarios outside their training domain (Yang et al., 2023). By contrast, we propose a compositional manipulation world model that integrates cuttiig-edge photorealism with Gaussian Splats and realistic dynamics with physics engines, showing we can tailor them for efficient imitation learning in real-worlh scenarios.
2.2 SIMULATION-BASEI RECONSTRUCTIONS
Thanks to recent advances in reconstruction (Kerbl et al., 2023) and in physics simula- tions (Coumans & Bai, 2016; Li al.. 2019), in the last year many have proposed to combine both for realistic simulations of actual environments (Xie et al., 2024; Jiang et al., 2024; Meyer et al., 2024; Ruan al., 2024; Torne al., 2024). Particle-based simulators are capable of modeling more complex dynamics than traditional rigid body simulators (Li al., 2019), making them beneficial for robotics (Chen et al., 2024a). Recently, neural radiance fields have been applied to implicitly learn dynamic systems (Li et al., 2023; Whitney et al., 2024). The introduction of Gaussian Splatting (Kerbl et al., 2023) has opened avenues for novel approaches that bridge these simulation techniques with high -ugh-qualtty rendering methods (Xie et al., 2024; Jiang et al., 2024). However, they are not compositional (Xie et al., 2024) object-centric (Xie et al., 2024; Jiang al., 2024). PEGASUS (Meyer et al., 2024) combines Gaussian representations with mesh simulations to improve the 6D pose estimation, while Abou-Chakra al. (2024) and Lou et al. (2024) show that Gaussian representations can be adapted to track moving objects. Both approaches do not allow agents to predict future states, rather focusing more on tracking in observed past trajectories. In contrast, we empirically demonstrate that our approach predicts useful estimation of future states. While simulating replica of the environment can benefit policy evaluation (Li et al., 2024), two recent studies showed that scene reconstruction can enhance robot policies (Ruan et al., 2024; Torne et al., 2024). Ruan et al. (2024) reconstruct the environment to deform robot trajectories, while our object-centric approach focuses on the interaction between objects and the agent. Although being close to our approach, Torne et al. (2024) requires to manually model the object-centric representation and needs adaptation between the simulation and the real world. Instead, we leverage the zero-shot capabilities of foundational models (Kirillov et al., 2023; Ren et al., 2024; Cheng et al., 2023) to automatically create a compositional world model and show the effectiveness of using transformed demonstrations for task learning without requiring extensive interaction with the world model.
al., 2024; Lou et al., 2024; Abou-Chake
2.3 IMITATION LEARNING AND DATA AUGMENTATION
Imitation learning involves acquiring task knowledge from dataset of demonstrations (James et al., 2022). Vision-basee imitation learning focuses on deriving policies directly from RGB or RGB-D images (Zeng al., 2021; Shridhar et al., 2023; 2022; Goyal et al., 2024). PerAct (Shridhar et al., 2023), for instance, voxelizes the environment to infer actions in 3D. However, despite advancements in generalization, many approaches (Ze et al., 2023; Goyal et al., 2024) still require hundreds of demonstrations to learn robust policies. MIRA (Lin et al., 2023) employs NeRF's novel view synthesis to simulate new camera views of the environment, while Lu et al. (2025) applied Gaussian Splatting for action inference. In contrast, our approach leverages Gaussian Splatting to generate not only novel views but also new environmental configurations, enhancing generalization.
Data augmentation is an established technique (Krizhevsky et al., 2012; He et al., 2016). However, its application in robotics challenging, as altering robot actions can unpredictably affect envi- ronments and tasks (Pitis et al., 2022). Despite this, Laskin et al. (2020) demonstrated that even simple augmentations can improve robot policies. Approaches such as generating counterfactual transitions (Pitis et al., 2022) and applying invariant transformations (Corrado & Hanna, 2024) enhance learning by expanding training datasets. However, these methods primarily recombine existing trajectories rather than generate entirely new configurations. In contrast, our approach imagines novel equivariant configurations of objects and actions, enabling more diverse and effec- tive data augmentation. A related method, MimicGen (Mandlekar et al., 2023), uses object-centric
3

---

Published as a conference paper at ICLR 2025
Sconc decompositio::
Ractttciiooon extractions: Training
Future prodiction: Physics-powered Imagination
Open-ocabulary segmentation tracking
Objects
Raplay
Imaggnee Demonstration
Dem Demenstration
Figure 2:
Steps create the compositional world model with DREMA: observation of the environment and scene decomposition, representation extraction and future predictions.
segments for compositional data generation but requires multiple demonstrations. Our method achieves comparable benefits while working effectively with as few as one demonstration. Whereas some authors argue that imagining various object configurations unrealistic (Zeng et al., 2021), we demonstrate that this is achievable in real-world scenarios.
3
WITH GAUSSIAN SPLATS & PHYSICS ENGINES
We want to construct compositional manipulation world model M, that is (i) explicitly composi- tional, (ii) object-centric, and (iii) controllable, allowing the agent to control all the objects inside M. The model .M- = (O,, Ac, D, TT) comprises set of K object assets in 3D, O2 = {ox.e) k== time the model of the agent Ar that manipulates the objects in the world; the dynamics of the world in the form of an operator function, (Or+1, actions of the agent A at time (- transforms the states of object assets O and agent's state A the manipulation tasks T expected from the agent. For clarity, drop the time index from ok,t when this can be inferred from the context. The input to the algorithm that constructs the compositional manipulation world model is simply RGB-D video X (xx) of N frames.
D Ar), which given the
In the following, we first discuss our compositional 3D scene representation based on Gaussian Splats, which forms the basis for world model. Afterward, we detail how we obtain each component of the world model given our 3D scene representation.
3.1 REPRESENTING 3D OBJECTS WITH GAUSSIAN SPLATTING
Central to the compositional manipulation world model the inverse graphics function g = (xX using Gaussian Splatting (Kerbl et al., 2023). The inverse graphics function returns a 3D rep- resentation the world G given a sequence of n images video (Kerbl al., 2023). The key idea behind Gaussian Splatting is to ground many Gaussian blobs g = (p, r, ar, c). Vg E g physically in the 3D space. Per Gaussian, we have p = [P..Py Py, P=] as its center its orientation r and scale s (Su, s.) (the covariance), the encoded color c. and the opacity of the color. That is, given any arbitrary 3D location and angle in space (corresponding to 2D camera plane), Gaussian Splats can render projected 2D image that resembles the true observed object (or scene). For rendering the appearance onto the 2D image, one simply a-blends the colors of the Gaussians along each of the rays that connect the 3D point to the points on the 2D projected image. the end, the Gaussian Splatting algorithm optimizes the Gaussian parameters in g to ren- der the objects and the scene photorcalistically for both known and unknown 3D locations (named novel view synthesis). Gaussian Splats have taken the graphics community by storm (Huang et al., 2024; Xie al., 2024; Meyer et al., 2024) not only for their exceptional photorealism but also for their real-time computation. For further details, we refer to the seminal work of Kerbl et al. (2023).
(II,,..,,,,, collected as
Object-centric Gaussian Splats. Gaussian Splats generate photorealistic renderings of the 3D scene that are unstructured in terms of object semantics, Instead, we are interested in object-centric representations. To obtain object-centric Gaussian Splats, together with the input 2D images in for our K objects in our world, we introduce sets of segmentation masks, Y; E y) k 1,..., K. We consider & to be varying so that not constrain in advance the number of objects in our world. Each object-set of segmentations Y, (VI, ..., Vo,k) contains the 2D masks for the k-th object in all N images in X. With segmentation masks Y₁ (y1,8)--. .r.kk) for each object & and for all n frames,
4

---

Published as a conference paper at ICLR 2025
can simply mask the respective video frames, Yn. O Zn,k- to obtain the object-centric instances, which we can then feed to the Gaussian Splatting. Since in our setting the robot is also equipped with depth sensors, we further strengthen the original Gaussian Splatting optimization by including the depth measurements x dd ε deppth for additional supervision, L = Lree + XnCn + Adepth Lueeun free and C,, are the reconstruction and normal losses defined in the original 2DGG.. The extra depth regularization, Luepth | 2d 111 uses the & norm to penalize Gaussians that have outlier depth. We set the hyperparameters Anoemal and Adepth for all experiments empirically as in Huang et al. (2024). Our object-centric inverse graphics function takes the form GE = h (X, X deeph: Y&)--
Deformable and articulated objects. is important to note that since the standard Gaussian Splats only work with (static) scenes, we cannot account for deformable articulated objects. An exciting future direction is to bootstrap on the significant progress in Dynamic Gaussian Splat- ting (Luiten et al., 2023) to automatically model also articulation (see App. G) and deformation.
Next, we describe how to decompose the scene into k sets of regions defined by segmentation masks Y' to obtain object assets and model dynamics.
3.2 DECOMPOSING THE SCENE INTO OBJECT- CENERIC REGIONS
Decomposing the high- nigh-dimensional scene representation into objects (Locatello et al., 2020; Seitzer al., 2023; Wu al., 2024) that are independently controllable greatly simplifies control and the manipulation of the dynamics in the environment (Zadaianchuk et al., 2021; 2022). Inspired by re- cent work on embeddable 3D assets (Deitke al., 2023; Wu et al., 2024). we ideally want to learn Gaussian Splats that belong to the same object. We take advantage of the zero-shot capabilities of open-vocabulary tracking models (Cheng et al., 2023; Wang et al., 2023; Ravi et al., 2024) to ensure that we obtain consistent predictions. Our compositional world model uses the DEVA open- vocabulary tracker to extract the objects and their segmentation masks 3%2 in the video X' of the scene. Since DEVA discovers objects given prompts, for all the objects and table for segmenting the table of our table-top robot. Having computed the segmentation masks 3 (yI,k: ..., kn, object & and for all frames,
use general and minimal prompts: ob ject
learn object-centric
Gaussian Splats as described in Sec. 3.1. Advances in the object-centric (Didolkar et al., 2024; Zadaianchuk et al., 2024; Wu al., 2024) and open-vocabulary (Ravi al., 2024; Bianchi et al., 2024) literature can be seamlessly integrated into our compositional manipulation world model.
3.3 MODELLING ROBOT-OBJECT DYNAMICS WITH OBJECT ASSETS
To be able to predict the consequences of the actions of the robot, we must model the dynamics D. We can directly use external physics engines (Todorov et al., 2012: Coumans & Bai, 2016) for the dy- namic operator D. For our model, we rely on PyBullet (Coumans Bai, 2016) for its simplicity.
While object-centric Gaussian Splat representations GR are good for photorealistic renderings of objects, they are not designed for physics simulations. Instead, physics engines -including PyBullet- typically operate object meshes. We thus need to convert G& into object mesh grids, M₁- = (Hk.., Pk.,,, uk, ). with mesh centers jkk,, orientations Pk.., and appearances Uk,,, to obtain our world object assets, o₁- = (Gx, My. w). With we explicitly denote the relevant physical parameters of the world, such as the masses of objects or the friction coefficient. In the robot simula- tor, these are in the URDF files (Villasevil et al., 2024). During learning the policies, the parameters w can be either given or inferred (Memmel et al., 2024). The physies engine acts as an operator on the positional and orientation parameters of the mesh grid of all objects in the world, that is
(M.41)
= (u., +4.C+AA1.... - p,,+1) (AAp.... p..
using standard Newtonian mechanics as the dynamics operator. Note that any object in the world, including the robot arm, can interact with any other object. This is why in Equation we drop the object subscript k, and why in the differential A we add all objects 1...K. K. Knowing the forces that the robot exerts on the environment at any moment, our compositional manipulation world model can easily imagine the future states of the world and the consequences of the robot's 's actions. Finally, we apply the change in mesh center Ap4; and mesh orientation Ap to all Gaussians g E G by updating their centers p' p + Aps and their orientations r' Apk
5

---

Published as a conference paper at ICLR 2025
Modelling the agent. In the above, we focused the objects or world. The robot is also an object that is selfbvi thus it can be just as easily represented as an object-centric mesh grid from the corresponding Gaussian Splat reconstruction. A difference of the robot arm compared to other objects is that it is articulated. However, the rohot joints are known in advance, and we can use their associated segmentation masks to link the different robot parts, ensuring accurate Gaussian representation for rendering purposes. Regarding the dynamics of the robot arm, our assumption is that the robot can model its own dynamics perfectly, similarly as humans have an accurate internal model ( (Wolpert et al., 1998). This is a reasonable assumption, since most manipulation robots have the corresponding URDF file, allowing for the instantiation of a simulator with the robot model.
World model verification from demonstrations. So far we assumed that the physical parameters fixed. With this assumption, we cannot guarantee that the manipulation world model will lead to correct predictions for arbitrary actions, as errors in the Gaussian Splat reconstructions and their mesh grids. misalignment in the physical parameters, could result in cumulative emrors when applying the dynamics. That said, our undertaking in this work is that by adopting explicitly grounded appearance and dynamics models for the world and its objects, useful model. Suppose that the model predictions for replaying the sequence of actions in the demonstration result in a final state similar the final state of the demonstration. In that case, we can still use the world model as an interactive and compositional representation of the demonstration itself. In the next section, we show that such an interactive representation allows us to augment the original demonstrations with imagined ones to train more robust imitation learning policy.
can still obtain
3.4 ENGINEERING
The proposed compositional manipulation world model is significant undertaking, requiring extensive engincering to ultimately be effective with real-world robots. However, one of the key advantages of this approach is that it enables real robot to learn novel policies using only a single example. We describe the most critical engineering components in detail, with additional information provided in App. J. While Gaussian Splatting is primarily designed for novel view synthesis, we observed unexpected behavior at low resolutions, leading to incorrect action pre- dictions by the robot. Additionally, depth information rendered was inaccurate, as the original Gaussian Splatting method does not ensure that Gaussians align with surfaces. To address this, implemented 2DGS Gaussian Splatting (Huang et al., 2024). For successful real-world deployment of DREMA, properly aligning the reconstruction and the simulator is crucial, which we achieved using a calibrated camera mounted on the robot (Allegro et al., 2024). Also, estimating physical parameters is a complex field in its own right (Mehra, 1974; Chebotar et al., 2019; Gao et al., 2022; Huang et al., 2023), and methods such as ASID (Memmel et al., 2024) or AdaptSim (Ren al., 2023) can be applied for this purpose. In this paper, for simplicity we pick constant physical parameters that seemed to work well in the validation sets.
4 EQUIVARIANT TRANSFORMATTIONS FOR IMITATION LEARNING
In computer vision, many data augmentations preserve the semantic information of an image, enabling robust recognition and self-supervised representation learning. However, augmenting demonstration data for imitation learning presents a greater challenge (Pitis et al., 2020; Urpi et al., 2024). When altering an observation sequence (e.g.. by changing an object's initial position), it is necessary to also modify the corresponding action sequence to maintain semantically consistent demonstration-oon that still solves the intended task. Therefore, generating valid augmentations in imitation learning requires equivariant trunsformations that simultaneously adjust both the high-dimensional observations and the associated actions.
In imitation learning, the agent learns from dataset 2 = (6I...,, consisting of M demonstra- tions, each paired with a corresponding task encoding T = {t,,..., J. When tasks are described by language-based goals, the task i'm is represented by a language embedding (e.g., pre-trained CLIP embeddings) corresponding to the phase that describes the task. Each demonstration Sm consists of a sequence of continuous actions A- (a1...., at), which are encoded end-effector poses and gripper states. These actions are paired with observations Q = (qI...., ..)... which in our case can be the object assets from the compositional manipulation world model.
6

---

Published as a conference paper at ICLR 2025
Original demonstration
Translation
Rotation
Object Rotation
Transformed
Transformed
Transformed
Environment:
Environment:
Environment:
Objects:
Objects:
Objects:
Original trajectory:
Transformed trajectory:
The effect of equivariant translation, equivariant rotation, and the object rotation transfor- mations. Top row: start of demonstration. Bottom row: target of demonstration.
Figure 3:
Equivariant transformations for imagining novel demonstrations. Finding equivariant transformations for high-dimensional observations Q and actions A is typically challenging, but DREMA simplifies this by using explicit representations for object assets O. our world model, both the robot's actions A (the 3D positions of the end-effector and gripper) and the 3D positions of the objects, including initial and target locations, are transformed equivariantly under any rigid transformation. As result, these transformations preserve task success in the imagined demon- strations. This allows us to generate novel "imaginations" by applying geometric transformations to the objects' positions and the robot's actions, without changing the robot's internal state. By executing these transformed actions A and rendering the corresponding observations Q in DREMA, we generate new demonstrations Znew not present in the original dataset Z. In the following, we describe several such transformations applicable to manipulation tasks (see 3).
Roto-translation of the environment and objects. Given the roto-translation R, and the (the beginning of the
point of application Px, we transform each object asset os time demonstration) and the Gaussians of the environment. To adapt the same demonstration to the transformed environment R2 applied to the end-effector translation et, and the end-effector rotation eT; of each action a; (et,, eT,, ga.). Consequently, the agent (fixed) is required to execute the transformed set of actions in the new environment.
Object rotation. This transformation involves rotating the objects around the final position of the action sequence A. The transformation R is applied to both the object poses and the action trajectory A, keeping the demonstration consistent with the original task. In this second transformation
fixed P. et, (the last position of the end-effector). We then rotate the objects assets Og at time t = 0 and the actions a; similarly to the previous approach. In this transformation do not change the environment since the trajectory brings the object similar position.
Verification of transformed demonstrations. Although these transformations can generate novel demonstrations, there is a risk of producing invalid outcomes due to factors like the robot's 's inability execute the entire sequence or inaceuracies in the world model dynamics. Therefore, it is crucial to verify that the generated demonstrations still solve the intended task. In a general setting, a success detector would be required ensure task completion. However, for goal-based tasks where cquivariant transformations are applied, it is sufficient to confirm that the final positions of the objects in the transformed trajectory remain close to their expected transformed goal positions i,; ||4 (4, - Pr.)) + Px.il |l < T where R2 represents the transformation applied, Px, is the center of the transformation, and = 0.015 meters is the empirical threshold for positional accuracy.
5 EXPERIMENTS
We evaluate the performance of DREMA in both simulated and real-world environments, comparing the effectiveness of imagination- generated data to imitation learning agents trained only on original demonstrations. For this evaluation, imitation learning. Our study addresses the following key questions: How does DREMA perform with only few demonstrations available, in single- and multi-task settings? Does manipulation
PerAct (Shridhar al., 2023), common baseline for

---

Published as a conference paper at ICLR 2025
Comparison of trained original, augmented, and DREMA with imagined and all demonstrations, reporting mean + std and max success over 5 runs on 50 test environments.
Table 1:
Single-task
close jur
pig
fift
pick cup =t std
shape
mean 1: std
max
mecn
2: std
max
mecn
2: std
max
macan
max
mean
1 sid
max
PerAcc (Original data))
38.4 + 0.80
40
0.0 + 0.00
0.0
22.8
1.6
26
13.2
2.04
16
6.4
1,50
8
Random patches (l.skii al., 2020)
3.49
48
2.0
1.79
20.4
2.65
3.25
Random table color (Chen et al., 2023)
45.0 d: 1.73
46
1.6 =: 1.50
23.6 +: 0.80
24
25.6
2.94
30
8.0 = 1.26
10
Distractors (Bharadhwaj et al., 2024)
36.1 0.80
38
0,4 0.80
22.8 1: 1.60
24
41.2
2.99
44
4.4 2.35
8
DRFM (imaginatios)
41.2
46
0.98
1.6
4.20
12
DREMA Original (All data)
$1.2 1.60
54
2 2.33
23.6 +: 1,50
26
34.4
3,88
40
11.2 :: 1.60
12
place wine
put cuphoard
slide black
stack blocks
avg Aingle-tasi
mean std
max
mean std max
mean +
max
mean
max
mean
max
PerAct (Original data)
10.0 2: 1.79
12
2.0 1: 0.00
48.4 1: 3.20
50
2.8 =: 0.98
4
16.0
17.6
Random patches (laskin al., 2020)
2,40
16
2.4 1,50
36.8 01.98
0.00
0.0
17.8
204
Random tablc color (Chen et al., 2023)
10.0 2.83
14
1.6 0.80
10.8 1: 4.32
46
7.2 : 2.04
10
18.2
20.7
Distractors (Bharadlwaj et al., 2024)
14.8 1: 4.49
22
3.2 L: 0.98
49.2 1: 0.98
30
8.4 4: 2.94
12
20.1
22.7
DREMA fimugination)
2.19
18
0.4 0.80
54.4 2.15
4,0
22.4
DREMA Original (All data)
26.8 1: 4.30
32
2.8 =: 0.98
62.0 1: 2.19
66
11.6 1.96
14
25.1
28.2
world model still provide an advantage when more demonstrations are available? 2. How do the different components of DREMA contribute overall performance? 3. How well does DREMA scale to real-world tasks, where both imitation learning and world model training are more complex?
5.1 ONE-SHOT POLICY LEARNING: SINGLE- AND MULTI-TASK
We test DREMA with a minimal number of demonstrations per task, i.e., one example per task variation up to five maximum examples (five examples for each task except for slide block and place wine that have and 3 variations). We compare PerAct trained with original demonstrations, PerAct trained with DREMA's 's imagined ones, and PerAct trained with imagined and original combined. All simulation experiments are conducted on 50 environment configurations, with the test repeated five times to account for variability introduced by the motion planner.
Data collection in simulation. We replicate the experimental setup of PerAct (Shridhar et al., 2023) using RLBench (James et al., 2020), which involves Franka Emika Panda robot with parallel jaw gripper, placed on a table with objects that vary in color and position. We focus on tasks excluding articulated objects, as our current object-centric priors do not handle articulations. The experiments cover 1) non non-prehensile tasks and 2) pick-and-place tasks (nine tasks in total). Following Ze et al. (2023), we collect a sequence, X, of 200 images by simulating a rotating camera around the scene. PerAct supports multi-task training, where a single policy adapts to multiple tasks based on prompts. We train both PerAct and DREMA on three tasks using the same demonstrations as in the single-task setup. Hyperparameters are consistent with the original paper (Shridhar et al., 2023), except for the training iterations and batch size (see App. L). To prevent overfitting, we validate all models performance on separate set for each 10k training iterations for multi-task and 5k for single-task and select the best-performing model.
One-shot single- and multi-task results (Table 1 & Table A.5). Since the limited number of demonstrations is insufficient learn a general policy, the PerAct model trained solely on original demonstrations performs the worst. In contrast, DREMA, exploiting around 800 demonstrations (see App. L), significantly outperforms PerAct with the original data in single-task settings (Table 1), achieving an average accuracy improvement of +9.1%. This result highlights that additional imagination-bass
demonstrations help create more robust policies, even when the original dataset is small. When DREMA is compared to invariant augmentations -namely, random patches, random table color, and random distractors- it achieves the best performance, with a 5.0% average improvement over the best invariant augmentation (see App. A). Additionally, DREMA improves performance in all tasks compared to PerAct, differently from invariant augmentations. In multi-task settings (Table A.5), DREMA is tested close jar, slide block, and sort shape simultancously. It achieves an average accuracy of 35.5%, representing a +13 +13.1% 1% improvement over PerAct, demonstrating that even this setting can significantly benefit from imagination-base
8

---

Published as a conference paper at ICLR 2025
The mean + std and maximum success rate of DREMA trained single-task demonstra- tions from different types of transformations over 5 test runs on 50 examples.
Table 2:
close jar
sort shape std
slide block
mean
max
mean
max
memm + std
max
Replay Object Rotation
10.0 1: 3.10
16
1.2 1 0.98
2
26.0 1 0.00
26
25.2 1 0.98
10,4
1. 2.65
14
42.0 0.00
Roto-translation
41.2 + 2.40
10.0
+ 2.53
14
50.4
52
DREMA Original (All data)
51.2 + 1.60
54
11.2
1.60
12
6.19
66
demonstrations. DREMA trained on only imagination demonstrations and evaluated the original environment still performs significantly better than PerAct (+3.1 on single task +9.1 9..1ol on multi-task).
5.2 ANALYSIS AND ABLATIONS
Varying dataset size (Figure & Figure C.6). We analyze the effectiveness of DREMA's imag- ined demonstrations based on the number of original demonstrations provided. Specifically, vary the dataset size (D| E (4, 10, 203 for the slide block and and [D] E {5, 10, 20} sort shape tasks (see App. for training details). Figure 4 shows the average success rate of DREMA and PerAct as a function of the original dataset size. Increasing the training examples to 20 improves performance both tasks. PerAct trained on 20 examples matches DREMA trained with imagined data from just 5 demonstrations (81 and 61 imagined demos for sort shape and slide block, respectively). However, when using the same examples to generate imagined data (437 and 391 imagined demonstrations for sort shape and slide block, respectively), the model improves by 6.4% and 10.4%. This shows that generating novel configurations and actions remains beneficial, even with more demonstrations. Figure C.6 in the appendix shows how augmented data cover wider portion of the robot's 's workspace, enabling the agent to handle configurations beyond those the original demonstrations.
Ablating transformations (Table 2). We compare the equivariant transforma- tions and include baseline called Replay, which replays the original demonstrations.
Place Shape Task
Slide Block Task
17.5
70
15.0
The Object Rotation and Roto-translation augmentations generate entirely imagined data. Both Object Rotation and Roto- translation increase the coverage of the state space and improve the generaliza- tion of the learned policy, with Roto- translation having the greatest impact. Combining these transformations results in even more general policies, yielding 10.4% improvement in mean accuracy. This highlights the importance of learning good world model and using explore and apply valid data augmentations.
Accurayy
12.5
Accuracy
50
10.0
40
7.5
20
2.5
0.0
10
10
20
Demonstrations
Demonstrations
(original only)
DreMa (imagination only)
Figure 4: Imagined demonstrations keep improving imitation learning even with increasing number of original data.
5.3 EXPERIMENTS WITH A REAL ROBOT
Real robot setup. We verify the effectiveness of our approach in real-world environments. Our setup includes Franka Emika Panda on table, with a calibrated Kinect Azure mounted on end-effector and a Kinect V2 positioned front of the table, replicating the setup used in PerAct. We collect five tasks (number of demonstrations in Table 4). We train two models per task -one using only the original demonstrations and another combining original and imagined demonstrations- and do model selection in PerAct (see App. L for details). The models tested in 20 different environment configurations: 10 with object positions similar to the training set and 10 with positions outside the distribution (OOD), assess generalization.
Verification of the learned world model (Table 3). To assess the accuracy of DREMA in modeling demonstrations, compare the final object positions at the end of the demonstration

---

Published as a conference paper at ICLR 2025
Table 3: Localization
Table 4: In- and out-of-distribution evaluation with real robots.
rors.
pick block
pick shupe
push
place obj
ctase cumpples
Aserge
exumpies
caunples
3 Eeaampla
4 exsampicr
Error (m)
In dist.
00D
11 disr.
00D
In diver
00D
In distr.
000
dist. 0OD
ang.
0OD
pick block
0,010
PerAet
55
50
31.7
25.0
pick shape
0.050
DREMA (AII)
90
90
38
62.9
$8.1
push
0.049
Average
Figure 5: Original (top) and imagined demonstration (bottom) after a 90 rotation transformation.
with those predicted by the world model. Results show that the world model aligns well with the collected data for simpler demonstrations. For more complex examples, errors could increase. This discrepancy could be addressed using real demonstrations to adapt the world model's 's parameters.
Real robot results (Table 4 & Table B.6). The imagined demonstrations double the model's task performance (62.9% for DREMA 31.7% for PerAct on average). Moreover, PerAct, trained only on original demonstrations, struggled when the object distribution differed from the training data. In contrast, DREMA performs robustly in- and out-of-diitribution environment configurations. The pick shape task had the lowest accuracy, likely due to the small size of the shapes and the model's difficulty in accurately recognizing them. Finally, in Table B.6, also confirm that DREMA with only imagined trajectories can outperform original PerAct training (+25% on pick block task) when more examples are available, showing the usefulness of the imagined trajectories on their own.
CONCLUSIONS
We introduce novel framework for constructing manipulation world models. By integrating real-time photorealism with Gaussian Splatting and physics simulators, DREMA enables robots to accurately predict the consequences their actions and learn from minimal data. Our evaluations show significant improvements in policy learning, demonstrating the model's robustness and efficiency in real-world applications. We believe DREMA can influence other research areas, such as augmenting task variations with VLMss enabling policy evaluation through imagination.
Limitations. Although DREMA demonstrates several strengths, it has limitations. First, requires full observability of the environment to model physics and appearances. Second, it currently handles only simple objects, but recent advances in automatic URDF estimation could address this issue (see App. G). Third, construct it with open-loop approach without feedbacks; incorporating them into the world model construction could improve both prediction accuracy and policy learning. Finally, scaling the world model large and diverse environments may increase simulation complexity. Overall, we believe that future research can address these issues, particularly the second and the third. Another interesting future direction is study the effect of DREMA with more recent models (such as RVT2 (Goyal et al., 2024) or Diffusion Policy (Chi et al., 2023)) or approaches that directly exploit RGB images show the benefit of Gaussian Splatting.
10

---

Published as a conference paper at ICLR 2025
Reproducibility statement. We are committed to ensuring the reproducibility of our research findings. In our paper and in the appendix, we have provided comprehensive details about our methodologies, experimental setups, and hyperparameter configurations. In addition, details about inference and training runtime of both DREMA and corresponding agent trained with data generated by DREMA are presented in Appendix D. Furthermore, will share the code and the synthetic datasets used for our simulation experiments. This will enable researchers to replicate our experi- ments and verify our results.
Acknowledgments. This work is supported by ERC Starting Grant EVA 950086. Andrii Zada- ianchuk is funded by the European Union (ERC, EVA, 950086). Leonardo Barcellona was supported by scholarship from Fondazione Ing. Aldo Gini (University of Padova) during part of the project.
REFERENCES
Jad Abou-Chakra Krishan Rana, Feras Dayoub, and Niko Sünderhauf. Physically embodied gaus- sian splatting: A realtime correctable world model for robotics. In 8th Annual Conference Robot Learming, 2024.
Davide Allegro, Matteo Terreran, and Stefano Ghidoni. Multi-camera hand-eye calibration for human-robot collaboration industrial robotic workcells. IEEE Roboties and Automation Let- ters, pp. 1-8, 2024. doi: 10.1
Homanga Bharadhwaj, Jay Vakil, Mohit Sharma, Abhinav Gupta, Shubham Tulsiani, and Vikash Kumar. Roboagent: Generalization and efficiency in robot manipulation via semantic augmenta- tions and action chunking. 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 4798. IEEE, 2024.
Lorenzo Bianchi, Fabio Carrara, Nicola Messina, Claudio Gennaro, and Fabrizio Falchi. The devil is in the fine-grained details: Evaluating open-vocabulary object detectors for fine-grained under- standing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recogni- tion, pp. 22520-22529, 2024.
Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, et al. Genie: Generative inter- active environments. In Forty-first International Conference on Machine Learning, 2024.
Berk Calli, Aaron Walsman, Arjum Singh, Siddhartha Srinivasa, Pieter Abbeel, and Aaron M Dol- lar. Benchmarking in manipulation research: The ycb object and model and benchmarking protocols. arXiv preprint /5ti:::502.3///// 2015.
Yevgen Chebotar, Ankur Handa, Viktor Makoviychuk, Miles Macklin, Jan Issac, Nathan Ratliff, and Dieter Fox. Closing the sim-to-real loop: Adapting simulation randomization with real world experience. In 2019 International Conference on Robotics and Automation (ICRA), pp. 8973- 8979. IEEF, 2019.
Siwei Chen, Yiqing Xu, Cunjun Yu, Linfeng Li, and David Hsu. Differentiable particles for general- purpose deformable object manipulation. arXiv preprint arXiv:2405.01044 2024a.
Zoey Chen, Sho Kiami, Abhishek Gupta, and Vikash Kumar. Genaug: Retargeting behaviors to unseen situations via generative augmentation. arXiv preprint urXiv:2302.0660 066677 2023.
Zoey Chen, Aaron Walsman, Marius Memmel, Kaichun Mo. Alex Fang, Karthikeya Vemuri, Alan Wu, Dieter Fox. and Abhishek Gupta. Urdformer: A pipeline for constructing articulated simula- tion environments from real-world images. arXiv preprint arXiv:2405.11656 2024b.
Ho Kei Cheng, Seoung Wug Oh, Brian Price, Alexander Schwing, and Joon- Joon-Young Lee. Tracking anything with decoupled video segmentation. In ICCV, 2023.
Cheng Chi, Zhenjia Xu, Siyuan Feng. Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The Inter- national Journal of Robotics Research, pp. 0278364924127366.. 2023.

---

Published as a conference paper at ICLR 2025
Nicholas Corrado and Josiah P. Hanna. Understanding when dynamics-invariant data augmenta- tions benefit model-free reinforcement learning updates. In The Twelfth International Confer- Learning Representations, 2024. URL sVEu225577.
net/ forummidd- wtwp:r/wwwwweeeeoooutmmaaaaaeeee
Erwin Coumans and Yunfei Bai. Pybullet, a python module for physics simulation for games, robotics and machine learning. http::/pybbullet.orgg 2016.
Matt Deitke, Dustin Schwenk, Jordi Salvador, Luca Weihs, Oscar Michel, Eli VanderBilt, Ludwig Schmidt, Kiana Ehsani, Aniruddha Kembhavi, and Ali Farhadi. Objaverse: A universe of anno- tated 3d objects. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 13153, 2023.
Aniket Didolkar, Andrii Zadaianchuk, Anirudh Goyal, Mike Mozer, Yoshua Bengio, Georg Mar- tius, and Maximilian Seitzer. Zero-shot object-centric representation learning. arXiv preprint arXiv:2408.09100 2024.
Martin A Fischler and Robert C Bolles. Random sample consensus: paradigm for model fitting with applications to image analysis and automated cartography. Communications of the ACM, 24 (6):381-395, 1981.
Ziyan Gao, Armagan Elibol, and Nak Young Chong. Estimating the center of mass of an unknown object for nonprehensile manipulation. In 2022 IEEE International Conference on Mechatronics and Automation (ICMA), pp. 1755-1760. IEEE, 2022.
Ankit Goyal, Jie Xu, Yijie Guo, Valts Blukis, Yu-Wei Chao, and Dieter Fox. Rvt: Robotic view transformer for 3d object manipulation. In Conference on Robot Learning, pp. 694-710. PMLR, 2023.
Ankit Goyal, Valts Blukis, Jie Xu, Yijie Guo, Yu-Wei Chao, and Dieter Fox. Rvt-2: Learning precise manipulation from few demonstrations, arXiv preprint 2en::2406.05555 2024.
Antoine Guédon and Vincent Lepetit. Sugar: Surface-aligned gaussian splatting for efficient 3d mesh reconstruction and high-quality mesh rendering. In Proceedings of the IÉEE/CVF Confer- ence on Computer Vision and Pattern Recognition, pp. 5354-5363, 2024.
Christian Gumbsch, Noor Sajid, Georg Martius, and Martin V. Butz. Learning hicrarchical world models with adaptive temporal abstractions from discrete latent dynamics. The Twelfth In- ternational Conference Learning Representations, 2024. URL https: https::/ /pwprevieeeeeeeeeee i dd=
David Ha and Jürgen Schmidhuber. World models. arXiv preprint
2018. 201..
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:191.001603 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.02/9. 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 041004, 2023.
Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for contin- uous control. In The Twelfth International Conference on Learning Representations, 2024. URI. net/ forum @==0xhhcstD.....
https: /ttp:////peeeeeeeeeeeeee
Jeff Hawkins, Subutai Ahmad, and Yuwei Cui. theory how columns in the neocortex enable learning the structure of the world. Frontiers in Neural Circuits, 11, 2017. ISSN 1662-51 10. doi: 10.3389/fincir20117.0000 00081. URL https: /tt://wwww /twwww frontiersin org/ jes/joounaa///// neural-circuits/are
300.3389//55 fncir.ici.....i. .010222171 00081.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
12

---

Published as a conference paper at ICLR 2025
Anthony Hu, Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie Shot- ton, and Gianluca Corrado. Gaia-l: A generative world model for autonomous driving. arXiv preprint ar arXiv:2309.11000 2023.
Binbin Huang, Zehao Yu, Anpei Chen, Andreas Geiger. and Shenghua Gao. 2d gaussian splatting for geometrically accurate radiance fields. In ACM SIGGRAPH 2024 Conference Papers, pp- 2024.
Peide Huang, Xilun Zhang, Ziang Cao, Shiqi Liu, Mengdi Xu, Wenhao Ding, Jonathan Francis, Bingqing Chen, and Ding Zhao. What went wrong? closing the sim-to-real gap via differentiable causal discovery. In Conference Robot Learning, pp. 734-760. PMLR, 2023.
Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew Davison. Rlbench: The robot learning benchmark & learning environment. IEEE Robotics and Automation Letters, Letters, 5(2):3019- 3026, 2020.
Stephen James, Kentaro Wada, Tristan Laidlow, and Andrew J Davison. Coarse-to-fine q-attention: Efficient learning for visual robotic manipulation via discretisation. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 19739-13743 2022.
Fan Jia, Weixin Mao, Yingfei Liu, Yucheng Zhao, Yuqing Wen, Chi Zhang, Xiangyu Zhang, and Tiancai Wang. Adriver-i: A general world model for autonomous driving. arXiv preprint arXiv:2311. 13549, 2023.
Ying Jiang. Chang Yu, Tianyi Xie, Xuan Li, Yutao Feng, Huamin Wang, Minchen Li, Henry Lau, Feng Gao, Yin Yang, et al. Vr-gs: a physical dynamics-aware interactive gaussian splatting system in virtual reality. In ACM SIGGRAPH 2024 Conference Papers, pp. 1-1, 2024.
Bernhard Kerbl, Georgios Kopanas, Thomas Leimkühler, and George Drettakis. 3d gaussian splat- ting for real-time radiance field rendering. ACM Trans. Graph., 42(4):139-1, 2023.
Justin Kerr, Chung Min Kim, Mingxuan Wu, Brent Yi, Qianqian Wang, Ken Goldberg, and Angjoo Kanazawa. Robot see robot do: Imitating articulated object manipulation with monocular 4d reconstruction. arXiv preprint arXiv:2409. 18121, 2024.
Thomas Kipt, Elise Van der Pol, and Max Welling. Contrastive learning of structured world models. 2019.
arXiv preprint
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloc Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander Berg, Wan-Yen Lo, et al. Segment anything. In Proceed- ings of the IEEE/CVF International Conference Computer Vision, pp. 4015-4020 2023.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Picter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. Advances in neural information processing systems, 33: 19895, 2020.
Xuan Li, Yi-Ling Qiao, Peter Yichen Chen, Krishna Murthy Jatavallabhula, Ming Lin, Chenfanfu Jiang, and Chuang Gan. PAC-neRF: Physics augmented continuum neural radiance fields for aoometry-agnostie system identification. In The Eleventh International Conference on Learning Representations, 2023. URL
Xuanlin Li, Kyle Hsu, Jiayuan Gu, Karl Pertsch, Oier Mees, Homer Rich Walke, Chuyuan Fu, Ishikaa Lunawat, Isabel Sieh, Sean Kirmani, et al. Evaluating real-world robot manipulation policies in simulation. arXiv preprint 2xXiv:2405.05941, 2024.
Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B. Tenenbaum, and Antonio Torralba. Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids. In International Confer- ence on Learning Representations, 2019. URL https: //tp///eeee/eee/e////////// net forum?id rJgbSn09Ym.
13

---

Published as a conference paper at ICLR 2025
Yen-Chen Lin, Pete Florence, Andy Zeng, Jonathan T Barron, Yilun Du, Wei-Chiu Ma, Anthony Simconov, Alberto Rodriguez Garcia, and Phillip Isola. Mira: Mental imagery for robotic affor- dances. In Conference on Robot Learning, pp. 1916-1927. PMLR, 2023.
Jiayi Liu, Ali Mahdavi-Amiri, and Manolis Savva. Paris: Part-level reconstruction and motion analysis for articulated objects. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 352-363, 2023.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-Centric Learning with Slot At- tention. In NeurlPS, 2020. URL https::tp/... http::::///opeedddee 0t11://81.adfff22222aeeaeee22222
neurips. wwtrriis.cccaaapeee wwtppppr//02....../ pdf.
Haozhe Lou, Yurong Liu, Yike Pan, Yiran Geng, Jianteng Chen, Wenlong Ma, Chenglong Li, Lin Wang. Hengzhen Feng, Lu Shi, al. Robo-gs: physics consistent spatial-temporal model for robotic arm with hybrid representation. arXiv preprint 1xxxi::240..188777 2024.
Guanxing Lu, Shiyi Zhang, Ziwei Wang, Changliu Liu, Jiwen Lu, and Yansong Tang. Manigaussian: Dynamic gaussian splatting for multi-task robotic manipulation. In European Conference Computer Vision, pp. 349-366. Springer, 2025.
Jonathon Luiten, Georgios Kopanas, Bastian Leibe, and Deva Ramanan. Dynamic 3d gaussians: Tracking by persistent dynamic view synthesis. arXiv preprint arXiv:2308.09713. 2023.
Ajay Mandlekar, Soroush Nasiriany, Bowen Wen, Iretiayo Akinola, Yashraj Narang, Linxi Fan, Yuke Zhu, and Dieter Fox. Mimicgen: A data generation system for scalable robot learning using human demonstrations. arXiv preprint arXiv:2310.17596 2023.
Raman Mchra. Optimal input signals for parameter estimation in dynamic systems-survey and new results. IEEE transactions on automatic control, 19(6):753-768, 1974.
Marius Memmel, Andrew Wagenmaker, Chuning Zhu, Dieter Fox, and Abhishek Gupta. Asid: Active exploration for system identification in robotic manipulation. In The Twelfth International Conference on Learning Representations, 2024.
Lukas Meyer, Floris Erich, Yusuke Yoshiyasu, Marc Stamminger, Noriaki Ando, and Yukiyasu Domae. Pegasus: Physically enhanced gaussian splatting simulation system for 6dof object pose dataset generation. arXiv preprini 21X0v:2401.02281, 2024.
Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances Neural Information Processing Systems, 33:3976-3990, 2020.
Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-basee counterfac- tual data augmentation. Advances in Neural Information Processing Systems, 35:18145 11118844-18555 2022.
Morgan Quigley, Ken Conley, Brian Gerkey, Josh Faust, Tully Foote, Jeremy Leibs, Rob Wheeler, Andrew Y Ng. et al. Ros: an open-source robot operating system. In ICRA workshop on open source software. Kobe, Japan, 2009.
Nikhila Ravi, Valentin Gabeur, Yuan-Ting Hu, Ronghang Hu, Chaitanya Ryali, Tengyu Ma, Haitham Khedr, Roman Rädle, Chloe Rolland, Laura Gustafson, al. Sam 2: Segment anything in images and videos. arXiv preprint AXXi:2408..07111 2024.
Allen Ren, Hongkai Dai, Benjamin Burchfiel, and Anirudha Majumdar. Adaptsim: Task-driven simulation adaptation for sim-to-real transfer. In Conference on Robot Learning, pp. 3434-3452. PMLR, 2023.
Tianhe Ren, Shilong Liu, Ailing Zeng, Jing Lin, Kunchang Li, He Cao, Jiayu Chen, Xinyu Huang, open-world
Yukang Chen, Feng Yan, et al. Grounded sam: Assembling open- world models for diverse visual tasks. arXiv preprint arXiv:2401. 2ete:240..4444 14159, 2024.
14

---

Published as a conference paper at ICLR 2025
Sipu Ruan, Weixiao Liu, Xiaoli Wang. Xin Meng, and Gregory Chirikjian. Primp: Probabilistically-nnooe -infomeedd motion primitives for efficient affordance learning from demonstra- tion. IEEE Transactions Robotics, 2024.
Maximilian Seitzer, Max Horn, Andrii Zadaianchuk, Dominik Zietlow, Tianjun Xiao, Carl-Johann Simon-Gabriel, Tong He, Zheng Zhang, Bernhard Schölkopf, Thomas Brox, and Francesco Locatello. Bridging the gap to real-world object-centric learning. In ICLR, 2023. URL wtt.://efoommtmoottttttto di-tttdd--bbbteeeeeeeeee
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Cliport: What and where pathways for robotic manipulation. In Conference on robot learning, pp. 894-906. PMLR, 2022.
Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Conference on Robot Learning, pp. 785-799. PMLR, 2023.
Elizabeth S. Spelke. Principles of object perception. Cognitive Science, 1990. URL https: mfonninelibraat wiley www. doo/ aolo//abs/000 120.1 htt.:/1/01/111000o...oo 140__34
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 JEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012. doi: 10110000..20012.36888
Marcel Torne Anthony Simeonov, Zechu Li, April Chan, Tao Chen, Abhishek Gupta, and Pulkit Agrawal. Reconciling reality through simulation: A real-to-simn-o-real approach for robust ma- nipulation. arXiv preprint arXiv:2403.03949 2024.
Núria Armengol Urpi, Marco Bagatella, Marin Vlastelica, and Georg Martius. Causal action influ- ence aware counterfactual data augmentation. arXiv preprint arXiv:2405. 100005..///11 /89 2024.
Marcel Tome Villasevil, Arhan Jain, Vidyaaranya Macha, Jiayi Yuan, Lars Lien Ankile, Anthony Simeonov, Pulkit Agrawal, and Abhishek Gupta. Scaling robot-learning by crowdsourcing simu- lation environments. In RSS 2024 Workshop: Data Generation for Robotics, 2024.
Haochen Wang, Cilin Yan, Shuai Wang. Xiaolong Jiang, Xu Tang, Yao Hu, Weidi Xie, and Efstratios Gavves. Towards open- vecaalaccyyl video instance segmentation. In proceedings of the IEEF/CVF international conference on computer vision, pp- 4057-4066, 2023.
Yijia Weng, Bowen Wen, Jonathan Tremblay, Valts Blukis, Dieter Fox, Leonidas Guibas, and Stan Birchfield. Neural implicit representation for building digital twins of unknown articulated ob- jects. In Proceedings of the IEFF//VV Conference Computer Vision and Pattern Recognition, pp. 3 3141-3150, 2024.
William F Whitney, Tatiana Lopez-Guevara, Tobias Pfaff, Yulia Rubanova, Thomas Kipf, Kim Stachenfeld, and Kelsey R Allen. Learning 3d particle-based simulators from RGB-d videos. In The Twelfth International Conference Learning Representations, 2024. URL https: / oppereeveeww net feret/formmiiiie tow.:ommdmd--
Daniel M Wolpert, R.Chris Miall, and Mitsuo Kawato. Internal models in the cerebellum. Trends in Cognitive Sciences, 2(9):338-347, 1998. ISSN 1364-6613. doi: https://doi.or/// 00166113364-661119999222332
URL https:/ //ww://wwwwwww// www.ssssiieeeeceeeeeeeeeeeeeeeee com/ science/e
art.icleep11/9118466118888
Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Daydreamer: World models for physical robot learning. In Conference robot learning, pp. 2240..240.
PMLR, 2023.
Ziyi Wu, Yulia Rubanova, Rishabh Kabra, Drew A Hudson, Igor Gilitschenski, Yusuf Aytar, Sjoerd van Steenkiste, Kelsey R Allen, and Thomas Kipf. Neural assets: 3d-aware multi-object scene synthesis with image diffusion models. arXiv preprint arXiv:2406.0922 2024.
Tianyi Xie, Zeshum Zong. Yuxing Qiu, Xuan Li. Yutao Feng, Yin Yang, and Chenfanfu Jiang. Physgaussian: doysics-integrated 3d gaussians for generative dynamics. Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 3389-4395. 2024.
15

---

Published as a conference paper at ICLR 2025
Mengjiao Yang. Yilun Du, Kamyar Ghasemipour, Jonathan Tompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. arXiv preprint arXiv:2310.06114, 2023.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius, Self-supervised visual reinforcement learning with object-centric representations. In International Conference on Learning Represen- tations, 2021. URL https: www./pvieeeeee...eooe. net/ forum? (immmmiddxxppe -itpummmdddxxpppxccbboo
Andrii Zadaianchuk, Georg Martius, and Fanny Yang. Self-supervised reinforcement learning with independently controllable subgoals. In Conference on Robot Learning, pp. 384-33433 394. PMLR, 2022.
Andrii Zadaianchuk, Maximilian Seitzer, and Georg Martius. Object-centric learning for real-world videos by predicting temporal feature similarities. Advances in Neural Information Processing Systems, 36, 2024.
Yanjie Ze, Ge Yan, Yuch-Hua Wu, Annabella Macaluso, Yuying Ge. Jianglong Ye, Nicklas Hansen, Li Erran Li, and Xiaolong Wang. Gnfactor: Multi-task real robot learning with generalizable neural feature fields. Conference on Robot Learning, pp. 284-301. PMLR, 2023.
Andy Zeng, Pete Florence, Jonathan Tompson, Stefan Welker, Jonathan Chien, Maria Attarian, Travis Armstrong, Ivan Krasin, Dan Duong, Vikas Sindhwani. et al. Transporter networks: Re- arranging the visual world for robotic manipulation. In Conference on Robot Learning, pp. 726- 747. PMLR, 2021.
Qian Yi Zhou, Jaesik Park, and Vladlen Koltun. Open3d: A modern library for 3d data processing. arXiv preprint
2018.
Siyuan Zhou, Yilun Du, Jiaben Chen, Yandong Li, Dit- Yan Yeung, and Chuang Gan. Robodreamer: Learning compositional world models for robot imagination. In arXiv preprint arXiv:2404. 12377, 2024.
16

---

Published as a conference paper at ICLR 2025
APPENDIX
A EXTENDED SINGLE TASK AND MULTI-TASK RESULTS
Comparing with other augmentations. In the single-task setting, compare DREMA to Per- Act trained with three invariant augmentation technique:: generating random patches on RGB-D images (Laskin al., 2020), randomly altering the table color (Chen et al., 2023), and introducing random objects into the scene (Chen et al., 2023; Bharadhwaj et al., 2024). For each approach, we generate 80 augmented examples and combine them with the original data to train the policy.
To mask the table, we use ground-truth images provided by RLBench. Random distractors are introduced by embedding the point clouds of objects into the original re-projected point cloud and projecting them back onto the original views. The objects are randomly selected from 18 scans of the YCB dataset (Calli al.. 2015). We believe that combining invariant augmentations with equivariant augmentations has the potential to yield more robust and generalizable policies.
Multi-fask results. In multi-task settings, DREMA tested on close jar, slide block, and sort shape simultaneously. The agent trained with both imagined and real data achieves the highest aceuracy, except in the sort shape task, likely because shape matching is more refined and requires additional original training data.
Consistent with findings by Shridhar et (2023), we observe that while multi-task learning enables the agent to handle diverse tasks, it comes at the cost of lower accuracy for each individual task. The same three tasks achieved an average accuracy 41.5% in the single-task setting.
Table A.5: Comparison of PerAct (Shridhar et al., 2023) trained original demonstrations to DREMA trained on only imagination demonstrations and the combination of both in the multi-task setting. The table reports the mean + std and maximum success rate over test runs.
Multi-task
close jar
sort shape
slide block
avg mudti-task mean
mean
std
max
meun +: std
max
mean + std
max
max
PerAct (Original data)
26.0 3.10
28
7.2 1 1.60
10
34.0 1 5.06
38
22.4
25.3
DREMA (PerAct with imagined data)
28.0 3.35
18.0 + 2.83
48.0 + 1.79
31.3
34.7
DREMA (PerAct with all data)
46.0 3.58
6.4 3.20
54.0
2.19
58
35.5
40.7
B IMAGINATION RESULTS IN THE REAL EXPERIMENT
In Table B.6, we present the complete results for the pick block task in the real-world experiment. This experiment aimed to verify whether the robot can learn purely through imagination, demon- strated by the simulation results. We observe that with only four original examples, the learned policy performed worse than the initial policy. However, its performance improved significantly as more original examples were added.
Table B.6: Results of the pick block task with original and imagination demonstrations.
examples
8 examples
Real Imagination
55
55
50
80
Imagination + Real
C DATA GENERATION ANALYSIS
As mentioned in the work, a common problem of imitation learning, and particularly behavior cloning, is generalizing outside the given configuration of the environment. The ability of DREMA
17

---

Published as a conference paper at ICLR 2025
to generate high-quality data by imagining configurations of the environment allows us in- crease the training distribution. Figure C.6 shows this effect by plotting the location of the keypoints (robot actions) in the original demonstration and in the generated imagination. This results in bet- ter and more robust policies. Gaussian Splatring guarantees that the domain gap small while the simulator allows us recreate reasonable interaction between the robot and the objects.
-1.2
-02
-52
-0.4
-04
-24
Figure C.6: Keypoints distribution of original (blue dots) and generated data (red cross) for the sort shape task. Data generated from one examples (left), five examples (middle) and twenty examples (right).
D RUNTIME ANALYSIS
The data generation approach is performed offline, as its runtime scales linearly with the number of objects. Training PerAct in single-task configurations takes approximately two days on an Nvidia A40 for 100k iterations with batch size 4, so adopting an online approach would have minimal im- pact. Using an Nvidia RTX 4090, DEVA-tracking segments images per second, while generating the Gaussian model and mesh for single object takes about 4 minutes. Replaying trajectories is significantly faster, as only involves applying roto-translations to the Gaussians. average time to reach a waypoint in tasks like close jar, slide block, and sort shape was 1.715 sec- onds without rendering and 1.883 seconds with rendering. Rendering required 0.168 seconds on average update Gaussian positions, render five 128 x 128 RGB-D images, and filter them. The PerAct inference time is comparable to the one reported in the original paper (around 2 frames per second), DREMA does not modify the model architecture but only generates training data.
our tests, the
E EXPLICIT WORLD MODELS AND LEARNABLE DIGITAL TWINS
Typically, an implicit world model (Ha & Schmidhuber, 2018; Hafner al., 2019) would consist of the encoder that maps observations into a latent state, latent dynamics prediction (including reward prediction) and a decoder to reconstruct original observations for latent state. Using such model after training, an agent can imagine the possible future trajectory and use it to learn policy.
In our work, the encoder is neural network masking the RGB inputs, mapping onto a Gaussian Splat Object Asset representation. Gaussian Splats can be seen as shallow (1-layer) Neural Field networks. The decoder is the Gaussian Splatting reconstruction, per object.
The representation function of the world relies on (i) the Gaussian Splat Object Assets, and (ii) the physics simulator to generate, aka "imagine", new RGB-D representations novel future world states never seen before. We can use the imagined states to train our policies. This aligns with Ha and Schmidhuber's idea of an agent being trained entirely within "its own hallucinated dream" (Ha & Schmidhuber, 2018).
Similarly, in our experiments, the PerAct agent was trained using imagined sensory data (RGB and depth images rendered via Gaussian Splatting). The results in Table 1 demonstrate the agent's ability to use this generated data, while Table highlights its effectiveness in a real-world setting. By integrating this generated data with the original data, demonstrate improved performance, further validating the predictive power of the constructed world model.
In summary, approach satisfies key components of explicit world model:
18

---

Published as a conference paper at ICLR 2025
Represented implicit latent vectors in traditional models, while in case, it corresponds to explicit compositional representations that consist of the position of Gaussian splats, meshes, and other environment parameters.
Latent State:
Observations Reconstructions: Traditionally inferred by a neural network, here derived from Gaussian Splatting renderings.
Dynamics Prediction: Traditionally learned by reconstruction of the observation, while in our work, constructed using explicit representation and physics engine.
World models and learnable digital twins. We further highlight the connection between world models and learnable digital twins, as in Abou-Chakra et al. (2024). When an agent learns a digital twin from observations and exploits it for decision decision-making or policy learning, the boundary between the two becomes blurred. Learnable digital twins can be seen as explicit models of the environment. While standard world models typically implicit, both they and the more explieit digital twins serve the shared purpose of enabling agents to simulate and reason about their surroundings. This connection is further supported by Yang et al. (2023), which uses implicit and interactive world models as simulators.
In conclusion, while our work incorporates elements of real2sim and data augmentation, we be- lieve it fundamentally adheres to and extends the principles of world models. The growing overlap between these areas underscores convergence of these methodologies, and we believe our con- tribution aligns with this evolution.
F TASKS DESCRIPTIONS
We give more details about the imitation learning tasks described in the experiments. The simula- tion tasks are chosen from those defined PerAct by Shridhar al. (2023) without modifications (see Figure F.7 and Figure F.8). The real tasks have been designed to be similar to those used in the simulation (see Figure F.9).
F.1 SIMULATION TASKS
Close jar. Place the lid on the jar with the specified color, sampled from set of 20 color instances. In each scene, there a target and a distractor lid. The task is considered correct when the robot releases the lid on the jar.
Insert peg. Insert small square inside colored peg. The size of the square has been slightly increased to improve reconstruction. The task is successful when the square is inside the correct peg.
Lift. Pick a colored cube identified by the prompt and move it to the red spot. Since the colored spot is not collidable in RLBench, it was manually set non-collidable. The task is successful when the correct cube reaches the target.
Pick cup. Pick a target cup from two color variations and lift it. The task is successful when the correct cup is lifted.
Sort shape. Pick the prompted shape and place inside the correct hole in the sorter box. In each scene, there four distractor shapes and one correct shape place. The task is considered correct when the correct shape is inside the box.
Place wine. Pick a bottle and place it on the rack. The task is successful when the bottle is released the correct location on the rack.
Put cupboard. Pick an object identified by the prompt and place it inside a cupboard. Since the cupboard is floating in RLBench, it was manually set non-collidable.
19

---

Published as a conference paper at ICLR 2025
Close jar
Insert peg
Life
Pick cup
Sort shape
Figure F.7: Visualization of the simulation tasks. From top to bottom: close jar, insert peg, lift, pick cup, sort shape
Slide block. The task involves sliding red block onto one colored spot. The target colors are limited red, blue, pink, and yellow (four variations). The task is considered successful when part of the block is inside the specified colored area.
Stack blocks. Stack given number of blocks with colors selected by the user. The task is suc- cessful when the desired number of blocks are correctly identified and stacked.
F.2 REAL WORLD TASKS
Pick Block. The task consists of picking a green block and putting it on top of one bigger blue or block. each scene, both the possible targets are on the table. The task is considered correct the robot releases the green cube on the right block and touches the correct place. We do not require the green block to stay on the block since we do not wish to show high accuracy position, robot learned to distinguish the two correct positions. Training demonstrations: the cube is never moved from the position while the two target blocks are equally distant from it and symmetric to the axis of the robot base. We swapped the two blocks in each recorded demonstration. Test of domain: the blocks are placed asymmetrically with respect to the r axis, and the cube is placed in random position.
the
20

---

Published as a conference paper at ICLR 2025
Place wine
Put in cupboard
Slide block
Stack blocks
Figure F.8: Visualization of the simulation tasks. From top bottom: place wine, put in cupboard, slide block, stack blocks
Pick Shape. The task consists of picking a colored shape and putting it inside red box. In each scene, there is an orange star, blue cube and pink cross. All the shapes are visible in the scene. The task is considered correct if the robot releases the right shape inside the red box. Training demonstrations: The shapes are randomly placed in the scene, while the red box is always left side of the robot, close to the base. Test out of domain: the shapes are randomly placed in the scene, but the box positioned the other side in position far from the (fixed) training one.
the
Push. The agent is required to pick a screwdriver and use it to push a green cube. We considered the task correct if the robot touched the green cube with the screwdriver. Training demonstrations: the cube is never moved from the position while the screwdriver is placed on the left, on the right and in front of the cube. Test out of domain: the blocks and the screwdriver are randomly placed in the scene.
Place Object. The task consists of picking an object selected by the user and put it inside box. In each scene, there is tape, stapler and yellow duck (that is a distractor). All the shapes are visible the scene. The task is considered correct if the robot releases the right object inside the brown box. Training demonstrations: The objects
randomly placed in the scene, while box is always on the left side of the robot. Test out of domain: the objects are randomly placed in the scene, but the box is positioned on the other side in position far from the training one.
Erase. The agent is required to pick an eraser and use it to erase colored spot among two vari- ations. We considered the task correct if the robot erase slide the object on the right color. Training demonstrations: eraser is in the middle and the two colored spot in fixed position. Test of domain: the eraser is randomly placed.
21

---

Published as a conference paper at ICLR 2025
Place block
Sort shape
Push block
Place object
Erase
Figure F99: Visualization of the real tasks. From top bottom: place block, sort shape, push block, place object, erase.
G EXTENDING THE WORLD MODEL TO ARTICULATED OBJECTS
Automatic modeling of articulated objects remains an open challenge (Kerr et al., 2024; Liu et al., 2023; Weng al., 2024). While DREMA does currently identify articulated objects, does support their simulation (for instance, the robotic arm itself is an articulated object). To model articulated objects, three key steps are required: (1) recognizing the object's components, (2) iden- tifying joint types and parameters (e.g., rotational about point or translational along an axis), and (3) determining the joint's location. Once these aspects are defined, URDF can be generated for simulation (Chen et al., 2024b).
Several promising directions could help address these challenges. One possible approach would leverages a robot's 8 human motion to infer object movement and joint properties, as demonstrated by Kerr et al. (2024) that used foundation models such as SAM (Kirillov et al., 2023) to identify object parts, while an optimization-base approach improves joint movement. Extracting the meshes of the parts and encoding their movement within a URDF would enable the automatic creation of articulated objects for DreMa. Alternatively, the articulated structure could be inferred directly from semantic information, as demonstrated by Chen et al. (2024b), where a model could trained to reconstruct articulated objects directly.
22

---

Published as a conference paper at ICLR 2025
In the current version, robot links segmented manually, while the model is assumed to be given. Automating this process could be achieved by segmenting links using a foundation model (Kirillov et al., 2023) or trained segmentation model for this task.
H DISCUSSION ON FAILURE SOURCES
The current implementation of DREMA is highly dependent on the segmentation mask. As shown in Figure H.10, if the segmentation mask for the objects is incorrect, the robot cannot interact with them. Similar issues may arise from incorrect mesh predictions, which PyBullet uses to detect collisions. Inaccurate meshes can lead to unexpected collisions or missed collisions.
The verification process prevents generating incomplete or incorrect demonstrations. Incomplete demonstrations can result from singularities in the trajectory or from trajectories that require unattainable robot configurations. Incorrect demonstrations often occur when object rotations move them outside the table, causing them fall.
When PerAct is trained with such faulty demonstrations, its performance is significantly impacted. Even small number of problematic demonstrations can make the model incapable of executing the task.
Figure H.10: From top to bottom consequences wrong segmentation, consequences of wrong mesh prediction, consequence of check of correct execution. In the first two rows the robot cannot pick the objects, while in the last row, can't perform the actions they requires unreachable configurations.
I EXTENDED RELATED WORKS FOR IMITATION LEARNING
I.1 IMITATION LEARNING
Imitation learning involves collecting a dataset demonstrations and learning a task from this data. Recent research has focused on vision-based imitation learning, where the policy is learned and ex- ecuted directly from RGB or RGB-D images (Zeng et al., 2021; Shridhar et al., 2023; 2022; Goyal al., 2024). A significant advancement in this area Transporter (Zeng et al., 2021), which in- troduced a network for executing conk-oonditioned placing. The authors analyze data augmentation and demonstrate its benefits for learning tasks, but argue that imagining various object configura-
23

---

Published as a conference paper at ICLR 2025
tions is unrealistic. However, show that generating demonstrations is feasible, thanks to recent advances in foundation models, Gaussian splatting, and physics simulation.
CLIPort (Shridhar et al., 2022) adapts Transporter for multi-task object manipulation by incorpo- rating CLIP features into the action computations. However, it is limited to simple pick-and-plae place tasks in 2D top-down settings. C2F-ARM (James et al., 2022) and PerAct (Shridhar et al., 2023) voxelized the environment and inferred actions directly in 3D. RVT (Goyal et al., 2023) builds on PerAct and uses the point cloud to render images from an optimal viewpoint to predict subsequent actions. All of these models require around one hundred demonstrations to learn a task. Recently, GNFactor (Ze al., 2023) and RVT2 (Goyal et al., 2024) reduced the number of examples needed using NeRF and a coarse-to- to-fine inference, respectively. However, they still require many examples to perform tasks accurately.
this work, we demonstrate that by reconstructing a simulation of the scene and imposing equivariant constraints on the data, we can generate new data and further reduce the number of demonstrations needed.
1.2 DATA AUGMENTATION FOR IMITATION LEARNING
Data augmentation a well-established practice Computer Vision. Early milestone works (Krizhevsky et al., 2012; He et al., 2016) applied random flipping and cropping improve model robustness. However, data augmentation in robotics presents unique challenges, as the robot's actions directly affect the environment. Thus, augmenting data is not always straightforward. Laskin et al. (2020) demonstrated that image augmentations, such as cropping or color jittering, can help agents learn better policies. However, their approach does not address augmenting robot actions. In contrast, our approach generates new trajectories.
In MoCoDa (Pitis al., 2022), data generation is achieved by creating counterfactual transitions between states and actions. Recently, Corrado & Hanna (2024) proposed dynamic invariant aug- mentations, showing that expanding the state-action coverage of the training data is more effective than increasing transitions.
The recently introduced MimicGen (Mandlekar et al., 2023) demonstrated that imitation learning can benefit from data augmentation by composing different object-centric segments. The main lim- itation of these approaches that they can only generate data by composing known states. In contrast, our approach generates new data by imagining novel equivariant configurations of objects and trajectories.
J ENGINEERING THE WORLD MODEL.
The proposed compositional manipulation world model is significant undertaking, requiring care- ful engineering to ensure its utility with real-world robots. However, the advantage of this approach is that, with the compositional manipulation world model we have developed, learn novel policies on a real robot with just a single example. For transparency and reproducibility, we provide detailed description of the key engineering efforts involved.
Low-resolution rendering. While Gaussian Splatting is designed for novel view synthesis, we observed unexpected behavior with low-resolution images, such as those in RLBench (James al., 2020). Initially, we collected images at the original resolution of 1280 x 720 and rendered them at 128 128, the resolution required by PerAct. However, this change in resolution led incorrect renderings of RGB images, and, more notably, depth images with inaccurate distance estimates from the cameras. a result, the robot incorrectly predicted the consequences of its actions. To address this issue, we simply render the images at higher resolution and then down-sample them to the desired resolution.
Point cloud re-projection from 2D Gaussian Spatting. PerAct uses depth images to re-project the point cloud in 3D. Unfortunately, we noticed that the depths rendered from the Gaussians repre- sentation along the edges of the robot and the objects were not accurate enough. The consequence was to have noisy point clouds that affected the learning of the agent. We solve this problem by fil- tering out Gaussians near edges using the radius outlier removal from Open3D (Zhou et al., 2018).
24

---

Published as a conference paper at ICLR 2025
Mesh extraction and physics engines. As motivated earlier, to model the object assets and their dynamics, physics engines require objects to be represented by mesh grids and not Gaussian blobs. We therefore adopt the 2DGS Gaussian Splatting (Huang et al., 2024) with TSDF to extract high- quality meshes, since the original Gaussian Splatting does not guarantee that the Gaussians align with the surfaces. Sugar (Guédon & Lepetit, 2024) proposes additional regularizations improve mesh extraction. However, we found that their meshes were not accurate enough for reliable simu- lations. As a physics engine, we use PyBullet as a core component.
Flat surfaces. To avoid falling objects in the simulator, the table is extracted by filtering the depth by using the segmentation mask. RANSAC (Fischler & Bolles, 1981) is used to estimate a plane necessary to create fixed convex hull that functions playground for the robot.
Camera Calibration. Unlike the graphics literature, which primarily focuses on the visual quality of reconstruction, and physics simulators designed for general or manually specified scenes, our goal is to integrate automated scene reconstruction with physics simulators. This requires aligning the geometry of the scene reconstruction with the geometry of the physics simulator. We achieve this by leveraging the calibrated camera mounted on top of the robot (Allegro et al., 2024). Additionally, since the robot agent is aware of its own dynamics, we can use its kinematics to estimate the extrinsic camera parameters as the robot moves. For more details on the real experiment, see Sec. 5.3, and for further information on camera calibration, refer to App. M.
End-effector manipulations. To move the robot inside the simulator, we specify a target pose that uniquely identifies the position and orientation of the end-effector. The robot is then controlled in position until the target is reached. Using the robot's ' kinematic model, expressed in the URDF, we can locate the different links and update the Gaussian accordingly.
Background. Once we have the Gaussian representation of the objects, then extract the repre- sentation of the table surface and the surrounding environment. This is done by removing the object Gaussians from the rest of the scene to avoid duplicates.
Physical parameters. For the physics simulator, the geometry of the objects is not enough; various physical parameters are also required. Some physical parameters can be assumed to be constant, such as the gravitational constant or air resistance, while others are not, including the masses of objects, object elasticity, or the friction coefficient. In general, estimating physical parameters is whole exciting field on its own (Mehra, 1974; Chebotar et al., 2019; Gao et al., 2022; Huang et al., 2023) and some of these methods, such as ASID (Memmel et al., 2024) or AdaptSim (Ren al., 2023), can be applied to identify them. In this paper, we focus on learning a photo-realistic representation of the world powered by a physics simulator, and its usefulness for compositional manipulation world models and imitation learning. Keeping in mind the overall complexity, we consider physical parameter identification to be out of scope, yet an important direction for future work.
K OPEN-VOCABULARY SEGMENTATION IN SIMULATION AND IN REAL WORLD
Segmenting simulation images with Grounded-SAM (Ren al., 2024) proved challenging due to the domain gap. While the typical gap occurs from simulation to the real world, in our case, we encountered the opposite issue. The apen-vocabulary segmenter, which was primarily trained on real data, struggled to accurately detect objects. the table, and the robot in the simulation. We observed that while SAM could identify objects, grounded-DINO misclassified instances, assigning incorrect labels. Figure K.11 shows examples of predicted masks in the simple environment with a cube in the center. In contrast, using only simple prompts such as "table" and "'bjecct,, and filtering out masks not on the table, the segmenter could correctly identify objects in the real data. Figure K. 12 displays examples from the real setting. One possible solution to improve recognition in the simulation is to fine-tune the segmenter to align the two domains. However, since our goal is to develop world model with real- world applicability, we chose to focus on enhancing the zero capabilities for the real world.
25

---

Published as a conference paper at ICLR 2025
Figure K.11: Segmentation masks in simulation. The middle column was generated with the prompt "object", whereas the right column was generated with the prompt "green cube".
Figure K.12: Examples of segmentation masks from the real images. These were generated by prompting "object" and "table".
L WORLD MODEL LEARNING DETAILS
To avoid falling objects in the simulator, the table is extracted by filtering the depth using the seg- mentation mask. RANSAC (Fischler & Bolles, 1981) is used estimate the plane necessary to create fixed convex hull. Finally, the entire scene is reconstructed, and the Gaussians of the objects are removed from the reconstruction. The procedure produced set of Gaussians paired with the mesh for each object and the Gaussians of the scene without the extracted objects.
World Model Learning. The set of images X collected using a similar approach of Ze et al. (2023) collecting 200 images from rotating camera. We set the resolution of the camera to 1280 x 720, mimicking real one. To avoid falling objects in the simulator, the table is extracted by re-projecting the depth filtered by the segmentation mask of this particular object. RANSAC (Fis- chler & Bolles, 1981) is used to estimate plane, necessary to create fixed convex hull used as workspace. The Gaussians of each object are trained for 7000 iterations. The meshes are extracted starting from the Gaussians using TSDF and the hull. The table is used to remove the parts of the mesh penetrating the table. Without this filtering, PyBuller may detect collisions creating problems in the simulation. A similar operation should be also done in case of touching objects. The entire scene is reconstructed. The Gaussians of each object are matched with their counterpart the entire scene using KNN and then removed to avoid duplications. The procedure produced a set of Gaus- sians paired with the mesh for each object, and the Gaussians of the scene without the extracted objects.
In imitation learning, the agent learss from dataset Z (G....,, SM. consisting of demonstra- tions, each paired with corresponding task encoding T = (4, (4,,..., 1M}. When tasks described
26

---

Published as a conference paper at ICLR 2025
by language- based goals, the task fm is represented by language embedding (e.g., pre-trained CLIP embeddings) corresponding to the phase that describes the task. Each demonstration Gm consists of sequence of continuous actions A = (AI..... .,)) which are encoded end-effector poses and gripper states. These actions are paired with observations Q = (g.....,0) ..)., which in our case can be the object assets from the compositional manipulation world model.
Data Generation. Following PerAct (Shridhar et al., 2023), we use a dataset Z C1, C2........ Sn of demonstrations, each paired with corresponding task encoding T = &1, t2, .......... ty. For language- based goals, the task & is the language embedding (e.g., a pre-trained CLIP embedding) of the phrase describing the task. Each demonstration is sequence of continuous actions A d1, a2...... at paired with observations Q - (g)...., gt). An action a consists of the 6-DoF pose, gripper open state, and whether collision avoidance was used by the motion planner reach intermediate pose: a "poss: "open- Gcollide- An observation & consists of RGB-D images from any number of cameras. We use three cameras for simulated experiments, Isim 9trent: Peet: Gright. and a single camera for real-world experiments, Greal = Qrroat. Each image in Q has dimensions 128 x 128.
To generate the data, the robot executes A in PyBullet and renders new set of observations using simulated cameras located at the same positions as 7si-- When the robot successfully completes a, the corresponding simulated observation is rendered, and the robot proceeds to execute a;+1. is important to note that the set of cameras Gfrom Qeett Fright is not used for optimizing the Gaussian representations. Instead, the novel-view synthesis capability of Gaussian Splatting is utilized. At the end of the execution, the objects' locations are saved. This data is then used to verify if the actions in the new environmental configurations can successfully complete the task.
The new configuration of the environment is generated by translating and rotating objects and the trajectory, while keeping the robot in the same position. We translate the environment along the and y axes using combinations of [0.0, 0.15, 0.155. The z-axis is not augmented, the test environments do not involve different table heights. However, such augmentation could be beneficial if the test environment includes varying heights. The rotation of the table and objects is performed around the position [0.30, 0.0, 0.0). with a complete rotation around the z-axis in steps of 30" For the trajectory rotation, the step size is reduced to 20". For object picking, no specific operations are applied since the demonstration provides the correct picking locations. However, reconstruction errors could impact the simulation and prevent correct object picking.
In single-task settings, PerAct is evaluated on nine tasks. For these tasks, DREMA generates n valid examples, where n is as follows: 112 for close jar, 136 for insert peg, 142 for lift. 142 for pick cup, 81 for sort shape, 36 for place wine, 58 for put in cupboard, 61 for slide block, and 63 for stack blocks. These demonstrations correspond to cases where the goal position was correctly reached in the augmented environment, as verified in Sec. 4. Each original demonstration allows for a maximum of 37 generated trajectories, with 37 original demonstrations used across the nine tasks. This results in theoretical maximum of 1369 demonstrations, but filtering reduces this to 831, retaining approximately 60% of the generated data.
We noticed that the jars and shape sorter are fixed in the simulator, experiments. Similarly, the cupboard, which is floating in the simulator, was adjusted accordingly in the reconstruction. Our goal is to create representation that the robot can manipulate and use to learn tasks that visually resemble the real-world scene.
also fixed them in the
Training. In the single-task setting, PerAct is trained using the same parameters as the original implementation, except for a reduced batch size of 4 and a decreased number of iterations to 100k. We also trained PerAct slide block, close jar, and sort shape tasks with the original batch size and found minimal differences in accuracy compared to the reduced batch size (close jar: +1.5, slide block: +1.2, sort shape -2.8). For the multi-task setting, we follow the single task setting while increasing the number of iterations to 600k. All training and testing experiments with PerAct are conducted on an Nvidia A40 GPU.
Validating and Testing. Model selection crucial to avoid deploying suboptimal policies (Shrid- har et al., 2023). Compared to the original study, the number of validation examples was increased from 25 to 40 to improve the selection of models. Similarly, the size of the test set was increased from 25 to 50 to accommodate wider range of evaluation scenarios. In addition, the test was
27

---

Published as a conference paper at ICLR 2025
peated five times to incorporate the variability introduced by motion planning and environmental interactions.
Rendering problems. Gaussian Splatting is designed for novel view synthesis. However, ticed strange renderings when it is trained with high-resolution images and deployed to generate low-resolution ones, that needed for the experiments in RLBench (James et al., 2020). Fig- ure L.13 shows how these problems affected the rendering of RGB and depth images. In particular, depth images must be carefully estimated since PerAct uses them create voxel representation of the environment. The rendering problem caused the robot to predict the actions incorrectly. To overcome this problem, we rendered higher-resolution image and down-sampled it to the desired resolution.
Point Cloud re-projection from 2D Gaussian Spatting. PerAct uses depth images to re- re-project the point clouds in 3D. We noticed that the depth images rendered by the Gaussians inaccurate along the edges of the robot and the objects. This created noisy point clouds that affected the agent. Figure I., 14 shows point cloud re-projected from the generated depth. We solve it by filtering using the radius outlier removal from Open3D (Zhou et al., 2018). This was particularly effective with far cameras but less effective with the wrist camera. Therefore, we used three cameras instead of four for the experiments.
Figure L.13: Three left images show the rendered RGB, the same RGB rendered with a higher resolution, and the original image. On the right, the original point cloud (red) does not align with the rendered one (green) when the image is directly rendered in low resolution.
Figure L.14: Re-projected point clouds before (left) and after (right) the filter.
M MORE INFORMATION ON THE REAL EXPERIMENT
In this section, we provide a more detailed description of how we learned the world model and deployed imitation learning in real robot setup. Moreover, we add some brief discussions on problems we encountered.
Collecting demonstration in the real world. One of the most essential parts of imitation learn- ing is data collection. At the beginning and the end of the demonstration, we recorded the images used to reconstruct the world model and checked that with the physical parameters used, the robot could replay the same trajectory. The demonstrations are collected only from the external Kinect V2 camera, positioned to resemble the front camera of RLBench as in the real experiment of Per- Act (Shridhar et al., 2023). The keypoints to execute were manually saved through the RVIZ inter- face in ROSI (Quigley et al., 2009). When the keypoints were established, we used sampling- -baed motion planning to generate the execution trajectory. During the execution, recorded Ros bag
28

---

Published as a conference paper at ICLR 2025
later processed to create data in RLBench format'. In collecting the data, noticed that quality is essential for correctly training the model. Sometimes formatted data contained more keypoints in similar positions, creating loops in the agent that could not exit. We encountered problems when the same task in a similar position had a different number of keypoints. We believe that these problems reduced by increasing the batch size, but due to time constraints, we could not verify this sumption. To avoid these issues, we collected demonstrations with a constant number of keypoints, and manually cleaned the data when we noticed unwanted keypoints.
World model creation and data generation. The images needed to construct the world model were collected with the camera on the robot. The camera was calibrated, and the inverse kinematics noticed that Gaussian Splatting requires accurate
was used to retrieve the camera pose. However, poses. Sometimes, the movements of the robot or the interactions with the objects made the camera move. Consequently,
had calibrate the camera again. The alignment between the world model and the real world played another important role. We observe that this should be accurately estimated. Otherwise, the replayed trajectory would miss the objects.
Training and testing PerAct. We trained PerAct with the original parameters but with batch size of 4 for 100k iterations. This was constrained by the training time required since we wanted to test policies with different training data. The original implementation saves a model every 10k iterations, We reduced this to 5k to keep more models close to the final part of the training. Model selection in the real world was not straightforward. We noticed that choosing the wrong model highly impacts the accuracy. While collecting a validation set and checking the loss is possible, this is not the best solution because the only way to understand the best model is by making the agent interact with the environment. Consequently, we tested the last five models with a minimal number of configurations. We think that model selection in the real world is crucial problem that opens the way to possible future works for our world model.
'https:///o wht.. com/ba hwtp://ww////c/////caaacccaoccaa rosbag data aicila/tree tsttte/maaaan
29

---

Published as a conference paper at ICLR 2025
N
EXAMPLES OF GENERATED DATA
Figure N.15: Visualization of the close jar task
30

---

Published as a conference paper at ICLR 2025
Figure N.16: Visualization of the sort shape task
31

---

Published as a conference paper at ICLR 2025
Figure N.17: Visualization of the slide block task
32

---

Published as a conference paper at ICLR 2025
Figure N.18: Visualization of the pick block task
33

---

Published as a conference paper at ICLR 2025
Figure N.19: Visualization of the pick shape task
34

---

Published as a conference paper at ICLR 2025
Figure N.20: Visualization of the push task
35


Reviews:
[{"review_id": "4KOgEDiWh6", "reviewer": "Reviewer_zmyw", "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.", "questions": "Figure 2: It would be helpful to include how “open-vocabulary tracking models” fit into the pipeline. \n\nL107: “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.\n\nL128: Why is (Cheng et al., 2023) used as a reference for the zero-shot capabilities of foundation models?\n\nL157: 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. \n\nL174: Set $\\mathcal{X}$ notation is used for a sequence.\n\nL267-268: Other objects in the world that the robot arm may interact with could also be articulated? ie. cabinets \n\nL409: 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?\n\nL412-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?\n\nL418: Why “112, 61, and 81” demonstrations for the three tasks? How were these number of “imagined” demonstrations chosen? \n\nL485: How were OOD test conditions chosen? Would they be within the set of equivariant transformations evaluated in Table 2? \n\n[minor editing]\n\nFigure 2: part 3, typo in Gaussian\n\nL105: “gaussian” -> “Gaussian”\n\nL233: $x_{n,k}$ should be $y$?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper tackles the data efficient regime, and their proposed approach uses a single demonstration\n\n- 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. \n\nThe 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? \n\n- 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?\n\n- There are no metrics on runtime performance of the proposed method, while the introduction mentions the “real-time” performance of Gaussian splatting. \n\n- Table 1: Comparisons to PerAct (Shridhar et al., 2023) are done using only 5 episodes per task, while PerAct uses 25 episodes per task. \n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8x7TEC7oue", "reviewer": "Reviewer_kqVG", "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.", "questions": "1. What’s the relationship with digital twin line of work?\n2. How to learn 3d models for parts of articulated objects and how to imagine demonstrations that manipulate articulated objects?\n3. Can you replay imagined trajectories in simulator to verify correctness? Will that help improve imitation success?\n4. 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! \n5. Can you add one more baseline method of data augmentation to compare to?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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. \n2. 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. \n3. Existing approach of verify imagined demonstration is rudimentary. \n\n[1] MIRA: Mental Imagery for Robotic Affordances", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oKXoMqvhkR", "reviewer": "Reviewer_Mne8", "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.\n\n**Strengths:**\n\n1. The paper introduces an innovative approach by leveraging world models to generate robotic data for imitation learning, which is a contribution to the field.\n2. The experiments are detailed, covering both simulation environments and real-world robot demonstrations, providing a robust evaluation of the approach.\n3. A creative method for augmenting data used in imitation learning is introduced, which could lead to improved learning efficiency.\n\n**Weaknesses:**\n\n1. The absence of publicly available source code limits the reproducibility of the results. It is suggested to release the code during the rebuttal stage.\n2. Some figures in the paper need improvement, as the text in several instances is too small to read clearly.\n3. 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.\n\nIn 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.", "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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n2. The experiments are detailed, covering both simulation environments and real-world robot demonstrations, providing a robust evaluation of the approach.\n3. 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.\n2. Some figures in the paper need improvement, as the text in several instances is too small to read clearly.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)