SelfWAM: A Self-Grounded Unified World Action Model for Fast Robot Control

2026-08-01

AI Summary

Existing world models struggle to accurately capture the visual dynamics induced by actions when predicting future observations. This work proposes SelfWAM—a self-anchored, unified world-action model that explicitly links visual prediction with the robot’s visible body and its action effects by jointly forecasting actions, action-conditioned future RGB frames, and a robot self-mask. The model employs a modality-specific Mixture-of-Transformers architecture, integrating action-conditioned generation, prompt-aware self-mask supervision, and joint training to significantly enhance action sensitivity in future-frame prediction while maintaining efficient action inference. Experimental results demonstrate that SelfWAM effectively improves policy performance on both the RoboTwin 2.0 benchmark and real-world manipulation tasks.

Abstract

World Action Models (WAMs) improve robot policy learning by jointly modeling actions and future observations. However, conditioning future prediction only on the task prompt and observation context risks capturing generic task progression rather than the action-specific consequences of the executed action. We introduce SelfWAM, a unified self-grounded WAM built on a modality-specialized Mixture-of-Transformers (MoT) architecture that jointly predicts actions, action-conditioned future RGB frames, and robot self-masks, thereby grounding future prediction in the robot's visible body and its action-induced motion. During joint training, SelfWAM allows future visual queries to attend to a clean copy of the demonstrated action, turning the video branch into an action-specific consequence model while leaving the fast action-only inference path unchanged. To focus video learning on action-relevant visual changes, we use prompt-specific objectives for future robot self-mask prediction, which removes appearance details and provides a target whose temporal evolution is tightly coupled with the conditioning action. Together, clean-action conditioning and future self-mask supervision make future predictions more directly reflect how the executed action changes the robot's visible motion and the surrounding scene. Experiments on RoboTwin 2.0 and real-world manipulation tasks show that SelfWAM produces more action-sensitive futures and preserves fast policy inference, while improving policy performance.

Problem

• World Action Models
• action-conditioned prediction
• robot control
• future observation grounding
• action-specific consequences

Innovation

• Self-Grounded World Model
• Mixture-of-Transformers
• Action-Conditioned Future Prediction
• Robot Self-Mask
• Fast Policy Inference