Learning Visual Feature-Based World Models via Residual Latent Action

Xinyu Zhang1 Zhengtong Xu2 Yutian Tao3 Yeping Wang3 Yu She2 Abdeslam Boularias1
1Rutgers University 2Purdue University 3University of Wisconsin-Madison

NeurIPS 2026

arXiv Paper Code

scroll

TL;DR We introduce the RLA World Model that can predict future DINO tokens accurately and efficiently, enabling a minimalist world action model and visual RL inside our world model.

Summary

  1. 1. RLA is surprisingly easy to learn, yet ideal for modeling state evolution.
    • 1.1 RLA is predictive. Decoder $f_\text{dec}$ accurately reconstructs $s_{t+h}$ from RLA $z$ and $s_t$ with just transformers without diffusion.
    • 1.2 RLA is generalizable, even when trained on limited data.
    • 1.3 RLA encodes temporal progression in latent space.
  2. 2. We build RLA World Model (RLA-WM), a simple and efficient state-of-the-art.
  3. 3. Applications
    • 3.1 RLA extends a BC policy to a minimalist world action model with just a linear layer, so it can learn directly from videos, and also improves a VLA on LIBERO and a real robot.
    • 3.2 We perform visual RL inside our RLA-WM, which is learned from offline videos, with a video-aligned reward and no online interaction. It can also learn from failed trajectories.
We use image $512\times 512$, which means the DINO token dimension is $1024\times 1024$, yet with a RLA dimension of 2048 or just 64, we reconstruct $s_{t+h}$ with high fidelity.
These robot-object interactions were never observed during RLA autoencoder training. Yet, RLA still reconstructs them well.
Interpolating between a Gaussian noise vector and RLA produces frames that approximate temporally intermediate states. This indicates that the RLA latent space inherently captures temporal progression.
We predict RLA to estimate state evolution, instead of absolute states.

Future Prediction

Given an input frame (from validation set) at $t=0$, each world model predicts future frames at $t=h, t=2h, \dots$ from the input action chunk $a_{t:t+h}, a_{t+h:t+2h}, \dots$. We use the official action space for each task in IWS. The action space of all ManiSkill tasks is simply the robot's joint angles.

RLA World Model

RLA-WM predicts future states by generating the residual latent action $z$ via flow matching, and decodes $\hat{s}_{t+h}$ with pre-trained RLA decoder $f_\text{dec}$.

Learning from Actionless Videos

We add a linear layer on top of a BC policy to predict RLA $z$, which learns from videos whose action or proprioceptive states are not available. As RLA enables accurate prediction of $s_{t+h}$, this simple framework becomes a minimalist world action model, which improves performance, without inference cost or coupling to a video backbone.

| Method | PushT | Roll Ball | Pull Cube | Pull Cube w. Tool | Poke Cube | Rank ↓ | Avg SR ↑ | | --- | --- | --- | --- | --- | --- | --- | --- | | BC-ResNet | 3.6 | 42.0 | 33.6 | 7.6 | 49.2 | 3.8 | 27.2 | | DINO CLS | 7.6 | 39.6 | 40.4 | 4.4 | 44.8 | 4.0 | 27.4 | | UniVLA | 6.0 | 37.6 | 42.8 | 7.2 | 50.0 | 3.8 | 28.7 | | AdaWorld | 9.2 | 38.4 | *48.4 | 10.8 | 61.6 | 2.2 | 33.7 | | **+RLA (Ours)** | *15.2 | *43.8 | 43.6 | *12.0 | *63.6 | *1.2 | *35.6 | Success rates (%) of imitation policies trained with actionless videos via different latent actions. RLA gives the highest average success rate and rank. All latent actions share the same auxiliary head, so the gain comes from the RLA target itself.

Scaling to VLAs and Real Robots

RLA also works with a moving wrist camera. We extend the RLA autoencoder to the concatenated DINO tokens of the front and wrist cameras, and let VLA-Adapter predict RLA chunks alongside robot actions. On a real SO-101 robot, we train BC policies for cup stacking from 120 teleoperated demonstrations.

| Method | Spatial | Object | Goal | Long | Avg ↑ | | --- | --- | --- | --- | --- | --- | | LAPA | 73.8 | 74.6 | 58.8 | 55.4 | 65.7 | | UniVLA | 96.5 | 96.8 | 95.6 | 92.0 | 95.2 | | VLA-Adapter | 97.0 | 97.4 | *96.6 | *93.4 | 96.1 | | **VLA-Adapter + RLA (Ours)** | *97.6 | *99.2 | *96.6 | *93.4 | *96.7 | Latent-action VLAs on LIBERO (success rate, %). RLA improves VLA-Adapter on the Spatial and Object suites and achieves the highest average.
| Real-robot cup stacking | Success ↑ | | --- | --- | | BC | 3/20 (15%) | | **BC + RLA (Ours)** | *8/20 (40%) | 120 teleoperated demonstrations on an SO-101 robot, evaluated over 20 rollouts.

Real-Robot Demos

Cup stacking on a real SO-101 robot with a BC policy extended by RLA.

Visual Reinforcement Learning inside RLA-WM

We show the first demonstration of visual RL inside a world model learned only from offline videos, which we refer to as World Model-based Reinforcement Learning (WMRL).

| Method | XArm Poke Cube | UR10e Roll Ball | UR10e PushT | Panda Pull Cube | Panda Pull Cube w. Tool | Avg ↑ | | --- | --- | --- | --- | --- | --- | --- | | BC-ResNet | 89.7 | 65.3 | 15.1 | 85.3 | 40.0 | 59.1 | | WMRL (Ours, avg) | 94.2† ± 1.1 | 65.6 ± 4.4 | 17.3† ± 1.6 | 83.7 ± 1.7 | 37.7 ± 4.1 | 59.7 | | **WMRL (Ours, best)** | *95.9 | *71.3 | *20.5 | *86.1 | *44.7 | *63.7 | We run WMRL from the best BC checkpoint with 15 seeds and evaluate over 750 episodes. 'best' is the best checkpoint over all seeds, and 'avg' is the mean ± std of each seed's best checkpoint. † WMRL (avg) outperforms BC under a one-sided paired t-test with p < 10⁻⁴.

Learning from Failures

Unlike BC, WMRL can learn from failed trajectories, since they still provide valid dynamics for the world model. Adding 8,000 failed PushT trajectories with noisy actions hurts BC, whereas WMRL improves further.

| PushT (UR10e) | Before | + Failure data | | --- | --- | --- | | BC-ResNet | 15.1 | 13.6 | | WMRL (best) | 20.5 | *22.5 | | WMRL (avg) | 17.3 ± 1.6 | *18.5 ± 2.3 |

Citation

@inproceedings{zhang2026learning,
  title={{Learning Visual Feature-Based World Models via Residual Latent Action}},
  author={Zhang, Xinyu and Xu, Zhengtong and Tao, Yutian and Wang, Yeping and She, Yu and Boularias, Abdeslam},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2026},
  eprint={2605.07079},
  archivePrefix={arXiv},
  primaryClass={cs.CV}
}