I fixed it. And to be upfront about how I work now: I did it pairing with an AI coding agent (Claude Code) — I drove the diagnosis and the decisions, it accelerated the grind. That collaboration is part of the story.

Why leave Mario at all?

Short version: Mario 1-1 is a single hand-built level, so you can’t separate learning from memorizing, and it runs on an emulator that caps how many environments you can parallelize. Procgen fixes both — procedurally-generated levels with a real train/test split, batched natively in one process. I unpack all of that in the previous post, From Mario to Procgen. This post is about the wall I hit getting there.

The wall

There’s no prebuilt Procgen wheel for Apple Silicon, so pip tries to build from source — and on an M1 that build simply fails. Three separate root causes, none of them obvious from the error spew:

1. An x86-only compiler flag. Procgen’s CMakeLists.txt compiles packaged builds with -march=ivybridge — an Intel microarchitecture. arm64 Clang rejects it outright. (The fix: on arm64/aarch64, use the base armv8-a ISA instead.)
2. Qt looked for at the wrong address. The build probes one hardcoded path for Qt — /usr/local/opt/qt5, the Intel Homebrew location. Apple Silicon Homebrew installs under /opt/homebrew, so CMake never finds it. (The fix: also probe /opt/homebrew/opt/qt@5/lib/cmake.)
3. Warnings that became errors. Modern AppleClang (21) promotes two “set-but-unused variable” warnings in starpilot.cpp and chaser.cpp to hard errors under -Werror. (The fix: mark the two locals [[maybe_unused]] — zero change to any game logic.)

None of these is deep. All three are the kind of paper cut that makes a person give up: buy a cloud box, switch benchmarks, or paste a hack into their local copy and never speak of it again.

The last mile

Here’s the part I actually care about. A local hack would have unblocked me in ten minutes. Instead:

• I traced each failure to its root cause rather than silencing symptoms.
• The fix is four files, +15/−5 lines, behavior-preserving — the smallest change that restores the source build without touching a single game’s determinism.
• I upstreamed it as a pull request to openai/procgen — and it turned out to close issue #69, “Apple Silicon Support,” which had been open since 2021 with a string of people hitting the same wall. So the fix doesn’t just help me; it helps everyone who tries Procgen on a Mac next.
• Then I proved it holds under load — not a 50-step smoke test, but a full 100M-step PPO run (IMPALA-CNN, 64 environments batched natively, training on Metal/MPS) plus an off-policy world-model experiment on top. Stable start to finish.

PR: openai/procgen#107 — Fix building from source on Apple Silicon (arm64) and recent Clang. My first open-source contribution.

That’s the difference between “I tried Procgen and it didn’t work” and “Procgen works on Apple Silicon now, for everyone, because I fixed it.” The distance between those two sentences is the last mile.

The payoff

With the build working, the question that justified the whole detour became answerable. Trained on 200 Miner levels and evaluated on the unseen full distribution:

The agent genuinely learned to play — 91% solved on levels it had never seen — but there’s a real 27% generalization gap it still leaves on the table. That gap is the entire point: a single, honest number that Mario could never have shown me, and the starting line for the representation and transfer experiments this project is actually about.

The compiler flag was never the interesting part. Refusing to route around it was.

1. Stop being CPU-bound — a genuine waste in deep RL, and
2. Be able to test generalization at all.

If you remember one sentence: Mario is one hand-built level behind an emulator; Procgen is a factory of levels running natively in one process — and that difference is preprocessing you don’t write, environments you can actually parallelize, and a generalization test you couldn’t run before.

What each agent actually sees. Mario’s rich 240×256 frame gets reduced to an 84×84 grayscale stack before the network touches it; Procgen’s 64×64 frame is fed in as-is.

Every new term is defined the first time it appears.

Part 0 — What an “environment” gives you

An RL environment is the world the agent acts in. Each step it takes an action and returns an observation (here: an image of the screen), a reward, and whether the episode ended. The agent never sees the game’s internal state — only pixels. So the shape and cost of those pixels, and how many copies of the world you can run at once, are first-order design facts. That’s what changes between Mario and Procgen.

Part 1 — Preprocessing: the frames you must fix vs. the frames you don’t

Mario hands you a raw NES frame: (240, 256, 3) RGB, at the console’s frame rate. You can’t feed that to a CNN as-is — three problems, three fixes, the standard Atari/DQN pipeline:

• Too big → resize to 84×84 (less compute per frame).
• Color is mostly decoration → grayscale (drop 3 channels to 1).
• A single frame isn’t Markov. From one still image you can’t tell if Mario is rising or falling — velocity is invisible. → stack the last 4 frames so the network can infer motion. And skip frames: repeat each chosen action for 4 game frames (cheaper, and one decision then maps to one meaningful change).

Markov: a state is Markov if it contains everything needed to predict the future — no memory of the past required. One Mario frame isn’t (no velocity); four stacked frames approximately are.

So before the agent learns anything, you write and tune four wrappers (resize, grayscale, stack, skip) that turn (240,256,3) into (4, 84, 84).

Procgen hands you (64, 64, 3) RGB — and you feed it straight to the CNN. No resize (already small), no grayscale (color is semantic here: keys, walls, and enemies are color-coded — graying it out destroys information), no stacking and no frame-skip (the games are designed to be ~fully observable from a single frame). The preprocessing you skip isn’t laziness — it’s the standard, so your results stay comparable to the published Procgen numbers.

Part 2 — The emulator tax: why I could run 8 Mario envs but 64 Procgen envs

RL is data-hungry, so you run many copies of the environment in parallel (a vectorized environment) to feed the network a big batch each step. How you parallelize depends entirely on what the environment is.

Mario runs on an NES emulator. Under the hood, gym-super-mario-bros drives nes-py, a program that simulates the console — a CPU-bound chunk of work, one instance per environment. To run many in parallel, “parallel” has to mean N separate operating-system processes (a subprocess vectorized env): each environment is its own process, and on every step you pickle the observation (a whole image) and push it through an OS pipe back to the learner.

IPC (inter-process communication): moving data between separate processes. Pickling an image and sending it over a pipe every step is pure overhead — it does no learning, it just moves bytes.

That overhead is brutal: on my M1, roughly two-thirds of the CPU was spent in the kernel shuffling observations between processes, not stepping games or training. It also caps how many envs you can afford. I ran 8.

Procgen is a native C++ environment that batches internally. ProcgenGym3Env(num=64) steps 64 levels inside one process, writing all 64 observations straight into a single array — no per-env process, no pickling, no pipes. One call in, a batch out. I jumped straight to 64 environments feeding one batched GPU forward pass, and throughput roughly tripled.

Why care? Deep RL alternates two phases: step the environment (CPU) and run the network (GPU). If env-stepping is slow, your expensive GPU sits idle waiting for data — you’re CPU-bound, paying for a GPU that twiddles its thumbs. The whole point of modern RL infra is to keep the GPU fed; an emulator-one-process-per-env starves it, a native batched env feeds it. Changing the environment did more for my throughput than any code optimization could.

Part 3 — Train and test: the question Mario can’t answer

Here’s the scientific reason, and it’s the big one.

Mario 1-1 is a single, fixed, hand-authored level. You train on it and evaluate on it — the same level. So when the agent finally clears it, you genuinely cannot tell whether it learned to play or memorized this exact layout. There is no held-out level to check against. (Mario has a “random-stages” mode, but it’s a small hand-made pool, not a real protocol.)

Procgen procedurally generates its levels from a seed. So you can train on a fixed set — say 200 levels — and then test on levels the agent has never seen (the rest of an effectively infinite distribution).

Nine CoinRun levels from nine different seeds — same rules, different layouts. Train on some, test on the rest, and the score difference tells you how much the agent actually learned versus memorized.

generalization gap: train score − test score. How much of the performance is transferable skill versus memorization of the specific training levels.

That gap is a single, honest number, and Mario simply can’t produce it. Measuring generalization is the entire reason Procgen exists (Cobbe et al., 2020). It’s what turns “my agent got good at the game” into “my agent learned something that transfers” — which is the question the rest of this project is actually about.

So: two goals, one environment change

• Not CPU-bound — native in-process batching took me from 8 → 64 environments and kept the GPU fed instead of idling behind an emulator.
• Test generalization — a train/test split over procedurally-generated levels gives a real generalization gap, which a single fixed level never could.

There was one catch: none of this mattered until Procgen actually built on my Apple-Silicon laptop — which, out of the box, it didn’t. Getting past that was its own small adventure in refusing to route around a broken build. That’s the next post: the last mile of fixing Procgen on Apple Silicon.

Code: rl_ablations/models/knn, built on the rl-factory engine.

This is the sequel to rnd_on_features_and_byol.md. There we tried rndf — RND novelty measured on the policy’s learned feature φ instead of raw pixels — and it collapsed: the predictor MLP fit the frozen random target across the compact 512-d φ-manifold almost instantly (distill_loss 36 → 0.007 within 200k steps), so prediction error → 0 everywhere and the novelty signal died by step ~20k.

