Plan With a World Model: JEPA + CEM From Scratch (MuJoCo)

A world model you can plan with, from scratch: the frozen JEPA encoder from Episode 3, plus an action-conditioned predictor that guesses the next embedding from the current one and a move. Then planning with the cross-entropy method: imagine 64 plans of 5 moves, score each by the distance to the embedding of a goal picture, keep the best 8, refit, repeat, and make only the first move (model-predictive control), the same loop V-JEPA 2-AC runs on real arms at a vastly larger scale. We test it honestly on 100 unseen MuJoCo scenes against random moves and the Episode 2 VLA. Within 2 cm of the goal (the line we chose before running anything): JEPA planner 15, random moves 9, VLA 19, and a planner on an untrained encoder 40. The JEPA planner mostly gets close and stalls 2–4 cm short (58 scenes); within 4 cm it is level with the untrained encoder (73 vs 75). We measure why: distance in embedding space follows real distance only loosely (rank correlation 0.41). About the setup: a simulated tabletop with two cubes, not a real robot or factory. Our models are under half a million parameters and train in minutes; they teach the mechanism, not the results. We did not run OpenVLA, π0 or V-JEPA 2; they need GPUs (numbers below are theirs, not ours). What you'll learn An action-conditioned predictor in embedding space, trained on 5-step imagined rollouts Why world-model data needs bad moves too, and why you must log the executed action CEM planning toward a goal image, and model-predictive control Reading results as a distribution, not one threshold Why "close in embedding space" is not "close in the world" Planning cost: 960 predictor calls, 70 ms per decision against a 50 ms budget What the real models need: OpenVLA (64 A100s × 14 days to train, 15 GB GPU memory to run), openpi π0 8 GB to run, 22.5 GB LoRA, 70 GB full fine-tune), V-JEPA 2-AC (~62 h of robot video) Chapters 0:00 Intro 0:36 Mental model 2:05 The code 3:55 How it runs 4:13 Results 5:18 Gotchas 5:55 Real models · recap This is the last episode of the playlist. All four episodes' code is linked above. #robotics #worldmodels #jepa #robotlearning #mujoco #pytorch #modelpredictivecontrol #physicalai