Teaching Mario to speedrun

Reinforcement learning is getting a lot more attention lately, mostly because of its role in post-training LLMs. The models that “think” before they answer got that way through RL.

It's also a field that's hard to learn from the ground up. The literature takes quite a journey to get from its beginnings to what people actually use today: bandits, Markov decision processes, Bellman equations, dynamic programming, Q-learning, then deep Q-networks, policy gradients, TRPO, and finally PPO. A lot of that history doesn't translate to what's important now. You can get a long way on a few ideas, and this post tries to give you exactly those.

It comes in two parts. The first is the basic formulation and the algorithm underneath nearly everything, with a tiny agent you can train right here in the page. The second is a real project: puff-rs, which trains agents to speedrun stars in Super Mario 64, the actual cartridge, recompiled to native code. Along the way it covers PufferLib, the library it's built on, and Go-Explore, the trick that got it to its first stars.

If neural networks themselves are new to you, How do neural networks work? covers them. Here they're just a function with knobs that gradient descent can turn.

1The loop

Like the rest of machine learning, RL is about learning a function from data. The function we want is called a policy. It takes in an observation of some environment and puts out an action:

action = policy(observation)

What makes RL its own thing is that the policy isn't labeling a fixed dataset. Its actions go back into an environment, which changes because of them, and hands back the next observation:

observation = environment(action)

Put the two together and you get a loop that looks a lot like playing a game:

observation = environment.reset()
while True:
    action = policy(observation)
    observation = environment(action)

That's the heart of RL. Everything else is about making the policy good.

For Mario, an observation is a list of numbers: where he is, how fast he's going, which way he faces, whether he's in the air, how high the floor is around him. An action is what you'd do with a controller: push the stick some direction, hold A or don't, hold B or don't, hold Z or don't. One trip around the loop is one step. A run from start to finish, whether that's grabbing the star, falling off the course or running out of time, is an episode.

2Copying someone

How does a policy learn which action to take given which observation? The simplest approach is imitation learning. Record someone good playing the game, save every (observation, action) pair, and you have a labeled dataset:

# a few hours of a human playing
dataset = [(obs_0, action_0), (obs_1, action_1), ...]

# ordinary supervised learning: predict the human's action
for obs, action in dataset:
    loss = cross_entropy(policy(obs), action)
    loss.backward(); optimizer.step()

This works, and it's essentially how an LLM is pretrained: copy the next word a human wrote. But it has a ceiling built in. A policy trained to copy a player is at best as good as the player. Worse, it's only seen the situations the player got into. Once its own small mistakes push it somewhere the player never went, it has no data for it, and the mistakes compound.

3Past the teacher

The magic of RL is that it isn't capped by anyone's examples. It can find things no human would.

• TD-Gammon (1992) taught itself backgammon by playing itself, and changed how top humans played some opening moves.
• DQN (2013–15) learned dozens of Atari games from the pixels and the score alone, with the same network and settings for every game.
• AlphaGo (2016) beat Lee Sedol, and its move 37 in game two was one that professionals first took for a mistake. It started from imitating human games and then improved by playing itself. Its successors, AlphaGo Zero and AlphaZero, skipped the human games entirely and got stronger.
• OpenAI Five and AlphaStar (2019) reached the top level of Dota 2 and StarCraft II.
• Reasoning LLMs (2024 on) are trained with RL on problems whose answers can be checked, and learn to work through long chains of thought nobody wrote down for them.

What all of these have in common is that learning wasn't confined to a dataset. It came from interacting with the environment itself. For that to work, the environment has to say whether things are going well.

4Reward

That's where reward comes in. Along with each observation, the environment hands back a single number saying how good the last step was. We decide what good and bad mean by deciding where that number comes from.

observation = environment.reset()
while True:
    action = policy(observation)
    observation, reward, done = environment(action)
    if done:
        observation = environment.reset()

For a star in Mario the reward can be very simple: +1 for touching the star, and a small cost for every frame that passes, so sooner is better. Nobody tells the policy how to get there. Its only job is to make the total reward of an episode, the return, as large as possible.

Rewards a few steps from now usually count for a bit less than reward right now. The return from step t is

G_t = r_t + γ·r_(t+1) + γ²·r_(t+2) + …       # γ (gamma) around 0.97 to 0.999