The diagnosis pointed at a fix the literature already uses (NGU / Agent57): stop predicting, start counting. This tutorial explains the knn model (rl_ablations/models/knn/) — why a memory/count signal structurally cannot saturate the way distillation did, how the NGU novelty kernel works, and what to watch for. Every new term is defined the first time it appears.

One-sentence version: novelty = how empty the neighbourhood around this state is, in a memory of where you’ve been this episode — and since each episode starts with an empty memory, novelty is always available, no matter how long training has run.

Reference: Badia et al., Never Give Up: Learning Directed Exploration Strategies, ICLR 2020 (papers/ — the episodic-novelty half of NGU).

Part 0 — Why distillation saturated, in one picture

RND-style novelty is prediction error: a predictor chases a target, and error = “I haven’t learned this yet.” That works only while the predictor can’t catch up. On raw pixels (84×84×4, a huge space, CNN target) it never fully catches up, so novelty persists. On a 512-d learned feature, the predictor fits the whole manifold in a few thousand steps → error ≈ 0 everywhere → no novelty left. A learned predictor is a moving target that eventually stops moving.

The escape: use a signal whose “freshness” doesn’t depend on a predictor that can converge.

Part 1 — Counting instead of predicting

Imagine you could count how many times you’d visited each state. A perfect visit count N(s) gives novelty 1/√N(s) — high for rarely-seen states, low for common ones — and it never saturates: a state you’ve never visited has N(s)=0 forever, however long you train. Count-based exploration is the gold standard; the problem is you can’t literally count continuous, high-dimensional states (you never see the exact same pixel frame twice).

kNN episodic novelty is an approximate, soft count: instead of “how many times have I been at exactly s”, ask “how crowded is the neighbourhood around φ(s) in my memory of recent states?” Many close neighbours ≈ high visit count (familiar); few or far neighbours ≈ low count (novel).

memory here is a running set M of the feature vectors φ(s) the agent has visited. Novelty reads M; it doesn’t train anything. In our code it’s a per-env ring buffer, EpisodicMemory (rl_ablations/models/knn/memory.py).

Part 2 — The “episodic” part is what makes it not saturate

The memory is reset at every episode boundary. So M only ever holds states from the current episode, and novelty measures “have I been near here this episode”, not ever.

This is the crux. A brand-new episode starts with M empty → the first states are maximally novel → the agent is paid to reach fresh ground every episode, forever. There is no predictor to converge, no target to match, nothing that decays with training age. Contrast the two failure clocks:

In the code this is the dones handling in EpisodicKNNModule.intrinsic_reward: on a done step the env’s memory is reset(), so the next episode scores against an empty buffer.

Part 3 — The NGU novelty kernel (the actual formula)

For an arrived state’s feature φ, novelty against memory M:

1. find the k nearest neighbours of φ in M           (k = 10)         → squared distances d²_1..d²_k
2. normalise each by a running mean of kNN distances  d̄²             → d²_i / d̄²
3. kernel per neighbour:   K_i = ε / (d²_i/d̄² + ε)   (ε = 1e-3)      → ~1 if very close, ~0 if far
4. similarity:             s   = √(Σ_i K_i)
5. reward:                 r   = 1 / (s + c)          (c = 1e-3), capped at reward_max

Read it as: Σ K_i is a soft neighbour count. Crowded neighbourhood → large Σ K → large s → small reward (familiar). Empty/far neighbourhood → Σ K → 0 → s → 0 → large reward (novel). The running d̄² (an EMA of the kNN distances) makes the kernel scale-free, so it works whatever the raw magnitude of φ — that normaliser, not per-dimension whitening, is how NGU handles φ’s scale. The c and the cap keep a near-empty memory from producing an unbounded spike. Code: EpisodicKNNModule._novelty.

As with the other curiosity modules, the raw reward is then η-scaled and EMA-normalised to ~unit scale (RewardNormalizer) before it’s folded into the reward ahead of GAE — so knn drops into the exact same PPO pipeline as rnd/rndf, only the bonus differs.

Part 4 — What’s different in the plumbing (and one honest caveat)

Two things set knn apart from rnd/rndf mechanically:

• No trainable curiosity network. Novelty is a distance query, so there’s no optimizer, no distill_loss. learn() does no gradient step; it reports mem_size (mean memory occupancy) as a health signal instead. The policy encoder is still only read (φ is detached) — PPO owns it.
• Stateful and order-dependent. intrinsic_reward must be called once per rollout, in temporal order, because it accumulates each visited φ into its env’s memory and resets at dones. The memory persists across rollouts (an episode spans many 128-step rollouts). That’s why the module reshapes the flattened (T·N,) transitions back to (T, N) and loops per-env, per-step — unlike the stateless, shuffled batching rnd/rndf can use.

The caveat to watch for: pure episodic novelty can reward resetting the episode. Dying clears the memory, and a fresh memory means fresh novelty — so an agent can learn that dying, then re-exploring the same early states, is “novel” and pays. NGU guards against this with a lifelong novelty modulator (an RND-on-pixels term that stays low in globally-familiar regions), which we’ve deliberately left out to isolate the episodic mechanism. If the pilot shows the agent farming deaths (short episodes, mem_size never growing, ep_max_x stuck low), that’s this failure — and the fix is to add the lifelong modulator, i.e. full NGU.

Part 5 — What we expect, and how we’ll read it

Matched pure-curiosity A/B, third arm alongside rnd_pc (pixels) and rndf (learned-feature distillation): same knobs (extrinsic 0.0, entropy 0.02, η 0.5, seed 0, 1M steps, 8 actors), novelty mechanism the only difference.

• The hypothesis: intrinsic_reward stays alive for the whole run (no crater to ~0), because a count can’t saturate — directly fixing the rndf failure.
• mem_size should grow within episodes and drop at resets (memory is doing its job).
• Success looks like healthy entropy and ep_max_x that climbs as curiosity drives coverage.
• The two failure modes to distinguish: (a) death-farming (Part 4 caveat) → needs the lifelong modulator; (b) distances-uninformative-in-φ → novelty noisy but not dead.

Code map

Prior tutorials in this arc: curiosity_icm_on_ppo.md → icm_vs_rnd.md → rnd_on_features_and_byol.md → this. The empirical thread (ICM saturates → RND fixes it → rndf collapses → knn) is in the curiosity-exploration-findings project note.

Results on Mario 1-1

Here are all four novelty mechanisms side by side in pure-curiosity mode — RND (distill on pixels), rndf (distill on learned φ), knn (count on the critic’s φ), and knn_idf (count on an inverse-dynamics φ):

The two distillation signals behave as predicted — rndf’s intrinsic reward flatlines near zero (its target got easy), and even RND-on-pixels drifts. The count-based kNN signals stay alive the whole run, because each episode starts with an empty memory: novelty is always available, no matter how long training has run. That’s the structural win — a novelty signal that can’t die.

Where the series pauses

That closes the arc I set out to walk — DQN → PPO → ICM → RND → kNN — each model earning its place by fixing a limitation of the last, all sharing the one rl-factory spine. The natural next step, learning the representation φ on purpose rather than borrowing the critic’s, is where this turns into open research — a story for another day.

Code: rl_ablations/models/rnd, built on the rl-factory engine.

You’ve read the ICM tutorial (curiosity_icm_on_ppo.md), so you know the big idea: to make an agent explore, hand it a bonus reward for surprise — for reaching states it couldn’t predict. Then we added a second curiosity method, RND, and the natural question is: is RND just ICM with a tweak?

No. They share one sentence of philosophy (“prediction error = novelty”) and nothing else of substance. This tutorial explains — for someone who has seen ICM but not RND — exactly where they diverge, why RND came after ICM but is simpler, and why that simplicity is what saved our Mario runs. Every new term is defined the first time it appears. The code lives in rl_ablations/models/icm/ and rl_ablations/models/rnd/; both plug into the same PPO training loop.

If you remember one sentence: ICM asks “can I predict what my action does here?”; RND asks “have I seen a state like this before?” — a dynamics question versus a visitation question.

• ICM: Pathak, Agrawal, Efros & Darrell, Curiosity-driven Exploration by Self-supervised Prediction, ICML 2017 (papers/curiosity_driven_exploration_pathak2017/).
• RND: Burda, Edwards, Storkey & Klimov, Exploration by Random Network Distillation, ICLR 2019 (papers/random_network_distillation_burda2018/).

Part 0 — The one thing they agree on

Both methods compute an intrinsic reward (a bonus the agent invents for itself, on top of the game’s extrinsic reward) as a prediction error: build a predictor, measure how wrong it is, pay the agent proportionally to the wrongness. High error = “I don’t understand this yet” = novel = go there.

In both of our implementations that idea shows up as the same two-method interface:

intrinsic_reward(transitions) -> per-transition bonus, (T·N,)   # score novelty
learn(transitions)            -> train the predictor            # make today's states cheap tomorrow

