PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Leonardo Barcellona, Andrii Zadaianchuk, Davide Allegro, Samuele Papa, Stefano Ghidoni, Efstratios Gavves

arXiv:2412.14957v2cs.ROcs.CV

TL;DR

Existing robot world models often lack explicit, realistic representations of the observed environment and its dynamics, limiting grounded imagination for robotics. DREMA constructs a compositional digital twin with Gaussian Splatting and physics simulation, then transforms demonstrations to generate imitation-learning data. The authors report improved policy accuracy, robustness, and generalization, including one-shot learning of novel physical tasks on a real Franka Emika Panda robot.

  • Problem

    Current robot world models typically cannot explicitly imitate the observed environment with realistic behavior and dynamics for real-world robotics.

    Source: §1 · pp. 1–2

  • Method

    DREMA integrates object-centric Gaussian Splatting with physics simulation to model scenes compositionally and generate transformed demonstrations for imitation learning.

    Source: §1 · p. 2

  • Results

    DREMA improves policy learning, robustness, and generalization across evaluated settings, including one-shot learning of novel physical tasks on a real robot.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    DREMA enables robots to learn from minimal demonstrations by using grounded imagination to create and verify novel training data.

    Source: §1 · p. 2

    DREMA requires full environment observability, currently handles only simple objects, uses open-loop world-model construction, and may face increased complexity when scaled to large diverse environments.

    Source: §6 · p. 10

Abstract

from arXiv · show

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

1 Introduction

DREMA addresses the need for compositional, object-centric, and controllable robot world models by combining object-centric Gaussian Splatting with physics simulation. It uses this grounded model to transform demonstrations into verified, realistic training data for imitation learning.

1 Introduction

  • Current robot world models struggle to represent configurations beyond their training distributions, motivating models that are compositional, object-centric, and easy to manipulate.These requirements support novel valid combinations of learned concepts and controllable action-consequence modeling.Source: §1 · pp. 1–2
  • DREMA combines object-centric Gaussian Splatting with a physics simulator to construct a compositional manipulation world model that supports grounded imagination.The model reconstructs scenes as manipulable Gaussian-based objects and embeds them in physics simulation for physically reliable predictions.Source: §1 · p. 2
  • DREMA generates new imitation-learning demonstrations by applying equivariant roto-translation transformations to demonstrated objects and actions, then verifying successful task execution.The world model checks transformed goal positions and renders the resulting observations to retain valid demonstrations.Source: §1 · p. 2
  • The paper reports improved agent generalizability in simulation and real-robot experiments while requiring only a minimal number of demonstrations.The contribution explicitly targets policy learning from sparse demonstrations through compositional imagination.Source: §1 · p. 2

2 Related Work

Related work spans predictive world models, simulation-based scene reconstructions, and imitation-learning augmentation. DREMA differs by combining object-centric physical prediction with automatic compositional scene construction and novel object-action configurations.

2.1 World Models

  • World-model approaches predict future states from actions but often fail to capture robot-specific environment dynamics outside their training domains.DREMA instead combines photorealistic Gaussian representations with physics-engine dynamics for real-world imitation learning.Source: §2.1 · pp. 2–3

2.2 Simulation-based reconstructions

  • Simulation-based reconstruction methods combine visual scene representations with physics, but prior approaches are not consistently compositional or object-centric and often track observed trajectories rather than predict future states.DREMA automatically constructs a compositional world model using foundational models rather than requiring manual object-centric modeling or simulation-to-real adaptation.Source: §2.2 · p. 3

2.3 Imitation learning and data augmentation

  • Vision-based imitation-learning methods can generalize better but commonly require hundreds of demonstrations, whereas DREMA generates new environmental configurations rather than only novel camera views.Its Gaussian-Splatting representation supports configuration changes involving objects and actions.Source: §2.3 · pp. 3–4
  • Robotics augmentation methods expand datasets through trajectory recombination or invariant transformations, while DREMA imagines novel equivariant object-action configurations.The paper states that this approach works with as few as one demonstration, unlike MimicGen’s requirement for multiple demonstrations.Source: §2.3 · pp. 3–4

3 Dream-to-Manipulate with Gaussian Splats & Physics Engines

DREMA constructs a compositional manipulation world model from explicit object-centric scene representations, physics-based dynamics, and controllable robot-object assets. It combines Gaussian Splatting with physics simulation to render and predict grounded future states, while verification and engineering choices support practical deployment.

3 Dream-to-Manipulate with Gaussian Splats & Physics Engines

  • DREMA represents a world as object assets, a robot agent, dynamics, and manipulation tasks, using an RGB-D video as input.The dynamics operator transforms object and agent states under robot actions.Source: §3 · p. 4