which answers the question “from here, doing what I did, how did things turn out?” That is the signal everything learns from. And notice the hard part: if the policy has never once touched the star, every episode returns the same bad number, and there's nothing to learn from at all. Keep that in mind for part 2.

5A whole example, live

Here's all of it in one small place. The environment is a grid. The agent starts in the bottom-left corner and the star is in the top-right. The red squares are pits. Four actions: up, right, down, left.

• Observation: which square the agent is on.
• Reward: +1 for the star, −0.3 for falling in a pit, −0.01 for every step. An episode ends at the star, in a pit, or after 80 steps.
• Policy: a table with four numbers for each square, turned into probabilities with a softmax. It starts out all zeros, so every action is equally likely: a random walk.

Press train. Each episode plays the current policy, then nudges each square's numbers toward the actions that turned out better than expected and away from the ones that turned out worse. The arrows show the policy: the longer an arrow, the more likely that move.

success rate, last 100 episodes

train 100 episodes play one episode reset show value

episodes0 reached the star— steps, last win—

The whole learning rule fits in a few lines. This is the code running above:

def train_one_episode():
    trajectory = play(policy)                  # [(square, action, reward), ...]
    G = 0
    for square, action, reward in reversed(trajectory):
        G = reward + gamma * G                 # the return from this step
        advantage = G - value[square]          # better or worse than usual from here?
        value[square] += 0.1 * advantage       # update what's usual
        probs = softmax(logits[square])
        for a in range(4):                     # the gradient of log probs[action]
            logits[square][a] += lr * advantage * ((a == action) - probs[a])

This is a policy gradient (REINFORCE, with a baseline). Two ideas are doing all the work:

• Try things, and do more of what worked. The term (a == action) - probs[a] is the gradient of the log-probability of the action taken. Following it makes that action more likely. Multiplying by the advantage decides which way and how far.
• “Worked” means “better than expected.” The value table is a running guess of the return from each square. Subtracting it means an action is only reinforced if it did better than the agent usually does from there. Without it, every action in a winning episode gets credit, including the dumb ones.

Replace the table with a neural network and the square with a real observation, and that's the core of modern RL.

6PPO, which is just a neural network

The workhorse algorithm in RL today, for robots, for games and (with variations) for LLM post-training, is PPO: proximal policy optimization (Schulman et al., 2017). It's the grid example above with three upgrades.

The policy and the value are one network

The tables become a neural network. The observation goes in. Out come logits, one per possible action, which a softmax turns into probabilities to sample from, plus one extra number: the value, the network's guess of the return from here. That's the same value table, now a second head on the same network.

logits, value = network(observation)
action = sample(softmax(logits))

When the action has parts, like a stick direction and three buttons, there's one group of logits per part, called a head, and each is sampled separately.

Collect a batch, then learn from it

Instead of learning after every episode, PPO runs many copies of the environment at once for a fixed number of steps, the horizon, and stores everything: observations, actions, the probability each action had, rewards, values and where episodes ended. Then it learns from that whole batch.

Better advantages

The full return G is noisy: it lumps together everything that happened for the rest of the episode. PPO usually uses generalized advantage estimation (GAE), which leans on the value head to judge each step more locally:

δ_t = r_t + γ·V(s_(t+1)) - V(s_t)            # how much better this step went than expected
A_t = δ_t + (γλ)·δ_(t+1) + (γλ)²·δ_(t+2) + …  # λ around 0.95 blends near and far

Don't move too far at once

This is the “proximal” part. The batch was collected by the old policy, but we want to take several gradient steps on it. If the policy changes a lot, the batch stops describing what the policy now does, and learning can fall off a cliff. So PPO compares the new probability of each action to the old one, and clips that ratio to stay within about 20%:

ratio  = new_prob(action) / old_prob(action)
loss_π = -min(ratio * A, clip(ratio, 1 - 0.2, 1 + 0.2) * A)
loss_V = (value - return)²
loss   = loss_π + 0.5 * loss_V - 0.01 * entropy  # entropy: keep some randomness

The whole algorithm, then:

while training:
    # 1. collect: run the current policy in N environments for `horizon` steps
    for t in range(horizon):
        logits, values[t] = network(obs)
        actions[t] = sample(logits); old_probs[t] = prob(logits, actions[t])
        obs, rewards[t], dones[t] = envs.step(actions[t])

    # 2. judge: how much better than expected was each action?
    advantages = gae(rewards, values, dones)

    # 3. learn: a few passes of gradient descent over minibatches of the batch
    for minibatch in shuffle(batch):
        loss = ppo_loss(network, minibatch, advantages)
        loss.backward(); optimizer.step()

