Search papers, labs, and topics across Lattice.
This paper introduces the Spatially Aware World Action Model (SA-WAM), which enhances traditional World Action Models by integrating RGB and depth information for improved action and future state predictions. By utilizing a nonlinear encoding to adapt a pretrained video model without requiring 3D-specific fine-tuning, SA-WAM effectively incorporates geometric data while leveraging existing visual and physical priors. The model achieves state-of-the-art performance on RoboCasa and LIBERO-Plus benchmarks and demonstrates significant improvements in real-world robotic applications, particularly with a UR5 robotic arm in randomized environments.
SA-WAM not only integrates 3D awareness into action models but also sets new performance benchmarks in robotic policy learning.
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.
