LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC

Shashank Hegde; Alexander Popov; Elie Aljalbout; Nikolai Smolyanskiy

cs.ROOct 8, 2026

LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC

Authors: Shashank Hegde, Alexander Popov, Elie Aljalbout, Nikolai Smolyanskiy

Organizations: NVIDIA

Abstract

World action models (WAMs) predict actions and future observations, typically from a reconstruction-based representation that carries noisy, redundant information which can complicate downstream predictions. We introduce LeWAM, a bidirectional transformer for forward, backward, inverse dynamics and policy prediction, on a decoder-free JEPA latent trained end-to-end through all four modes. We see the following benefits: 1) Alignment: linear probes read robot and object state from LeWAM's latent better than from a regular Le World Model (a forward-only JEPA world model), while the latent ignores visual distractors as well as LeWM does and far better than a reconstruction-based WAM. 2) Acting: Closed-loop evaluations of LeWAM match a regular flow-matching policy trained on the same encoder at matched size, while also providing a world model. 3) Planning: Sampling raw actions when planning with WAMs lets MPC exploit dynamics-model inaccuracies; planning in the noise space of the policy head instead improves the closed-loop performance of these WAMs.

Figures & tables

	decoder-	action loss	policy	forward	inverse	backward	multimodal	one
	free	to encoder	head	dyn.	dyn.	dyn.	head	backbone
Dreamer	–	–	✓	✓	–	–	–	–
TD-MPC2	✓	–	✓	✓	–	–	–	–
DINO-WM / V-JEPA 2-AC	✓	–	–	✓	–	–	–	–
LeWM	✓	–	–	✓	–	–	–	–
UWM	–	–	✓	✓	✓	–	✓	✓

Table 1: World models and world action models as published: what shapes the latent and what one network does. LeWM has no policy head. Dreamer [ 13 ] ; TD-MPC2 [ 15 ] ; DINO-WM and V-JEPA 2-AC [ 40 , 2 ] ; LeWM [ 24 ] ; UWM [ 42 ] ; Cosmos Policy [ 19 ] .

Figure 1: The four training modes on one window: the mode is set by what is masked at the transformer’s input. Nothing is decoded to pixels.

Figure 2: LeWAM: a shared ViT-tiny encoder with SIGReg on its single latent per frame, and one bidirectional transformer that serves the four training modes by masking. No pixel decodeding.

Panel of Figure 3 — Figure 3: The six tasks as the models see them, with their distractors: painted dot and bar on PushT, two wobbling boxes on the robomimic tasks (cameras and resolutions in Appendix B ).

PushT
	$R^{2}=1-\mathrm{MSE}/\mathrm{Var}$		Pearson $r$		MSE
Model	Task-state $\uparrow$	Distractors $\downarrow$	Task-state $\uparrow$	Distractors $\downarrow$	Task-state $\downarrow$	Distractors $\uparrow$
Dreamer	0.56 $\pm$ 0.01	0.03 $\pm$ 0.03	0.78 $\pm$ 0.00	0.30 $\pm$ 0.01	0.39 $\pm$ 0.01	0.96 $\pm$ 0.03
LeWM + diff. head	0.54 $\pm$ 0.05	0.00 $\pm$ 0.02	0.77 $\pm$ 0.03	0.06 $\pm$ 0.04	0.40 $\pm$ 0.05	1.07 $\pm$ 0.02
RecWAM	0.64 $\pm$ 0.02	0.79 $\pm$ 0.03	0.82 $\pm$ 0.02	0.89 $\pm$ 0.02	0.32 $\pm$ 0.02	0.21 $\pm$ 0.03
LeWAM (ours)	0.81 $\pm$ 0.01	0.00 $\pm$ 0.02	0.92 $\pm$ 0.00	0.06 $\pm$ 0.02	0.16 $\pm$ 0.01	1.19 $\pm$ 0.02

Table 2: Linear probe metrics across all tasks, 5 seeds (Negative

R^{2}

clipped to 0).

Panel of Figure 4 — Figure 4: Imagined rollouts with LeWAM using ground truth actions. The top image is for PushT and the bottom one is for the Lift Robomimic task. The first 3 frames are context frames that is passed as input to the model. The 5 frames that follow are imagined frames by rolling out the LeWAM’s forward mode using ground truth actions.

Figure 5: Closed-loop success of all models and planners, over 5 training seeds $\times$ 50 evaluation seeds (250 rollouts).

Task	Variant	$R^{2}$ state $\uparrow$	Policy $\uparrow$	DS-CEM $\uparrow$
Lift	All four modes	$\mathbf{0.423\pm 0.019}$	$\mathbf{94.0\pm 4.2}$	$\mathbf{97.6\pm 0.8}$
	No forward	$0.323\pm 0.045$	$91.6\pm 2.3$	$95.6\pm 2.9$
	No inverse	$0.386\pm 0.031$	$92.0\pm 5.5$	$95.6\pm 2.9$
	No policy	$0.359\pm 0.019$	–	$0.8\pm 1.6$
	No backward	$0.370\pm 0.056$	$88.0\pm 4.6$	$93.6\pm 2.3$
Can	All four modes	$\mathbf{0.839\pm 0.004}$	$\mathbf{82.8\pm 6.1}$	$\mathbf{78.8\pm 4.5}$

