How does LeWAM combine JEPAs for end-to-end world and action modeling?

One encoder serves prediction and action

Earlier JEPA pipelines build a visual representation and action-conditioned dynamics first, then attach a control policy or an online solver. LeWAM instead trains representation, latent dynamics, and a stochastic policy in one stage. Each training sample supplies clean actions for next-latent-state prediction and noisy actions for flow matching, and a Mixture-of-Transformers predictor routes the two token groups through separate parameters while one masked self-attention spans the whole sequence. Because both objectives backpropagate into the same ResNet-18-style encoder, action supervision is not disconnected from visual representation learning.

The scale is modest: the visual-planning model has roughly 16.9M trainable parameters, and the paper describes it as approximately 17M. The predictor uses a 192-dimensional latent state, four blocks, and four attention heads. The dynamics branch sees five clean actions to predict the latent five raw steps ahead, while the flow policy models ten actions. Only frames are subsampled, so action indices stay at the raw environment rate.

A Gaussian latent constraint keeps planning informative

The second ingredient is SIGReg, a regularizer that constrains the latent state distribution and prevents representation collapse. The paper reports that this constraint is what makes the latent space useful for planning, not just for avoiding collapse. On Push-T, adding SIGReg produces a 14.7 percentage-point planning gain, compared with 2.7 points without it. Removing SIGReg from the full model lowers planning success most strongly on Push-T (25.3 points), then Puzzle (5.3) and Scene (4.0).

Linear probes give a concrete view of what the representation retains. On Push-T, LeWAM with SIGReg reaches R² of 0.90 for object state, 0.86 for proprioception, and 0.40 for action, or 0.75 when the probe uses both the current and future latent. The same model without SIGReg keeps object and proprioception scores near 0.88 and 0.85 but drops action decodability to 0.38. A latent-only LeWM without SIGReg scores 0.00 on every reported Push-T probe, which the authors read as evidence that observed actions act as an anti-collapse signal anchored to targets independent of the encoder.

Planning gains concentrate on contact-rich tasks

Across ten simulated navigation and manipulation environments, LeWAM averages 93.5% goal-reaching success versus 50.9% for LeWM with a cross-entropy-method planner, a 42.7 percentage-point gap, and it exceeds that baseline on every task. The largest gaps appear where contact transitions make local action optimization hard: ToolHang reaches 94.7% versus 35.4%, and Transport reaches 83.3% versus 2.8%. These are goal-conditioned visual-planning results with a goal image and a fixed step budget, not task-completion results.

The abstract's headline pair uses a different aggregation. The 28.6% to 89.7% contact-rich comparison and the 10.9% to 46.2% long-horizon comparison are separate summaries from the ten-environment sweep and the horizon sweep, so they should not be read as the same average as the 93.5% versus 50.9% figure. In the long-horizon sweep, goals are moved farther along each recorded trajectory while the starting frame stays fixed on Push-T, Cube, Scene, and Puzzle; at a 100-step goal offset, LeWAM reaches 46.2% mean success across those four tasks versus 10.9% for LeWM.

A learned action prior narrows the search

Speed comes from two mechanisms the paper names: an expressive distribution prior that narrows the search space, and gradient-based optimization that uses GPU parallelism. Across four visual-goal tasks, LeWAM averages 52.50 plans per second versus 1.62 for LeWM with CEM, a 32.3-fold speedup. During refinement, model parameters stay frozen and only the normalized action plans receive gradients; a quadratic proximity penalty keeps updates near each sampled plan, while candidate retention and final selection use the unpenalized goal cost.

The action head matters for more than speed. With the same latent-dynamics objective, flow matching yields larger planning gains than an MSE head whose point predictions are perturbed with Gaussian noise. Without SIGReg, flow matching also retains more linearly decodable object information (R² 0.88 versus 0.66), with comparable proprioception scores (0.85 versus 0.86). The authors explain this by noting that a deterministic MSE target fits the conditional mean action, whereas flow matching models the conditional action distribution and can preserve distinctions between states whose action distributions differ despite similar means.

Joint training beats frozen visual features

Frozen DINOv2 features underperform end-to-end features in this setup. Across four tasks, direct success averages 86.7% with the jointly trained encoder, versus 79.0% with frozen CLS features and 75.8% with frozen patch features. The paper separates the two contributions: joint training improves the representation used for action selection, and the additional gain from planning shows that the learned dynamics also support action refinement.

A separate ablation isolates the dynamics. Removing the action prediction head at evaluation time and planning with CEM over only the action-conditioned latent dynamics gives visual predictions comparable to LeWM and competitive planning success on Push-T. On Cube, a 3D manipulation task, LeWAM shows clearer predictions of control-relevant features such as the gripper and object, and higher planning success. The authors conclude that joint action learning is compatible with world modeling and can improve the control-relevant parts.

Direct execution extends the same model

LeWAM also runs without goal context, sampling action chunks from observation history and bypassing test-time planning. In simulation, task completion on DrawerCleanup, Transport, and ToolHang is competitive with Diffusion Policy on two tasks and improves success by 16 percentage points on Transport. These are behavior-cloning results from demonstrations, a different setting from the goal-image planning comparisons, so the two success numbers answer different questions.

On physical robots, the paper evaluates Pick and Place T, Push-T, and Tape Stack through direct policy execution, comparing end-to-end LeWAM, a variant using frozen V-JEPA 2.1 features, Diffusion Policy, and the vision-language-action policy π0.5. One full LeWAM action chunk uses 270.2 GFLOPs versus 3603.6 GFLOPs for π0.5, which the paper expresses as a 13.3-fold difference in operation-count efficiency. The reported physical-robot metric is successful trials out of 30 per task.

Where the evidence stops

The task suite combines expert demonstrations and non-expert interactions at limited scale, and the authors note that broader play data and exploration on contact-rich tasks would extend coverage of difficult contact transitions. Models are trained or adapted separately for each task and embodiment, so the shared dynamics and policy have not been tested across robots or language-conditioned tasks. Generalization tests on unfamiliar objects, camera viewpoints, and contact states are named as future work rather than reported results.

The planning model is also small at roughly 17M parameters, and the paper frames larger models, larger datasets, larger candidate populations, and more refinement steps as untested extensions. Physical deployment requires independent safety mechanisms. These boundaries qualify the scale of the claims rather than the direction of the comparisons within the tested environments.

What to watch in the comparison

The most decision-relevant split in this paper is between planning and direct execution. The large margins over LeWM appear in goal-conditioned visual planning, especially on contact-rich ToolHang and Transport and at long goal offsets, where the learned action distribution and gradient refinement replace broad CEM sampling. The behavior-cloning results are described as competitive rather than dominant, with a 16-point gain only on Transport.

A reader comparing LeWAM with prior JEPAs should therefore check which setting a reported number comes from, which baseline and planner it uses, and what goal offset or step budget applies. The paper's own priority is clear: joint training of latent dynamics and a stochastic action policy is what makes one model usable for both fast planning and direct control, and the evidence for that priority is strongest where contact and horizon make search hardest.

References

Minghao Fu, Tavis Siebert, Eryk Halicki, Randall Balestriero. LeWAM: End-to-End World and Action Modeling with JEPAs. arXiv, 2026. arXiv:2610.12818.

Sources used in this answer