# /// script # requires-python = ">=3.11" # dependencies = [ # "marimo", # "matplotlib", # "numpy", # ] # /// import marimo __generated_with = "0.23.3" app = marimo.App(width="medium") @app.cell(hide_code=True) def _(mo): mo.Html(""" """) return @app.cell(hide_code=True) def _(mo): mo.md(r""" # Curiosity Killed the Snake: When Intrinsic Motivation Helps and When It Doesn't *An interactive exploration of "Curiosity-driven Exploration by Self-supervised Prediction" (Pathak et al., 2017 · [arXiv](https://arxiv.org/abs/1705.05363) · [alphaxiv](https://alphaxiv.org/abs/1705.05363)) and analyzing the impact of curiosity on the game of Snake.* [![Open in marimo](https://marimo.io/shield.svg)](https://molab.marimo.io/github/Saheb/rl-snake/blob/main/notebooks/curiosity.py/wasm) ### The Problem: Sparse Reward Trap In reinforcement learning, an agent learns by maximizing reward. Usually that reward is **external**: it comes from the environment, like `+1` for eating an apple in Snake or `-1` for dying. But what happens when the environment is vast, the board is empty, and the external reward is incredibly hard to find by pure chance? What happens if the environment gives no useful reward signal for a long time? When an $\epsilon$-greedy agent, such as a standard **Deep Q-Network (DQN)**, faces a sparse environment, its exploration is random — with no signal to guide it, it stumbles across the state space rather than systematically searching. The natural fix is to give the agent an intrinsic drive toward novelty. But whether that drive actually helps depends entirely on what is bottlenecking learning — exploration, or something else. **Play with the agent below to see how it performs when the reward is 13 steps away.** """) return @app.cell(hide_code=True) def _(json, mo, random, wasm_iframe): grid_size = 14 start = (grid_size // 2, grid_size // 2) reward = (0, grid_size - 1) n_steps = 120 def clamp(v): return max(0, min(grid_size - 1, v)) # --- Random walk --- rng1 = random.Random(42) hist_random = [start] for _ in range(n_steps): y, x = hist_random[-1] moves = [(0,1),(0,-1),(1,0),(-1,0)] rng1.shuffle(moves) for dy, dx in moves: ny, nx = clamp(y+dy), clamp(x+dx) if (ny, nx) != (y, x): break hist_random.append((ny, nx)) # --- Curiosity-Driven Walk (seeks unvisited/surprising tiles) --- rng2 = random.Random(42) hist_curious = [start] visit_counts = {start: 1} for _ in range(n_steps): y, x = hist_curious[-1] moves = [(0,1),(0,-1),(1,0),(-1,0)] valid_neighbors = [] for dy, dx in moves: ny, nx = clamp(y+dy), clamp(x+dx) if (ny, nx) != (y, x): valid_neighbors.append((ny, nx)) rng2.shuffle(valid_neighbors) valid_neighbors.sort(key=lambda pos: visit_counts.get(pos, 0)) next_pos = valid_neighbors[0] hist_curious.append(next_pos) visit_counts[next_pos] = visit_counts.get(next_pos, 0) + 1 html = f"""

Random Walk

Illustrative baseline - no reward signal, pure chance

Illustrative Count-Based Curiosity Proxy

Heuristic stand-in - seeks lowest-visited tiles

Step 0 / {n_steps}

● Agent ■ Start ★ Reward

""" mo.vstack([ wasm_iframe(html, height="500px"), mo.callout( mo.md( "**Illustration note:** The right-hand agent uses a **count-based heuristic** " "(always move to the lowest visit-count neighbour), not a live neural network. " "In this discrete finite grid, visit counts serve as a practical stand-in for " "the Intrinsic Curiosity Module (ICM)'s forward-model prediction error - both signal how *novel* a state is, " "and the outward-seeking behaviour looks similar. However, this equivalence " "**only holds in tabular settings** with an enumerable state space. " "The Intrinsic Curiosity Module's core contribution is generalizing curiosity to **continuous, high-dimensional " "observation spaces** (raw pixels, sensor arrays) where visit counts are undefined - " "and where a learned forward model is the only tractable novelty signal." ), kind="info" ) ]) return @app.cell(hide_code=True) def _(mo): eta_slider = mo.ui.slider( start=0.01, stop=1.0, step=0.01, value=0.5, label="Curiosity Weight η (Eta)", show_value=True ) return (eta_slider,) @app.cell(hide_code=True) def _(eta_slider, mo, wasm_iframe): # 1. The Narrative Text (Explaining the Forward Model) intro_text = mo.md( """ ### The Concept: Mathematically Defining "Curiosity" If the environment is not giving the agent useful rewards, the agent has to generate its own. Pathak et al. achieved this by giving the agent a **Forward Model**: a neural network inside the agent that learns to predict what the next observation should look like after the agent takes an action. Before the agent takes a step, the Forward Model predicts the next state. The agent then takes the step, observes the actual next state, and compares prediction to reality. The difference is the **Prediction Error** (Mean Squared Error). $$Intrinsic\\ Reward = \\frac{\\eta}{2} (Predicted\\ State - Actual\\ State)^2$$ Where **$\\eta$ (Eta)** is the **Curiosity Weight**. It scales how much intrinsic reward the agent gets from being surprised. The key insight of this equation is that **Prediction Error = Surprise = Reward**. The squared error here is not arbitrary. It is the standard regression loss for a continuous target and corresponds to a Gaussian negative log-likelihood assumption: large prediction mistakes are penalized quadratically, while small residual errors fade smoothly. The factor of $\\frac{1}{2}$ is conventional - it cancels with the gradient of the square during backprop, so $\\eta$ becomes the clean multiplier on the residual. The Forward Model is trained from the agent's own experience: each transition `(state, action, next_state)` becomes a supervised learning example. After many updates on familiar transitions, its predictions become accurate and the intrinsic reward shrinks. * If the agent repeatedly visits the same tile and observes the same transition, its Forward Model learns that transition. The prediction error becomes small. The agent is bored. * If the agent reaches a new tile or a hard-to-predict transition, the Forward Model is wrong. The error is large. The agent experiences a spike of surprise. **Try the illustrative simulator below:** Click the tiles and watch how "surprise" spikes when you explore new areas, then decays as transitions become familiar. The η slider underneath rescales the reward magnitude. """ ) # 2. The Interactive "Boredom Simulator" (Native HTML5 Canvas) eta_value = eta_slider.value html2 = """

The Environment (Click to Explore)

Click tiles to generate intrinsic reward.

Forward Model: Prediction Error (Reward)

Total Intrinsic Reward: 0.00

""".replace("__ETA__", f"{eta_value:.2f}") # 3. Stack the markdown and the widget together mo.vstack([ intro_text, wasm_iframe(html2, height="400px"), mo.callout( mo.md( "**Illustration note:** The reward curve above uses a hardcoded 3-step decay " "(1st visit → 1.0, 2nd → 0.4, 3rd → 0.1) as a pedagogical proxy, not a real forward model. " "A live ICM computes $\\frac{\\eta}{2}(\\hat{\\phi}_{t+1} - \\phi_{t+1})^2$ in latent space - " "the qualitative shape is the same but the exact decay depends on network training. " "See the **Mini-Experiment** section for real gradient-descent-based decay. " "The slider below now rescales this illustrative reward curve as well." ), kind="info" ), ]) return @app.cell(hide_code=True) def _(eta_slider, mo): mo.hstack( [eta_slider, mo.md(f"**η = {eta_slider.value:.2f}** → max intrinsic reward on first visit ≈ **{eta_slider.value / 2:.3f}** (η/2 × 1²)")], justify="start", align="center", gap=2 ) return @app.cell(hide_code=True) def _(json, mo, wasm_iframe): # 1. The Narrative Text (Explaining the Trap and the Filter) act2_text = mo.md( """ ### A fatal flaw: The Noisy TV Problem If you try to predict raw pixels, the agent gets distracted by irrelevant stochasticity (a flickering leaf, a TV screen with random noise) that has nothing to do with the agent's actions. Nor does it affect the agent in any way. It's novel, but it's not relevant. This is sometimes called the "noisy TV problem" How do we remove this noise, so agent can focus on what matters? How do we teach the agent to ignore this noise? Rather than making predictions in raw sensory space (e.g. pixels), the paper transforms sensory input into a feature space where only information relevant to the agent's actions is represented. The core insight is to predict only those changes in the environment that could possibly be due to the actions of the agent or could affect the agent - and ignore everything else. > "We learn this feature space using self-supervision – training a neural network on a proxy inverse dynamics task of predicting the agent’s action given its current and next states. Since the neural network is only required to predict the action, it has no incentive to represent within its feature embedding space the factors of variation in the environment that do not affect the agent itself." > > - Pathak et al., *"Curiosity-driven Exploration by Self-supervised Prediction"*, ICML 2017. [arXiv:1705.05363](https://arxiv.org/abs/1705.05363) The Inverse Model is trained to predict the *agent’s own actions*. Because the agent's actions cannot control the random TV static, the neural network learns to completely ignore the TV when creating the latent vector. The noise is mathematically filtered out. **Run the illustrative agents below** to see why raw-pixel prediction fails, and how the latent filter in the Intrinsic Curiosity Module (ICM) saves the agent. """ ) # 2. Hardcode the pedagogical paths with UNIQUE variables grid_h_tv, grid_w_tv = 7, 8 start_tv_rc = (1, 1) goal_tv_rc = (1, 6) noisy_tv_rc = (5, 1) hallway_path_rc = [(1,1), (2,1), (3,1), (4,1), (5,1)] # RAW AGENT: Gets stuck at the TV forever (infinite prediction error) hist_raw_tv_rc = hallway_path_rc + [noisy_tv_rc] * 40 # ICM AGENT: Stares for 5 frames, learns to filter it, turns around, finds goal hist_icm_tv_rc = hallway_path_rc + [noisy_tv_rc] * 5 + [(4,1), (4,2), (4,3), (4,4), (4,5), (4,6), (3,6), (2,6), goal_tv_rc] + [goal_tv_rc] * 20 # 3. Build the Canvas HTML (Saved to html_tv) html_tv = f"""

Raw Pixel Prediction

Unable to predict the TV's static, the agent gets permanently stuck by infinite "surprise".

Latent Feature Prediction (ICM)

The Inverse Model filters out the TV noise. The agent gets bored and moves on.

Step: 0

■ Start ★ Goal ▒ Noisy TV ● Agent