3.1 Representing 3D Objects with Gaussian Splatting

  • Gaussian Splatting converts video observations into photorealistic 3D representations that can render objects from known and novel viewpoints.Each Gaussian stores position, orientation, scale, color, and opacity, enabling real-time novel-view synthesis.Source: §3.1 · p. 4
  • Object-centric Gaussian Splats are obtained by segmenting individual objects across video frames and applying depth-supervised optimization.The decomposition supports independently controllable object assets, while depth regularization penalizes outlier Gaussian depths.Source: §3.1 · pp. 4–5, §3.2 · p. 5
  • The representation cannot currently model deformable or articulated objects because standard Gaussian Splats assume static scenes.Dynamic Gaussian Splatting is identified as a future direction for articulation and deformation.Source: §3.1 · p. 5

3.3 Modelling Robot-Object Dynamics with Object Assets

  • DREMA converts Gaussian object representations into meshes for PyBullet, whose physics dynamics update object positions and orientations under robot interaction.Updated mesh states are propagated back to the Gaussians so the simulator can predict and render future consequences of actions.Source: §3.3 · p. 5
  • World-model predictions are not guaranteed for arbitrary actions because reconstruction, meshing, parameter, and alignment errors can accumulate; demonstration replay is therefore used as a practical verification criterion.The robot arm is modeled using known joints and an assumption of perfectly modeled self-dynamics.Source: §3.3 · p. 6

3.4 Engineering

  • Real-world deployment required engineering for low-resolution rendering, inaccurate depth, reconstruction-simulator alignment, and physical parameters, including a switch to 2DGS and calibrated camera alignment.The experiments use constant physical parameters selected on validation sets.Source: §3.4 · p. 6

4 Equivariant Transformations for Imitation Learning

DREMA augments imitation-learning demonstrations by transforming object observations and robot actions together, preserving task semantics under geometric changes. The generated demonstrations are rendered and checked in the world model before use.

4 Equivariant Transformations for Imitation Learning

  • Equivariant augmentation must transform observations and corresponding actions together so altered demonstrations remain semantically consistent and solve the intended task.This differs from standard image augmentation because imitation-learning actions must remain aligned with changed observations.Source: §4 · p. 6
  • DREMA uses explicit object representations to apply rigid transformations equivariantly to object positions, end-effector actions, and target locations without changing the robot’s internal state.Executing transformed actions and rendering transformed observations produces novel demonstrations absent from the original dataset.Source: §4 · p. 7
  • Roto-translation transforms the environment, object assets, and each action’s end-effector translation and rotation around a chosen application point.The fixed robot executes the transformed action sequence in the transformed environment.Source: §4 · p. 7
  • Object rotation rotates initial object poses and the action trajectory around the final end-effector position while leaving the environment unchanged.The unchanged environment is appropriate because the trajectory brings the object to a similar final position.Source: §4 · p. 7
  • Transformed demonstrations are verified by checking whether final object positions remain within 0.015 meters of their expected transformed goals.This goal-based test replaces a general success detector when equivariant transformations define the expected outcome.Source: §4 · p. 7

5 Experiments

DREMA improves imitation learning from minimal demonstrations in simulation and on a real robot by generating imagined data through compositional world-model transformations. Its benefits persist with more original data and extend to in- and out-of-distribution real-world configurations.

5 Experiments

  • The study compares PerAct trained with original, imagined, or combined demonstrations across nine simulated tasks, using 50 configurations and five repeated runs.In real-world evaluation, models use original-only or combined demonstrations and are tested in 20 configurations, including 10 out-of-distribution object-position configurations.Source: §5.1 · p. 8, §5.3 · p. 9

5.1 One-shot Policy Learning: Single- and Multi-task

  • DREMA improves average accuracy by 9.1% over original-data PerAct in single-task settings and reaches 35.5% in multi-task settings, a 13.1% improvement.It also outperforms PerAct when trained only on imagined demonstrations and evaluated in the original environment.Source: §5.1 · p. 8, Table 1 · p. 8

5.2 Analysis and ablations

  • With 20 original examples, imagined data improve success by 6.4% for sort shape and 10.4% for slide block over training on the original examples alone.DREMA trained from five demonstrations matches PerAct trained from 20 examples on the reported tasks.Source: §5.2 · p. 9, Figure 4 · p. 9
  • Combining object rotation and roto-translation transformations yields a 10.4% mean-accuracy improvement, with roto-translation having the greatest individual impact.The transformations increase state-space coverage and support more general policies.Source: §5.2 · p. 9, Table 2 · p. 9

