LeMario: Training a JEPA World Model on Super Mario Bros
LeMario: Super Mario Bros trained on a JEPA Model
I wanted to reproduce LeWorldModel , a small Joint-Embedding Predictive Architecture (JEPA) that learns world dynamics from pixels and actions. The original paper used it for reward-free planning in Push-T. But, since I loved video games, and at the same time wanted to learn more deeply about LeCun's JEPA architecture, I decided to write the whole architecture from scratch and train it on Super Mario Bros.
The model passed every test I initially thought mattered. It generalized to held-out episodes, used the actions, and predicted five-step futures better than strong baselines. Raw reward-free planning could move Mario toward nearby image goals and finish within two and five pixels of the targets. :D
For a moment, it looked like the model had learned to play. Then I moved the goal farther into the level... Mario could not reliably jump over the first major obstacle or navigate toward a single distant goal image.
The model had learned to predict the game, but that did not mean it had learned how to make progress through it. D:
This post is both a technical walkthrough and a postmortem, what I built, how I tested it, the mistakes I made, and the experiments that gradually exposed the real problem. (Most of the lessons seem obvious in hindsight T^T )