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

There has been a heated discussion in the AI community about the nature of world models. Many, like Yann LeCun, believe that generative models focusing on predicting the future by generating it, especially in pixel detail, are a waste of resources. An alternative, called JEPA (Joint-Embedding Predictive Architecture), suggests a different approach. JEPA consists of three components: an encoder, a predictor, and a planner.

The encoder converts each observation frame into a learned embedding, a concise vector that describes the scene. The predictor, given this embedding and an action, forecasts the embedding of the subsequent frame. The planner then attempts numerous action sequences within the predictor's predictions to select the one that brings the final predicted state closest to the embedding of the goal.

All this occurs within embedding space, meaning the model primarily needs to learn the crucial aspects of the scene for decision-making.

One JEPA world model, LeWM, was recently released, with its primary benchmark being Push-T. Push-T is a simulation where a robot must push a gray T-shaped block across a table until it reaches a designated green target area on a goal image.

In recent months, the author ran 600 experiments with LeWM. Initially, they tested how much LeWM is constrained by its predictions by replacing the predictor with the simulator's perfect physics while maintaining everything else constant. The success rate for both versions on long-horizon tasks (75 steps away) was 78% and 77%, respectively. Intriguingly, LeWM's predictions appear as accurate as perfect physics for planning along paths it was trained on.

Subsequent experiments led to the discovery of four new planning methods that improved LeWM's success rate on long-horizon tasks. The first two methods increased the success rate from 22% to 85%, comparable to using expert demonstration subgoals. The third method, involving a change in goal scoring, boosted the success rate by 60 points, three times the improvement achieved by using perfect physics.

This suggests that LeWM's useful planning horizon is approximately 25 steps. To extend its capabilities, hierarchical planning can be employed, where a high-level component sets subgoals, and a low-level component handles reaching them.

This hierarchical planning method proved to be the most significant factor in enhancing LeWM's performance, with the added benefit of reducing computational costs significantly. However, the current test design using a demonstration has limitations, as most goals align with paths already traveled in the demonstrations. To address this, running the planner with a different random seed in fresh episodes yielded a success rate of 44%, highlighting the potential of JEPA world models under improved conditions.