iMoWM: Taming Interactive Multi-Modal World Model for Robotic Manipulation

TL;DR

This work tackles the limitation of RGB-only world models in robotic manipulation by introducing iMoWM, a geometry-aware, interactive multi-modal world model that predicts color images, depth maps, and robot-arm masks conditioned on actions. Central to the approach is MMTokenizer, which unifies multi-modal inputs into compact tokens to enable efficient use of large pretrained video models within an autoregressive transformer framework. The method demonstrates state-of-the-art performance in action-conditioned video generation, as well as improved model-based RL rollouts and real-world imitation learning on benchmarks like BAIR, RoboNet, and Meta-World. By explicitly incorporating 3D information and multi-modal signals, iMoWM provides a scalable neural simulator that enhances robotic manipulation training and data efficiency for both RL and IL scenarios.

Abstract

Learned world models hold significant potential for robotic manipulation, as they can serve as simulator for real-world interactions. While extensive progress has been made in 2D video-based world models, these approaches often lack geometric and spatial reasoning, which is essential for capturing the physical structure of the 3D world. To address this limitation, we introduce iMoWM, a novel interactive world model designed to generate color images, depth maps, and robot arm masks in an autoregressive manner conditioned on actions. To overcome the high computational cost associated with three-dimensional information, we propose MMTokenizer, which unifies multi-modal inputs into a compact token representation. This design enables iMoWM to leverage large-scale pretrained VideoGPT models while maintaining high efficiency and incorporating richer physical information. With its multi-modal representation, iMoWM not only improves the visual quality of future predictions but also serves as an effective simulator for model-based reinforcement learning (MBRL) and facilitates real-world imitation learning. Extensive experiments demonstrate the superiority of iMoWM across these tasks, showcasing the advantages of multi-modal world modeling for robotic manipulation. Homepage: https://xingyoujun.github.io/imowm/

Figures (7)

• Figure 1: Given color images, depth maps, and robot-arm masks, iMoWM generates future 3D multi-modal fields conditioned on actions. By incorporating geometric information, this multi-modal design enhances the representation of dynamic real-world scenes and improves the applicability of world models to robotic manipulation, supporting both real-world imitation learning and visual model-based reinforcement learning across diverse tasks.
• Figure 2: Framework of proposed iMoWM. Our method takes a color image as input and extracts the robot-arm mask. The metric depth map is obtained either from RGB-D camera observations or from our depth module. These inputs are then encoded into discrete tokens using the proposed MMTokenizer. By injecting action-conditioned slot tokens, the autoregressive transformer generates future tokens sequentially, which are finally decoded into multi-modal outputs.
• Figure 3: Overview of MMTokenizer. We introduce MMTokenizer to generate unified tokens from multi-modal inputs, which not only reduces computational cost but also incorporates geometric information into the tokens, enhancing the capabilities of iMoWM. Specifically, we repeat the first and last layers of the original VAE and combine different embeddings to mitigate distributional differences across modalities.
• Figure 4: Comparisons on BAIR dataset. We present action-conditioned multi-modal video generation results on the BAIR dataset, comparing with iVideoGPT. Our method produces better generations and captures richer geometric information.
• Figure 5: Comparison on real data. We showcase qualitative results on real-world, high-resolution data. The results demonstrate the superiority of iMoWM in controllable multi-modal video generation with interaction.