LeWAM: End-to-End World and Action Modeling with JEPAs

AI summary

Overview

Research area: Robotics and visual control — specifically latent world models (Joint-Embedding Predictive Architectures, or JEPAs) that learn action-conditioned dynamics for planning and policy learning.

Technical level: Advanced. The paper assumes familiarity with self-supervised representation learning, flow matching, transformer attention masking, and model-predictive control with gradient-based trajectory optimization.

Scope: The paper introduces LeWAM, a single end-to-end JEPA that jointly trains a visual encoder, action-conditioned latent dynamics, and a stochastic flow-matching policy, and evaluates it on ten simulated environments plus three physical-robot tasks.

What This Paper Is About

Existing JEPA-based world models learn visual representations and action-conditioned latent dynamics first, and only afterward bolt on a control mechanism — either an online sampler like the cross-entropy method (CEM) or a separately trained action predictor on frozen features. This splits action supervision from representation learning and makes planning expensive, especially for contact-rich manipulation and long horizons. LeWAM's goal is to train representation, dynamics, and policy together in one stage, so that a single small model can both execute actions directly and plan efficiently by refining actions through its own learned latent dynamics.

Key Contributions

1. A single jointly trained JEPA for control. To the authors' knowledge, this is the first JEPA that jointly learns the visual representation, action-conditioned latent dynamics, and a stochastic policy in one training stage, using a joint prediction objective, one weight for latent-space regularization, and approximately 17M parameters.

2. A dual-branch Mixture-of-Transformers predictor. A shared masked self-attention is used by two branches with separate query/key/value and feedforward projections: a state branch that predicts the next latent state from clean actions, and an action branch that performs flow matching on noisy actions conditioned on goal information. Both objectives therefore shape the shared encoder.

3. A planning procedure that samples then refines. Candidate action sequences are drawn from the learned action distribution by alternating action sampling with latent rollouts, then refined by gradient descent through the frozen learned dynamics on a latent cost-to-goal, with the lowest-cost iterate selected.

4. Broad empirical evaluation. Ten heterogeneous simulated environments for visual planning, horizon sweeps on four tasks, behavior-cloning comparisons in simulation, and closed-loop evaluation on three physical-robot tasks against Diffusion Policy and π0.5.

Main Findings

• Large planning gains over prior JEPAs. Across the ten environments, LeWAM averages 93.5% goal-reaching success versus 50.9% for LeWM (CEM), a 42.7 percentage-point gap, exceeding that baseline on every task. The abstract summarizes the improvement as 28.6% to 89.7% on contact-rich tasks.

• Contact-rich manipulation is where the gap is widest. ToolHang reaches 94.7% versus 35.4%, and Transport reaches 83.3% versus 2.8%.

• On a logically structured task, it also beats a goal-conditioned behavior-cloning baseline. Puzzle reaches 80.7% versus 44.0% for goal-conditioned GCBC.

• Long horizons degrade more gracefully. At a 100-step goal offset on Push-T, Cube, Scene, and Puzzle, LeWAM achieves 46.2% mean success versus 10.9% for LeWM.

• Planning is much faster. Across four visual-goal tasks, LeWAM averages 52.50 plans/s versus 1.62 plans/s for LeWM (CEM), a 32.3× speedup.

• SIGReg matters both for anti-collapse and for planning quality. Without it, the latent-only LeWM baseline scores R² = 0.00 on every reported Push-T probe, while LeWAM with either action head retains decodable state and action information. On Push-T the planning gain attributable to SIGReg is 14.7 percentage points, compared with 2.7 without it; removing SIGReg lowers planning success by 25.3 points on Push-T, 5.3 on Puzzle, and 4.0 on Scene.

• Action supervision improves what the latents contain. Adding the action-modeling objective raises average linear-probe R² for proprioception, action, and IDM by +0.21, +0.20, and +0.19 respectively relative to LeWM.

• End-to-end encoder training beats frozen features. Direct success averages 86.7% with the jointly trained encoder, versus 79.0% with frozen DINOv2 CLS features and 75.8% with frozen patch features.

• Flow matching beats a deterministic MSE action head. Without SIGReg, flow matching retains more linearly decodable object information (R² = 0.88 versus 0.66), with comparable proprioception (0.85 versus 0.86); a two-door navigation experiment illustrates the multi-modal action distribution the flow head captures.

• Dynamics remain usable in isolation. Removing the action head at test time and planning with CEM over LeWAM's dynamics gives comparable visual predictions and competitive Push-T success versus LeWM, with clearer control-relevant predictions and higher success on Cube.

• Behavior cloning is competitive. On goal-agnostic simulated tasks (DrawerCleanup, Transport, ToolHang), LeWAM is competitive with Diffusion Policy on two tasks and improves success by 16 percentage points on Transport.

