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

arXiv:2610.12407 · cs.RO, cs.AI, cs.CV · Submitted 2026-10-08 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: Robotics Radio. Generated commentary on the latest robotics and control papers.

Rosa: I'm Rosa, and with me are Dev and Taro, guest researcher.

Dev: Today's paper: "LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC".

Rosa: The gist LeWAM introduces a unified, reconstruction-free world action model that combines forward, backward, inverse dynamics, and policy prediction through a bidirectional transformer trained end-to-end across all four modes.

Dev: First, who's behind it and why it matters.

Paper summary: Rosa: So we're looking at LeWAM today, which is this new idea called a JEPA World Action Model with Diffusion-Steering Based MPC. The main thing here is they built a model that does four things at once: forward dynamics, backward dynamics, inverse dynamics, and policy prediction.

Dev: That’s right. It’s not just one model; it's this bidirectional transformer that learns all those modes end-to-end across all four of them. They claim this system works on a decoder-free JEPA latent space trained through all four modes simultaneously.

Taro: What I find interesting is how they handled the representation—they used a single latent vector, z t in R384, which they get by projecting the CLS token through a BatchNorm MLP projector. Rosa They argue this latent representation captures dynamic robot and object states as effectively as reconstruction-based baselines, but crucially, it ignores visual distractors just like LeWM did.

Dev: Exactly. So, in terms of alignment, they show that this latent space reads the robot and object state better than a regular forward-only JEPA world model. But the key is that it ignores the visual stuff that isn't relevant to the task dynamics.

Rosa: And for acting, they validated their performance by showing that this integrated flow-matching policy matches standard behavioral cloning policies when you do closed-loop evaluations, while still providing a world model for dynamics prediction.

Taro: That means it doesn’t just predict what will happen next; it actually acts on that prediction in a way that matches what you'd expect from a standard policy trained the same way. It shows they've got the dynamics and control loop working together properly.

Dev: And then there’s planning, which is where things get tricky. They establish a stable planning paradigm by shifting the decision variable of sampling-based MPC from raw actions to the deterministic noise space of this flow-matching policy.

Rosa: So, instead of sampling raw actions—which often fails in practice because they exploit model inaccuracies—the planner searches strictly within the world model’s imagination, evaluating candidates decoded at imagined states.

Taro: That diffusion steering approach means the search distribution stays Euclidean; only proposals are projected, which they say is what lowers the success rate of diffusion steering-based CEM on all three tasks.

Dev: And they found that sampling in this diffusion noise space gives both LeWAM and RecWAM a performance gain at inference time, whereas action-space MPC collapses every model to around three to eight percent success.

Rosa: So, what this means for someone just listening is that if you’re building a robot system and you try to plan using raw actions from a dynamics model, it often fails because the model has small errors. But by planning in the noise space of this policy head, you improve closed-loop performance across tasks.

Taro: It suggests that having these four training modes—forward, backward, inverse dynamics—actually adds value to both forming an informative latent space and achieving good closed-loop performance on average across different tasks.

Dev: The paper shows that LeWAM encodes task-relevant state better than LeWM and even reconstruction-based models, while it also ignores those visual features that don't affect the task dynamics.

Rosa: And they matched the model sizes in terms of trainable parameters, with LeWAM having forty-one point zero million trainable parameters. That’s a pretty compact size for what it achieves.

Taro: The paper flags one limitation: the projection ablation shows that removing that projection step in Section four point four lowers the diffusion steering-based CEM success on all three tasks, so you need that specific step to get the best planning performance <ref:2610.12407#pg2>.

Dev: So, to wrap up this summary of LeWAM: it’s a decoder-free JEPA world action model where every mode—forward, backward, inverse dynamics and policy prediction—is trained end-to-end on task pixels under a strict regularizer. The results show its latent space is more informative than LeWM’s and void of visual features that don't affect dynamics.

Rosa: And by sampling in the diffusion noise space during inference, we can effectively plan with this model and improve average closed-loop performance across tasks, which is a significant result for autonomous systems.

Taro: The next step they suggest is applying this to real-world data and improving planning algorithms using LeWAM as the foundation.

Dev: So, if you’re working on robotics control loops right now, you should keep an eye on this paper because it shows a way to get better planning stability by changing how you sample during model predictive control.

Conclusion: Rosa: So we’re looking at LeWAM today, which is this new model combining dynamics, policy prediction, and planning using diffusion steering.

Dev: That’s right, a bidirectional transformer trained on four different modes to build a world model that can actually act on it.

Taro: The core idea is this latent space captures the robot and object states well enough to ignore visual clutter while still knowing how things move.

Rosa: It seems they’ve managed to get the dynamics prediction and action generation working together in one unified architecture.

