Learning Visual Feature-Based World Models via Residual Latent Action

Xinyu Zhang; Zhengtong Xu; Yutian Tao; Yeping Wang; Yu She; Abdeslam Boularias

Source-linked AI summary

Learning Visual Feature-Based World Models via Residual Latent Action

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

arXiv:2605.07079v1May 8, 2026cs.CVcs.AIcs.LGcs.RO

TL;DR

Visual feature-based world models seek efficient, less hallucination-prone alternatives to pixel generation, but direct regression struggles with complex interactions. This paper introduces RLA-WM, which predicts compact Residual Latent Actions via flow matching and outperforms feature-based and video-diffusion world models while being orders of magnitude faster than video diffusion.

Problem
Visual feature-based world models need compact transition representations because direct regression struggles with complex interactions and generative modeling in high-dimensional feature spaces remains challenging.
Method
RLA-WM learns Residual Latent Actions from DINO token residuals and predicts them via flow matching before decoding future visual features.
Results
RLA-WM outperforms state-of-the-art feature-based and video-diffusion world models on simulation and real-world datasets while remaining orders of magnitude faster than video diffusion.
Takeaways & Limitations
RLA and RLA-WM support a minimalist world action model for actionless videos and visual reinforcement learning entirely inside an offline-video-trained world model.
Takeaways & Limitations
Task-irrelevant background motion or workspace randomness can waste RLA representation capacity and degrade the latent space.

Abstract

from arXiv · show

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

1 Introduction

The paper addresses the limitations of pixel-based and direct-regression world models by introducing Residual Latent Actions (RLA), a compact representation of DINO residuals. RLA-WM predicts these latent actions through flow matching and supports two downstream robot-learning techniques.

Motivation: Video-generation world models can hallucinate and impose heavy computational overhead, motivating more efficient alternatives.They predict future trajectories in pixel-aligned VAE latent spaces.
Motivation: Direct regression of future DINO tokens is efficient for 2D manipulation but struggles to extend feature-based world models to complex interactions.The paper frames visual features as high-dimensional while valid physical transitions may lie on a lower-dimensional manifold.
RLA: Residual Latent Action encodes the residual between DINO tokens of two frames into a compact latent vector.This representation is introduced to capture the lower-dimensional structure of valid physical transitions.
RLA-WM: RLA-WM predicts RLA via flow matching from the current DINO tokens and actions, then predicts future DINO tokens from the current tokens and RLA.Operating flow matching in the compact RLA space makes the model more efficient than feature-space generative pipelines.
Robot learning: The paper applies RLA and RLA-WM to a minimalist world action model and a visual reinforcement learning framework.The world action model learns from actionless videos, while the visual reinforcement learning framework optimizes policy inside the learned world model.

2 Related Work

Related work spans offline world models for robotics, visual feature-space prediction, and latent actions learned from videos. These approaches motivate RLA by highlighting the computational cost of video diffusion, task limitations of 3D models, and blurred or collapsed feature predictions from direct regression.

World Models for Robotics: Offline robotic world models primarily predict future videos or 3D geometry, but video diffusion is computationally expensive and 3D structural assumptions can restrict task applicability.These approaches learn world models from offline datasets for future-state prediction in robotics.
World Models in Visual Feature Space: Visual feature-space models predict future embeddings instead of pixels, with DINO-WM directly regressing future DINO tokens to reduce reliance on heavy generative models.V-JEPA predicts future features for self-supervised learning, while DINO-WM applies feature prediction to 2D robot manipulation.
World Models in Visual Feature Space: Direct regression in visual feature space can produce blurred or collapsed estimates, whereas the proposed approach enables efficient and accurate multimodal DINO-token prediction.The stated mechanism is avoiding regression-to-the-mean.
Latent Actions: Video-based latent-action methods commonly use compact latents either as proxy controls for imitation from actionless videos or as weak labels for video diffusion.The paper instead learns Residual Latent Action from DINO residuals rather than raw pixels.

3 Method

RLA-WM learns compact residual latent actions from DINO token differences, then predicts them with lightweight flow matching and decodes future features from the current state. RLA is predictive, generalizes to unseen scenes, and encodes temporal progression despite training only on frame pairs.

Problem Formulation: The method learns a dynamics function that maps DINO patch tokens and action chunks to future tokens using only offline videos, without online rollouts or handcrafted rewards.DINO tokens represent image patches, and the model directly predicts multi-step future features in feature space.
Learning Latent Actions on DINO Residuals: RLA is learned from DINO token residuals st+h − st, which represent transitions and correspond to flow-matching velocities between current and future features.The RLA autoencoder uses learnable queries, self-attention, and a single regression loss.
Predictive Sufficiency: RLA decoding reconstructs future DINO tokens with high fidelity in a single feedforward pass, without requiring iterative generation.This makes RLA a predictive representation rather than merely weak conditioning for diffusion.
Generalizability and Temporal Topology: RLA generalizes from task-agnostic training videos to unseen task-relevant interactions and naturally encodes temporal progression through intermediate latent states.Interpolating between Gaussian noise and RLA produces frames corresponding to temporally intermediate states, although training uses only frame pairs.
RLA World Model: RLA-WM predicts compact RLA z instead of high-dimensional future features, then decodes z with the current state st to reconstruct st+h.Flow matching starts from Gaussian noise and integrates the predicted velocity from τ = 0 to 1; iterative generation remains within compact RLA space.

4 Experiments

The experiments evaluate RLA-WM on simulated and real-world robot manipulation videos for multi-step visual-feature prediction, and test RLA-based techniques for improving robot policies. RLA-WM outperforms feature-based and video-diffusion baselines, while RLA improves imitation learning and enables offline reinforcement learning inside the learned world model.