That's it. Everything that makes it work in practice is in the details: how big the batch is, how the network carries memory between steps, how fast the environment can be stepped. Which brings us to Mario.

7PufferLib

PufferLib is a super sensible and very fast RL library. Its environments are written in plain C, its trainer is one tight program, and its defaults work without much tuning. On a single GPU it trains simple environments at millions of steps a second. Most RL code spends its time in Python glue between the game and the network. PufferLib's whole design is about not doing that.

The core of it is a data contract. For every agent, in every copy of every environment, there are four flat arrays the trainer and the environments share:

observations  float[num_agents][obs_size]    # env writes, trainer reads
actions       float[num_agents][num_heads]   # trainer writes, env reads
rewards       float[num_agents]              # env writes, trainer reads
terminals     float[num_agents]              # env writes: 1 where an episode just ended

Stepping every environment is then a loop over indices. It's the same observation = environment(action) from part 1, done for every agent at once, with no allocation and no copying:

// vecenv.c, simplified: one step of every env, in parallel on CPU threads
memset(rewards, 0, num_agents * sizeof(float));
memset(terminals, 0, num_agents * sizeof(float));
#pragma omp parallel for
for (int i = 0; i < num_envs; i++) {
    puf_step(&envs[i]);   // reads its actions, writes its observations, rewards, terminals
}

And the trainer's side of it, collecting one horizon:

for t in range(horizon):                # 128 steps
    logits, value = policy(observations)    # all agents at once, on the GPU
    actions[:] = sample(logits)             # written straight into the shared buffer
    vec_step()                              # every env steps, on CPU threads
    store(observations, actions, rewards, terminals, logits, value)

What it means to be a Puffer env

A PufferLib environment is a C header that defines a struct called Env and a handful of functions. The struct holds whatever the game needs, plus an Agent for each agent it runs, whose pointers the vecenv aims at that agent's rows of the four buffers. The functions are:

void puf_init(Env* env, Dict* config);   // read settings from the .ini, allocate
void puf_reset(Env* env);                // start an episode, write the first observation
void puf_step(Env* env);                 // read actions, advance, write obs/rewards/terminals
void puf_render(Env* env);               // draw it, for watching
void puf_close(Env* env);
// plus a Log struct of floats, summed over finished episodes for the dashboard

Here's the step function of the 2D platformer in puff-rs, the trainer's smoke test, trimmed a little. It's an ordinary game update that happens to read its controller from agent->actions and write its score to agent->rewards:

void puf_step(Platformer* env) {
    Agent* agent = &env->agents[0];
    int move = agent->actions[0];          // head 0: none, left, right
    int vert = agent->actions[1];          // head 1: none, jump, drop
    for (int i = 0; i < env->frameskip; i++) {
        int outcome = physics_tick(env, move, vert);
        // pay for new ground covered toward the flag; the whole level is worth 1
        if (player_x > env->best_x) {
            agent->rewards[0] += (player_x - env->best_x) * progress_scale;
            env->best_x = player_x;
        }
        if (outcome == TICK_GOAL)  { agent->rewards[0] += 1.0f; end_episode(env); return; }
        if (outcome == TICK_DEATH) { agent->rewards[0] -= 0.1f; end_episode(env); return; }
    }
    compute_observations(env);             // write agent->observations
}

There's no Python, no gym wrapper and no inter-process messaging between the game and the trainer. It's why a Puffer env can step millions of times a second, and why an environment as heavy as a whole Nintendo 64 still fits the same mold.

8The Puffer net, and the trainer

PufferLib 5.0's default policy is small, and that's on purpose. When an environment steps millions of times a second, a big network becomes the bottleneck, and in RL more experience usually beats more parameters.

• Encoder. One bias-free linear layer from the observation to 128 numbers. If the observation includes a picture, three convolutions from the classic Atari DQN go in front.
• MinGRU. Four layers of a minimal recurrent unit. This is the policy's memory: it carries a hidden state from step to step, so it can know it pressed jump a moment ago, or that it's been falling for a second. Its state carries across horizons and clears when an episode ends. A MinGRU's gates depend only on the input, not the previous state, so a whole training segment can be computed as a parallel scan instead of one step at a time.
• Decoder. One matrix from the hidden state to every head's logits plus the value.

