PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.01397v1cs.ROcs.CV

TL;DR

Existing World Action Models lack a future-modeling space that is simultaneously aligned with action generation and sufficiently geometry-aware. SG-WAM learns such dynamics in policy representations through self-guided prediction and geometric supervision, achieving 98.5% average success on LIBERO and strong in-distribution and out-of-distribution real-world performance.

  • Problem

    Existing WAMs use future targets or auxiliary representations that are not explicitly aligned with action generation and may not preserve manipulation-relevant spatial information.

  • Method

    SG-WAM predicts action-conditioned future dynamics tokens in the policy representation space, using EMA self-supervision and geometric grounding with joint end-to-end action generation.

  • Results

    SG-WAM achieves 98.5% average success on LIBERO and outperforms strong baselines in both in-distribution and out-of-distribution real-world evaluations.

  • Takeaways & Limitations

    The results support policy-coupled future prediction as a geometry-aware approach for reliable manipulation under familiar conditions and visual distribution shifts.

  • Takeaways & Limitations

    Existing predictive targets are not guaranteed to coincide with the representations directly used to condition action generation.

Abstract

from arXiv · show

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.

1 Introduction

SG-WAM addresses limitations of existing WAMs by learning action-conditioned future dynamics in policy-derived representations that are jointly action-relevant and geometry-aware. It uses self-guided latent prediction, geometric supervision, and achieves strong LIBERO and real-world results with a 0.9B model.

  • Motivation: Existing WAMs couple action generation with future-state prediction, but their supervision may emphasize perceptual fidelity rather than manipulation-relevant dynamics.Latent approaches reduce appearance-reconstruction burden, yet externally defined or auxiliary targets may mismatch the representations used for action generation.
  • Motivation: Effective future modeling must preserve both direct coupling to action generation and fine-grained geometry describing where and how actions change the scene.Action-oriented compression can suppress appearance variation without retaining the spatial information required for manipulation.
  • Method: SG-WAM learns geometry-aware, action-conditioned dynamics directly in policy representations by forecasting learnable dynamics tokens conditioned on intervening robot actions.A Self-Guided World Predictor operates on dynamics tokens inserted into the VLM sequence.
  • Method: EMA targets from the same policy backbone provide stable supervision within the representation family used for action generation, while geometric supervision structures policy image tokens spatially.The geometry-grounded image tokens provide spatial context for future alignment.
  • Results: 98.5% average success on LIBERO and 73% on LIBERO-Plus are achieved by SG-WAM using a 0.9B model without large-scale embodied pretraining.The method also shows clear gains over strong baselines in in-distribution and out-of-distribution real-world evaluations.

2 Related Work

Related WAMs either model future observations explicitly or predict compact latent dynamics, but their future supervision is not necessarily constructed from the same policy representation used for action generation. SG-WAM addresses this distinction by constructing its predictive target from the policy used for action generation.

  • WAMs with Explicit Future Modeling: Explicit-future WAMs predict environment evolution under robot actions using future images or videos as intermediate targets for action generation.UniPi predicts task-conditioned videos and recovers actions through inverse dynamics; later methods target longer-horizon reasoning and controllability.
  • WAMs with Implicit Latent Modeling: Implicit-future WAMs model scene evolution in compact latent spaces, avoiding detailed observation reconstruction while integrating prediction with policy learning.Joint-embedding predictive models demonstrate that future information can be learned without pixel-level reconstruction.
  • Limitations of Existing Future Supervision: Joint prediction and control do not ensure that future targets come from the same policy used for action generation.Existing approaches use future-observation or target encoders, pretrained feature spaces, latent visual subgoals, posterior branches, compact condition spaces, or spatial-temporal latents jointly modeled with actions.

3 Method

SG-WAM performs action-conditioned world modeling directly in a geometry-aware policy representation space. It combines self-guided latent prediction with geometric supervision and flow-matching action generation in one jointly optimized framework.

  • Policy representation: SG-WAM jointly contextualizes visual, language, and learnable dynamics tokens within a shared Qwen3.5-0.8B VLM backbone.The resulting contextualized states support both future prediction and action generation.
  • Deployment: During inference, SG-WAM removes the frozen geometry teacher, SGWP prediction branch, and EMA target pathway, retaining the online VLM, dynamics tokens, and action expert.This leaves the deployed policy with only the components needed for action generation.
  • Geometry-aware policy states: A frozen VGGT teacher geometrically shapes main-view policy visual tokens, providing spatial context for the dynamics-token states without defining their future-prediction targets.Geometry supervision enters the policy-derived dynamics representation through jointly contextualized visual states.
  • Self-guided world predictor: SGWP predicts future dynamics-token states from the current policy representation and intervening actions, while an EMA copy of the same policy generates aligned future targets.The EMA pathway is detached from backpropagation and updated only through the EMA rule.
  • Action generation: The flow-matching action expert consumes the complete policy context, keeping dynamics-token states optimized by future prediction in the action-generation pathway.At inference, the action chunk is obtained by integrating the learned velocity field conditioned on the policy context.

