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

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

Paper: arXiv:2608.10780 | Authors: Xiao Liu, Yuguang Yang, Xi Wang, Kai Jiang, Cheng Chi, Yong Xu, Wenchao Ding, Yilun Chen, Yan Wang | Code: Not publicly available

One-Line Summary

StageWAM augments a Motus-based World Action Model (WAM) with a goal-conditioned JEPA predictor that estimates the next task stage's latent representation and injects it into the WAM's video stream via a bounded scalar gate, lifting overall success from 87.85% to 90.25% across 50 bimanual manipulation tasks while reducing execution steps in successful rollouts by 5.97%.

Background and Motivation

Generalist robot policies aim to convert multimodal visual observations, robot states, and task instructions into executable actions across diverse tasks. Two dominant paradigms have emerged: Vision-Language-Action models (VLAs) and World Action Models (WAMs). VLAs directly generate actions from observations and instructions, inheriting strong semantic understanding and instruction-following from pretrained vision-language models, but they do not explicitly model how the environment should change as the task progresses. WAMs jointly predict future video frames and action chunks, providing dense supervision over object motion, contact, and local scene evolution, yet predicting a plausible video-action trajectory does not mean the policy understands what state the task should reach next.

This distinction is critical. Existing WAMs typically represent the future as a fixed-length dense temporal chunk that captures local motion and contact dynamics but leaves the next task-relevant object-state or relation change implicit. For example, in a "pick up the cup and place it on the shelf" task, a WAM can predict the cup's motion trajectory over future frames, but it does not explicitly know that the cup should end up on the shelf. The policy may execute correct local motions while drifting off-course on overall task progress.

This raises a central question: how can a robot policy predict the next meaningful stage of task progress before generating local actions? Such a stage future should encode task-relevant object-state and relation changes (e.g., spatial alignment or robot-object interaction) without committing to pixel details or transition duration. Directly predicting a distant target image is poorly matched to this goal, since pixel-level prediction emphasizes appearance over task completion.

StageWAM's starting point is to distinguish two complementary futures: a short-term physical future describing immediate video-action evolution, and a stage-level semantic future specifying the next object or relation state to reach. The former is handled by the WAM, the latter by a JEPA predictor. Their combination gives the policy both "where to go next" and "how to get there."

Joint-Embedding Predictive Architectures (JEPAs) provide a natural mechanism for modeling this stage-level future. Instead of reconstructing every pixel, a JEPA predicts the latent representation of a target state from the current context, preserving task-relevant semantic and physical structure while remaining less sensitive to low-level appearance. StageWAM uses a goal-conditioned JEPA predictor to estimate the latent representation of the next task stage inferred from robot demonstrations, serving as an internal progress target to condition short-term world and action generation.

Preliminaries

V-JEPA2 is Meta's self-supervised video representation learning model that learns universal video representations through masked prediction of video clips. StageWAM uses a frozen V-JEPA2 encoder to extract latent representations of observations. This representation is sensitive to semantic and physical structure in video while being insensitive to pixel-level appearance changes, exactly what stage-level prediction requires. The model facebook/vjepa2-vitl-fpc64-256 accepts 64-frame inputs.

Motus is a Mixture-of-Transformers world-action model comprising three streams: a video-generation expert, an action expert, and a vision-language understanding expert. The video expert is initialized from Wan2.2 and models local visual evolution; the action expert processes robot states and dense action chunks; the understanding expert uses a frozen Qwen3-VL encoder to extract instruction- and observation-dependent features. The three streams exchange information through tri-modal joint attention while preserving their modality-specific normalization, feed-forward layers, and output heads.

Method

StageWAM's training proceeds in two stages. Stage I trains a goal-conditioned JEPA predictor to infer the next-stage future representation from the current observation and task goal. Stage II freezes this predictor and trains a WAM to generate the short-horizon visual future and action chunk conditioned on the predicted stage. At inference, the two components run sequentially in a receding-horizon loop.

Stage I: Training the JEPA Predictor

