World models · JEPA · Planning 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.
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.
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.
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.
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?