Robotics · 2026-10-09
LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC
Shashank Hegde, Alexander Popov, Elie Aljalbout, Nikolai Smolyanskiy
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
LeWAM is a decoder-free JEPA world action model that trains a bidirectional transformer end-to-end for forward, backward, inverse dynamics and policy prediction. It uses a single non-reconstruction latent from a ViT-tiny encoder with SIGReg. For planning, it replaces action-space MPC decisions with diffusion-steering: it samples in the policy’s flow/noise space and projects proposals onto the sampler’s typical set.
Why it matters
The authors argue that reconstruction-based WAM latents encode noisy, redundant visual detail that can hinder downstream predictive tasks and control. LeWAM’s reconstruction-free, end-to-end JEPA latent is designed to ignore visual distractors while still encoding task-relevant robot/object state, and diffusion-steering-based MPC is proposed to avoid failure modes of action-space planning that exploit model inaccuracies.
Method
- Decoder-free JEPA world action model: bidirectional transformer with four training modes (forward/backward/inverse/policy) over a single encoder latent; trained end-to-end with SIGReg regularization and no pixel decoder or pixel/token losses.
- Action generation via rectified flow: train velocity prediction for noisy actions and integrate from Gaussian noise using eight Euler steps to produce action chunks.
- Diffusion steering MPC: shift the planner decision variable from actions to the policy head’s deterministic noise input (flow/epsilon), decode candidates in imagined states, and project proposals back onto the typical set before decoding.
Abstract (from arXiv)
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.
Related papers
- LeWAM: End-to-End World and Action Modeling with JEPAs
- {\Delta}WAM: Distilling Action Tangent Fields into World Action Models
- De-Diffusing WAMs: Action-Sufficient Direct Prediction for Robot Control
- MARS: Systematic Multi-Level Acceleration for World Action Models
- MiniWAM: Learning Compact Future Targets for Efficient World-Action Modeling
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-08 · Indexed by Pub-AI: 2026-10-09 UTC
LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC
Preprint posted: 2026-10-08 · Original record ↗
- arxiv
- 2610.12407
Explicit source versions (1)
- v1 · 2026-10-08
- Preprint available
Source date: 2026-10-08 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
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.