LeMario：在《超级马里奥兄弟》上训练JEPA世界模型

作者从头实现了一个Joint-Embedding Predictive Architecture （JEPA）世界模型LeMario，在32个《超级马里奥兄弟》关卡、737，134帧数据上训练。

模型在五步预测中比"无变化"基线好45.5%，能通过交叉熵方法（CEM）规划短距离移动，但无法可靠跳过第一个主要障碍或导航到远处目标--学会预测游戏不等于学会通关。

I 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.

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 )

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 ):

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:

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:

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:

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:

These are compared with the latents produced by the three real next frames:

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 SIGReg 1 ：#sigreg! It prevents this collapse by encouraging the real frame latents to remain varied and informative. So our new loss function becomes:

So that's the whole architecture! Now moving on to the actual results.

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:

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.

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.

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 x t ​ and goal image x g x_g x g ​ , the encoder produces z t z_t z t ​ and z g z_g z g ​ . CEM then:

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.

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 .

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?

The latent is only 192 numbers. I needed a way to ask what information was inside without changing the encoder.

I 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.

The 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.

I 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.

For a target at x=72 , probe-scored CEM moved Mario from x=40 to x=71 . With local replanning, it later reached x=176 for a goal at x=177 .

This was the first rollout that worked! The JEPA model could imagine useful horizontal motion, and the probe could find it.

This 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.

The 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.

Taking the example for level 3-1: This time Mario moved much farther. Instead of going from x=40 to x=44 , the three runs reached roughly x=290–307 . They still died around the first meaningful hazard, but raw latent planning was no longer doing nothing.

CEM 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.

But Mario was still 1,442 world pixels from the goal, while the encoder assigned his final scene a latent distance of only 0.164 . CEM had predicted 0.153 , 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).

Mario’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.

So 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 , the most probe-free progress in the project.

But still shorter goals did not meaningfully fix the representation.

Mario reached the first target within two pixels. For the second, he overshot to x=283 , corrected backward, and stopped at x=239 five pixels from the reference.

Visually, 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).

This 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.

The 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.

By this point, the failures no longer looked unrelated. They pointed to three main problems.

The encoder is rewarded for representing whatever helps predict future images. Camera position, enemy phase, animation, and timer state can all be useful.

The controller needs something different, a state where distance corresponds to controllable progress.

CEM did exactly what I asked, find actions that LeMario believed would reach the goal. It could not recognize when the model was wrong.

This made the model’s weaknesses obvious. It treated visually similar locations as the same place and struggled to reason about vertical movement.

Push-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.

LeMario 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.

I 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.

Despite LeMario not learning to play Super Mario Bros. perfectly, it was still able to capture short-horizon, action-conditioned dynamics.

Looking 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.