ICM: giving Mario curiosity (and the bug that ate it)

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.