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

arXiv:2608.013972026-08-01

Teaching a robot to predict its own future using the very same policy that generates its actions works better than using a separate predictor

SG-WAM is a robot policy that learns to predict how a scene will change while it generates the next action, instead of using a separate future-prediction module. It uses a slowly-updated copy of itself (an EMA target) as the prediction goal and a frozen 3D geometry model to inject spatial structure into its visual tokens. It reaches 98.5% average success on the LIBERO simulation benchmark and 73.0% on the harder LIBERO-Plus benchmark, and outperforms baselines in real-world robot tests.

METAL MEDIA explanatory visual

SG-WAM architecture: predicting the future from the policy's own representation

Evidence statusMeasured results reported

Observation + language + dynamics tokensCurrent multi-view camera images, the task instruction, and 8 learnable dynamics tokens are jointly processed by a shared VLM backbone (Qwen3.5-0.8B)
Geometric supervision (VGGT teacher)A frozen VGGT 3D model shapes the main-view visual tokens via a cosine-similarity loss, injecting spatial structure into the policy's representation
Self-Guided World Predictor (SGWP)Takes the current dynamics-token states plus the intervening robot action sequence and predicts how the dynamics tokens will evolve
EMA target pathwayA slowly-updated copy of the same policy processes the future observation to produce the prediction target, with gradients blocked
Flow-matching action expertGenerates the continuous robot action chunk conditioned on the full policy context including dynamics tokens; at inference the teacher, predictor, and EMA pathway are all removed

An explanatory diagram made by METAL MEDIA, not a figure supplied by the paper's authors.

What they did

Existing World Action Models (WAMs) either predict future images/video directly, which forces the model to waste capacity on irrelevant visual details like texture and lighting, or predict features in a separate latent space that may not match what the action-generation part of the model actually needs.
SG-WAM inserts 8 learnable 'dynamics tokens' into a vision-language model (VLM) and trains a Self-Guided World Predictor (SGWP) to forecast how these tokens will change under the robot's upcoming actions. The prediction target is generated by an EMA (exponential moving average) copy of the same policy looking at the future observation, so both the prediction and its target come from the same representation family.
In parallel, a frozen 3D foundation model called VGGT is used to shape the policy's main-view visual tokens with geometric structure via a cosine-similarity loss, so the dynamics tokens inherit spatial awareness of where things are in the scene.
All components (geometric shaping, future prediction, and a flow-matching action generator that produces the actual robot actions) are trained jointly in one stage; at inference time the geometry teacher, the SGWP branch, and the EMA target pathway are all removed, leaving only the lightweight deployed policy.
SG-WAM achieved 98.5% average success across four LIBERO suites, 73.0% on the more challenging zero-shot LIBERO-Plus benchmark (best among compared methods), and outperformed VPP and VLA-JEPA baselines on three real-world UR5e robot tasks (pick-and-place, towel folding, toolbox organization) under both familiar and novel conditions.

Figure 1: Conceptual comparison of existing WAMs and SG-WAM. Whereas explicit and auxiliary latent targets may introduce perceptual burden or target–policy mismatch, SG-WAM learns intervening-action-conditioned dynamics in geometry-structured, policy-derived representations.

Table 1: Simulation results on LIBERO. Success rates are reported for the four standard LIBERO suites. Bold indicates the best result, and underlining indicates the second-best result.
Method	Params	Embodied PT.	Spatial	Object	Goal	Long	Avg.
OpenVLA-OFT [21]	7B	✓	97.6	98.4	97.9	94.5	97.1
π0 [6]	3.3B	✓	98.0	96.8	94.4	88.4	94.4
π0-FAST [31]	3.3B	✓	96.4	96.8	88.6	60.2	85.5
π0.5 [5]	3.3B	✓	98.8	98.2	98.0	92.4	96.9
GR00T N1.6 [4]	3B	✓	97.7	98.5	97.5	94.4	97.0
Spatial Forcing [24]	7B	✓	99.4	99.6	98.8	96.0	98.5
WorldVLA [10]	7B	✗	87.6	96.2	83.4	60.0	81.8
LAPA [44]	7B	✓	55.4	58.8	74.6	73.8	65.7
RynnVLA-002 [9]	7B	✗	99.0	99.8	96.4	94.4	97.4
Mantis [42]	5.8B	✓	98.8	99.2	94.4	94.2	96.7
UniVLA [8]	7B	✓	96.5	96.8	95.6	92.0	95.2
Fast-WAM [45]	6B	✗	98.2	100.0	97.0	95.2	97.6
VLA-JEPA [34]	2B	✗	94.8	99.6	95.8	94.0	96.1
SG-WAM	0.9B	✗	99.4	99.8	98.6	96.2	98.5

