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

Shashank Hegde, Alexander Popov, Elie Aljalbout, Nikolai Smolyanskiy · N/A · 2026

Summary

Abstract Summary

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.

Key Points

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.

LeWAM

We use a multi-mode skeleton following UWM and Cosmos Policy. This model operates like the predictor in LeWM ( Figure 1 ). A training window contains four frames, $o_{t-2:t+1}$, at a frameskip of five, and the three 5-step action chunks between them, $a_{t-2:t}\in\mathbb{R}^{A}$.

The chunk dimension $A$ is five steps times the per-step action dimension ($A=10$ on PushT, $35$ on robomimic, $70$ on Transport), z-scored per dimension.

LeWAM — LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC — 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.

Diffusion steering: planning in the policy’s noise space

Given context $\mathbf{z}_{t}$ and a goal latent $z_{g}$ retrieved from a demonstration, following [ 40 , 24 ] , we score a chunk $a$ by with $\hat{f}$ the forward mode, and minimize it over $H$ chunks by a sampling-based planner.

The latter proposes $S$ plans, scores each on its rolled-out final latent, updates the distribution, repeats this process $I$ times, and executes the chosen plan’s first chunk before replanning. The update rule defines the planner. Sampling-based MPCs originally search over raw actions $a_{1:H}\in\mathbb{R}^{AH}$.

Diffusion steering: planning in the policy’s noise space — LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC — 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.

Experiments and Results

PushT [ 10 , 7 ] : a circular agent pushes a T-shaped block to a fixed target pose; actions are absolute target positions at 10 Hz in chunks of five. We use the 206 human demonstrations of Chi et al. [7] and paint two distractors that are independent of the actions: a dot on a Lissajous curve and a bar sliding along the top edge.

robomimic [ 25 ] : Lift, Can, Square, ToolHang and Transport, using the benchmark’s 200 proficient-human demonstrations replayed in robosuite.

Experiments and Results — LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC — 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 ).

Sources and Demonstrations

Figures are reproduced from Hegde et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.