4 Experiments

SG-WAM is evaluated on LIBERO, LIBERO-Plus, and real-world manipulation tasks against representative generalist, explicit-WAM, and implicit-WAM baselines. It achieves strong benchmark and real-world performance, while ablations show benefits from self-guided world modeling, geometric supervision, and an appropriate number of dynamics tokens.

  • Simulation Setup: SG-WAM is evaluated across LIBERO and LIBERO-Plus, which test standard manipulation capabilities, distribution-shift robustness, and generalization.LIBERO contains four suites, while LIBERO-Plus adds more challenging configurations and visual variations.
  • Simulation Results: 98.5% average success on LIBERO matched the strongest baseline using a 0.9B model without additional embodied pretraining.Most high-performing baselines rely on larger backbones or large-scale embodied pretraining.
  • Real-World Results: SG-WAM consistently outperformed VPP and VLA-JEPA under both in-distribution and out-of-distribution real-world conditions across three manipulation tasks.It retained its advantage under changes in background, light, and object, including the long-horizon Toolbox Organization task.
  • Ablations: 98.5% average success with both geometric supervision and self-guided world modeling exceeded 95.3% without either component.Removing world modeling reduced average success by 1.9 percentage points, versus 0.9 percentage points after removing geometric supervision.
  • Ablations: Increasing dynamics tokens from one to eight improved average success from 96.1% to 98.5%, while 16 tokens reduced it to 97.2%.On LIBERO-Long, success increased from 90.2% to 96.2% when token count rose from one to eight; the final model uses eight tokens.
  • Ablations: With geometric supervision, dynamics tokens attended more consistently to interaction-relevant regions such as the end effector, receptacle, and surrounding spatial context.The comparison uses middle-layer attention from dynamics tokens to main-view image tokens from the same frame.

5 Conclusion · Appendix · A Details of Learnable Dynamics Tokens

SG-WAM learns action-conditioned future dynamics in a policy-derived latent space using learnable dynamics tokens, EMA targets, and geometric supervision. The appendix specifies how these tokens are integrated into the multimodal backbone to support both dynamics prediction and action prediction.

  • 5 Conclusion: SG-WAM learns action-conditioned future dynamics in a latent prediction space derived from the policy itself.
  • 5 Conclusion: Learnable dynamics tokens predict future states from the current policy state and intervening robot actions, while an EMA policy copy supplies stable future targets.
  • 5 Conclusion: Geometric supervision structures policy visual tokens, enabling latent future prediction and action generation to be jointly optimized over spatially grounded representations.
  • 5 Conclusion: During inference, the geometry teacher, SGWP, and EMA target pathway are removed, while the learned dynamics-token representation pathway remains active.
  • A Details of Learnable Dynamics Tokens: The model introduces Q ∈ R^Nq×D as learnable continuous embeddings representing the observation’s action-relevant latent state, with Nq = 8.D is the vision-language backbone’s hidden dimension.
  • A Details of Learnable Dynamics Tokens: Q is appended to the image-language sequence with valid attention positions and rotary indices continuing from the original sequence.
  • A Details of Learnable Dynamics Tokens: Because Q participates in every transformer layer, it aggregates visual observations and language instructions while providing context for both auxiliary dynamics learning and action prediction.

B Details of Self-Guided World Predictor

The Self-Guided World Predictor forecasts future latent scene states by combining temporally ordered intervening actions with current dynamics-token representations. Stable future-state targets come from an EMA copy of the visual-language backbone, projector, and dynamics tokens, updated without gradient optimization.

  • Action-Conditioning: Eight seven-dimensional intervening actions are projected into temporally positioned action tokens, preserving the full sequence to match the action expert’s eight-step chunk.Each action is embedded with an MLP-based projector, and learned temporal position embeddings are added before prediction.
  • Future-State Prediction: Current scene-state slots first exchange information through self-attention, then cross-attend to action tokens to predict action-relevant future latent changes.Self-attention contextualizes the latent scene, while cross-attention lets each slot select the most relevant action steps.
  • EMA target pathway: The target future states are generated by EMA copies of the online backbone, projector, and learnable dynamics tokens.The EMA pathway computes target representations from the future visual observation, language instruction, and EMA dynamics token.
  • EMA target pathway: The EMA coefficient is µ = 0.999, retaining fraction µ of prior target parameters and incorporating fraction 1 − µ of their online counterparts at each optimization step.The online and EMA parameter sets are updated according to the stated moving-average rule.
  • EMA target pathway: The target pathway is excluded from gradient optimization, remains in evaluation mode, and runs without gradient tracking when computing target representations.Its parameters are updated exclusively through the EMA rule.

