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

There is an ongoing debate in AI about the nature of world models. Most focus on generative models that predict the future by creating it in pixels, suggesting it's a waste of effort. The alternative is JEPA (joint-embedding predictive architecture), with three components: an encoder, a predictor, and a planner. The encoder converts each frame of observation into a learned embedding, a compact vector describing the scene.

The predictor takes the embedding and action to predict the next frame's embedding. The planner attempts numerous action sequences inside the predictor and selects the one whose predicted end state is closest to the goal's embedding. All this occurs in embedding space, learning only the important parts of the scene for action decisions.

The author ran 600 experiments on JEPA world model LeWM, primarily to understand its limitations. They began by replacing the predictor with perfect simulator physics, keeping everything else constant. Tasks requiring long horizons (75 steps) were given, along with subgoals from expert demonstrations every 25 steps. The version with true physics succeeded in 77% of cases, while vanilla LeWM succeeded in 78%. Most of the author's experiments focused on improving these planner methods.

The author discovered four planning methods that raised LeWM's success rate on long-horizon tasks from 22% to 85%, matching the performance of a planner with ideal subgoals from an expert demonstration. The first two methods together boosted LeWM's success rate from 22% to 85%. One method significantly improved performance on harder goals, raising success by 60%, three times the benefit of swapping in perfect physics.

Using the planner as published, LeWM succeeded on 22% of long-horizon tasks, while 25-step goals achieved success 89.7% of the time, almost matching the exact-physics version at 90.5%.

To extend LeWM's useful planning horizon, the author employed hierarchical planning. A high-level planner sets subgoals, while a low-level planner executes them. This technique, widely considered one of the main open problems for world models, proved valuable, yielding three times the improvement compared to swapping in perfect physics.

However, extra search didn't compensate for this limitation. The subgoals used came from expert demonstrations, which deployed robots wouldn't have. The second method eliminated the dependency on expert subgoals by training LeWM on expert demonstrations of the task. Its encoder created a graph of 460,000 embedded states, linking to subsequent states in its own demonstrations and close states in other demonstrations.

This embedding space serves as a map of what the robot can reach, with subgoals ranked by their embedding distances. With this method, success on long-horizon tasks (goal 75 steps away) increased from 62% to 85%. However, the standard test took a demonstration and set the goal a fixed number of steps further along the same demonstration, creating a situation where 84% of the standard test goals already held nearly identical states from another demonstration.

The test sample also consisted of only 50 episodes, while 250 fresh episodes drew from the same source yielded only 44% success.