The total reward the policy optimizes is r = r_extrinsic + η · r_intrinsic in both (η is intrinsic_scale in each config). And in both, the raw error is EMA-normalized by a RewardNormalizer so the bonus stays ~unit-scale as the predictor improves.

That’s the entire overlap. Everything below is a difference.

Part 1 — The core split: what do you predict?

This is the whole ballgame. Both predict something and pay you the error. They predict different things.

ICM predicts the future — a dynamics model

ICM’s predictor is a forward model: given the current state’s features and the action you took, predict the next state’s features.

forward model: φ̂(s’) = f(φ(s), a). “If I’m here and press A, where do I end up?” ICM intrinsic reward: ‖φ̂(s’) − φ(s’)‖² — how wrong that guess was.

The signal answers: “Are the consequences of my actions still surprising here?” It’s inherently about transitions — it needs (state, action, next_state). Look at rl_ablations/models/icm/module.py: intrinsic_reward encodes both states and next_states, takes actions, and computes _forward_error(phi_s, actions, phi_next).

RND predicts a fixed random function — a visitation counter

RND has no future, no action, no dynamics. It sets up two networks over the current state only:

• a target network: randomly initialized weights, then frozen forever. It’s an arbitrary, meaningless-but-fixed function f(s).
• a predictor network: trained to reproduce the target’s output, f̂(s) ≈ f(s), on the states the agent actually visits.

RND intrinsic reward: ‖f̂(s) − f(s)‖² — how badly the predictor matches the frozen target at this state.

Why is that novelty? Because a neural network only gets good at inputs it’s been trained on. If the agent has visited states like s many times, the predictor has had many gradient steps there, so it matches the target closely → low error → familiar. If s is somewhere new, the predictor has never been fit there → high error → novel. The error is, in effect, a soft visit counter: “how many times have I been near here?” See rl_ablations/models/rnd/model.py — the target’s parameters get requires_grad_(False) and the predictor chases it.

This is the extension question, answered concretely: RND does not extend ICM’s forward model — it deletes it. Gone: the action input, the next-state prediction, the whole notion of dynamics. RND is strictly less machinery than ICM, not more.

Part 2 — The encoder: ICM’s clever bit that RND doesn’t need

ICM’s subtlety was where it predicts. Predicting raw pixels invites the noisy-TV problem: a patch of random flicker (TV static, swaying leaves) is forever unpredictable, so a pixel-predictor pays curiosity forever for staring at noise. ICM’s fix is to predict in a learned feature space φ trained by an inverse model — a second head that predicts which action took you from s to s'. To do that job, φ must keep only what the agent can control and discard uncontrollable noise. So the flicker never enters φ, and the noisy-TV trap closes.

That inverse model is real, load-bearing complexity in rl_ablations/models/icm/: an extra network head, a cross-entropy loss, and (as the ICM tutorial’s war story tells) a mandatory detach so the forward loss doesn’t collapse φ’s scale to zero.

RND sidesteps all of it. Its target is a fixed random function of the state, so there is no representation to shape, no inverse model, no encoder-collapse bug to guard against. RND’s only required preprocessing is observation normalization (whiten the pixels with a RunningMeanStd before the nets) — because a frozen random target is unlearnable on raw, arbitrarily-scaled pixels. That’s it. Compare the two module.py files: ICM juggles two losses and a beta blend; RND runs a single distillation MSE.

Part 3 — Why the difference mattered on Mario

This isn’t academic — it’s exactly why we have both in the repo. On our Atari-style Mario (4-stacked frames + frameskip), ICM’s forward model saturates in ~20k steps: consecutive frames are trivially predictable, so forward_loss → ~0.003 everywhere, even in genuinely new regions. The curiosity signal dies, and no amount of tuning entropy_coef revives a reward that has gone to zero. Pure-curiosity ICM gets Mario stuck around x≈400.

RND doesn’t have a “next frame” to find predictable, so it can’t saturate that way. Its error tracks visitation, which keeps rising as new areas appear. In our 80k-step pure-curiosity run (game reward off), RND’s intrinsic_reward stayed 0.04–0.13 and rose on reaching new areas, entropy stayed healthy, and Mario reached x≈1200 on curiosity alone. The mechanism that saved us is precisely the ICM machinery RND removed — dynamics prediction was the thing that saturated.

Takeaway for the exam: “RND is simpler than ICM” isn’t a footnote — the simplification (drop the forward model, count visitation instead) is the fix for the failure mode that sank ICM here.

Part 4 — So, sibling or successor?

Chronologically RND (2019) came after ICM (2017), from an overlapping group of authors, and was partly motivated by ICM’s known weaknesses (noisy-TV, and the saturation we hit). But it is not built on ICM — it’s a different, leaner answer to the same question. If you must draw a family tree:

• Shared parent: “intrinsic reward = prediction error” (curiosity as self-supervised surprise).
• ICM’s branch: predict dynamics in a learned, control-aware feature space. Powerful, but more moving parts and prone to saturating when dynamics are easy.
• RND’s branch: predict a fixed random function of the raw (normalized) state. A soft visit counter. Fewer moving parts, robust to easy dynamics, but tells you nothing about why a state is novel (no controllability signal).

Neither is strictly better. ICM’s controllability filter is genuinely useful when uncontrollable noise dominates the observation; RND’s visitation signal is more robust when the environment’s dynamics are too easy to stay surprising. In this repo they’re peers — two PPOLearner wrappers behind one interface — which is exactly how you should think of them.

Where to look in the code

Both are training-only (they never touch acting or evaluation — the policy is a plain PPO PolicyAgent) and neither is checkpointed. Run either with its config in configs/icm.yaml / configs/rnd.yaml; set extrinsic_reward_scale: 0.0 (RND) or reward_scale: 0 (ICM) to see pure curiosity, which is the only honest test that the intrinsic signal actually works.

Results on Mario 1-1

Tested as pure curiosity, the design question becomes empirical: put the novelty split on raw pixels (RND) versus on the policy’s learned features φ (rndf). The feature-space variant collapses — the predictor fits the frozen random target across the compact φ-manifold almost instantly, so prediction error → 0 everywhere and the novelty signal dies within ~20k steps. RND-on-pixels stays alive:

That collapse is the whole lesson: any distillation-based novelty can saturate once its target becomes easy to predict. Which raises the obvious question — is there a novelty signal that structurally cannot saturate? There is: stop predicting, start counting → kNN episodic novelty.

Code: rl_ablations/models/icm, built on the rl-factory engine.

This tutorial explains the Intrinsic Curiosity Module (ICM) we added to the project — what curiosity means for a reinforcement-learning agent, how we bolted it onto our existing PPO agent, and the story of a real bug where the curiosity signal quietly collapsed to zero and how we fixed it. It assumes you’ve read the PPO tutorial (policy_gradients_reinforce_to_ppo.md) or at least know that PPO learns a policy π(a|s) by nudging it toward actions that beat expectations. Every new term is defined the first time it appears. The code lives in rl_ablations/models/icm/ (the curiosity model) and rl_factory/training/rollout.py + rl_factory/training/on_policy.py (where it plugs into training).

If you only remember one sentence: we give the agent a bonus reward for reaching states it couldn’t predict, computed in a learned feature space, and the whole trick is making sure that “surprise” signal doesn’t accidentally shrink to nothing.

The paper is Pathak, Agrawal, Efros & Darrell, Curiosity-driven Exploration by Self-supervised Prediction, ICML 2017 (papers/curiosity_driven_exploration_pathak2017/).

Part 0 — Why an agent needs curiosity

Our PPO agent learns from the game’s reward — roughly, how far right Mario moved. That works because the reward is dense: almost every step gives some signal. But many interesting problems have sparse rewards: you get nothing until you reach a far-off goal. Imagine Mario only scored at the flagpole. A policy that explores randomly would almost never reach it, so it would almost never see a reward, so it would never learn. This is the exploration problem.

Curiosity is one answer. The idea: manufacture your own reward for exploring. On top of the game’s reward (the extrinsic reward, “from outside”), add an intrinsic reward (“from inside the agent”) that is high when the agent encounters something novel — something it hasn’t figured out yet. The agent then seeks out novelty on its own, which spreads it across the level even when the game is stingy with points.

The famous demo from the paper: with the extrinsic reward turned completely off, a curiosity-only Mario still learns to run right, jump gaps, and dodge enemies — purely because doing so keeps producing states it can’t predict.

The total reward the policy learns from is r = r_extrinsic + η · r_intrinsic, where η (eta) is a knob for how much the curiosity bonus counts. In our code η is ICMConfig.intrinsic_scale.

Part 1 — What counts as “novel”? Prediction error.

How does the agent know a state is novel? It tries to predict what happens next, and measures how wrong it is. If the agent can accurately predict the next state, that situation is familiar — boring. If its prediction is badly wrong, the situation is surprising — worth exploring. So:

intrinsic reward = how badly the agent predicted the next state.