5.3 Experiments with a Real Robot

  • World-model predictions align well with simpler demonstrations, while errors increase for more complex examples and may require adaptation using real demonstrations.The real-robot setup compares localization errors between collected and predicted final object positions.Source: §5.3 · p. 9, Table 3 · p. 10
  • DREMA achieves 62.9% average real-robot task performance versus 31.7% for PerAct, while remaining robust to in- and out-of-distribution object configurations.Pick shape has the lowest accuracy, likely because small shapes are difficult to recognize.Source: §5.3 · p. 10, Table 4 · p. 10

6 Conclusions

DREMA constructs manipulation world models by integrating photorealistic Gaussian Splatting with physics simulation. The resulting model supports action prediction and policy learning from minimal data, while remaining limited by observability, object complexity, feedback, and scaling requirements.

6 Conclusions

  • DREMA integrates Gaussian Splatting and physics simulators to provide photorealistic, physics-constrained manipulation world models.It is presented as a framework for predicting action consequences and learning from minimal data.Source: §6 · p. 10
  • DREMA requires full environmental observability, currently handles only simple objects, uses open-loop construction without feedback, and may become computationally complex at scale.The paper identifies automatic URDF estimation and feedback incorporation as possible future directions.Source: §6 · p. 10

Reproducibility statement

The paper supports reproducibility by documenting its methods, experiments, and configurations and by committing to release code and synthetic simulation datasets.

Reproducibility statement

  • The authors provide methodological, experimental, hyperparameter, runtime, and inference details and plan to share code and synthetic simulation datasets.These materials are intended to enable replication and verification of the experiments.Source: Reproducibility statement · p. 11

Contributions

The project involved shared leadership, implementation, experimentation, data collection, writing, revision, and figure creation across the author team.

Contributions

  • LB implemented most code and ran simulation experiments, while DA managed data collection and real-world experiments with LB's support.EG led the project; LB and AZ shaped its direction and drafted the paper, with contributions from the remaining authors.Source: Contributions · p. 11

A Extended single task and multi-task results

DREMA augments demonstrations with invariant and equivariant transformations, and combining imagined with real data generally improves multi-task accuracy, except for the more demanding sort shape task.

Comparing with other augmentations

  • DREMA is compared against random RGB-D patches, table-color changes, and randomly added objects, using 80 augmented examples per method.These invariant augmentations are combined with the original demonstrations for PerAct training.Source: A Extended single task and multi-task… · p. 17
  • Combining invariant augmentations with equivariant augmentations is presented as having potential to produce more robust and generalizable policies.The supplied comparison describes the invariant baselines, while the proposed equivariant strategy changes object and action configurations.Source: A Extended single task and multi-task… · p. 17

Multi-task results

B Imagination results in the real experiment

The real-world pick block experiment tests whether a robot can learn through imagination, finding that performance depends on the number of original demonstrations.

Imagination results in the real experiment

  • With four original examples, the learned pick block policy performs worse than the initial policy, but performance improves significantly as more original examples are added.Table B.6 reports the complete real-world results for original and imagination demonstrations.Source: B Imagination results in the real… · p. 17, Table B.6 · p. 17

C Data Generation Analysis

DREMA expands the training distribution by imagining new object configurations and robot-action keypoint locations, supporting more robust policies.

Data Generation Analysis

  • DREMA increases the training distribution by generating high-quality data for new environment configurations, producing better and more robust policies.Gaussian Splatting keeps the visual domain gap small, while the simulator recreates robot-object interactions.Source: C Data Generation Analysis · pp. 17–18
  • Figure C.6 compares original keypoints with generated keypoints using one, five, and twenty examples for the sort shape task.Blue dots denote original data and red crosses denote generated data.Source: Figure C.6 · p. 18

D Runtime analysis

DREMA combines explicit Gaussian Splat representations with physics simulation to imagine rendered future states for imitation learning, extending world-model principles to robot demonstrations and tasks.

D Runtime analysis

  • DREMA’s runtime analysis reports 1.715 seconds per waypoint without rendering and 1.883 seconds with rendering for tasks including close jar, slide block, and sort shape.Rendering adds 0.168 seconds on average to update Gaussians, render five 128 × 128 RGB-D images, and filter them.Source: D Runtime analysis · p. 18

