kNN episodic novelty: the count that can't saturate

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.