""" # 4. Render everything together mo.vstack([ act2_text, wasm_iframe(html_tv, height="380px"), mo.callout( mo.md( "**Illustration note:** The agent paths above are **scripted pedagogical illustrations**, " "not live neural network runs. The Raw agent is hardcoded to stay fixed at the TV tile; " "the ICM agent is hardcoded to move on after 5 steps. The qualitative behavior is " "grounded in the paper's theoretical predictions - a real ICM agent learns to filter " "action-independent noise through its inverse model training objective - but the paths " "themselves are hand-authored for clarity." ), kind="info" ) ]) return @app.cell(hide_code=True) def _(mo): # The PyTorch Architecture Walkthrough architecture_text = mo.md( r""" ### The Architecture: Under the Hood of ICM Before we look at how this breaks down in late-stage training, we need to understand how the Intrinsic Curiosity Module is actually constructed. The ICM is not a standalone agent; it is an auxiliary subsystem that runs *inside the training loop* of a standard reinforcement-learning algorithm such as Proximal Policy Optimization (PPO) or Deep Q-Network (DQN). The policy still chooses actions and learns from rewards. ICM adds an extra reward term and trains its own predictive models from the same transitions. 1. **The Feature Encoder ($\phi$):** Compresses raw pixels/states into a dense, latent vector. 2. **The Inverse Model:** Takes the current state $\phi(s_t)$ and the next state $\phi(s_{t+1})$ to predict the action $a_t$. This is the filter that ignores unpredictable noise. 3. **The Forward Model:** Takes the current state $\phi(s_t)$ and the action $a_t$ to predict the next state $\hat{\phi}(s_{t+1})$. Formally, the inverse model is a classifier and the forward model is a regressor: $$L_I = -\log p(a_t \mid \phi(s_t), \phi(s_{t+1}))$$ $$L_F = \frac{1}{2}\left\|\hat{\phi}(s_{t+1}) - \phi(s_{t+1})\right\|_2^2$$ $$L_{ICM} = (1 - \beta)L_I + \beta L_F$$ $L_I$ trains the encoder to preserve action-relevant information: if a visual feature does not help infer which action moved the agent from $s_t$ to $s_{t+1}$, the inverse objective has little reason to keep it. $L_F$ trains the dynamics predictor in that filtered feature space. The intrinsic reward is proportional to $L_F$, but the encoder is shaped by both losses. In a DQN-style loop, the call site looks like this: ```python action = dqn_policy.select_action(state) next_state, external_reward, done = env.step(action) inverse_loss, forward_loss, intrinsic_reward = icm(state, next_state, action) reward_for_dqn = external_reward + eta * intrinsic_reward replay_buffer.add(state, action, reward_for_dqn, next_state, done) dqn_loss = train_dqn_from_replay(replay_buffer) icm_loss = (1 - beta) * inverse_loss + beta * forward_loss ``` So the ICM is called immediately after the environment step. Its prediction error becomes an exploration bonus, while its own networks are trained on the observed transition. Here is the core PyTorch implementation of the module: ```python import torch import torch.nn as nn import torch.nn.functional as F class IntrinsicCuriosityModule(nn.Module): def __init__(self, state_dim, action_dim, latent_dim=256): super().__init__() # 1. Feature Encoder: Compresses state into latent space self.encoder = nn.Sequential( nn.Linear(state_dim, 128), nn.ReLU(), nn.Linear(128, latent_dim), nn.ReLU() ) # 2. Inverse Model: Predicts action from (s_t, s_{t+1}) self.inverse_model = nn.Sequential( nn.Linear(latent_dim * 2, 128), nn.ReLU(), nn.Linear(128, action_dim) ) # 3. Forward Model: Predicts next state from (s_t, action) self.forward_model = nn.Sequential( nn.Linear(latent_dim + action_dim, 128), nn.ReLU(), nn.Linear(128, latent_dim) ) def forward(self, state, next_state, action_one_hot): # Encode states into the latent space phi_t = self.encoder(state) phi_t_plus_1 = self.encoder(next_state) # Inverse Model Prediction (for training the encoder to filter noise) inverse_input = torch.cat([phi_t, phi_t_plus_1], dim=1) pred_action_logits = self.inverse_model(inverse_input) # Forward Model Prediction (for generating Intrinsic Reward) forward_input = torch.cat([phi_t, action_one_hot], dim=1) pred_phi_t_plus_1 = self.forward_model(forward_input) # Calculate Intrinsic Reward (Mean Squared Error in Latent Space) intrinsic_reward = 0.5 * (pred_phi_t_plus_1 - phi_t_plus_1).pow(2).sum(dim=1) return pred_action_logits, pred_phi_t_plus_1, phi_t_plus_1, intrinsic_reward ``` Notice the `intrinsic_reward` calculation at the very bottom. It is decoupled from the game's actual score, then added to the environment reward before the policy update. The agent trains this module simultaneously with its policy, constantly trying to minimize both the Forward and Inverse loss, while using the current Forward Model error as an exploration bonus. """ ) architecture_text return @app.cell(hide_code=True) def _(mo): _state_dim, _action_dim, _batch, _latent_dim = 20, 4, 8, 256 mo.callout( mo.md(f""" **WASM-compatible forward-pass shape check** - same tensor contract as the PyTorch module above: | Tensor | Shape | Role | |--------|-------|------| | `pred_action_logits` | `{[_batch, _action_dim]}` | Inverse model output - predicted action | | `pred_phi_next` | `{[_batch, _latent_dim]}` | Forward model output - predicted next latent | | `phi_next` | `{[_batch, _latent_dim]}` | Encoder output - actual next latent | | `intrinsic_reward` | `{[_batch]}` | Per-step curiosity signal (MSE in latent space) | The executable PyTorch demo was removed from the notebook runtime because `torch` has no Pyodide/WASM build. The architecture above remains the implementation used in the repository; the browser cells below use NumPy proxies to preserve the same math and interactivity. """), kind="info", ) return @app.cell(hide_code=True) def _(mo): _train_intro = mo.md(r""" ### Mini-Experiment: Does the Forward Model Actually Learn? The cell above defined the ICM tensor contract. This experiment shows the core learning signal directly. We train a tiny **NumPy linear forward model** on a **5×5 grid** with a **purely random policy** (ε = 1.0, zero extrinsic reward). It predicts the next grid position from the current position and action, then uses its own prediction error as the intrinsic reward. This is "boredom by gradient descent." Same phenomenon as the boredom simulator - now produced by a browser-compatible linear model that can run inside Pyodide. The full ICM paper uses a neural network; this simplified example is intentionally linear because that is sufficient to learn deterministic 5×5 grid transitions. """) train_btn = mo.ui.button( label="▶ Train NumPy linear forward model on 5×5 grid", kind="success", value=0, on_click=lambda v: v + 1, ) mo.vstack([_train_intro, train_btn]) return (train_btn,) @app.cell(hide_code=True) def _(mo, np, plt, random, train_btn): mo.stop( np is None or plt is None, mo.callout( mo.md("**NumPy/Matplotlib are unavailable in this runtime.** In WASM they are installed from the notebook metadata."), kind="warn", ), ) if not train_btn.value: _fig0, _ax0 = plt.subplots(figsize=(9, 3.5)) _ax0.set_xlabel("Episode", fontsize=11) _ax0.set_ylabel("Mean Intrinsic Reward per Step (log scale)", fontsize=11) _ax0.set_title("Linear Forward Model Learning Curve - 5×5 Grid, Random Policy", fontweight="bold", fontsize=13) _ax0.set_xlim(1, 300) _ax0.spines["top"].set_visible(False) _ax0.spines["right"].set_visible(False) _ax0.grid(axis="y", alpha=0.3) plt.tight_layout() _out = mo.as_html(_fig0) else: class _GridEnv: def __init__(self, size=5): self.size = size self.pos = [size // 2, size // 2] def reset(self): self.pos = [self.size // 2, self.size // 2] return np.array(self.pos, dtype=np.float32) / self.size def step(self, action): dy, dx = [(-1, 0), (1, 0), (0, -1), (0, 1)][action] self.pos[0] = max(0, min(self.size - 1, self.pos[0] + dy)) self.pos[1] = max(0, min(self.size - 1, self.pos[1] + dx)) return np.array(self.pos, dtype=np.float32) / self.size, 0.0, False _env = _GridEnv(5) _rng = np.random.default_rng(7) _w = _rng.normal(0.0, 0.15, size=(6, 2)) _b = np.zeros(2) _lr = 0.2 _ep_rewards = [] for _ep in range(300): _s = _env.reset() _features, _targets = [], [] for _ in range(25): _a = random.randint(0, 3) _ns, _, _ = _env.step(_a) _action_one_hot = np.eye(4, dtype=np.float32)[_a] _features.append(np.concatenate([_s, _action_one_hot])) _targets.append(_ns) _s = _ns _x = np.asarray(_features) _y = np.asarray(_targets) _pred = _x @ _w + _b _err = _pred - _y _r_i = 0.5 * np.sum(_err * _err, axis=1) _ep_rewards.append(float(np.mean(_r_i))) _grad_pred = _err / len(_x) _w -= _lr * (_x.T @ _grad_pred) _b -= _lr * np.sum(_grad_pred, axis=0) _episodes = list(range(1, 301)) _floor = 1e-6 _rewards_floored = [max(r, _floor) for r in _ep_rewards] _smooth = np.maximum(np.convolve(_rewards_floored, np.ones(10) / 10, mode='valid'), _floor) _fig, _ax = plt.subplots(figsize=(9, 3.5)) _ax.plot(_episodes, _rewards_floored, color='#93c5fd', linewidth=1, alpha=0.5, label='Per-episode') _ax.plot(_episodes[9:], _smooth, color='#2563eb', linewidth=2.5, label='10-episode rolling avg') _ax.set_yscale('log') _ax.set_xlabel("Episode", fontsize=11) _ax.set_ylabel("Mean Intrinsic Reward per Step (log scale)", fontsize=11) _ax.set_title("Linear Forward Model Learning Curve - 5×5 Grid, Random Policy", fontweight='bold', fontsize=13) _ax.legend(fontsize=10) _ax.spines['top'].set_visible(False) _ax.spines['right'].set_visible(False) _ax.grid(axis='y', alpha=0.3, which='both') _ax.set_xlim(1, 300) plt.tight_layout() _first5 = sum(_ep_rewards[:5]) / 5 _last5 = sum(_ep_rewards[-5:]) / 5 _drop = int(_first5 / (_last5 + 1e-9)) _out = mo.vstack([ mo.as_html(_fig), mo.callout( mo.md( f"Intrinsic reward decayed from **{_first5:.4f}** (first 5 episodes) → " f"**{_last5:.5f}** (last 5 episodes) - a **{_drop:,}× drop**. " "The linear forward model has learned the grid's transition structure. There is little left to be curious about. " "The same phenomenon the boredom simulator illustrated - now shown by gradient descent in a minimal linear setting." ), kind="success" ) ]) _out return @app.cell(hide_code=True) def _(json, mo, wasm_iframe): # 1. Narrative Text train_text = mo.md( """ ### Live Training: Real ICM Data from Snake Below are **real training logs** from our Snake DQN agent (Dueling DQN + 3-Step Returns + PER + ICM) trained on a 10×10 board for 16,000 games. The `Intrinsic` value is the **actual forward model prediction error** - not a simulation. Two signals tell the story: * **Mean Score** (blue) - rolling average game score over the last 100 games * **Intrinsic Reward** (orange) - the forward model's live prediction error (the curiosity signal) * **Epsilon** (grey dashed) - exploration rate, decaying 0.95 → 0.01 Watch what happens after **Game ~2,000** (when ε hits its floor and the agent goes fully greedy): intrinsic reward keeps spiking on every novel death. The agent never stops being "surprised" - but that persistent surprise doesn't translate into better performance. The next sections show why. Methodological caveat: this is a mechanistic case study, not a benchmark claim. A research-grade evaluation would run multiple random seeds and report confidence intervals for at least four ablations: DQN + PER, DQN + ICM, DQN + PER + ICM, and DQN + PER + ICM with terminal intrinsic rewards masked. The single-run logs here are still useful because they expose a concrete failure mode in the replay distribution. """ ) # 2. Inlined training data (parsed from dqn_per_icm_10x10.log; inlined so the notebook # runs without local files in molab/WASM environments) _games = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900, 2000, 2100, 2200, 2300, 2400, 2500, 2600, 2700, 2800, 2900, 3000, 3100, 3200, 3300, 3400, 3500, 3600, 3700, 3800, 3900, 4000, 4100, 4200, 4300, 4400, 4500, 4600, 4700, 4800, 4900, 5000, 5100, 5200, 5300, 5400, 5500, 5600, 5700, 5800, 5900, 6000, 6100, 6200, 6300, 6400, 6500, 6600, 6700, 6800, 6900, 7000, 7100, 7200, 7300, 7400, 7500, 7600, 7700, 7800, 7900, 8000, 8100, 8200, 8300, 8400, 8500, 8600, 8700, 8800, 8900, 9000, 9100, 9200, 9300, 9400, 9500, 9600, 9700, 9800, 9900, 10000, 10100, 10200, 10300, 10400, 10500, 10600, 10700, 10800, 10900, 11000, 11100, 11200, 11300, 11400, 11500, 11600, 11700, 11800, 11900, 12000, 12100, 12200, 12300, 12400, 12500, 12600, 12700, 12800, 12900, 13000, 13100, 13200, 13300, 13400, 13500, 13600, 13700, 13800, 13900, 14000, 14100, 14200, 14300, 14400, 14500, 14600, 14700, 14800, 14900, 15000, 15100, 15200, 15300, 15400, 15500, 15600, 15700, 15800, 15900, 16000] _mean_scores = [0.11, 0.13, 0.13, 0.14, 0.15, 0.16, 0.19, 0.23, 0.27, 0.3, 0.34, 0.37, 0.42, 0.46, 0.49, 0.54, 0.6, 0.66, 0.74, 0.84, 0.96, 1.07, 1.23, 1.37, 1.49, 1.61, 1.71, 1.79, 1.88, 1.95, 2.05, 2.17, 2.27, 2.4, 2.53, 2.65, 2.76, 2.88, 2.95, 3.03, 3.12, 3.19, 3.26, 3.38, 3.45, 3.53, 3.61, 3.68, 3.77, 3.85, 3.91, 4.05, 4.15, 4.21, 4.3, 4.39, 4.43, 4.48, 4.6, 4.72, 4.84, 4.88, 4.97, 5.04, 5.13, 5.18, 5.25, 5.3, 5.35, 5.4, 5.47, 5.54, 5.59, 5.63, 5.67, 5.71, 5.75, 5.8, 5.84, 5.92, 5.96, 6.0, 6.04, 6.1, 6.15, 6.19, 6.21, 6.25, 6.3, 6.34, 6.38, 6.42, 6.44, 6.48, 6.51, 6.54, 6.6, 6.64, 6.66, 6.67, 6.67, 6.7, 6.73, 6.77, 6.81, 6.85, 6.86, 6.89, 6.92, 6.95, 6.99, 7.04, 7.08, 7.12, 7.15, 7.18, 7.19, 7.23, 7.25, 7.29, 7.31, 7.33, 7.36, 7.37, 7.39, 7.4, 7.42, 7.45, 7.48, 7.49, 7.51, 7.53, 7.57, 7.6, 7.63, 7.65, 7.66, 7.68, 7.68, 7.69, 7.72, 7.74, 7.76, 7.79, 7.8, 7.84, 7.86, 7.89, 7.9, 7.92, 7.94, 7.93, 7.96, 7.98, 7.99, 8.01, 8.03, 8.03, 8.05, 8.06] _epsilons = [0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65, 0.6, 0.55, 0.5, 0.45, 0.4, 0.35, 0.3, 0.25, 0.2, 0.15, 0.1, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01] _intrinsics = [0.003, 0.003, 0.001, 0.001, 0.002, 0.003, 0.002, 0.003, 0.002, 0.001, 0.007, 0.005, 0.002, 0.014, 0.004, 0.001, 0.002, 0.011, 0.003, 0.003, 0.045, 0.011, 0.002, 0.047, 0.004, 0.004, 0.001, 0.001, 0.002, 0.001, 0.043, 0.002, 0.047, 0.001, 0.008, 0.054, 0.02, 0.012, 0.007, 0.013, 0.015, 0.006, 0.006, 0.013, 0.008, 0.006, 0.003, 0.004, 0.066, 0.107, 0.002, 0.004, 0.127, 0.008, 0.034, 0.003, 0.011, 0.067, 0.026, 0.02, 0.082, 0.004, 0.119, 0.11, 0.005, 0.034, 0.017, 0.158, 0.016, 0.002, 0.068, 0.012, 0.006, 0.085, 0.007, 0.042, 0.003, 0.003, 0.011, 0.034, 0.036, 0.112, 0.005, 0.064, 0.082, 0.009, 0.004, 0.125, 0.196, 0.009, 0.083, 0.007, 0.005, 0.011, 0.085, 0.011, 0.072, 0.007, 0.105, 0.089, 0.021, 0.09, 0.124, 0.014, 0.091, 0.158, 0.095, 0.005, 0.11, 0.132, 0.047, 0.044, 0.044, 0.096, 0.008, 0.01, 0.006, 0.058, 0.006, 0.169, 0.133, 0.03, 0.016, 0.04, 0.01, 0.007, 0.005, 0.064, 0.121, 0.007, 0.019, 0.07, 0.003, 0.143, 0.009, 0.168, 0.008, 0.01, 0.107, 0.026, 0.014, 0.135, 0.005, 0.015, 0.006, 0.092, 0.004, 0.106, 0.002, 0.081, 0.129, 0.022, 0.078, 0.108, 0.002, 0.079, 0.046, 0.002, 0.067, 0.008] _max_score = max(_mean_scores) if _mean_scores else 1.0 _max_intrinsic = max(_intrinsics) if _intrinsics else 1.0 _epsilon_floor_game = next((g for g, e in zip(_games, _epsilons) if e <= 0.011), None) # 3. Code accordion code_accordion = mo.accordion({ "🔍 View Original Log Format": mo.md( """ Each checkpoint line in `dqn_per_icm_10x10.log` looked like: ``` Game 2000 | Score: 1 | Record: 12 | Mean Score: 0.84 | Epsilon: 0.01 | Intrinsic: 0.003 | Eta: 0.01000 ``` Parsed with: ```python import re m = re.match( r"Game (\\d+).*?Mean Score:\\s*([\\d.]+).*?Epsilon:\\s*([\\d.]+).*?Intrinsic:\\s*([\\d.]+)", line ) ``` The 160 checkpoint values (every 100 games) are inlined above so the notebook runs without local files in molab / WASM environments. """ ) }) # 4. Inject real data into the animated chart _json_data = json.dumps({ "games": _games, "scores": [s / _max_score for s in _mean_scores], "intrinsics": [min(i / _max_intrinsic, 1.0) for i in _intrinsics], "epsilons": _epsilons, "max_score": _max_score, "epsilon_floor_game": _epsilon_floor_game, "last_game": _games[-1] if _games else 16000, }) html_train = f"""

