💻 computer science

Spatially Aware World Action Model via Geometric Latent Diffusion

This paper introduces SA-WAM, a spatially aware world action model that repurposes pretrained video diffusion models to jointly predict actions, RGB, and depth via a novel nonlinear depth encoding, achieving state-of-the-art performance in both simulation and real-world robotic tasks.

Original authors: Javier Alejandro Lopetegui Gonzalez, Paul Pacaud, Cordelia Schmid

Published 2026-09-03

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Javier Alejandro Lopetegui Gonzalez, Paul Pacaud, Cordelia Schmid

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

To teach a robot to move through the real world, scientists have long relied on a strategy that mimics how humans learn: watching and doing. In recent years, a powerful new approach has emerged where robots are trained on vast libraries of video, learning to predict what happens next in a scene and how to act within it. These systems, known as world action models, are like students who have watched millions of hours of human activity; they can guess the future of a moving object or a pouring liquid based on patterns they have seen before. However, for all their visual sophistication, these models have a significant blind spot. They typically see the world only as a flat, two-dimensional picture, much like a photograph. They lack an innate sense of depth, struggling to understand the true distance between a robot's hand and a cup, or how far a door is from a wall. This limitation becomes a major hurdle when a robot tries to perform delicate tasks, such as picking up a fragile object or navigating a cluttered room, where knowing the exact three-dimensional shape of the environment is critical.

A team of researchers has now bridged this gap by creating a new system that gives these video-trained models a sense of space. They developed a method to feed the robot not just the standard color images it is used to, but also a precise map of distances, known as depth information, directly into its learning process. The challenge was that the existing models were built to understand flat images, and depth data behaves very differently, with distances stretching out infinitely. To solve this, the researchers devised a clever way to compress this vast range of distances into a format the model could digest without needing to be completely rebuilt. They used a mathematical trick that preserves the fine details of objects close to the robot—where precision matters most—while still keeping track of things further away. This allowed them to reuse a powerful, pre-trained video model, essentially giving it a new pair of eyes that can see in three dimensions without losing the knowledge it had already gained from watching millions of hours of video.

The result is a system called SA-WAM, which stands for Spatially Aware World Action Model. In tests, this new approach proved to be significantly more capable than previous methods. When evaluated on a complex simulation of a kitchen environment, where a robot arm had to perform tasks like opening drawers, turning on faucets, and moving objects between counters, the new model succeeded in 76.6% of attempts. This was a substantial improvement over the best existing systems, which managed success rates in the low 70s or lower, even when those systems were trained with far more data. The researchers found that by adding this geometric awareness, the robot became much better at predicting the future state of the world. It could foresee exactly where an object would be after a movement, leading to more accurate and reliable actions. The model didn't just guess; it understood the physical space, allowing it to handle tasks that required precise alignment, such as placing a bottle into a sink or stacking cups, with a level of confidence that flat-image models lacked.

The power of this spatial awareness was further demonstrated when the team tested the system on a real robot arm in a physical laboratory. They set up a series of manipulation tasks, such as putting items into boxes or hanging mugs on a rack, and then introduced a layer of difficulty by randomizing the environment. They moved objects to new positions, added distracting items that looked similar to the target, and changed the lighting. In these chaotic conditions, the old models often failed, confused by the visual clutter. The new spatially aware model, however, remained robust. It successfully completed 77.5% of the randomized tasks, while the previous best systems dropped to less than half that rate. The researchers observed that when the visual appearance of an object was ambiguous, the depth information acted as a reliable guide, helping the robot distinguish the target from the background noise. This suggests that for robots to operate safely and effectively in the unpredictable real world, they need more than just a good eye; they need a true understanding of the space they occupy.

The study also revealed a fascinating link between how well the robot predicts the future and how well it performs the task. The researchers analyzed the robot's internal predictions and found that when the model accurately forecasted the three-dimensional position of an object, the task was almost always successful. Conversely, when the prediction of the object's location was slightly off, the task tended to fail. This correlation suggests that the ability to visualize the future in three dimensions is not just a nice bonus, but a fundamental requirement for success. By measuring the error in these geometric predictions, the team could even diagnose why a robot failed, pinpointing the exact moment the spatial understanding broke down. This insight offers a clear path forward for improving these systems, suggesting that the key to better robot policies lies in refining their ability to model the physical world with geometric precision.

Ultimately, this work demonstrates that the next generation of robot intelligence will likely come from combining the pattern recognition of massive video datasets with the hard facts of physical geometry. By teaching these models to see depth, the researchers have given them a more complete picture of reality. The findings show that when a robot can truly understand the distance and shape of the world around it, it becomes far more capable of navigating the complexities of human environments. This is not a magic solution that solves every problem, as the researchers note that challenges remain in predicting inconsistent futures and improving the speed of these calculations. However, the progress is clear: by adding a simple, yet profound, layer of spatial understanding, robots are taking a significant step toward becoming reliable partners in our daily lives.

Technical Summary: Spatially Aware World Action Model via Geometric Latent Diffusion

