🤖 AI

Learning Visual Feature-Based World Models via Residual Latent Action

This paper introduces the RLA World Model, which leverages a novel Residual Latent Action representation learned from DINO residuals and predicted via flow matching to achieve efficient, high-fidelity visual feature prediction and enable advanced robot learning from offline videos without online interaction or handcrafted rewards.

Original authors: Xinyu Zhang, Zhengtong Xu, Yutian Tao, Yeping Wang, Yu She, Abdeslam Boularias

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xinyu Zhang, Zhengtong Xu, Yutian Tao, Yeping Wang, Yu She, Abdeslam Boularias

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to move a cup, a ball, or a tool. To do this, the robot needs a "world model"—a mental simulation that lets it guess what will happen next if it moves its arm in a certain way.

For a long time, scientists tried to build these mental simulations by predicting the exact pixels of the next video frame (like trying to predict every single grain of sand in a sandcastle before it's built). This is computationally heavy, slow, and often leads to the robot "hallucinating" things that don't exist, like a cup suddenly turning into a cat.

Other researchers tried to predict visual features (like the general shape and color of objects) instead of raw pixels. This is faster, but it often results in "blurry" predictions, like looking at a photo through a foggy window. In complex 3D tasks, these predictions would sometimes collapse entirely, making the robot lose track of where things are.

This paper introduces a new solution called RLA-WM (Residual Latent Action World Model). Here is how it works, using simple analogies:

1. The "Delta" Trick (Residual Latent Action)

Instead of trying to predict the entire next picture of the world, the authors realized the robot only needs to learn the difference between "now" and "later."

The Analogy: Imagine you are watching a movie. Instead of trying to memorize every single frame of the next scene from scratch, you just need to remember what changed. Did the ball move left? Did the cup tilt?
The Innovation: The authors created a compact "change code" called Residual Latent Action (RLA). They take the difference between the current image and the future image, compress it into a tiny, efficient vector (a short list of numbers), and call this the "action."
Why it helps: This is like giving the robot a "change log" instead of a whole new book. It's much smaller, faster to process, and avoids the "blurry" problem because it focuses only on the movement.

2. The "Flow" Machine

Once the robot has this "change code," it needs to predict what that code will look like in the future.

The Analogy: Think of a river. You don't need to simulate every single water molecule to know where the river flows. You just need to know the current and the direction of the flow.
The Innovation: They use a technique called Flow Matching. It's like a GPS that calculates the smoothest path from "now" to "future" based on the robot's actions. Because they are calculating this path in the tiny "change code" space (not the huge video space), it is incredibly fast and accurate.

3. Two Superpowers for Robots

The paper shows that this new model helps robots learn in two specific, powerful ways:

A. Learning from "Silent" Videos (No Action Labels Needed)

The Problem: Usually, to teach a robot, you need videos where you know exactly what buttons were pressed (action labels). But most videos on the internet (like YouTube) just show the robot moving without telling you how it moved.
The Solution: The RLA model can look at a silent video, figure out the "change code" (RLA) just by watching the movement, and then teach a robot to mimic that movement. It's like learning to dance just by watching a video of someone else dancing, without needing a choreographer to tell you the steps.

B. Training Inside a "Dream" (Visual Reinforcement Learning)

The Problem: Training robots in the real world is slow and risky. If a robot drops a cup, you have to pick it up and try again.
The Solution: The authors built a system where the robot practices entirely inside the simulation (the world model).
- The robot "dreams" of a task.
- It tries a move in its dream.
- The RLA-WM predicts the result instantly.
- The robot gets a "reward" if the dream outcome looks like a successful video it saw earlier.
- It repeats this millions of times in its head, learning the perfect strategy without ever touching a real object or needing a human to manually grade its performance.

The Results

The paper claims that this method is:

More Accurate: It predicts future states better than previous methods that tried to generate full videos or blurry features.
Much Faster: It is orders of magnitude faster than video-generation models because it works with the tiny "change codes" instead of massive video files.
More Reliable: It doesn't "hallucinate" weird objects; it sticks to the physics of the scene.

In short, this paper teaches robots to stop trying to "paint" the future and start simply calculating the "steps" needed to get there, allowing them to learn faster, smarter, and from a wider variety of videos.

Technical Summary: Learning Visual Feature-Based World Models via Residual Latent Action

Problem Statement

World models aim to predict future state transitions from observations and actions, serving as a foundation for policy learning and reasoning. The current dominant paradigm relies on video generation in pixel-aligned latent spaces (e.g., VAEs). While visually compelling, these approaches suffer from high computational overhead and a propensity for hallucination (generating physically implausible states).

Conversely, visual feature-based world models (e.g., DINO-WM) predict future visual features (such as DINO tokens) rather than raw pixels. This approach is more efficient and less prone to hallucination. However, existing feature-based methods rely on direct regression to predict future tokens. In complex 3D interactions, this leads to "regression-to-the-mean," resulting in blurry or collapsed predictions. A seemingly viable alternative—using generative models (like diffusion) directly in the feature space—is hindered by the curse of dimensionality; DINO tokens have significantly higher dimensionality (e.g., ~~1M dimensions for a 512×512 image) compared to standard VAE latents (~~16k), making generative modeling computationally prohibitive and defeating the efficiency goals of feature-based models.

The paper addresses two key questions:

How to develop an efficient yet accurate world model in a visual feature space that scales to complex 3D manipulation?
How to leverage such world models to improve downstream robot policies?

Methodology

1. Residual Latent Action (RLA)

The authors propose a new latent action representation called Residual Latent Action (RLA). Instead of predicting the absolute future state $s_{t+h}$ or generating it via heavy diffusion, RLA encodes the residual between the current DINO tokens ( $s_t$ ) and future DINO tokens ( $s_{t+h}$ ) into a compact latent vector $z$ .

Architecture: An autoencoder is trained where an encoder $f_{enc}$ processes the residual $(s_{t+h} - s_t)$ along with learnable queries to produce a compact latent $z$ . A decoder $f_{dec}$ reconstructs $s_{t+h}$ using $s_t$ and $z$ .
Key Properties:
- Predictive Sufficiency: Unlike prior latent actions used as weak conditioning for iterative generation, RLA allows for the accurate reconstruction of future tokens in a single feedforward pass.
- Generalizability: The RLA encoder generalizes to novel scenes and motion patterns even when trained on limited, task-agnostic data.
- Temporal Topology: The latent space naturally encodes temporal progression. Linear interpolation between Gaussian noise and an RLA vector yields representations corresponding to intermediate temporal states.

2. RLA World Model (RLA-WM)

Building on RLA, the authors propose RLA-WM, a world model that predicts the compact RLA vector $z$ rather than the high-dimensional $s_{t+h}$ directly.

Flow Matching: The model uses flow matching to predict the velocity required to transform Gaussian noise into the target RLA $z$ $z$ .
- Input: Current state $s_t$ and an action chunk $a_{t:t+h}$ .
- Process: A condition network processes $s_t$ and actions to generate condition tokens. These are concatenated with a noisy latent $z_\tau$ . A flow network predicts the velocity $\hat{v}$ to iteratively transform the noise into the final RLA $z_1$ .
- Output: The predicted RLA $z_1$ is decoded with $s_t$ to reconstruct the future state $s_{t+h}$ .
Efficiency: By performing flow matching in the compact RLA space (e.g., 2048 dimensions) rather than the full DINO token space, the model achieves high fidelity with significantly lower computational cost (FLOPs) compared to video diffusion models.

3. Robot Learning Applications

The paper demonstrates two applications of RLA and RLA-WM:

Minimalist World Action Model (WAM) for Actionless Videos:
- Extends a standard Behavior Cloning (BC) policy (ResNet-18) by adding a single linear layer to predict RLA $z$ from observations.
- Trained on a mix of videos with and without action labels. For actionless videos, the action loss is masked, and the network learns to predict the RLA target (derived from the video pair) to guide the policy.
- This allows learning from large-scale, unlabeled videos without coupling the policy to heavy video generation backbones.
Visual Reinforcement Learning (WMRL) inside RLA-WM:
- The first demonstration of visual RL trained entirely inside a world model learned from offline videos only.
- No online interactions, handcrafted rewards, or auxiliary BC loss are used during the RL phase.
- Video Aligned Reward (VAR): The reward is defined as the negative L1 distance between the DINO tokens of the model's rollout and the ground-truth DINO tokens of a reference offline video (time-synchronized).
- The policy (initialized from BC) is optimized via PPO using rollouts generated entirely within RLA-WM.

Key Contributions

Residual Latent Action (RLA): A simple, compact latent action representation learned from DINO residuals that is predictive, generalizable, and encodes temporal topology.
RLA-WM: A state-of-the-art visual feature-based world model that predicts RLA via flow matching. It outperforms existing feature-based and video diffusion models in accuracy while being orders of magnitude faster.
Novel Applications:
- A minimalist WAM that improves imitation learning from actionless videos.
- A visual RL framework (WMRL) that optimizes policies via rollouts in a learned world model without online interaction or manual reward engineering.

Experimental Results

Prediction Quality

Evaluated on ManiSkill (simulation) and IWS (real-world bimanual manipulation) datasets:

Metrics: RLA-WM achieved superior performance in LPIPS, SSIM, and DINO L1 distance compared to baselines including DINO-WM (regression), RAE (diffusion in feature space), FM-WM (flow matching in feature space), and Vid2World (video diffusion).
Efficiency: RLA-WM required 3.5T FLOPs per inference, significantly lower than Vid2World (1.1P FLOPs) and RAE/FM-WM (14.3T FLOPs).
Qualitative: RLA-WM produced physically grounded predictions with minimal hallucination, whereas video diffusion models (Vid2World) often generated sharp but physically inconsistent trajectories.

Robot Learning Performance

Actionless Video Learning: In imitation learning tasks using only 5% action-labeled data, the RLA-based WAM achieved a 35.6% average success rate, outperforming the BC baseline (27.2%) and other latent action methods (e.g., AdaWorld at 33.7%). The improvement was most pronounced on the challenging Push-T task (+15.2% over baseline).
Visual RL (WMRL): Across 1500 evaluation seeds, WMRL significantly improved policy performance on XArm and UR10e robots compared to the initial BC policy. On average, it achieved a +1.1% gain over BC across five tasks without using any additional real-world interaction data.

Significance and Claims

The paper claims that RLA-WM represents a strong step forward for visual feature-based world models. By shifting from direct regression or heavy generative modeling to residual latent action modeling, the authors achieve a balance of efficiency, accuracy, and scalability that was previously unattainable in high-dimensional feature spaces.

The significance lies in:

Overcoming the Dimensionality Curse: Demonstrating that generative modeling in feature space is feasible if the target is a compact residual representation rather than the full feature vector.
Enabling Offline-Only RL: Proving that high-quality visual RL is possible using only offline video data and a learned world model, eliminating the need for expensive online exploration or handcrafted reward functions.
Scalability to Unlabeled Data: Providing a pathway to leverage the vast amount of unlabeled robot video data for policy improvement through the minimalist WAM approach.

The authors acknowledge limitations, such as the difficulty of handling task-irrelevant background motion (e.g., camera shake) and partial observability (occlusions), suggesting future work in 3D view-independent RLA and multi-frame conditioning. However, they emphasize that their results on small-scale datasets isolate the method-driven gains from mere data scaling, establishing a solid foundation for future open-world applications.