Upload 29 files

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.py[cod]
+*$py.class
+.pytest_cache/
+.mypy_cache/
+.ipynb_checkpoints/
+*Scaffold.html
+*.html

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Wayfinder Team
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,61 @@
----
-license: apache-2.0
----
+# ARC-AGI-3 Wayfinder Agent
+
+Competition agent for the **ARC Prize 2026 — ARC-AGI-3 track** (Kaggle Code Competition).
+
+## Architecture
+
+The agent uses a **hybrid world-model + planning** architecture combining four
+modules mapped to the benchmark's core capabilities:
+
+| Module | Role | Capability |
+|--------|------|------------|
+| Perception Encoder | CNN over one-hot 64×64×16 frames → compact latent | — |
+| World/Transition Model | Self-supervised P(frame changes | state, action) + forward model | Modeling |
+| State Memory Graph | Hash-deduplicated directed graph of observed states | Exploration |
+| Intrinsic Reward | Extrinsic Δscore + graph novelty + prediction-error curiosity | Goal-setting |
+| Planner | Short-horizon tree search using world model as simulator | Planning |
+| Action Head | Hierarchical: action-type softmax + conv coordinate head for ACTION6 | — |
+
+## Quick Start
+
+```bash
+# Install
+uv pip install -e ".[dev]"
+
+# Run against a public game
+uv run main.py --agent=wayfinder --game=ls20
+
+# Run tests
+uv run pytest
+
+# Offline evaluation
+uv run python eval/run_local_eval.py --agent=wayfinder --games=ls20,ls21,ls22
+```
+
+## Repository Layout
+
+```
+agents/wayfinder/     # Core agent modules (perception, world_model, memory_graph, ...)
+training/             # Replay buffer, training loops, configs
+eval/                 # Offline evaluation harness, metrics
+notebooks/            # Kaggle submission notebook
+tests/                # Unit + integration tests
+```
+
+## Key Constraints
+
+- **No internet at inference time** — Kaggle scoring sessions disable network access.
+- **MIT/CC0 license** — all authored code; third-party deps must be permissively licensed.
+- **Action budget** — agent self-terminates stuck levels (~5× human median actions).
+- **No per-game hardcoding** — same code runs against all unseen games.
+
+## Reproducing Results
+
+1. Install dependencies: `uv pip install -e ".[dev]"`
+2. Download public games via the SDK's local mode.
+3. Run evaluation: `uv run python eval/run_local_eval.py --agent=wayfinder`
+4. Results are logged to `eval/results/` with per-game/level breakdowns.
+
+## License
+
+MIT — see [LICENSE](LICENSE).

agents/__init__.py

@@ -0,0 +1 @@
+"""ARC-AGI-3 agent package."""

agents/wayfinder/__init__.py

@@ -0,0 +1,15 @@
+"""Wayfinder agent for ARC-AGI-3.
+
+A hybrid world-model + planning agent that combines:
+- A learned perception encoder (CNN)
+- A self-supervised world/transition model
+- A hash-deduplicated state memory graph
+- An intrinsic reward module (extrinsic + novelty + curiosity)
+- A short-horizon tree-search planner
+- A hierarchical action head (type selection + coordinate prediction)
+"""
+
+from agents.wayfinder.agent import WayfinderAgent
+
+__all__ = ["WayfinderAgent"]
+__version__ = "0.1.0"

agents/wayfinder/action_head.py

@@ -0,0 +1,256 @@
+"""Action head — Module F.
+
+Hierarchical action selection: first choose an action type (RESET,
+ACTION1–ACTION5, ACTION6), then for ACTION6 choose a click coordinate
+via a conv-based coordinate head producing a 64×64 heatmap.
+
+The action type is selected via softmax over per-type change-probabilities
+(from the world model) with epsilon-greedy exploration. The coordinate
+head uses a small CNN to produce a spatial heatmap over the 64×64 grid.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from agents.wayfinder.world_model import WorldModel
+
+logger = logging.getLogger(__name__)
+
+GRID_SIZE = 64
+NUM_COLORS = 16
+
+
+class CoordinateHead(nn.Module):
+    """Conv-based coordinate prediction head for ACTION6.
+
+    Takes the one-hot frame + latent as input and produces a 64×64
+    heatmap over click targets. This is more sample-efficient than a
+    flat 4096-way softmax because it exploits spatial structure.
+
+    Architecture:
+        Input: (one_hot[16, 64, 64] + latent_tile[1, 64, 64]) → 17 channels
+        Conv2d(17, 32, 3, padding=1) → ReLU
+        Conv2d(32, 16, 3, padding=1) → ReLU
+        Conv2d(16, 1, 1) → Sigmoid → 64×64 heatmap
+    """
+
+    def __init__(self, latent_dim: int = 256) -> None:
+        """Initialize the coordinate head.
+
+        Args:
+            latent_dim: Dimension of the input latent vector.
+        """
+        super().__init__()
+        self.latent_dim = latent_dim
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(NUM_COLORS + 1, 32, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(32, 16, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(16, 1, kernel_size=1),
+        )
+
+    def forward(self, one_hot: torch.Tensor, latent: torch.Tensor) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            one_hot: (B, 16, 64, 64) one-hot frame.
+            latent: (B, latent_dim) latent vector.
+
+        Returns:
+            (B, 64, 64) heatmap with values in [0, 1].
+        """
+        # Broadcast latent to spatial dimensions
+        latent_map = latent[:, :1].unsqueeze(-1).unsqueeze(-1)  # (B, 1, 1, 1)
+        latent_map = latent_map.expand(-1, -1, GRID_SIZE, GRID_SIZE)  # (B, 1, 64, 64)
+
+        x = torch.cat([one_hot, latent_map], dim=1)  # (B, 17, 64, 64)
+        heatmap = self.conv(x).squeeze(1)  # (B, 64, 64)
+        return torch.sigmoid(heatmap)
+
+
+class ActionHead:
+    """Hierarchical action selection head.
+
+    Phase 1: Choose action type via softmax over per-type change-probabilities.
+    Phase 2: For ACTION6, use the coordinate head to pick a click target.
+
+    Attributes:
+        latent_dim: Input latent dimension.
+        device: Torch device.
+        coord_head: Coordinate prediction CNN.
+    """
+
+    def __init__(self, latent_dim: int = 256, device: str = "cpu") -> None:
+        """Initialize the action head.
+
+        Args:
+            latent_dim: Latent dimension from the perception encoder.
+            device: Torch device.
+        """
+        self.latent_dim = latent_dim
+        self.device = torch.device(device)
+        self.coord_head = CoordinateHead(latent_dim=latent_dim).to(self.device)
+        self.coord_head.eval()
+
+        logger.info("ActionHead initialized (latent_dim=%d, device=%s)", latent_dim, device)
+
+    def select(
+        self,
+        latent: np.ndarray,
+        diff_mask: np.ndarray,
+        available_actions: list[str],
+        world_model: WorldModel,
+        epsilon: float = 0.1,
+    ) -> dict[str, object]:
+        """Select an action hierarchically.
+
+        Args:
+            latent: Current state latent vector.
+            diff_mask: 64×64 boolean diff mask (where frame changed).
+            available_actions: Actions available this step.
+            world_model: World model for change prediction.
+            epsilon: Exploration probability for epsilon-greedy.
+
+        Returns:
+            Action dict with "action" (str) and optional "data" (dict).
+        """
+        # --- Epsilon-greedy: random exploration ---
+        if np.random.random() < epsilon:
+            action_name = str(np.random.choice(available_actions))
+            if action_name == "ACTION6":
+                return self._select_coordinate(latent)
+            return {"action": action_name}
+
+        # --- Phase 1: Choose action type ---
+        # Compute P(frame changes) for each available action
+        change_probs = []
+        for action_name in available_actions:
+            if action_name == "RESET":
+                change_probs.append(0.1)  # RESET always changes the frame
+                continue
+            prob = world_model.predict_change(latent, {"action": action_name})
+            change_probs.append(prob)
+
+        change_probs_arr = np.array(change_probs, dtype=np.float32)
+
+        # Softmax with temperature
+        temperature = 0.5
+        logits = np.log(change_probs_arr + 1e-7) / temperature
+        probs = np.exp(logits - np.max(logits))
+        probs = probs / probs.sum()
+
+        selected_idx = np.random.choice(len(available_actions), p=probs)
+        action_name = available_actions[selected_idx]
+
+        # --- Phase 2: For ACTION6, select coordinates ---
+        if action_name == "ACTION6":
+            return self._select_coordinate(latent)
+
+        return {"action": action_name}
+
+    def _select_coordinate(self, latent: np.ndarray) -> dict[str, object]:
+        """Select a click coordinate for ACTION6.
+
+        Uses the coordinate head to produce a 64×64 heatmap, then samples
+        from it as a categorical distribution. Falls back to uniform random
+        if the heatmap is degenerate.
+
+        Args:
+            latent: Current state latent vector.
+
+        Returns:
+            Action dict with "action": "ACTION6" and "data": {"x": int, "y": int}.
+        """
+        # Build a dummy one-hot from the latent (in practice, we'd use the
+        # actual frame; here we create a placeholder since the action head
+        # receives latents, not raw frames — in the full implementation,
+        # we'd pass the frame through or cache it)
+        one_hot = torch.zeros(1, NUM_COLORS, GRID_SIZE, GRID_SIZE, device=self.device)
+        # Place the latent-derived signal in channel 0 as a spatial broadcast
+        lat_tensor = torch.from_numpy(latent).unsqueeze(0).to(self.device)
+        one_hot[:, 0] = lat_tensor[:, :1].unsqueeze(-1).expand(-1, -1, GRID_SIZE).squeeze(0).unsqueeze(0)
+
+        with torch.no_grad():
+            heatmap = self.coord_head(one_hot, lat_tensor)  # (1, 64, 64)
+            heatmap = heatmap.squeeze(0).cpu().numpy()  # (64, 64)
+
+        # Flatten and sample
+        flat = heatmap.flatten()
+        flat = flat - flat.min()
+        total = flat.sum()
+
+        if total < 1e-7:
+            # Degenerate heatmap — uniform random
+            x = int(np.random.randint(0, GRID_SIZE))
+            y = int(np.random.randint(0, GRID_SIZE))
+        else:
+            probs = flat / total
+            idx = np.random.choice(GRID_SIZE * GRID_SIZE, p=probs)
+            y, x = divmod(idx, GRID_SIZE)
+
+        return {"action": "ACTION6", "data": {"x": int(x), "y": int(y)}}
+
+    def train_coordinate_head(
+        self,
+        frames: np.ndarray,
+        latents: np.ndarray,
+        targets: np.ndarray,
+        lr: float = 1e-3,
+        epochs: int = 10,
+    ) -> float:
+        """Train the coordinate head on labeled click targets.
+
+        Used during offline training when click data is available.
+
+        Args:
+            frames: (B, 64, 64) uint8 frames.
+            latents: (B, latent_dim) latent vectors.
+            targets: (B, 2) array of (x, y) click coordinates.
+            lr: Learning rate.
+            epochs: Number of training epochs.
+
+        Returns:
+            Final training loss.
+        """
+        self.coord_head.train()
+        optimizer = torch.optim.Adam(self.coord_head.parameters(), lr=lr)
+
+        # Convert frames to one-hot
+        one_hot = np.zeros((len(frames), NUM_COLORS, GRID_SIZE, GRID_SIZE), dtype=np.float32)
+        for i, frame in enumerate(frames):
+            for c in range(NUM_COLORS):
+                one_hot[i, c] = (frame == c).astype(np.float32)
+
+        one_hot_t = torch.from_numpy(one_hot).to(self.device)
+        latents_t = torch.from_numpy(latents.astype(np.float32)).to(self.device)
+
+        # Create target heatmaps (Gaussian around target)
+        target_maps = torch.zeros(len(frames), GRID_SIZE, GRID_SIZE, device=self.device)
+        for i, (x, y) in enumerate(targets):
+            for dy in range(-3, 4):
+                for dx in range(-3, 4):
+                    ny, nx = int(y) + dy, int(x) + dx
+                    if 0 <= ny < GRID_SIZE and 0 <= nx < GRID_SIZE:
+                        target_maps[i, ny, nx] = np.exp(-(dx**2 + dy**2) / 2.0)
+
+        loss_fn = nn.BCELoss()
+        final_loss = 0.0
+
+        for epoch in range(epochs):
+            optimizer.zero_grad()
+            pred = self.coord_head(one_hot_t, latents_t)
+            loss = loss_fn(pred, target_maps)
+            loss.backward()
+            optimizer.step()
+            final_loss = loss.item()
+
+        self.coord_head.eval()
+        return final_loss

agents/wayfinder/agent.py