Experimental Setup: Experiments use ManiSkill simulation and the IWS real-world dataset across multiple robot arms, manipulation tasks, and a task-agnostic play environment.ManiSkill includes Panda, XArm, and UR10 arms and five built-in tasks, plus task-agnostic interactions with primitive shapes.
World Model Evaluation: Qualitatively, RLA-WM predicts sharp, physically faithful future frames, whereas DINO-WM blurs over longer horizons and Vid2World hallucinates physically inconsistent trajectories.Direct diffusion or flow matching in DINO token space also performs worse than RLA-WM.
World Model Evaluation: RLA-WM significantly outperforms DINO-WM, RAE, FM-WM, and Vid2World across all measured metrics on both ManiSkill and IWS.The comparison covers feature-based methods and video diffusion, with qualitative evaluation on validation episodes unseen during training.
Learning from Actionless Videos: RLA improves imitation policies trained on actionless videos by +8.5% over the BC baseline and +1.9% over AdaWorld, the second-best method.On Push-T, RLA improves by 15.2% over the baseline’s 3.6%, while AdaWorld achieves 9.2%.
World-Model Reinforcement Learning: WMRL trains policies entirely inside the learned world model without additional real-world interactions, and consistently improves the best-performing models over BC across all tasks.The evaluation uses 15 trials with seeds 1–15 for 2,400 steps each and evaluates performance over 50 episodes with standard seeds 42–91.

5 Limitations and Conclusion

The paper introduces Residual Latent Action (RLA), a compact representation of visual state dynamics, and RLA-WM, a state-of-the-art visual feature-based world model. It also presents two robot learning techniques enabled by this framework while acknowledging remaining limitations, including task-irrelevant background motion.

Conclusion: RLA is introduced as a compact representation of visual state dynamics.
Conclusion: RLA-WM is proposed as a state-of-the-art visual feature-based world model.
Conclusion: The framework enables a minimalist world action model and a visual RL framework in RLA-WM.
Limitations: Task-irrelevant background motion can cause visual changes between s_t and s

A Appendix · A.1 Code

The supplementary folder includes the source code and instructions for installing dependencies, setting up datasets, downloading pretrained models, running the demo notebook, and training.

A.1 Code: The supplementary Code folder provides the source code and README instructions for installation, dataset setup, pretrained-model downloading, demo execution, and training.The instructions are in Code/README.md and include a demo Jupyter notebook.

A.2 Implementation Details

The experiments consistently use DINOv3-Large features with 1024 channels and AdamW, while RLA is trained to reconstruct future features and encode temporal progression. The implementation also addresses simulation-to-world-model observation gaps through image decoding during policy training and evaluation.

All experiments use DINOv3-Large with channel size 1024 and the AdamW optimizer.
RLA World Model: The RLA World Model uses a 1024-channel condition and flow-matching architecture with 32 query tokens and a 2048-dimensional RLA representation.The flow-matching network operates on 64 tokens total, comprising 32 condition tokens and 32 noisy RLA tokens.
RLA Autoencoder: RLA reconstructs future DINO tokens accurately despite their 1024 × 1024 dimensionality, outperforming blurred predictions from AdaWorld and UniVLA.RLA enables accurate future-token decoding in a single feedforward pass, whereas AdaWorld and UniVLA use latent dimensions 2048 and 256, respectively.
RLA Autoencoder: Interpolating normalized RLA with Gaussian noise produces approximate intermediate-frame reconstructions, indicating that the latent space captures temporal progression.For example, (ϵ + z̄)/2 approximately reconstructs st+h/2, while direct interpolation of DINO tokens performs worse.
Visual RL within RLA-WM: Decoded RGB observations are used during policy pre-training and evaluation to reduce the neural-to-simulation gap caused by differences between rendered and ray-traced images.The authors note that this correction increases computational cost and may lower the performance ceiling.

A.3 Limitations and Future Directions

The paper identifies limitations involving background motion, partial observability, missing proprioception, dataset scale, and Panda-specific WMRL performance. Future directions include 3D and multi-frame representations, proprioceptive prediction, larger datasets, and addressing robot-specific data and viewpoint factors.

Background and random motion: Background motion and workspace randomness can consume RLA capacity and degrade its latent space, motivating 3D token projections and learning.These task-irrelevant visual changes arise from sources such as humanoid robots and eye-in-hand cameras.
Memory and partial observability: Single-frame RLA cannot represent history-dependent occlusion under partial observability, motivating conditioning on multiple frames.The world model predicts s_t+h from s_t and a_t:t+h, so occluded objects may need to be memorized in the latent space.
Proprioceptive world model: RLA-WM predicts visual evolution but not future proprioceptive states, motivating a world model that predicts both modalities.Proprioceptive input has been shown useful for policy learning, so joint prediction would broaden applicability.
Scaling to larger datasets: The evaluation deliberately uses small-scale ManiSkill and IWS datasets, leaving scaling to massive real-world datasets as a future step.This design isolates method-driven gains from improvements attributable merely to data volume while demonstrating RLA and RLA-WM’s core properties.
Panda WMRL Results: WMRL underperforms BC on Panda’s Pull Cube and Pull Cube with Tool tasks, unlike the other robots’ results.The paper attributes the consistent Panda drop to kinematic structure, camera viewpoint, and insufficient action diversity in demonstrations.
Panda WMRL Results: Panda’s 8-DoF action space, equal demonstration count, limited pulling diversity, and frequent front-top-view occlusion make learning more difficult.The XArm has 7 DoF and the UR10e has 5; higher-dimensional action spaces naturally require more data, while Panda pulling motions often cause occlusion.