Mario world model learned the game, not how to beat it

A homegrown JEPA world model trained on Super Mario Bros. could predict what happens next with surprising accuracy — yet still failed at the thing that mattered most: getting Mario through the level.

In a detailed write-up, the project’s author rebuilt LeWorldModel from scratch to better understand Yann LeCun’s Joint-Embedding Predictive Architecture and adapted it from Push-T to Nintendo’s platformer. The resulting system, dubbed LeMario, learned short-horizon, action-conditioned dynamics from pixels and button presses. It generalized to held-out episodes, used actions rather than ignoring them, and beat strong baselines at predicting five-step futures.

The setup used a vision encoder to compress frames into 192-dimensional latents, an action encoder to turn 5 × 6 button sequences into matching vectors, and a causal predictor with six transformer blocks. Actions were injected through AdaLN-Zero, which modulated the predictor with learned shift, scale, and gate values for both attention and MLP branches. Training combined prediction loss with SIGReg, a regularizer meant to prevent representation collapse.

What LeMario got right

LeMario trained on 737,134 frames from 280 episodes across 32 Mario levels. On held-out episodes, it beat both a persistence baseline and a shuffled-action baseline:

• LeMario: one-step error 0.013773, five-step error 0.077717
• Predict no change: one-step error 0.014472, five-step error 0.142473
• Shuffle the actions: one-step error 0.016555, five-step error 0.114648

According to the author, shuffling actions increased one-step error by 20.2%. Over five recursive steps, LeMario outperformed persistence by 45.5%, while shuffled actions were 47.5% worse. The farther the model predicted, the more the controls mattered.

Recommended reading

Open models are overtaking AI’s front line

That suggested the model had genuinely learned local game dynamics, not just a cheap “nothing changes” shortcut.

Why planning still broke down

The trouble started when the author tried reward-free planning with the Cross-Entropy Method. In principle, the system could encode a current frame and a goal frame, imagine many action sequences, and choose the ones whose predicted final latent landed closest to the goal.

In practice, it barely moved. For a simple target where Mario started at x=40 and the goal frame showed x=72, raw JEPA + CEM ended at x=44.

A frozen probe revealed part of the issue. When trained to recover Mario’s coordinates from the latent, it could infer:

• Horizontal position: MAE = 9.30 px, R² = 0.997
• Vertical position: MAE = 21.62 px, R² = 0.188

So the latent clearly encoded useful horizontal information, but much less about vertical state. When the author used the probe to score imagined futures by horizontal position, planning suddenly worked far better: for a target at x=72, Mario moved from x=40 to x=71. With local replanning, it later reached x=176 for a goal at x=177.

That result showed the world model could imagine useful motion. The bigger problem was that latent distance was a poor measure of actual progress.

When the author tried a farther target roughly halfway through levels 1-1, 2-1, and 3-1, Mario moved much farther — in 3-1, runs reached roughly x=290–307 — but still died at the first major hazard. Worse, Mario could still be 1,442 world pixels from the goal while the encoder judged the scene quite similar, with a latent distance of just 0.164. CEM had predicted 0.153, so the predictor was not wildly hallucinating; the representation itself treated visually similar but genuinely distant locations as close.

That mismatch came from Mario’s scrolling camera. Different places in the level can look alike even when they are far apart in game coordinates. Breaking the task into smaller image checkpoints helped somewhat — raw latent planning reached x=314 — but did not fix the underlying issue. The planner still struggled with jumps, and even when Mario visually reached a checkpoint, small differences like the HUD could keep the final embedding from matching closely enough.

The author’s conclusion is blunt: a state representation that is good for predicting future frames is not necessarily good for control. LeMario succeeded as a short-horizon predictive model, but failed as a navigation policy because the latent space did not align with controllable progress through the game. That gap matters even more in Mario than in Push-T, where the original setup used nearby goals, a fixed camera, smooth movement, and 20,000 expert episodes over ten epochs. LeMario, by contrast, trained for one epoch on 280 episodes spread across 32 levels.