The training algorithm is PPO as described above, with a few 5.0 choices: advantages are computed per minibatch from fresh values (GAE, or V-trace), and the optimizer is Muon, which orthogonalizes each weight matrix's update instead of scaling it per parameter as Adam does.

puff-rs: the trainer on any GPU

PufferLib 5.0 trains with one CUDA program, which only runs on NVIDIA cards. I wanted it on a Mac. So puff-rs is that trainer ported to Rust on wgpu, which runs on Metal, Vulkan or DX12. There's no CUDA, no Python and no autograd: every forward and backward pass is a hand-written WGSL compute kernel, the way the CUDA original hand-writes its own. A matmul kernel stands in for cuBLAS, the MinGRU gets a scan forward and backward, the PPO loss computes its gradient in the same pass, and Muon does five Newton-Schulz iterations on the GPU.

Every shape is known when training starts, so every GPU dispatch is built once. A minibatch is 146 dispatches, which the CPU encodes in 0.2 ms. All the weights live in one buffer in PufferLib's checkpoint layout, so saving a checkpoint is writing that buffer to disk, and PufferLib's own tools can read it. Tests check every gradient against golden numbers from an autograd implementation.

On an M4 Pro, the platformer trains at 445,000 steps a second with 1,024 agents, and reaches the flag in 99.7% of episodes after 45 seconds of training. Super Mario 64 is a different story, because each “environment” is a whole game.

9The game as an environment

The game isn't emulated, and it isn't reimplemented. N64Recomp showed that an N64 cartridge's MIPS machine code can be statically recompiled: translated, ahead of time, into C, and compiled into a native program that runs the original game logic directly on your CPU. N64Bundler, a separate project of mine, packages that up: give it your ROM, and it recompiles it and runs it.

For RL, its important feature is a gym mode, which hands control of the console to another process. Three properties make training possible:

• The game runs off the clock. It advances one frame when the environment asks, as fast as the machine can manage. Headless, that's about 5,000 frames a second per game, around 170 times real time.
• Steps are deterministic. A frame is over when every thread of the game is parked waiting on the hardware, and that's checked rather than waited out. The same inputs from the same state give the same Mario, to the bit.
• Savestates. The console's 8 MB of memory plus the registers of each of the game's threads. Every episode starts by loading one.

So the Mario environment is a Puffer env like any other. puf_step turns the four action numbers into a controller state, asks the game for two frames, then reads Mario out of RAM at the addresses the US cartridge keeps him at, and writes the observation and reward. Eight games at once train at about 3,500 to 5,000 steps a second on a laptop. That's a hundred times slower than the platformer, and it shapes every decision after this.

A goal is one named star, like pss-2: Peach's Secret Slide's second star, the one that appears only if Mario reaches the bottom within 21 seconds. The first time a star is asked for, the environment plays through the intro, walks Mario into the course, picks the act, and saves a savestate. Every episode loads that and starts the clock. The episode ends the frame Mario touches the star, checked by that star's own bit flipping in the save file.

10What the policy sees, and does

The observation is 292 floats, almost all of it read straight out of the game's memory every step:

A few of these needed real work. The floor probes walk the course's collision geometry in RAM the way the game's own find_floor does, and they agree with the game to a tenth of a unit. They're how the policy can know the slide's track ends in a 6,750-unit drop before it jumps. Everything angular is relative to Mario's facing, so “the star is ahead and to the left” looks the same wherever he is on the course.

The action has four heads: a stick direction (none, or one of sixteen), and A, B and Z, each held or not. The stick is also relative to the way Mario faces, and that took fixing. The game adds the camera's angle to whatever the stick says, so a policy pushing “up” would run wherever the camera happened to point, and would have to learn the camera too. Instead, the environment measures the offset between the stick it sent and the direction the game took from it every frame, and corrects for it. One decision lasts two frames, so the policy acts fifteen times a second, which is fast enough for a long jump's A-then-Z.

There's an optional mode that appends the game's actual picture, shrunk to 80×60, read from the console's video interface. On the slide it roughly paid for the half of the throughput it costs, and elsewhere it was worth nothing measurable, so the runs below use the numbers alone.