Stage-level prediction provides guidance about task evolution rather than its detailed realization. It captures semantic changes in object states and relations, leaving local appearance and motion to the WAM. Therefore, this stage's learning uses representation-space supervision before introducing visual and action objectives.

Stage-pair construction. Stage I constructs supervision from robot trajectories without manual stage labels. For an episode, candidate centers $\mathcal{C}=\{c_j\}_{j=1}^{M}$ are sampled every eight frames, and a centered 64-frame clip at each center is encoded with frozen V-JEPA2, yielding token representations $U_j=\{u_{j,n}\}_n$. The clip-level representation $v_j$ is obtained by pooling these tokens. Representation change is scored across a local window of $w=2$ candidate positions:

$$q_j = \max\left\{\mathcal{N}\left(\lVert v_{j^+} - v_{j^-} \rVert_2\right),\; \mathcal{N}\left(\max_n \lVert u_{j^+,n} - u_{j^-,n} \rVert_2\right)\right\}$$

where $\mathcal{N}$ is per-episode min-max normalization, $j^- = \max(1, j-w)$, and $j^+ = \min(M, j+w)$. The score $q_j$ considers both clip-level and token-level representation changes, capturing stage transition signals from coarse to fine.

Temporal NMS with radius $r=3$ retains at most $K=5$ transition keyframes:

$$\mathcal{K} = \{c_j \mid j \in \operatorname{NMS}(\{q_j\}_{j=1}^{M};\; r,\; K)\}$$

Keyframes together with the initial and terminal frames define boundaries $\mathcal{B} = \operatorname{sort}(\{0\} \cup \mathcal{K} \cup \{T-1\})$. For each stage interval $[b_i, b_{i+1})$, the current frame $t$ is paired with its next boundary $b_{i+1}$, constructing dense training pairs:

$$\mathcal{P} = \bigcup_{i=0}^{L-1}\{(t, b_{i+1}) \mid b_i \leq t < b_{i+1}\}$$

Keyframes only define boundaries; Stage I is trained with dense current-to-next-stage pairs $(t, b_{i+1})$. This automatic stage detection avoids manual annotation while leveraging V-JEPA2's semantic representation to identify meaningful task transition points.

Goal-conditioned stage prediction. StageWAM realizes the future prediction model $f_\theta$ as Stage-JEPA, using a frozen V-JEPA2 encoder and a trainable goal-conditioned predictor initialized from the V-JEPA2 predictor. Since V-JEPA2 requires 64-frame inputs, a dense pair $(t, b_{i+1})$ is converted into a 64-frame current observation slice $\mathcal{S}_t$ centered at $t$ and a 64-frame target observation slice $\mathcal{S}_t^{\text{target}}$ centered at $b_{i+1}$. The frozen encoder $E$ independently produces the current latent $z_t = E(\mathcal{S}_t)$ and target latent $z_t^{\text{target}} = E(\mathcal{S}_t^{\text{target}})$. The predictor produces token-level next-stage predictions $\widehat{Z}_t^{\text{target}}$, pooled and normalized as $\widehat{z}_t^{\text{target}}$ for supervision.

For a batch of $B$ normalized prediction-target pairs, the cosine loss is optimized:

$$\mathcal{L}_{\text{stage}} = \frac{1}{B}\sum_{i=1}^{B}\left(1 - (\widehat{z}_i^{\text{target}})^\top z_i^{\text{target}}\right)$$

This loss directly aligns the predicted latent with the target latent in representation space, avoiding pixel-reconstruction appearance bias and focusing the predictor on the semantic goal of task progress.

Stage II: Training the Stage-Conditioned WAM

The WAM complements stage guidance with the fine-grained visual dynamics and actions needed for local execution. Since joint optimization could let detail-oriented losses alter the semantic target, Stage II freezes JEPA and trains only the WAM and conditioning interface. The current observation slice and instruction produce $\widehat{z}_t^{\text{target}} \equiv \widehat{Y}_t^{\text{stage}}$.