The naive version: predict the next screen of pixels from the current screen and the action. This fails badly, for a reason the paper calls the noisy-TV problem. Suppose there’s a patch of the screen full of random flicker (TV static, swaying leaves, a particle effect). Pixel prediction there is always wrong, so the agent gets permanent curiosity reward for staring at noise. It becomes a couch potato hypnotized by static, never exploring the actual level.

The fix is the heart of ICM: don’t predict pixels. Predict in a learned feature space that only keeps the parts of the world the agent can control or affect. If the feature space throws away the flickering noise (because the agent can’t influence it), then noise contributes no prediction error, and the noisy-TV trap disappears.

Call that feature space φ (phi): a neural network — an encoder — that maps a raw state s (four stacked 84×84 game frames) to a vector of 512 numbers, φ(s). The rest of ICM is about (a) shaping φ to keep only controllable stuff, and (b) using φ to measure surprise.

Part 2 — The three networks of ICM

ICM is one network with three parts (rl_ablations/models/icm/model.py, class ICMNet):

1. The encoder φ. Convolutional layers (the same “Nature-DQN” torso our PPO agent uses) ending in a 512-number feature vector. This is the learned feature space.

2. The inverse model. Given the features of two consecutive states, φ(sₜ) and φ(sₜ₊₁), it tries to predict which action aₜ the agent took to get from one to the other. Think about what this forces φ to learn: to name the action, φ must encode the things that changed because of the action — Mario’s position, velocity, whether he jumped. It has no reason to encode the swaying background, because the background doesn’t tell you which button was pressed. Training the inverse model is what shapes φ to keep only controllable, action-relevant information. This is the anti-noisy-TV mechanism. Its loss is a standard classification loss (cross-entropy) over the 7 possible actions.

3. The forward model. Given the current features φ(sₜ) and the action aₜ, it predicts the next features, φ̂(sₜ₊₁) (the hat means “predicted”). The error between the prediction and the real next features,

forward error = ½ · ‖ φ̂(sₜ₊₁) − φ(sₜ₊₁) ‖²          # squared distance between two 512-vectors

is the curiosity reward. High error = the agent couldn’t predict this = novel = explore it.

Notice the two opposite “directions,” which trip everyone up at first:

The combined training loss for the ICM is (1−β)·inverse_loss + β·forward_loss, with β (beta) = 0.2 tilting most of the weight onto shaping φ well. This all lives in CuriosityModule (rl_ablations/models/icm/module.py).

Part 3 — How ICM rides on our PPO, without touching the PPO loop

A design goal in this codebase: the training loop shouldn’t know or care which algorithm it’s running. So we added curiosity without changing a single line of the PPO update or the training strategy’s logic. Here’s the layering.

• The agent (the thing that picks actions) is unchanged — it’s the same PPO PolicyAgent. Curiosity changes only the reward, not how Mario acts. So evaluation and checkpoints use the plain PPO policy; the ICM is a training-time scaffold that’s thrown away afterward.

• The strategy OnPolicyStrategy (rl_factory/training/on_policy.py) just collects experience: it steps the 8 parallel environments, records each transition into a RolloutBuffer, and hands the raw rollout to a learner. It does not compute rewards, advantages, or losses. It talks to the learner through a tiny interface (OnPolicyLearner): “here’s the rollout, give me back an update.” That’s the whole contract.

• For plain PPO, the learner is PPOLearner. For curiosity, the learner is ICMLearner (rl_ablations/models/icm/learner.py), which wraps a PPOLearner and a CuriosityModule. Its one job is to insert the curiosity bonus before handing off to PPO’s math:

  def learn(self, rollout, last_values):
      transitions = rollout.transitions()                      # raw (s, a, s') tuples
      intrinsic = self._curiosity.intrinsic_reward(transitions)  # the bonus, per step
      rollout.compute_advantages(..., intrinsic=intrinsic)     # fold bonus into the reward, then GAE
      ppo_metrics = self._ppo.update(rollout)                  # PPO's clipped update — reused as-is
      icm_metrics = self._curiosity.learn(transitions)         # train the ICM for next time
      return merged_metrics
  

Because the bonus is added to the reward before GAE (the advantage calculation), PPO treats it exactly like any other reward — no special cases. Adding a different exploration method later (say RND) is “write a new learner,” never “edit the training loop.” That’s the payoff of keeping the strategy generic.

One subtlety about “next state.” The 8 environments auto-reset when Mario dies, so the frame returned after a death is actually the first frame of a new life, not a real successor. Feeding that bogus (s, a, s’) to the ICM would be garbage, so those transitions are masked out of both the reward and the ICM’s training. (Search dones / valid in module.py.)

Part 4 — We ran it, and the curiosity was… zero

We launched a full training run (logs/icm_training.log, config configs/icm.yaml). Every 10,000 steps the trainer logs the metrics. Here is what the curiosity numbers did:

The intrinsic_reward printed as 0.0000 — the curiosity bonus had shrunk to essentially nothing. The agent was, in effect, running plain PPO with a useless dangling term. Curiosity was dead.

Two clues in that table:

• inverse_loss was fine (~0.7–1.3, well below the ~1.95 you’d get from random guessing over 7 actions). So φ was not completely broken — it still encoded enough to name actions.
• But forward_loss fell all the way to ~0. The forward model had become a perfect predictor. If you can predict the next state perfectly, your surprise — and thus your curiosity reward — is zero, forever.

So the question became: why did the forward model’s error collapse to zero?

Part 5 — Three reasons the signal died

Reason 1 (the big one): the forward model was allowed to cheat by shrinking the features.

Here’s the trap. The forward error is the distance between two feature vectors, φ̂(s') and φ(s'). Our first implementation trained the encoder φ from both the inverse loss and the forward loss. But think about how you can make the forward error small. One way is to get better at predicting. A lazier way is to make the features themselves tiny — if φ outputs numbers close to zero for every state, then φ̂(s') and φ(s') are both ≈0, their distance is ≈0, and the forward loss is ≈0 without predicting anything at all. The forward loss, given a gradient path into the encoder, quietly pushed φ to collapse toward a small, low-variance blob. The distances vanished, and so did the curiosity. This is a classic ICM failure mode.

The paper’s own logic already told us the fix: the inverse model is supposed to be the only thing that shapes φ. The forward model should be a customer of φ, not an author of it. So we detach the features before they enter the forward model — “detach” means “use these numbers but don’t send training gradients back through them.” Now the forward model can only improve by predicting better, never by shrinking φ, and φ’s scale is anchored by the inverse loss. In rl_ablations/models/icm/module.py:

pred_phi_next = self._net.predict_next(phi_s=phi_s.detach(), actions=actions)   # .detach() = don't train φ here
return 0.5 * (pred_phi_next - phi_next.detach()).pow(2).sum(dim=1)

Reason 2: we divided the signal by 512. Our first version averaged the squared error over the 512 feature dimensions (.mean()). The paper sums them (.sum()). Averaging made the raw number ~512× smaller than it should be — starting us off at a tiny 1e-4 before anything else went wrong. We switched to .sum(dim=1).

Reason 3: raw prediction error is the wrong unit, and it’s not comparable to the game reward. Even a “healthy” forward error is some arbitrary number in feature-distance units. Meanwhile the extrinsic reward, after scaling, is around ~0.8 per step. A curiosity bonus of 0.0001 next to an extrinsic reward of 0.8 is a rounding error — the policy can’t even feel it. And as the ICM learns, the raw error keeps drifting, so no fixed η would keep it balanced.

The fix is normalization, and it comes straight from this repo’s companion paper (Burda et al. 2018, Large-Scale Study of Curiosity-Driven Learning, papers/large_scale_study_curiosity_burda2018/): divide the intrinsic reward by a running estimate of its standard deviation. This throws away the meaningless absolute scale and keeps only the relative structure — “which states are more surprising than the others right now” — rescaled to roughly unit size. Then η multiplies something scale-stable, so it actually means what you think it means. We keep that running standard deviation with a small streaming statistics helper, RunningMeanStd, updated every rollout:

error = self._forward_error(...).cpu().numpy()          # raw surprise, per step
self._reward_rms.update(error[valid > 0])               # update running std (skip auto-reset steps)
reward = self.config.intrinsic_scale * error / (self._reward_rms.std + 1e-8)   # normalized × η

A wrinkle we found while testing: the first rollout has almost no history, so the running std is tiny and the first normalized reward spikes (and briefly inflates PPO’s value loss). It self-corrects within a few rollouts as the statistics warm up. If it ever mattered, you’d warm up the std or clip the reward — for us it settles on its own.

Part 6 — After the fix

Same short run, with all three fixes in place:

The curiosity reward is now alive — three-plus orders of magnitude larger than before — and, crucially, it moves: when the agent pushes into unfamiliar territory the forward error and the bonus rise together, which is exactly the behavior we want. inverse_loss stays healthy, meaning φ is still learning controllable features rather than collapsing.

With the extrinsic reward left on (our default), curiosity is a secondary bonus — a nudge toward unexplored areas layered on the main “go right” signal. To reproduce the paper’s headline no-reward experiment, set PPOConfig.reward_scale = 0: the game gives zero points and curiosity becomes the only reward. That’s the real test of whether the intrinsic signal alone can drive Mario down the level.