The reward

• The star pays 1 and ends the episode.
• The clock costs a little every frame, adding up to 1 over the whole episode. Sooner is better.
• Progress pays for each new closest approach to the star, worth less the later it comes. It's measured along the paths episodes have actually taken, not as the crow flies. The slide winds, so for long stretches of the fastest route the straight-line distance to the star doesn't shrink at all.
• Death, or warping out of the course, costs whatever is left of the clock at once, so dying is never a way to stop it early.

11Go-Explore

Remember the problem from part 1: if the policy never touches the star, every episode returns the same number and there's nothing to learn. Mario makes this about as bad as it gets. pss-2 only appears if he reaches the bottom within 21 seconds, so a careful, slow run down the slide pays exactly as much as falling off at the top. This env's first goal was just opening the castle's front door, and 12 million steps of a goal and a clock alone never found it once.

Rewarding novelty, being somewhere new, helps a bit. Cut the course into cubes 500 units on a side and pay the first time an episode enters each one. But novelty alone still gets stuck: one run swam the castle moat for 10 million steps, because there were always new cubes in the water.

The fix is Go-Explore (Ecoffet et al., 2019), the method that set records on Montezuma's Revenge, the Atari game that long stumped RL for exactly this reason. Its idea is simple: remember the places you've been, go back to the promising ones, and explore from there.

archive = {}   # cell -> the shortest list of pad inputs that reached it

def reset():
    load(start_savestate)
    if random() < go_explore:
        cell = pick(archive)               # rare cells, or ones on the way to the star
        for pad in archive[cell]:
            game.step(pad)                 # replay: deterministic, so this lands exactly there
    start_clock()

def step(action):
    ...
    cell = cell_of(mario)
    if cell not in archive or len(inputs) < len(archive[cell]):
        archive[cell] = inputs.copy()      # a new place, or a faster way to an old one

Because the game is deterministic, a list of pad inputs is a place. It's a few kilobytes, against a savestate's 8 MB, so the archive can hold over a hundred thousand cells. A share of episodes start from the archive instead of the real start, which turns one impossible problem, getting all the way down the slide fast, into many small ones: getting a little further than anyone has from here.

Most of the work went into choosing which cells to start from:

• Cells are a 500-unit cube plus a speed bucket, whether Mario is in the air, and one of eight headings. With cubes alone, a slow approach and a fast one to the same spot were the same cell, and whichever got there first was kept forever.
• Rarity. Some starts are weighted toward cells few episodes have entered. That's the exploring.
• The frontier. Others come from the cells most on pace for the star: reached early, and close to the star measured along the way episodes have actually gone. With fine cells, the archive grows to tens of thousands of them, and rarity alone spreads starts evenly over the whole course.
• The clock. An episode from the frontier starts with its clock at the frames its way there took, so it races the star's real time limit. Other archive starts get a clock of their own to explore with, so the archive can still work its way down the slide in stages.
• The ghost. Entering a cell on the way to the star sooner than the archive's best way into it pays a little per frame of difference, and then the archive keeps the faster way. Only a run faster than every run before it gets paid.

The policy and the archive help each other. The archive puts the policy in useful places it couldn't reach on its own, and the policy, trained by PPO from those starts, gets better at moving fast, which adds faster ways to the archive. Each new record run is then a sequence of inputs that can be replayed, checked and recorded as video.

A bug worth mentioning: the archive's cubes originally stopped at a height of −4,096. The bottom of the slide, star included, is below that. For several runs Go-Explore had never kept a single place at the bottom. After extending them, the first star from a blank slate came 26 minutes in.

12Results

Here's where things stand. More to come.

Times are from the moment the course starts to the frame Mario touches the star. About a hundred of those frames on the slide are a cutscene while the star comes out, which no run can shorten.

Every time in the table is a route: a list of pad inputs the archive put together from episodes that began partway there, which replays from the real start to the star. That isn't the same as a policy that gets there on its own, and none does yet. The log keeps the two apart. perf counts stars from every start, archive starts included, and start_perf counts only episodes from the real start. It has read 0 in every run, on both stars. Folding them into one success rate would hide the part that's still unsolved.

The code, the record inputs and the full experiment log are at github.com/bigmah/puff-rs. You need your own copy of the cartridge, and then make sm64-train builds everything and starts training.