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.

  • Problem

    World Action Models mainly operate on RGB observations, while 3D-aware policies have shown the value of geometry for manipulation-intensive tasks.

  • Method

    SA-WAM jointly predicts actions, RGB, and depth by injecting normalized depth into a pretrained video model while reusing its frozen VAE tokenizer.

  • Results

    SA-WAM achieves the highest reported success on RoboCasa at 76.6% with 50 demonstrations per task, reaches 86.6% weighted-average success on LIBERO-Plus, and shows lower future 3D prediction error in successful rollouts.

  • Takeaways & Limitations

    Explicit geometry improves benchmark policy performance and links task-relevant 3D prediction quality with rollout success.

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.
  • 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.
  • 3D-aware policies have shown value for manipulation-intensive tasks where object geometry matters, but spatial awareness remains largely unexplored in WAMs.
  • SA-WAM augments a pretrained video diffusion model with depth encoded as additional latent frames processed by the same frozen VAE tokenizer as RGB.
  • Nonlinear depth normalization maps unbounded metric depth into the frozen tokenizer’s bounded feature space while preserving near-field resolution important for manipulation.
  • SA-WAM also reports strong LIBERO-Plus and real-world UR5 results, including robustness to randomized environments.

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.
  • World Models for Robot Policy Learning: Video-Action Models repurpose pretrained video diffusion transformers for robot-specific modalities such as proprioception and actions without modality-specific modules.
  • 3D-Aware Robot Policies: 3D-aware robot policies incorporate point clouds, spatial encodings, or geometric representations to support manipulation reasoning and action prediction.
  • 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.
  • Metric Geometry Encoding for Pretrained Visual Models: Log-scale depth encoding distributes representational capacity across large depth ranges while emphasizing densely represented near-depth regions.

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.

  • SA-WAM jointly denoises an action chunk with future RGB-D observation latents conditioned on task instructions, proprioception, and current multi-view RGB-D observations.
  • The model represents wrist, left, and right camera observations with RGB frames and corresponding metric depth maps encoded by a frozen VAE.
  • RGB and depth use separate latent-frame positions while sharing the same pretrained tokenizer, whereas proprioception and actions occupy dedicated sequence positions.
  • The compared normalizers are linear, inverse-depth, and log-scale mappings with different allocations of representational resolution across distance.
  • Camera-specific depth bounds are estimated from training data because wrist and external cameras cover different depth ranges.
  • Clamped metric depth is transformed by a monotone normalizer, rescaled to the [−1,1] tokenizer range, and replicated across three channels.

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.
  • Gaussian noise corrupts the target sequence, and an EDM denoiser is optimized with noise-dependent weighting.
  • The same objective jointly supervises action, future RGB, and future depth slots within one generative sequence.
  • The design avoids dedicated prediction heads and modality-specific losses.

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.
  • RoboCasa contains 24 kitchen-manipulation tasks spanning pick-and-place, articulated-object interaction, appliance control, and faucet manipulation.
  • RoboCasa provides two third-person cameras, a wrist-mounted camera, and proprioception, with average task success reported over three random seeds.
  • LIBERO-Plus introduces controlled perturbations across seven dimensions and is evaluated zero-shot using only original LIBERO demonstrations, without adaptation.
  • Training fine-tunes the 2B-parameter Cosmos-Predict2 DiT while keeping the Wan2.1 VAE tokenizer frozen.

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: Ablations evaluate depth encodings through frozen-tokenizer metric reconstruction and downstream policy performance.
  • Metric reconstruction: Log normalization achieves the lowest overall reconstruction error at 0.63% AbsRel while maintaining comparatively low error across wrist-camera depth ranges.
  • Metric reconstruction: 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.
  • Policy ablation: Log-scale normalization improves policy performance over linear normalization by 5.2 points and over inverse normalization by 2.2 points.
  • Policy ablation: The study retains log-scale normalization for subsequent SA-WAM experiments because it achieves the highest policy success with balanced metric precision.

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: SA-WAM achieves 76.6% overall success on RoboCasa using only 50 demonstrations per task.
  • RoboCasa: SA-WAM improves over Cosmos-Policy by 9.5 points overall on RoboCasa, with gains of 16.7 points on PnP and 11.2 points on Turn/Toggle tasks.
  • RoboCasa: A representative pick-and-place example shows SA-WAM completing the manipulation while Cosmos-Policy fails, with predictions better aligned to the simulator rollout.
  • LIBERO: SA-WAM reaches 98.4% average success on standard LIBERO, remaining on par with the strongest policies on that benchmark.
  • LIBERO-Plus: SA-WAM achieves 86.6% weighted-average success on zero-shot LIBERO-Plus, exceeding π0.5 by 2.0 points and Cosmos-Policy by 5.2 points.

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.
  • 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.
  • Stronger future-state modeling is associated with better policy performance in WAMs.

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: Object-centric 3D prediction error is lower for successful rollouts than failed rollouts after the grasp attempt.The error compares predicted and realized future object geometry from the wrist camera at a fixed short horizon.
  • Prediction quality versus rollout success: A threshold on object-centric 3D error detects 80% of failures at a 13% false-alarm rate, with AUC ≈0.88.
  • 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.
  • Real-world experiments: 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.
  • Real-world experiments: Qualitative randomized-environment examples show SA-WAM completing manipulation while Cosmos-Policy fails on visually confusing distractors.

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.
  • SA-WAM achieves state-of-the-art results on RoboCasa and LIBERO-Plus and transfers strong performance to a real-world UR5 setting.
  • The paper reports improved RGB prediction quality over Cosmos-Policy on RoboCasa, highlighting explicit spatial grounding for WAM-based robot policies.
  • Inconsistent future predictions and actions remain, while geometrically consistent training, verification, and more efficient inference are proposed as future directions.
Loading 2609.02531v1…