JEPA-WAM: Stage-Level Joint-Embedding Prediction for World-Action Models in Robot Manipulation

2026-08-11

AI Summary

Existing general-purpose robotic policies model only short-term physical futures and lack explicit representations of stage-level semantic futures, hindering efficient planning for multi-stage manipulation tasks. This work proposes JEPA-WAM, which integrates a Stage-JEPA module into the Motus architecture to jointly model both short-term physical dynamics and stage-level semantic futures for the first time. The approach leverages a frozen V-JEPA2 encoder to extract stage-state representations and employs goal-conditioned joint embedding to predict latent representations of subsequent stages, thereby enhancing task-level planning capabilities. Evaluated on 50 tasks in RoboTwin 2.0, the method achieves an overall success rate of 90.25% and reduces action steps by an average of 5.97% among successful executions.

Abstract

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 JEPA-WAM, 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, JEPA-WAM achieves 90.25% overall success and reduces the mean number of execution steps in successful rollouts by 5.97% relative to the strongest baseline.

Problem

• robot manipulation
• stage-level future
• generalist robot policies
• task progress
• future prediction

Innovation

• Stage-Level Prediction
• Joint-Embedding Predictive Architecture
• World-Action Model
• Robot Manipulation
• Semantic Future