LaDi-WM: A Latent Diffusion-based World Model for Predictive Manipulation

TL;DR

LaDi-WM introduces a diffusion-based world model predicting future states in a latent space aligned with pre-trained visual models, significantly improving robot policy performance and generalizability in predictive manipulation tasks.

Contribution

The paper presents LaDi-WM, a novel latent diffusion-based world model that predicts future states in a learned latent space, enhancing accuracy and generalization over pixel-level predictions.

Findings

Improves policy performance by 27.9% on LIBERO-LONG benchmark.

Achieves 20% improvement in real-world scenarios.

Demonstrates strong generalizability in experiments.

Abstract

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.