PaperlayerView exact revision on arXiv

Source-linked AI summary

JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling

Yihan Lin, Jiawei He, Shifeng Bao, Chen Zhao, Yang Li, Xiaobo Wang, Yan Wang, Cheng Chi, Jing Zhang

arXiv:2608.09381v1cs.RO

TL;DR

Video-generation WAMs are costly to deploy, and latent WAMs can compress predictive representations or separate prediction from action generation. JEPA-WAM uses a spatially structured transition target and shared predictor; it reports strong generalization across simulation and real-world manipulation, with π0.5 reaching 86.3% on LIBERO-Plus.

  • Problem

    Video-generation WAMs incur substantial deployment cost, while latent WAMs may compress predictive representations or separate predictive modeling from the representations used for action generation.

    Source: §1 · p. 1

  • Method

    JEPA-WAM predicts a spatially structured joint current–future target in pretrained V-JEPA space and uses a shared predictor to shape the action-generating backbone.

    Source: §1 · p. 2

  • Results

    JEPA-WAM generalizes under visual and spatial shifts across LIBERO-Plus, RoboTwin 2.0, and real-world manipulation; pretrained π0.5 reaches 86.3% from 84.5%, the best overall result.

    Source: §1 · p. 2, §4.2 · p. 5

  • Takeaways & Limitations

    The transition-target formulation also benefits pretrained VLA policies while preserving their original perception and action pathways.

    Source: §1 · p. 2

    Task-shared visual transition supervision may be less expressive when the same observation leads to substantially different transitions under different instructions.

    Source: §6 · p. 8

Abstract

from arXiv · show

Robust robot control benefits from explicitly modeling state transitions, but video-generation world action models (WAMs) introduce substantial deployment cost. Existing latent WAMs avoid explicit future generation, but often compress predictive representations or separate predictive modeling from the representations used for action generation. We introduce JEPA-WAM, a latent WAM built in a pretrained V-JEPA space, which couples latent transition prediction with continuous action generation through a shared predictor. JEPA-WAM predicts a spatially structured joint current-future target that captures task-shared visual temporal structure between current and future observations, while preserving dense patch-level correspondence. Through the shared predictor, transition supervision directly shapes the backbone, from which dedicated representations are extracted for action prediction. The same design can also be instantiated in pretrained VLA policies while preserving their original perception and action pathways. On LIBERO-Plus, JEPA-WAM achieves 79.2%, the best result without large-scale robot-policy pretraining, while its pretrained $\pi_{0.5}$ instantiation reaches 86.3%, achieving the best overall performance. Experiments on RoboTwin 2.0 and real-world bimanual manipulation further demonstrate strong generalization under visual and spatial shifts.

1 Introduction

Video-generation WAMs are costly to deploy, while latent WAMs face challenges in target design and in connecting predictive supervision to action generation. JEPA-WAM addresses these challenges with a spatially structured transition target and a shared predictor, and reports strong generalization across simulation and real-world manipulation.

1 Introduction

  • Video-generation WAMs model future states alongside actions but incur substantial deployment cost from iterative future prediction.Latent WAMs retain predictive modeling without generating future observations.Source: §1 · p. 1
  • 79.2% average success makes JEPA-WAM the best LIBERO-Plus method without large-scale robot-policy pretraining; π0.5 improves from 84.5% to 86.3%, the best overall result.The paper also reports generalization under visual and spatial shifts on RoboTwin 2.0 and in real-world manipulation.Source: Figure 1 · p. 1, §1 · p. 2
  • Existing latent WAM targets may lose fine spatial structure, while separate prediction objectives may weakly influence the representations used for action generation.The paper also notes that adding predicted future representations as action context may expose the action module to redundant information.Source: §1 · pp. 1–2
  • JEPA-WAM builds on pretrained V-JEPA and predicts a spatially structured joint current–future target that captures temporal relations while preserving dense patch-level information.The formulation is also applied to pretrained VLA policies, including π0.5.Source: §1 · p. 2
  • A shared predictor models transitions and produces dedicated action-conditioning representations, so transition supervision shapes the action-generating backbone without requiring the action expert to use the full predicted transition.At deployment, transition prediction is removed and only action generation remains.Source: §1 · p. 2, Figure 2 · p. 2