Part 7 — What to watch, and the knobs

Watch these columns in logs/icm_training.log (all live in ICMMetrics / ICMLearnerMetrics):

• intrinsic_reward — the average curiosity bonus handed to the policy. Should be nonzero and wobble as exploration ebbs and flows. If it’s 0.0000, curiosity is dead (that was the bug).
• inverse_loss — how well the inverse model names the action. Should sit clearly below ~1.95 (random). If it climbs back to ~1.95, φ has collapsed.
• forward_loss — the ICM’s own prediction error while training. Healthy is “small but not exactly zero.”
• ep_max_x / ep_return — how far Mario gets and his score. The bottom line: is curiosity helping him explore further than plain PPO? (Compare against logs/ppo_training.log.)

The tuning knobs, all in rl_ablations/models/icm/config.py (ICMConfig) unless noted:

• intrinsic_scale (η) — how loud curiosity is. Now that the reward is normalized, this is meaningful: ~0.1 for a gentle bonus, higher to make exploration dominate.
• beta — inverse-vs-forward loss mix. Higher = train the forward model harder; lower = spend more on shaping φ.
• feature_dim, lr, epochs, num_minibatches — the ICM network’s size and how hard we train it each rollout.
• reward_scale (in PPOConfig) — set to 0 for pure, no-extrinsic curiosity.

The one-paragraph recap

Curiosity gives the agent a bonus for surprise, measured as how badly it predicts the next state — but done in a learned feature space φ that (thanks to an inverse model predicting the agent’s own actions) keeps only what the agent can control, dodging the noisy-TV trap. A forward model’s prediction error in that space is the bonus. We layered this onto PPO purely as a new learner (ICMLearner) that folds the bonus into the reward before PPO’s normal math, leaving the training loop untouched. Our first run looked fine but the curiosity silently died, because the forward loss was collapsing the features, the error was divided down by 512, and its raw scale was dwarfed by the game reward. Detaching the features (so only the inverse model shapes φ), summing instead of averaging, and normalizing the bonus by its running standard deviation brought the signal back to life.

Results on Mario 1-1

An honest result: on 1-1’s fairly dense reward, curiosity is not a free win. Measuring how far right Mario gets, plain PPO (extrinsic only) pulls ahead, while PPO+ICM and PPO+RND sit lower. Curiosity’s payoff is an exploration mechanism for sparse reward, not a leaderboard boost on an easy level:

So we study curiosity where it’s actually testable — in pure-curiosity mode, with the extrinsic reward switched off, asking only: does the intrinsic signal stay alive and keep driving exploration? ICM has a failure mode there — its forward-prediction signal can collapse as the feature space becomes trivially predictable. That raises the real question — is prediction error even the right novelty signal? — which pushes us to RND → RND.

Code: rl_ablations/models/ppo, built on the rl-factory engine.

This tutorial explains how our PPO agent actually learns — the RolloutBuffer, the advantage calculation, and the loss function — starting from the very first policy-gradient idea (REINFORCE, 1992) and building up historically through A2C to PPO. It assumes you know basic calculus and a little probability, and nothing else. Every term is defined the first time it appears. Wherever we describe code, it lives in rl_factory/training/rollout.py (the buffer

• advantage math) and rl_ablations/models/ppo/learner.py (the loss).

If you only remember one sentence: we nudge the network to make actions that turned out better-than-expected more likely, and we do it carefully enough that we can squeeze several gradient steps out of each batch of experience without the training blowing up.

Part 0 — The problem we’re solving

An agent interacts with an environment (for us, Super Mario Bros). At each timestep it sees a state s (four stacked game frames), picks an action a (press a button), and receives a reward r (roughly, how far right Mario moved). This repeats until Mario dies or finishes the level — one episode.

The agent’s behaviour is a policy, written π(a | s) (“pi of a given s”): a function that, given a state, outputs a probability for each possible action. Our policy is a neural network. “Learning” means adjusting the network’s weights θ (theta) so that the policy collects more total reward.

We measure “total reward” with the return. The return from timestep t, written G_t, is the sum of all rewards from t onward, with later rewards discounted by a factor γ (gamma, between 0 and 1, e.g. 0.99):

G_t = r_t + γ·r_{t+1} + γ²·r_{t+2} + γ³·r_{t+3} + ...

Discounting says “a reward now is worth a bit more than the same reward later.” It also keeps the sum finite for long episodes. The goal of the whole enterprise:

Adjust θ to maximize the expected return — the average G we’d get if we let the policy play many times.

Everything below is a story about how to compute a good gradient — a direction to push θ — for that goal.

Part 1 — REINFORCE (1992): the original idea

1.1 The intuition

Suppose Mario plays a whole episode and gets a great return. Whatever actions he took, they probably weren’t all brilliant — but on average they were good, so let’s make all of them a little more likely next time. If the episode went badly, make those actions less likely. Scale the nudge by how good the outcome was.

That’s REINFORCE in one sentence. Now the math that makes it precise.

1.2 The policy-gradient formula

We want to increase expected return J(θ). The famous policy-gradient theorem says its gradient can be written as an average over experience the policy itself generated:

∇J(θ)  =  E[  Σ_t  ∇ log π(a_t | s_t) · G_t  ]

Let’s decode every piece:

• E[ … ] — “the average over many episodes played by the current policy.”
• Σ_t — sum over every timestep in the episode.
• π(a_t | s_t) — the probability the policy assigned to the action it actually took.
• log π(a_t | s_t) — the natural log of that probability. (Logs show up because of a calculus identity, ∇π = π · ∇log π; they also turn the tiny products of probabilities into stable sums. You don’t need to derive this — just know ∇log π(a_t|s_t) is “the direction in weight space that makes action a_t more likely in state s_t.”)
• G_t — the return that followed. This is the crucial weight: a number, not something we differentiate through.

Read the formula in plain English: for each action taken, step in the direction that raises its probability, scaled by the return that followed. Good returns → push those actions up. Bad returns → the scaling is small (or, once we add a baseline in Part 2, negative) → push down.

1.3 From gradient to “loss”

Deep-learning libraries minimize a loss by gradient descent. We want to maximize J, so we minimize its negative. The quantity whose gradient equals the formula above is:

Loss = − Σ_t  log π(a_t | s_t) · G_t

We compute this number, call .backward(), and the optimizer does the rest. Note G_t is treated as a fixed constant (we detach it from the graph): we are not trying to change the rewards, only the probabilities.

1.4 Why REINFORCE isn’t enough

Two problems, and the rest of history is about fixing them:

1. It’s extremely noisy (high variance). G_t depends on everything that happened after t — every later action, every bit of randomness in the game. The same good action can be attached to a huge return in one episode and a tiny one in the next, purely by luck. Training on such a jittery signal is slow and unstable.

2. It needs complete episodes and fresh data. You can’t compute G_t until the episode ends, and because the average is “over episodes the current policy generated,” the moment you update θ the data you collected is stale and must be thrown away. This is what on-policy means: you may only learn from data your current policy produced. (Contrast DQN, which is off-policy and reuses a giant replay buffer.)

Part 2 — Baselines and the birth of the critic

2.1 Subtracting a baseline

Here’s a small idea with a big payoff. Take the policy gradient and subtract a baseline b(s_t) — any number that depends on the state but not on the action:

∇J(θ)  =  E[  Σ_t  ∇ log π(a_t | s_t) · ( G_t − b(s_t) )  ]

Mathematically this changes nothing about the average (the baseline term averages out to zero, because probabilities sum to 1). But it can dramatically reduce the variance — the jitter — if we choose b well. Intuitively: instead of asking “was the return big?”, we ask “was the return bigger than what we normally get from this state?” That comparison is far more stable than the raw return.

2.2 The best baseline is the value function

The natural choice of baseline is the value function V(s): the expected return if you start in state s and follow the current policy. “From this spot, how well do we usually do?”

We don’t know V, so we learn it too, with a second neural network (or, in our code, a second head on the same network — see ActorCriticNet in rl_ablations/models/ppo/model.py). This second learner is called the critic; the policy is the actor. The critic doesn’t choose actions; it just predicts how good states are, to serve as the baseline.

Now G_t − V(s_t) has a name and a meaning: the advantage.

Advantage A(s_t, a_t)  ≈  G_t − V(s_t)
                       =  (what actually happened)  −  (what we expected)

• A > 0: this action did better than average from this state → make it more likely.
• A < 0: it did worse than average → make it less likely.

The advantage is the honest “was this action a good idea?” signal. From here on, the policy gradient is:

Loss_policy = − Σ_t  log π(a_t | s_t) · A_t

Same shape as REINFORCE, but with the advantage in place of the raw return. This is the heart of every method that follows.

Part 3 — A2C (Advantage Actor-Critic): the practical loop