• Physical-robot results are mixed but competitive. Out of 30 trials each: Diffusion Policy scores 19/30 on Pick and Place T, 16/30 on Push-T, 0/30 on Tape Stack (38.9% average); π0.5 scores 14/30, 3/30, 9/30 (28.9% average); LeWAM with a frozen V-JEPA 2.1 encoder scores 15/30, 12/30, 5/30 (35.6%); LeWAM from scratch scores 16/30, 8/30, 6/30 (33.3%). One LeWAM action chunk costs 270.2 GFLOPs versus 3603.6 GFLOPs for π0.5, which the authors express as a 13.3× difference in action chunks per TFLOP.

Methodology in Plain English

The model takes offline trajectories of pixels and actions — no rewards — and learns three things at once.

Encoding. A ResNet-18 encoder with spatial softmax pooling turns each frame into a 192-dimensional latent vector. The same encoder processes history frames, future target frames, and goal images. A distributional regularizer (SIGReg, built on the Epps–Pulley test statistic over random projections) pushes the batch of latents toward an isotropic Gaussian, which prevents collapse and, the authors argue, keeps the latent geometry meaningful for measuring distance to a goal.

Two branches, one attention. A small Mixture-of-Transformers predictor (width 192, 4 blocks, 4 heads) runs two heads over a shared, carefully masked attention. The state branch takes a learned query token plus history latents and the five clean recorded actions, and predicts the latent state five steps ahead. The action branch takes noisy action tokens and predicts the flow-matching velocity field that transports Gaussian noise to a real ten-action chunk. The attention mask blocks clean action tokens from leaking into the action branch, so the two predictions are genuinely separate tasks that both backpropagate into the encoder. The loss is simply next-latent-state MSE plus flow-matching MSE plus a weighted SIGReg term (λ = 0.04).

Control, two ways. For direct control, integrate the learned velocity field from noise to a full action chunk and execute the first five actions. For planning, sample many candidate plans by repeatedly drawing action chunks and rolling the latent dynamics forward, then improve each candidate by gradient descent on the MSE between the predicted final latent and the goal latent (with a small penalty for drifting far from the sample), and execute the best one — a receding-horizon loop.

Evaluation. Planning is compared against LeWM with CEM, LeWM with a post-trained inverse dynamics model, DINO-WM, and GCBC across ten simulated tasks. Separate experiments ablate SIGReg, the frozen-vs-end-to-end encoder, and the flow-vs-MSE action head; linear probes test what physical state and action information the latents encode; and the same architecture minus goal conditioning is tested as a plain behavior-cloning policy in simulation and on a real robot.

Why This Matters

Research impact. The paper argues that action supervision and world modeling should not be staged sequentially. It shows empirically that a policy objective acts as an anti-collapse signal for JEPA representations while still improving control-relevant dynamics, and it offers a planning recipe that is both more accurate and orders of magnitude cheaper than population-based search on the same latent model. It also provides a probing-based account of why the gains appear (better object-state decodability, sharper latent distance-to-goal) rather than reporting success rates alone.

Real-world applications:

• Contact-rich manipulation such as tool hanging, transport, and drawer cleanup, where dynamics switch abruptly between contact modes and local optimizers struggle.
• Real-time robot control on modest compute, where a 270.2 GFLOP action chunk versus 3603.6 GFLOPs for a vision-language-action policy matters for onboard deployment.
• Long-horizon goal-image tasks such as pick-and-place, pushing, and tape stacking specified by a picture rather than a language prompt.
• Behavior cloning from teleoperation demonstrations, where no goal image is available and the policy must infer progress and react to object pose changes itself.

Industry relevance. The result that a ~17M-parameter model trained from scratch on task-specific demonstrations is competitive with a much larger vision-language-action policy on three physical tasks, at roughly a thirteenth of the inference cost, is directly relevant to teams deploying manipulation policies under latency and compute budgets. The finding that frozen pretrained features (DINOv2, V-JEPA 2.1) are not automatically better than end-to-end training is also practically useful for anyone choosing a backbone.

Future Directions

• Data coverage. The task suite mixes expert demonstrations with non-expert interactions at limited scale; broader play data and exploration on contact-rich tasks would cover difficult contact transitions and a wider range of behaviors.
• Generalization. Tests on unfamiliar objects, camera viewpoints, and contact states would probe whether learned action-conditioned dynamics generalize better than policy-only approaches.
• Multi-task and multi-embodiment learning. Models are currently trained or adapted separately per task and per embodiment; multitask training across robots and language-conditioned tasks is posed as an extension.
• Hierarchical and scaled-up control. LeWAM is positioned as a low-level planner/controller that a high-level planner could drive via subgoals, and both training scale (larger models and datasets) and test-time search scale (larger candidate populations, more refinement steps) are flagged as untested levers.

Target Audience

Robotics and embodied-AI researchers working on world models, model-based control, and imitation learning; self-supervised learning researchers interested in JEPAs and in how auxiliary prediction objectives shape representations; and practitioners building manipulation policies who care about inference cost and closed-loop performance on real hardware. Readers without background in flow matching, transformer attention masking, or trajectory optimization will find the methodology sections dense and should start with the introduction and the findings in Section 4.

Authors’ abstract

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.