SG-WAM: Self-Guided World Modeling in Geometry-Aware Policy Space
Ruiteng Zhao, Zhengshen Zhang, Yue Su, Wenshuo Wang, Jiahui Li, Zhiyuan Yang, Francis E. H. Tay, Marcelo H. Ang, Haiyue Zhu
cs.RO, cs.CV
2026-08-03
SG-WAM predicts futures inside a 0.9B policy's own representation space, with an EMA copy of itself as the target. It scores 98.5% on LIBERO and 73% on LIBERO-Plus, matching 7B.
Robot manipulation policies, the ones that drive an arm to pick, place, or fold, increasingly come in a family called World Action Models (WAMs). A WAM does not just emit actions; it also predicts what the scene will look like after the action. Predicting the future pushes the policy to learn how its actions change the environment, rather than only mimicking demonstration trajectories.
But existing WAMs split on where that future is predicted, and both sides have a flaw. One side predicts future images, videos, or depth maps: dense supervision, but the model is forced to reconstruct texture, lighting, and background that barely matter for whether the task succeeds, burning capacity on looking pretty. The other side predicts in latent space, skipping reconstruction, but the representation that defines the future target usually comes from an external encoder or a pretrained feature space, not from the same representation the action expert actually uses. So prediction target and action expert are misaligned: what the model labors to predict is not necessarily what the action needs. The paper asks how to build a future-prediction space that is both native to action generation and rich enough in the spatial geometry manipulation requires.
SG-WAM stitches three things into a single 0.9B VLM (Qwen3.5-0.8B) and trains them end to end.
First, self-guided world modeling. The model inserts several learnable dynamics tokens into the visual and language token sequence, and a Self-Guided World Predictor forecasts their future latent states from the current policy state and an intervening action chunk. The trick is the target: instead of a separate encoder encoding the future observation, the future observation is fed to an exponential moving average (EMA) copy of the same policy backbone, which produces a slowly updated target. That guarantees the prediction space and the action expert's space are the same representation family, with no mismatch. The EMA self-supervision is the same lineage as BYOL and DINO.
Second, geometry supervision. Being native to the action is not enough; manipulation needs fine-grained spatial information, where things are and how they are grasped. During training a frozen 3D foundation model, VGGT, acts as a teacher, pulling the main-view image-token representations toward geometric features by cosine similarity. This teacher only shapes the visual context; it never defines the prediction target, which always comes from the EMA policy itself.
Third, a flow-matching action expert generates a continuous action chunk from the full context, dynamics tokens included.
All three objectives (action, geometry, future prediction) are optimized jointly. At inference, the geometry teacher, the world predictor branch, and the EMA target pathway are all removed; only the online VLM, dynamics tokens, and action expert deploy. World modeling is a training-time aid that adds no inference cost.
On the LIBERO simulation benchmark, averaged over four task suites, SG-WAM reaches 98.5%, tied for best with Spatial Forcing (7B), but at 0.9B parameters and without large-scale embodied pretraining.
| Method | Params | Embodied PT | LIBERO avg |
| OpenVLA-OFT | 7B | Yes | 97.1% |
| π0 | 3.3B | Yes | 94.4% |
| Fast-WAM | 6B | No | 97.6% |
| VLA-JEPA | 2B | No | 96.1% |
| Spatial Forcing | 7B | Yes | 98.5% |
| SG-WAM | 0.9B | No | 98.5% |
Zero-shot transfer tells a clearer story. LIBERO-Plus perturbs camera viewpoint, robot embodiment, language, illumination, background, observation noise, and layout. SG-WAM scores 73.0% overall, the highest, and ranks first under camera, language, illumination, and layout shifts. Methods that looked strong on standard LIBERO (Spatial Forcing, Fast-WAM) drop noticeably out of distribution.
In the real world (UR5e arm, three tasks), SG-WAM beats the explicit VPP and the implicit VLA-JEPA both in and out of distribution. Ablations show removing world modeling costs 1.9 points (most on the long-horizon LIBERO-Long, 96.2% down to 92.2%), removing geometry supervision costs 0.9, and the two are complementary. Scaling dynamics tokens from 1 to 8 raises success (96.1% to 98.5%); 16 tokens instead falls back to 97.2%, so enough is enough and more is not better.
The signal is direct: a top manipulation policy does not require 7B parameters or large-scale embodied pretraining. Put future prediction in the same representation space the policy acts in, pin the spatial geometry down, and 0.9B matches the big models while generalizing better. For robotics practitioners that means smaller deployable models and a lower data bar. The EMA self-supervision is also a reminder that the target-policy mismatch of earlier latent WAMs can be sidestepped with off-the-shelf self-supervised tools.
Honestly, this is a gradual advance on LIBERO, a benchmark that is already near saturation, so 98.5% is no longer the bottleneck. The real signal is its edge over larger models out of distribution and in the real world, and the viability of the small-model-plus-good-structure route.
The authors acknowledge validation on a single embodiment (UR5e) so far, with scaling to larger backbones and cross-embodiment datasets left as future work. The real-world study covers only three tasks with small sample sizes (pick-and-place and towel folding 20 trials each, toolbox organization 10). Absolute success is not always pretty: in-distribution towel folding is 45%, and under novel objects every task falls to 25% to 40%, so genuine novel-object generalization is still hard.
One more caveat: the 0.9B-versus-7B comparison is not fully apples-to-apples, since most 7B baselines benefit from large-scale embodied pretraining that SG-WAM does without. That is exactly its selling point, but it means you cannot conclude "structure beats scale"; only that without pretraining, structure can compensate for scale. And LIBERO-Plus is still simulation; real-world out-of-distribution covers only background, light, and novel object, far from open-ended settings.