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

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
• Scale: 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:

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¹! 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.

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:

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$

1. Samples hundreds of action sequences.
2. Rolls each sequence forward through LeMario.
3. Scores the predicted final latent against $z_g$
4. Keeps the best candidates.
5. 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.

Did it forget Mario?

I froze the JEPA and trained a small 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.

horizontal position: MAE = 9.30 px, R² = 0.997
vertical position:   MAE = 21.62 px, R² = 0.188

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.

The probe that “fixed” everything

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.

Try a goal half a level away

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.

Fine, make the goals smaller

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.

Smaller goals helped with the planning horizon. They did not make latent distance a reliable measure of progress.

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

Predictive state is not control state

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 searches for model weaknesses

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.

Mario changed the assumptions behind Push-T

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.

Where I landed

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.

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

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

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

Appendix

¹ SIGReg

SIGReg prevented the encoder from mapping every frame to the same latent. It checked whether the real frame latents remained varied and approximately Gaussian.

192-dimensional JEPA latents
    → project onto 1,024 random directions
    → evaluate each projection at 17 points
    → compare with a standard Gaussian
    → average the errors

SIGReg was applied independently to each time step across the batch.

² The position probe

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

The MLP was deliberately small:

192-dimensional JEPA latent
    → Linear(192, 128)
    → GELU
    → Linear(128, 2)
    → predicted (x, y)