2 Related Work

Prior work combines vision-language policies with action prediction, while WAMs also model future states and latent WAMs do so in representation space. JEPA-WAM differs from prior latent approaches by integrating predictive modeling and action conditioning through a shared predictor.

2 Related Work

  • Latent WAMs predict future states in representation space without explicit frame generation, unlike WAMs that commonly predict future observations.WAMs additionally model future states alongside action prediction.Source: §2 · p. 2
  • Prior latent WAMs reuse generative-model representations, predict compressed future tokens, or use dedicated latent dynamics to condition action generation.These approaches are described by Kim et al. 2026; Ye et al. 2026a; Li et al. 2026c, Zheng et al. 2025; Zhao et al. 2026; Luo et al. 2026a, and Chen et al. 2026, respectively.Source: §2 · p. 2
  • JEPA methods learn predictive representations by matching target embeddings rather than reconstructing pixels; robot-policy variants include separate world-model supervision and V-JEPA history inputs.VLA-JEPA uses a separate JEPA world model (Sun et al. 2026), while JEPA-VLA provides V-JEPA history representations as additional policy inputs (Miao et al. 2026).Source: §2 · pp. 2–3
  • JEPA-WAM uses V-JEPA for the current state and spatially structured transition target, with transition prediction directly supervising the shared action-generating backbone.This contrasts with approaches that add V-JEPA history representations as policy inputs or use a separate world model.Source: §2 · p. 3

3 Method

JEPA-WAM uses pretrained V-JEPA features and a shared predictor to couple dense transition prediction with action generation. Its joint current–future target preserves patch-level correspondence, and the same supervision can be transferred to pretrained VLA policies.

3.1 V-JEPA Representation Space and Joint Current–Future Target

  • Dense V-JEPA patch tokens form a joint current–future target that emphasizes visual relations across time rather than reconstructing a unique future observation.The target preserves camera and spatial-token ordering, providing task-shared temporal supervision without compressing transitions into a few global tokens.Source: §3.1 · pp. 3–4

3.2 Shared Predictor for Transition Prediction and Action Generation

  • A shared predictor produces patch-aligned transition predictions and dedicated action-conditioning representations, so transition supervision updates the backbone used for action generation.The transition loss uses mean patch-level cosine distance; the action representations condition the action expert.Source: §3.2 · p. 4

3.3 Joint Training and Deployment

  • Joint transition-prediction and action-generation losses update the shared predictor during training, while deployment removes the target branch and prediction head.The transition-loss weight balances the two objectives.Source: §3.3 · p. 4

3.4 Transfer to Pretrained VLA Policies

  • Pretrained VLA policies receive patch-level supervision from the same joint target through added future tokens, without changes to their original perception or action pathways.Future-token features are reshaped, projected, and spatially upsampled to align with the target.Source: §3.4 · pp. 4–5

4 Experiments

The experiments evaluate JEPA-WAM across simulation benchmarks, design ablations, and real-world bimanual manipulation, focusing on in-distribution performance and generalization under shifts.

4.1 Experimental Setup

  • The benchmarks comprise LIBERO, LIBERO-Plus, RoboTwin 2.0, and a real-world bimanual platform; the same policy is evaluated on LIBERO-Plus without OOD fine-tuning.RoboTwin training uses Clean demonstrations from 20 tasks and evaluates those tasks under Clean and Random settings.Source: §4.1 · p. 5

