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

TL;DR

LeWAM (Fu, Siebert, Halicki, Balestriero) is a Joint-Embedding Predictive Architecture (JEPA) trained end-to-end from pixels under two losses — next-latent prediction and action flow matching — with a single SIGReg regularizer forcing the latent state toward an isotropic Gaussian. The result is a 17M-parameter model (one shared ResNet-18 encoder + a Mixture-of-Transformers predictor with separate state and action branches) that serves as both a direct behavior-cloning policy and a planner (sample candidate actions → roll out latents → refine against a goal latent). On goal-conditioned visual planning across ten simulated environments, LeWAM lifts success from 28.6% → 89.7% on contact-rich manipulation and 10.9% → 46.2% on long-horizon tasks vs the predecessor LeWM (CEM) baseline, at 32.3× faster planning (~19 ms/plan). On behavior cloning it matches Diffusion Policy / π0.5 across three real-robot tabletop tasks on a 7-DoF arm.

Key claims

Joint training over next-latent prediction + action flow matching + SIGReg produces one checkpoint that supports both direct policy execution and latent-space planning, trained in parallel on 1 GPU [§Architecture details].
An isotropic-Gaussian constraint on the latent state (SIGReg) is load-bearing for planning — removes an outright failure mode and lifts object-state decodability (effective rank 42.9 vs 10.9 without SIGReg) [§Latent representation probing].
End-to-end joint training beats a frozen DINOv2 CLS encoder for direct control (86.7% vs 79.0% direct-mode success across four matched tasks) — the visual-encoder-is-what-matters position carries a measured counter-example here [§Latent representation probing].
Flow-matching action heads beat MSE action heads when planning gains matter: MSE regresses to a conditional mean, flow matching models multiple action modes, and the planning uplift over direct control is larger in all three reported tasks [§Latent representation probing].
Removing LeWAM’s action head and planning with CEM over its latent dynamics beats LeWM (CEM) by +22.2 pp on Cube (86.0 ± 3.2 vs 63.8 ± 2.2) — the dynamics representation itself is a lever, not just the end-to-end action head [§Isolating action-conditioned dynamics].
Action-prediction supervision acts as an anti-collapse signal: without SIGReg, LeWAM still retains decodable object/proprioception/action information while LeWM collapses to R²≈0 on Push-T probes [§Latent representation probing].
On real 7-DoF hardware with scene + wrist cameras, LeWAM runs within the same 270 GFLOPs/action-chunk budget as Diffusion Policy (vs π0.5 at 3603.6) and reaches 33–36% avg success across Pick-and-Place T / Push-T / Tape Stack (Diffusion Policy 38.9%, π0.5 28.9%) [§Physical-robot control].

Method

LeWAM uses a shared ResNet-18 encoder mapping history, target, and goal images to 192-dim latent states, then a Mixture-of-Transformers (MoT) predictor with two branches: a state predictor f_z that predicts future latents from clean actions and prior states, and an action predictor f_a that outputs a velocity field for flow matching, conditioned on current state, task context, and normalized remaining horizon.

Training jointly optimizes L_LeWAM = L_dyn + L_act + λ_sig · L_SIGReg:

L_dyn supervises predicted latent ẑ_{t+K} against the encoder’s own encoding of the observed future (next-latent regression).
L_act is standard flow-matching regression on action chunks.
L_SIGReg pushes the latent state distribution toward an isotropic Gaussian — borrowed from the LeJEPA / LeWorldModel line.

At inference the model supports two modes: (1) Direct control reads actions straight from the action head; (2) Latent joint imagination samples candidate action chunks, rolls them forward through the frozen dynamics, scores each terminal latent against the goal latent with a proximity penalty, selects the lowest-cost refined candidate, executes the first K steps, replans. The planning loop is 32.3× faster than LeWM’s CEM baseline at ~19 ms/plan.

Results

Visual planning (10 sim envs): contact-rich manipulation 28.6% → 89.7% success; long-horizon tasks 10.9% → 46.2% vs LeWM baseline.
BC in simulation (three Robomimic tasks): LeWAM 64.7% avg vs Diffusion Policy 61.6%; wins Transport 61.3 ± 9.9 vs 45.3 ± 9.0, essentially ties Diffusion Policy on DrawerCleanup and ToolHang.
Physical robot (3 tasks × 30 trials): LeWAM end-to-end 33.3% avg (16/8/6 of 30), LeWAM frozen V-JEPA 2.1 35.6%, Diffusion Policy 38.9% (19/16/0), π0.5 28.9% (14/3/9); LeWAM runs at 270.2 GFLOPs vs π0.5’s 3603.6.
Dynamics-only ablation (CEM over LeWAM latents, action head removed): Cube 86.0 ± 3.2 (vs LeWM 63.8 ± 2.2); Push-T 84.0 ± 1.3 (vs LeWM 92.4 ± 1.2 — LeWM wins this one).
Representation probing: SIGReg + joint action training yields effective rank 42.9 and object R² 0.90 / proprio R² 0.86 / action R² 0.40 (IDM 0.75); LeWM without SIGReg collapses to effective rank 16.2 and R²=0 across all three probes on Push-T.
CEM vs flow-matching sampler: flow matching gives larger planning uplift over direct-mode on all three reported tasks because MSE action regression predicts conditional means that are not representative under multimodal action distributions.

Why it’s interesting

LeWAM is the sharpest filed instance of the LeJEPA / LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels line arguing that end-to-end JEPA WFMs trained from pixels under a two-loss recipe can be competitive policies and planners simultaneously — now extended by making the action channel first-class (flow-matching head, joint training) and keeping the single-GPU 17M-parameter research-budget shape. It directly addresses the central weakness The Obsessed Encoder — Latent-Space Misallocation in JEPA-Style Self-Supervised Models flagged for the pure latent-predictive JEPA pole (predictable-feature collapse under episode-constant goals) by using the action-regression branch as an independent anti-collapse signal — observed-action targets don’t depend on the encoder, so even without SIGReg the model retains decodable object/action information while LeWM alone collapses on Push-T probes. It also sits on the same end-to-end-joint-training side of the WAM recipe board as Causal World Modeling for Robot Control (LingBot-VA) LingBot-VA, but at ~1000× smaller scale and with a planning interface over the latent dynamics (CEM-over-LeWAM-latents already beats LeWM-CEM on Cube by +22.2 pp with the action head removed), suggesting that the latent-dynamics representation itself — not just the action head — is a lever.