LeMario: Super Mario Bros trained on a JEPA Model<br>‹back to projectsLeMario: Super Mario Bros trained on a JEPA Model<br>Benjamin Bai·July 2026·LeWorldModel paper·GitHub<br>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.<br>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<br>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.<br>The model had learned to predict the game, but that did not mean it had learned how to make progress through it. D:<br>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 )

The whole architecture<br>Before introducing each equation separately, it helps to see the whole machine at once:<br>Let’s start with the green path. Each training sample contains four Mario frames. The vision encoder compresses every frame into a 192-number representation called a latent (zzz):<br>zt=Eθ(xt),zt∈R192z_t = E_\theta(x_t), \qquad z_t \in \mathbb{R}^{192}zt​=Eθ​(xt​),zt​∈R192<br>You can think of a latent as the model’s private description of a screenshot.<br>The red path contains the controller inputs. Each pair of observations is separated by five emulator frames, and every frame contains six possible button states:<br>frames: [batch, 4, 3, 224, 224]<br>actions: [batch, 4, 5, 6] # Left, Right, Up, Down, A, BThe action encoder compresses each 5 × 6 button sequence into another 192-number vector.<br>The frame and action latents are then piped into the causal predictor . Its job is to answer:<br>Given what the previous frames looked like and which buttons were pressed, what should the next frame’s latent look like?<br>The predictor contains six transformer blocks. Where each of the frames will attend to the previous frames. But however, how would we inject action?<br>The actions enter these transformer blocks through Adaptive LayerNorm Zero (AdaLN-Zero).<br>Rather than simply attaching the action vector to the frame vector, AdaLN-Zero turns each action into three kinds of controls:<br>Shift: adds an action-dependent offset to the frame features<br>Scale: turns particular features up or down.<br>Gate: controls how strongly the transformer updates the current state.<br>Now how do these affect the transformer block? Usually our normal attention would give each frame its previous context then pass it through the feedforward (MLP) to synthesize that information, however AdaLN modifies both stages according to the action.<br>For example, a jump action might scale up latent features related to vertical motion and scale down features that matter less for predicting the jump. Shift moves the normalized features toward a different action-dependent baseline. Finally, the gate decides how strongly the attention or MLP update should affect the predicted state.<br>These controls are produced separately for the attention and MLP branches, giving the block six values in total: a shift, scale, and gate for each branch.<br>The “Zero” just means their weights begin at zero, so the predictor starts without random action effects and gradually learns which gates to open during training.<br>Now, during training after the six transformer blocks, a small projection head produces three predicted future latents:<br>z^1,z^2,z^3\hat z_1,\hat z_2,\hat z_3z^1​,z^2​,z^3​<br>These are compared with the latents produced by the three real next frames:<br>Lpred=MSE⁡([z^1,z^2,z^3],[z1,z2,z3])\mathcal L_{\text{pred}} = \operatorname{MSE} \left( [\hat z_1,\hat z_2,\hat z_3], [z_1,z_2,z_3] \right)Lpred​=MSE([z^1​,z^2​,z^3​],[z1​,z2​,z3​])<br>Now we just want to lower this loss to 0... but clearly there is one easy way for the model to cheat, we can just make all the latent vectors the same! Prediction would become perfect because Mario, a pipe, a new world would all look identical (representation collapse).<br>So to prevent it from cheating, we use SIGReg 1! It prevents this collapse by encouraging the real frame latents to remain varied and informative. So our new loss function becomes:<br>L=Lpred+0.1LSIGReg\mathcal L = \mathcal L_{\text{pred}} + 0.1\mathcal L_{\text{SIGReg}}L=Lpred​+0.1LSIGReg​<br>So that's the whole architecture! Now moving on to the actual results.

But did it Actually Learn?<br>LeMario...