SG-WAM: Self-Guided World Modeling in Geometry-Aware Policy Space

2026-08-02

AI Summary

Existing world models struggle to simultaneously capture action relevance and geometric awareness, leading to a disconnect between action generation and state prediction. This work proposes the SG-WAM framework, which constructs a geometry-aware, action-conditioned dynamics model within a policy-derived representation space. It leverages learnable dynamic tokens, a self-guided world predictor, and an EMA copy of the policy backbone to generate stable supervision signals, enabling end-to-end joint optimization with explicit geometric constraints. Notably, SG-WAM achieves strong performance without large-scale embodied pretraining, attaining average success rates of 98.5% on LIBERO and 73% on LIBERO-Plus—significantly outperforming existing strong baselines—and demonstrates robust generalization across both in-distribution and out-of-distribution tasks.

Abstract

World Action Models (WAMs) couple action generation with prediction of future states. Their effectiveness depends on whether future dynamics are modeled in a space that is both aligned with action generation and sufficiently geometry-aware to capture where and how actions change the scene. Existing WAMs typically satisfy only part of this requirement, relying on either perceptually heavy observation-space targets or auxiliary latent spaces that are not jointly structured for action relevance and geometry. We propose SG-WAM, a self-guided framework that learns geometry-aware action-conditioned dynamics directly in the policy-derived representation space. SG-WAM introduces learnable dynamics tokens and a Self-Guided World Predictor that forecasts their future latent states conditioned on intervening robot actions. Prediction targets are generated by an exponential moving average copy of the same policy backbone, providing stable supervision within the representation family used by the action expert. Geometric supervision further structures the policy image-token representations, providing spatially grounded context for the dynamics tokens and yielding a future-alignment space that is both action-relevant and geometry-aware. Latent future prediction, geometric grounding, and flow-matching action generation are jointly optimized end-to-end in a unified framework. Built on a 0.9B model without large-scale embodied pretraining, SG-WAM achieves 98.5% average success on LIBERO and 73% on LIBERO-Plus, while outperforming strong baselines in both in-distribution and out-of-distribution real-world evaluations.

Problem

• World Action Models
• geometry-aware
• action-conditioned dynamics
• policy representation space
• future state prediction

Innovation

• Self-Guided World Modeling
• Geometry-Aware Policy Space
• Action-Conditioned Dynamics
• Latent Future Prediction
• Geometric Grounding