Mean Score

Intrinsic Reward (scaled)

Epsilon

Orange spikes = real forward-model surprise events on novel states and deaths. This is real training data, but still from a single-run case study rather than a multi-seed benchmark.

""" mo.vstack([train_text, code_accordion, wasm_iframe(html_train, height="430px")]) return @app.cell(hide_code=True) def _(mo): # 1. The Narrative Text act3_text = mo.md( """ ### The Edge Case: The Death Oversampling Trap (PER + ICM) ICM is brilliant in sparse environments, but in this Snake setup it creates a specific failure mode once terminal states and replay interact. When an agent dies, the environment resets abruptly. The Forward Model cannot predict a "Game Over" screen from a standard movement action. This results in a massive spike in Prediction Error, which translates to a massive Intrinsic Reward. **The agent learns that suicide is incredibly rewarding.** This becomes catastrophic if you are using **Prioritized Experience Replay (PER)**. PER samples memories from the replay buffer based on their error magnitude. Because "deaths" have the highest error, PER oversamples them. **The Trap:** the replay distribution becomes poisoned. Training batches become flooded with death sequences, drowning out ordinary movement and successful apple captures. The agent is no longer just exploring poorly; it is repeatedly training on the transition that says "dying was salient." Ablation intuition: | Variant | Expected behavior | |---|---| | PER only | Replays high-TD-error events, but death transitions are not additionally inflated by curiosity. | | ICM only | Death may be intrinsically surprising, but uniform replay limits how often that transition dominates training. | | PER + ICM | Death transitions get both high TD error and high intrinsic error, so replay probability compounds. | | PER + ICM + terminal mask | Keeps curiosity for non-terminal novelty while preventing terminal resets from becoming a reward source. | *Adjust the toggle below to see an illustrative replay-distribution sketch of how adding ICM to a PER buffer can poison the training batch.* """ ) _mechanism_detail = mo.md( r""" The failure mode is not just that terminal states are surprising — it is that **curiosity changes what replay considers important**. The mechanism in three steps: 1. **TD target inflation:** the Q-learning target contains `r_ext + η r_int`. If a terminal transition produces a large intrinsic reward spike, its TD error jumps even when the extrinsic outcome is bad. 2. **Replay amplification:** PER samples transitions with probability $p \propto |\delta|^\alpha$. Elevated $r_i$ inflates $\delta$, which inflates $p$, so novel (and terminal) transitions are over-sampled. 3. **Terminal surprise concentration:** game-over resets are unusually hard for the forward model to predict, so the same death transitions are both highly surprising *and* highly replayed — a compounding loop. **ICM raises the surprise of terminal events; PER turns that surprise into repeated training signal.** """ ) act3_mechanism_accordion = mo.accordion({"▶ Why PER + ICM specifically amplifies death (causal chain)": _mechanism_detail}) # 2. The Interactive UI (FIXED: Using a flat list to prevent KeyErrors) agent_type = mo.ui.radio( options=["Standard Agent (PER Only)", "Curious Agent (PER + ICM)"], value="Standard Agent (PER Only)", label="**Select Agent Architecture:**" ) sample_btn = mo.ui.button( label="🎲 Sample Training Batch (n=32)", kind="success" ) return act3_mechanism_accordion, act3_text, agent_type, sample_btn @app.cell(hide_code=True) def _( act3_mechanism_accordion, act3_text, agent_type, mo, np, plt, sample_btn, ): mo.stop( np is None or plt is None, mo.callout( mo.md("**NumPy/Matplotlib are unavailable in this runtime.** In WASM they are installed from the notebook metadata."), kind="warn", ), ) # 3. The Visualization Logic def render_batch_viz(agent_val, btn_click): # Base Replay Buffer Composition (100 memories total) # 80 Normal Steps, 15 Apples, 5 Deaths types = ['Normal'] * 80 + ['Apple'] * 15 + ['Death'] * 5 colors = ['#B0BEC5'] * 80 + ['#4CAF50'] * 15 + ['#F44336'] * 5 # TD error ordering approximates empirical training: normal < apple < death. # Exact values are ordinal stand-ins; the qualitative point holds for any # ordering where deaths > apples > normal steps. base_td = np.array([0.1] * 80 + [0.4] * 15 + [0.8] * 5) if agent_val == "Curious Agent (PER + ICM)": # ICM adds a forward-model MSE spike at terminal transitions. # In the training logs, per-step intrinsic reward at game-over frames # ran ~10-50x normal-step values before eta scaling (η=0.01). # The 4.0 approximates raw latent-space MSE on terminal transitions. intrinsic_surprise = np.array([0.0] * 80 + [0.0] * 15 + [4.0] * 5) priorities = base_td + intrinsic_surprise else: priorities = base_td probabilities = priorities / np.sum(priorities) # Sample the Batch batch_size = 32 sampled_indices = np.random.choice( len(types), size=batch_size, p=probabilities, replace=True ) sampled_types = [types[i] for i in sampled_indices] counts = { 'Normal Steps': sampled_types.count('Normal'), 'Apples (Reward)': sampled_types.count('Apple'), 'Deaths (Penalty)': sampled_types.count('Death') } # Render fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4), gridspec_kw={'width_ratios': [1.2, 1]}) # Left Plot: Memory Pool Bubble Chart np.random.seed(42) # Fixed seed for positions so bubbles don't jump around x_pos, y_pos = np.random.rand(100), np.random.rand(100) np.random.seed(None) sizes = priorities * 150 ax1.scatter(x_pos, y_pos, s=sizes, c=colors, alpha=0.8, edgecolors='black', linewidth=0.5) ax1.set_title("Replay Buffer (Bubble Size = Selection Priority)", fontweight='bold') ax1.axis('off') # Right Plot: Bar Chart labels = list(counts.keys()) values = list(counts.values()) bar_colors = ['#B0BEC5', '#4CAF50', '#F44336'] bars = ax2.bar(labels, values, color=bar_colors, edgecolor='black') ax2.set_title(f"Sampled Training Batch (Size: {batch_size})", fontweight='bold') ax2.set_ylabel("Number of Samples in Batch") ax2.set_ylim(0, batch_size) for bar in bars: yval = bar.get_height() ax2.text(bar.get_x() + bar.get_width()/2, yval + 0.5, int(yval), ha='center', va='bottom', fontweight='bold') plt.tight_layout() plt.gca().spines['top'].set_visible(False) plt.gca().spines['right'].set_visible(False) return mo.as_html(fig) # 4. Stack the UI elements ui_stack = mo.vstack([ act3_text, act3_mechanism_accordion, mo.hstack([agent_type, sample_btn], justify="start", align="center", gap=2), mo.callout( mo.md( "**Interpretation note:** This widget is an **illustrative sampling model**, not a replay log. " "The bubble sizes and priorities are hand-set to mirror the qualitative ordering argued above: " "normal transitions < apples < terminal deaths under PER + ICM. The real evidence appears in the training-log sections before and after this sketch." ), kind="info" ) ]) # 5. Bind the render function to the UI state # Marimo will automatically re-run this when the button or radio changes chart = render_batch_viz(agent_type.value, sample_btn.value) mo.vstack([ui_stack, chart]) return @app.cell(hide_code=True) def _(json, mo, wasm_iframe): # 1. Narrative fix_text = mo.md( """ ### The Paper's Stated Limitations Before we reveal the fix, it's worth noting what Pathak et al. themselves acknowledged. The Inverse Model perfectly filters noise that is *independent* of the agent's actions (the TV with random static). But it fails when the noise is **action-dependent**: > If the environment contains a remote control that changes the TV to a random channel, > the Inverse Model *cannot* filter it out. Because the agent's action directly caused > the unpredictable state change, the agent will sit and press the button forever. The paper flags this as an open problem. Our Snake failure is a different but related variant: **terminal states** are action-dependent (the death is caused by the agent's move), and the physics of a game-over reset is maximally unpredictable. The paper's own framework predicts this should be a problem - and it is. --- ### The Fix: One Line Changes Everything The death trap has a clean solution - **zero out the intrinsic reward on terminal steps**. When the game ends (`done=True`), `(1 − done) = 0`, so the massive death-surprise is suppressed before it ever enters the replay buffer. PER can no longer oversample deaths with inflated priorities. ```python intrinsic_reward = agent.train_icm(state_old, final_move, state_new) intrinsic_reward *= (1.0 - float(done)) # ← The entire fix reward += intrinsic_reward ``` The chart below overlays the **actual Mean Score curves** from two training runs on the same 10×10 board for 16,000 games. Watch the red curve plateau near **8.1** while the green one climbs to **10.1**. In this run, the poisoned agent does **not** visibly collapse or regress. The failure mode is subtler: it learns a worse policy and then stalls there. That still matters - the terminal mask lifts final performance by roughly **25%** while changing only one line of code. This comparison strengthens the case that the failure mode is real in this setup, but it still does **not** replace a full robustness study across seeds and hyperparameters. """ ) # 2. Inlined score curves (parsed from training logs; inlined for molab/WASM compatibility) _g_poi = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900, 2000, 2100, 2200, 2300, 2400, 2500, 2600, 2700, 2800, 2900, 3000, 3100, 3200, 3300, 3400, 3500, 3600, 3700, 3800, 3900, 4000, 4100, 4200, 4300, 4400, 4500, 4600, 4700, 4800, 4900, 5000, 5100, 5200, 5300, 5400, 5500, 5600, 5700, 5800, 5900, 6000, 6100, 6200, 6300, 6400, 6500, 6600, 6700, 6800, 6900, 7000, 7100, 7200, 7300, 7400, 7500, 7600, 7700, 7800, 7900, 8000, 8100, 8200, 8300, 8400, 8500, 8600, 8700, 8800, 8900, 9000, 9100, 9200, 9300, 9400, 9500, 9600, 9700, 9800, 9900, 10000, 10100, 10200, 10300, 10400, 10500, 10600, 10700, 10800, 10900, 11000, 11100, 11200, 11300, 11400, 11500, 11600, 11700, 11800, 11900, 12000, 12100, 12200, 12300, 12400, 12500, 12600, 12700, 12800, 12900, 13000, 13100, 13200, 13300, 13400, 13500, 13600, 13700, 13800, 13900, 14000, 14100, 14200, 14300, 14400, 14500, 14600, 14700, 14800, 14900, 15000, 15100, 15200, 15300, 15400, 15500, 15600, 15700, 15800, 15900, 16000] _s_poi = [0.11, 0.13, 0.13, 0.14, 0.15, 0.16, 0.19, 0.23, 0.27, 0.3, 0.34, 0.37, 0.42, 0.46, 0.49, 0.54, 0.6, 0.66, 0.74, 0.84, 0.96, 1.07, 1.23, 1.37, 1.49, 1.61, 1.71, 1.79, 1.88, 1.95, 2.05, 2.17, 2.27, 2.4, 2.53, 2.65, 2.76, 2.88, 2.95, 3.03, 3.12, 3.19, 3.26, 3.38, 3.45, 3.53, 3.61, 3.68, 3.77, 3.85, 3.91, 4.05, 4.15, 4.21, 4.3, 4.39, 4.43, 4.48, 4.6, 4.72, 4.84, 4.88, 4.97, 5.04, 5.13, 5.18, 5.25, 5.3, 5.35, 5.4, 5.47, 5.54, 5.59, 5.63, 5.67, 5.71, 5.75, 5.8, 5.84, 5.92, 5.96, 6.0, 6.04, 6.1, 6.15, 6.19, 6.21, 6.25, 6.3, 6.34, 6.38, 6.42, 6.44, 6.48, 6.51, 6.54, 6.6, 6.64, 6.66, 6.67, 6.67, 6.7, 6.73, 6.77, 6.81, 6.85, 6.86, 6.89, 6.92, 6.95, 6.99, 7.04, 7.08, 7.12, 7.15, 7.18, 7.19, 7.23, 7.25, 7.29, 7.31, 7.33, 7.36, 7.37, 7.39, 7.4, 7.42, 7.45, 7.48, 7.49, 7.51, 7.53, 7.57, 7.6, 7.63, 7.65, 7.66, 7.68, 7.68, 7.69, 7.72, 7.74, 7.76, 7.79, 7.8, 7.84, 7.86, 7.89, 7.9, 7.92, 7.94, 7.93, 7.96, 7.98, 7.99, 8.01, 8.03, 8.03, 8.05, 8.06] _g_fix = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900, 2000, 2100, 2200, 2300, 2400, 2500, 2600, 2700, 2800, 2900, 3000, 3100, 3200, 3300, 3400, 3500, 3600, 3700, 3800, 3900, 4000, 4100, 4200, 4300, 4400, 4500, 4600, 4700, 4800, 4900, 5000, 5100, 5200, 5300, 5400, 5500, 5600, 5700, 5800, 5900, 6000, 6100, 6200, 6300, 6400, 6500, 6600, 6700, 6800, 6900, 7000, 7100, 7200, 7300, 7400, 7500, 7600, 7700, 7800, 7900, 8000, 8100, 8200, 8300, 8400, 8500, 8600, 8700, 8800, 8900, 9000, 9100, 9200, 9300, 9400, 9500, 9600, 9700, 9800, 9900, 10000, 10100, 10200, 10300, 10400, 10500, 10600, 10700, 10800, 10900, 11000, 11100, 11200, 11300, 11400, 11500, 11600, 11700, 11800, 11900, 12000, 12100, 12200, 12300, 12400, 12500, 12600, 12700, 12800, 12900, 13000, 13100, 13200, 13300, 13400, 13500, 13600, 13700, 13800, 13900, 14000, 14100, 14200, 14300, 14400, 14500, 14600, 14700, 14800, 14900, 15000, 15100, 15200, 15300, 15400, 15500, 15600, 15700, 15800, 15900, 16000] _s_fix = [0.12, 0.13, 0.14, 0.18, 0.19, 0.2, 0.23, 0.26, 0.29, 0.35, 0.4, 0.43, 0.46, 0.52, 0.59, 0.71, 0.82, 0.99, 1.17, 1.47, 1.83, 2.24, 2.63, 2.87, 3.15, 3.44, 3.73, 3.92, 4.15, 4.36, 4.57, 4.81, 4.94, 5.1, 5.27, 5.43, 5.57, 5.74, 5.86, 5.94, 6.07, 6.19, 6.34, 6.46, 6.55, 6.65, 6.73, 6.85, 6.94, 7.01, 7.11, 7.19, 7.27, 7.35, 7.44, 7.5, 7.56, 7.58, 7.65, 7.71, 7.81, 7.84, 7.9, 7.95, 8.03, 8.07, 8.12, 8.18, 8.23, 8.25, 8.29, 8.32, 8.37, 8.4, 8.44, 8.47, 8.5, 8.55, 8.56, 8.62, 8.67, 8.72, 8.75, 8.79, 8.81, 8.83, 8.85, 8.88, 8.91, 8.95, 8.96, 8.97, 9.0, 9.04, 9.06, 9.1, 9.12, 9.16, 9.18, 9.21, 9.23, 9.25, 9.27, 9.29, 9.33, 9.33, 9.36, 9.38, 9.41, 9.43, 9.44, 9.44, 9.44, 9.47, 9.48, 9.49, 9.5, 9.49, 9.52, 9.54, 9.56, 9.57, 9.58, 9.6, 9.61, 9.62, 9.65, 9.65, 9.66, 9.66, 9.68, 9.7, 9.69, 9.71, 9.72, 9.72, 9.73, 9.75, 9.78, 9.8, 9.81, 9.83, 9.83, 9.85, 9.85, 9.88, 9.88, 9.89, 9.91, 9.93, 9.94, 9.94, 9.96, 9.99, 10.02, 10.03, 10.03, 10.05, 10.06, 10.06] _max_score = max(max(_s_poi), max(_s_fix)) _fix_json = json.dumps({ "games_poi": _g_poi, "scores_poi": [s / _max_score for s in _s_poi], "games_fix": _g_fix, "scores_fix": [s / _max_score for s in _s_fix], "max_score": _max_score, "poi_final": _s_poi[-1], "fix_final": _s_fix[-1], "last_game": _g_poi[-1], }) # 3. Animated two-phase chart html_fix = f"""