A2C is the algorithm that turns the above into something you can actually run efficiently. It contributes three practical ideas, and all three are visible in our RolloutBuffer and OnPolicyStrategy.

3.1 Bootstrapping: don’t wait for the episode to end

Waiting for a full episode to compute G_t is slow and noisy. Instead, use the critic to estimate the tail. The simplest advantage estimate uses just one real reward plus the critic’s guess for the rest:

δ_t  =  r_t  +  γ·V(s_{t+1})  −  V(s_t)

This δ_t (delta) is the TD error (“temporal-difference error”). Read it as: “the reward I actually got, plus what the critic thinks the next state is worth, minus what the critic thought this state was worth.” If that’s positive, things went better than the critic expected. Using a prediction of the future to improve an estimate of the present is called bootstrapping. It lets us learn from short slices of experience instead of whole episodes.

3.2 Fixed-length rollouts across many parallel environments

Rather than one episode at a time, A2C runs N environments in parallel and collects a fixed number of steps T from each — a rollout of shape T × N. This is exactly what our RolloutBuffer stores, and what OnPolicyStrategy.step() fills by stepping all N envs in lockstep (rl_factory/training/on_policy.py). More envs = more, less-correlated data per update = steadier gradients. We’ll come back to how episodes of different lengths fit into a fixed T × N grid in Part 5 — it’s the question that trips everyone up.

3.3 An entropy bonus to keep exploring

A policy can prematurely become over-confident — always pressing the same button — and stop discovering better strategies. To fight this, A2C adds an entropy bonus. Entropy is a measure of how spread-out a probability distribution is: high when the policy is unsure (spreads probability across actions), low when it’s certain (all probability on one action). We add a small reward for keeping entropy high, which gently encourages the policy to keep its options open. In the loss it appears as a term we subtract (because we minimize loss but want to increase entropy).

3.4 Training the critic

The critic is trained by plain regression: make V(s_t) predict the actual return G_t (in practice, the bootstrapped target). The loss is the mean squared error between prediction and target — literally average of (V(s_t) − target_t)². As the critic gets better, the advantages get more accurate, which makes the actor’s updates better — a virtuous cycle.

At this point we essentially have A2C. PPO changes only two things: a smarter advantage estimate (GAE, Part 4) and a safety mechanism that lets us reuse each rollout several times (clipping, Part 6).

Part 4 — GAE: dialing the bias/variance knob

4.1 The trade-off

We now have two ways to estimate the advantage, and they sit at opposite extremes:

• One-step TD (δ_t from §3.1): uses one real reward and trusts the critic for everything after. Low variance (only one step of randomness), but biased — if the critic is wrong, the estimate is wrong.
• Full Monte Carlo (G_t − V(s_t), the REINFORCE-style version): uses all the real rewards to the episode’s end. Unbiased, but high variance — it drags in all that future noise.

“Bias” = systematically off in one direction (a bad-but-consistent ruler). “Variance” = jittery, different each time. You’d like a knob to trade one for the other. Generalized Advantage Estimation (GAE) is that knob.

4.2 The GAE formula

GAE blends all the multi-step estimates with an exponentially decaying weight controlled by a new parameter λ (lambda, between 0 and 1):

A_t  =  δ_t  +  (γλ)·δ_{t+1}  +  (γλ)²·δ_{t+2}  +  (γλ)³·δ_{t+3}  +  ...

• λ = 0: only δ_t survives → one-step TD (low variance, more bias).
• λ = 1: everything survives with only γ decay → Monte Carlo (unbiased, high variance).
• λ ≈ 0.95 (our default): mostly trusts a few real steps, leans on the critic for the far future. In practice this is the sweet spot.

4.3 How the code computes it — the backward loop

You would never sum that infinite series directly. There’s a beautiful trick: the same series can be computed with a single backward pass, because each step’s advantage is just its own TD error plus a discounted copy of the next step’s advantage. Here is the actual code from compute_gae in rl_factory/training/rollout.py, annotated:

for t in reversed(range(horizon)):          # walk backward through the T steps
    nonterminal = 1.0 - dones[t]            # 0 if the episode ended at step t, else 1
    delta = rewards[t] + gamma * next_values * nonterminal - values[t]   # δ_t
    last_advantage = delta + gamma * gae_lambda * nonterminal * last_advantage  # A_t
    advantages[t] = last_advantage
    next_values = values[t]                 # for the next (earlier) iteration

Walking backward, last_advantage carries the accumulated (γλ)·δ + (γλ)²·δ + … from all later steps, so each line adds one fresh δ_t and decays the rest — the whole infinite sum, in one cheap loop.

The nonterminal = 1 − dones[t] factor is the important safety catch. On the step where an episode ended, it becomes 0, which does two things at once: it drops the γ·V(s_{t+1}) term (you must not bootstrap the value of a state in a different episode) and it zeroes the carried-over last_advantage (rewards from after a reset must not leak backward into the previous episode). This is what lets many independent episodes share one array safely.

4.4 The value target falls out for free

At the end, compute_gae also returns the returns, computed as:

returns = advantages + values

Rearranged, that’s return_t = A_t + V(s_t) — an improved estimate of G_t. These are the regression targets the critic is trained against (§3.4). So one backward pass produces both things the loss needs: the advantages (for the actor) and the returns (for the critic).

Part 5 — The RolloutBuffer, concretely

Now we can read the buffer (rl_factory/training/rollout.py) end to end. It is the on-policy counterpart of a replay buffer: where DQN keeps a big pool of old transitions to sample randomly, the RolloutBuffer holds exactly one fresh batch of experience, learns from it, and throws it away — because on-policy methods may only learn from current-policy data (Part 1).

It has a strict three-phase life:

Phase 1 — Accumulate (add)

During collection, OnPolicyStrategy steps all N envs and calls add(...) once per timestep, appending that step’s slice to per-field lists. Per step it records, for all N envs:

Phase 2 — Finalize (compute_advantages)

Once T steps are collected, this stacks the lists into (T, N) arrays, runs compute_gae (Part 4) to get advantages and returns, and flattens everything from (T, N, …) to (T·N, …). After this the buffer is just a flat bag of T·N independent training examples, each a tuple (state, action, old_log_prob, advantage, return) — that’s the RolloutBatch NamedTuple.

Phase 3 — Serve (minibatches)

Yields shuffled minibatches (small random subsets) as tensors on the GPU. The learner asks for these repeatedly — several passes over the same rollout — which is where PPO’s clipping becomes essential (Part 6).

5.1 The question everyone asks: different episode lengths

“Each environment’s episode has a different length. How does a fixed T × N buffer cope?”

The answer is that the buffer never stores episodes — it stores a fixed-size time window. Every env is stepped exactly T times per rollout, in lockstep, so the buffer is always a full T × N rectangle, never ragged. Episodes simply don’t line up with the window or with each other, and that’s fine:

• When an env’s episode ends partway through the window, the environment worker auto-resets and immediately returns the first frame of the next episode (see VectorEnv in rl_factory/vector_env.py). That env just keeps producing steps — now belonging to a new episode. One env’s column can contain several episode boundaries; another’s may contain none.
• The only record that a boundary happened is the done flag at that step. And GAE’s nonterminal = 1 − done factor (§4.3) walls off the advantage math at each boundary, so reward never leaks from one episode into another, and each env’s column is handled independently.

Picture T = 5, N = 3 (· = ordinary step, D = episode ended and reset here):

        env0   env1   env2
 t=0     ·      ·      ·
 t=1     ·      D      ·
 t=2     D      ·      ·
 t=3     ·      ·      D
 t=4     ·      ·      ·     ← at the window's edge, these episodes are still unfinished

This whole grid is one rollout. It contains several partial episodes of different lengths, yet it’s a perfect rectangle. Two edges get special handling:

• Episodes still running at t = 4 (the window’s end): we can’t compute their tail, so we ask the critic for one last estimate, V(s_T), and seed the backward pass with it (the last_values argument to compute_gae, obtained via agent.bootstrap_values). That’s why we can learn from a slice of an unfinished episode.
• Episodes that ended inside the window: the done mask zeroes the bootstrap there, as described.

After Phase 2 flattens and Phase 3 shuffles, episode identity is gone entirely — every sample is an independent (state, action, old_log_prob, advantage, return). That’s safe because everything temporal was already baked into advantage and return during GAE. The gradient step doesn’t need to know which env or episode a sample came from.

Part 6 — PPO: reusing each rollout without blowing up

6.1 The motivation

Plain A2C is wasteful: collect a rollout, take one gradient step, throw it away. Collection (stepping the game) is expensive; we’d love to take several gradient steps per rollout. But there’s a catch rooted in Part 1: the policy-gradient average is only valid for the policy that collected the data. After one update, the policy has moved, and taking another step on the same data is technically using the wrong distribution. Push too far and the update becomes nonsense — the policy can collapse. PPO (“Proximal Policy Optimization,” 2017) is the mechanism that makes reusing a rollout safe.

6.2 The probability ratio

To reuse old data with an updated policy, we reweight each sample by how much more (or less) likely the new policy is to take the action than the old one did:

