Why pixel reconstruction makes a weaker latent
World action models predict both actions and future observations, and most of them shape their internal representation through pixel-level reconstruction. The paper argues that this forces the latent to encode every visual detail a decoder might render, whether or not that detail affects the task. In visually busy scenes, capacity is spent on noisy, task-irrelevant information, which the authors say complicates downstream prediction.
LeWAM takes the opposite route. It builds on decoder-free joint-embedding predictive architectures, which model dynamics purely in latent space and drop the pixel decoder entirely. The paper's stated motivation is that removing the decoder removes the pressure to store appearance details that do not change how the scene evolves. The cost of that choice is that the latent must be shaped by something else: accurate dynamics prediction and robust action generation.
Four prediction modes share one bidirectional transformer
LeWAM is a bidirectional transformer conditioned on actions, with a single backbone handling forward dynamics, backward dynamics, inverse dynamics and policy prediction. The encoder is trained end-to-end directly from task pixels under a strict regularizer, with no pixel decoder, tokenizer or pixel-space loss in the loop. The paper's comparison table places LeWAM as the only listed model that is both decoder-free and equipped with a policy head, forward, inverse and backward dynamics, and a multimodal one-backbone design.
The practical meaning of the four modes is that the same latent supports imagining what comes next, inferring what came before, recovering the action between two states, and proposing actions directly. Because all four objectives train the same representation, the latent is not optimized for reconstruction fidelity but for being useful to prediction and control at once.
Linear probes read task state and ignore distractors
The alignment evidence comes from linear probes: for each model and task, the authors fit a ridge regression from the latent to agent, object, gripper and distractor state, with cross-validation, and evaluate on a held-out set. They report MSE, Pearson r and explained variance R² separately for task-relevant state and for task-irrelevant distractors, averaged over six tasks and five seeds.
On average across tasks, LeWAM reaches R² 0.80 on task-relevant state against 0.62 for LeWM with a diffusion head and 0.78 for the reconstruction-based world action model. Its distractor R² is 0.00, matching LeWM's 0.00 and far below the reconstruction model's 0.73. The pattern is not uniform per task: on Lift, LeWAM's task-state R² of 0.42 is below the reconstruction model's 0.56, while on ToolHang it reaches 0.93 against 0.88. The authors' reading is that LeWAM encodes task-relevant state better than LeWM and even the reconstruction-based models on average, while ignoring distractors as well as LeWM.
A separate check supports the same reading. The authors train a plain feed-forward convolutional decoder on a frozen, already-trained LeWAM, so decoder gradients cannot reshape the latent. In imagined rollouts for PushT and Robomimic Lift, the T block's location is consistently identified but its rotation is blurry, and distractor blocks are almost entirely missing from imagined frames. The paper treats the occasional misplaced distractor as decoder hallucination rather than latent content.
Control matches a same-encoder flow policy
The acting result is a matched comparison: closed-loop evaluations of LeWAM match a regular flow-matching policy trained on the same encoder at matched size, while LeWAM additionally provides a world model. That framing matters because it isolates what the world-model structure costs or buys relative to a policy-only baseline sharing the same representation, rather than comparing against an unrelated system.
The paper reports this as a match, not an improvement, so the claim to carry away is that adding forward, backward and inverse dynamics prediction to the same encoder did not degrade closed-loop control in these evaluations. The evaluations are averaged across tasks and over five training seeds, with per-task curves reported separately in the appendix figures.
Raw-action planning lets MPC exploit model error
Planning with a learned world model invites sampling-based model predictive control, but the paper states that executing such control over raw actions often fails in practice. The mechanism it gives is exploitation: when the planner optimizes directly over action space, it tends to find plans that exploit inevitable inaccuracies in the learned dynamics model, producing trajectories that look good against a hallucinated future and then fail during deployment.
LeWAM's response is to change the decision variable. Instead of searching over actions, the planner searches the deterministic noise space of the policy head. A plan is a sequence of noise vectors; each chunk is decoded at the state imagined after the preceding chunks, the predicted latent is appended while the oldest context latent is dropped, and the plan is scored on its final imagined latent. The search starts from the policy's prior, so the first iteration of every diffusion-steering planner is best-of-N proposal ranking over projected noise.
Projection keeps proposals inside the trained distribution
The policy head was trained to start from draws of a standard normal distribution, which concentrate on a sphere of radius equal to the square root of the action dimension. As a planner narrows its search, proposals drift toward a mean of much smaller norm, which for a rectified-flow head behaves like sampling at reduced temperature and makes the search mode-seeking.
Projected diffusion steering maps every proposal back onto that sphere before decoding. The search distribution itself stays Euclidean; only the proposals are projected. The paper's stated consequences are that the search stays over directions, every decoded candidate is one the head was trained to produce, and the temperature effect is given up. The paper reports this projection as ablated in its Table 4, and the three diffusion-steering planners it compares — DS-CEM, DS-MPPI and DS-CMA-ES — share one budget of 64 samples across 3 iterations at horizon 5, differing only in their update rule.
The noise-space decision variable is the transferable result
The planning finding is stated as a comparison within the paper: sampling raw actions when planning with world action models lets MPC exploit dynamics-model inaccuracies, whereas planning in the noise space of the policy head improves the closed-loop performance of these world action models. The procedure is inference-only, so it changes how an already-trained model is used rather than retraining it.
The authors name two next steps rather than claiming completion: applying LeWAM to real-world data, which they describe as containing an abundance of dynamically inert features, and improving planning algorithms, which they call slow and dependent on a notion of a goal at inference time. Those two conditions mark where the current evidence stops. The most decision-relevant takeaway is therefore a priority ordering: the latent result is an average-over-tasks claim with per-task exceptions such as Lift, the control result is a match rather than a gain, and the clearest actionable change is to plan in the policy head's noise space with projection rather than over raw actions.
Sources used in this answer
LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC
A decoder-free JEPA world action model trained through forward, backward, inverse dynamics and policy prediction reads task-relevant state better than LeWM and a reconstruction-based world action model on average, ignores distractors, matches a same-encoder flow-matching policy in closed-loop control, and improves closed-loop performance when planning in the policy head's noise space instead of raw action space.