Acting Through Registers
Organizations: AutoLab, School of Artificial Intelligence, Shanghai Jiao Tong University · Institute of Automation, Chinese Academy of Sciences · Humanoid Robot (Shanghai) Co., Ltd.
Abstract
World Action Models (WAMs) use generative visual representations to condition robot actions, yet how these representations organize action-relevant information remains poorly understood. We identify persistent high-norm spatial outliers in WAM representations and find that their action relevance persists even after normalization and K/V projection attenuate their extreme magnitude. Motivated by this observation, we introduce RegisterWAM, which adds learnable non-spatial registers to the world model and uses them directly for action conditioning. Registers alone do not fully eliminate spatial outliers; pairing them with observation flow-matching (FM) loss supervision suppresses spatial outliers while concentrating high-norm responses in registers. Beyond this reorganization, observation FM loss supervision also improves action learning without registers, motivating an additional noised-observation FM objective that learns current-scene recovery from corrupted inputs. The resulting model achieves 81.64% success on LIBERO-Plus without additional embodied pretraining data or foundation models, reaches 93.0% average success on RoboTwin 2.0, and improves real-world manipulation performance.
Figures & tables
| Method | Emb. PT. | Extra Prior | Camera | Robot | Lang. | Light | Backg. | Noise | Layout | Avg. |
|---|---|---|---|---|---|---|---|---|---|---|
| [ 4 ] | – | 75.4 | 77.5 | 85.6 | 96.9 | 94.6 | 89.7 | 85.7 | 85.7 | |
| FastWAM [ 38 ] | – | – | 16.4 | 44.5 | 68.9 | 78.2 | 53.7 | 37.7 | 60.7 | 51.5 |
| Faster-WAM [ 22 ] | – | – | 67.9 | 49.0 | 92.1 | 94.3 | 57.0 | 82.3 | 82.7 | 75.0 |
| Faster-WAM [ 44 ] | – | – | 53.8 | 71.6 | 94.7 | 96.3 | 61.3 | 63.6 | 79.1 | 73.6 |
| 4D-WAM [ 34 ] | – | 45.2 | 64.3 | 90.6 | 94.3 | 57.7 | 69.1 | 79.2 | 71.0 | |
| ST-WAM [ 30 ] | – | 55.4 | 60.1 | 79.3 | 93.0 | 74.2 | 79.5 | 74.3 | 72.8 |
| In-distribution | Out-of-distribution | ||||||
|---|---|---|---|---|---|---|---|
| Method | Batteries | Stacking | Handover | Height | Tablecloth | Distractor | Avg. |
| FastWAM | 0.5000 | 0.0000 | 0.6200 | 0.5500 | 0.0000 | 0.4200 | 0.3483 |
| RegisterWAM | 0.6875 | 0.2500 | 0.7550 | 0.7300 | 0.0000 | 0.7200 | 0.5238 |
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Action-Conditioned World Models
Action-conditioned world models predict how environments change under actions to support planning, control, and simulation.
See all 346 papers