PER + ICM - unmasked (poisoned)

PER + ICM + terminal mask - fixed

Both curves trained identically - the only difference is one line of code. In this run, the unmasked agent plateaus near 8.1 while the masked agent reaches 10.1 (~25% higher).

""" source_accordion = mo.accordion({ "📂 Source code & training scripts": mo.md(""" **GitHub:** [github.com/Saheb/rl-snake](https://github.com/Saheb/rl-snake) | File | Purpose | |------|---------| | [`scripts/train_dqn.py`](https://github.com/Saheb/rl-snake/blob/main/scripts/train_dqn.py) | DQN + PER + ICM training loop - terminal mask fix at line 547 | | [`utils/icm.py`](https://github.com/Saheb/rl-snake/blob/main/utils/icm.py) | ICM module (encoder, inverse model, forward model) | """) }) mo.vstack([fix_text, wasm_iframe(html_fix, height="430px"), source_accordion]) return @app.cell(hide_code=True) def _(mo): repro_text = mo.md( """ --- ### Following Up: Does the Failure Reproduce? The "death oversampling" story above is mechanistically clean, and the score gap in the chart (8.1 → 10.1) is real - measured on an earlier code revision. But when we re-ran the comparison against the present implementation with **direct buffer-composition instrumentation**, we found something more nuanced. Both the unmasked ("poisoned") and masked ("fixed") versions show roughly the **same** terminal-state oversampling. Single seed, 8×8, 3,000 games, η = 0.1, no priority cap, no foundation memory: | Metric | Unmasked ICM | Masked ICM | Difference | |--------|-------------:|-----------:|----------:| | Buffer terminal % | 17.4 % | 19.3 % | −1.9 pp | | Sampled terminal % | 24.8 % | 28.4 % | −3.6 pp | | **Terminal oversample factor** | **1.44 ×** | **1.48 ×** | **−0.04** | | Mean score (cumulative) | 3.50 | 3.21 | +0.29 | The masked version doesn't sample fewer terminals - and on this seed it scores *slightly worse*. The dramatic plateau-vs-fix gap from the original log does not reproduce reliably. **Why the difference?** The likely culprit is the **state representation** evolved from 14 dims to 24 dims (adding tail-direction one-hot, normalized snake length, normalized Manhattan food distance, and four normalized wall distances on top of the original danger / heading / food-direction / flood-fill features). Richer features mean the ICM forward model has lower MSE on novel transitions, so intrinsic rewards never grow large enough to dominate PER's priorities. The original log was real; the failure mode it captured was an artifact of a specific (now-superseded) configuration. The terminal mask remains a sensible defensive measure but is **not load-bearing** in the present setup. --- ### A Better Question: When Does Curiosity *Actually* Help on Snake? That is the thesis of everything that follows. The original ICM paper paired curiosity with on-policy A3C on Mario, where extrinsic reward is rare and exploration *is* the bottleneck. Snake is a different animal: | Event | Reward | |-------|-------:| | Move closer to food | +0.1 | | Move further from food | −0.1 | | Waffling / loop | −0.5 | | Step | −0.01 | | Eat food | +1 | | Die | −1 | | Win (fill board) | +5 | With distance shaping firing every single step, the extrinsic gradient already screams the answer at the agent. **In this regime, curiosity has no problem to solve.** To put ICM on harder experimental ground, we strip the shaping out and ask: does curiosity rescue learning when the extrinsic reward actually becomes sparse? | Reward Mode | Step Penalty | Distance Shaping | Anti-Loop Penalty | Food / Death | |------------|:------------:|:----------------:|:-----------------:|:------------:| | **dense** *(current)* | −0.01 | ±0.1 | −0.5 | +1 / −1 | | **sparse** | −0.01 | - | - | +1 / −1 | | **pure_sparse** | - | - | - | +1 / −1 | Crossed with `{DQN, DQN + ICM}`, run on 8×8 (and 10×10 for the most extreme `pure_sparse` regime), 3 seeds × 5,000 games each. """ ) repro_text return @app.cell(hide_code=True) def _(mo): # 3-seed means of cumulative `Final Mean Score` (mean score over all 5,000 # games of training) per condition. Parsed from pilot_logs/* by # scripts/parse_pilot_logs.py. Format: (mean, std). All cells at n=3. sparsity_scoreboard = { "8x8": { "dense": {"DQN": (4.69, 0.11), "DQN+ICM": (4.73, 0.11)}, "sparse": {"DQN": (4.08, 0.12), "DQN+ICM": (4.17, 0.07)}, "pure_sparse": {"DQN": (4.20, 0.14), "DQN+ICM": (4.06, 0.18)}, }, "10x10": { "dense": {"DQN": (5.19, 0.18), "DQN+ICM": (5.21, 0.09)}, "sparse": {"DQN": (4.36, 0.06), "DQN+ICM": (4.24, 0.21)}, "pure_sparse": {"DQN": (4.38, 0.08), "DQN+ICM": (4.40, 0.13)}, }, } def _fmt(cell): if cell is None: return "_running_" m, s = cell return f"{m:.2f} ± {s:.2f}" def _delta(_cells): a, b = _cells.get("DQN"), _cells.get("DQN+ICM") if a is None or b is None: return "-" d = b[0] - a[0] sign = "+" if d >= 0 else "" return f"{sign}{d:.2f}" _rows = [] for _board, _modes in sparsity_scoreboard.items(): for _mode, _cells in _modes.items(): _rows.append( f"| **{_board}** | `{_mode}` | {_fmt(_cells.get('DQN'))} | " f"{_fmt(_cells.get('DQN+ICM'))} | {_delta(_cells)} |" ) table_md = ( "| Board | Reward mode | DQN | DQN + ICM | Δ (ICM − base) |\n" "|-------|-------------|----:|----------:|---------------:|\n" + "\n".join(_rows) ) scoreboard_heading = mo.md( "### Scoreboard: Final Mean Score (3-seed mean ± std, 5,000 games)" ) scoreboard_table = mo.md(table_md) scoreboard_text = mo.md( r""" **What we expected vs. what we got:** - **H1 - ICM is irrelevant in dense Snake** → ✓ **Confirmed.** The 8×8 and 10×10 dense cells show |Δ| < 0.05 - well inside seed noise. With distance shaping firing every step, the policy gradient swamps the intrinsic signal. - **H2 - ICM rescues sparse learning for DQN** → ✗ **Falsified.** Across `sparse` and `pure_sparse` on 8×8, |Δ| stays under 0.15. ICM doesn't help DQN even when the extrinsic signal is gone. - **H3 - Dense shaping upper-bounds everything** → ✓ **Confirmed.** Stripping shaping costs ~0.6 score on 8×8 (4.69 → 4.08) and ~0.8 on 10×10 (5.19 → 4.38). Shaping is doing real work for DQN. **Experiment scope note.** PPO was run only on the 10×10 board and therefore does not appear in the 8×8 rows above. The 10×10 board is where state-space coverage starts to matter: at 8×8 every DQN condition saturates at ~98.4% within 5,000 games, so there is no coverage gap for curiosity to exploit; at 10×10 the same saturation holds for DQN, but the larger board is the natural setting for an on-policy learner whose credit-assignment horizon can extend further. **Mechanism in one sentence.** Intrinsic reward is added to the TD target, raising the TD error for surprising transitions; in sparse environments this promotes exploration, but in dense environments the existing reward shaping dominates and the intrinsic signal becomes negligible. **Why DQN can't consume ICM's signal - the causal chain.** ICM emits a per-transition intrinsic reward (forward-model latent-prediction error): $$r_i = \eta \cdot \lVert \phi(s') - \hat{\phi}_\theta(s, a) \rVert^2$$ Three things are *supposed* to happen to that scalar: 1. **It enters the TD target.** The target becomes $y = (r_e + r_i) + \gamma \max_{a'} Q(s', a')$, so a novel transition looks like it has higher value - exactly the intended bias. 2. **PER amplifies it.** Prioritized replay samples transitions with probability $p \propto |\delta|^\alpha$ where $\delta = y - Q(s, a)$. Elevated $r_i$ inflates $\delta$, which inflates $p$, so novel transitions are over-sampled relative to uniform replay. PER turns the intrinsic signal into a *sampling pressure* on top of a value-target shift. 3. **The breakdown.** $r_i$ is computed *once*, at the moment the transition is collected, by the ICM forward model that existed at that moment. By the time PER selects that transition for a gradient step (often thousands of updates later), the forward model has trained further. Its current estimate of novelty for the same $(s, a, s')$ is much lower than what's stored. The Q-network is being fit to *yesterday's* surprise - a value target that the world model has since revised downward. The intrinsic gradient is structurally there, but it's pointing at a moving target that has already moved. PPO has no such gap: each rollout produces transitions whose $r_i$ is computed by the same ICM that the value function then updates against, in the same step. Intrinsic and extrinsic estimates stay consistent; the advantage function sees a coherent reward signal. There is no stale-$r_i$ problem because there is no buffer. To test that interpretation, we need an on-policy learner - see the PPO results below. """ ) scoreboard_callout = mo.callout( mo.md( "**Coverage note.** Every DQN condition above visits a near-identical " "fraction of the reachable `(head, food)` state space within 5,000 games " "- 8×8 conditions all land at exactly **4,032 / 4,096 ≈ 98.4%**, and on " "10×10 every condition (DQN, DQN+ICM, every reward mode) sits in a " "5-state band around 98.9%. See the coverage chart further down for the " "full picture. Exploration *coverage* is not the bottleneck at this horizon " "for DQN; sample-efficient *learning* from explored states is. That's why " "ICM, which targets coverage, can't move the needle for DQN here." ), kind="info", ) mo.vstack([scoreboard_heading, scoreboard_table, scoreboard_text, scoreboard_callout]) return @app.cell(hide_code=True) def _(mo): # PPO 10x10 final mean score across seeds. Parsed from # pilot_logs/ppo_*_10x10_g5000_seed*.log by scripts/parse_pilot_logs.py. # Format: (mean, std, n_seeds). All cells at n=3. ppo_scoreboard = { "dense": {"PPO": (0.12, 0.04, 3), "PPO+ICM": (0.14, 0.04, 3)}, "pure_sparse": {"PPO": (5.36, 0.76, 3), "PPO+ICM": (6.63, 0.79, 3)}, } def _fmt_ppo(cell): if cell is None: return "_running_" m, s, n = cell if n <= 1: return f"{m:.2f} (n=1)" return f"{m:.2f} ± {s:.2f}" def _delta_ppo(_cells): a = _cells.get("PPO") b = _cells.get("PPO+ICM") if a is None or b is None: return "-" d = b[0] - a[0] rel = (d / a[0] * 100) if a[0] > 0.05 else None sign = "+" if d >= 0 else "" if rel is None: return f"{sign}{d:.2f}" return f"{sign}{d:.2f} ({sign}{rel:.0f}%)" _rows = [] for _mode, _cells in ppo_scoreboard.items(): _rows.append( f"| `{_mode}` | {_fmt_ppo(_cells['PPO'])} | " f"{_fmt_ppo(_cells['PPO+ICM'])} | {_delta_ppo(_cells)} |" ) ppo_table = ( "| Reward mode | PPO | PPO + ICM | Δ |\n" "|-------------|----:|----------:|--:|\n" + "\n".join(_rows) ) ppo_intro = mo.md( """ --- ### Switching to On-Policy: PPO + ICM on 10×10 Snake If our DQN-stale-signal interpretation is right, swapping the off-policy DQN for an on-policy PPO should let the curiosity bonus actually steer the policy. Same environment, same 24-dim state, same ICM module, same η = 0.1, same 5,000 games - only the learner changes. See `scripts/train_ppo.py`. """ ) ppo_table_md = mo.md(ppo_table) ppo_dense_callout = mo.callout( mo.md( "**Dense reward: PPO stalls at ~0.12 (vs DQN's ~5.2 in the same environment).**" " PPO updates on full rollout episodes; with −0.1/step shaping, a typical 10–20" " step episode returns ~−1 to −2 cumulative reward before death. The value" " function learns all states are deeply negative and the sparse +1 food signal" " can never flip the advantage positive. DQN sidesteps this via off-policy TD:" " each food transition bootstraps directly regardless of surrounding negatives." " ICM doesn't rescue it — this is a value-bootstrapping failure, not an" " exploration failure." ), kind="neutral", ) ppo_text = mo.md( """ **Key finding:** on `pure_sparse`, PPO + ICM beats PPO baseline by +24% across 3 seeds (6.63 ± 0.79 vs. 5.36 ± 0.76). The effect is robust — the *worst* ICM seed (5.75) still beats the median baseline seed (4.95), and the best ICM run (7.27) is well outside the baseline distribution. **And the mechanism is not what you'd expect from the textbook ICM story** — see the coverage chart that follows. Both PPO baseline and PPO + ICM saturate at ~98.9% of the 10×10 state space, so the +24% score gain is *not* exploration: both agents already see the whole board. What ICM contributes is **per-step reward densification** — a continuously non-zero novelty signal that PPO's advantage estimator can credit-assign over, in the regime where the extrinsic signal is one sparse `+1` per food. Functionally, ICM here behaves less like an exploration bonus and more like a *learned shaping function*, a self-supervised replacement for the distance shaping we deliberately removed. PPO consumes it cleanly because the gradient is on-policy and fresh; DQN's replay buffer does not get the same benefit (see the DQN null above). This sharpens — rather than confirms — the H4 thesis: ICM substitutes for *shaping*, not for *exploration*. """ ) mo.vstack([ppo_intro, ppo_table_md, ppo_dense_callout, ppo_text]) return @app.cell(hide_code=True) def _(mo): try: import matplotlib.pyplot as _plt except ImportError: _plt = None _COVERAGE_BARS = [ {"label": "DQN dense", "color": "tab:blue", "mean": 98.94, "std": 0.02645751311064333, "n": 3}, {"label": "DQN dense + ICM", "color": "tab:red", "mean": 98.95333333333333, "std": 0.02516611478423459, "n": 3}, {"label": "DQN sparse", "color": "tab:blue", "mean": 98.93333333333334, "std": 0.011547005383790217, "n": 3}, {"label": "DQN sparse + ICM", "color": "tab:red", "mean": 98.89, "std": 0.052915026221294716, "n": 3}, {"label": "DQN pure_sparse", "color": "tab:blue", "mean": 98.91, "std": 0.04358898943540784, "n": 3}, {"label": "DQN pure_sparse +ICM", "color": "tab:red", "mean": 98.92, "std": 0.036055512754636664, "n": 3}, {"label": "PPO dense", "color": "tab:orange", "mean": 51.79666666666667, "std": 2.0867279011249504, "n": 3}, {"label": "PPO dense + ICM", "color": "tab:red", "mean": 50.67, "std": 1.5998749951168063, "n": 3}, {"label": "PPO pure_sparse", "color": "tab:orange", "mean": 98.87, "std": 0.05567764362830701, "n": 3}, {"label": "PPO pure_sparse +ICM", "color": "tab:red", "mean": 98.89666666666666, "std": 0.08386497083606279, "n": 3}, ] def _render_coverage_chart(): if _plt is None: return mo.md("_(matplotlib unavailable; coverage chart skipped)_") rows = _COVERAGE_BARS fig, ax = _plt.subplots(figsize=(9, 5)) labels = [r["label"] for r in rows] means = [r["mean"] for r in rows] stds = [r["std"] for r in rows] colors = [r["color"] for r in rows] ns = [r["n"] for r in rows] y = list(range(len(labels))) ax.barh( y, means, xerr=stds, color=colors, edgecolor="black", linewidth=0.5, capsize=3, alpha=0.85, ) ax.set_yticks(y) ax.set_yticklabels(labels) ax.invert_yaxis() ax.set_xlim(0, 108) ax.axvline(100, color="gray", lw=0.7, ls=":") ax.set_xlabel("Cumulative state coverage at game 5000 (% of 10,000 reachable states)") ax.set_title( "Coverage saturation: DQN flatlines, PPO splits by reward mode, ICM never moves it", fontsize=11, ) for i, (m, s, n) in enumerate(zip(means, stds, ns)): tag = f"{m:.1f}%" + (f" ±{s:.1f}" if n > 1 else "") ax.text(m + max(s, 0.5) + 1, i, tag, va="center", fontsize=9) ax.grid(axis="x", alpha=0.3) fig.tight_layout() return fig coverage_caption = mo.md( """ --- ### Visual: State Coverage Across All 10 Final-Cell Conditions (10×10) One picture, three findings: - **DQN (blue / red, top six bars)**: identical 98.9% across every reward mode and every algorithm variant. ICM never moves coverage for DQN. - **PPO `dense` (orange / red, middle two bars)**: stuck at ~51% because episodes terminate within 10-20 steps under the negative shaping. ICM changes nothing here either - the agent dies before novelty has time to act. - **PPO `pure_sparse` (orange / red, bottom two bars)**: jumps back to 98.9% once the shaping is removed and episodes can run. ICM and baseline match to within 0.1 percentage points. Read across: the only thing that moves coverage on this problem is **whether the agent stays alive long enough to walk the board**, which is governed by the reward structure, not by ICM. The +24% score gain in `PPO pure_sparse + ICM` therefore cannot be explained by exploration. It must come from somewhere else - the per-step reward densification effect described above. """ ) coverage_chart = _render_coverage_chart() mo.vstack([coverage_caption, coverage_chart]) return @app.cell(hide_code=True) def _(mo): try: import matplotlib.pyplot as _plt import numpy as _np except ImportError: _plt = None _np = None _GAMES = list(range(100, 5001, 100)) _LEARNING_CURVES = { "DQN_dense_baseline": {"games": _GAMES, "n": 3, "mean": [0.11,0.12,0.123,0.127,0.127,0.13,0.137,0.163,0.177,0.187,0.203,0.22,0.26,0.303,0.36,0.4,0.45,0.5,0.567,0.64,0.727,0.837,0.963,1.133,1.317,1.537,1.76,1.99,2.16,2.37,2.543,2.73,2.917,3.09,3.297,3.47,3.637,3.787,3.927,4.087,4.22,4.337,4.477,4.583,4.703,4.79,4.913,5.03,5.12,5.187], "std": [0.062,0.017,0.032,0.021,0.012,0.0,0.006,0.006,0.012,0.021,0.015,0.035,0.035,0.04,0.04,0.046,0.052,0.044,0.068,0.066,0.091,0.111,0.144,0.142,0.136,0.155,0.182,0.184,0.168,0.177,0.19,0.187,0.194,0.19,0.197,0.192,0.237,0.245,0.238,0.211,0.217,0.211,0.206,0.206,0.205,0.208,0.179,0.191,0.182,0.185]}, "DQN_dense_icm": {"games": _GAMES, "n": 3, "mean": [0.113,0.12,0.12,0.107,0.103,0.113,0.127,0.137,0.157,0.177,0.197,0.24,0.273,0.31,0.35,0.407,0.463,0.54,0.61,0.677,0.757,0.867,0.97,1.147,1.307,1.53,1.773,1.983,2.223,2.433,2.637,2.84,3.037,3.203,3.353,3.513,3.67,3.79,3.937,4.073,4.22,4.357,4.477,4.577,4.693,4.79,4.897,5.0,5.123,5.213], "std": [0.061,0.05,0.03,0.015,0.021,0.031,0.032,0.032,0.035,0.045,0.031,0.035,0.049,0.052,0.044,0.038,0.032,0.044,0.061,0.055,0.055,0.055,0.078,0.055,0.055,0.069,0.074,0.064,0.06,0.101,0.116,0.13,0.148,0.136,0.112,0.136,0.164,0.182,0.161,0.118,0.106,0.11,0.096,0.074,0.067,0.113,0.103,0.125,0.107,0.085]}, "DQN_pure_sparse_baseline": {"games": _GAMES, "n": 3, "mean": [0.113,0.11,0.11,0.13,0.14,0.167,0.17,0.203,0.23,0.257,0.293,0.33,0.357,0.387,0.43,0.477,0.527,0.57,0.643,0.713,0.8,0.89,1.017,1.137,1.34,1.47,1.637,1.813,1.99,2.15,2.307,2.44,2.59,2.707,2.833,2.95,3.067,3.197,3.307,3.41,3.51,3.613,3.693,3.797,3.887,3.99,4.087,4.18,4.28,4.38], "std": [0.035,0.017,0.01,0.01,0.02,0.015,0.026,0.031,0.026,0.025,0.038,0.036,0.035,0.035,0.046,0.045,0.035,0.046,0.045,0.032,0.04,0.046,0.025,0.032,0.036,0.046,0.055,0.06,0.046,0.075,0.04,0.062,0.061,0.032,0.031,0.046,0.04,0.031,0.057,0.026,0.01,0.029,0.055,0.085,0.107,0.121,0.112,0.087,0.075,0.079]}, "DQN_pure_sparse_icm": {"games": _GAMES, "n": 3, "mean": [0.107,0.117,0.123,0.127,0.147,0.163,0.173,0.193,0.223,0.237,0.257,0.29,0.337,0.373,0.413,0.453,0.507,0.563,0.637,0.7,0.77,0.87,1.017,1.17,1.377,1.56,1.74,1.87,2.027,2.157,2.307,2.45,2.573,2.723,2.877,3.027,3.153,3.273,3.337,3.46,3.557,3.677,3.777,3.873,3.98,4.07,4.157,4.22,4.317,4.397], "std": [0.021,0.025,0.021,0.015,0.006,0.006,0.006,0.015,0.032,0.029,0.031,0.017,0.029,0.032,0.025,0.015,0.006,0.021,0.015,0.01,0.03,0.01,0.023,0.066,0.076,0.062,0.056,0.02,0.051,0.042,0.045,0.062,0.067,0.072,0.068,0.09,0.117,0.111,0.11,0.115,0.098,0.099,0.116,0.097,0.114,0.13,0.133,0.13,0.124,0.133]}, "PPO_dense_baseline": {"games": _GAMES, "n": 3, "mean": [0.13,0.143,0.133,0.12,0.117,0.117,0.113,0.117,0.113,0.11,0.11,0.107,0.113,0.113,0.107,0.113,0.113,0.107,0.107,0.103,0.107,0.103,0.103,0.1,0.1,0.1,0.1,0.097,0.1,0.1,0.1,0.1,0.097,0.097,0.097,0.097,0.1,0.1,0.1,0.103,0.107,0.107,0.107,0.11,0.11,0.11,0.117,0.117,0.117,0.12], "std": [0.056,0.031,0.023,0.026,0.021,0.006,0.012,0.006,0.006,0.0,0.0,0.006,0.006,0.006,0.006,0.006,0.006,0.006,0.006,0.006,0.006,0.006,0.006,0.01,0.01,0.01,0.01,0.006,0.01,0.01,0.01,0.01,0.012,0.012,0.012,0.012,0.017,0.017,0.017,0.015,0.021,0.021,0.021,0.03,0.03,0.03,0.035,0.035,0.035,0.04]}, "PPO_dense_icm": {"games": _GAMES, "n": 3, "mean": [0.16,0.147,0.137,0.13,0.12,0.12,0.117,0.12,0.113,0.11,0.11,0.11,0.11,0.11,0.107,0.103,0.103,0.1,0.1,0.097,0.097,0.093,0.093,0.093,0.097,0.1,0.1,0.1,0.1,0.103,0.103,0.107,0.107,0.107,0.113,0.113,0.113,0.117,0.12,0.12,0.123,0.127,0.13,0.133,0.133,0.137,0.143,0.143,0.143,0.143], "std": [0.026,0.006,0.021,0.01,0.01,0.01,0.012,0.01,0.015,0.01,0.01,0.01,0.01,0.01,0.012,0.006,0.006,0.01,0.01,0.006,0.012,0.006,0.006,0.006,0.012,0.01,0.01,0.01,0.017,0.015,0.015,0.021,0.021,0.021,0.025,0.025,0.025,0.031,0.03,0.03,0.035,0.035,0.04,0.04,0.04,0.035,0.04,0.04,0.04,0.04]}, "PPO_pure_sparse_baseline": {"games": _GAMES, "n": 3, "mean": [0.13,0.143,0.147,0.153,0.167,0.2,0.23,0.27,0.323,0.373,0.437,0.49,0.547,0.593,0.647,0.747,0.857,0.983,1.123,1.25,1.39,1.53,1.69,1.833,1.96,2.127,2.28,2.41,2.553,2.68,2.773,2.9,3.023,3.183,3.303,3.42,3.51,3.653,3.783,3.897,4.01,4.15,4.263,4.417,4.573,4.737,4.88,5.037,5.193,5.363], "std": [0.056,0.055,0.042,0.038,0.023,0.035,0.046,0.044,0.055,0.05,0.051,0.062,0.065,0.071,0.074,0.122,0.18,0.237,0.301,0.348,0.347,0.401,0.404,0.437,0.469,0.482,0.468,0.459,0.479,0.451,0.463,0.433,0.445,0.402,0.393,0.39,0.373,0.35,0.289,0.284,0.291,0.25,0.232,0.25,0.284,0.297,0.395,0.491,0.613,0.76]}, "PPO_pure_sparse_icm": {"games": _GAMES, "n": 3, "mean": [0.16,0.157,0.167,0.173,0.197,0.24,0.283,0.34,0.407,0.477,0.54,0.597,0.663,0.733,0.813,0.93,1.05,1.183,1.34,1.497,1.667,1.76,1.927,2.103,2.297,2.473,2.66,2.797,2.933,3.12,3.28,3.467,3.687,3.887,4.053,4.247,4.397,4.563,4.737,4.877,5.02,5.213,5.37,5.54,5.69,5.877,6.033,6.233,6.417,6.627], "std": [0.026,0.015,0.012,0.021,0.025,0.036,0.047,0.05,0.055,0.06,0.075,0.096,0.133,0.176,0.219,0.295,0.338,0.395,0.471,0.501,0.545,0.54,0.539,0.483,0.528,0.532,0.512,0.425,0.415,0.387,0.382,0.358,0.272,0.215,0.195,0.188,0.186,0.199,0.232,0.232,0.248,0.277,0.306,0.4,0.44,0.501,0.532,0.587,0.688,0.786]}, } def _render_chart(): if _plt is None or _np is None: return mo.md("_(matplotlib/numpy unavailable; chart skipped - see scoreboards above)_") data = _LEARNING_CURVES panels = [ ("DQN_dense_baseline", "DQN_dense_icm", "DQN, 10×10 dense", "tab:blue"), ("DQN_pure_sparse_baseline", "DQN_pure_sparse_icm", "DQN, 10×10 pure_sparse", "tab:blue"), ("PPO_dense_baseline", "PPO_dense_icm", "PPO, 10×10 dense", "tab:orange"), ("PPO_pure_sparse_baseline", "PPO_pure_sparse_icm", "PPO, 10×10 pure_sparse", "tab:orange"), ] fig, axes = _plt.subplots(2, 2, figsize=(11, 7), sharex=True) for ax, (k_base, k_icm, title, base_color) in zip(axes.flat, panels): for key, label, ls, color in [ (k_base, "baseline", "-", base_color), (k_icm, "+ ICM", "--", "tab:red"), ]: if key not in data: continue d = data[key] gs = _np.array(d["games"]) mu = _np.array(d["mean"]) sd = _np.array(d["std"]) n = d["n"] seed_label = f"{label} (n={n})" ax.plot(gs, mu, ls=ls, color=color, lw=1.8, label=seed_label) if n > 1: ax.fill_between(gs, mu - sd, mu + sd, color=color, alpha=0.15, linewidth=0) ax.set_title(title, fontsize=11) ax.grid(True, alpha=0.3) ax.legend(loc="upper left", fontsize=9, framealpha=0.9) ax.set_ylabel("Mean Score (cumulative)") for ax in axes[-1, :]: ax.set_xlabel("Training Game") fig.suptitle( "Learning curves: ICM helps PPO under sparse reward, but not DQN under any reward", fontsize=12, y=1.00, ) fig.tight_layout() return fig chart_caption = mo.md( """ --- ### Visual: Learning Curves Across the 4 Headline Cells Solid blue/orange = baseline (no ICM). Dashed red = ICM-augmented. Shaded band is ±1 σ across seeds (no band drawn when n = 1). All panels are 10×10 boards, 5,000 training games. Read top-to-bottom: **DQN curves track each other** in both reward modes - the dashed-red ICM line sits inside the seed band of the solid baseline, so any difference is noise. **PPO curves diverge in the bottom-right panel** - under `pure_sparse`, PPO + ICM separates from PPO baseline within the first ~2,000 games and stays above through training. (PPO `dense` panel: both curves stuck at ~0.1 - credit-assignment failure under negative shaping, as discussed above.) """ ) chart_obj = _render_chart() mo.vstack([chart_caption, chart_obj]) return @app.cell(hide_code=True) def _(mo): takeaway_text = mo.md( """ --- ### What This Investigation Actually Showed We started with a clean story (terminal mask fixes a death-oversampling bug), couldn't reproduce it reliably, and used that miss as an excuse to ask the harder question - *when does curiosity matter on Snake at all?* The 2 × 3 × 2 × 3 grid (algorithm × reward mode × board × seed) gave a sharper answer than we expected: | Question | Answer | Evidence | |---|---|---| | Does ICM help DQN with dense reward? | No | |Δ| < 0.05 across boards (H1 ✓) | | Does ICM rescue DQN under sparse reward? | **No** | |Δ| < 0.15 across `sparse` and `pure_sparse` (H2 ✗) | | Does dense shaping matter for DQN? | Yes | ~15% drop without it (H3 ✓) | | Can vanilla PPO learn dense 10×10 Snake? | **No** | Mean 0.12 apples after 5k games | | Does ICM help PPO under sparse reward? | **Yes** | +24% on `pure_sparse` 10×10 (6.63 vs 5.36, n=3) | | Does ICM increase state coverage? | **No** | PPO and PPO+ICM both reach ~98.9% on `pure_sparse` 10×10 | **The mechanistic takeaway.** Two effects, both contrary to the standard "ICM = exploration bonus" framing. First, ICM is not a free upgrade across algorithms - DQN's replay buffer dilutes the intrinsic signal across stale transitions, while PPO consumes it fresh on every rollout. Second, even where ICM helps (PPO `pure_sparse`), it does not help by *expanding coverage* - both PPO and PPO + ICM reach the same ~98.9% of the state space. ICM's contribution on this problem is **per-step reward densification**: turning a one-`+1`-per-food sparse signal into a continuously non-zero novelty signal that PPO's advantage estimator can credit-assign over. ICM behaves less like an exploration bonus and more like a learned, self-supervised replacement for the hand-engineered distance shaping we removed. **What we'd want to test next.** The +24% PPO + ICM effect holds across 3 seeds with the worst ICM seed (5.75) beating the best baseline seed within the seed band (5.36 ± 0.76 baseline vs 6.63 ± 0.79 ICM). What the investigation does *not* yet pin down: whether the effect generalizes to (a) larger boards (12×12, 15×15) where coverage actually fails to saturate at 5,000 games, (b) longer horizons (10,000+ games - does ICM keep helping or does the gap close?), and (c) a fairer comparison against the curriculum-trained 10×10 PPO baseline that already lives in `scripts/train_curriculum_10x10.py`. If ICM matches curriculum without a hand-engineered training schedule, the H4 thesis ("curiosity substitutes for curriculum") becomes the cleanest one-line claim of the project. """ ) takeaway_text return @app.cell(hide_code=True) def _(mo): _conclusion = mo.md( """ --- ### Conclusion: Curiosity, Conditioned The Intrinsic Curiosity Module (Pathak et al., 2017) reframed exploration in reinforcement learning by equating *surprise with reward* - a single elegant move that let agents learn in environments with no extrinsic feedback. What this notebook stress-tested is the *boundary* of that elegance. The original finding it was built around - that a terminal-mask fix unlocks ICM's performance on Snake - does not reproduce reliably under the present 24-dim state representation. Rather than paper over the failed reproduction, we used it as motivation to ask the more interesting question: **when does curiosity actually help on this problem?** The 2 × 3 × 2 × 3 grid (algorithm × reward mode × board × seed) gave a sharper answer than the textbook ICM story would predict: - **Algorithm matters.** ICM gave DQN nothing - across every reward mode and board size, |Δ| stayed inside seed noise. The same module on PPO `pure_sparse` delivered +24% (n = 3, robust per seed). ICM's value is bottlenecked by whether the learner can consume the intrinsic signal *fresh*, before replay-buffer staleness washes it out. - **Reward structure matters more.** Where ICM did help, the mechanism was not exploration - both PPO and PPO + ICM saturated state coverage at ~98.9%. ICM's actual contribution was *per-step reward densification*: a continuous novelty signal that filled in for the hand-engineered distance shaping we deliberately removed. Functionally, ICM behaved as a self-supervised replacement for reward shaping, not as an exploration bonus. **Curiosity is not universally helpful - it is highly dependent on the reward structure of the environment.** --- ### Beyond ICM: Where the Research Went Next The lineage of papers below can be read as systematic responses to the same boundary conditions this investigation traced - each one repairing a different failure mode of pure forward-prediction curiosity. """ ) _lit_tabs = mo.ui.tabs({ "ICM (2017)": mo.md(""" **Pathak et al., ICML 2017** - [arXiv:1705.05363](https://arxiv.org/abs/1705.05363) **Problem solved:** Dense exploration in sparse-reward environments without hand-coded reward shaping. **Key result:** Self-supervised agents matched or exceeded hand-shaped baselines on VizDoom and Super Mario Bros - with zero domain knowledge injected. **Relation to this notebook:** The foundation. Every section above builds on or stress-tests this paper's core mechanism. """), "Large-Scale Curiosity (2018)": mo.md(""" **Burda, Edwards, Pathak et al., ICLR 2019** - [arXiv:1808.04355](https://arxiv.org/abs/1808.04355) **Problem solved:** Does intrinsic curiosity *alone* (zero extrinsic reward) produce meaningful learning? **Key result:** Yes - agents trained purely on curiosity learned 54 Atari environments and navigated 3D mazes, with no game score signal whatsoever. **Relation:** Validates ICM's hypothesis at scale. Also found the Noisy TV problem is endemic, motivating the next entry. """), "RND (2018)": mo.md(""" **Burda et al., ICLR 2019** - [arXiv:1810.12894](https://arxiv.org/abs/1810.12894) **Problem solved:** The Noisy TV problem - unpredictable environment stochasticity corrupts ICM's forward model prediction. **Key result:** State-of-the-art on Montezuma's Revenge by replacing forward prediction with *random network distillation* - a frozen random network whose output a second network learns to predict. Novelty = gap between the two. Stochasticity cannot corrupt a random target. **Relation:** RND discards ICM's forward model entirely. No physics prediction, no noise vulnerability. """), "Episodic Curiosity (2018)": mo.md(""" **Savinov et al., ICLR 2019** - [arXiv:1810.02274](https://arxiv.org/abs/1810.02274) **Problem solved:** Within-episode reward collapse - an agent that finds one novel state and farms it forever instead of exploring further. **Key result:** Outperformed ICM on long-horizon navigation tasks by gating rewards through *episodic reachability*: only states that are hard to reach from anything seen this episode earn reward. **Relation:** ICM's reward decays only as the forward model learns. Episodic Curiosity forces continued movement regardless of model quality. """), "DreamerV3 (2023)": mo.md(""" **Hafner et al., DeepMind 2023** - [arXiv:2301.04104](https://arxiv.org/abs/2301.04104) **Problem solved:** Sample efficiency and generalization across radically different environments without any task-specific hyperparameter tuning. **Key result:** Human-level performance on 150+ tasks including Minecraft diamond collection - with a single fixed hyperparameter set. **Relation:** ICM's forward model predicts one step ahead in latent space. Dreamer learns a complete *world model* - dynamics, rewards, and terminations - and trains its policy entirely inside imagined rollouts. The lineage from ICM's single forward model to Dreamer's full generative world model is direct and unbroken. """), "Agent57 (2020)": mo.md(""" **Badia et al., DeepMind 2020** - [arXiv:2003.13350](https://arxiv.org/abs/2003.13350) **Problem solved:** Achieving above-human performance on *all* 57 classic Atari games simultaneously - including the hard-exploration games that had resisted every prior method. **Key result:** First agent to surpass human baselines on the complete Atari-57 suite. Used a meta-controller to balance intrinsic vs. extrinsic reward, combining episodic novelty (short-term) with a lifelong novelty metric (long-term). **Relation:** Agent57 is the convergence point of the ICM lineage. It synthesizes episodic curiosity, RND-style novelty, and adaptive exploration weighting into a single agent that never stops exploring but knows when to exploit. """), }) mo.vstack([ _conclusion, _lit_tabs, ]) return @app.cell(hide_code=True) def _(): import base64 import marimo as mo import json import random try: with mo.capture_stderr(): import matplotlib.pyplot as plt except ImportError: plt = None try: import numpy as np except ImportError: np = None def wasm_iframe(html: str, *, width: str = "100%", height: str = "400px"): encoded = base64.b64encode(html.encode("utf-8")).decode("ascii") return mo.Html( f'' ) return json, mo, np, plt, random, wasm_iframe if __name__ == "__main__": app.run()