StageWAM augments the WAM before its transformer layers. The video branch forms current-conditioned noisy video tokens $X_t^v \in \mathbb{R}^{N_v \times d_v}$, and the Stage-JEPA output $\widehat{Z}_t^{\text{target}} \in \mathbb{R}^{N_j \times d_j}$ contains $N_j$ predicted stage positions. This output is mean-pooled and an MLP $P_\psi$ matches the Wan video-token dimension before additive injection:

$$h_t = P_\psi\left(\operatorname{Pool}(\widehat{Z}_t^{\text{target}})\right) \in \mathbb{R}^{d_v}$$

$$\alpha = 0.2\,\sigma(\beta)$$

$$\widetilde{X}_t^v = X_t^v + \alpha\, h_t$$

where $\beta$ is a zero-dimensional trainable parameter and $\alpha$ is a global bounded scalar gate shared across all samples, layers, video tokens, and channels. The initial value is set to 0.02, ensuring that the stage-condition injection is a gentle residual update early in training. The condition vector $h_t$ remains sample-specific and is broadcast to all video tokens of that sample. The downstream WAM then predicts the local visual future and action chunk.

Training retains the WAM's native visual and action objectives and optionally includes a regularizer for the conditioning interface:

$$\mathcal{L} = \mathcal{L}_{\text{video}} + \mathcal{L}_{\text{action}} + \mathcal{L}_{\text{reg}}$$

$$\mathcal{L}_{\text{reg}} = \lambda_g\,\alpha^2 + \lambda_\rho(\rho - \rho_0)^2,\quad \rho = \frac{\lVert \alpha\, h_t \rVert}{\lVert X_t^v \rVert}$$

where $\rho$ measures the magnitude of the injected stage feature relative to the original video tokens, and $\rho_0$ is a manually selected target ratio (set to 0.08). The first term $\lambda_g \alpha^2$ biases the bounded gate toward a small residual update, while the second $\lambda_\rho(\rho - \rho_0)^2$ discourages the relative update from departing from $\rho_0$. The regularizer weights are $\lambda_g = 10^{-4}$ and $\lambda_\rho = 10^{-3}$.

The two stages thus learn complementary information: Stage I specifies what task-level change should occur, and Stage II learns how to realize it through detailed visual evolution and actions.

Closed-Loop Inference

At deployment, the policy maintains a causal 64-frame observation buffer $\mathcal{H}_t$. Before 64 observations are available, the earliest observed frame is repeated to pad the temporal slice. At each policy query, it performs:

$$\widehat{z}_t^{\text{target}} = f_\theta(\mathcal{H}_t, l)$$

$$(\widehat{V}_t^{\text{short}}, \widehat{A}_t) = W_\phi(o_t, s_t, l, \widehat{z}_t^{\text{target}})$$

The predicted stage latent passes through the same conditioning interface as in Stage II, after which the WAM predicts the finite-horizon visual future and action chunk. The environment executes the action chunk, appends the resulting observations to the buffer, and queries the policy again. Stage boundaries and target observation slices are training-time supervision only; inference repeatedly predicts its own stage condition from observed history.

The following flowchart illustrates StageWAM's complete data flow from stage detection through inference:

flowchart LR
    A["Robot Demo
Trajectory"] --> B["V-JEPA2
Feature Extract"]
    B --> C["Stage Boundary
Detection (NMS)"]
    C --> D["Stage-Pair
Construction"]
    D --> E["Stage-JEPA
Training (Cosine Loss)"]
    E --> F["Frozen JEPA
Predictor"]
    F --> G["Gated Injection
alpha=0.2*sigma(beta)"]
    H["Current
Observation"] --> G
    G --> I["Motus WAM
Video+Action Gen"]
    I --> J["Execute
Action Chunk"]
    J --> H

Figure 1: StageWAM framework overview. Frozen Stage-JEPA guidance is injected into the WAM video stream, where predicted next-stage latent tokens condition short-term video-flow and action-chunk prediction.

Experimental Results

Experiments are conducted on RoboTwin 2.0, a bimanual manipulation benchmark with 50 tasks evaluated under clean and randomized environment configurations. Each task-configuration pair runs 100 closed-loop episodes with early stopping disabled, totaling 10,000 episodes. The primary metric is task success rate.

Baselines include GO-1 (AgiBot's embodied foundation model based on a vision-language-latent-action formulation), $\pi_{0.5}$ (a generalist VLA that directly predicts action chunks), X-VLA (which strengthens VLA control through cross-task and cross-embodiment pretraining), and Motus (a WAM jointly modeling understanding, future video, and action prediction). Baseline numbers are published results from the Motus paper, and StageWAM and Motus are evaluated with the same closed-loop protocol.

Table 1 shows success rates grouped by semantic manipulation skill. StageWAM achieves 91.42% in clean settings and 89.08% in randomized settings, outperforming Motus in 5 of 7 semantic categories under the clean setting and 6 under randomization.

Semantic Category (Tasks)	GO-1 Clean	GO-1 Rand	pi0.5 Clean	pi0.5 Rand	X-VLA Clean	X-VLA Rand	Motus Clean	Motus Rand	StageWAM Clean	StageWAM Rand
Acquisition & Lifting (5)	76.40	76.60	36.80	36.20	80.80	74.20	94.20	94.60	92.60	92.00
Handover (2)	10.50	10.00	23.00	18.50	36.50	18.50	82.00	68.00	96.50	93.50
Targeted Placement (13)	25.00	23.23	42.46	46.00	78.85	79.77	82.62	85.08	88.23	86.54
Container Packing (8)	36.88	40.25	34.75	37.00	74.38	73.25	88.62	87.25	91.12	90.00
Arrangement & Stacking (8)	17.62	17.25	43.62	42.38	63.00	64.75	90.12	88.00	91.88	89.62
Articulated/Device (7)	61.43	54.71	47.86	47.43	78.43	82.43	93.57	89.86	90.71	90.00
Tool Use (7)	42.71	37.71	57.86	58.43	70.57	73.71	91.29	86.43	95.57	87.86
All Tasks (50)	37.86	36.24	42.98	43.84	72.88	72.84	88.66	87.02	91.42	89.08

Table 1: RoboTwin 2.0 success rates (%) grouped by dominant semantic manipulation skill. Bold denotes the highest rate in each row.

The strongest results occur on tasks whose execution depends on an explicit change in task progress. Handover, targeted placement, and container packing require the policy to maintain a desired object ownership or object-target relation across transfer and alignment. Arrangement and stacking require ordered intermediate relations, while tool use and dynamic manipulation require recognizing whether the robot should approach, align, or interact. These tasks directly match StageWAM's two-scale design: Stage-JEPA specifies what task-relevant state should be reached next, and the WAM models how to realize it through local visual dynamics and actions.

In contrast, acquisition and lifting depend primarily on immediate grasp geometry and contact stability, leaving less complementary stage structure for JEPA conditioning. Articulated/device interaction contains meaningful state changes, but success also requires mechanism-specific contact, constrained motion, and precise timing, for which the current indirectly conditioned stage signal can be insufficient. This indicates that stage-level guidance yields the largest gains for tasks requiring explicit task-progress reasoning, and limited gains for tasks where immediate perceptual-motor control suffices.

Table 2's ablation study isolates the four main components of StageWAM:

Configuration	Clean (%)	Randomized (%)	Overall (%)
w/o JEPA	87.80	86.66	87.23
Current-frame tokens only	91.10	88.68	89.89
w/o Stage-I training	91.34	88.62	89.98
w/o gate regularizers	91.00	88.76	89.88
StageWAM (full)	91.42	89.08	90.25

Table 2: Ablation of StageWAM on RoboTwin 2.0. Success rates averaged over 50 tasks.

Removing JEPA produces the largest degradation (from 90.25% to 87.23%), indicating that the stage-conditioning pathway provides information beyond the short-horizon WAM backbone. Keeping the JEPA branch but preventing future-latent injection also lowers success, supporting injecting the predicted next-stage latent into the WAM representation rather than using it only to modify the current observation. The model without Stage-I training remains competitive (89.98%), suggesting that the pretrained V-JEPA2 predictor already supplies useful generic visual structure; nevertheless, training on current-to-next-stage pairs gives the best clean, randomized, and overall results. Removing both gate regularizers produces a smaller drop, indicating that the sigmoid-bounded gate already stabilizes the residual update while the two regularization terms provide additional constraints on its scale.

The qualitative analysis (Figure 2) compares StageWAM and Motus on a representative successful rollout. Although both eventually satisfy the success condition, their execution patterns differ substantially. In the execution trace, StageWAM reaches success after a short action sequence, while Motus continues for a much longer horizon and repeatedly moves the manipulated object up and down. This illustrates the role of the predicted stage latent as an internal progress target that biases local actions toward task progress and reduces redundant corrective motions.

Figure 2: Qualitative case study of stage-guided execution. StageWAM uses the predicted next-stage latent to guide local action generation and reaches the success condition with fewer executed steps.

Figure 3: Visual rollout comparison. StageWAM quickly advances from initial interaction to the next task-relevant state, while Motus performs repeated local corrections around the object.

For execution efficiency, StageWAM reduces the mean episode length from 156.28 to 132.80 steps across all rollouts, a 15.02% reduction relative to Motus. This aggregate comparison combines successful and failed episodes. Conditioning on successful rollouts provides a more direct measure: StageWAM reduces the mean length from 79.57 to 74.82 steps, a 5.97% reduction while preserving task completion. This is consistent with stage guidance helping the policy progress more directly toward intermediate task targets and reducing redundant local corrections.

Limitations

Author-acknowledged limitations. The paper's conclusion explicitly identifies two future directions that imply current limitations. First, WAM inference latency limits closed-loop control frequency. While StageWAM improves success rates significantly, the two-stage sequential inference (JEPA prediction followed by WAM generation) adds latency overhead that may become a bottleneck in precision manipulation requiring high-frequency control. Second, the current multi-scale latent representation does not yet jointly encode coarse-grained task progress and fine-grained contact, geometry, and motion information, meaning that the stage condition is currently indirect and yields limited gains for tasks requiring mechanism-specific contact and precise timing (e.g., articulated/device interaction).

Independent assessment of limitations. First, the automatic stage boundary detection relies on a V-JEPA2 representation-change heuristic with fixed parameters (NMS radius $r=3$, maximum keyframes $K=5$, candidate interval of 8 frames). Different task structures may require different boundary granularity, and the paper does not examine the robustness of these parameters. Second, the gated injection acts only on the video token stream; the action expert does not directly receive the stage condition, meaning that stage information must indirectly influence action prediction through joint attention. The theoretical justification for this design choice is not thoroughly discussed. Third, all experiments are limited to the simulated RoboTwin 2.0 environment, with no real-robot validation or sim-to-real gap assessment.

Conclusion and Outlook

StageWAM's core contribution is distinguishing two complementary temporal scales of future modeling for robot manipulation and injecting the stage-level semantic prediction directly into the world-action model's video stream via a gated interface. This design does not alter the WAM's inference architecture, only adding a bounded residual injection before transformer layers, allowing the pretrained Motus model to be inherited seamlessly. Experiments across 50 RoboTwin 2.0 tasks demonstrate that stage-level future guidance provides substantial gains for manipulations requiring explicit task-progress reasoning, while ablations confirm the individual contributions of the JEPA predictor, latent injection, Stage-I training, and gate regularization.

From a broader perspective, StageWAM represents a trend of combining predictive ideas from representation learning (JEPA) with generative world models (WAM), using semantic-level future targets to guide local generation. This differs from methods that directly predict target images or visual subgoals, which are susceptible to pixel-level appearance bias. StageWAM's latent-space prediction preserves task-relevant structure while avoiding sensitivity to low-level appearance. Future work will proceed along two robotics-centered axes: reducing WAM inference latency to enable higher-frequency closed-loop control, and learning multi-scale latent representations that jointly encode coarse-grained task progress and the fine-grained contact, geometry, and motion information required for precise manipulation.

Golden quote: Knowing where to go next matters more than knowing how to get there, but in robot manipulation you need both, and you need them to work together.