{"slug": "lemario-training-a-jepa-world-model-on-super-mario-bros", "title": "LeMario: Training a JEPA World Model on Super Mario Bros", "summary": "A developer trained a Joint-Embedding Predictive Architecture (JEPA) world model on Super Mario Bros, finding that while the model could predict future frames and perform short-range planning, it failed to navigate distant goals or overcome obstacles, revealing a gap between prediction and goal-directed behavior.", "body_md": "## LeMario: Super Mario Bros trained on a JEPA Model\n\nI wanted to reproduce [LeWorldModel](https://arxiv.org/abs/2603.19312), 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.\n\nThe 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\n\nFor 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.\n\nThe model had learned to predict the game, but that did not mean it had learned how to make progress through it. D:\n\nThis 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 )\n\n### The whole architecture\n\nBefore introducing each equation separately, it helps to see the whole machine at once:\n\nLet’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** ():\n\nYou can think of a latent as the model’s private description of a screenshot.\n\nThe red path contains the controller inputs. Each pair of observations is separated by five emulator frames, and every frame contains six possible button states:\n\n```\nframes:  [batch, 4, 3, 224, 224]\nactions: [batch, 4, 5, 6]  # Left, Right, Up, Down, A, B\n```\n\nThe **action encoder** compresses each `5 × 6`\n\nbutton sequence into another 192-number vector.\n\nThe frame and action latents are then piped into the **causal predictor**. Its job is to answer:\n\nGiven what the previous frames looked like and which buttons were pressed, what should the next frame’s latent look like?\n\nThe predictor contains six transformer blocks. Where each of the frames will attend to the previous frames. But however, how would we inject action?\n\nThe actions enter these transformer blocks through **Adaptive LayerNorm Zero** (AdaLN-Zero).\n\nRather than simply attaching the action vector to the frame vector, AdaLN-Zero turns each action into three kinds of controls:\n\n**Shift:** adds an action-dependent offset to the frame features**Scale:** turns particular features up or down.**Gate:** controls how strongly the transformer updates the current state.\n\nNow 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.\n\nFor 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.\n\nThese 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.\n\nThe “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.\n\nNow, during training after the six transformer blocks, a small projection head produces three predicted future latents:\n\nThese are compared with the latents produced by the three real next frames:\n\nNow 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).\n\nSo to prevent it from cheating, we use [ SIGReg1](#sigreg)! It prevents this collapse by encouraging the real frame latents to remain varied and informative. So our new loss function becomes:\n\nSo that's the whole architecture! Now moving on to the actual results.\n\n### But did it Actually Learn?\n\nLeMario 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.\n\nOn held-out episodes, I compared LeMario with that persistence baseline and with the real frame history paired with shuffled actions:\n\n| Method | One-step error | Five-step error |\n|---|---|---|\n| LeMario | 0.013773 | 0.077717 |\n| Predict no change | 0.014472 | 0.142473 |\n| Shuffle the actions | 0.016555 | 0.114648 |\n\nShuffling 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.\n\nLeMario had learned short-horizon Mario dynamics conditioned on the player’s actions!\n\n### Letting it touch the controller\n\nNow for the fun part. Once the model can imagine futures, we can search through those futures and let it choose what Mario should do.\n\n### Searching through imagination\n\nNow 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!\n\nGiven a current image and goal image , the encoder produces and . CEM then:\n\n- Samples hundreds of action sequences.\n- Rolls each sequence forward through LeMario.\n- Scores the predicted final latent against .\n- Keeps the best candidates.\n- Resamples around them and repeats.\n\nCEM 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.\n\n### Then Mario barely moved\n\nI began with a tiny goal. Mario started at `x=40`\n\n; the goal frame showed him at `x=72`\n\n. Raw JEPA+CEM ended at `x=44`\n\n.\n\nIn layman's terms, it sucked.\n\nAt 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?\n\n### What is inside 192 numbers?\n\nThe latent is only 192 numbers. I needed a way to ask what information was inside without changing the encoder.\n\n### Did it forget Mario?\n\nI froze the JEPA and trained a small [probe 2](#the-position-probe) to recover Mario’s emulator coordinates from its latent. Because the encoder could not change, any position the probe recovered had to be information the JEPA had already learned.\n\n```\nhorizontal position: MAE = 9.30 px, R² = 0.997\nvertical position:   MAE = 21.62 px, R² = 0.188\n```\n\nThe probe worked! Mario’s horizontal position was almost perfectly recoverable. Vertical state was much weaker, but the encoder had clearly learned useful information about the player.\n\n### The probe that “fixed” everything\n\nI temporarily scored CEM’s imagined futures using horizontal position predicted by the probe. LeMario still imagined the futures and CEM still chose the actions, the probe only changed how those futures were ranked.\n\nFor a target at `x=72`\n\n, probe-scored CEM moved Mario from `x=40`\n\nto `x=71`\n\n. With local replanning, it later reached `x=176`\n\nfor a goal at `x=177`\n\n.\n\nThis was the first rollout that worked! The JEPA model could imagine useful horizontal motion, and the probe could find it.\n\nThis convinced me that latent planning should work in theory. If Mario’s position was already inside the representation and the learned dynamics could move it, maybe the first goal was simply too similar to the starting frame.\n\n### Try a goal half a level away\n\nThe planner’s job was to find actions connecting two embeddings. If the starting frame and nearby goal already had similar embeddings, doing almost nothing could look like success. So I removed the supervised probe and tried raw latent planning again with reachable goal images roughly halfway through Worlds 1-1, 2-1, and 3-1.\n\nTaking the example for level 3-1: This time Mario moved much farther. Instead of going from `x=40`\n\nto `x=44`\n\n, the three runs reached roughly `x=290–307`\n\n. They still died around the first meaningful hazard, but raw latent planning was no longer doing nothing.\n\n| Start | Goal | JEPA + CEM |\n|---|---|---|\n\nCEM had done exactly what I asked: search for actions that connected two embeddings. A more distant goal apparently created enough pressure in latent space to make Mario move.\n\nBut Mario was still 1,442 world pixels from the goal, while the encoder assigned his final scene a latent distance of only `0.164`\n\n. CEM had predicted `0.153`\n\n, so the predictor was not wildly hallucinating. The encoder itself considered the wrong scene fairly close (as we can see from the start and goal frame above).\n\nMario’s scrolling camera explains why. Two distant locations can look very similar despite being 2 distinct locations within the game! So the model in fact was doing its job, it believed that it found a path to the \"goal\", but in reality it was just stuck at a location that looked similar to the goal embedding.\n\n### Fine, make the goals smaller\n\nSo I then theorized that if we split it into smaller checkpoints then the frames would have more differences, so I split the human run into intermediate image goals. This helped, raw latent planning reached `x=314`\n\n, the most probe-free progress in the project.\n\nBut still shorter goals did not meaningfully fix the representation.\n\nMario reached the first target within two pixels. For the second, he overshot to `x=283`\n\n, corrected backward, and stopped at `x=239`\n\nfive pixels from the reference.\n\nVisually, Mario had reached the right place. But the benchmark still marked the checkpoint as a failure because the final embedding was not close enough to the goal embedding (probably due to the HUD, or other small details).\n\nThis showed that distance was only part of the problem. Smaller goals helped Mario move farther, but latent distance was still a brittle way to decide whether he had arrived.\n\nThe next goal required Mario to jump, and the planner failed again. That matched the earlier probe result, horizontal position was represented well, while vertical position was much weaker.\n\nSmaller goals helped with the planning horizon. They did not make latent distance a reliable measure of progress.\n\nBy this point, the failures no longer looked unrelated. They pointed to three main problems.\n\n### Predictive state is not control state\n\nThe encoder is rewarded for representing whatever helps predict future images. Camera position, enemy phase, animation, and timer state can all be useful.\n\nThe controller needs something different, a state where distance corresponds to controllable progress.\n\n### CEM searches for model weaknesses\n\nCEM did exactly what I asked, find actions that LeMario believed would reach the goal. It could not recognize when the model was wrong.\n\nThis made the model’s weaknesses obvious. It treated visually similar locations as the same place and struggled to reason about vertical movement.\n\n### Mario changed the assumptions behind Push-T\n\nPush-T used nearby goals from expert trajectories, a fixed camera, smooth movement, and a scene whose visual similarity aligned with progress. Its model trained for ten epochs on 20,000 expert episodes.\n\nLeMario trained for one epoch on 280 episodes spread across 32 levels. Nearby goals became half a level away. A fixed camera became a scrolling one. Smooth movement became momentum, jumping, pits, enemies, animation, and death.\n\nI copied the architecture while changing several conditions that made its planning rule work, failing to recognize that they were also vital to the method’s success.\n\n### Where I landed\n\nDespite LeMario not learning to play *Super Mario Bros.* perfectly, it was still able to capture short-horizon, action-conditioned dynamics.\n\nLooking back, many of the mistakes seem obvious in hindsight. I placed too much emphasis on reproducing the architecture itself and overlooked the fact that dataset, environment, evaluation, and underlying assumptions are equally important.\n\nI also learned that unexpected results do not mark the end of an experiment. When Mario barely moved, I had to isolate each component of the system (the predictor, representation, planner, and evaluation metric) and test them individually. Each failure and follow-up experiment deepened my understanding of both the problem and the model.\n\nFor future research projects, I want to validate the central assumptions much earlier, establish simple baselines before trusting any metric, and design evaluations around the behaviors I actually care about rather than solely optimizing for training loss.\n\nLeMario did not become the Mario agent I imagined at the beginning. But it gave me a nice introduction to the field of world model research.\n\n### Appendix\n\n1 SIGReg\n\nSIGReg prevented the encoder from mapping every frame to the same latent. It checked whether the real frame latents remained varied and approximately Gaussian.\n\n```\n192-dimensional JEPA latents\n    → project onto 1,024 random directions\n    → evaluate each projection at 17 points\n    → compare with a standard Gaussian\n    → average the errors\n```\n\nSIGReg was applied independently to each time step across the batch.\n\n2 The position probe\n\nThe probe did not fine-tune the JEPA. I froze the encoder, collected 60 World 1-1 trajectories, and split complete trajectories into training and validation sets. That produced 3,203 training latents and 859 held-out latents.\n\nThe MLP was deliberately small:\n\n```\n192-dimensional JEPA latent\n    → Linear(192, 128)\n    → GELU\n    → Linear(128, 2)\n    → predicted (x, y)\n```\n\n", "url": "https://wpnews.pro/news/lemario-training-a-jepa-world-model-on-super-mario-bros", "canonical_source": "https://www.benjamin-bai.com/projects/lemario", "published_at": "2026-07-14 22:30:47+00:00", "updated_at": "2026-07-14 22:47:43.541439+00:00", "lang": "en", "topics": ["machine-learning", "artificial-intelligence", "neural-networks", "ai-research"], "entities": ["LeWorldModel", "Super Mario Bros", "JEPA", "LeCun"], "alternates": {"html": "https://wpnews.pro/news/lemario-training-a-jepa-world-model-on-super-mario-bros", "markdown": "https://wpnews.pro/news/lemario-training-a-jepa-world-model-on-super-mario-bros.md", "text": "https://wpnews.pro/news/lemario-training-a-jepa-world-model-on-super-mario-bros.txt", "jsonld": "https://wpnews.pro/news/lemario-training-a-jepa-world-model-on-super-mario-bros.jsonld"}}