Source-linked AI summary
LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC
Shashank Hegde, Alexander Popov, Elie Aljalbout, Nikolai Smolyanskiy
TL;DR
Reconstruction-based world models and WAMs encode visual details regardless of task relevance, while existing JEPA approaches typically separate dynamics and policy. LeWAM trains a decoder-free latent model across four modes; it reports alignment and acting results, and improved closed-loop planning with noise-space steering.
Problem
Reconstruction-based world models encode visual details regardless of task relevance, while JEPA approaches typically lack the unified multi-modal structure of WAMs.
Source: §1 · p. 1Method
LeWAM uses a decoder-free JEPA latent and a bidirectional transformer trained for forward, backward, inverse-dynamics, and policy prediction.
Source: §1 · pp. 1–2Results
LeWAM captures dynamic states while ignoring distractors, matches standard behavioral-cloning policy performance, and noise-space steering improves closed-loop planning.
Source: §1 · p. 2Takeaways & Limitations
LeWAM provides a world model alongside a policy, and inference-time noise-space planning improves average closed-loop performance across tasks.
Source: §1 · p. 2, §6 · p. 9The authors identify planning speed and the need for a goal at inference time as limitations of their planning algorithms.
Source: §6 · p. 9
Abstract
from arXiv · showhide
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.
1 Introduction
Reconstruction-based world models and WAMs encode visual details that may be irrelevant to task dynamics, while JEPA approaches often lack a unified policy-and-dynamics structure. LeWAM combines four prediction modes in a decoder-free model, with reported benefits for alignment, acting, and planning.
1 Introduction
- LeWAM unifies forward, backward, inverse-dynamics, and policy prediction in a reconstruction-free world action model.Its representations are shaped by dynamics-prediction and action-generation objectives rather than pixel-space losses.Source: §1 · pp. 1–2
- Evaluations use human-generated demonstration datasets augmented with visual distractors.Source: §1 · p. 2
- LeWAM’s latent captures dynamic robot and object states as effectively as reconstruction-based baselines while ignoring visual distractors.Source: §1 · p. 2
- LeWAM’s integrated flow-matching policy matches standard behavioral-cloning policies while also providing a world model for dynamics prediction.Source: §1 · p. 2
- Steering planning in policy noise space improves closed-loop performance during inference and prevents exploitation of dynamics errors.The planner searches noise inputs decoded by the policy rather than raw actions.Source: §1 · p. 2
2 Related work
Related work addresses task-irrelevant visual information in world models and connects learned dynamics to action generation. LeWAM’s approach sits among decoder-free latent models, multi-function WAMs, and methods for steering generative policies.
2 Related work
- Reconstruction-based world models shape latents through image reconstruction, while decoder-free methods resist task-irrelevant appearance and some JEPA models plan using latent goals.LeWAM adds an action head and multi-mode structure to LeWM’s latent; LeWM itself has no policy head.Source: §2 · p. 2
- WAMs couple future-state prediction with action generation, with prior approaches using video generation, separate video and action diffusion, or pretrained video priors.Source: §2 · p. 3
- Policy steering methods include guidance, candidate ranking, action-space refinement, and optimization of the noise that generates an action.Diffusion steering searches at each decision within world-model imagination, scoring decoded action chunks by goal-latent distance without search-time environment interaction or a learned value function.Source: §2 · p. 3
3 LeWAM
LeWAM uses one decoder-free latent representation and a bidirectional transformer to support four masked training modes. Its training combines dynamics prediction and action generation, with rectified flow generating action chunks.
3 LeWAM
- A training window contains four frames and three intervening five-step action chunks; the first three frames provide context and the last frame is the prediction target.The final action chunk is the policy output.Source: §3 · p. 3
- Masking inputs makes the backbone perform forward, backward, inverse-dynamics, or joint policy-and-future prediction.In policy mode, past action tokens are replaced with an unknown-action token, preventing the policy from copying recent demonstrator motion.Source: §3 · p. 4
- Every training batch runs all four modes; forward and backward prediction with SIGReg impose JEPA structure, while inverse and policy modes make the latent more task-aligned.Source: §3 · p. 4
- Rectified flow trains the backbone to predict noisy-action velocity and generates test-time action chunks by integrating Gaussian noise over eight Euler steps.The integration is a deterministic mapping from noise and the current latent to an action chunk.Source: §3 · p. 4
- LeWAM has no pixel decoder, tokenizer, or pixel- or token-space loss; its latent is shaped by dynamics prediction and action generation.A shared ViT-tiny encoder produces one latent per frame, and a bidirectional transformer serves the four masked modes.Source: §3 · p. 4, Figure 2 · p. 4
4 Diffusion steering: planning in the policy’s noise space
The section shifts sampling-based MPC from raw actions to a flow-matching policy’s noise space, addressing action-space plans that exploit dynamics errors. It also projects proposals onto the sampler’s typical set.
4.1 Goal-conditioned planning with a latent distance
- Goal-conditioned planning scores action chunks using the forward mode and a demonstration goal latent, then minimizes the score over H chunks with sampling-based MPC.The planner proposes S plans, updates its distribution over I iterations, executes the first chunk, and replans.Source: §4.1 · p. 5
4.2 Motivation: Failure of Action-Space MPC with LeWM
- On the smaller human-sourced datasets, action-space sampling-based MPC fails with LeWM as rollout errors lead it to plan for hallucinated futures.The authors hypothesize that the dynamics model is imperfect with less data and compounds errors during planning rollouts.Source: §4.2 · p. 5
4.3 Diffusion steering
- Diffusion steering optimizes policy noise rather than raw actions, decoding each chunk at the imagined state and scoring plans by their final latent.The search starts from the policy prior; its first iteration is best-of-N proposal ranking over projected noise.Source: §4.3 · p. 5
4.4 Projection onto the sampler’s typical set
- Projected DS maps proposals onto the sphere where the sampler was trained, so candidates remain in its trained set but the search gives up the temperature effect.The search distribution remains Euclidean; only proposals are projected.Source: §4.4 · p. 5
5 Experiments and Results
Experiments evaluate LeWAM on six visual-distractor tasks using latent probes, decoded rollouts, closed-loop control and planning, and mode and projection ablations. LeWAM’s policy matches the reference, noise-space planning improves performance, and the ablations report benefits from the modes and projection step.
5.1 Tasks and data.
- The six-task evaluation uses visual distractors, and models train on pixels alone with distractors present at evaluation.PushT uses action-independent moving distractors; robomimic scenes add wobbling boxes that do not affect dynamics.Source: §5.1 · pp. 5–6
5.3 Latent Space analysis
- Across tasks, LeWAM encodes task-relevant state better than LeWM and reconstruction-based models, while ignoring distractors as well as LeWM.The linear probes use ridge regression from latents to task and distractor state; the table reports metrics across all tasks and five seeds.Source: §5.3 · pp. 6–7, Table 2 · p. 7
- Decoded rollouts locate PushT’s T block but blur its rotation, while distractors are mostly absent or misplaced in PushT and Lift.The authors suggest the misplaced Lift distractor may be a decoder hallucination.Source: §5.3 · p. 7, Figure 4 · p. 8
5.4 Closed-loop evaluation.
- LeWAM’s policy maintains the reference’s performance, while noise-space sampling improves LeWAM and RecWAM; action-space MPC yields only 3–8% success for every model.The closed-loop evaluation covers 250 rollouts across six tasks with visual distractors.Source: Figure 5 · p. 8, §5.4 · p. 8
5.5 Ablation
- Removing modes generally reduces latent-state and closed-loop performance; forward and backward modes consistently help, while inverse is neutral on Square.Table 3 ablates one mode at a time from the four-mode model.Source: Table 3 · p. 9, §5.5 · p. 9
- Removing the projection step lowers DS-CEM success on Lift, Square, and Can by about 5, 12, and 56 points, respectively.The projection ablation keeps the budget and goals unchanged.Source: Table 4 · p. 9, §5.5 · p. 9
6 Conclusion and Future work
LeWAM combines a decoder-free JEPA world action model with diffusion-noise-space planning. The paper reports stronger latent alignment and improved average closed-loop performance, while identifying real-world data and faster, goal-independent planning as next steps.
6 Conclusion and Future work
- LeWAM’s latent captures dynamic robot and object states while ignoring visual distractors, and noise-space planning improves average closed-loop performance across tasks.The authors identify real-world data and faster planning that does not require an inference-time goal as next steps.Source: §6 · p. 9
A Results per Task
Figure 6 reports per-task closed-loop evaluation and MPC results. It averages performance over five training seeds and 50 evaluation seeds, yielding 250 unique rollouts.
A Results per Task
- The per-task results complement Figure 5, which reports closed-loop evaluation averaged across tasks.Figure 6 includes both closed-loop evaluation and MPC results.Source: A Results per Task · p. 12
- Figure 6 reports average closed-loop performance for all models and MPC on each task.The figure covers 250 unique rollouts: five training seeds and 50 evaluation seeds.Source: Figure 6 · p. 12, A Results per Task · p. 12
- The figure averages results over five training seeds and 50 evaluation seeds per task.This corresponds to 250 unique rollouts.Source: Figure 6 · p. 12
B Configuration of the tasks
The appendix specifies camera views and resolutions for the tasks, and gives one LeWAM configuration for five robomimic tasks. Training context is randomly truncated rather than masked, with the newest frame always included.
B Configuration of the tasks
- Transport uses stacked shoulder-camera views at 256 px, while the other tasks use task-specific rendered views and resolutions before resizing frames to 224 px.PushT is rendered at 96 px on a 512 px table; Lift, Can, and Square at 84 px; and ToolHang at 168 px from the side camera.Source: B Configuration of the tasks · pp. 12–13
- Training context is randomly truncated to one to three left-aligned frames, always retaining the newest frame, rather than masked.This supervises the first episode decision, when the simulator repeats the reset frame.Source: B Configuration of the tasks · p. 13
- Table 5 gives one LeWAM configuration for Lift, Can, Square, ToolHang, and Transport.Evaluation uses 50 seeded simulator resets per seed, with moving boxes.Source: Table 5 · p. 13, B Configuration of the tasks · p. 13
C Model sizes
The model-size comparison matches four models in trainable parameters, with LeWAM, LeWM, RecWAM, and Dreamer each near 41–43M. Their architectures and training setups differ, including RecWAM’s additional frozen tokenizer.
C Model sizes
- LeWAM has 41.0M trainable parameters, LeWM 42.2M, RecWAM 42.6M, and Dreamer 43.4M.RecWAM also uses a pretrained, frozen 3.9M VAE tokenizer.Source: C Model sizes · p. 13
- LeWAM uses a ViT-tiny/14 encoder and a width-512 bidirectional transformer, while LeWM uses a widened width-512 autoregressive predictor.LeWAM’s transformer has depth 8; LeWM’s has depth 6.Source: C Model sizes · p. 13
- Dreamer’s 43.4M RSSM is trained by image reconstruction at 96 px with balanced KL and enlarged components.Its deterministic state is widened from 512 to 1536 and its stochastic state from 32 to 64.Source: C Model sizes · p. 13
D Protocol and training hyperparameters
The appendix specifies state-probe evaluation and matched planning protocols, including model-specific goal retrieval and robomimic horizons.
D Protocol and training hyperparameters
- Probes predict robot and object states plus distractor positions from latent representations, using held-out-episode R2, five-fold penalty selection, and a shuffled-target control.PushT probes agent and block location and block angle; robomimic probes end-effector position, gripper opening, task-object position, and box positions.Source: D Protocol and training hyperparameters · p. 13
- Planners share 64 samples, 3 iterations, and 8 elites; action-space CEM, MPPI, and CMA-ES search sequences of 5 action chunks.Dreamer ranks its policy chunk alongside 63 Gaussian perturbations using its decoded-rollout pixel error to the goal frame; the policy chunk is always included, and Dreamer has no DS variant.Source: D Protocol and training hyperparameters · p. 13
- Planning goals come from each model’s own latent space, using the observation 25 steps after the closest latent in a demonstration.Source: D Protocol and training hyperparameters · p. 13
- Robomimic horizons are 100 steps for Lift, 400 for Can and Square, and 700 for ToolHang and Transport; LeWAM uses Table 5’s configuration across tasks.Source: D Protocol and training hyperparameters · p. 14