ratio  =  π_new(a | s)  /  π_old(a | s)  =  exp( new_log_prob − old_log_prob )

(The exp of a difference of logs is just the ratio of the probabilities — cheaper and more stable to compute this way.) This is why the buffer stored old_log_probs back at collection time: it’s the “before” snapshot the ratio needs. A ratio of 1 means “the policy hasn’t changed its mind about this action”; above 1 means “now more likely”; below 1, “now less likely.”

The reweighted objective for one action is ratio · advantage. If we optimized just this, we could still push the ratio arbitrarily far and break things — which is exactly the danger.

6.3 The clip: PPO’s one real idea

PPO caps how far the ratio is allowed to help. It compares the plain objective against a clipped version where the ratio is confined to [1 − ε, 1 + ε] (ε ≈ 0.2), and keeps the more pessimistic (smaller) of the two. Here’s the actual code, from _clipped_policy_loss in rl_ablations/models/ppo/learner.py:

ratio     = torch.exp(new_log_probs - old_log_probs)
unclipped = ratio * advantages
clipped   = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * advantages
return -torch.min(unclipped, clipped).mean()

Why min of the clipped and unclipped versions? It removes any incentive to move the policy too far in the helpful direction:

• Good action (advantage > 0): we want to raise its probability, so the ratio wants to grow. But once the ratio passes 1 + ε, the clipped term stops increasing, and min picks it — so there’s no extra reward for pushing further. The update is content to nudge, not lunge.
• Bad action (advantage < 0): we want to lower its probability. The clip floors the ratio at 1 − ε, and min (of two negative numbers) keeps the penalty from being softened by a runaway ratio — again capping how far one update moves.

The − and .mean() turn it into a loss we minimize, averaged over the minibatch. The net effect: each rollout can be reused for several epochs of minibatch updates (our defaults: epochs passes × num_minibatches each) because no single step is allowed to move the policy far from the one that collected the data. That safety is the entire reason PPO is sample- efficient and stable.

6.4 A drift monitor: approximate KL

As a diagnostic, the learner also computes approx_kl = mean(old_log_prob − new_log_prob), a cheap estimate of how far the policy has drifted from the collector (KL is a standard “distance between two probability distributions”). We only log it here, but it’s the standard signal for “we’ve reused this rollout enough; stop before we drift too far.”

Part 7 — Putting the whole loss together

One PPO gradient step (_gradient_step in rl_ablations/models/ppo/learner.py) combines three terms we’ve now met. In plain words:

total_loss  =        clipped_policy_loss          # Part 6 — make good actions likelier, safely
            +  value_coef  · value_loss           # Part 3.4 — critic predicts returns better
            −  entropy_coef · entropy             # Part 3.3 — keep exploring

and in the actual code (_total_loss):

return policy_loss + config.value_coef * value_loss - config.entropy_coef * entropy

Reading the three pieces:

1. policy_loss — the clipped surrogate from Part 6. Drives the actor. Uses the advantages from GAE and the stored old_log_probs.
2. value_loss — F.mse_loss(values, returns): mean squared error between the critic’s current predictions and the GAE returns. Drives the critic. value_coef (≈ 0.5) sets how much we care about critic accuracy relative to the policy.
3. entropy — the policy’s entropy, subtracted so that minimizing the loss increases it. entropy_coef (a small number like 0.01) sets how hard we push for exploration.

The step itself (_descend) is standard deep learning with one extra guardrail:

self._optimizer.zero_grad()          # clear old gradients
loss.backward()                      # compute new gradients
clip_grad_norm_(..., max_grad_norm)  # cap the gradient size (a second safety belt)
self._optimizer.step()               # nudge the weights

clip_grad_norm_ scales the gradient down if it’s too large — a blunt but effective way to stop a single freak batch from taking a giant, destabilizing step. It complements PPO’s probability clipping: one bounds how far the policy moves in probability, the other bounds how big the weight update is.

One more detail you’ll see in _normalize: before use, advantages are often whitened (shifted and scaled to mean 0, standard deviation 1) across the minibatch. This doesn’t change their signs (good stays good), just keeps their scale consistent from batch to batch, which makes the learning rate behave predictably. It’s a practical trick, not part of the theory.

Part 8 — The historical arc, in one table

Each method kept everything before it and added one idea. PPO is A2C-with-GAE plus a clip — which is why, once you understand the RolloutBuffer (collection + advantages) and the three- term loss, you understand the whole lineage.

Part 9 — Trace it yourself

The best way to lock this in is to follow one rollout through the code:

1. Collect — OnPolicyStrategy.step() in rl_factory/training/on_policy.py: the loop that steps all N envs T times and calls buffer.add(...).
2. Bootstrap + advantages — the same method calls agent.bootstrap_values(...) then buffer.compute_advantages(...), which runs compute_gae in rl_factory/training/rollout.py. Print advantages and returns and sanity-check a few by hand.
3. Learn — PPOLearner.learn(rollout) in rl_ablations/models/ppo/learner.py loops over buffer.minibatches(...) calling _gradient_step. Read _clipped_policy_loss, the value_loss, and _total_loss, and match each to Parts 6–7 above.

