GWM: Towards Scalable Gaussian World Models for Robotic Manipulation

TL;DR

This paper introduces GWM, a novel 3D Gaussian world model for robotic manipulation that improves future scene prediction and policy training by integrating a latent Diffusion Transformer with a 3D autoencoder, demonstrating superior performance in experiments.

Contribution

The paper presents GWM, a new Gaussian world model that reconstructs future states with Gaussian primitives, combining a latent Diffusion Transformer and 3D autoencoder for enhanced scene understanding.

Findings

GWM accurately predicts future scenes conditioned on robot actions.

Policies trained with GWM outperform state-of-the-art methods.

GWM demonstrates strong data scaling potential for 3D world modeling.

Abstract

Training robot policies within a learned world model is trending due to the inefficiency of real-world interactions. The established image-based world models and policies have shown prior success, but lack robust geometric information that requires consistent spatial and physical understanding of the three-dimensional world, even pre-trained on internet-scale video sources. To this end, we propose a novel branch of world model named Gaussian World Model (GWM) for robotic manipulation, which reconstructs the future state by inferring the propagation of Gaussian primitives under the effect of robot actions. At its core is a latent Diffusion Transformer (DiT) combined with a 3D variational autoencoder, enabling fine-grained scene-level future state reconstruction with Gaussian Splatting. GWM can not only enhance the visual representation for imitation learning agent by self-supervised future prediction training, but can serve as a neural simulator that supports model-based reinforcement learning. Both simulated and real-world experiments depict that GWM can precisely predict future scenes conditioned on diverse robot actions, and can be further utilized to train policies that outperform the state-of-the-art by impressive margins, showcasing the initial data scaling potential of 3D world model.