Latent Dynamics · 2026-10-12
LeWAM: End-to-End World and Action Modeling with JEPAs
Minghao Fu, Tavis Siebert, Eryk Halicki, Randall Balestriero
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 an end-to-end joint-embedding predictive architecture for control that learns (i) next latent-state prediction and (ii) action flow matching with an isotropic Gaussian latent constraint (SIGReg). A shared Mixture-of-Transformers predicts latent dynamics and learns a stochastic action policy, enabling direct policy execution and faster visual planning via gradient-based refinement through learned latent dynamics.
Why it matters
The authors argue prior JEPA control pipelines typically train representations and action-conditioned dynamics before adding control, and often rely on online solvers like CEM or separate action predictors. LeWAM jointly trains representation, latent dynamics, and stochastic policy in one stage and reports large planning success improvements on contact-rich and long-horizon tasks, plus planning speedups using its learned action proposals.
Method
- End-to-end JEPA training from reward-free offline pixel trajectories: shared encoder + Mixture-of-Transformers with two objectives—latent next-state prediction (MSE) and action flow matching (flow matching) plus SIGReg isotropic Gaussian latent constraint.
- Action-conditional latent dynamics and a stochastic policy: the MoT action branch uses clean actions for dynamics prediction and noisy actions conditioned on goal for flow matching/policy learning.
- Test-time control: direct control by sampling actions from the learned policy, and planning by sampling candidate action sequences from the learned distribution and refining them with gradient descent through frozen latent dynamics to minimize a goal-cost.
Limitation
Though a rigorous verification is beyond the paper, the authors state their solution to ill-posed action solving in contact-rich manipulation is heuristic and consider rigorous verification future work.
Abstract (from arXiv)
Joint-embedding predictive architectures (JEPAs) learn compact latent world models without pixel reconstruction. Existing JEPA pipelines typically train visual representations and action-conditioned dynamics before adding a control policy or solver. We introduce LeWAM, an end-to-end JEPA for control with: (1) a joint training objective over next latent state prediction and action flow matching, and (2) an isotropic Gaussian constraint on the latent state, enabling a single model to support direct policy execution and efficient planning by evaluating sampled actions with its learned latent dynamics. Compared with prior JEPAs in visual planning, LeWAM improves success rate from 28.6% to 89.7% on contact-rich tasks and from 10.9% to 46.2% on long-horizon tasks, and plans 32.3 times faster. On behavior cloning tasks in simulation and physical robots, LeWAM is also competitive with established methods. Interpretability analyses show that LeWAM learns control-relevant representations and action distributions effectively guiding test-time planning and control.
Related papers
- LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC
- Hamiltonian JEPA: Action-Conditioned World Models with an Inherited Control State
- H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning
- LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
- V-JEPA Policy: Building Effective World-Action Models on Predictive Visual Latents
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-12 UTC
LeWAM: End-to-End World and Action Modeling with JEPAs
Preprint posted: 2026-10-08 · Original record ↗
- arxiv
- 2610.12818
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
Joint-embedding predictive architectures (JEPAs) learn compact latent world models without pixel reconstruction. Existing JEPA pipelines typically train visual representations and action-conditioned dynamics before adding a control policy or solver. We introduce LeWAM, an end-to-end JEPA for control with: (1) a joint training objective over next latent state prediction and action flow matching, and (2) an isotropic Gaussian constraint on the latent state, enabling a single model to support direct policy execution and efficient planning by evaluating sampled actions with its learned latent dynamics. Compared with prior JEPAs in visual planning, LeWAM improves success rate from 28.6% to 89.7% on contact-rich tasks and from 10.9% to 46.2% on long-horizon tasks, and plans 32.3 times faster. On behavior cloning tasks in simulation and physical robots, LeWAM is also competitive with established methods. Interpretability analyses show that LeWAM learns control-relevant representations and action distributions effectively guiding test-time planning and control.