C Implementation Details of Geometric Supervision

Geometric supervision uses a frozen VGGT-1B teacher to construct spatial targets from main-view observations. These targets supervise the policy’s existing main-view image-token states through token-wise projection, without adding VGGT tokens to the VLM sequence.

  • Teacher and target construction: A frozen VGGT-1B model processes only the main-view observation to construct geometric training targets.The teacher is used exclusively for target construction.
  • Teacher and target construction: Main-view frames are resized to 518 × 518, processed as single-frame sequences, and converted from 1,369 patch tokens into an 8 × 8 geometric target grid.The 1,369 tokens are reshaped into a 37 × 37 grid before adaptive average pooling.
  • Policy-token supervision: Supervision acts directly on the policy’s existing main-view image-token states and does not add VGGT tokens to the VLM sequence.Because main-view tokens precede wrist-view tokens, the method selects the first main-view image-token states.
  • Policy-token supervision: Each policy image token is independently projected into the VGGT feature dimension using a lightweight token-wise projector Gγ.The projector maps policy image-token representations into the teacher’s feature space for geometric supervision.

D Training Objectives

SG-WAM jointly optimizes action generation, geometric alignment, and action-conditioned future-state prediction in a single-stage objective. Prediction uses normalized mean-squared error with gradient updates to online components while the EMA target pathway remains gradient-free.

  • Future-state prediction: Lpred uses normalized mean-squared error averaged across all learnable dynamics-token slots and feature dimensions.Normalization by Nqd averages prediction error over Nq dynamics tokens and d feature dimensions.
  • Future-state prediction: Gradients from Lpred update the online VLM, dynamics tokens, projectors, action encoder, and Self-Guided World Predictor, while the EMA target pathway remains gradient-free.The stop-gradient operation is applied to the target pathway.
  • Geometric alignment: Directional alignment with frozen VGGT features encourages policy image-token states to encode manipulation-relevant 3D cues without over-constraining their magnitude or distribution.This preserves the VLM’s original token structure and pretrained-representation compatibility while retaining spatial information for action-dependent future transitions.
  • Joint optimization: The online VLM representations are shared by geometry alignment, future-state prediction, and the conditional flow-matching action expert during end-to-end training.All online components are optimized jointly in a single-stage procedure, while the frozen VGGT teacher receives no gradient updates.
  • Joint optimization: The training objective combines action loss, geometry alignment, and future-state prediction as L = Lact + 0.1Lgeo + 0.1Lpred.The coefficients λgeo = 0.1 and λpred = 0.1 weight geometry alignment and action-conditioned future-state prediction.

E Simulation Experiments

SG-WAM is evaluated under standard LIBERO task-suite protocols and LIBERO-Plus zero-shot distribution shifts. Experiments use a 0.8B pretrained vision-language backbone, eight dynamics tokens, and end-to-end training without adaptation at evaluation.

  • Evaluation settings: LIBERO evaluates spatial reasoning, object-centric manipulation, goal-conditioned behavior, and long-horizon execution through LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, and LIBERO-Long.The four suites are evaluated under the standard LIBERO protocol.
  • Evaluation settings: LIBERO-Plus tests zero-shot robustness under shifts in camera viewpoint, robot embodiment, language, illumination, background, observation noise, and scene layout.These conditions assess whether policy representations remain stable when visual, semantic, and physical conditions differ from training.
  • Training setup: SG-WAM uses a pretrained Qwen3.5-0.8B vision-language backbone, eight dynamics tokens, and end-to-end visuomotor training.Self-guided world modeling and geometric supervision are applied only during training.
  • Training setup: 40k training steps use a global batch size of 96, cosine learning-rate scheduling, 5k-step linear warmup, and a peak learning rate of 1e-5.The model is trained on all four standard LIBERO suites.
  • Evaluation settings: The same checkpoint trained on original LIBERO simulation data is directly evaluated on LIBERO-Plus without additional fine-tuning or adaptation.Success rate is computed as the percentage of successful rollouts in each perturbation setting, with the final score based on overall successful rollouts.

F Real-World Experiments