@@ -0,0 +1,280 @@
+"""Main agent class — glue connecting all modules.
+
+Subclasses the official ARC-AGI-3 SDK Agent base class, wiring together
+perception, world model, memory graph, intrinsic reward, planner, and
+action head into a single play loop.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+
+from agents.wayfinder.action_head import ActionHead
+from agents.wayfinder.intrinsic_reward import IntrinsicReward
+from agents.wayfinder.memory_graph import MemoryGraph
+from agents.wayfinder.perception import PerceptionEncoder
+from agents.wayfinder.planner import Planner
+from agents.wayfinder.world_model import WorldModel
+
+logger = logging.getLogger(__name__)
+
+
+class WayfinderAgent:
+    """Hybrid world-model + planning agent for ARC-AGI-3.
+
+    The agent maintains:
+    - A perception encoder that converts raw 64×64×4-bit frames into latents.
+    - A world model that predicts frame changes and forward dynamics.
+    - A memory graph that tracks visited states for novelty.
+    - An intrinsic reward module combining extrinsic score, novelty, curiosity.
+    - A planner that does short-horizon tree search using the world model.
+    - An action head with hierarchical action-type + coordinate selection.
+
+    The main loop per step:
+    1. Encode current frame → latent + diff mask
+    2. Update memory graph with the new state
+    3. Train world model on the latest transition (online)
+    4. Compute intrinsic reward for the current state
+    5. If world model confidence is high, use planner; else use reactive policy
+    6. Select action via action head (hierarchical)
+    7. Return the action and reasoning blob for audit
+
+    Attributes:
+        max_actions: Maximum actions before self-terminating a level.
+        action_count: Actions taken in the current level attempt.
+        _encoder: Perception encoder module.
+        _world_model: World/transition model.
+        _memory: State memory graph.
+        _reward: Intrinsic reward calculator.
+        _planner: Short-horizon tree search planner.
+        _action_head: Hierarchical action selection head.
+    """
+
+    def __init__(
+        self,
+        max_actions: int = 1000,
+        latent_dim: int = 256,
+        buffer_size: int = 200_000,
+        device: str = "cpu",
+    ) -> None:
+        """Initialize the Wayfinder agent.
+
+        Args:
+            max_actions: Safety cap on actions per level attempt.
+            latent_dim: Size of the perception encoder's output latent.
+            buffer_size: Maximum transitions stored in the replay buffer.
+            device: Torch device ("cpu" or "cuda").
+        """
+        self.max_actions = max_actions
+        self.action_count = 0
+        self.current_score: float = 0.0
+        self.win_threshold: float = 1.0
+        self._device = device
+
+        # Module A: Perception encoder
+        self._encoder = PerceptionEncoder(latent_dim=latent_dim, device=device)
+
+        # Module B: World/transition model
+        self._world_model = WorldModel(
+            latent_dim=latent_dim,
+            buffer_size=buffer_size,
+            device=device,
+        )
+
+        # Module C: State memory graph
+        self._memory = MemoryGraph()
+
+        # Module D: Intrinsic reward
+        self._reward = IntrinsicReward()
+
+        # Module E: Planner
+        self._planner = Planner(
+            world_model=self._world_model,
+            reward_module=self._reward,
+            max_depth=5,
+            max_simulations=50,
+        )
+
+        # Module F: Action head
+        self._action_head = ActionHead(
+            latent_dim=latent_dim,
+            device=device,
+        )
+
+        # Previous state tracking
+        self._prev_latent: np.ndarray | None = None
+        self._prev_frame_hash: str | None = None
+        self._is_done = False
+
+        logger.info("WayfinderAgent initialized (device=%s, max_actions=%d)", device, max_actions)
+
+    def is_done(self, frames: list[np.ndarray], state: str, **kwargs: Any) -> bool:
+        """Check if the agent should stop playing the current level.
+
+        The agent self-terminates when:
+        - The game state is WIN or GAME_OVER
+        - The action budget is exhausted
+        - The session is NOT_STARTED (needs RESET)
+
+        Args:
+            frames: List of 64×64 frames from the latest step.
+            state: Current game state string.
+            **kwargs: Additional context (score, etc.).
+
+        Returns:
+            True if the agent should stop, False to continue.
+        """
+        if state in ("WIN", "GAME_OVER", "NOT_STARTED"):
+            self._is_done = True
+            return True
+
+        if self.action_count >= self.max_actions:
+            logger.warning(
+                "Action budget exhausted (%d/%d) — self-terminating level",
+                self.action_count,
+                self.max_actions,
+            )
+            self._is_done = True
+            return True
+
+        return False
+
+    def act(
+        self,
+        frames: list[np.ndarray],
+        state: str,
+        score: float,
+        win_score: float,
+        available_actions: list[str],
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Select the next action given the current observation.
+
+        This is the main per-step entry point called by the SDK harness.
+
+        Args:
+            frames: List of 64×64 numpy arrays (4-bit color indices 0–15).
+            state: Game state string ("NOT_FINISHED", etc.).
+            score: Current running score.
+            win_score: Score threshold to win the level.
+            available_actions: List of action names available this step.
+            **kwargs: Additional SDK fields (levels_completed, etc.).
+
+        Returns:
+            Dict with "action" (str), "data" (dict, for ACTION6 x/y),
+            and "reasoning" (dict, for audit/logging).
+        """
+        self.action_count += 1
+        self.current_score = score
+        self.win_threshold = win_score
+
+        # Use the last frame if multiple are returned (animation settled)
+        current_frame = frames[-1] if frames else np.zeros((64, 64), dtype=np.uint8)
+
+        # --- Step 1: Encode frame ---
+        latent, diff_mask = self._encoder.encode(current_frame)
+        frame_hash = self._memory.hash_frame(current_frame)
+
+        # --- Step 2: Update memory graph ---
+        score_delta = score - self.current_score if self._prev_latent is not None else 0.0
+        if self._prev_frame_hash is not None:
+            # We don't know the action yet — we'll record the edge after
+            # the action is chosen. For now, register the node.
+            pass
+        self._memory.add_node(frame_hash, latent, score)
+
+        novelty = self._memory.novelty(frame_hash)
+
+        # --- Step 3: Online world-model update ---
+        if self._prev_latent is not None and self._last_action is not None:
+            self._world_model.add_transition(
+                state_latent=self._prev_latent,
+                action=self._last_action,
+                next_latent=latent,
+                frame_changed=bool(np.any(diff_mask)),
+            )
+            self._world_model.train_step()
+
+        # --- Step 4: Intrinsic reward ---
+        prediction_error = self._world_model.prediction_error(
+            self._prev_latent, self._last_action, latent
+        ) if self._prev_latent is not None and self._last_action else 0.0
+
+        utility = self._reward.compute(
+            extrinsic_delta=score_delta,
+            novelty=novelty,
+            prediction_error=prediction_error,
+        )
+
+        # --- Step 5: Plan or react ---
+        model_confidence = self._world_model.confidence()
+        if model_confidence > 0.65 and self.action_count > 10:
+            # Use planner when the world model is confident
+            action = self._planner.plan(
+                latent=latent,
+                available_actions=available_actions,
+                action_budget_remaining=self.max_actions - self.action_count,
+                utility_fn=self._reward.compute,
+            )
+        else:
+            # Reactive policy via action head
+            action = self._action_head.select(
+                latent=latent,
+                diff_mask=diff_mask,
+                available_actions=available_actions,
+                world_model=self._world_model,
+                epsilon=0.15 if novelty > 0.5 else 0.05,
+            )
+
+        # --- Step 6: Record transition in memory graph ---
+        if self._prev_frame_hash is not None:
+            self._memory.add_edge(
+                from_hash=self._prev_frame_hash,
+                to_hash=frame_hash,
+                action=action["action"],
+                action_data=action.get("data", {}),
+            )
+
+        # --- Step 7: Update previous state ---
+        self._prev_latent = latent
+        self._prev_frame_hash = frame_hash
+        self._last_action = action
+
+        # Build reasoning blob for audit (≤16 KB)
+        reasoning = {
+            "step": self.action_count,
+            "score": score,
+            "win_score": win_score,
+            "model_confidence": model_confidence,
+            "novelty": novelty,
+            "prediction_error": prediction_error,
+            "utility": utility,
+            "mode": "plan" if model_confidence > 0.65 else "react",
+            "frame_hash": frame_hash[:16],
+        }
+
+        logger.debug("Step %d: action=%s, reasoning=%s", self.action_count, action["action"], reasoning)
+
+        return {"action": action["action"], "data": action.get("data", {}), "reasoning": reasoning}
+
+    def reset(self) -> None:
+        """Reset agent state for a new level attempt.
+
+        Clears per-level state: action count, memory graph, previous latents.
+        The world model and perception encoder retain their learned weights
+        across resets (they generalize across levels within a game).
+        """
+        self.action_count = 0
+        self.current_score = 0.0
+        self._prev_latent = None
+        self._prev_frame_hash = None
+        self._last_action: dict[str, Any] | None = None
+        self._is_done = False
+        self._memory.reset()
+        logger.info("Agent reset for new level attempt")
+
+    # Expose _last_action with a default for the first step
+    _last_action: dict[str, Any] | None = None

agents/wayfinder/cli.py

@@ -0,0 +1,101 @@
+"""CLI entry point for the Wayfinder agent.
+
+Wraps the SDK's main.py interface so the agent can be run via:
+    uv run main.py --agent=wayfinder --game=ls20
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+
+logger = logging.getLogger(__name__)
+
+
+def main() -> int:
+    """CLI entry point.
+
+    Returns:
+        Exit code (0 for success).
+    """
+    parser = argparse.ArgumentParser(
+        description="Wayfinder agent for ARC-AGI-3"
+    )
+    parser.add_argument(
+        "--agent", default="wayfinder",
+        help="Agent name (default: wayfinder)"
+    )
+    parser.add_argument(
+        "--game", default="ls20",
+        help="Game ID to play (default: ls20)"
+    )
+    parser.add_argument(
+        "--max-actions", type=int, default=1000,
+        help="Max actions per level (default: 1000)"
+    )
+    parser.add_argument(
+        "--device", default="cpu",
+        help="Torch device (default: cpu)"
+    )
+    parser.add_argument(
+        "--verbose", "-v", action="store_true",
+        help="Enable verbose logging"
+    )
+
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.DEBUG if args.verbose else logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    )
+
+    # Try to use the SDK's runner
+    try:
+        from arc_agi_3 import main as sdk_main  # type: ignore[import]
+        logger.info("Using SDK runner with agent=%s, game=%s", args.agent, args.game)
+        # The SDK's main.py handles the game loop
+        sdk_main(agent_name=args.agent, game_id=args.game)
+    except ImportError:
+        logger.error(
+            "arc-agi-3 SDK not found. Install with: uv pip install arc-agi-3"
+        )
+        logger.info("Running in standalone mode (no SDK)...")
+        _run_standalone(args)
+
+    return 0
+
+
+def _run_standalone(args: argparse.Namespace) -> None:
+    """Run the agent without the SDK (for testing/scaffolding).
+
+    Args:
+        args: Parsed CLI arguments.
+    """
+    import numpy as np
+
+    from agents.wayfinder.agent import WayfinderAgent
+
+    agent = WayfinderAgent(
+        max_actions=args.max_actions,
+        device=args.device,
+    )
+
+    # Simulate a few steps with random frames
+    logger.info("Running standalone simulation (5 steps)...")
+    for step in range(5):
+        frame = np.random.randint(0, 16, size=(64, 64), dtype=np.uint8)
+        result = agent.act(
+            frames=[frame],
+            state="NOT_FINISHED",
+            score=0.0,
+            win_score=1.0,
+            available_actions=["ACTION1", "ACTION2", "ACTION3", "ACTION4", "ACTION5"],
+        )
+        logger.info("Step %d: %s", step + 1, result)
+
+    logger.info("Standalone simulation complete.")
+
+
+if __name__ == "__main__":
+    sys.exit(main())

agents/wayfinder/intrinsic_reward.py

@@ -0,0 +1,176 @@
+"""Intrinsic reward module — Module D.
+
+Combines extrinsic score delta, graph novelty, and world-model
+prediction error into a single scalar utility. Weights are
+configurable via a YAML config file (not hard-coded).
+
+The utility function guides both the reactive policy (via action
+selection) and the planner (as the value function for rollouts).
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RewardConfig:
+    """Configuration for the intrinsic reward module.
+
+    Attributes:
+        w_extrinsic: Weight for extrinsic score delta.
+        w_novelty: Weight for graph novelty signal.
+        w_curiosity: Weight for world-model prediction error (curiosity).
+        novelty_decay: Multiplicative decay per visit (reduces novelty
+            of repeatedly-visited states faster).
+        curiosity_clip: Maximum curiosity value (prevents runaway).
+        score_baseline: Baseline score for normalizing extrinsic signal.
+    """
+
+    w_extrinsic: float = 1.0
+    w_novelty: float = 0.3
+    w_curiosity: float = 0.2
+    novelty_decay: float = 0.95
+    curiosity_clip: float = 10.0
+    score_baseline: float = 0.0
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> "RewardConfig":
+        """Load config from a YAML file.
+
+        Args:
+            path: Path to the YAML config file.
+
+        Returns:
+            RewardConfig instance.
+        """
+        with open(path) as f:
+            data = yaml.safe_load(f)
+        return cls(**data.get("intrinsic_reward", {}))
+
+    @classmethod
+    def default(cls) -> "RewardConfig":
+        """Return default configuration."""
+        return cls()
+
+
+class IntrinsicReward:
+    """Computes intrinsic reward combining extrinsic + novelty + curiosity.
+
+    The utility function is:
+        utility = w_extrinsic * Δscore + w_novelty * novelty + w_curiosity * curiosity
+
+    This maps to the "Goal-setting" benchmark capability — the agent
+    learns to value states that increase score, explore new territory,
+    and reduce its model's prediction error.
+
+    Attributes:
+        config: Reward configuration with tunable weights.
+    """
+
+    def __init__(self, config: RewardConfig | None = None) -> None:
+        """Initialize the intrinsic reward module.
+
+        Args:
+            config: Reward configuration. If None, uses defaults.
+        """
+        self.config = config or RewardConfig.default()
+        self._running_baseline = self.config.score_baseline
+        self._update_count = 0
+
+        logger.info(
+            "IntrinsicReward initialized (w_ext=%.2f, w_nov=%.2f, w_cur=%.2f)",
+            self.config.w_extrinsic,
+            self.config.w_novelty,
+            self.config.w_curiosity,
+        )
+
+    def compute(
+        self,
+        extrinsic_delta: float = 0.0,
+        novelty: float = 0.0,
+        prediction_error: float = 0.0,
+        **kwargs: float,
+    ) -> float:
+        """Compute the intrinsic utility for a state transition.
+
+        Args:
+            extrinsic_delta: Change in game score (score_t - score_{t-1}).
+            novelty: Novelty score from the memory graph (0–1).
+            prediction_error: World model's prediction error (curiosity).
+            **kwargs: Additional signals that could be incorporated.
+
+        Returns:
+            Scalar utility value.
+        """
+        # Clip curiosity to prevent runaway values
+        curiosity = min(prediction_error, self.config.curiosity_clip)
+
+        # Normalize extrinsic by running baseline (helps across games
+        # with different score scales)
+        normalized_extrinsic = extrinsic_delta
+        if self._running_baseline > 0:
+            normalized_extrinsic = extrinsic_delta / max(abs(self._running_baseline), 1.0)
+
+        # Update running baseline (exponential moving average)
+        if extrinsic_delta != 0:
+            self._running_baseline = (
+                0.99 * self._running_baseline + 0.01 * abs(extrinsic_delta)
+            )
+        self._update_count += 1
+
+        utility = (
+            self.config.w_extrinsic * normalized_extrinsic
+            + self.config.w_novelty * novelty
+            + self.config.w_curiosity * curiosity
+        )
+
+        return float(utility)
+
+    def compute_for_latent(
+        self,
+        novelty: float,
+        prediction_error: float,
+        score_delta: float = 0.0,
+    ) -> float:
+        """Compute utility for a latent state (used by the planner).
+
+        Args:
+            novelty: Novelty of the state.
+            prediction_error: Model's prediction error for reaching this state.
+            score_delta: Score change when entering this state.
+
+        Returns:
+            Utility value.
+        """
+        return self.compute(
+            extrinsic_delta=score_delta,
+            novelty=novelty,
+            prediction_error=prediction_error,
+        )
+
+    def update_weights(self, **new_weights: float) -> None:
+        """Dynamically update reward weights (e.g. via meta-learning).
+
+        Args:
+            **new_weights: Keyword arguments matching RewardConfig fields.
+        """
+        if "w_extrinsic" in new_weights:
+            self.config.w_extrinsic = new_weights["w_extrinsic"]
+        if "w_novelty" in new_weights:
+            self.config.w_novelty = new_weights["w_novelty"]
+        if "w_curiosity" in new_weights:
+            self.config.w_curiosity = new_weights["w_curiosity"]
+
+        logger.debug("Updated reward weights: %s", new_weights)
+
+    def reset(self) -> None:
+        """Reset per-level state (running baseline)."""
+        self._running_baseline = self.config.score_baseline
+        self._update_count = 0

