GoKawiil‹ back to projects
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 )
The whole architecture Before introducing each equation separately, it helps to see the whole machine at once: 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 ( z z z): z t = E θ ( x t ) , z t ∈ R 192 z_t = E_\theta(x_t), \qquad z_t \in \mathbb{R}^{192} z t = E θ ( x t ) , z t ∈ R 192 You can think of a latent as the model’s private description of a screenshot. 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: frames: [batch, 4, 3, 224, 224] actions: [batch, 4, 5, 6] # Left, Right, Up, Down, A, B The action encoder compresses each 5 × 6 button sequence into another 192-number vector. The frame and action latents are then piped into the causal predictor. Its job is to answer: Given what the previous frames looked like and which buttons were pressed, what should the next frame’s latent look like? The predictor contains six transformer blocks. Where each of the frames will attend to the previous frames. But however, how would we inject action? The actions enter these transformer blocks through Adaptive LayerNorm Zero (AdaLN-Zero). Rather than simply attaching the action vector to the frame vector, AdaLN-Zero turns each action into three kinds of controls: Shift: adds an action-dependent offset to the frame features
adds an action-dependent offset to the frame features Scale: turns particular features up or down.
turns particular features up or down. Gate: controls how strongly the transformer updates the current state. 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. 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. 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. 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. Now, during training after the six transformer blocks, a small projection head produces three predicted future latents: z ^ 1 , z ^ 2 , z ^ 3 \hat z_1,\hat z_2,\hat z_3 z ^ 1 , z ^ 2 , z ^ 3 These are compared with the latents produced by the three real next frames: L pred = MSE ( [ z ^ 1 , z ^ 2 , z ^ 3 ] , [ z 1 , z 2 , z 3 ] ) \mathcal L_{\text{pred}} = \operatorname{MSE} \left( [\hat z_1,\hat z_2,\hat z_3], [z_1,z_2,z_3] \right) L pred = MSE ( [ z ^ 1 , z ^ 2 , z ^ 3 ] , [ z 1 , z 2 , z 3 ] ) 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). So to prevent it from cheating, we use SIGReg1! It prevents this collapse by encouraging the real frame latents to remain varied and informative. So our new loss function becomes: L = L pred + 0.1 L SIGReg \mathcal L = \mathcal L_{\text{pred}} + 0.1\mathcal L_{\text{SIGReg}} L = L pred + 0.1 L SIGReg So that's the whole architecture! Now moving on to the actual results.
But did it Actually Learn? LeMario trained on 737,134 frames from 280 episodes across 32 Mario levels. To verify that the model has learned we couldn't just look for a lower loss. Lower loss was not enough to prove it had learned dynamics, nearby frames often look so similar that predicting “nothing changes” is a strong baseline. On held-out episodes, I compared LeMario with that persistence baseline and with the real frame history paired with shuffled actions: Method One-step error Five-step error LeMario 0.013773 0.077717 Predict no change 0.014472 0.142473 Shuffle the actions 0.016555 0.114648 Shuffling the actions raised one-step error by 20.2%. Across five recursive steps, LeMario beat persistence by 45.5%, while shuffled actions were 47.5% worse. The farther it predicted, the more the buttons mattered. LeMario had learned short-horizon Mario dynamics conditioned on the player’s actions!
Letting it touch the controller Now for the fun part. Once the model can imagine futures, we can search through those futures and let it choose what Mario should do. Searching through imagination Now to turn the model that only predicts a frame ahead given an action, into something that could predict multiple actions and steps into the future, I use the Cross-Entropy Method! Given a current image x t x_t xt and goal image x g x_g xg, the encoder produces z t z_t zt and z g z_g zg. CEM then: Samples hundreds of action sequences. Rolls each sequence forward through LeMario. Scores the predicted final latent against z g z_g z g . Keeps the best candidates. Resamples around them and repeats. CEM found action sequences with predicted goal distances far below random candidates. I had a model that could imagine, an optimizer that could search its imagination, and a goal image that required no reward engineering. It was magnificent.
Then Mario barely moved I began with a tiny goal. Mario started at x=40 ; the goal frame showed him at x=72 . Raw JEPA+CEM ended at x=44 . In layman's terms, it sucked. At this point I did not know which part had failed. Maybe the predictor was wrong, maybe CEM was broken, or maybe the encoder had ignored Mario entirely. I needed to start with the simplest question, did those 192 numbers even contain Mario’s position?
What is inside 192 numbers? The latent is only 192 numbers. I needed a way to ask what information was inside without changing the encoder.
... continue reading