LaDi-WM: A Latent Diffusion-based World Model for Predictive Manipulation
This work addresses a well-known bottleneck in Embodied AI for robotics by enhancing predictive manipulation through improved world modeling, with incremental advancements in method integration.
The paper tackles the challenge of generating accurate future visual states for robot-object interactions in predictive manipulation by proposing LaDi-WM, a world model that predicts latent spaces using diffusion modeling, resulting in policy performance improvements of 27.9% on the LIBERO-LONG benchmark and 20% in real-world scenarios.
Predictive manipulation has recently gained considerable attention in the Embodied AI community due to its potential to improve robot policy performance by leveraging predicted states. However, generating accurate future visual states of robot-object interactions from world models remains a well-known challenge, particularly in achieving high-quality pixel-level representations. To this end, we propose LaDi-WM, a world model that predicts the latent space of future states using diffusion modeling. Specifically, LaDi-WM leverages the well-established latent space aligned with pre-trained Visual Foundation Models (VFMs), which comprises both geometric features (DINO-based) and semantic features (CLIP-based). We find that predicting the evolution of the latent space is easier to learn and more generalizable than directly predicting pixel-level images. Building on LaDi-WM, we design a diffusion policy that iteratively refines output actions by incorporating forecasted states, thereby generating more consistent and accurate results. Extensive experiments on both synthetic and real-world benchmarks demonstrate that LaDi-WM significantly enhances policy performance by 27.9\% on the LIBERO-LONG benchmark and 20\% on the real-world scenario. Furthermore, our world model and policies achieve impressive generalizability in real-world experiments.