SA-WAM

Spatially Aware World Action Model
via Geometric Latent Diffusion

Javier Alejandro Lopetegui Gonzalez Paul Pacaud Cordelia Schmid

Inria, École normale supérieure, CNRS, PSL Research University

Paper Code BibTeX

Abstract

World Action Models (WAMs) leverage the capabilities of large-scale pretrained video diffusion models to jointly predict future observations and actions, inheriting rich visual and physical priors from internet-scale video. This has made them a promising paradigm for robot policy learning, yet the prevailing models operate exclusively on RGB observations and do not leverage 3D information. To bridge this gap, we introduce a Spatially Aware World Action Model (SA-WAM), which repurposes a pretrained video model for joint action, RGB and depth prediction, enabling 3D-aware world modeling and action prediction within a single diffusion backbone. We use a nonlinear encoding that maps the unbounded depth signal into the bounded input domain expected by the frozen VAE tokenizer. This allows us to reuse the encoder without 3D-specific fine-tuning, incorporating geometric information without sacrificing the pretrained priors. SA-WAM achieves state-of-the-art results on the RoboCasa and LIBERO-Plus benchmarks, while simultaneously improving future-state predictions. Furthermore, SA-WAM outperforms strong baselines in real world evaluation using a UR5 robotic arm, with strong gains in randomized environments. We analyze the correlation between world model prediction quality and rollout success, providing insights into WAM performance and avenues for its improvement.

Method

SA-WAM architecture
SA-WAM architecture. SA-WAM injects geometry into the WAM latent sequence as tokenizer-compatible frames paired with each RGB view, while proprioception and the action chunk are inserted directly into latent sequence positions. The conditioning signals include the task description, one RGB frame per view interleaved with its corresponding 3D modality, and the proprioception at time t. The action token as well as the future-state RGB and 3D frames go through the denoising process. The decoding phase is optional at inference time.

The frozen video tokenizer expects a bounded input, while metric depth is unbounded — so the choice of normalization decides how much geometric contrast survives encoding. We compare linear, inverse and log-scale normalization and retain the log-scale mapping, which allocates greater resolution to the near field where manipulation precision matters while keeping usable precision at range. It reaches the lowest overall reconstruction error (0.63 % AbsRel) and the highest policy success rate, improving over linear normalization by 5.2 points.

Results

RoboCasa

Category-wise success rate (%) on the 24 RoboCasa tasks. SA-WAM obtains the highest overall success rate while training on 50 demonstrations per task — 6–20× fewer than most baselines.

MethodTrain demos / taskPnPOpen/Close Turn/ToggleCoffeeAvg.
π0300––––62.5
VideoPolicy300––––66.0
UWM100035.682.0––60.8
FLARE1000––––71.3
GR00T-N1.530053.885.761.754.064.1
GR00T-N1.5 + HAMLET30048.587.769.764.066.4
GR00T-N1.6-ft + World2Act300––––72.6
Cosmos-Policy5051.891.166.361.667.1
SA-WAM (ours)5068.592.6 77.564.776.6

LIBERO-Plus

Weighted-average success rate (%) on LIBERO-Plus under zero-shot transfer — trained only on the original LIBERO demonstrations, with no adaptation to the seven perturbation axes.

MethodAvg. SR (%)
π053.6
π0-FAST61.6
OpenVLA-OFT69.6
RIPT-VLA68.4
π0.584.6
HoloBrain-0-GD74.0
ABot-M080.5
Fast-WAM50.0
Cosmos-Policy81.4
SA-WAM (ours)86.6

SA-WAM also reaches 98.4 % average success on standard LIBERO, on par with the strongest policies on this largely saturated benchmark.

Real-world UR5

Task-completion scores on the real UR5 arm, clean environments
(a) Clean environments
Task-completion scores on the real UR5 arm, randomized environments
(b) Randomized environments

Real-world UR5 results. Task-completion scores over 10 trials per category for π0, Cosmos-Policy and SA-WAM. SA-WAM reaches an aggregate completion score of 90.0 % in the clean setting against 75.0 % for Cosmos-Policy, and retains 77.5 % under randomization where Cosmos-Policy drops to 48.8 % — consistent with depth providing a complementary spatial cue in the presence of visually confusing distractors.

Qualitative results

RoboCasa

RoboCasa pick-and-place rollouts and predictions for Cosmos-Policy and SA-WAM
Pick the condiment bottle from the counter and place it in the sink. The upper block shows Cosmos-Policy and the lower one SA-WAM, with predictions compared against their corresponding simulator rollouts; for SA-WAM we also show depth predictions. Red boxes mark rollout inconsistencies, the green box task completion.

Real-world UR5, randomized environments

Two manipulation examples under domain randomization where SA-WAM reaches task completion while Cosmos-Policy fails to disambiguate colour distractors. Rollout errors are highlighted with red squares, task completion with green squares.

UR5 randomized example: take the pink mug and put it on the middle part of the hanger
(a) take the pink mug and put it on the middle part of the hanger
UR5 randomized example: put the grapes in the yellow plate, then put the banana in the pink plate
(b) put the grapes in the yellow plate, then put the banana in the pink plate

BibTeX

@misc{lopetegui2026sawam,
      title={Spatially Aware World Action Model via Geometric Latent Diffusion},
      author={Javier Alejandro Lopetegui Gonzalez and Paul Pacaud and Cordelia Schmid},
      year={2026},
      eprint={2609.02531},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2609.02531},
}

Acknowledgement

This work was granted access to HPC resources of IDRIS under the allocation AD011017145 made by GENCI. It was funded in part by the French government under management of Agence Nationale de la Recherche as part of the “France 2030” program, reference ANR-23-IACL-0008 (PR[AI]RIE-PSAI projet) and the ANR project VideoPredict (ANR-21-FAI1-0002-01). Cordelia Schmid would like to acknowledge the support by the Körber European Science Prize. The authors thank Peteris Kulits, Federica Spinola and Zeeshan Khan for their valuable contributions to this project.