4.2 Generalization in Simulation

  • 96.7% average success on LIBERO keeps JEPA-WAM competitive, while transition supervision raises pretrained π0.5 from 96.9% to 97.8%.The LIBERO results indicate that transition modeling preserves strong in-distribution performance.Source: Table 1 · p. 5, §4.2 · p. 5
  • 79.2% on LIBERO-Plus is best without large-scale robot-policy pretraining; π0.5 reaches 86.3%, best overall, while JEPA-WAM scores 79.9% Clean and 36.9% Random on RoboTwin 2.0.The LIBERO-Plus evaluation uses no OOD fine-tuning; the RoboTwin evaluation covers Clean and Random settings.Source: §4.2 · p. 5, Table 2 · p. 6, Table 3 · p. 6

4.3 Design Analysis

  • V-JEPA features alone raise average success from 73.2% to 77.0%, with larger gains under camera, lighting, and noise perturbations.The comparison disables transition prediction and contrasts V-JEPA with DINOv2+SigLIP.Source: Table 4 · p. 6, §4.3 · p. 6
  • The joint current–future target scores 79.2% versus 77.3% for future-only prediction, and frozen probes find additional temporal-relation information that generalizes beyond probe training conditions.The joint target directly makes the relation between both temporal endpoints available to prediction.Source: Table 4 · p. 6, §4.3 · p. 7
  • Preserving patch structure scores 79.2% versus 74.7%; intermediate-layer supervision scores 76.5%, while conditioning on all hidden states scores 73.1%.The authors suggest local feature mixing may weaken spatial correspondence, and that action placeholders provide a dedicated action readout.Source: Table 4 · p. 6, §4.3 · p. 7

4.4 Real-World Experiments

  • Across five real-world bimanual manipulation tasks, JEPA-WAM averaged 59.8% ID and 54.2% OOD, versus π0’s 51.8% and 22.5%; transition supervision raised pretrained π0.5 from 77.5%/72.5% to 90.3%/84.7%.OOD shifts involved changes in scene appearance and object configurations.Source: Figure 5 · p. 7, §4.4 · p. 7

5 Conclusion

JEPA-WAM couples latent transition modeling and continuous action generation in pretrained V-JEPA space. Its experiments show generalization under visual and spatial distribution shifts.

5 Conclusion

  • JEPA-WAM couples latent transition modeling with continuous action generation through a shared predictor in pretrained V-JEPA space.The model predicts a spatially structured joint current–future target that preserves patch-level structure without reconstructing a unique future observation.Source: §5 · p. 8
  • Transition supervision directly shapes the backbone, which supplies task-conditioned representations for action generation.The supervision emphasizes localized changes, object relations, and spatial reconfiguration.Source: §5 · p. 8
  • The same formulation can be used in pretrained VLA policies, and experiments show generalization under visual and spatial shifts.Source: §5 · p. 8

6 Limitations

JEPA-WAM targets general visual temporal structure rather than task-specific future reconstruction, with a stated boundary when instructions imply different transitions from the same observation.

6 Limitations

  • JEPA-WAM learns general visual temporal structure rather than reconstructing a task-specific future.Source: §6 · p. 8
  • The shared visual supervision may be less expressive when one observation leads to substantially different transitions under different instructions.Source: §6 · p. 8
  • Language-conditioned or multimodal transition targets are identified as a future direction.Source: §6 · p. 8

A.1 JEPA-WAM Architecture

The architecture uses dense V-JEPA patch features and a shared predictor for transition prediction and action conditioning, then adapts transition supervision to pretrained VLA policies and removes transition prediction at deployment.

A.1 JEPA-WAM Architecture

  • A frozen V-JEPA encoder produces per-view patch features; a shared predictor maps visual-token states back to V-JEPA space for patch-wise cosine supervision, while 64 dedicated action placeholders provide action conditioning.Source: §A.1 · p. 11

A.2 Vision–Language Initialization

  • Vision-language initialization keeps the V-JEPA encoder frozen while jointly fine-tuning the visual projector and Qwen2.5-0.5B backbone, then robot-policy training freezes those base weights.The robot-policy stage optimizes Qwen LoRA adapters, the transition prediction head, and the action expert.Source: §A.2 · p. 11