Problem Statement
World Action Models (WAMs) have emerged as a promising paradigm for robot policy learning by leveraging large-scale pretrained video diffusion models to jointly predict future observations and actions. These models inherit rich visual and physical priors from internet-scale video data. However, prevailing WAMs operate exclusively on RGB observations, failing to leverage 3D geometric information. This limitation is significant because 3D-aware policies have demonstrated superior performance in manipulation-intensive tasks where object geometry and spatial relations are critical. The challenge lies in integrating explicit 3D information into WAMs without sacrificing the powerful priors of frozen, pretrained video models or requiring modality-specific 3D encoders that would necessitate extensive retraining.

Methodology: SA-WAM
The authors propose the Spatially Aware World Action Model (SA-WAM), which repurposes a pretrained video diffusion model to jointly predict actions, RGB observations, and depth maps within a single diffusion backbone.

Architecture and Latent Injection: SA-WAM follows a latent-frame injection strategy. It augments the latent sequence of a Diffusion Transformer (DiT) with depth signals. While robot proprioception and action chunks occupy dedicated positions in the latent sequence, RGB and depth observations are encoded into latent frames using the same frozen VAE tokenizer used for RGB. This design avoids introducing dedicated action heads or geometry-specific encoders.
Nonlinear Depth Normalization: A core technical challenge is mapping unbounded metric depth into the bounded input domain expected by the frozen VAE tokenizer. SA-WAM employs a nonlinear encoding strategy to transform metric depth into a tokenizer-compatible, three-channel image representation. The authors compare three normalization strategies:
1. Linear: Distributes resolution uniformly but loses precision in the near field.
2. Inverse-depth: Concentrates resolution on near-field geometry but degrades at medium and far distances.
3. Log-scale: Provides a balanced profile, compressing the far-field while preserving near-field resolution where manipulation precision is most critical.
  The log-scale normalization is selected as the optimal approach, as it aligns best with the frozen tokenizer's distribution and maintains geometric precision across the workspace.
Training Objective: The model is initialized from the Cosmos-Predict2 DiT backbone (trained for physical AI) with the Wan2.1 VAE tokenizer kept frozen. The training objective minimizes the denoising loss (EDM formulation) across a unified sequence containing the action chunk, future RGB frames, and future depth frames. This allows all modalities to be learned jointly without modality-specific loss functions.

Key Contributions

Integration of 3D Geometry into WAMs: SA-WAM introduces explicit 3D geometric information (metric depth) into World Action Models, enabling 3D-aware world modeling and action prediction within a single diffusion backbone. While prior works like 3D-VLA have connected 3D perception and action through generative world models, SA-WAM specifically adapts pretrained video diffusion transformers for this purpose using a frozen tokenizer.
Efficient Encoding Strategy: The paper introduces a nonlinear depth normalization technique that allows the reuse of a frozen video tokenizer for depth data, avoiding the need for dedicated 3D encoders or fine-tuning the tokenizer for 3D data.
Empirical Validation: The authors demonstrate that log-scale normalization yields the best policy performance and metric precision. They further analyze the correlation between world model prediction quality and rollout success, showing that lower geometric prediction errors in the near field strongly correlate with successful task execution.

Results

Simulation Benchmarks:
- RoboCasa: SA-WAM achieves a state-of-the-art success rate of 76.6% using only 50 demonstrations per task. This outperforms the matched Cosmos-Policy baseline by 9.5 points and surpasses methods trained with 6–20× more data. Improvements are particularly notable in pick-and-place (16.7% gain) and turn/toggle tasks (11.2% gain).
- LIBERO-Plus: Under zero-shot evaluation with controlled perturbations, SA-WAM achieves a weighted-average success rate of 86.6%, outperforming the strongest VLA baseline ( $\pi0.5$ ) by 2.0 points and the Cosmos-Policy baseline by 5.2 points.
World Modeling Quality: SA-WAM improves future-state prediction quality (measured by PSNR, SSIM, and LPIPS) over the Cosmos-Policy baseline, particularly in the dynamic wrist camera view and within masked regions corresponding to the robot arm and target objects.
Real-World Evaluation: Evaluated on a UR5 robotic arm, SA-WAM demonstrates strong robustness. In randomized environments with visual distractors, SA-WAM retains a 77.5% completion score, whereas Cosmos-Policy drops to 48.8% and $\pi0$ to 11.3%. Qualitative results show SA-WAM successfully disambiguating visually similar distractors where baselines fail.

Significance and Claims
The paper claims that SA-WAM establishes a new state-of-the-art for robot policy learning by bridging the gap between video generative models and 3D spatial awareness. The primary significance lies in demonstrating that explicit geometric conditioning significantly enhances policy robustness, particularly in randomized environments and tasks requiring precise object localization. The authors posit that the correlation between improved future-state prediction quality (specifically geometric accuracy) and policy rollout success suggests that stronger world modeling is a key driver for better policy performance. The work highlights that spatial grounding is essential for WAM-based policies to generalize effectively in complex, real-world manipulation scenarios.

Limitations
The authors acknowledge that while SA-WAM improves upon existing WAMs, the problem of predicting inconsistent futures and actions persists. They suggest that geometrically consistent training and verification, as well as improvements in inference efficiency, are necessary avenues for future work.

Technical Summary: Spatially Aware World Action Model via Geometric Latent Diffusion

More like this