LeMario: Training a JEPA World Model on Super Mario Bros — Technical Walkthrough and Postmortem

Benjamin Bai's LeMario project reproduces LeWorldModel — a small Joint-Embedding Predictive Architecture (JEPA) — from scratch and trains it on Super Mario Bros. The detailed technical walkthrough covers the full architecture, training setup, and a frank postmortem of what worked and what didn't.

Architecture Overview

The model processes four consecutive Mario frames. A vision encoder compresses each frame into a 192-dimensional latent vector: z_t = E_θ(x_t), z_t ∈ R^192. Actions (button states for Left, Right, Up, Down, A, B over five emulator frames) are encoded separately into another 192-dimensional vector via an action encoder.

The frame and action latents feed into a causal predictor with six transformer blocks. Actions are injected via Adaptive LayerNorm Zero (AdaLN-Zero), which produces shift, scale, and gate controls per branch (attention and MLP). The gate controls how strongly updates affect the predicted state. Weights start at zero so the predictor learns gradually.

The predictor outputs three predicted future latents: ẑ₁, ẑ₂, ẑ₃. These are compared with the latents of the three real next frames using MSE loss: L_pred = MSE([ẑ₁,ẑ₂,ẑ₃], [z₁,z₂,z₃]). To prevent representation collapse, SIGReg regularization is added: L = L_pred + 0.1 * L_SIGReg.

What Worked

• Model generalized to held-out episodes.
• Used action information effectively.
• Predicted five-step futures better than strong baselines.
• Raw reward-free planning could move Mario within 2-5 pixels of nearby image goals.

What Failed

• When goals were placed farther into the level, Mario couldn't reliably jump over the first major obstacle.
• Could not navigate toward a single distant goal image.
• Conclusion: predicting the game ≠ learning to make progress. The model learned pixel-level dynamics but not long-horizon planning or goal-directed behavior.

Key Lessons

Bai notes that most lessons seem obvious in hindsight. The gap between short-term prediction and long-horizon planning is critical — and not captured by standard prediction loss. The post details experiments that gradually exposed this limitation.