Figure 2: Overview of the SG-WAM framework. The VLM jointly contextualizes multi-view observations, language, and learnable dynamics tokens for action generation. During training, a frozen VGGT teacher shapes main-view image tokens, while SGWP predicts future dynamics-token states conditioned on intervening actions and aligns them with an EMA policy target. All auxiliary branches are removed at inference.

Table 2: Zero-shot transfer results on LIBERO-Plus. Success rates are reported under different perturbation settings. Bold indicates the best result, and underlining indicates the second-best result.
Method	Params	Camera	Robot	Language	Light	Background	Noise	Layout	Overall
WorldVLA [10]	7B	0.1	27.9	41.6	43.7	17.1	10.9	38.0	25.0
Spatial Forcing [24]	7B	20.1	13.4	40.9	29.1	33.4	25.7	39.3	29.1
Mantis [42]	5.8B	15.7	41.8	45.9	45.1	28.9	39.2	62.5	39.8
UniVLA [8]	7B	4.3	50.3	71.8	59.1	80.0	25.3	34.3	41.5
Fast-WAM [45]	6B	16.4	44.5	68.9	78.2	53.7	37.7	60.7	50.0
π0 [6]	3.3B	13.8	6.0	58.8	85.0	81.4	79.0	68.9	53.6
VLA-JEPA [34]	2B	40.3	55.7	72.9	88.2	70.5	38.2	74.6	62.9
OpenVLA-OFT [21]	7B	56.4	31.9	79.5	88.7	93.3	75.8	74.2	69.6
SG-WAM	0.9B	58.6	48.9	81.4	89.8	86.1	80.7	74.2	73.0

Figure 3: Left: Visualization of the real-world platform. We use a UR5e robot arm as the manipulation platform, the Kinect Azure camera as the main camera and the RealSense D405 as the gripper camera. Right: Visualization of 3 real-world tasks. 1) Top: Pick and Place. 2) Middle: Towel Folding. 3) Bottom: Toolbox Organization.

Table 3: Success rates under different visual perturbations. Bold indicates the best result.
Model	Pick and Place	Towel Folding	Toolbox Organization
ID	Background	Light Change	Novel Object	ID	Background	Light Change	Novel Object	ID
VLA-JEPA	35%	20%	25%	20%	20%	10%	10%	15%	20%
VPP	30%	15%	10%	10%	35%	15%	15%	10%	30%
SG-WAM	75%	55%	60%	40%	45%	25%	35%	25%	50%

Figure 4: Representative middle-layer attention maps from dynamics tokens to main-view image tokens. The task is Pick up the black bowl between the plate and the ramekin and place it on the plate. With geometric supervision, the selected tokens attend more consistently to the robot–object interaction regions.

Table 4: Ablation study of geometric supervision and self-guided world modeling on LIBERO. Geo. denotes geometric supervision, and WM. denotes self-guided world modeling. Success rates (%) are reported. Bold indicates the best result.
Geo.	WM.	Spatial	Object	Goal	Long	Avg.
✗	✗	97.0	97.6	95.6	91.0	95.3
✓	✗	98.2	97.8	98.0	92.2	96.6
✗	✓	97.8	99.8	98.2	94.4	97.6
✓	✓	99.4	99.8	98.6	96.2	98.5

Figure 5: Overview of the Self-Guided World Predictor (SGWP). The projected dynamics-token states are first contextualized through self-attention and then attend to the encoded intervening actions through cross-attention to predict future latent states.

Table 5: Ablation study on the number of learnable dynamics tokens on LIBERO. Success rates (%) are reported. Bold indicates the best result.
Number of Token	Spatial	Object	Goal	Long	Avg.
1	97.2	98.8	98.2	90.2	96.1
4	99.0	98.4	97.4	95.2	97.5
8	99.4	99.8	98.6	96.2	98.5
16	99.4	99.4	97.6	92.4	97.2

Figure 6: Complete visualization of the Pick and Place task.

Table 6: Subtasks success rates of Towel Folding and Toolbox Organization tasks in ID settings. Bold indicates the best result.
Model	Towel Folding	Toolbox Organization
First Folding	Second Folding	Pick the Screwdriver	Pick the First Gear	Pick the Second Gear	Close the Toolbox
VLA-JEPA	50%	20%	50%	40%	20%	20%
VPP	60%	35%	70%	50%	40%	30%
SG-WAM	75%	45%	80%	60%	50%	50%

