# StageWAM: Joint-Embedding Stage Prediction for World-Action Models in Robot Manipulation

arXiv: https://arxiv.org/abs/2608.10780 · HTML: https://mlsanity.com/paper/2608.10780/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 26 of its 40 references that are in the map.
- First posted on arXiv 11 Aug 2026.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)
- Vision-language-action models / [World model-based VLA](https://mlsanity.com/vla/world-model-based-vla/)

## Builds on (most cited first, 26 of 26)

- [OpenVLA: An Open-Source Vision-Language-Action Model](https://mlsanity.com/paper/2406.09246/) (2024)
- [Qwen3-VL Technical Report](https://mlsanity.com/paper/2511.21631/) (2025)
- [Wan: Open and Advanced Large-Scale Video Generative Models](https://mlsanity.com/paper/2503.20314/) (2025)
- [π0.5: a Vision-Language-Action Model with Open-World Generalization](https://mlsanity.com/paper/2504.16054/) (2025)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [Revisiting Feature Prediction for Learning Visual Representations from Video](https://mlsanity.com/paper/2404.08471/) (2024)
- [RoboTwin 2.0: A Scalable Data Generator and Benchmark with Strong Domain Randomization for Robust Bimanual Robotic Manipulation](https://mlsanity.com/paper/2506.18088/) (2025)
- [X-VLA: Soft-Prompted Transformer as Scalable Cross-Embodiment Vision-Language-Action Model](https://mlsanity.com/paper/2510.10274/) (2025)
- [Motus: A Unified Latent Action World Model](https://mlsanity.com/paper/2512.13030/) (2025)
- [Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets](https://mlsanity.com/paper/2504.02792/) (2025)
- [Unified Video Action Model](https://mlsanity.com/paper/2503.00200/) (2025)
- [VLA-JEPA: Enhancing Vision-Language-Action Model with Latent World Model](https://mlsanity.com/paper/2602.10098/) (2026)
- [Motubrain: An Advanced World Action Model for Robot Control](https://mlsanity.com/paper/2604.27792/) (2026)
- [LaWAM: Latent World Action Models for Efficient Dynamics-Aware Robot Policies](https://mlsanity.com/paper/2606.15768/) (2026)
- [OA-WAM: Object-Addressable World Action Model for Robust Robot Manipulation](https://mlsanity.com/paper/2605.06481/) (2026)
- [Scaling World Model for Hierarchical Manipulation Policies](https://mlsanity.com/paper/2602.10983/) (2026)
- [Learning 4D Geometric Priors for Inference-Efficient World Action Models](https://mlsanity.com/paper/2607.05468/) (2026)
- [HarmoWAM: Harmonizing Generalizable and Precise Manipulation via Adaptive World Action Models](https://mlsanity.com/paper/2605.10942/) (2026)
- [H-WM: Robotic Task and Motion Planning Guided by Hierarchical World Model](https://mlsanity.com/paper/2602.11291/) (2026)
- [Unifying Object-Centric World Models and Diffusion Policy: A Hierarchical Framework for Multi-Stage Robotic Tasks](https://mlsanity.com/paper/2606.08775/) (2026)
- [Improving Vision-Language-Action Model Fine-Tuning with Structured Stage and Keyframe Supervision](https://mlsanity.com/paper/2606.26801/) (2026)
- [DexFuture: Hierarchical Future-State Visuomotor Targeting for Bimanual Dexterous Tool Use](https://mlsanity.com/paper/2606.05699/) (2026)

## Built on by (most cited first, 3 of 3)

- [PAVE: Predictive Alignment and Value-Guided Evolution for World-Action Policies](https://mlsanity.com/paper/2608.30378/) (2026)
- [CoRe-WAM: Correspondence-Aligned Temporal Residuals for World Action Models](https://mlsanity.com/paper/2609.27314/) (2026)
- [Copper-Policy: Focus on the Representation for Robust Robot Manipulation](https://mlsanity.com/paper/2609.32779/) (2026)

## Abstract (arXiv)

Generalist robot policies aim to map multimodal observations and linguistic task instructions to actions across diverse tasks. However, existing methods typically represent the future as a fixed, short video-action chunk. This short-term future captures local scene evolution for action execution, but it does not explicitly describe the stage-level future that specifies how a task should progress from its current stage to the next. We therefore distinguish two complementary futures for robot manipulation: a short-term physical future to capture local scene evolution and a stage-level semantic future to represent task progress. We introduce StageWAM, which augments a Motus-based World Action Model (WAM) with Stage-JEPA, a goal-conditioned Joint-Embedding Predictive Architecture (JEPA) predictor. Given the current observation and task instruction, Stage-JEPA uses a frozen V-JEPA2 encoder to extract the current-state representation and predicts the latent target of the next inferred stage. Across 50 RoboTwin 2.0 tasks in clean and randomized environments, StageWAM achieves 90.25% overall success and reduces the mean number of execution steps in successful rollouts by 5.97% relative to the strongest baseline.