Table 3: LeWAM mode ablation. One mode removed at a time from the four-mode model (41M, object-distractor datasets, 300 epochs).

R^{2}

: ridge probe from the frozen latent to task state, computed as

1-\mathrm{MSE}/\mathrm{Var}

on z-scored targets. Policy / DS-CEM: closed-loop success (%), 50 episodes per cell. The no-policy variant has no policy head, so its DS-CEM column steers the inverse-trained flow head.

Task	With $\sqrt{A}$ shell	Without shell
Lift	$\mathbf{97.6\pm 0.8}$	$93.0\pm 3.0$
Can	$\mathbf{78.8\pm 4.5}$	$23.0\pm 8.0$
Square	$\mathbf{87.2\pm 4.7}$	$75.0\pm 3.0$

Table 4: Projection ablation. DS-CEM closed-loop success (%) with LeWAM, with and without projecting proposals onto the

\sqrt{A}

shell before decoding ( Section 4.4 ). mean

\pm

std over seeds 0–4, 50 episodes per cell.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Average closed-loop performance of all models/MPC, for each task, over 5 training seeds, 50 evaluation seeds: 250 unique rollouts.

data	200 demonstrations per task, frames resized once to 224 px, frameskip 5, windows of 3 context frames + 1 target, 5-step action chunks (35-D; 70-D on Transport), z-scored
schedule	300 epochs, batch 128, AdamW (lr $5\times 10^{-5}$ , weight decay $10^{-3}$ ), linear warm-up then cosine, bf16, gradient clip 1.0
encoder	ViT-tiny/14 at 224 px, [CLS] token, projector MLP (hidden 2048) to a 384-D latent
regularizer	SIGReg on the frame latent, weight 0.36 (0.09 $\times$ the four modes), 17 knots, 1024 projections
backbone	bidirectional transformer, depth 8, width 512, 16 heads $\times$ 64, MLP 2048
modes	all four on every batch; forward mode weight 2, inverse, policy and backward weight 1
context	random context length (1 to 3 frames, left-aligned truncation); no context-frame masking

Appendix

Table 5: LeWAM on the robomimic tasks: one configuration for Lift, Can, Square, ToolHang and Transport.

Explore similar work

Latent evolving World Action Model
Sep 23, 2026Xueji Fang, Boqiang Duan, Hua Wu +2Joint-Embedding Predictive Architecture Action-Conditioned World Models
Zhejiang University · Westlake University · Baidu Inc.
JEPA-WAM: Learning Vision-Language-Action Policies with Joint-Embedding World Modeling
Aug 10, 2026Yihan Lin, Jiawei He, Shifeng Bao +6Joint-Embedding Predictive Architecture Robotic Manipulation
School of Information, Renmin University of China, Beijing, China · XYZ Embodied AI, Beijing, China · Key Laboratory of Data Engineering and Knowledge Engineering, Beijing, China +3
LeapBot-WA: World-Anchor Action Models via Predictive Latent Alignments
Jul 27, 2026Pei Liu, Nan Zheng, Lang Zhang +8Robot Policy Learning World Models for Robotics
1The Hong Kong University of Science and Technology (Guangzhou) · 2The Hong Kong University of Science and Technology · 3Leapmotor +1

Sep 23, 2026, cs.CVJ/K move · Enter open · S save

Latent evolving World Action Model

Xueji Fang, Boqiang Duan, Hua Wu, Jingdong Wang+1

Zhejiang University · Westlake University · Baidu Inc.

Joint-Embedding Predictive Architecture Action-Conditioned World Models Robotic Manipulation World Action Models Latent World Models Offline Preference Optimization

Open paper AskPDF arXiv

World Action Models (WAMs) jointly model action generation and environment dynamics and are mostly built on pretrained Video Diffusion Models (VDMs). In VDM-based WAMs, observations are first encoded by a VAE, and the resulting compressed latents are then processed by large video diffusion backbones to extract effective features for action generation. However, this paradigm ties WAM performance and training cost to large-scale video generation pretraining, limiting WAM efficiency and scalability. In this paper, we theoretically and empirically investigate how visual representations affect action generation in WAMs. Our results show that predictive embeddings from Joint-Embedding Predictive Architecture (JEPA) encoders better support action generation than compressed VAE latents, with I-JEPA performing best in our encoder comparison. Based on these findings, we propose LeWAM, which conditions action generation on JEPA embeddings and models environment evolution by predicting future embeddings in the same space, without relying on a video diffusion backbone. We further find that imitation learning matches demonstrated actions but does not distinguish better actions from worse ones, even though small action deviations can greatly affect task success. To address this limitation without additional environment interaction or the human oversight required for resets and safety, we introduce Demonstration-Guided DPO (DemoDPO), an offline preference refinement stage that derives preference supervision directly from demonstrations. With only 0.4B trainable parameters, LeWAM achieves an average success rate of 92.28% on RoboTwin 2.0, comparable to that of state-of-the-art VLAs and WAMs, and maintains practical effectiveness on real-world manipulation tasks.

Abstract

Explore similar work

Latent evolving World Action Model

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

LeapBot-WA: World-Anchor Action Models via Predictive Latent Alignments