After 600 experiments, here's what I discovered about JEPA world models

Solving long-horizon tasks without any extra training 0 / 24 on target · step 0 There's a big debate in AI right now about what a world model should be. Most of the attention goes to generative world models, which predict the future by generating it, usually in pixels. Many scientists, including Yann LeCun, argue that this is wasted effort. Most of the pixels don't matter for deciding what to do, and many of them can't be predicted at all. [1] The proposed alternative is JEPA, the joint-embedding predictive architecture. It has three parts: an encoder, a predictor and a planner. The encoder turns each frame of observation into a learned embedding, a compact vector that describes the scene. The predictor takes that embedding plus an action and predicts the embedding of the next frame. The planner then tries hundreds of action sequences inside the predictor and picks the one whose predicted end state lands closest to the embedding of the goal. The whole loop runs in embedding space, so the model only has to learn the parts of the scene that are important for deciding actions. This year a team including LeCun released LeWM [2] , a JEPA world model whose main benchmark is Push-T [3] , a simulation where a robot pushes a gray T-shaped block across a table until the T sits in the green target position shown in a goal image. I've spent the last few months running 600 experiments on it. A goal image from Push-T The first thing I wanted to know was how much LeWM is held back by its predictions, so I started by swapping out the predictor for the simulator's perfect physics, while keeping everything else the same. I then gave both versions long-horizon tasks, with the goal 75 steps away, along with the same subgoals: frames from an expert demonstration, one every 25 steps. Vanilla LeWM succeeded 78% of the time. The version with true physics succeeded 77%. Along the kinds of paths it was trained on, LeWM's predictions are already as good as perfect physics for planning. Success on 250 long-horizon tasks 78 % Vanilla LeWM learned predictor ≈ 77 % LeWM real physics the simulator itself That means the next gains for a JEPA world model come from how the trained architecture is used, and that's where most of my 600 experiments went. Without retraining the model once, I found four new planning methods. The first two together took LeWM's success rate on long-horizon tasks from 22% to 85%, about the same as a planner that is handed the ideal subgoals from an expert demonstration. The other two methods extend LeWM to harder goals, and one of them, a single change to how the goal is scored, raised the success rate by 60 points, three times as much as swapping in perfect physics did. 1. Plan through subgoals 25 steps apart Used exactly as published, LeWM succeeds on 22% of long-horizon tasks, where the goal is 75 steps away. [2] When the goal is 25 steps away, it reaches it 89.7% of the time, within a point of the exact-physics version at 90.5%. So the model is accurate over 25 steps and struggles over 75. The usual explanation is that small prediction errors compound over a long plan, which would call for a better model. I tested that against the alternative on the same 250 long-horizon tasks. 250 long-horizon tasks, goal 75 steps away LeWM, all 75 steps at once as published 22% Exact physics, all 75 steps at once the simulator itself 42% +20 LeWM given a new subgoal every 25 steps same model, shorter hops 80% +58 0 25 50 75 100% Exact physics added 20 points. Splitting the task into 25-step segments, each the length LeWM already handles almost perfectly, added nearly 60. LeWM's useful planning horizon is about 25 steps, so it was clear that the way to go further is to chain segments of that length together. LeWM's planning horizon is about 25 steps Plan all 75 steps at once ... one search, far goal 22% A subgoal every 25 steps ... three short searches 80% This is hierarchical planning: a high level picks subgoals and a low level reaches them. It's widely considered one of the main open problems for world models, and on Push-T it's worth three times as much as swapping in perfect physics. Extra search doesn't substitute for it. Running the planner four times and keeping the lowest-cost plan took four times the compute and added less than one point, because the four runs almost always succeeded or failed together. The subgoals in this experiment came from an expert demonstration, which a deployed robot won't have. The second method removes that dependence. 2. A latent memory graph with a cost on every jump LeWM was trained on expert demonstrations of the task, and its encoder can embed every frame of them, giving about 460,000 embedded states that I connected into a graph. Each state links to the states that follow it in its own demonstration, and to states in other demonstrations that are close in embedding space. To solve a long-horizon task, the planner searches this graph for a path from the current state to the goal and uses the states along that path as subgoals. This only works if distance in embedding space means something physical, and it does. Take one subgoal the robot can reach within 25 steps and one it can't, and LeWM's embedding distance ranks the reachable one as closer 87% of the time. So the embedding space doubles as a map of what the robot can reach. The first version had a flaw. A jump between nearby states in two different demonstrations cost the same as a real step, so the search strung jumps together into paths the block could never follow (318 of 1,015 plans tried to use six or more jumps). The fix was adding a cost on every jump between demonstrations in the graph. Using nothing but LeWM's own embeddings of past demonstrations, success on the same 250 long-horizon tasks (goal 75 steps away) went from 62% to 85%, which is the same score received by the planner using expert subgoals. Can't get much better than that, right? Planning through a memory of demonstrations start goal Jumps are free · 8 jumps · 62% Each jump has a cost · 1 jump · 85% The standard Push-T test is easier than it looks The caveat is that the standard test takes a demonstration, starts the robot partway through it, and sets the goal a fixed number of steps further along the same demonstration. That's a reasonable design, but in Push-T every demonstration ends with the block in the same place. So every test goal sits on a path that demonstrations already travel, heading for the spot they all end at. For 84% of the standard test's goals, the memory graph already held a nearly identical state from another demonstration. demonstrations, all ending in about the same place goals in the standard test goals from other demonstrations The usual sample is also small. On the official 50 episodes, LeWM succeeded 68% of the time with the goal 50 steps away. On 250 fresh episodes drawn the same way, it succeeded only 44%. Changing nothing but the planner's random seed flipped 38 of 250 episodes between success and failure. That's where I decided to build a harder test, where each task takes its start from one demonstration and its goal from another, so the robot has to travel between states that no single demonstration connects. It only counts as a success if the block stays on target for the last 10 steps. Under these rules the memory-graph planner succeeds 75% of the time on standard tasks and 45% on the new ones. The third and fourth methods I found work on that gap. 3. Fill gaps in memory with transformed memories A memory graph is only as good as its coverage. When I deleted every stored state near the goal, success fell from 93% to 32%. Push-T has a useful symmetry here. The physics of pushing a block doesn't change if you shift, rotate or mirror the whole table. So when the graph has no states near a goal, it can take demonstration segments from elsewhere, transform them to fit the gap, and plan th