Dream-State Mario — Alexander Xu

Can a neural network be the game? That’s the idea behind a world model. There are no sprites or level files, and nothing in it knows what a Goomba is. You press right, and the model guesses what the next frame should look like.

I’d already poked at this with a small World Models replication on CarRacing, which taught me how to squash frames into a latent space and not much else. Mario felt like the real test because it has rules you can see breaking. If Mario walks through a pipe, you notice.

TLDR: a 46M-parameter latent world model in two parts. A convolutional VAE squashes each frame into a small latent, and a rectified-flow diffusion U-Net predicts the next latent from the last few plus the buttons you’re holding. I trained it on 205 minutes of human gameplay, exported it to ONNX, and it runs in the browser on WebGPU, drawing every frame from your inputs in real time.

The data

I didn’t record my own gameplay. There’s a public dataset of human Mario runs (rafaelcp/smbdataset, CC-BY) with the frames, the buttons, and a RAM dump for every frame. After converting it I had 737,134 frames: 280 attempts, about 205 minutes of play.

I plotted it before anything else, which was lucky, because the data is wildly lopsided. World 8 has almost 109 minutes. Level 1-1, the one everyone knows, has five runs and about three minutes.

Bar chart of minutes of recorded play per level, with world 8 far taller than every other world — Minutes of recorded play per level, runs above each bar. 1-1 gets three minutes. World 8 gets almost two hours.

The buttons are lopsided too. People hold run 83% of the time and press left about 9% of the time, so the model barely ever sees Mario walk backwards.

Tokenizer

NES frames are drawn from a fixed palette of 64 colours. That sounds like trivia, but it shaped the whole tokenizer. Instead of asking the decoder for RGB values (a regression, which blurs everything into brown), I made it pick one of 64 colours per pixel. It’s a classification problem now, trained with cross-entropy, and the edges stay hard.

It gets 99.77% of held-out pixels exactly right. That number is a bit flattering, because most of a Mario frame is flat sky, so I also tracked accuracy on non-background pixels only: 99.22%.

Three rows of Mario frames: original, reconstruction, and the wrong pixels shown in white — The worst held-out frame per world, as original, reconstruction and wrong pixels. Even the worst ones only miss a few specks.

Dynamics

The dynamics model is a U-Net trained with rectified flow: start from noise, learn the straight-line velocity toward the real next latent, take a handful of Euler steps at inference. The buttons and the noise level go in through adaptive GroupNorm, so they act as global facts about the frame rather than something the network has to find in the pixels.

The denoiser has 46.2M parameters, with channel widths of 128, 256 and 384 down the U-Net. That’s too big to train on a laptop (my Mac managed about half a step per second), so it trained on a rented GPU: 300,000 steps at batch 64, in bfloat16.

Where it broke

Single-step predictions looked great. Long rollouts didn’t, and that’s where I learned the most.

Feed the model its own output for a while and the frames stay sharp, but the content drifts. Mario shrinks or disappears, the screen stops scrolling, and scenery shows up that isn’t in any level. Measured by pixels, a long rollout can score about as well as a frozen screenshot, which tells you pixel accuracy alone isn’t the right test.

Four pairs of strips comparing real frames above generated frames at steps 1 to 100; the generated Mario fades out after a few dozen steps — Held-out rollouts, real frames above generated ones, steps 1 to 100. Look at Mario: he drifts out of the frame as the model feeds on its own guesses.

So I changed what the model is responsible for. A hand-written game engine keeps track of the level (physics, collisions, enemies), and the network redraws every frame from a noised copy of the engine’s frame instead of from its own last guess. With the guide off, the dream drifts away like before. Turn it up and the dream stays on the level.

Grid of frames from a bot playing level 5-1 under seven settings; the model-alone row invents a different level while the guided rows match the top row — Same 200 steps of 5-1 under seven settings. Row 1 is the real game. Row 2 is the model alone, which loses Mario and invents a different level. Rows 5 to 7 are guided and stay on track.

Getting it into a browser

The model exports to ONNX and runs on WebGPU. At half precision the denoiser is 91 MB. In headless Chrome on my M5, two denoising steps plus the decode take 30.7 ms, inside a 50 ms frame budget. Four steps take 58 ms, which is too slow, so the site runs at two. The sampler on the web side is a straight port of the Python one.

Python source for the Euler sampler with an optional guide and a looser mask for the score and timer area — The sampler in mmodel/flow.py. The guide decides how much noise to start from; the HUD (score, coins, timer) gets a looser mask so the digits can change.

Every piece has an assert-based check, and the JavaScript and Python versions of the engine produce identical frames: 1,600 frames over four levels, checksum for checksum.

Terminal output of four test commands, each ending in all tests passed — The checks I ran before trusting anything. The denoiser check confirms 46.1M parameters; the last block compares the browser engine against the Python one.

A bug I haven’t fixed yet

While plotting the RAM dump I noticed byte value 10 never appears, not even once in about 1.5 billion bytes. Level page 10 has zero frames, and page 13 has 56,336.

Two bar charts: one byte value has zero occurrences, and level page 10 has zero frames while page 13 is inflated — Value 10 is missing and 13 is inflated, in both charts. That's \r\n being collapsed to \r when the RAM is read.

It’s a newline bug in the reader: \r\n gets collapsed to \r, so every 10 that follows a 13 vanishes. It doesn’t break the visuals, but anything that reads game state from RAM is slightly wrong, so it’s first on the list.

What’s next

Fix the RAM reader, distil the model so it can take fewer denoising steps, and see how far the model alone can go before it needs the guide. I’d also love more 1-1 data, because three minutes of the most famous level in gaming isn’t much to learn from.