A.3 Robot Policy Training

  • Robot-policy training jointly optimizes action generation and transition prediction using patch-wise cosine distance to the detached joint current–future target.The LIBERO setup trains jointly on four suites for 60K optimization steps.Source: §A.3 · p. 11

A.4 Transfer to Pretrained VLA Policies

  • For pretrained π0.5, 64 future tokens predict the transition while preserving the original perception and action pathways; the auxiliary transition loss trains alongside the flow-matching action loss.The future tokens can attend to the image and language prefix, but action tokens cannot attend to the new queries.Source: §A.4 · p. 12

A.5 Deployment

  • At deployment, JEPA-WAM needs only the current observation and instruction, and its transition target and prediction head are not run; π0.5 retains its original action pathway.For JEPA-WAM, the action expert generates each chunk from Gaussian noise using four Euler integration steps.Source: §A.5 · p. 12

B Experimental Details

The experiments evaluate transfer and robustness across LIBERO-Plus, simulated bimanual tasks, and real-world manipulation. They specify separate training and evaluation conditions for each setting.

B.1 LIBERO and LIBERO-Plus

  • LIBERO policies are evaluated on the standard test tasks, while policies trained on LIBERO demonstrations transfer to perturbed LIBERO-Plus environments without fine-tuning.Success rates follow the standard evaluation protocol.Source: §B.1 · p. 13

B.2 RoboTwin 2.0

  • On RoboTwin 2.0, policies train only on Clean demonstrations and are evaluated on both Clean and Random scene and object configurations across 20 tasks.The reported measure is average task success over all 20 tasks; the bimanual policy predicts 14-dimensional actions over a 50-step horizon.Source: §B.2 · p. 13

B.3 Real-World Setup

  • Real-world evaluation uses five bimanual tasks, with 100 training demonstrations per task and ID and OOD tests without additional fine-tuning.OOD tests change the background or initial object configurations while keeping the objective and language instruction unchanged; task completion is scored on a normalized scale.Source: §B.3 · p. 13, Table 5 · p. 13

C.1 Design Ablations

The ablations compare transition-target constructions and action-conditioning designs. The joint current–future target outperforms future-only prediction and explicit endpoint differencing on LIBERO-Plus.

C.1 Design Ablations

  • The endpoint-difference comparison tests whether joint encoding's benefit can be recovered by independently encoding the two observations and subtracting their features.The ablations also compare future-only targets, per-view convolutional alignment, lower-layer supervision, and conditioning on the full predictor hidden sequence.Source: §C.1 · p. 14
  • 79.2% for the joint current–future target exceeds 77.3% for future-only prediction and 70.9% for explicit endpoint differencing on LIBERO-Plus.The table reports LIBERO-Plus success rates for different transition-target constructions.Source: §C.1 · p. 14, Table 6 · p. 14
  • Direct patch-level supervision improves the average by 4.5 points over iREPA-style alignment, while final-layer supervision gains 2.7 points over the lower-layer variant.Dedicated action-conditioning representations improve by 6.1 points over conditioning on the full hidden sequence.Source: §C.1 · p. 14

C.2 Transition Information in the Joint Current–Future Target

The probes examine what temporal and spatial information the joint current–future target makes accessible. They find temporal structure beyond endpoint differencing, while deployment omits the transition-prediction branch.