E Explicit World Models and Learnable Digital Twins

  • The explicit world model represents objects with Gaussian Splat assets and uses a physics simulator to generate novel RGB-D future states for policy training.Gaussian Splatting provides reconstruction renderings, while the physics engine constructs dynamics predictions.Source: E Explicit World Models and Learnable… · pp. 18–19
  • PerAct trained with imagined sensory data, and combining imagined with original data improves performance in simulation and real-world experiments.The imagined inputs are RGB and depth images rendered through Gaussian Splatting.Source: E Explicit World Models and Learnable… · p. 18
  • The work frames explicit digital twins and world models as converging approaches that let agents simulate and reason about their surroundings.DREMA incorporates real2sim and data augmentation while extending world-model principles.Source: E Explicit World Models and Learnable… · p. 19

F Tasks descriptions

  • DREMA is evaluated on simulation tasks defined by PerAct and on real tasks designed to be similar to those simulations.The simulation tasks are used without modifications.Source: F Tasks descriptions · p. 19

F.1 Simulation tasks

This section details DreMa’s evaluation and deployment across simulated and real-world manipulation tasks, covering data collection, world-model construction, data generation, and PerAct deployment. It emphasizes the requirements and challenges of reliable imitation learning, including reconstruction, segmentation, mesh accuracy, calibration, alignment, physical validity, and model selection.

F.1 Simulation tasks

  • DreMa evaluates nine simulated manipulation tasks spanning grasping, placing, sorting, sliding, lifting, and stacking, each defined by task-specific success conditions.Tasks include close jar, insert peg, lift, pick cup, sort shape, place wine, put in cupboard, slide block, and stack blocks.Source: §F.1 · p. 19

F.2 Real world tasks

  • Real-world evaluation tests pick-and-place and tool-use generalization by shifting object, target, or tool configurations beyond the training demonstrations.The tasks are Pick Block, Pick Shape, Push, Place Object, and Erase, with out-of-domain placements or object arrangements used at test time.Source: §F.2 · pp. 20–21

G Extending the world model to articulated objects

  • DREMA can simulate articulated objects but does not yet identify them automatically, requiring component recognition, joint estimation, and joint localization before URDF generation.Robot links are currently segmented manually, while future approaches could infer articulated structure from motion, foundation-model segmentation, optimization, or semantic reconstruction.Source: G Extending the world model to… · pp. 21–22

H Discussion on failure sources

  • Incorrect segmentation or mesh prediction can prevent robot-object interaction through missed or unexpected collisions, while faulty demonstrations can make PerAct incapable of executing the task.Verification filters incomplete trajectories caused by singularities or unreachable configurations and incorrect trajectories caused by objects falling outside the table.Source: H Discussion on failure sources · pp. 22–23

I.1 Imitation learning / I.2 Data augmentation for Imitation learning

  • Unlike prior imitation-learning methods that require many demonstrations or compose known states, DreMa generates new trajectories by imagining equivariant object configurations in a physics-powered reconstructed scene.The approach augments robot actions as well as observations, extending data beyond known states and reducing the demonstrations needed for accurate policies.Source: §I.1 · p. 23, §I.2 · p. 24

J Engineering the world model

  • The world model combines Gaussian Splatting with mesh extraction, PyBullet physics, camera calibration, kinematic end-effector control, table reconstruction, and object-specific physical modeling.High-resolution rendering followed by down-sampling fixes low-resolution RGB-D errors, while filtering edge Gaussians reduces noisy point clouds; physical-parameter identification remains out of scope.Source: J Engineering the world model · pp. 24–25

K Open-vocabulary segmentation in simulation and in real world

  • Grounded-SAM struggles in simulation because its real-data training creates a domain gap, whereas simple prompts and mask filtering can identify objects in real scenes.The authors prioritize improving zero-shot recognition in real-world data rather than fine-tuning the segmenter for simulation.Source: K Open-vocabulary segmentation in… · p. 25

L World model learning details

  • DreMa reconstructs object Gaussians and meshes, generates validated demonstrations through PyBullet execution, and augments translations and rotations while preserving the robot pose.Across nine single-task settings, filtering leaves 831 valid demonstrations from a theoretical maximum of 1369, approximately 60% of generated data; experiments use three simulated cameras and one real camera.Source: L World model learning details · pp. 26–27

M More information on the real experiment

  • Demonstrations used an external Kinect V2 camera, manually recorded RVIZ keypoints, sampling-based motion planning, and constant-keypoint cleaning to avoid trajectory loops.Images at the demonstration’s beginning and end were used for world-model reconstruction, while Ros bags were converted into RLBench-format data.Source: M More information on the real… · p. 28
  • Accurate camera poses and world-to-real alignment were essential because robot or object motion could disturb calibration and cause replayed trajectories to miss objects.The camera was calibrated using inverse kinematics, but Gaussian Splatting required recalibration when the camera moved.Source: M More information on the real… · p. 28

N Examples of generated data

Related papers

Loading 2412.14957v2…