agents/wayfinder/memory_graph.py

@@ -0,0 +1,233 @@
+"""State memory graph — Module C.
+
+A hash-deduplicated directed graph of observed game states. Nodes are
+hashed frames; edges are actions taken. Tracks visit counts for novelty
+detection and supports backtracking/loop avoidance.
+
+Inspired by Blind Squirrel's graph-based exploration, but integrated
+with the learned perception encoder for efficient similarity checks.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from collections import defaultdict
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryGraph:
+    """Directed graph of observed game states.
+
+    Nodes are identified by frame hashes (MD5 of the raw 64×64 frame).
+    Each node stores:
+    - A latent vector (from the perception encoder) for similarity checks.
+    - The score observed at that state.
+    - A visit count (incremented each time the state is seen).
+
+    Edges store the action that caused the transition.
+
+    The graph resets per RESET (per level attempt), but the agent can
+    optionally carry summary statistics across attempts.
+
+    Attributes:
+        nodes: Dict mapping frame_hash → node data.
+        edges: Dict mapping (from_hash, to_hash) → edge data.
+        adjacency: Dict mapping from_hash → list of (action, to_hash).
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty memory graph."""
+        self.nodes: dict[str, dict] = {}
+        self.edges: dict[tuple[str, str], dict] = {}
+        self.adjacency: dict[str, list[tuple[str, str]]] = defaultdict(list)
+
+    @staticmethod
+    def hash_frame(frame: np.ndarray) -> str:
+        """Compute a stable hash for a 64×64 frame.
+
+        Args:
+            frame: 64×64 uint8 array of color indices.
+
+        Returns:
+            MD5 hex string of the frame data.
+        """
+        return hashlib.md5(frame.tobytes()).hexdigest()
+
+    def add_node(
+        self,
+        frame_hash: str,
+        latent: np.ndarray,
+        score: float = 0.0,
+    ) -> None:
+        """Add or update a node in the graph.
+
+        If the node already exists, increment its visit count.
+
+        Args:
+            frame_hash: Hash of the frame.
+            latent: Latent vector from the perception encoder.
+            score: Game score at this state.
+        """
+        if frame_hash in self.nodes:
+            self.nodes[frame_hash]["visit_count"] += 1
+            self.nodes[frame_hash]["score"] = score
+        else:
+            self.nodes[frame_hash] = {
+                "latent": latent.copy(),
+                "score": score,
+                "visit_count": 1,
+            }
+
+    def add_edge(
+        self,
+        from_hash: str,
+        to_hash: str,
+        action: str,
+        action_data: dict | None = None,
+    ) -> None:
+        """Add a directed edge to the graph.
+
+        Args:
+            from_hash: Source node hash.
+            to_hash: Destination node hash.
+            action: Action name that caused this transition.
+            action_data: Optional action data (e.g. coordinates for ACTION6).
+        """
+        edge_key = (from_hash, to_hash)
+        if edge_key in self.edges:
+            self.edges[edge_key]["count"] += 1
+        else:
+            self.edges[edge_key] = {
+                "action": action,
+                "action_data": action_data,
+                "count": 1,
+            }
+            self.adjacency[from_hash].append((action, to_hash))
+
+    def novelty(self, frame_hash: str) -> float:
+        """Compute novelty score for a frame.
+
+        Novelty is high for unvisited or rarely-visited states, low for
+        frequently-visited ones. Used by the intrinsic reward module.
+
+        Formula: novelty = 1.0 / (1.0 + visit_count)
+
+        Args:
+            frame_hash: Hash of the frame to evaluate.
+
+        Returns:
+            Novelty score in (0, 1]. New states return 1.0.
+        """
+        if frame_hash not in self.nodes:
+            return 1.0
+        visit_count = self.nodes[frame_hash]["visit_count"]
+        return 1.0 / (1.0 + visit_count)
+
+    def shortest_path_to_unexplored(self, from_hash: str) -> list[str] | None:
+        """Find shortest path from a node to any unexplored action.
+
+        Uses BFS to find the nearest node that has untried actions
+        (actions not yet taken from that node).
+
+        Args:
+            from_hash: Starting node hash.
+
+        Returns:
+            List of action names forming the path, or None if no
+            unexplored action is reachable.
+        """
+        if from_hash not in self.nodes:
+            return None
+
+        # BFS
+        from collections import deque
+
+        queue: deque[tuple[str, list[str]]] = deque([(from_hash, [])])
+        visited: set[str] = {from_hash}
+
+        all_actions = {"RESET", "ACTION1", "ACTION2", "ACTION3",
+                       "ACTION4", "ACTION5", "ACTION6"}
+
+        while queue:
+            current_hash, path = queue.popleft()
+
+            # Check if current node has unexplored actions
+            tried_actions = {
+                edge_action
+                for edge_action, _ in self.adjacency.get(current_hash, [])
+            }
+            unexplored = all_actions - tried_actions
+
+            if unexplored and len(path) > 0:
+                return path
+            if unexplored:
+                # We're at the start node with unexplored actions
+                return []
+
+            # Expand neighbors
+            for edge_action, neighbor_hash in self.adjacency.get(current_hash, []):
+                if neighbor_hash not in visited:
+                    visited.add(neighbor_hash)
+                    queue.append((neighbor_hash, path + [edge_action]))
+
+        return None
+
+    def get_visited_states(self) -> set[str]:
+        """Return the set of all visited node hashes.
+
+        Returns:
+            Set of frame hashes that have been visited.
+        """
+        return set(self.nodes.keys())
+
+    def get_transition_count(self) -> int:
+        """Return total number of unique transitions (edges).
+
+        Returns:
+            Number of edges in the graph.
+        """
+        return len(self.edges)
+
+    def stats(self) -> dict:
+        """Return summary statistics about the graph.
+
+        Returns:
+            Dict with node_count, edge_count, avg_visits, max_visits.
+        """
+        visit_counts = [n["visit_count"] for n in self.nodes.values()]
+        return {
+            "node_count": len(self.nodes),
+            "edge_count": len(self.edges),
+            "avg_visits": np.mean(visit_counts) if visit_counts else 0.0,
+            "max_visits": max(visit_counts) if visit_counts else 0,
+        }
+
+    def reset(self) -> None:
+        """Clear all nodes and edges.
+
+        Called when the agent RESETs a level or starts a new one.
+        """
+        self.nodes.clear()
+        self.edges.clear()
+        self.adjacency.clear()
+        logger.debug("Memory graph reset")
+
+    def serialize(self) -> dict:
+        """Serialize the graph for checkpointing.
+
+        Returns:
+            Dict representation of the graph.
+        """
+        return {
+            "nodes": {
+                h: {"score": n["score"], "visit_count": n["visit_count"]}
+                for h, n in self.nodes.items()
+            },
+            "edges": {
+                f"{k[0]}->{k[1]}": v for k, v in self.edges.items()
+            },
+        }

agents/wayfinder/perception.py

@@ -0,0 +1,158 @@
+"""Perception encoder — Module A.
+
+Converts a raw 64×64 grid of 4-bit color indices (0–15) into a compact
+latent representation using a small from-scratch CNN.
+
+The encoder also computes a binary diff mask against the previous frame,
+which is used by the world model to predict frame changes.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+GRID_SIZE = 64
+NUM_COLORS = 16
+LATENT_DIM_DEFAULT = 256
+
+
+class PerceptionEncoder(nn.Module):
+    """CNN encoder for 64×64×16 one-hot frames.
+
+    Architecture (3 conv layers + 1 FC head):
+        Conv2d(16, 32, 3, padding=1) → ReLU → MaxPool2d(2)  # 64→32
+        Conv2d(32, 64, 3, padding=1) → ReLU → MaxPool2d(2)  # 32→16
+        Conv2d(64, 128, 3, padding=1) → ReLU → MaxPool2d(2) # 16→8
+        Flatten → Linear(128*8*8, latent_dim)
+
+    No pretrained weights exist for this domain — trained from scratch.
+
+    Attributes:
+        latent_dim: Output latent vector dimensionality.
+        device: Torch device for inference.
+    """
+
+    def __init__(self, latent_dim: int = LATENT_DIM_DEFAULT, device: str = "cpu") -> None:
+        """Initialize the encoder.
+
+        Args:
+            latent_dim: Output latent dimension.
+            device: Torch device ("cpu" or "cuda").
+        """
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.device = torch.device(device)
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(NUM_COLORS, 32, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 64 → 32
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 32 → 16
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 16 → 8
+        )
+        self.fc = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(128 * 8 * 8, 512),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, latent_dim),
+        )
+
+        self.to(self.device)
+        self.eval()
+        logger.info("PerceptionEncoder initialized (latent_dim=%d, device=%s)", latent_dim, device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass: one-hot frame → latent vector.
+
+        Args:
+            x: Tensor of shape (B, 16, 64, 64) — one-hot encoded frames.
+
+        Returns:
+            Latent tensor of shape (B, latent_dim).
+        """
+        features = self.conv(x)
+        latent = self.fc(features)
+        return latent
+
+    def encode(self, frame: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        """Encode a single frame into a latent vector + diff mask.
+
+        Args:
+            frame: 64×64 numpy array of uint8 color indices (0–15).
+
+        Returns:
+            Tuple of (latent, diff_mask):
+            - latent: float32 array of shape (latent_dim,).
+            - diff_mask: 64×64 bool array — True where frame differs from previous.
+        """
+        one_hot = self._to_one_hot(frame)
+        with torch.no_grad():
+            tensor = torch.from_numpy(one_hot).unsqueeze(0).to(self.device)
+            latent = self.forward(tensor).squeeze(0).cpu().numpy()
+
+        # Compute diff mask against previous frame
+        diff_mask = self._compute_diff(frame)
+
+        # Store current frame for next step's diff
+        self._prev_frame = frame.copy()
+
+        return latent.astype(np.float32), diff_mask
+
+    def _to_one_hot(self, frame: np.ndarray) -> np.ndarray:
+        """Convert a 64×64 integer frame to 16×64×64 one-hot float.
+
+        Args:
+            frame: 64×64 uint8 array with values in [0, 15].
+
+        Returns:
+            16×64×64 float32 one-hot array.
+        """
+        one_hot = np.zeros((NUM_COLORS, GRID_SIZE, GRID_SIZE), dtype=np.float32)
+        for c in range(NUM_COLORS):
+            one_hot[c] = (frame == c).astype(np.float32)
+        return one_hot
+
+    def _compute_diff(self, frame: np.ndarray) -> np.ndarray:
+        """Compute binary diff mask against the previous frame.
+
+        Args:
+            frame: Current 64×64 frame.
+
+        Returns:
+            64×64 bool array — True where pixels changed.
+        """
+        if not hasattr(self, "_prev_frame"):
+            return np.zeros((GRID_SIZE, GRID_SIZE), dtype=bool)
+        return frame != self._prev_frame
+
+    def encode_batch(self, frames: np.ndarray) -> np.ndarray:
+        """Encode a batch of frames into latent vectors.
+
+        Used during offline training of the world model.
+
+        Args:
+            frames: (B, 64, 64) uint8 array.
+
+        Returns:
+            (B, latent_dim) float32 array.
+        """
+        one_hot = np.zeros((len(frames), NUM_COLORS, GRID_SIZE, GRID_SIZE), dtype=np.float32)
+        for i, frame in enumerate(frames):
+            for c in range(NUM_COLORS):
+                one_hot[i, c] = (frame == c).astype(np.float32)
+
+        with torch.no_grad():
+            tensor = torch.from_numpy(one_hot).to(self.device)
+            latents = self.forward(tensor).cpu().numpy()
+
+        return latents.astype(np.float32)

agents/wayfinder/planner.py

@@ -0,0 +1,290 @@
+"""Planner — Module E.
+
+Short-horizon tree search using the world model as a cheap simulator
+and the intrinsic reward module as the value function.
+
+Adapts the rollout/backup structure from the Topdeck ISMCTS agent,
+simplified since ARC-AGI-3 frames are fully observed (no information-set
+handling needed). The search budget scales with the remaining action
+allowance — more budget early in a level, less when running low.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from typing import Any, Callable
+
+import numpy as np
+
+from agents.wayfinder.world_model import WorldModel
+
+logger = logging.getLogger(__name__)
+
+# Type alias for the utility function
+UtilityFn = Callable[..., float]
+
+
+class PlannerNode:
+    """A node in the search tree.
+
+    Attributes:
+        latent: State latent vector at this node.
+        action: Action that led to this node (None for root).
+        parent: Parent node (None for root).
+        children: List of child PlannerNodes.
+        visits: Number of times this node has been visited in search.
+        total_value: Accumulated utility from rollouts through this node.
+        is_terminal: Whether this is a terminal state.
+    """
+
+    __slots__ = ("latent", "action", "action_data", "parent", "children",
+                 "visits", "total_value", "is_terminal")
+
+    def __init__(
+        self,
+        latent: np.ndarray,
+        action: str | None = None,
+        action_data: dict | None = None,
+        parent: "PlannerNode | None" = None,
+        is_terminal: bool = False,
+    ) -> None:
+        self.latent = latent
+        self.action = action
+        self.action_data = action_data
+        self.parent = parent
+        self.children: list[PlannerNode] = []
+        self.visits = 0
+        self.total_value = 0.0
+        self.is_terminal = is_terminal
+
+    @property
+    def mean_value(self) -> float:
+        """Mean utility value from rollouts through this node."""
+        return self.total_value / self.visits if self.visits > 0 else 0.0
+
+    def ucb1(self, exploration: float = 1.414) -> float:
+        """UCB1 selection score.
+
+        Args:
+            exploration: Exploration constant (sqrt(2) by default).
+
+        Returns:
+            UCB1 value for node selection.
+        """
+        if self.visits == 0:
+            return float("inf")
+        parent_visits = self.parent.visits if self.parent else self.visits
+        exploit = self.mean_value
+        explore = exploration * math.sqrt(math.log(parent_visits) / self.visits)
+        return exploit + explore
+
+    def best_child(self, exploration: float = 1.414) -> "PlannerNode | None":
+        """Select the child with the highest UCB1 score.
+
+        Args:
+            exploration: UCB1 exploration constant.
+
+        Returns:
+            Best child node, or None if no children.
+        """
+        if not self.children:
+            return None
+        return max(self.children, key=lambda c: c.ucb1(exploration))
+
+
+class Planner:
+    """Short-horizon tree search planner.
+
+    Uses the world model to simulate forward, the intrinsic reward to
+    evaluate states, and UCB1-based selection (like MCTS/ISMCTS) to
+    balance exploration and exploitation in the search tree.
+
+    The search depth and number of simulations scale with the remaining
+    action budget — more budget means deeper search.
+
+    Attributes:
+        world_model: The world/transition model (used as simulator).
+        max_depth: Maximum search depth per simulation.
+        max_simulations: Number of MCTS simulations per planning step.
+    """
+
+    ALL_ACTIONS = ["ACTION1", "ACTION2", "ACTION3", "ACTION4", "ACTION5"]
+
+    def __init__(
+        self,
+        world_model: WorldModel,
+        reward_module: Any,
+        max_depth: int = 5,
+        max_simulations: int = 50,
+    ) -> None:
+        """Initialize the planner.
+
+        Args:
+            world_model: World model for forward simulation.
+            reward_module: Intrinsic reward module for value estimation.
+            max_depth: Max rollout depth.
+            max_simulations: Number of MCTS simulations per plan() call.
+        """
+        self.world_model = world_model
+        self.reward_module = reward_module
+        self.max_depth = max_depth
+        self.max_simulations = max_simulations
+
+        logger.info(
+            "Planner initialized (max_depth=%d, max_sims=%d)",
+            max_depth,
+            max_simulations,
+        )
+
+    def plan(
+        self,
+        latent: np.ndarray,
+        available_actions: list[str],
+        action_budget_remaining: int,
+        utility_fn: UtilityFn,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Plan the next action using tree search.
+
+        Args:
+            latent: Current state latent.
+            available_actions: Actions available this step.
+            action_budget_remaining: Remaining action budget.
+            utility_fn: Function to compute utility of a state.
+            **kwargs: Additional context.
+
+        Returns:
+            Action dict with "action" and optional "data" keys.
+        """
+        # Scale simulations with remaining budget
+        budget_factor = min(1.0, action_budget_remaining / 100.0)
+        num_sims = max(5, int(self.max_simulations * budget_factor))
+        depth = max(2, int(self.max_depth * budget_factor))
+
+        # Filter to simple actions for tree search (ACTION6's 4096-position
+        # space is too large for full tree search — handled by action head)
+        search_actions = [a for a in available_actions if a in self.ALL_ACTIONS]
+        if not search_actions:
+            search_actions = self.ALL_ACTIONS[:3]
+
+        root = PlannerNode(latent=latent)
+
+        # Run simulations
+        for _ in range(num_sims):
+            self._simulate(root, search_actions, depth, utility_fn)
+
+        # Select the best action from root's children
+        best = root.best_child(exploration=0.0)  # Greedy at root
+        if best is None or best.action is None:
+            # Fallback: pick first available action
+            return {"action": available_actions[0]}
+
+        result: dict[str, Any] = {"action": best.action}
+        if best.action_data:
+            result["data"] = best.action_data
+        return result
+
+    def _simulate(
+        self,
+        root: PlannerNode,
+        actions: list[str],
+        max_depth: int,
+        utility_fn: UtilityFn,
+    ) -> float:
+        """Run one MCTS simulation: select → expand → rollout → backup.
+
+        Args:
+            root: Root of the search tree.
+            actions: Available actions for expansion.
+            max_depth: Maximum rollout depth.
+            utility_fn: Utility function for leaf evaluation.
+
+        Returns:
+            Accumulated utility from the rollout.
+        """
+        # --- Selection ---
+        node = root
+        depth = 0
+        while node.children and not node.is_terminal and depth < max_depth:
+            node = node.best_child()
+            if node is None:
+                break
+            depth += 1
+
+        # --- Expansion ---
+        if node is not None and not node.is_terminal and depth < max_depth:
+            # Expand: add children for all actions
+            for action in actions:
+                action_dict = {"action": action}
+                # Use world model to predict next latent
+                next_latent = self.world_model.predict_next_latent(
+                    node.latent, action_dict
+                )
+                child = PlannerNode(
+                    latent=next_latent,
+                    action=action,
+                    parent=node,
+                )
+                node.children.append(child)
+            # Pick a child for rollout
+            if node.children:
+                node = np.random.choice(node.children)
+
+        # --- Rollout (random forward simulation) ---
+        if node is not None:
+            total_utility = self._rollout(
+                node.latent, actions, max_depth - depth, utility_fn
+            )
+        else:
+            total_utility = 0.0
+
+        # --- Backup ---
+        while node is not None:
+            node.visits += 1
+            node.total_value += total_utility
+            node = node.parent
+
+        return total_utility
+
+    def _rollout(
+        self,
+        latent: np.ndarray,
+        actions: list[str],
+        depth: int,
+        utility_fn: UtilityFn,
+    ) -> float:
+        """Random rollout from a latent state.
+
+        Args:
+            latent: Starting latent.
+            actions: Available actions.
+            depth: Remaining rollout depth.
+            utility_fn: Utility function.
+
+        Returns:
+            Accumulated utility.
+        """
+        total = 0.0
+        current = latent
+
+        for _ in range(max(depth, 1)):
+            action = np.random.choice(actions)
+            action_dict = {"action": action}
+
+            # Simulate forward
+            next_latent = self.world_model.predict_next_latent(current, action_dict)
+
+            # Estimate utility (novelty ≈ 0 in simulation since we're
+            # predicting latents, not real frames; use prediction error
+            # as a proxy for curiosity)
+            pred_error = float(np.linalg.norm(next_latent - current))
+            utility = utility_fn(
+                extrinsic_delta=0.0,
+                novelty=0.0,
+                prediction_error=pred_error,
+            )
+            total += utility * (0.95 ** _)  # Discount
+            current = next_latent
+
+        return total

agents/wayfinder/world_model.py

@@ -0,0 +1,320 @@
+"""World/transition model — Module B.
+
+A self-supervised model that learns two things from played transitions:
+1. P(frame changes | state, action) — a binary change-prediction head.
+2. A forward model ŝ_{t+1} = f(s_t, a_t) predicting the next latent.
+
+Trained online from a deduplicated replay buffer. The change-prediction
+head is the primary signal for the reactive policy (like StochasticGoose),
+while the forward model enables the planner's tree search.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from collections import deque
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Transition:
+    """A single (state, action, next_state, changed) transition.
+
+    Attributes:
+        state_latent: Latent vector of the state before the action.
+        action: Action name string (e.g. "ACTION1", "ACTION6").
+        action_data: Optional action data (e.g. {"x": 32, "y": 32} for ACTION6).
+        next_latent: Latent vector of the state after the action.
+        frame_changed: Whether the frame changed as a result of the action.
+    """
+
+    state_latent: np.ndarray
+    action: str
+    action_data: dict | None
+    next_latent: np.ndarray
+    frame_changed: bool
+
+
+class WorldModel(nn.Module):
+    """World/transition model for ARC-AGI-3.
+
+    Architecture:
+    - Action embedding: 7 action types → 32-dim embedding.
+    - Change predictor: MLP(latent_dim + action_emb) → 1 (sigmoid).
+    - Forward model: MLP(latent_dim + action_emb) → latent_dim.
+
+    The change predictor is trained with BCE loss + entropy regularization.
+    The forward model is trained with MSE loss on the latent difference.
+
+    Attributes:
+        latent_dim: Dimension of the perception encoder's output.
+        buffer_size: Max transitions in the replay buffer.
+        device: Torch device.
+    """
+
+    ACTION_TYPES = ["RESET", "ACTION1", "ACTION2", "ACTION3", "ACTION4", "ACTION5", "ACTION6"]
+    ACTION_TO_IDX = {a: i for i, a in enumerate(ACTION_TYPES)}
+
+    def __init__(
+        self,
+        latent_dim: int = 256,
+        buffer_size: int = 200_000,
+        device: str = "cpu",
+        lr: float = 1e-4,
+    ) -> None:
+        """Initialize the world model.
+
+        Args:
+            latent_dim: Input latent dimensionality.
+            buffer_size: Maximum transitions stored.
+            device: Torch device.
+            lr: Learning rate for the optimizer.
+        """
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.buffer_size = buffer_size
+        self.device = torch.device(device)
+        self.lr = lr
+
+        # Action embedding
+        self.action_embed = nn.Embedding(len(self.ACTION_TYPES), 32)
+
+        # Change prediction head: P(frame changes | state, action)
+        self.change_head = nn.Sequential(
+            nn.Linear(latent_dim + 32, 128),
+            nn.ReLU(inplace=True),
+            nn.Linear(128, 64),
+            nn.ReLU(inplace=True),
+            nn.Linear(64, 1),
+        )
+
+        # Forward model: predict next latent
+        self.forward_head = nn.Sequential(
+            nn.Linear(latent_dim + 32, 256),
+            nn.ReLU(inplace=True),
+            nn.Linear(256, 256),
+            nn.ReLU(inplace=True),
+            nn.Linear(256, latent_dim),
+        )
+
+        self.optimizer = torch.optim.Adam(self.parameters(), lr=lr)
+        self.to(self.device)
+
+        # Replay buffer (deduplicated by state hash)
+        self._buffer: deque[Transition] = deque(maxlen=buffer_size)
+        self._seen_hashes: set[str] = set()
+
+        # Training stats
+        self._train_steps = 0
+        self._change_loss_avg = 0.0
+        self._forward_loss_avg = 0.0
+
+        logger.info("WorldModel initialized (latent_dim=%d, buffer=%d)", latent_dim, buffer_size)
+
+    def forward(self, latent: torch.Tensor, action_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass.
+
+        Args:
+            latent: (B, latent_dim) state latent.
+            action_idx: (B,) action index tensor.
+
+        Returns:
+            Tuple of (change_prob, next_latent_pred):
+            - change_prob: (B, 1) sigmoid probability of frame change.
+            - next_latent_pred: (B, latent_dim) predicted next latent.
+        """
+        act_emb = self.action_embed(action_idx)
+        x = torch.cat([latent, act_emb], dim=-1)
+        change_prob = torch.sigmoid(self.change_head(x))
+        next_latent = self.forward_head(x)
+        return change_prob, next_latent
+
+    def add_transition(
+        self,
+        state_latent: np.ndarray,
+        action: dict[str, Any],
+        next_latent: np.ndarray,
+        frame_changed: bool,
+    ) -> None:
+        """Add a transition to the replay buffer (with deduplication).
+
+        Args:
+            state_latent: Latent before the action.
+            action: Action dict with "action" key and optional "data".
+            next_latent: Latent after the action.
+            frame_changed: Whether the frame visually changed.
+        """
+        # Deduplicate by hashing the state+action combination
+        h = hashlib.md5(
+            state_latent.tobytes() + action["action"].encode()
+        ).hexdigest()
+
+        if h in self._seen_hashes:
+            return  # Skip duplicate
+        self._seen_hashes.add(h)
+
+        self._buffer.append(
+            Transition(
+                state_latent=state_latent.copy(),
+                action=action["action"],
+                action_data=action.get("data"),
+                next_latent=next_latent.copy(),
+                frame_changed=frame_changed,
+            )
+        )
+
+    def train_step(self, batch_size: int = 64) -> float:
+        """Perform one gradient step on a random batch from the buffer.
+
+        Uses BCE loss for change prediction + MSE for forward model
+        + entropy regularization on the change prediction.
+
+        Args:
+            batch_size: Mini-batch size.
+
+        Returns:
+            Total loss value.
+        """
+        if len(self._buffer) < batch_size:
+            return 0.0
+
+        # Sample random batch
+        indices = np.random.choice(len(self._buffer), size=batch_size, replace=False)
+        batch = [self._buffer[i] for i in indices]
+
+        latents = torch.from_numpy(np.stack([t.state_latent for t in batch])).to(self.device)
+        next_latents = torch.from_numpy(np.stack([t.next_latent for t in batch])).to(self.device)
+        action_indices = torch.tensor(
+            [self.ACTION_TO_IDX.get(t.action, 0) for t in batch],
+            dtype=torch.long,
+            device=self.device,
+        )
+        changed = torch.tensor(
+            [float(t.frame_changed) for t in batch],
+            dtype=torch.float32,
+            device=self.device,
+        ).unsqueeze(1)
+
+        self.train()
+        self.optimizer.zero_grad()
+
+        change_prob, next_latent_pred = self.forward(latents, action_indices)
+
+        # BCE loss for change prediction
+        bce_loss = nn.functional.binary_cross_entropy(change_prob, changed)
+
+        # Entropy regularization (encourage calibrated probabilities)
+        eps = 1e-7
+        entropy = -(change_prob * torch.log(change_prob + eps) +
+                    (1 - change_prob) * torch.log(1 - change_prob + eps))
+        entropy_reg = -0.01 * entropy.mean()
+
+        # MSE loss for forward model
+        fwd_loss = nn.functional.mse_loss(next_latent_pred, next_latents)
+
+        total_loss = bce_loss + entropy_reg + 0.5 * fwd_loss
+        total_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.parameters(), 1.0)
+        self.optimizer.step()
+
+        self._train_steps += 1
+        self._change_loss_avg = 0.99 * self._change_loss_avg + 0.01 * bce_loss.item()
+        self._forward_loss_avg = 0.99 * self._forward_loss_avg + 0.01 * fwd_loss.item()
+
+        self.eval()
+        return total_loss.item()
+
+    def predict_change(
+        self, latent: np.ndarray, action: dict[str, Any]
+    ) -> float:
+        """Predict P(frame changes | state, action).
+
+        Args:
+            latent: State latent vector.
+            action: Action dict.
+
+        Returns:
+            Probability of frame change (0–1).
+        """
+        with torch.no_grad():
+            lat = torch.from_numpy(latent).unsqueeze(0).to(self.device)
+            act_idx = torch.tensor(
+                [self.ACTION_TO_IDX.get(action["action"], 0)],
+                dtype=torch.long,
+                device=self.device,
+            )
+            change_prob, _ = self.forward(lat, act_idx)
+            return change_prob.item()
+
+    def predict_next_latent(
+        self, latent: np.ndarray, action: dict[str, Any]
+    ) -> np.ndarray:
+        """Predict the next latent given current state and action.
+
+        Used by the planner as a cheap simulator.
+
+        Args:
+            latent: Current state latent.
+            action: Action dict.
+
+        Returns:
+            Predicted next latent vector.
+        """
+        with torch.no_grad():
+            lat = torch.from_numpy(latent).unsqueeze(0).to(self.device)
+            act_idx = torch.tensor(
+                [self.ACTION_TO_IDX.get(action["action"], 0)],
+                dtype=torch.long,
+                device=self.device,
+            )
+            _, next_latent = self.forward(lat, act_idx)
+            return next_latent.squeeze(0).cpu().numpy()
+
+    def prediction_error(
+        self,
+        state_latent: np.ndarray | None,
+        action: dict[str, Any] | None,
+        actual_next_latent: np.ndarray,
+    ) -> float:
+        """Compute prediction error (curiosity signal).
+
+        Args:
+            state_latent: Latent before action (None if first step).
+            action: Action taken (None if first step).
+            actual_next_latent: Actual next latent.
+
+        Returns:
+            L2 distance between predicted and actual next latent.
+        """
+        if state_latent is None or action is None:
+            return 0.0
+        predicted = self.predict_next_latent(state_latent, action)
+        return float(np.linalg.norm(predicted - actual_next_latent))
+
+    def confidence(self) -> float:
+        """Estimate the world model's confidence.
+
+        Returns a value in [0, 1] based on training progress and
+        average change-prediction loss. Used by the agent to decide
+        whether to use the planner or the reactive policy.
+
+        Returns:
+            Confidence score in [0, 1].
+        """
+        if self._train_steps < 10:
+            return 0.0
+        # Lower loss → higher confidence
+        confidence = 1.0 / (1.0 + self._change_loss_avg * 5)
+        return min(confidence, 1.0)
+
+    @property
+    def buffer_size_current(self) -> int:
+        """Current number of transitions in the buffer."""
+        return len(self._buffer)

eval/metrics.py

@@ -0,0 +1,198 @@
+"""Evaluation metrics for ARC-AGI-3 agent performance.
+
+Tracks per-game/level scores, action efficiency vs. human baseline,
+and world-model prediction quality (AUC, calibration).
+"""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class EvalMetrics:
+    """Tracks and computes evaluation metrics.
+
+    Records per-step data during gameplay and computes aggregate
+    metrics: score progression, action distribution, efficiency,
+    state coverage, and model quality.
+
+    Attributes:
+        steps: List of per-step records.
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty metrics tracker."""
+        self.steps: list[dict] = []
+
+    def record_step(
+        self,
+        score: float,
+        action: str,
+        **extra: object,
+    ) -> None:
+        """Record a single step's metrics.
+
+        Args:
+            score: Current game score.
+            action: Action taken.
+            **extra: Additional metrics (novelty, confidence, etc.).
+        """
+        self.steps.append({
+            "score": score,
+            "action": action,
+            **extra,
+        })
+
+    def compute(self) -> dict:
+        """Compute aggregate metrics.
+
+        Returns:
+            Dict with:
+            - total_steps: Number of actions taken.
+            - final_score: Last recorded score.
+            - action_distribution: Dict of action → count.
+            - unique_actions: Number of distinct actions used.
+            - score_delta: Final score - initial score.
+        """
+        if not self.steps:
+            return {
+                "total_steps": 0,
+                "final_score": 0.0,
+                "action_distribution": {},
+                "unique_actions": 0,
+                "score_delta": 0.0,
+            }
+
+        action_counts: dict[str, int] = defaultdict(int)
+        for step in self.steps:
+            action_counts[step["action"]] += 1
+
+        scores = [s["score"] for s in self.steps]
+        initial_score = scores[0] if scores else 0.0
+        final_score = scores[-1] if scores else 0.0
+
+        return {
+            "total_steps": len(self.steps),
+            "final_score": final_score,
+            "action_distribution": dict(action_counts),
+            "unique_actions": len(action_counts),
+            "score_delta": final_score - initial_score,
+        }
+
+
+def compute_efficiency(metrics: dict, max_actions: int) -> dict:
+    """Compute action efficiency relative to the action budget.
+
+    ARC-AGI-3 scores are squared and action budgets are capped at
+    ~5× human median. This function computes how efficiently the agent
+    used its budget.
+
+    Args:
+        metrics: Output of EvalMetrics.compute().
+        max_actions: The action budget for this level.
+
+    Returns:
+        Dict with:
+        - budget_used: Fraction of budget used (0–1).
+        - actions_per_score: Actions per unit of score gained.
+        - efficiency_score: Composite efficiency metric (0–1, higher is better).
+    """
+    total_steps = metrics.get("total_steps", 0)
+    score_delta = metrics.get("score_delta", 0.0)
+
+    budget_used = total_steps / max(1, max_actions)
+    actions_per_score = total_steps / max(abs(score_delta), 0.001)
+
+    # Efficiency: high score with low budget usage is best
+    if score_delta > 0:
+        efficiency_score = score_delta * (1.0 - 0.5 * budget_used)
+    else:
+        efficiency_score = 0.0
+
+    return {
+        "budget_used": budget_used,
+        "actions_per_score": actions_per_score,
+        "efficiency_score": min(efficiency_score, 1.0),
+    }
+
+
+def compute_change_prediction_auc(
+    predictions: np.ndarray,
+    labels: np.ndarray,
+) -> float:
+    """Compute AUC for the world model's change-prediction head.
+
+    Args:
+        predictions: Predicted probabilities (float array).
+        labels: Binary labels (0 or 1).
+
+    Returns:
+        ROC AUC score (0–1).
+    """
+    if len(predictions) == 0 or len(np.unique(labels)) < 2:
+        return 0.5
+
+    # Sort by prediction descending
+    order = np.argsort(-predictions)
+    labels_sorted = labels[order]
+
+    # Compute ROC AUC via rank-based formula
+    n_pos = labels.sum()
+    n_neg = len(labels) - n_pos
+
+    if n_pos == 0 or n_neg == 0:
+        return 0.5
+
+    # Rank sum
+    ranks = np.zeros(len(predictions))
+    for i, idx in enumerate(order):
+        ranks[idx] = len(predictions) - i
+
+    sum_ranks_pos = ranks[labels == 1].sum()
+    auc = (sum_ranks_pos - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)
+
+    return float(auc)
+
+
+def compare_agents(
+    baseline: dict,
+    candidate: dict,
+) -> dict:
+    """Compare two agents' evaluation results.
+
+    Args:
+        baseline: Baseline agent results.
+        candidate: Candidate agent results.
+
+    Returns:
+        Dict with per-game and aggregate comparisons.
+    """
+    comparison = {
+        "games": {},
+        "aggregate": {},
+    }
+
+    for game_id in baseline.get("games", {}):
+        if game_id in candidate.get("games", {}):
+            b = baseline["games"][game_id]
+            c = candidate["games"][game_id]
+            comparison["games"][game_id] = {
+                "score_delta": c["score"] - b["score"],
+                "action_delta": c["total_actions"] - b["total_actions"],
+                "win_delta": c["levels_won"] - b["levels_won"],
+            }
+
+    b_agg = baseline.get("aggregate", {})
+    c_agg = candidate.get("aggregate", {})
+    comparison["aggregate"] = {
+        "score_delta": c_agg.get("total_score", 0) - b_agg.get("total_score", 0),
+        "win_rate_delta": c_agg.get("win_rate", 0) - b_agg.get("win_rate", 0),
+        "action_delta": c_agg.get("total_actions", 0) - b_agg.get("total_actions", 0),
+    }
+
+    return comparison

eval/run_local_eval.py

@@ -0,0 +1,233 @@
+"""Offline evaluation harness — runs the agent against local games.
+
+Uses the toolkit's local-execution mode so iteration doesn't burn API
+quota. Produces per-game/level score breakdowns and action-efficiency
+metrics compared against random and human baselines.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import time
+from pathlib import Path
+
+import numpy as np
+
+from agents.wayfinder.agent import WayfinderAgent
+from eval.metrics import EvalMetrics, compute_efficiency
+
+logger = logging.getLogger(__name__)
+
+OUTPUT_DIR = Path("eval/results")
+
+
+def run_local_eval(
+    agent_name: str = "wayfinder",
+    games: list[str] | None = None,
+    max_actions_per_level: int = 1000,
+    device: str = "cpu",
+    output_dir: Path | None = None,
+) -> dict:
+    """Run offline evaluation against local games.
+
+    Args:
+        agent_name: Name of the agent to evaluate.
+        games: List of game IDs to evaluate. If None, uses defaults.
+        max_actions_per_level: Max actions per level.
+        device: Torch device.
+        output_dir: Directory for output files.
+
+    Returns:
+        Evaluation results dict.
+    """
+    if games is None:
+        games = ["ls20", "ls21", "ls22"]
+
+    if output_dir is None:
+        output_dir = OUTPUT_DIR
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    agent = WayfinderAgent(
+        max_actions=max_actions_per_level,
+        device=device,
+    )
+
+    all_results = {}
+
+    for game_id in games:
+        logger.info("Evaluating game: %s", game_id)
+        game_result = _evaluate_game(agent, game_id, max_actions_per_level)
+        all_results[game_id] = game_result
+
+    # Compute aggregate metrics
+    total_score = sum(r["score"] for r in all_results.values())
+    total_actions = sum(r["total_actions"] for r in all_results.values())
+    total_levels = sum(r["levels_attempted"] for r in all_results.values())
+    levels_won = sum(r["levels_won"] for r in all_results.values())
+
+    summary = {
+        "agent": agent_name,
+        "games": all_results,
+        "aggregate": {
+            "total_score": total_score,
+            "total_actions": total_actions,
+            "total_levels": total_levels,
+            "levels_won": levels_won,
+            "win_rate": levels_won / max(total_levels, 1),
+            "avg_actions_per_level": total_actions / max(total_levels, 1),
+        },
+        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+    }
+
+    # Save results
+    results_path = output_dir / f"eval_{agent_name}_{int(time.time())}.json"
+    with open(results_path, "w") as f:
+        json.dump(summary, f, indent=2)
+    logger.info("Results saved to %s", results_path)
+
+    return summary
+
+
+def _evaluate_game(
+    agent: WayfinderAgent,
+    game_id: str,
+    max_actions: int,
+) -> dict:
+    """Evaluate the agent on a single game.
+
+    In production, this uses the SDK's local execution mode. For
+    scaffolding/testing, it simulates with random frames.
+
+    Args:
+        agent: The agent to evaluate.
+        game_id: Game identifier.
+        max_actions: Max actions per level.
+
+    Returns:
+        Dict with score, actions, levels, etc.
+    """
+    agent.reset()
+
+    metrics = EvalMetrics()
+    levels_attempted = 0
+    levels_won = 0
+    total_score = 0.0
+    total_actions = 0
+
+    try:
+        # Try to use the SDK's local execution
+        from arc_agi_3 import LocalEnvironment  # type: ignore[import]
+
+        env = LocalEnvironment(game_id=game_id)
+        frames, state, score, win_score, available = env.reset()
+
+        while state == "NOT_FINISHED":
+            result = agent.act(
+                frames=frames,
+                state=state,
+                score=score,
+                win_score=win_score,
+                available_actions=available,
+            )
+            frames, state, score, win_score, available = env.step(result)
+            metrics.record_step(score, result["action"])
+            total_actions += 1
+
+            if agent.is_done(frames, state):
+                break
+
+        levels_attempted = 1
+        levels_won = 1 if state == "WIN" else 0
+        total_score = score
+
+    except ImportError:
+        # SDK not available — simulate
+        logger.warning("SDK not available — running simulation for %s", game_id)
+
+        for level in range(3):  # Simulate 3 levels
+            agent.reset()
+            levels_attempted += 1
+            level_score = 0.0
+
+            for step in range(max_actions):
+                frame = np.random.randint(0, 16, size=(64, 64), dtype=np.uint8)
+                state = "NOT_FINISHED"
+
+                result = agent.act(
+                    frames=[frame],
+                    state=state,
+                    score=level_score,
+                    win_score=1.0,
+                    available_actions=["ACTION1", "ACTION2", "ACTION3", "ACTION4", "ACTION5"],
+                )
+                metrics.record_step(level_score, result["action"])
+                total_actions += 1
+
+                # Random chance of "winning" for simulation
+                if np.random.random() < 0.01:
+                    levels_won += 1
+                    level_score = 1.0
+                    break
+
+                if agent.is_done([frame], state):
+                    break
+
+            total_score += level_score
+
+    stats = metrics.compute()
+    stats["efficiency"] = compute_efficiency(stats, max_actions)
+
+    return {
+        "score": total_score,
+        "total_actions": total_actions,
+        "levels_attempted": levels_attempted,
+        "levels_won": levels_won,
+        "metrics": stats,
+    }
+
+
+def main() -> int:
+    """CLI entry point for evaluation."""
+    parser = argparse.ArgumentParser(description="Run local evaluation")
+    parser.add_argument("--agent", default="wayfinder")
+    parser.add_argument("--games", default="ls20,ls21,ls22", help="Comma-separated game IDs")
+    parser.add_argument("--max-actions", type=int, default=1000)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--output-dir", default="eval/results")
+    parser.add_argument("-v", "--verbose", action="store_true")
+
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.DEBUG if args.verbose else logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    )
+
+    results = run_local_eval(
+        agent_name=args.agent,
+        games=args.games.split(","),
+        max_actions_per_level=args.max_actions,
+        device=args.device,
+        output_dir=Path(args.output_dir),
+    )
+
+    # Print summary
+    agg = results["aggregate"]
+    print("\n" + "=" * 60)
+    print(f"Agent: {results['agent']}")
+    print(f"Games: {len(results['games'])}")
+    print(f"Levels: {agg['total_levels']} (won: {agg['levels_won']})")
+    print(f"Win rate: {agg['win_rate']:.1%}")
+    print(f"Total score: {agg['total_score']:.2f}")
+    print(f"Total actions: {agg['total_actions']}")
+    print(f"Avg actions/level: {agg['avg_actions_per_level']:.1f}")
+    print("=" * 60)
+
+    return 0
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(main())

main.py

@@ -0,0 +1,10 @@
+"""Entry point for running the Wayfinder agent via the SDK.
+
+Usage:
+    uv run main.py --agent=wayfinder --game=ls20
+"""
+
+from agents.wayfinder.cli import main
+
+if __name__ == "__main__":
+    main()

notebooks/kaggle_submission.ipynb

@@ -0,0 +1,72 @@
+"""Kaggle submission notebook (rendered as Python for scaffolding).
+
+In the actual Kaggle environment, this is a Jupyter notebook. Here we
+provide the cell-by-cell Python equivalent that can be pasted into
+a notebook.
+
+Cell 1: Setup and installation
+Cell 2: Import agent
+Cell 3: Run agent against all games
+Cell 4: Submit results
+"""
+
+# === Cell 1: Setup ===
+# Install dependencies (internet is disabled during scoring, but available
+# during notebook setup before committing)
+# !pip install -q arc-agi-3 torch numpy pyyaml
+
+import os
+import sys
+import logging
+import numpy as np
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("kaggle_submission")
+
+DEVICE = "cuda" if os.environ.get("KAGGLE_GPU_TYPE") else "cpu"
+logger.info("Device: %s", DEVICE)
+
+
+# === Cell 2: Import agent ===
+# The agent code is included in the notebook's environment (uploaded as
+# a dataset or pip package).
+sys.path.insert(0, "/kaggle/input/wayfinder-agent")
+
+from agents.wayfinder.agent import WayfinderAgent
+
+agent = WayfinderAgent(
+    max_actions=1000,
+    latent_dim=256,
+    buffer_size=200_000,
+    device=DEVICE,
+)
+logger.info("Agent initialized")
+
+
+# === Cell 3: Run against all games ===
+# The SDK provides the game list and harness.
+try:
+    from arc_agi_3 import run_agent  # type: ignore[import]
+
+    results = run_agent(
+        agent=agent,
+        games="all",  # or a specific list
+        local_mode=True,  # No internet needed
+    )
+    logger.info("Evaluation complete: %s", results)
+except ImportError:
+    logger.error("arc-agi-3 SDK not available. Ensure it's installed.")
+
+
+# === Cell 4: Output results ===
+# The Kaggle harness expects results in a specific format.
+# This cell ensures the output is written correctly.
+import json
+
+output_path = "/kaggle/working/results.json"
+try:
+    with open(output_path, "w") as f:
+        json.dump(results, f, indent=2)
+    logger.info("Results written to %s", output_path)
+except NameError:
+    logger.warning("No results to write (SDK not available)")

pyproject.toml

@@ -0,0 +1,89 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "arc-agi-3-wayfinder"
+version = "0.1.0"
+description = "ARC Prize 2026 — ARC-AGI-3 track competition agent"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.11"
+authors = [{ name = "Wayfinder Team" }]
+keywords = ["arc-agi", "reinforcement-learning", "game-ai"]
+
+dependencies = [
+    "arc-agi-3>=0.9.0",
+    "torch>=2.2.0",
+    "numpy>=1.26.0",
+    "pyyaml>=6.0",
+    "rich>=13.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=5.0.0",
+    "ruff>=0.4.0",
+    "mypy>=1.10.0",
+]
+training = [
+    "wandb>=0.17.0",
+    "tensorboard>=2.16.0",
+]
+
+[project.scripts]
+wayfinder = "agents.wayfinder.cli:main"
+
+[tool.hatch.build.targets.wheel]
+packages = ["agents"]
+
+[tool.ruff]
+target-version = "py311"
+line-length = 100
+src = ["agents", "training", "eval", "tests"]
+
+[tool.ruff.lint]
+select = [
+    "E",   # pycodestyle errors
+    "W",   # pycodestyle warnings
+    "F",   # pyflakes
+    "I",   # isort
+    "B",   # flake8-bugbear
+    "C4",  # flake8-comprehensions
+    "UP",  # pyupgrade
+    "SIM", # flake8-simplify
+    "TCH", # flake8-type-checking
+]
+ignore = [
+    "E501",  # line too long — handled by formatter
+    "B008",  # function call in default argument
+]
+
+[tool.ruff.lint.isort]
+known-first-party = ["agents", "training", "eval"]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+
+[tool.pytest.ini_options]
+minversion = "8.0"
+testpaths = ["tests"]
+addopts = "-v --tb=short --strict-markers"
+markers = [
+    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+    "integration: marks tests requiring the SDK environment",
+]
+
+[tool.mypy]
+python_version = "3.11"
+strict = true
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true
+
+[[tool.mypy.overrides]]
+module = "arc_agi_3.*"
+ignore_missing_imports = true

tests/__init__.py

@@ -0,0 +1 @@
+"""Test package for Wayfinder agent."""

tests/test_action_head.py

@@ -0,0 +1,77 @@
+"""Tests for the action head module."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from agents.wayfinder.action_head import ActionHead
+from agents.wayfinder.world_model import WorldModel
+
+
+class TestActionHead:
+    """Test cases for ActionHead."""
+
+    @pytest.fixture
+    def action_head(self) -> ActionHead:
+        """Create a test action head."""
+        return ActionHead(latent_dim=32, device="cpu")
+
+    @pytest.fixture
+    def world_model(self) -> WorldModel:
+        """Create a test world model."""
+        return WorldModel(latent_dim=32, device="cpu")
+
+    @pytest.fixture
+    def latent(self) -> np.ndarray:
+        """Create a sample latent."""
+        return np.random.randn(32).astype(np.float32)
+
+    def test_select_returns_valid_action(
+        self, action_head: ActionHead, world_model: WorldModel, latent: np.ndarray
+    ) -> None:
+        """Test that select returns a valid action."""
+        result = action_head.select(
+            latent=latent,
+            diff_mask=np.zeros((64, 64), dtype=bool),
+            available_actions=["ACTION1", "ACTION2", "ACTION3"],
+            world_model=world_model,
+            epsilon=0.0,
+        )
+        assert "action" in result
+        assert result["action"] in ["ACTION1", "ACTION2", "ACTION3"]
+
+    def test_select_action6_includes_coordinates(
+        self, action_head: ActionHead, world_model: WorldModel, latent: np.ndarray
+    ) -> None:
+        """Test that ACTION6 includes x, y coordinates."""
+        result = action_head.select(
+            latent=latent,
+            diff_mask=np.zeros((64, 64), dtype=bool),
+            available_actions=["ACTION6"],
+            world_model=world_model,
+            epsilon=0.0,
+        )
+        assert result["action"] == "ACTION6"
+        assert "data" in result
+        assert "x" in result["data"]
+        assert "y" in result["data"]
+        assert 0 <= result["data"]["x"] < 64
+        assert 0 <= result["data"]["y"] < 64
+
+    def test_epsilon_greedy_can_explore(
+        self, action_head: ActionHead, world_model: WorldModel, latent: np.ndarray
+    ) -> None:
+        """Test that epsilon > 0 allows random exploration."""
+        actions_taken: set[str] = set()
+        for _ in range(50):
+            result = action_head.select(
+                latent=latent,
+                diff_mask=np.zeros((64, 64), dtype=bool),
+                available_actions=["ACTION1", "ACTION2", "ACTION3"],
+                world_model=world_model,
+                epsilon=1.0,  # Always explore
+            )
+            actions_taken.add(result["action"])
+        # Should have tried multiple different actions
+        assert len(actions_taken) > 1

tests/test_agent.py

@@ -0,0 +1,104 @@
+"""Integration tests for the WayfinderAgent."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from agents.wayfinder.agent import WayfinderAgent
+
+
+class TestWayfinderAgent:
+    """Integration test cases for WayfinderAgent."""
+
+    @pytest.fixture
+    def agent(self) -> WayfinderAgent:
+        """Create a test agent with small dimensions."""
+        return WayfinderAgent(
+            max_actions=100,
+            latent_dim=32,
+            buffer_size=1000,
+            device="cpu",
+        )
+
+    @pytest.fixture
+    def frame(self) -> np.ndarray:
+        """Create a sample frame."""
+        return np.random.randint(0, 16, size=(64, 64), dtype=np.uint8)
+
+    def test_act_returns_valid_action(self, agent: WayfinderAgent, frame: np.ndarray) -> None:
+        """Test that act returns a valid action dict."""
+        result = agent.act(
+            frames=[frame],
+            state="NOT_FINISHED",
+            score=0.0,
+            win_score=1.0,
+            available_actions=["ACTION1", "ACTION2", "ACTION3", "ACTION4", "ACTION5"],
+        )
+        assert "action" in result
+        assert "reasoning" in result
+        assert result["action"] in ["ACTION1", "ACTION2", "ACTION3", "ACTION4", "ACTION5"]
+
+    def test_is_done_on_win(self, agent: WayfinderAgent) -> None:
+        """Test that is_done returns True on WIN state."""
+        assert agent.is_done([], "WIN")
+
+    def test_is_done_on_game_over(self, agent: WayfinderAgent) -> None:
+        """Test that is_done returns True on GAME_OVER state."""
+        assert agent.is_done([], "GAME_OVER")
+
+    def test_is_done_on_budget_exhausted(self, agent: WayfinderAgent) -> None:
+        """Test that is_done returns True when action budget is exhausted."""
+        agent.action_count = agent.max_actions
+        assert agent.is_done([], "NOT_FINISHED")
+
+    def test_is_done_false_during_play(self, agent: WayfinderAgent) -> None:
+        """Test that is_done returns False during normal play."""
+        agent.action_count = 5
+        assert not agent.is_done([], "NOT_FINISHED")
+
+    def test_reset_clears_state(self, agent: WayfinderAgent, frame: np.ndarray) -> None:
+        """Test that reset clears per-level state."""
+        agent.act(
+            frames=[frame],
+            state="NOT_FINISHED",
+            score=0.0,
+            win_score=1.0,
+            available_actions=["ACTION1"],
+        )
+        assert agent.action_count > 0
+
+        agent.reset()
+        assert agent.action_count == 0
+        assert agent._prev_latent is None
+
+    def test_multiple_steps_accumulate_transitions(
+        self, agent: WayfinderAgent, frame: np.ndarray
+    ) -> None:
+        """Test that multiple steps build up the world model buffer."""
+        for i in range(5):
+            f = np.random.randint(0, 16, size=(64, 64), dtype=np.uint8)
+            agent.act(
+                frames=[f],
+                state="NOT_FINISHED",
+                score=0.0,
+                win_score=1.0,
+                available_actions=["ACTION1", "ACTION2", "ACTION3"],
+            )
+        # After 5 steps, at least 3 transitions should be in the buffer
+        assert agent._world_model.buffer_size_current >= 3
+
+    def test_reasoning_blob_has_required_fields(self, agent: WayfinderAgent, frame: np.ndarray) -> None:
+        """Test that the reasoning blob contains required audit fields."""
+        result = agent.act(
+            frames=[frame],
+            state="NOT_FINISHED",
+            score=0.0,
+            win_score=1.0,
+            available_actions=["ACTION1"],
+        )
+        reasoning = result["reasoning"]
+        assert "step" in reasoning
+        assert "score" in reasoning
+        assert "model_confidence" in reasoning
+        assert "novelty" in reasoning

tests/test_intrinsic_reward.py

@@ -0,0 +1,63 @@
+"""Tests for the intrinsic reward module."""
+
+from __future__ import annotations
+
+import pytest
+
+from agents.wayfinder.intrinsic_reward import IntrinsicReward, RewardConfig
+
+
+class TestIntrinsicReward:
+    """Test cases for IntrinsicReward."""
+
+    @pytest.fixture
+    def reward(self) -> IntrinsicReward:
+        """Create a test intrinsic reward module."""
+        return IntrinsicReward(RewardConfig(w_extrinsic=1.0, w_novelty=0.3, w_curiosity=0.2))
+
+    def test_compute_returns_float(self, reward: IntrinsicReward) -> None:
+        """Test that compute returns a float."""
+        result = reward.compute(extrinsic_delta=1.0, novelty=0.5, prediction_error=2.0)
+        assert isinstance(result, float)
+
+    def test_higher_extrinsic_gives_higher_utility(self, reward: IntrinsicReward) -> None:
+        """Test that higher extrinsic delta increases utility."""
+        low = reward.compute(extrinsic_delta=0.0, novelty=0.5, prediction_error=1.0)
+        high = reward.compute(extrinsic_delta=1.0, novelty=0.5, prediction_error=1.0)
+        assert high > low
+
+    def test_higher_novelty_gives_higher_utility(self, reward: IntrinsicReward) -> None:
+        """Test that higher novelty increases utility."""
+        low = reward.compute(extrinsic_delta=0.5, novelty=0.0, prediction_error=1.0)
+        high = reward.compute(extrinsic_delta=0.5, novelty=1.0, prediction_error=1.0)
+        assert high > low
+
+    def test_curiosity_clip(self) -> None:
+        """Test that curiosity is clipped."""
+        reward = IntrinsicReward(RewardConfig(w_curiosity=1.0, curiosity_clip=5.0))
+        clipped = reward.compute(extrinsic_delta=0.0, novelty=0.0, prediction_error=100.0)
+        # Curiosity is clipped to 5.0, and w_curiosity=1.0, so utility should be 5.0
+        assert abs(clipped - 5.0) < 0.1
+
+    def test_update_weights(self, reward: IntrinsicReward) -> None:
+        """Test that weights can be updated."""
+        reward.update_weights(w_extrinsic=2.0)
+        assert reward.config.w_extrinsic == 2.0
+
+    def test_reset(self, reward: IntrinsicReward) -> None:
+        """Test that reset clears running baseline."""
+        reward.compute(extrinsic_delta=10.0)
+        reward.reset()
+        assert reward._running_baseline == 0.0
+
+    def test_config_from_yaml(self, tmp_path) -> None:
+        """Test loading config from YAML."""
+        import yaml
+        config_data = {"intrinsic_reward": {"w_extrinsic": 2.0, "w_novelty": 0.5}}
+        config_path = tmp_path / "config.yaml"
+        with open(config_path, "w") as f:
+            yaml.dump(config_data, f)
+
+        config = RewardConfig.from_yaml(config_path)
+        assert config.w_extrinsic == 2.0
+        assert config.w_novelty == 0.5

tests/test_memory_graph.py

@@ -0,0 +1,102 @@
+"""Tests for the memory graph module."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from agents.wayfinder.memory_graph import MemoryGraph
+from agents.wayfinder.perception import GRID_SIZE
+
+
+class TestMemoryGraph:
+    """Test cases for MemoryGraph."""
+
+    @pytest.fixture
+    def graph(self) -> MemoryGraph:
+        """Create a test memory graph."""
+        return MemoryGraph()
+
+    @pytest.fixture
+    def sample_frame(self) -> np.ndarray:
+        """Create a sample frame."""
+        return np.random.randint(0, 16, size=(GRID_SIZE, GRID_SIZE), dtype=np.uint8)
+
+    def test_hash_frame_stable(self, graph: MemoryGraph, sample_frame: np.ndarray) -> None:
+        """Test that hashing is deterministic."""
+        h1 = graph.hash_frame(sample_frame)
+        h2 = graph.hash_frame(sample_frame)
+        assert h1 == h2
+
+    def test_hash_frame_different_for_different_frames(self, graph: MemoryGraph) -> None:
+        """Test that different frames have different hashes."""
+        f1 = np.zeros((GRID_SIZE, GRID_SIZE), dtype=np.uint8)
+        f2 = np.ones((GRID_SIZE, GRID_SIZE), dtype=np.uint8)
+        assert graph.hash_frame(f1) != graph.hash_frame(f2)
+
+    def test_add_node_new(self, graph: MemoryGraph, sample_frame: np.ndarray) -> None:
+        """Test adding a new node."""
+        h = graph.hash_frame(sample_frame)
+        latent = np.random.randn(64).astype(np.float32)
+        graph.add_node(h, latent, score=0.5)
+        assert h in graph.nodes
+        assert graph.nodes[h]["visit_count"] == 1
+        assert graph.nodes[h]["score"] == 0.5
+
+    def test_add_node_existing_increments_visits(self, graph: MemoryGraph, sample_frame: np.ndarray) -> None:
+        """Test that re-adding a node increments its visit count."""
+        h = graph.hash_frame(sample_frame)
+        latent = np.random.randn(64).astype(np.float32)
+        graph.add_node(h, latent)
+        graph.add_node(h, latent)
+        graph.add_node(h, latent)
+        assert graph.nodes[h]["visit_count"] == 3
+
+    def test_novelty_new_node(self, graph: MemoryGraph) -> None:
+        """Test novelty is 1.0 for unseen nodes."""
+        assert graph.novelty("nonexistent") == 1.0
+
+    def test_novelty_decreases_with_visits(self, graph: MemoryGraph, sample_frame: np.ndarray) -> None:
+        """Test that novelty decreases as visit count increases."""
+        h = graph.hash_frame(sample_frame)
+        latent = np.random.randn(64).astype(np.float32)
+        graph.add_node(h, latent)
+        n1 = graph.novelty(h)
+        graph.add_node(h, latent)
+        n2 = graph.novelty(h)
+        assert n2 < n1
+        assert 0 < n2 < 1
+
+    def test_add_edge(self, graph: MemoryGraph, sample_frame: np.ndarray) -> None:
+        """Test adding an edge between nodes."""
+        h1 = graph.hash_frame(sample_frame)
+        h2 = graph.hash_frame(sample_frame + 1)
+        latent = np.random.randn(64).astype(np.float32)
+        graph.add_node(h1, latent)
+        graph.add_node(h2, latent)
+        graph.add_edge(h1, h2, "ACTION1")
+
+        assert (h1, h2) in graph.edges
+        assert graph.edges[(h1, h2)]["action"] == "ACTION1"
+        assert len(graph.adjacency[h1]) == 1
+
+    def test_reset_clears_graph(self, graph: MemoryGraph, sample_frame: np.ndarray) -> None:
+        """Test that reset clears all nodes and edges."""
+        h = graph.hash_frame(sample_frame)
+        latent = np.random.randn(64).astype(np.float32)
+        graph.add_node(h, latent)
+        graph.add_edge(h, h, "ACTION1")
+
+        graph.reset()
+        assert len(graph.nodes) == 0
+        assert len(graph.edges) == 0
+
+    def test_stats(self, graph: MemoryGraph, sample_frame: np.ndarray) -> None:
+        """Test stats returns correct summary."""
+        h = graph.hash_frame(sample_frame)
+        latent = np.random.randn(64).astype(np.float32)
+        graph.add_node(h, latent)
+        graph.add_node(h, latent)
+        stats = graph.stats()
+        assert stats["node_count"] == 1
+        assert stats["avg_visits"] == 2.0

tests/test_perception.py

@@ -0,0 +1,54 @@
+"""Tests for the perception encoder module."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from agents.wayfinder.perception import PerceptionEncoder, GRID_SIZE, NUM_COLORS
+
+
+class TestPerceptionEncoder:
+    """Test cases for PerceptionEncoder."""
+
+    @pytest.fixture
+    def encoder(self) -> PerceptionEncoder:
+        """Create a test encoder instance."""
+        return PerceptionEncoder(latent_dim=64, device="cpu")
+
+    @pytest.fixture
+    def sample_frame(self) -> np.ndarray:
+        """Create a sample 64×64 frame."""
+        return np.random.randint(0, NUM_COLORS, size=(GRID_SIZE, GRID_SIZE), dtype=np.uint8)
+
+    def test_encode_returns_correct_shapes(self, encoder: PerceptionEncoder, sample_frame: np.ndarray) -> None:
+        """Test that encode returns latent and diff_mask with correct shapes."""
+        latent, diff_mask = encoder.encode(sample_frame)
+        assert latent.shape == (64,)
+        assert diff_mask.shape == (GRID_SIZE, GRID_SIZE)
+        assert diff_mask.dtype == bool
+
+    def test_first_frame_diff_is_zero(self, encoder: PerceptionEncoder, sample_frame: np.ndarray) -> None:
+        """Test that the first frame has an all-zero diff mask."""
+        _, diff_mask = encoder.encode(sample_frame)
+        assert not np.any(diff_mask)
+
+    def test_second_frame_diff_detects_changes(self, encoder: PerceptionEncoder, sample_frame: np.ndarray) -> None:
+        """Test that diff mask detects changes between frames."""
+        encoder.encode(sample_frame)
+        modified = sample_frame.copy()
+        modified[0, 0] = (modified[0, 0] + 1) % NUM_COLORS
+        _, diff_mask = encoder.encode(modified)
+        assert diff_mask[0, 0]
+        assert not diff_mask[1, 1]  # Unchanged pixel
+
+    def test_encode_batch(self, encoder: PerceptionEncoder) -> None:
+        """Test batch encoding."""
+        frames = np.random.randint(0, NUM_COLORS, size=(4, GRID_SIZE, GRID_SIZE), dtype=np.uint8)
+        latents = encoder.encode_batch(frames)
+        assert latents.shape == (4, 64)
+
+    def test_latent_is_float32(self, encoder: PerceptionEncoder, sample_frame: np.ndarray) -> None:
+        """Test that the latent is float32."""
+        latent, _ = encoder.encode(sample_frame)
+        assert latent.dtype == np.float32

tests/test_planner.py

@@ -0,0 +1,91 @@
+"""Tests for the planner module."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from agents.wayfinder.planner import Planner, PlannerNode
+from agents.wayfinder.world_model import WorldModel
+from agents.wayfinder.intrinsic_reward import IntrinsicReward
+
+
+class TestPlannerNode:
+    """Test cases for PlannerNode."""
+
+    def test_ucb1_infinite_for_unvisited(self) -> None:
+        """Test that UCB1 is infinite for unvisited nodes."""
+        node = PlannerNode(latent=np.zeros(32))
+        assert node.ucb1() == float("inf")
+
+    def test_ucb1_finite_after_visit(self) -> None:
+        """Test that UCB1 is finite after a visit."""
+        parent = PlannerNode(latent=np.zeros(32))
+        parent.visits = 10
+        node = PlannerNode(latent=np.zeros(32), parent=parent)
+        node.visits = 1
+        node.total_value = 0.5
+        assert node.ucb1() < float("inf")
+
+    def test_mean_value(self) -> None:
+        """Test mean value computation."""
+        node = PlannerNode(latent=np.zeros(32))
+        node.visits = 4
+        node.total_value = 2.0
+        assert node.mean_value == 0.5
+
+    def test_mean_value_zero_visits(self) -> None:
+        """Test mean value is 0 for unvisited nodes."""
+        node = PlannerNode(latent=np.zeros(32))
+        assert node.mean_value == 0.0
+
+
+class TestPlanner:
+    """Test cases for Planner."""
+
+    @pytest.fixture
+    def planner(self) -> Planner:
+        """Create a test planner."""
+        world_model = WorldModel(latent_dim=32, device="cpu")
+        reward_module = IntrinsicReward()
+        return Planner(
+            world_model=world_model,
+            reward_module=reward_module,
+            max_depth=3,
+            max_simulations=10,
+        )
+
+    def test_plan_returns_action_dict(self, planner: Planner) -> None:
+        """Test that plan returns a valid action dict."""
+        latent = np.random.randn(32).astype(np.float32)
+        result = planner.plan(
+            latent=latent,
+            available_actions=["ACTION1", "ACTION2", "ACTION3"],
+            action_budget_remaining=100,
+            utility_fn=lambda **kw: 0.0,
+        )
+        assert "action" in result
+        assert result["action"] in ["ACTION1", "ACTION2", "ACTION3"]
+
+    def test_plan_fallback_on_no_actions(self, planner: Planner) -> None:
+        """Test that plan falls back when no valid actions."""
+        latent = np.random.randn(32).astype(np.float32)
+        result = planner.plan(
+            latent=latent,
+            available_actions=["ACTION1"],
+            action_budget_remaining=100,
+            utility_fn=lambda **kw: 0.0,
+        )
+        assert "action" in result
+
+    def test_plan_scales_with_budget(self, planner: Planner) -> None:
+        """Test that planning scales with remaining budget."""
+        latent = np.random.randn(32).astype(np.float32)
+        # Should not crash with very low budget
+        result = planner.plan(
+            latent=latent,
+            available_actions=["ACTION1", "ACTION2"],
+            action_budget_remaining=1,
+            utility_fn=lambda **kw: 0.0,
+        )
+        assert "action" in result

tests/test_world_model.py

@@ -0,0 +1,76 @@
+"""Tests for the world model module."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from agents.wayfinder.world_model import WorldModel, Transition
+
+
+class TestWorldModel:
+    """Test cases for WorldModel."""
+
+    @pytest.fixture
+    def model(self) -> WorldModel:
+        """Create a test world model."""
+        return WorldModel(latent_dim=32, buffer_size=1000, device="cpu")
+
+    @pytest.fixture
+    def sample_latent(self) -> np.ndarray:
+        """Create a sample latent vector."""
+        return np.random.randn(32).astype(np.float32)
+
+    def test_predict_change_returns_probability(self, model: WorldModel, sample_latent: np.ndarray) -> None:
+        """Test that predict_change returns a valid probability."""
+        prob = model.predict_change(sample_latent, {"action": "ACTION1"})
+        assert 0.0 <= prob <= 1.0
+
+    def test_predict_next_latent_shape(self, model: WorldModel, sample_latent: np.ndarray) -> None:
+        """Test that predict_next_latent returns correct shape."""
+        next_latent = model.predict_next_latent(sample_latent, {"action": "ACTION1"})
+        assert next_latent.shape == (32,)
+
+    def test_add_transition_increases_buffer(self, model: WorldModel, sample_latent: np.ndarray) -> None:
+        """Test that add_transition adds to the buffer."""
+        assert model.buffer_size_current == 0
+        model.add_transition(
+            state_latent=sample_latent,
+            action={"action": "ACTION1"},
+            next_latent=sample_latent + 0.1,
+            frame_changed=True,
+        )
+        assert model.buffer_size_current == 1
+
+    def test_add_transition_deduplicates(self, model: WorldModel, sample_latent: np.ndarray) -> None:
+        """Test that identical transitions are deduplicated."""
+        for _ in range(3):
+            model.add_transition(
+                state_latent=sample_latent,
+                action={"action": "ACTION1"},
+                next_latent=sample_latent + 0.1,
+                frame_changed=True,
+            )
+        assert model.buffer_size_current == 1
+
+    def test_train_step_returns_loss(self, model: WorldModel, sample_latent: np.ndarray) -> None:
+        """Test that train_step returns a loss value after enough data."""
+        # Add enough transitions for a batch
+        for i in range(70):
+            model.add_transition(
+                state_latent=sample_latent + i * 0.01,
+                action={"action": f"ACTION{i % 4 + 1}"},
+                next_latent=sample_latent + (i + 1) * 0.01,
+                frame_changed=(i % 2 == 0),
+            )
+        loss = model.train_step(batch_size=32)
+        assert loss >= 0.0
+
+    def test_confidence_starts_low(self, model: WorldModel) -> None:
+        """Test that confidence is 0 before training."""
+        assert model.confidence() == 0.0
+
+    def test_prediction_error_zero_on_first_step(self, model: WorldModel) -> None:
+        """Test that prediction error is 0 on the first step."""
+        latent = np.random.randn(32).astype(np.float32)
+        assert model.prediction_error(None, None, latent) == 0.0

training/configs/default.yaml

@@ -0,0 +1,54 @@
+# Default configuration for Wayfinder agent training
+
+agent:
+  max_actions: 1000
+  latent_dim: 256
+  buffer_size: 200000
+  device: cpu  # or cuda
+
+perception:
+  latent_dim: 256
+  conv_channels: [32, 64, 128]
+  conv_kernel: 3
+
+world_model:
+  latent_dim: 256
+  action_embed_dim: 32
+  hidden_dim: 256
+  lr: 0.0001
+  buffer_size: 200000
+  batch_size: 64
+  entropy_reg: 0.01
+  fwd_loss_weight: 0.5
+
+intrinsic_reward:
+  w_extrinsic: 1.0
+  w_novelty: 0.3
+  w_curiosity: 0.2
+  novelty_decay: 0.95
+  curiosity_clip: 10.0
+  score_baseline: 0.0
+
+planner:
+  max_depth: 5
+  max_simulations: 50
+  exploration_constant: 1.414
+  discount: 0.95
+
+action_head:
+  temperature: 0.5
+  epsilon: 0.15
+  coord_head_channels: [32, 16]
+
+training:
+  mode: online  # or offline
+  epochs: 50
+  batch_size: 64
+  lr: 0.0001
+  eval_interval: 10
+  save_interval: 20
+
+eval:
+  games: ["ls20", "ls21", "ls22"]
+  max_actions_per_level: 1000
+  output_dir: "eval/results"

training/replay_buffer.py

@@ -0,0 +1,199 @@
+"""Deduplicated replay buffer for offline + online training.
+
+Stores transitions with frame-level deduplication (hash-based) to
+avoid wasting training capacity on near-identical states. Supports
+random sampling, prioritized sampling, and buffer persistence.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Iterator
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BufferedTransition:
+    """A transition stored in the replay buffer.
+
+    Attributes:
+        frame: 64×64 uint8 frame before the action.
+        action: Action name string.
+        action_data: Optional action data (e.g. coordinates).
+        next_frame: 64×64 uint8 frame after the action.
+        reward: Extrinsic reward (score delta).
+        frame_changed: Whether the frame visually changed.
+        frame_hash: MD5 hash of the frame (for dedup).
+    """
+
+    frame: np.ndarray
+    action: str
+    action_data: dict | None
+    next_frame: np.ndarray
+    reward: float
+    frame_changed: bool
+    frame_hash: str = ""
+
+    def __post_init__(self) -> None:
+        if not self.frame_hash:
+            self.frame_hash = hashlib.md5(self.frame.tobytes()).hexdigest()
+
+
+class ReplayBuffer:
+    """Deduplicated replay buffer for transition storage.
+
+    Features:
+    - Frame-level deduplication (stores unique frames only once).
+    - Transition-level deduplication (same state+action → skip).
+    - Random and prioritized sampling.
+    - Configurable maximum size.
+
+    Attributes:
+        max_size: Maximum number of transitions.
+        frames: Dict mapping frame_hash → frame array (deduplicated storage).
+        transitions: Deque of BufferedTransition objects.
+    """
+
+    def __init__(self, max_size: int = 200_000) -> None:
+        """Initialize the replay buffer.
+
+        Args:
+            max_size: Maximum number of transitions to store.
+        """
+        self.max_size = max_size
+        self.frames: dict[str, np.ndarray] = {}
+        self.transitions: deque[BufferedTransition] = deque(maxlen=max_size)
+        self._seen_keys: set[str] = set()
+
+        logger.info("ReplayBuffer initialized (max_size=%d)", max_size)
+
+    def add(
+        self,
+        frame: np.ndarray,
+        action: str,
+        action_data: dict | None,
+        next_frame: np.ndarray,
+        reward: float,
+        frame_changed: bool,
+    ) -> bool:
+        """Add a transition to the buffer.
+
+        Args:
+            frame: Frame before action.
+            action: Action name.
+            action_data: Optional action data.
+            next_frame: Frame after action.
+            reward: Extrinsic reward.
+            frame_changed: Whether frame visually changed.
+
+        Returns:
+            True if the transition was added, False if deduplicated.
+        """
+        frame_hash = hashlib.md5(frame.tobytes()).hexdigest()
+        next_hash = hashlib.md5(next_frame.tobytes()).hexdigest()
+        dedup_key = f"{frame_hash}:{action}:{action_data}"
+
+        if dedup_key in self._seen_keys:
+            return False
+
+        self._seen_keys.add(dedup_key)
+
+        # Store unique frames
+        if frame_hash not in self.frames:
+            self.frames[frame_hash] = frame.copy()
+        if next_hash not in self.frames:
+            self.frames[next_hash] = next_frame.copy()
+
+        self.transitions.append(
+            BufferedTransition(
+                frame=frame.copy(),
+                action=action,
+                action_data=action_data,
+                next_frame=next_frame.copy(),
+                reward=reward,
+                frame_changed=frame_changed,
+                frame_hash=frame_hash,
+            )
+        )
+
+        return True
+
+    def sample(self, batch_size: int) -> list[BufferedTransition]:
+        """Sample a random batch of transitions.
+
+        Args:
+            batch_size: Number of transitions to sample.
+
+        Returns:
+            List of BufferedTransition objects.
+        """
+        if len(self.transitions) < batch_size:
+            return list(self.transitions)
+        indices = np.random.choice(len(self.transitions), size=batch_size, replace=False)
+        return [self.transitions[i] for i in indices]
+
+    def sample_prioritized(
+        self,
+        batch_size: int,
+        alpha: float = 0.6,
+    ) -> list[BufferedTransition]:
+        """Sample a batch with prioritization toward changed frames.
+
+        Prioritizes transitions where the frame changed (more informative
+        for training the world model).
+
+        Args:
+            batch_size: Number of transitions.
+            alpha: Prioritization exponent (0=uniform, 1=full priority).
+
+        Returns:
+            List of BufferedTransition objects.
+        """
+        if len(self.transitions) < batch_size:
+            return list(self.transitions)
+
+        priorities = np.array([
+            (1.0 if t.frame_changed else 0.1) ** alpha
+            for t in self.transitions
+        ])
+        probs = priorities / priorities.sum()
+        indices = np.random.choice(len(self.transitions), size=batch_size, p=probs, replace=False)
+        return [self.transitions[i] for i in indices]
+
+    def __len__(self) -> int:
+        """Return the number of transitions in the buffer."""
+        return len(self.transitions)
+
+    def __iter__(self) -> Iterator[BufferedTransition]:
+        """Iterate over all transitions."""
+        return iter(self.transitions)
+
+    @property
+    def num_unique_frames(self) -> int:
+        """Number of unique frames stored."""
+        return len(self.frames)
+
+    def stats(self) -> dict:
+        """Return buffer statistics.
+
+        Returns:
+            Dict with count, unique_frames, changed_ratio.
+        """
+        changed_count = sum(1 for t in self.transitions if t.frame_changed)
+        return {
+            "count": len(self.transitions),
+            "unique_frames": len(self.frames),
+            "changed_ratio": changed_count / max(len(self.transitions), 1),
+        }
+
+    def clear(self) -> None:
+        """Clear all stored data."""
+        self.frames.clear()
+        self.transitions.clear()
+        self._seen_keys.clear()

training/train_world_model.py

@@ -0,0 +1,211 @@
+"""Offline + online training loop for the world model.
+
+This script can run in two modes:
+1. Offline: Train on pre-collected transitions from a replay buffer file.
+2. Online: Play games and train the world model incrementally.
+
+Usage:
+    uv run python training/train_world_model.py --mode offline --buffer data/buffer.pkl
+    uv run python training/train_world_model.py --mode online --games ls20,ls21
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import pickle
+from pathlib import Path
+
+import numpy as np
+
+from agents.wayfinder.perception import PerceptionEncoder
+from agents.wayfinder.world_model import WorldModel
+from training.replay_buffer import ReplayBuffer
+
+logger = logging.getLogger(__name__)
+
+
+def train_offline(
+    buffer_path: str,
+    latent_dim: int = 256,
+    epochs: int = 50,
+    batch_size: int = 64,
+    lr: float = 1e-4,
+    device: str = "cpu",
+    save_path: str = "models/world_model.pt",
+) -> None:
+    """Train the world model offline on a pre-collected buffer.
+
+    Args:
+        buffer_path: Path to a pickled ReplayBuffer.
+        latent_dim: Latent dimension.
+        epochs: Number of training epochs.
+        batch_size: Training batch size.
+        lr: Learning rate.
+        device: Torch device.
+        save_path: Where to save the trained model.
+    """
+    logger.info("Loading replay buffer from %s", buffer_path)
+    with open(buffer_path, "rb") as f:
+        buffer: ReplayBuffer = pickle.load(f)
+
+    logger.info("Buffer loaded: %d transitions, %d unique frames", len(buffer), buffer.num_unique_frames)
+
+    encoder = PerceptionEncoder(latent_dim=latent_dim, device=device)
+    world_model = WorldModel(latent_dim=latent_dim, device=device, lr=lr)
+
+    # Pre-encode all unique frames
+    logger.info("Encoding unique frames...")
+    frame_hashes = list(buffer.frames.keys())
+    frame_arrays = np.stack([buffer.frames[h] for h in frame_hashes])
+    latents = encoder.encode_batch(frame_arrays)
+    hash_to_latent = {h: lat for h, lat in zip(frame_hashes, latents)}
+
+    logger.info("Training for %d epochs...", epochs)
+    for epoch in range(epochs):
+        batch = buffer.sample_prioritized(batch_size)
+        total_loss = 0.0
+
+        for t in batch:
+            state_latent = hash_to_latent.get(t.frame_hash)
+            next_hash = hash_to_latent.get(
+                __import__("hashlib").md5(t.next_frame.tobytes()).hexdigest()
+            )
+            if state_latent is None or next_hash is None:
+                continue
+
+            world_model.add_transition(
+                state_latent=state_latent,
+                action={"action": t.action, "data": t.action_data},
+                next_latent=next_hash,
+                frame_changed=t.frame_changed,
+            )
+            loss = world_model.train_step(batch_size=min(batch_size, len(world_model._buffer)))
+            total_loss += loss
+
+        avg_loss = total_loss / max(len(batch), 1)
+        if (epoch + 1) % 5 == 0:
+            logger.info(
+                "Epoch %d/%d: avg_loss=%.4f, buffer=%d, confidence=%.3f",
+                epoch + 1, epochs, avg_loss,
+                world_model.buffer_size_current,
+                world_model.confidence(),
+            )
+
+    # Save model
+    save_dir = Path(save_path).parent
+    save_dir.mkdir(parents=True, exist_ok=True)
+    import torch
+    torch.save({
+        "world_model": world_model.state_dict(),
+        "encoder": encoder.state_dict(),
+        "latent_dim": latent_dim,
+    }, save_path)
+    logger.info("Model saved to %s", save_path)
+
+
+def train_online(
+    games: list[str],
+    max_actions_per_game: int = 500,
+    latent_dim: int = 256,
+    device: str = "cpu",
+    save_path: str = "models/world_model_online.pt",
+) -> None:
+    """Train the world model online by playing games.
+
+    Args:
+        games: List of game IDs to play.
+        max_actions_per_game: Max actions per game.
+        latent_dim: Latent dimension.
+        device: Torch device.
+        save_path: Where to save the model.
+    """
+    from agents.wayfinder.agent import WayfinderAgent
+
+    agent = WayfinderAgent(
+        max_actions=max_actions_per_game,
+        latent_dim=latent_dim,
+        device=device,
+    )
+
+    for game_id in games:
+        logger.info("Playing game %s...", game_id)
+        agent.reset()
+
+        # In real usage, this would use the SDK to play the game.
+        # For now, we simulate with random frames.
+        for step in range(max_actions_per_game):
+            frame = np.random.randint(0, 16, size=(64, 64), dtype=np.uint8)
+            result = agent.act(
+                frames=[frame],
+                state="NOT_FINISHED",
+                score=0.0,
+                win_score=1.0,
+                available_actions=["ACTION1", "ACTION2", "ACTION3", "ACTION4", "ACTION5"],
+            )
+
+            if agent.is_done([frame], "NOT_FINISHED"):
+                break
+
+        logger.info(
+            "Game %s: %d actions, buffer=%d, confidence=%.3f",
+            game_id, agent.action_count,
+            agent._world_model.buffer_size_current,
+            agent._world_model.confidence(),
+        )
+
+    import torch
+    torch.save({
+        "world_model": agent._world_model.state_dict(),
+        "encoder": agent._encoder.state_dict(),
+        "latent_dim": latent_dim,
+    }, save_path)
+    logger.info("Online model saved to %s", save_path)
+
+
+def main() -> int:
+    """CLI entry point for training."""
+    parser = argparse.ArgumentParser(description="Train the world model")
+    parser.add_argument("--mode", choices=["offline", "online"], default="online")
+    parser.add_argument("--buffer", default="data/buffer.pkl", help="Path to replay buffer (offline mode)")
+    parser.add_argument("--games", default="ls20,ls21,ls22", help="Comma-separated game IDs (online mode)")
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--latent-dim", type=int, default=256)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--save-path", default="models/world_model.pt")
+    parser.add_argument("-v", "--verbose", action="store_true")
+
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.DEBUG if args.verbose else logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    )
+
+    if args.mode == "offline":
+        train_offline(
+            buffer_path=args.buffer,
+            latent_dim=args.latent_dim,
+            epochs=args.epochs,
+            batch_size=args.batch_size,
+            lr=args.lr,
+            device=args.device,
+            save_path=args.save_path,
+        )
+    else:
+        train_online(
+            games=args.games.split(","),
+            max_actions_per_game=500,
+            latent_dim=args.latent_dim,
+            device=args.device,
+            save_path=args.save_path,
+        )
+
+    return 0
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(main())