Dev: They show that by planning in a specific noise space within that model, you can get better results than just sampling raw actions from a standard dynamics predictor.

Taro: The implication for autonomy is that you can use this model’s imagination to test potential moves, which improves performance on closed-loop tasks significantly.

Rosa: So the authors are arguing this isn't just another prediction tool; it’s a way to structure the entire decision-making process around the world model.

Dev: They claim the resulting latent space is more informative than previous models like LeWM and even reconstruction-based ones, especially regarding task-relevant state information.

Taro: It means we might get better autonomy when dealing with complex environments because the model filters out noise that doesn't actually affect the physics of what we’re trying to do.

Rosa: The authors are showing that this framework works across different tasks, which is a big deal for real-world deployment outside of lab settings.

Dev: We need to keep watching how they handle things like the latency and the actual robustness when the model encounters states it hasn't seen during training.

Shashank Hegde Alexander Popov Elie Aljalbout Nikolai Smolyanskiy

NVIDIA

cs.RO, cs.AI, cs.CV

Submitted: 2026-10-08

Updated: 2026-10-08

Comments: 14 pages, 6 figures

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 85/100

The gist: The gist LeWAM introduces a unified, reconstruction-free world action model that combines forward, backward, inverse dynamics, and policy prediction through a bidirectional transformer trained

Key concepts

Bidirectional Transformer
A transformer architecture used in LeWAM that can simultaneously perform four tasks: forward dynamics, backward dynamics, inverse dynamics, and policy prediction. It processes input data through a depth 8 backbone to generate a single latent vector representing the state.
JEPA Structure
A training structure applied during training that imposes constraints on the latent space. The forward and backward modes enforce this structure, while inverse and policy modes make the space more aligned with task goals, resulting in a highly informative representation.
Diffusion Steering MPC
A planning method where decisions are made by sampling within the deterministic noise space of the flow-matching policy instead of raw actions. This diffusion steering approach searches strictly inside the world model's imagination, leading to better performance than traditional action-space MPC.
Decoder-Free Model
LeWAM operates without needing a pixel decoder or tokenizer to generate outputs. It uses the learned latent vector directly for all tasks (dynamics prediction, action generation), relying solely on the transformer backbone's representation.

Terminology

Summary

The gist

LeWAM introduces a unified, reconstruction-free world action model that combines forward, backward, inverse dynamics, and policy prediction through a bidirectional transformer trained end-to-end across all four modes.

Alignment

LeWAM demonstrates representational alignment by capturing dynamic robot and object states as effectively as reconstruction-based baselines while successfully ignoring visual distractors The latent space captures dynamic robot and object states as effectively as reconstruction-based baselines, while successfully ignoring visual distractors.

Acting

The acting capabilities of LeWAM are validated by showing that the integrated flow-matching policy matches the performance of standard behavioral cloning policies while simultaneously providing a world model for dynamics prediction This validation shows that the integrated flow-matching policy matches the performance of standard behavioral cloning policies while simultaneously providing a world model for dynamics prediction.

Planning

The paper establishes a stable planning paradigm by shifting the decision variable of sampling-based MPC from raw actions to the deterministic noise space of the flow-matching policy This diffusion steering approach searches strictly within the world model’s imagination, evaluating candidates decoded at imagined states.

Key aspects of LeWAM include:

**)&LeWAM introduces a bidirectional transformer capable of simultaneous forward dynamics, backward dynamics, inverse dynamics, and policy prediction. The predictor is conditioned on actions, while the encoder is trained end-to-end directly from task pixels under a strict regularizer. The architecture eschews pixel decoders, tokenizers, or any pixel- or token-space loss. The latent representations are shaped entirely by the dual objectives of dynamics prediction and action generation. Each frame is independently encoded at a resolution of 224 pixels using a ViT-tiny/14, and this vector serves as the model’s sole representation. The backbone is a bidirectional transformer of depth 8, width 512 and 16 heads that serves four modes by masking. The encoder’s [CLS] token passes through a BatchNorm MLP projector to produce a single latent vector, zt ∈ R384. The latent is read by the probes, compared by the planner, and regressed by the dynamics model. The training involves four modes: (i) Forward, (ii) Backward, (iii) Inverse, and (iv) Policy. During training, every mode is executed on every batch. The Forward and Backward modes, alongside SIGReg [14], impose a JEPA structure on the latent space, while the Inverse and Policy modes make the space more task-aligned. Action generation employs rectified flow [16], where during training, we sample σ ∼ U(0, 1) and construct a noisy action aσ = (1 − σ)a + σϵ, training the backbone to predict the velocity target ϵ − a. The model operates without a pixel decoder, tokenizer, or any pixel- or token-space loss. The search distribution stays Euclidean; only proposals are projected (ablated in Table 4). The paper shows that sampling in noise space gives both LeWAM and RecWAM a performance gain at inference time, whereas action-space MPC collapses every model to 3–8% success. The two logical next steps are applying LeWAM to real-world data and improving planning algorithms with LeWAM. The paper shows that on average, LeWAM encodes task-relevant state better than LeWM and even the reconstruction-based models, and ignores distractors as well as LeWM. The model sizes are matched in trainable parameters, with LeWAM having 41.0M trainable parameters. The planning protocol details state that every planner uses the same budget, 64 samples, 3 iterations and 8 elites. The robomimic horizons are 100 steps on Lift, 400 on Can and Square, and 700 on ToolHang and Transport. The results show that all modes add value to both forming an informative latent space as well as closed-loop performance on average over tasks. The projection ablation shows that removing the projection step of Section 4.4 lowers DS-CEM success on all three tasks. The paper presents LeWAM, a decoder-free JEPA-based world action model. The results show that LeWAM’s latent space is more informative than LeWM’s, and is void of the visual features that do not affect task dynamics. The paper shows that by sampling in diffusion noise space (an inference-only procedure), we can effectively plan with LeWAM and improve average closed-loop performance across tasks. The paper presents LeWAM, a decoder-free JEPA-based world action model. The results show that LeWAM’s latent space is more informative than LeWM’s, and is void of the visual features that do not affect task dynamics. The paper shows that by sampling in diffusion noise space (an inference-only procedure), we can effectively plan with LeWAM and improve average closed-loop performance across tasks. The paper presents LeWAM, a decoder-free JEPA-based world action model. The results show that LeWAM’s latent space is more informative than LeWM’s, and is void of the visual features that do not affect task dynamics. The paper shows that by sampling in diffusion noise space (an inference-only procedure), we can effectively plan with LeWAM and improve average closed-loop performance across tasks. The paper presents LeWAM, a decoder-free JEPA-based world action model. The results show that LeWAM’s latent space is more informative than LeWM’s, and is void of the visual features that do not affect task dynamics. The paper shows that by sampling in diffusion noise space (an inference-only procedure), we can effectively plan with LeWAM and improve average closed-loop performance across tasks. The paper presents LeWAM, a decoder-free JEPA-based world action model. The results show that LeWAM’s latent space is more informative than LeWM’s, and is void of the visual features that do not affect task dynamics. The paper shows that by sampling in diffusion noise space (an inference-only procedure), we can effectively plan with LeWAM and improve average closed-loop performance across tasks. The paper presents LeWAM, a decoder-free JEPA-based world action model. The results show that LeWAM’s latent space is more informative than LeWM’s, and is void of the visual features that do not affect task dynamics. The paper shows that by sampling in diffusion noise space (an inference-only procedure), we can effectively plan with LeWAM and improve average closed-loop performance across tasks<ref:2610.

Improvements for AI systems

  1. A decoder-free latent space that ignores visual distractors as well as LeWM does allows for a more robust representation of task dynamics, leading to better policy performance when dealing with visually complex environments. This improved system can achieve higher Task-state ↑ Distractors ↓ Task-state ↑ Distractors metrics in linear probes, indicating superior disentanglement of relevant versus irrelevant visual information.

  2. The introduction of Diffusion steering (DS) as a planning paradigm shifts the decision variable from raw actions to the deterministic noise space of the flowmatching policy, which prevents the exploitation of dynamics errors and significantly improves closed-loop performance during inference. This enables a planner to search strictly within the world model’s imagination, evaluating candidates decoded at imagined states, leading to a performance gain over action-space MPC that typically collapses every model to 3–8% success.

  3. The LeWAM architecture, which is a bidirectional transformer for forward, backward dynamics and policy prediction, allows the system to simultaneously predict future states and infer actions based on past states. This capability enables the model to achieve superior closed-loop evaluation compared to baselines, as evidenced by LeWAM's policy maintaining performance close to the reference BC-only flow-matching policy.

  4. The four distinct training modes—Forward, Backward, Inverse, and Policy—allow for multi-modal learning where the encoder is trained end-to-end directly from task pixels under a strict regularizer, resulting in representations shaped by the dual objectives of accurate dynamics prediction and robust action generation. This comprehensive training regimen maximizes the utility of the model across different control tasks.

  5. The use of Projected DS maps every proposal back onto the sphere, ϵ ← √A ϵ/∥ϵ∥, before decoding, ensuring that every decoded candidate is one the head was trained to produce, which eliminates temperature effects and keeps the search distribution Euclidean. This ensures that planning remains constrained to actions the policy head was explicitly trained to produce.

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.

Sources

Related papers