PaperlayerView exact revision on arXiv

Source-linked AI summary

Spatially Aware World Action Model via Geometric Latent Diffusion

Javier Alejandro Lopetegui Gonzalez, Paul Pacaud, Cordelia Schmid

arXiv:2609.02531v1cs.CVcs.RO

TL;DR

Existing World Action Models mainly use RGB observations, leaving 3D information underused for robot policy learning. SA-WAM incorporates depth into a pretrained video diffusion model through tokenizer-compatible normalization and reports stronger benchmark and rollout-related results.

Abstract

from arXiv · show

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 tokenizer 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.

I. INTRODUCTION

World Action Models primarily use RGB observations, leaving 3D information underused despite its relevance to manipulation. SA-WAM adds explicit geometry to pretrained video-based world-action modeling through tokenizer-compatible depth representations.

  • World Action Models mainly operate on RGB observations and do not leverage 3D information.Source: §I · p. 1
  • 76.6% success rate on RoboCasa with 50 demonstrations per task improves over the matched Cosmos-Policy baseline by 9.5 points.The result also exceeds prior approaches trained with 6–20× more data.Source: Figure 1 · p. 1
  • 3D-aware policies have shown value for manipulation-intensive tasks where object geometry matters, but spatial awareness remains largely unexplored in WAMs.Source: §I · p. 1
  • SA-WAM augments a pretrained video diffusion model with depth encoded as additional latent frames processed by the same frozen VAE tokenizer as RGB.Source: §I · p. 1
  • Nonlinear depth normalization maps unbounded metric depth into the frozen tokenizer’s bounded feature space while preserving near-field resolution important for manipulation.Source: §I · p. 1
  • SA-WAM also reports strong LIBERO-Plus and real-world UR5 results, including robustness to randomized environments.Source: §I · p. 2

II. RELATED WORK

Related work spans video-based world-action modeling, 3D-aware robot policies, and geometric encodings adapted to pretrained visual representations. SA-WAM combines these directions by bringing explicit depth into a pretrained video-model framework.

World Models for Robot Policy Learning

  • Video-based robot policies use pretrained video models to synthesize future visual trajectories or predict actions, but existing approaches primarily use RGB observations.Source: §II · p. 2
  • Video-Action Models repurpose pretrained video diffusion transformers for robot-specific modalities such as proprioception and actions without modality-specific modules.Source: §II · p. 2

3D-Aware Robot Policies

  • 3D-aware robot policies incorporate point clouds, spatial encodings, or geometric representations to support manipulation reasoning and action prediction.Source: §II · p. 2

Metric Geometry Encoding for Pretrained Visual Models

  • Geometric information can be represented with dedicated depth latent spaces or adapted to existing visual-model domains through nonlinear parameterizations.Source: §II · pp. 2–3
  • Log-scale depth encoding distributes representational capacity across large depth ranges while emphasizing densely represented near-depth regions.Source: §II · p. 3

A. 3D Modality Injection for World Action Modeling

SA-WAM injects RGB-D observations into a shared latent sequence while reserving dedicated positions for proprioception and actions. Camera-specific depth processing maps metric depth into the frozen tokenizer’s input domain using selectable monotone normalizers.

B. Training Objective

SA-WAM trains a single diffusion sequence to jointly model actions, future RGB, and future depth. The shared objective supervises all modalities without dedicated prediction heads or modality-specific losses.

  • The clean target sequence combines an action chunk with tokenizer latents for future RGB and geometric frames, conditioned on task, proprioception, and current multi-view RGB-D observations.Source: B. Training Objective · p. 3
  • Gaussian noise corrupts the target sequence, and an EDM denoiser is optimized with noise-dependent weighting.Source: B. Training Objective · p. 3
  • The same objective jointly supervises action, future RGB, and future depth slots within one generative sequence.Source: B. Training Objective · pp. 3–4
  • The design avoids dedicated prediction heads and modality-specific losses.Source: B. Training Objective · p. 4

A. Experimental setup

