JEPA-WAM

Abstract

Robust robot control benefits from explicitly modeling state transitions, but video-generation world action models introduce substantial deployment cost. We introduce JEPA-WAM, a latent world action model 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 represents the temporal relation between current and future observations while preserving dense patch-level correspondence. Transition supervision directly shapes the backbone used for action prediction, without requiring explicit future-frame generation at deployment.

Method

JEPA-WAM architecture with a shared predictor for transition prediction and action generation — A frozen V-JEPA 2.1 encoder constructs dense current-future transition targets. The shared predictor learns the transition while producing representations for the flow matching action expert.

Experiments

JEPA-WAM is evaluated on in-distribution manipulation, scene-level OOD shifts, randomized bimanual manipulation, and controlled design ablations.

LIBERO-Plus

Table 2 · OOD success rate (%)

Method	Params.	PT	Camera	Robot	Language	Light	Background	Noise	Layout	Avg.
Without robot-policy pretraining
VLA-Adapter	0.5	✗	36.2	37.9	74.6	70.6	76.1	58.0	69.7	60.4
RoVLA	2	✗	58.4	36.3	92.9	95.6	95.0	80.9	73.0	76.0
ResVLA	2	✗	49.8	59.9	88.5	90.5	94.9	76.8	79.0	77.1
JEPA-WAM (Ours)	0.5	✗	79.2	59.2	68.2	93.3	94.6	83.6	76.1	79.2
With robot-policy pretraining
VLA-JEPA	2	✓	63.3	67.1	85.4	95.6	93.6	66.3	85.1	79.5
PokeVLA	0.5	✓	84.7	46.1	84.8	94.6	82.6	89.8	77.2	80.0
ABot-M0	4	✓	60.4	67.9	86.4	96.2	91.6	86.4	82.6	81.6
Cosmos-Policy	2	✓	75.8	63.3	81.7	96.5	88.9	92.7	82.2	83.0
π0.5	3	✓	69.4	75.3	82.6	96.7	96.8	84.3	86.2	84.5
Being-H0.7	3	✓	82.0	59.0	82.8	97.8	90.0	93.5	88.5	84.8
π0.5 + JEPA Obj. (Ours)	3	✓	66.0	82.0	86.5	96.8	96.0	88.3	88.3	86.3

Success rates after training on LIBERO demonstrations, evaluated on seven out-of-distribution shift categories without OOD fine-tuning.

RoboTwin 2.0

Table 3 · Clean / Random success rate (%)

Method	Adjust Bottle		Dump Bin Bigbin		Shake Bottle		Press Stapler		Stack Bowls Two		Avg.
Method	Cle.	Ran.	Cle.	Ran.	Cle.	Ran.	Cle.	Ran.	Cle.	Ran.	Cle.	Ran.
Without robot-policy pretraining
ACT	97	23	68	1	74	10	31	6	82	0	51.2	4.0
DP	97	0	49	0	65	8	6	0	61	0	48.0	1.6
DP3	99	3	85	53	98	19	69	3	83	6	73.9	8.3
JEPA-WAM	99	87	94	63	94	55	88	70	94	76	79.9	36.9
With robot-policy pretraining
RDT-1B†	81	75	64	32	74	45	41	24	76	30	56.0	24.1
π0†	90	56	83	24	97	60	62	29	91	41	62.6	24.0
π0.5†	98	26	95	41	99	82	66	22	87	40	75.4	37.3
π0.5 + JEPA Obj.†	100	30	90	58	100	84	86	22	93	45	84.6	37.5

Selected tasks are shown; averages are computed over all 20 tasks. † denotes large-scale robot-policy pretraining.

Design Ablations

Table 4 · LIBERO-Plus success rate (%)

Method	Camera	Robot	Language	Light	Background	Noise	Layout	Avg.
a. DINO+SigLIP	60.0	61.9	74.1	88.7	88.0	64.2	75.7	73.2
b. V-JEPA only	78.7	40.9	70.9	96.7	84.1	88.3	79.3	77.0
c. Future only	75.1	47.1	69.6	96.0	93.4	81.5	78.4	77.3
d. iREPA align.	68.9	45.5	69.2	90.9	89.1	81.5	77.7	74.7
e. Lower-16 align.	77.5	41.6	75.0	95.5	86.3	82.2	77.2	76.5
f. Full hidden	62.5	49.9	70.1	89.3	88.6	75.6	76.0	73.1
JEPA-WAM	79.2	59.2	68.2	93.3	94.6	83.6	76.1	79.2

Category-wise analysis of the visual representation, transition target, spatial alignment, and action-conditioning design.

Real-World Results

Real-world tasks and success rates under in-domain and out-of-domain settings — JEPA-WAM achieves 59.8% under in-domain conditions and 54.2% under OOD shifts. Applying the same transition objective to a pretrained π0.5 policy reaches 90.3% ID and 84.7% OOD.