World models stop training a controller separately: LeWAM folds action prediction into one pass, lifting contact-rich success from 28.6% to 89.7%

The world model and the controller are no longer trained in two stages. Earlier JEPA pipelines learned visual representations and action-conditioned dynamics first, then added a control policy or an online solver. Prior JEPAs either built on pretrained encoders or learned representations from scratch with distributional regularizers alone, leaving action supervision disconnected from visual representation learning. Across ten heterogeneous simulated environments, LeWAM averages 93.5% success versus 50.9% for LeWM (CEM); on contact-rich ToolHang it reaches 94.7 versus 35.4, and on Transport 83.3 versus 2.8.

One predictor predicts both the next latent state and the action. Clean actions feed the dynamics branch, noisy actions feed the flow-matching branch, and the two branches share one masked self-attention. Flow matching integrates an action chunk from Gaussian noise along a learned velocity field, so the policy itself proposes candidate actions instead of sampling blindly from a broad Gaussian prior. Planning samples N candidate plans from the learned action distribution, refines them by gradient descent through frozen dynamics against a goal cost, and keeps the lowest-cost iterate; across four visual-goal tasks it averages 52.50 plans/s versus 1.62 plans/s for LeWM (CEM).

The latent state is constrained toward an isotropic Gaussian to prevent representation collapse. Without that constraint, latent-only LeWM yields R2 = 0.00 on every reported Push-T probe. Action prediction alone is not enough: on Push-T the planning gain is 14.7 percentage points with SIGReg versus 2.7 without it. Linear probes show joint training raises average proprioception, action, and inverse-dynamics R2 by +0.21, +0.20, and +0.19; the end-to-end encoder reaches 86.7% direct success versus 79.0% for frozen DINOv2 CLS and 75.8% for frozen patch features.

The report is for researchers and engineers working on robot learning, world models, and model-based planning. The method is evaluated on ten simulated navigation and manipulation environments, four of which adapt LeWM benchmarks while six add contact-rich manipulation and navigation; there are also three simulated task-completion experiments and three physical-robot tasks on a reBot Arm 102 teleoperated to a reBot B601-DM, with 30 trials per task. Visual planning assumes access to a goal image, while task completion provides none. The authors propose extending to broader play data, unfamiliar objects and viewpoints, multitask and multi-embodiment training, and integration with a high-level planner.

The authors list several limitations: the task suite combines expert demonstrations and non-expert interactions at limited scale, so coverage of difficult contact transitions is incomplete; generalization to unfamiliar objects, camera viewpoints, and contact states is untested; models are trained or adapted separately for each task and embodiment; LeWAM provides only low-level planning and control; and the visual-planning model has about 17M parameters, with larger models, datasets, candidate populations, and more refinement steps left as future work. The authors also note that the planning problem in 3D contact-rich manipulation is well known and their solution is heuristic, with rigorous verification left as future work.