C.2 Transition Information in the Joint Current–Future Target

  • 67.2% temporal-gap accuracy for the joint target exceeds 47.0% for endpoint differencing, with a 20.1-point paired improvement and 95% bootstrap interval [18.3, 21.9].The future observation is held fixed across gap classes; the caption identifies the controlled temporal-gap decoding table.Source: §C.2 · pp. 14–15, Table 7 · p. 15
  • 8.88 frames joint-target MAE compares with 13.32 for endpoint differencing on unseen gaps, a paired reduction of 4.44 frames; the joint target is not better at every gap.The probe trains on gaps {0, 20, 40} and tests on unseen gaps {10, 30, 50}; on interpolation gaps {10, 30}, MAE is 7.31 versus 9.51 frames.Source: §C.2 · pp. 15–16, Table 8 · p. 15
  • 0.582 mean R2 for the joint target exceeds 0.485 for endpoint differencing when predicting the trajectory residual after removing straight-line endpoint displacement.The probe predicts the complete 49 × 14 residual trajectory from frozen visual representations of the two endpoints, without observing intermediate robot states.Source: §C.2 · p. 16, Table 9 · p. 16
  • For direct endpoint-displacement prediction, endpoint differencing is slightly better, with a paired Joint-minus-Difference effect of −0.022 and 95% interval [−0.036, −0.009].This reverse result distinguishes strong first-order endpoint-change features from the joint target's temporal-relation and within-interval structure.Source: §C.2 · p. 16, Table 10 · p. 16
  • 79.2% average success for the joint target exceeds 70.9% for endpoint differencing, and frozen-representation probes show the joint target preserves temporal relations beyond first-order endpoint difference.The probes examine temporal-gap accessibility, generalization to unseen gaps, and trajectory structure after removing endpoint displacement.Source: §C.2 · pp. 14, 16
  • Patch-wise representation changes are positively associated with visually changing regions, but the correspondence is moderate rather than precise motion localization.The qualitative diagnostic compares patch-wise residuals with RGB change and interprets the target as capturing higher-level visual structure as well.Source: Spatial Diagnostic of the Joint Target · pp. 17–18, Table 11 · p. 17, Figure 7 · p. 17

C.3 Inference Efficiency

  • At deployment, JEPA-WAM omits target encoding, latent transition prediction, and iterative future-frame generation; these components are used only during training.Table 12 reports inference latency and frequency in the RoboTwin setting, but the supplied passages do not give its values.Source: §C.3 · p. 18, Table 12 · p. 18

D Complete RoboTwin 2.0 Results

The appendix provides task-wise RoboTwin 2.0 success rates for methods with and without large-scale robot-policy pretraining. The tables distinguish Clean and Random settings.

D Complete RoboTwin 2.0 Results

  • The appendix reports complete task-wise success rates for methods without large-scale robot-policy pretraining across RoboTwin 2.0 tasks.Table 13 covers Clean and Random settings and marks the best result for each task and setting.Source: D Complete RoboTwin 2.0 Results · p. 19, Table 13 · p. 19
  • A separate table reports task-wise success rates for methods with large-scale robot-policy pretraining.Table 14 covers Clean and Random settings and marks the best result for each task and setting.Source: D Complete RoboTwin 2.0 Results · p. 19, Table 14 · p. 20
  • The tables distinguish Clean from Random settings, with the best result identified for each task and setting.The supplied passages describe the tables' scope but do not provide individual task values.Source: Table 13 · p. 19, Table 14 · p. 20

E.1 Task-Level Summary

The section describes mean task-completion scores in Table 15 and rollout-level records in Tables 16–19, including the task abbreviations and score normalization.

E.1 Task-Level Summary

  • Table 15 reports mean normalized task-completion scores (%) across five tasks, averaged over 10 real-world rollouts per task and setting.Source: §E.1 · p. 21, Table 15 · p. 21

E.2 Per-Rollout Records

  • Tables 16–19 provide rollout-level results for π0, JEPA-WAM, π0.5, and π0.5 + JEPA Obj.Source: §E.2 · p. 21, Table 16 · p. 21, Table 17 · p. 22
  • Scores are normalized to [0, 1], with the final row giving each task’s mean over 10 rollouts; Br., Fr., Pl., St., and Dr. denote the five tasks.Source: §E.2 · p. 21

Related papers

Loading 2608.09381v1…