The evaluation covers RoboCasa, LIBERO, and LIBERO-Plus simulation benchmarks, with RoboCasa measuring performance across 24 kitchen-manipulation tasks. LIBERO-Plus evaluates zero-shot robustness to seven controlled perturbation dimensions, while training uses a fine-tuned 2B-parameter DiT and frozen tokenizer.

  • Evaluation uses RoboCasa, LIBERO, and LIBERO-Plus simulation benchmarks.Source: A. Experimental setup · p. 4
  • RoboCasa contains 24 kitchen-manipulation tasks spanning pick-and-place, articulated-object interaction, appliance control, and faucet manipulation.Source: A. Experimental setup · p. 4
  • RoboCasa provides two third-person cameras, a wrist-mounted camera, and proprioception, with average task success reported over three random seeds.Source: A. Experimental setup · p. 4
  • LIBERO-Plus introduces controlled perturbations across seven dimensions and is evaluated zero-shot using only original LIBERO demonstrations, without adaptation.Source: A. Experimental setup · p. 4
  • Training fine-tunes the 2B-parameter Cosmos-Predict2 DiT while keeping the Wan2.1 VAE tokenizer frozen.Source: A. Experimental setup · p. 4

B. Ablations on the depth normalization approach

The ablations compare linear, inverse-depth, and log-scale depth normalization through tokenizer reconstruction and policy performance. Log-scale normalization provides the best balance, yielding the lowest overall reconstruction error and the strongest policy improvement.

Ablation design

Metric reconstruction

  • Log normalization achieves the lowest overall reconstruction error at 0.63% AbsRel while maintaining comparatively low error across wrist-camera depth ranges.Source: B. Ablations on the depth… · p. 4
  • Linear normalization reaches 6.68% AbsRel below 0.3 m, whereas inverse-depth normalization records 0.62% in the near field but 4.16% and 13.19% at medium and far distances.Source: B. Ablations on the depth… · p. 4

Policy ablation

C. Comparison to the State of the Art

SA-WAM outperforms strong baselines on RoboCasa and LIBERO-Plus while remaining competitive on standard LIBERO. Its gains are largest in spatially demanding RoboCasa categories, and its future predictions better align with simulator rollouts in a representative task.

RoboCasa

LIBERO

LIBERO-Plus

D. World Modeling Quality Analysis

The analysis evaluates future RGB prediction quality and relates stronger future-state modeling to better policy performance. It reports comparisons across fixed-camera and wrist views using pixel, structural, and perceptual metrics.

  • World-model RGB quality is evaluated with PSNR, SSIM, and LPIPS across fixed-camera and wrist-camera views on RoboCasa.Higher PSNR and SSIM indicate better predictions, while lower LPIPS is better.Source: D. World Modeling Quality Analysis · p. 5, Table IV · p. 6
  • The reported evaluation separates averaged fixed-camera views from the wrist camera and also denotes masked regions.F denotes fixed-camera views, W the wrist camera, and subscript m masked regions.Source: Table IV · p. 6
  • Stronger future-state modeling is associated with better policy performance in WAMs.Source: D. World Modeling Quality Analysis · p. 6

E. Prediction Quality versus Rollout Success

The paper studies whether future prediction consistency tracks rollout success and evaluates this relationship through object-centric 3D error and real-world manipulation results. Lower geometric prediction error accompanies successful rollouts, while SA-WAM remains strongest under randomized conditions.

Prediction quality versus rollout success

Real-world experiments

  • SA-WAM reaches 90.0% aggregate completion in clean real-world trials, compared with 75.0% for Cosmos-Policy and 21.3% for π0.The evaluation uses 10 trials per category and compares four evaluated task categories.Source: §V · p. 7
  • Under randomization, SA-WAM retains 77.5% aggregate completion, while Cosmos-Policy and π0 decrease to 48.8% and 11.3%.The trend is consistent with explicit depth helping when appearance alone cannot distinguish targets from distractors.Source: §V · pp. 7–8
  • Qualitative randomized-environment examples show SA-WAM completing manipulation while Cosmos-Policy fails on visually confusing distractors.Source: Figure 6 · p. 7, §V · p. 8

VI. CONCLUSION

The paper concludes that SA-WAM adds explicit 3D information to a pretrained latent video diffusion model for geometry-aware world-action modeling. It reports strong simulation and real-world performance, while identifying inconsistent futures and inference efficiency as remaining challenges.

  • SA-WAM injects explicit 3D information into the DiT latent-frame sequence and uses log-scale depth normalization with a frozen VAE tokenizer.Source: §VI · p. 8
  • SA-WAM achieves state-of-the-art results on RoboCasa and LIBERO-Plus and transfers strong performance to a real-world UR5 setting.Source: §VI · p. 8
  • The paper reports improved RGB prediction quality over Cosmos-Policy on RoboCasa, highlighting explicit spatial grounding for WAM-based robot policies.Source: §VI · p. 8
  • Inconsistent future predictions and actions remain, while geometrically consistent training, verification, and more efficient inference are proposed as future directions.Source: §VI · p. 8

Related papers

Loading 2609.02531v1…