If you can point at the line where each row of the Part 8 table lives, you’re done. ```

Results on Mario 1-1

Plugged into the on-policy strategy, PPO overtakes DQN. Where DQN plateaued around a mean return of ~2774, PPO keeps climbing to ~4545 over 15M steps — the head-to-head curve is in the DQN post. Same Trainer, same environment; the win comes from optimizing the policy directly instead of routing every improvement through an argmax over Q-values.

Where PPO runs out of road

PPO is a strong baseline on 1-1 — but 1-1 hands out a fairly dense reward (move right, get reward). The moment reward is sparse — long stretches with no feedback until a distant payoff — a pure reward-maximizer has nothing to climb. The next posts add an intrinsic reward: a bonus for surprise, so the agent keeps exploring when the environment goes quiet. First up is ICM, curiosity from predicting the consequences of your own actions → ICM.

Code: rl_ablations/models/dqn, built on the rl-factory engine.

This is the second post in the series (rl-factory came first — the framework everything plugs into). It’s the first agent that actually learns. Before this, the only baseline was a random agent: on World 1-1 it reaches a mean max_x ≈ 665 and completes the level 0% of the time (the flagpole is at x ≈ 3160, so random Mario gets about a fifth of the way). That’s the floor. DQN’s job is to beat it.

If you remember one sentence: DQN learns a function Q(s, a) — “how good is pressing button a in state s?” — and acts by picking the button with the highest Q; the entire trick is making that estimate stable enough to train against itself.

Part 0 — where DQN sits in the factory

From the previous post: the one real axis of variation is on-policy vs. off-policy, and DQN is the off-policy family. It inherits OffPolicyStrategy unchanged — async-ish collection into a replay buffer, learning from sampled minibatches. So this post is not about a training loop (that’s shared); it’s about one Learner — DQNLearner — and the network it trains.

Part 1 — the Q-network

The network is the Nature-DQN conv stack: a stacked (4, 84, 84) frame (four grayscale frames, so the agent can see motion) in, one Q-value per action out.

self.conv = nn.Sequential(
    nn.Conv2d(in_channels, 32, kernel_size=8, stride=4), nn.ReLU(),  # 84x84 -> 20x20
    nn.Conv2d(32, 64, kernel_size=4, stride=2), nn.ReLU(),           # 20x20 -> 9x9
    nn.Conv2d(64, 64, kernel_size=3, stride=1), nn.ReLU(),           # 9x9   -> 7x7
    nn.Flatten(),
)
self.head = nn.Sequential(
    nn.Linear(64 * 7 * 7, 512), nn.ReLU(),
    nn.Linear(512, n_actions),   # one Q-value per action; no activation
)

No softmax, no value in [0,1] — the outputs are raw expected returns, one per button combination. Acting is argmax over them (with ε-greedy exploration: a small fraction of the time, press a random button so the agent keeps discovering).

Part 2 — the three ideas that make it trainable

Q-learning’s target is bootstrapped: you regress Q(s, a) toward r + γ·maxₐ' Q(s', a') — an estimate that uses the network’s own output. Do that naïvely and it diverges: you’re chasing a target that moves every time you take a step. DQN’s contribution is three stabilizers, all visible in DQNLearner:

1. Experience replay. Store transitions in a buffer and learn from random minibatches, not the latest step. This breaks the correlation between consecutive frames (which would otherwise make each gradient step redundant with the last) and lets each transition train the network many times.

2. A frozen target network. Keep a second copy of the Q-network, self._target, and compute the bootstrap target from it, not the live network. It’s a stationary goalpost:

self._target = copy.deepcopy(self._q)   # frozen goalpost
...
with torch.no_grad():
    q_next = self._target(next_states).max(dim=1).values
    td_target = rewards + config.gamma * q_next * (1.0 - dones)   # 0 bootstrap on terminal

Every target_sync_every steps the target is refreshed from the live network. Between syncs, the thing you’re regressing toward holds still — which is the difference between training and diverging.

3. Huber loss + reward clipping. The TD error is squashed with smooth_l1_loss (Huber) and rewards are clipped to [-1, 1], so a single surprising transition can’t produce a giant gradient. Mario’s raw reward has big spikes; this keeps the update well-behaved.

The learner self-gates the whole schedule — warm up the buffer (learning_starts), then a gradient step every train_every, and a target sync on its own cadence:

def learn(self):
    self._steps += 1
    metrics = self._gradient_step() if self._is_train_step() else None
    if self._is_sync_step():
        self._sync_target()
    return metrics

That’s DQN in one method: replay-sample, regress toward the frozen-target TD value, periodically move the goalpost.

Part 3 — does it beat the floor?

Yes — comfortably. Against the random baseline’s max_x ≈ 665, DQN climbs to a mean episode return around 2774 on Mario 1-1. It learns to run right, jump the first pits and Goombas, and make real progress into the level. Here it is against the PPO agent we build next:

Read the blue curve first: DQN rises off the floor fast, then plateaus around ~2800 by roughly 8M steps and stops improving. It’s a real agent — but it’s stuck.

Part 4 — the ceiling, and why we move on

DQN plateaus for reasons that are structural, not just tuning:

• It’s value-based on a big discrete action space. Every improvement has to route through an argmax over Q-values; the policy is implicit and changes in lurches.
• Off-policy replay mixes stale experience. The buffer is full of transitions from an older, worse policy, which is efficient but noisy for a fast-moving control problem like Mario.
• Sparse, delayed reward is hard for TD. Progress in Mario often needs a committed sequence of jumps before any reward arrives; bootstrapped value estimates propagate that credit slowly.

The orange curve above is the fix: PPO, an on-policy policy-gradient method that optimizes the policy directly and keeps climbing past DQN’s plateau to ~4500. That’s the next post — and, in factory terms, it’s the other strategy family. Same Trainer, same env, a different half of the split.

The claim I want to earn: adding a new RL algorithm to this codebase is registering a learner, not rewriting a training loop. Everything in the series after this one is an exercise in cashing that claim — each model drops into the same spine and the diff is tiny.

If you remember one sentence: there is exactly one real axis of variation in these algorithms — on-policy vs. off-policy — so the framework encodes that as a single choice and treats everything else (DQN vs. PPO vs. curiosity) as a swapped-in Learner.

The code. The engine lives in rl-factory; the five models that plug into it live in rl-ablations. This post walks the engine.

Part 0 — the trap this avoids

The naïve way to grow an RL repo is one script per algorithm: train_dqn.py, train_ppo.py, train_icm.py. Each re-implements the run loop (step budget, checkpointing, logging, the env-stepping), and they drift. By the fifth algorithm you have five subtly different training loops and no way to compare results fairly, because “the setup” isn’t one thing.

The opposite trap is a single train() with a growing thicket of if algo == "ppo": ... elif algo == "icm": .... That centralizes the loop but scatters each algorithm’s logic across every function it touches.

rl-factory takes a third path: one shared loop, and a small typed contract that each algorithm implements. The loop never mentions a specific algorithm; an algorithm never re-implements the loop.

Part 1 — the three pieces

The spine is three things:

1. Typed interfaces (rl_factory/data_classes/) — the contracts: Agent (something that acts), Learner (something that turns experience into a weight update), Environment, and the config objects. Nothing here knows about Mario or about DQN; they’re just shapes.

2. The Trainer (rl_factory/training/trainer.py) — owns only the shared concerns: the env-step budget, periodic checkpointing, and stats logging. It drives a strategy through four methods and is otherwise blind to what’s being trained:

   self._strategy.setup()
   while steps < config.steps:
       done, metrics = self._strategy.step()   # collect + learn one unit
       steps += done
       stats.record(metrics=metrics)
       # ... checkpoint / log on schedule ...
   self._strategy.teardown()
   

   That’s the entire training loop for every algorithm in the series. DQN and PPO and ICM all run this exact code.

3. The TrainStrategy (rl_factory/training/strategy.py) — the pluggable “how.” A four-method contract: setup / step / weights / teardown.

   class TrainStrategy(ABC):
       @abstractmethod
       def setup(self) -> None: ...
       @abstractmethod
       def step(self) -> tuple[int, dict[str, float]]: ...   # (env-steps done, metrics)
       @abstractmethod
       def weights(self) -> StateDict: ...
       @abstractmethod
       def teardown(self) -> None: ...
   

   A strategy owns how experience is collected and learned from; the Trainer owns when to stop, save, and log. Clean seam.

Part 2 — the one branch that’s real: on-policy vs off-policy

Here’s the design decision the whole thing pivots on. If you line up DQN, PPO, and the three curiosity methods, they differ in a hundred small ways — but those differences collapse onto one structural axis:

• Off-policy (DQN, the random baseline): collect into a replay buffer, learn from sampled batches with a TD target against a shared Q-network.
• On-policy (PPO, ICM, RND, kNN): collect a fresh rollout from a vectorized env, learn from it once with GAE, throw it away.

So there are exactly two concrete strategies — OffPolicyStrategy and OnPolicyStrategy — and which one a model uses is not code, it’s data: a family each model declares when it registers.

class StrategyFamily(StrEnum):
    OFF_POLICY = "off_policy"   # async actors + replay buffer (DQN / random)
    ON_POLICY  = "on_policy"    # vectorized env + on-policy rollout (PPO / ICM / RND / kNN)

# ...and each model just names its family when it registers itself (rl_ablations/register.py):
register_model("dqn", build_agent=dqn.build_agent, build_learner=dqn.build_learner, family=OFF)
register_model("ppo", build_agent=ppo.build_agent, build_learner=ppo.build_learner, family=ON)
register_model("icm", build_agent=ppo.build_agent, build_learner=icm.build_learner, family=ON)

Adding ICM did not add an if icm anywhere in the training loop. ICM is on-policy, so it inherits OnPolicyStrategy unchanged. The only thing that makes it ICM instead of PPO is a different learner (more on that next). That is the single most important sentence in this whole codebase:

ICM rides the on-policy strategy unchanged — it’s just a different build_learner.

Part 3 — a model, factored

Every model is the same four files, so once you’ve read one you can read them all:

The Agent (acting policy) and the Learner (weight update) are deliberately separate. That separation is why the curiosity models are so cheap: ICM, RND, and kNN all act with a plain PPO policy — they only change how the update is computed. You can read that straight off the registrations — they pass the same build_agent (PPO’s) and differ only in build_learner:

# rl_ablations/register.py — icm/rnd/knn share PPO's acting policy; only the learner differs
register_model("ppo", build_agent=ppo.build_agent, build_learner=ppo.build_learner, family=ON)
register_model("icm", build_agent=ppo.build_agent, build_learner=icm.build_learner, family=ON)
register_model("rnd", build_agent=ppo.build_agent, build_learner=rnd.build_learner, family=ON)
register_model("knn", build_agent=ppo.build_agent, build_learner=knn.build_learner, family=ON)

Each curiosity build_learner wraps PPO’s update with an intrinsic-reward term. Same policy, same rollout, same loop; a different learner. Novelty in, exploration out.

Part 4 — the factory tie

rl_factory.algorithm(name) is the one call that assembles a registered model into everything the Trainer needs — it looks the model up in the registry and closes over its family:

def algorithm(name):
    spec = _REGISTRY[name]                       # build_agent, build_learner, family

    def make_strategy(build_env, config):
        return _build_strategy(spec.family, build_env, spec.build_agent, spec.build_learner, config)

    return Algorithm(build_agent=spec.build_agent, make_strategy=make_strategy)

Note what make_strategy takes: a build_env passed in from outside. rl-factory never imports Mario — the environment is a parameter. That’s what keeps it a framework and not a Mario script, and it’s the seam this project splits along: rl-factory is the engine in this post, and rl-ablations is the separate repo that brings the Mario environment and registers the five models on top of it.

Part 5 — so what does “add a model” actually cost?

Concretely, to add an algorithm:

1. Write its four files (config / model / learner / factory).
2. Call register_model(name, build_agent=…, build_learner=…, family=…) — one line.

That’s the whole cost. No new training loop. No new checkpointing, logging, or env-stepping. If it’s on-policy, no new strategy either — you’re writing a Learner and nothing else structural. That’s the payoff, and it’s why the rest of this series can move fast: each post is one model dropped into this frame, and the interesting part is the idea, because the plumbing is already paid for.

What’s next

From here the series walks the models in the order I built them, and each one earns its place by fixing a limitation of the last:

1. DQN — off-policy value learning; beat the random floor.
2. PPO — on-policy policy gradients; a more stable learner.
3. ICM — add curiosity so Mario explores when the reward is sparse.
4. RND — a simpler curiosity signal (and why simpler was better).
5. kNN episodic novelty — a count-based signal that can’t saturate the way the others did.

Every one of them is the same spine from this post, plus a Learner. Let’s start beating the floor.