Figure 7: Complete visualization of the Towel Folding task.

Table 7: Ablation of the action information provided to the Self-Guided World Predictor on LIBERO. For the null-action variant, the ground-truth intervening action sequence is replaced with an all-zero sequence before being processed by the action encoder. The action-conditioning architecture, temporal positional embeddings, and all other model components and training settings remain unchanged. Success rates (%) are reported.
Variant	Spatial	Object	Goal	Long	Avg.
Null-Action Sequence	98.4	99.2	98.0	94.6	97.6
SG-WAM	99.4	99.8	98.6	96.2	98.5

Figure 8: Complete visualization of the Toolbox Organization task.

Findings

On the four standard LIBERO suites, SG-WAM reached 98.5% average success, matching the strongest baseline despite using a much smaller model without large-scale embodied pretraining.
On LIBERO-Plus, a zero-shot transfer test with shifts in camera viewpoint, lighting, background, language, and layout, SG-WAM achieved the highest overall success rate of 73.0%, notably best under camera, language, illumination, and layout shifts.
In real-world tests on pick-and-place, towel folding, and toolbox organization with a UR5e arm, SG-WAM outperformed the VPP (explicit video-prediction) and VLA-JEPA (implicit latent-prediction) baselines both in-distribution and under background shift, light change, and novel-object conditions.
Removing both geometric supervision and self-guided world modeling dropped average LIBERO success to 95.3%; combining both gave the best 98.5%. Removing world modeling alone cost 1.9 points, removing geometric supervision alone cost 0.9 points, with the largest gap on the long-horizon LIBERO-Long suite (96.2% to 92.2%).
Increasing the number of dynamics tokens from 1 to 8 raised average success from 96.1% to 98.5%, but increasing further to 16 tokens reduced it to 97.2%. Replacing the real intervening actions with an all-zero sequence lowered average success from 98.5% to 97.6%, showing the actual action information provides a measurable, though modest, benefit.

Figure 10: Visualization of the attention weight matrix of all latent dynamics tokens attending to main-view image tokens. w/ Geo. and w/o Geo. denote with geometric supervision and without geometric supervision, respectively.

Where it can be used

Designing robot learning pipelines that want both strong in-distribution manipulation performance and better robustness to unseen visual conditions, using a relatively small model.
Building robot policies that model how actions change the environment without generating expensive future images or video.
Reducing compounding failures in long, multi-step manipulation tasks such as sequentially organizing multiple objects.

Figure 11: Complete visualization of the Pick and Place task under the background shift.

Limits and open work

Real-world experiments are limited to a UR5e arm, a specific camera setup, and three tasks (pick-and-place, towel folding, toolbox organization); generalization to other robot embodiments or a much wider range of tasks is untested.
Results come from a comparatively small 0.9B-parameter model without large-scale embodied pretraining; the authors explicitly leave scaling to larger backbones and cross-embodiment datasets as future work.
The number of dynamics tokens shows a non-monotonic effect (performance drops at 16 tokens), so the optimal token count found here may not transfer directly to other tasks or model scales.
Real-world out-of-distribution evaluation covers only three specific shift types (background, lighting, novel object), so robustness to more extreme or different kinds of environmental changes remains unverified.

Figure 12: Complete visualization of the Pick and Place task under the light change.

Why it matters

This work suggests that robot policies can gain robustness and manipulation performance by predicting their own future internal states rather than reconstructing pixels or relying on a mismatched auxiliary latent space. It does this with a relatively small 0.9B-parameter model and no large-scale robot pretraining, which matters for teams with limited compute.

Terms in this paper

World Action Model (WAM) · A robot model that jointly generates actions and predicts how the environment will change as a result
dynamics tokens · Learnable vectors added to the model to encode how the scene changes as a result of robot actions
EMA (exponential moving average) target · A slowly-updated copy of the model used to produce stable prediction targets without back-propagating gradients into it
flow matching · A generation technique that transforms random noise into a target output (here, continuous robot actions) step by step
VGGT · A frozen 3D geometry foundation model used during training to inject spatial structure into the policy's visual features

Figures we cannot republish

Figure 9: Visualization of 3 OOD settings. 1) Left: Background Shift. 2) Middle: Light Change. 3) Right: Novel Object.

See the figures in the original paper →

Original abstract (English)

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

Authors · Ruiteng Zhao

Read on arXiv

Latest papers

All papers →

Latest from METAL MEDIA

NCT 127 drop BLINGY, their first full album in two years

Figures: Ruiteng Zhao et al., arXiv:2608.01397, arxiv-nonexclusive