Real-world experiments evaluate SG-WAM on three in-distribution tabletop tasks and three unseen test-time variations without additional adaptation. SG-WAM achieves the highest average success rate, with consistent gains across task types and intermediate execution stages, including challenging deformable and long-horizon subtasks.

  • In-Distribution Evaluation: The in-distribution evaluation uses three tabletop manipulation tasks with 100 expert demonstrations per task, covering obstacle-aware pick-and-place, deformable-object manipulation, and long-horizon multi-object organization.Demonstrations contain synchronized RGB observations, language instructions, and robot actions, while evaluation uses the same task families.
  • Out-of-Distribution Evaluation: The out-of-distribution evaluation introduces background shifts, lighting changes, and novel objects absent from demonstrations, using the same task definitions and success criteria without fine-tuning or test-time adaptation.All models are trained jointly on the three collected tasks for 40 epochs.
  • Overall Results: SG-WAM achieves the highest average success rate across both settings, with consistent in-distribution gains across all three tasks rather than improvements limited to one manipulation type.The results suggest benefits for spatially grounded and temporally coherent policy behavior in real-world manipulation.
  • Subtask Analysis: Across all Towel Folding and Toolbox Organization subtasks, SG-WAM outperforms VLA-JEPA and VPP, extending gains beyond final completion to intermediate execution stages.The fine-grained breakdown indicates improved performance throughout task execution.
  • Subtask Analysis: SG-WAM is more robust on sequential deformable-object manipulation and improves later stages of the long-horizon Toolbox Organization task, where errors accumulate over subtasks.The second towel fold is harder because it requires reasoning over the cloth state produced by the first interaction; toolbox performance declines from screwdriver picking through closing.

G Ablation Study

The ablations show that geometric supervision makes dynamics-token attention more interaction-relevant, while ground-truth intervening actions improve future latent-transition prediction. Removing action information reduces success modestly but leaves the model competitive.

  • Geometric Supervision: With geometric supervision, all eight dynamics tokens attend more consistently to interaction-relevant regions, including the end effector, manipulated object, target receptacle, and workspace.Without geometric supervision, several tokens attend more strongly to isolated visually salient regions less directly related to robot-object interaction.
  • Intervening Action Conditioning: 98.5% to 97.6%: replacing ground-truth intervening actions with an all-zero sequence lowers average LIBERO success.The null-action variant preserves the action-conditioning pathway, sequence length, temporal positional embeddings, model components, objectives, and training settings.
  • Intervening Action Conditioning: 1.0, 0.6, 0.6, and 1.6 percentage points: null actions reduce performance on LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, and LIBERO-Long, respectively.The reductions are reported for the four LIBERO suites in that order.
  • Intervening Action Conditioning: The null-action variant remains competitive, but ground-truth intervening actions consistently improve future latent-transition modeling, particularly for long-horizon tasks.The result indicates that specific action information contributes beyond temporal structure, while the current policy state already contains substantial predictive information.

H Real-World OOD Rollouts Visualization · I LIBERO Visualization · J LIBERO-Plus Visualization

The paper provides complete visualizations of real-world out-of-distribution rollouts under background, lighting, and object changes, alongside LIBERO and LIBERO-Plus task visualizations. These figures cover multiple benchmark suites and perturbations to viewpoints, initial states, instructions, lighting, backgrounds, sensor noise, and object layouts.

  • H Real-World OOD Rollouts Visualization: Complete Pick and Place rollouts are visualized under background shift, light change, and novel-object conditions.These real-world OOD conditions are shown separately in Figures 11–13.
  • H Real-World OOD Rollouts Visualization: Complete Towel Folding rollouts are visualized under background shift, light change, and novel-object conditions.These conditions are presented in Figures 14–16.
  • I LIBERO Visualization: LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, and LIBERO-Long tasks are visualized across bowl placement, alphabet-soup placement, stove operation, and moka-pot placement.The figures depict one representative task from each listed LIBERO suite.
  • J LIBERO-Plus Visualization: LIBERO-Plus camera-viewpoint perturbation is visualized for placing a wine bottle on top of the cabinet.This visualization isolates perturbation of camera viewpoints.
  • J LIBERO-Plus Visualization: LIBERO-Plus robot-initial-state and language-instruction perturbations are visualized for bowl placement and opening a drawer before placing a bowl inside.The corresponding figures show altered robot initial states and task instructions.
  • J LIBERO-Plus Visualization: LIBERO-Plus light-condition and background-texture perturbations are visualized for wine-bottle placement and opening a cabinet drawer.These figures represent visual-environment changes in lighting and background appearance.
  • J LIBERO-Plus Visualization: LIBERO-Plus sensor-noise and object-layout perturbations are visualized for pushing a plate toward the stove and placing cream cheese in a bowl.The final figures cover sensing degradation and rearranged object configurations.
Loading 2608.01397v1…