Learning Visual Feature-Based World Models via Residual Latent Action

TL;DR

This paper introduces Residual Latent Action (RLA), a novel predictive latent representation for visual feature-based world models, enabling faster and more accurate policy learning from offline videos.

Contribution

It proposes RLA-WM, a flow-matching world model using RLA, and demonstrates its effectiveness in simulation, real-world datasets, and robot policy learning from offline videos.

Findings

RLA-WM outperforms state-of-the-art models in accuracy and speed.

RLA enables learning from actionless demonstration videos.

Visual RL trained entirely offline achieves competitive performance.

Abstract

World models predict future transitions from observations and actions. Existing works predominantly focus on image generation only. Visual feature-based world models, on the other hand, predict future visual features instead of raw video pixels, offering a promising alternative that is more efficient and less prone to hallucination. However, current feature-based approaches rely on direct regression, which leads to blurry or collapsed predictions in complex interactions, while generative modeling in high-dimensional feature spaces still remains challenging. In this work, we discover that a new type of latent action representation, which we refer to as *Residual Latent Action* (RLA), can be easily learned from DINO residuals. We also show that RLA is predictive, generalizable, and encodes temporal progression. Building on RLA, we propose *RLA World Model* (RLA-WM), which predicts RLA values via flow matching. RLA-WM outperforms both state-of-the-art feature-based and video-diffusion world models on simulation and real-world datasets, while being orders of magnitude faster than video diffusion. Furthermore, we develop two robot learning techniques that use RLA-WM to improve policy learning. The first one is a minimalist world action model with RLA that learns from actionless demonstration videos. The second one is the first visual RL framework trained entirely inside a world model learned from offline videos only, using a video-aligned reward and no online interactions or handcrafted rewards. Project page: https://mlzxy.github.io/rla-wm