Publish full project code and checkpoint

.claude/settings.local.json

@@ -0,0 +1,8 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(hf auth *)",
+      "Bash(hf upload *)"
+    ]
+  }
+}

.gitignore

@@ -0,0 +1,29 @@
+# Python
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.Python
+*.egg-info/
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# Trained checkpoints -- regenerable via scripts/train.py, deliberately NOT
+# blanket-ignored: checkpoints/connectx_checkpoint.pt (the one actually
+# shipped/deployed) is explicitly kept, everything else here is not.
+checkpoints/*.pt
+!checkpoints/connectx_checkpoint.pt
+
+# Local run logs / scratch output
+*.log
+*.out
+
+# OS/editor cruft
+.DS_Store
+Thumbs.db
+.idea/
+.vscode/
+*.swp

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Alexandros Titonis
+
+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,211 @@
----
-license: mit
----
+# ConnectX-WorldModel
+
+A reinforcement-learning world model (encoder + latent dynamics + value head)
+combined with real adversarial search and an exact endgame solver, applied to
+Kaggle's **[ConnectX competition](https://kaggle.com/competitions/connectx)**
+(ranked Connect-4, 7-wide x 6-tall board, evergreen/Knowledge-only, no
+deadline).
+
+Instead of a hand-written Connect-4 bot, this trains a small model to predict
+what happens next in latent space, then searches over that prediction to pick
+a move — MuZero/AlphaZero-style planning [1][2], applied to a real,
+Kaggle-ranked adversarial game, and trained end to end on a single laptop (RTX
+4060 laptop GPU, Ryzen AI 9 HX 370, 32GB RAM — no cluster, no cloud run). Full
+write-up, architecture rationale, every bug found along the way, an attempted
+fix for a diagnosed zugzwang weakness, and honest limitations:
+**[`docs/WHITEPAPER.pdf`](docs/WHITEPAPER.pdf)**.
+
+This is also not a Connect-4-only architecture — it's one instance of a
+general latent-space-simulation pattern (the same encoder/dynamics/value/
+decoder + search code has been applied to other domains — equation solving,
+constraint-satisfaction logic puzzles, route planning — with zero core-code
+changes, just a new implementation of the small interface in
+`connectx/environment.py`). ConnectX is documented here because it's the
+first genuinely *adversarial* two-player domain this architecture was applied
+to — see the whitepaper's Section 1 for the full framing.
+
+## How it works, briefly
+
+1. **A learned world model** (`connectx/model.py`) — an encoder maps a board
+   to a latent vector, a dynamics model predicts the next latent + reward
+   given an action, a value head estimates cost-to-go, and a diagnostic
+   decoder checks the latent space isn't collapsing. Trained self-supervised
+   (`scripts/train.py`, `connectx/train_utils.py`) with **no oracle** — the
+   real 7x6 board's game tree is too large to brute-force, so the value head
+   trains via on-policy Monte Carlo returns (`connectx/verifier.py`) instead
+   of regression against exact labels.
+2. **Real adversarial search, not latent imagination**
+   (`connectx/adversarial_search.py`) — the board's rules are exactly known,
+   so rather than asking the trained model to *imagine* the opponent's reply,
+   the search enumerates the agent's real legal moves and the opponent's real
+   legal replies directly against the actual board simulator
+   (`connectx/env.py`), and uses the learned value head only as the leaf
+   evaluator. This is the single biggest lever: it alone took the agent from
+   struggling against even a fixed heuristic to reliably beating one it never
+   trained against.
+3. **An exact endgame solver**, folded into the same search — once a
+   position narrows to a handful of legal columns (which happens naturally
+   as a real game fills up), the remaining game tree is small regardless of
+   how many moves are left, and gets solved exactly instead of estimated.
+4. **Episodic memory + LoRA self-play fine-tuning**
+   (`connectx/episodic_memory.py`, `connectx/lora.py`,
+   `scripts/lora_selfplay_finetune.py`) layered on top, each confirmed to
+   help before being kept.
+5. **A first, honestly-reported attempt at a real diagnosed weakness**
+   (Connect-4 zugzwang/parity traps) — built, measured, and shipped as an
+   opt-in experimental parameter rather than oversold as solved. See the
+   whitepaper's Section 4.
+
+## Results
+
+Win rate against an independent, trusted test harness (`tests/submission_test.py`
+— an alternating-turn engine that does **not** reuse the environment's own
+bundled `step()`, deliberately avoiding the shortcut a real submission
+validator needs to avoid). Three opponents: random-legal play, the fixed
+weak heuristic trained against, and a stronger 1-ply-deeper heuristic never
+seen during training.
+
+| Configuration | vs random | vs weak heuristic | vs stronger heuristic |
+|---|---|---|---|
+| Latent beam search (original baseline) | 63.3% | 0.0% | 6.7% |
+| Real adversarial search (1 round) | 96.7% | 100.0% | 66.7% |
+| + episodic memory + online learner | 98.3% | 78.3%¹ | 50.0% |
+| + a real bug fixed (loss ≠ draw) | 100.0% | 100.0% | 80.0% |
+| + LoRA self-play fine-tune | — | — | 81.7% |
+| + curriculum self-play (best-ever) | — | — | **85.0%** |
+| + exact endgame solver (current, deployed) | 100.0% | 100.0% | 83.3% |
+
+¹ Explained, not a mystery — that harness ran all opponent blocks
+sequentially in one process, so the online learner had already drifted from
+60 preceding random-opponent games by the time it reached this block. See
+the whitepaper for the full table, the diagnosis, and every real bug found
+along the way (7 of them — a couple are worth knowing if you build on this).
+
+On Kaggle's own real rating: this is a TrueSkill-style score that starts
+uncertain and converges over dozens of real games — don't read a
+freshly-uploaded rating as a verdict. 7x6 Connect-4 is a mathematically
+**solved** game, so the real competitive pool likely includes near-perfect
+solvers; these results demonstrate the architecture works on this domain,
+not a leaderboard-rating prediction.
+
+## Running locally
+
+Requires Python 3.10+ and PyTorch (CPU is fine — nothing here needs a GPU;
+the reference results above were produced on a laptop GPU but nothing in the
+code requires one).
+
+```bash
+pip install torch
+
+# Self-test the environment (small board, exact BFS oracle) + a real-board smoke test
+python -m connectx.env
+
+# Verify the already-trained, shipped agent against the trusted harness
+python -m tests.submission_test
+
+# Train a fresh checkpoint from scratch (real 7x6 board, no oracle at this scale, ~1hr on a laptop)
+python -m scripts.train
+
+# LoRA self-play fine-tune an existing checkpoint (this is the step that produced the deployed one)
+python -m scripts.lora_selfplay_finetune
+
+# Package a checkpoint + offline self-play memory into a self-contained Kaggle submission.py
+python -c "from scripts.build_submission import main; main()"
+```
+
+Run everything from the repository root (not from inside `connectx/` or
+`scripts/`) — every command above is `python -m <package>.<module>`, matching
+the layout below.
+
+`submission.py` (repo root) is the actual file to upload to Kaggle as-is —
+it has zero dependencies beyond `torch`/`base64`/`io`/`time`, with the
+trained weights and episodic memory embedded directly in the file.
+
+## Training this for a different game
+
+Nothing here is Connect-4-specific beyond `connectx/env.py`. Every other
+file only depends on the small interface `connectx/environment.py` defines:
+
+- `state_dim`, `num_actions`, `always_legal_actions` (properties)
+- `is_solved(state)`, `is_legal(state, action_idx)`
+- `step(state, action_idx) -> (next_state, reward, done)`
+- `random_problem(rng) -> (state, answer)`
+- `bfs_solve(state, max_depth=8) -> path or None` (an exact oracle, if one
+  exists at a scale small enough to brute-force — return `None`
+  unconditionally if it doesn't, the way `env.py` does above `BFS_MAX_CELLS`)
+
+To point this at a new game or puzzle:
+
+1. Implement `Environment` for your domain (see `connectx/env.py` for a full
+   worked example, including the small-board-with-an-oracle / large-board-
+   without-one pattern).
+2. Run `connectx.train_utils.train_stage1` on your new environment to get a
+   working encoder/dynamics/decoder — no oracle needed for this stage in any
+   domain.
+3. **If your domain is single-agent** (a puzzle, not a two-player game):
+   the value head can train directly against `bfs_solve` labels if you have
+   an oracle, or via `connectx.verifier.train_mc_value_onpolicy` (drop
+   `unsolved_penalty`, since a single-agent domain never has a "loss," only
+   "unsolved") if you don't.
+4. **If your domain is adversarial** (two players, like this one): bake a
+   fixed opponent into your `step()` first (the "easy path" — see `env.py`'s
+   own module docstring) to get the pipeline working end to end, THEN set
+   `train_mc_value_onpolicy(..., unsolved_penalty=<something>)` — this is
+   the one setting that matters most for an adversarial domain and is
+   exactly what bug #2 in the whitepaper was about: without it, the value
+   head never sees a single example of "this leads to losing."
+5. If your domain's rules are exactly known (not something that needs to be
+   learned) and the state space is too large for `bfs_solve` to search
+   exhaustively at decision time, real search over the real environment
+   (`connectx/adversarial_search.py`'s pattern, or the plain single-agent
+   search in `connectx/search.py`) will almost always beat asking the
+   dynamics model to imagine ahead — that was this project's single biggest
+   result.
+
+## Layout
+
+```
+connectx-opensource/
+  connectx/                        The importable package -- everything domain-agnostic + ConnectX itself
+    environment.py                   The generic Environment interface everything else depends on
+    env.py                           ConnectX itself: board, rules, the fixed training opponent
+    model.py                         Encoder / dynamics / value head / decoder
+    train_utils.py                   Self-supervised stage-1 training (no oracle, no labels)
+    verifier.py                      On-policy Monte Carlo value-head training (no oracle)
+    search.py                        Latent-space search (comparison baseline) + load_checkpoint
+    adversarial_search.py            The REAL search actually deployed: minimax + exact endgame solver
+                                      + the experimental parity-heuristic attempt (Section 4.1)
+    episodic_memory.py               k-NN memory over real self-play trajectories (won + lost)
+    memory_build.py                  Builds that memory offline, at packaging time
+    lora.py                          Generic LoRA wrapper
+  scripts/                         Executable entry points (run as `python -m scripts.<name>`)
+    train.py                         Full training pipeline (stage 1 + stage 2 + self-play)
+    lora_selfplay_finetune.py        LoRA self-play fine-tune of an existing checkpoint
+    build_submission.py              Packages a checkpoint + memory into one self-contained submission.py
+  tests/
+    submission_test.py               The trusted, independent test harness
+  checkpoints/
+    connectx_checkpoint.pt           Trained weights
+  docs/
+    WHITEPAPER.pdf                   Full write-up: architecture, every bug, all results, limitations
+    build_whitepaper.py              Regenerates WHITEPAPER.pdf (requires `pip install fpdf2`)
+  submission.py                    THE deployed file — upload this to Kaggle as-is
+  LICENSE
+  .gitignore
+```
+
+## Citations
+
+Full reference list with page/venue detail is in `docs/WHITEPAPER.pdf`.
+Headline credits: latent-space planning follows MuZero [1] / AlphaZero [2];
+episodic memory follows Model-Free Episodic Control [3] and Neural Episodic
+Control [4]; LoRA fine-tuning follows Hu et al. [5]; the zugzwang/parity
+endgame theory referenced in the whitepaper's case study traces to Victor
+Allis's 1988 solution of Connect-4 [6]; the board/config schema and
+fixed-opponent-in-`step` convention are ported from Kaggle's own ConnectX
+competition [8] and the `kaggle_environments` package [9], not from any
+published agent's code.
+
+## License
+
+[MIT](LICENSE).

checkpoints/connectx_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2de769eb71c0c10e917f07fd609210502dc194304a0466630255d11200fbdb83
+size 1831357

connectx/adversarial_search.py

@@ -0,0 +1,387 @@
+"""
+Real (not latent-imagined) minimax search over the actual board, plus an
+exact endgame solver -- the search this project actually deploys, and the
+main reason it beats the latent-search baseline in `search.py` decisively.
+
+Why real board space, not latent imagination: `env.step()` bundles the
+agent's move and the fixed opponent's reply into ONE transition, so the
+trained dynamics model was only ever shown full round-trips as single
+training examples -- it structurally cannot represent "the board right
+after my move, before their reply" as a state, because it never saw that
+state shape. But Connect-4's rules are exactly known (there's a full plain-
+Python board simulator right here), so there's no need to make a neural
+network imagine something 40 lines of Python computes for free. This
+module does the adversarial ply in REAL board space (enumerate the agent's
+real legal moves; for each, enumerate the opponent's real legal replies and
+assume they pick whichever hurts the agent most -- a genuine minimax, not a
+guess) and uses the learned value head ONLY as the leaf evaluator, on a
+real, never-imagined state.
+
+The exact endgame solver (`_exact_endgame_solve`) goes one step further:
+once a position narrows down to a handful of legal columns -- which happens
+naturally as a real board fills up -- the remaining game tree is small
+regardless of how many plies are left, and can be solved exactly with no
+learned value head at all. Branching factor is controlled by how many
+columns are legal, not by how many cells are empty, which is what makes
+this cheap even fairly late in a real game.
+"""
+import time
+
+import torch
+
+from .env import EMPTY, AGENT, OPPONENT, _decode_board, _encode_board, _rc, _lowest_empty_row, _legal_columns, _wins_for, _board_full
+from .train_utils import states_to_tensor
+from .search import _evaluate_with_memory
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _apply_move(cells, col, mark, width, height):
+    row = _lowest_empty_row(cells, col, width, height)
+    new_cells = list(cells)
+    new_cells[_rc(row, col, width)] = mark
+    return new_cells
+
+
+class _EndgameTimeout(Exception):
+    pass
+
+
+def _exact_endgame_solve(cells0, mover, width, height, win_len, deadline):
+    """Memoized alpha-beta minimax to the TRUE end of the game, with
+    center-out move ordering for stronger pruning. Bounded by a hard
+    wall-clock `deadline` (not a node/depth budget), so a position outside
+    the calibrated safe zone (see the caller's `endgame_max_cols` gate)
+    fails loudly here (raises internally, caught below, reported as
+    `(None, None)`) rather than silently blowing the real move-time
+    budget -- the caller is expected to fall back to the round-based
+    search whenever this returns `(None, None)`.
+
+    Calibrated empirically against real Kaggle replay data: 5 legal
+    columns / 24 empty cells solves in ~0.4s; 6+ legal columns can take
+    5+ seconds with this plain-Python (no bitboard/transposition-table)
+    implementation -- callers should gate on `len(legal_columns) <= 5`
+    before calling this at all; the deadline is a second, independent
+    safety net, not the only guard.
+
+    Returns `(best_action, value)`, value from `mover`'s own perspective
+    (+1 win / -1 loss / 0 draw), or `(None, None)` if the deadline hit
+    before a definite answer was found."""
+    memo = {}
+    center = (width - 1) / 2
+
+    def solve(cells, to_move, alpha, beta):
+        if time.time() > deadline:
+            raise _EndgameTimeout()
+        key = (tuple(cells), to_move)
+        cached = memo.get(key)
+        if cached is not None:
+            return cached
+        other = OPPONENT if to_move == AGENT else AGENT
+        legal = sorted(_legal_columns(cells, width, height), key=lambda c: abs(c - center))
+        if not legal:
+            memo[key] = 0.0
+            return 0.0
+        if to_move == AGENT:
+            best = -2.0
+            for c in legal:
+                nxt = _apply_move(cells, c, to_move, width, height)
+                if _wins_for(nxt, to_move, width, height, win_len):
+                    val = 1.0
+                elif _board_full(nxt):
+                    val = 0.0
+                else:
+                    val = solve(nxt, other, alpha, beta)
+                best = max(best, val)
+                alpha = max(alpha, best)
+                if alpha >= beta:
+                    break
+        else:
+            best = 2.0
+            for c in legal:
+                nxt = _apply_move(cells, c, to_move, width, height)
+                if _wins_for(nxt, to_move, width, height, win_len):
+                    val = -1.0
+                elif _board_full(nxt):
+                    val = 0.0
+                else:
+                    val = solve(nxt, other, alpha, beta)
+                best = min(best, val)
+                beta = min(beta, best)
+                if alpha >= beta:
+                    break
+        memo[key] = best
+        return best
+
+    root_legal = _legal_columns(cells0, width, height)
+    if not root_legal:
+        return None, None
+    root_legal = sorted(root_legal, key=lambda c: abs(c - center))
+    other = OPPONENT if mover == AGENT else AGENT
+    try:
+        best_a, best_val = None, None
+        for c in root_legal:
+            nxt = _apply_move(cells0, c, mover, width, height)
+            if _wins_for(nxt, mover, width, height, win_len):
+                val = 1.0 if mover == AGENT else -1.0
+            elif _board_full(nxt):
+                val = 0.0
+            else:
+                val = solve(nxt, other, -1.0, 1.0)
+            if best_val is None or (mover == AGENT and val > best_val) or (mover == OPPONENT and val < best_val):
+                best_a, best_val = c, val
+            if (mover == AGENT and best_val == 1.0) or (mover == OPPONENT and best_val == -1.0):
+                break  # proven best-possible outcome -- no need to keep searching
+        return best_a, best_val
+    except _EndgameTimeout:
+        return None, None
+
+
+def _parity_cost(cells, width, height):
+    """A heuristic column-parity feature, in the spirit of classical
+    Connect-4 zugzwang/odd-even threat theory (Allis 1988) -- an
+    EXPERIMENTAL, opt-in attempt at the "avoid walking into a zugzwang"
+    problem this project's whitepaper documents as diagnosed-but-unsolved.
+
+    For each still-open column with `r` empty cells remaining, under
+    NAIVE same-column-only alternation starting with the player about to
+    move, the player who ends up placing the TOP piece is: the mover if
+    `r` is odd, the other player if `r` is even. Every leaf state this is
+    computed on follows a full agent-move-then-opponent-reply round (see
+    `real_adversarial_plan_action`'s two-phase leaf collection), so the
+    player "about to move" at every leaf is always AGENT -- no need to
+    track whose turn it is separately.
+
+    This is NOT a proof, NOT a guarantee, and NOT full Claimeven (a real
+    claimeven strategy requires REACTIVE move-pairing enforced across an
+    entire game, not a one-shot column count at a single position) -- it
+    is a cheap, directionally-motivated NUDGE: positions with more
+    even-parity open columns are scored as costlier (worse for the
+    agent), consistent with the theory's own prediction that parity
+    structure matters, without claiming this fully captures it. Returns
+    a COST (higher = worse for the agent), meant to be ADDED to the
+    leaf's existing value estimate, scaled by a small `parity_weight`
+    -- see `real_adversarial_plan_action`'s own docstring for the
+    honest, measured verdict on whether this actually helps."""
+    cost = 0.0
+    for c in range(width):
+        remaining = sum(1 for row in range(height) if cells[_rc(row, c, width)] == EMPTY)
+        if remaining == 0:
+            continue
+        cost += 1.0 if remaining % 2 == 0 else -1.0
+    return cost
+
+
+def _narrow_to_center(legal_cols, width, max_branching):
+    """Prunes a legal-column list down to `max_branching` columns closest
+    to center -- real Connect-4 domain knowledge (a center column touches
+    more potential 4-in-a-row lines than an edge one). `max_branching=None`
+    is a no-op; deeper (rounds=3+) search needs this to stay inside a real
+    time budget."""
+    if max_branching is None or len(legal_cols) <= max_branching:
+        return legal_cols
+    center = (width - 1) / 2
+    return sorted(legal_cols, key=lambda c: abs(c - center))[:max_branching]
+
+
+class _RoundSearchTimeout(Exception):
+    pass
+
+
+@torch.no_grad()
+def real_adversarial_plan_action(env, model, normalizer, real_state, memory=None,
+                                  memory_weight=0.25, memory_k=5, max_steps=None, rounds=1,
+                                  max_branching=None, endgame_max_cols=5, endgame_time_budget=1.2,
+                                  parity_weight=0.0, deeper_rounds=None, deeper_max_branching=4,
+                                  deeper_time_budget=0.6):
+    """`deeper_rounds`/`deeper_max_branching`/`deeper_time_budget`: a
+    real, mined-from-real-games gap this was built for -- `rounds` can
+    see zero danger on a position (every column looks equally safe) just
+    a couple of plies before a trap that one round DEEPER already
+    narrows down to exactly one safe column. A deeper search is provably
+    too slow to run on EVERY move (measured 8-13s at a 6-7-legal-column
+    branching factor) -- so this is a SAFE, opportunistic escalation, not
+    a blanket depth increase: `deeper_rounds=None` (default) reproduces
+    the original `rounds`-only behavior byte-for-byte. When set, AFTER
+    computing the normal-`rounds` answer (always -- the guaranteed-safe
+    fallback), a `deeper_rounds`-round search is attempted under a hard
+    `deeper_time_budget` deadline (`_RoundSearchTimeout`, same pattern as
+    `_exact_endgame_solve`'s own deadline -- checked at both exponential-
+    blowup recursion points AND immediately around the one batched NN
+    leaf-evaluation call, which is otherwise uninterruptible once
+    started). If it finishes in time, its answer is used instead
+    (strictly more information, never less); if it times out, the
+    original `rounds`-answer is returned completely unchanged -- this can
+    only ever help or be a no-op, never make the chosen move worse or
+    blow the real move-time budget by more than `deeper_time_budget`.
+    Calibrated to `deeper_time_budget=0.6s` against a 180-game regression
+    suite (random/weak/stronger opponents): zero win-rate regression, max
+    observed combined move time 1.641s -- comfortably under a 2s budget.
+
+    `parity_weight` (default 0.0, OFF -- byte-for-byte the original
+    behavior unless explicitly enabled): blends `_parity_cost` into every
+    leaf's value estimate, scaled by this weight. EXPERIMENTAL -- see
+    `_parity_cost`'s own docstring for exactly what this does and doesn't
+    claim. Measured (not just theorized) against the trusted test harness
+    at a few weights before shipping any non-zero default; see the
+    whitepaper for the honest result.
+
+    `rounds` real adversarial rounds (our move, then the opponent's
+    worst-case real reply, repeated) before falling back to the learned
+    value head as the leaf evaluator. `rounds=1` is the deployed default.
+    Root action never returns PASS.
+
+    `endgame_max_cols`/`endgame_time_budget`: before running the round-
+    based search at all, check whether the position already has few
+    enough legal columns for `_exact_endgame_solve` to solve it exactly,
+    well within budget. If it finishes in time, its answer is used
+    directly (provably optimal); otherwise this falls straight through to
+    the round-based search below, unchanged. `endgame_max_cols=0` disables
+    this path entirely.
+
+    Terminal-outcome convention: an opponent win scores `2 * max_steps`
+    (worse than merely running out of steps); a draw scores `max_steps` --
+    both far above any real achievable remaining-steps value, so they
+    never get confused with a genuine near-solved position.
+
+    Two-phase, globally batched leaf evaluation: rather than calling the
+    value head separately at every node in the search tree (expensive --
+    each call pays its own tensor-creation/memory-blend overhead), this
+    walks the tree TWICE: once (pure Python) to collect every non-
+    terminal leaf across the WHOLE tree into one deduplicated set
+    (transpositions collapse for free), then ONE batched value+memory
+    call, then a second walk doing the actual minimax from the
+    precomputed lookup."""
+    width, height, win_len = env.width, env.height, env.win_len
+    max_steps = max_steps if max_steps is not None else (width * height) // 2 + 2
+
+    cells_now = list(_decode_board(real_state))
+    now_legal = _legal_columns(cells_now, width, height)
+    if now_legal and endgame_max_cols and len(now_legal) <= endgame_max_cols:
+        exact_a, _exact_val = _exact_endgame_solve(
+            cells_now, AGENT, width, height, win_len, deadline=time.time() + endgame_time_budget,
+        )
+        if exact_a is not None:
+            return exact_a
+
+    def leaf_batch_values(states):
+        if not states:
+            return {}
+        states_t = states_to_tensor([env.observe(s) for s in states]).to(DEVICE)
+        norm_t = normalizer.normalize(states_t)
+        z = model.encode(norm_t)
+        vals = _evaluate_with_memory(model, z, memory, memory_weight, memory_k).tolist()
+        if parity_weight:
+            vals = [v + parity_weight * _parity_cost(list(_decode_board(s)), width, height)
+                    for v, s in zip(vals, states)]
+        return dict(zip(states, vals))
+
+    def run_search(search_rounds, search_max_branching, deadline):
+        """One full root-to-leaf search at a given (rounds, max_branching)
+        setting -- factored out so it can be called at two different
+        depths, see `deeper_rounds` above. `deadline` (optional): checked
+        at both exponential-blowup recursion points (leaf collection, our
+        own follow-up enumeration) AND immediately around the one batched
+        NN leaf-evaluation call (otherwise uninterruptible once started)
+        -- raises `_RoundSearchTimeout` the instant it's exceeded, letting
+        the caller safely abandon this attempt."""
+
+        def check_deadline():
+            if deadline is not None and time.time() > deadline:
+                raise _RoundSearchTimeout()
+
+        def collect_leaves(cells1, remaining_rounds, leaf_cache):
+            check_deadline()
+            if _board_full(cells1):
+                return
+            for opp_col in _legal_columns(cells1, width, height):
+                cells2 = _apply_move(cells1, opp_col, OPPONENT, width, height)
+                if _wins_for(cells2, OPPONENT, width, height, win_len) or _board_full(cells2):
+                    continue
+                if remaining_rounds <= 1:
+                    leaf_cache[tuple(_encode_board(cells2))] = None
+                else:
+                    for a2 in _narrow_to_center(_legal_columns(cells2, width, height), width, search_max_branching):
+                        cells3 = _apply_move(cells2, a2, AGENT, width, height)
+                        if _wins_for(cells3, AGENT, width, height, win_len):
+                            continue
+                        collect_leaves(cells3, remaining_rounds - 1, leaf_cache)
+
+        def score_after_our_move(cells1, remaining_rounds, leaf_cache):
+            """cells1: real board right after OUR move. Returns our worst-case
+            score -- the opponent picks whichever real reply hurts us most."""
+            if _board_full(cells1):
+                return float(max_steps)
+            vals = []
+            for opp_col in _legal_columns(cells1, width, height):
+                cells2 = _apply_move(cells1, opp_col, OPPONENT, width, height)
+                if _wins_for(cells2, OPPONENT, width, height, win_len):
+                    vals.append(float(2 * max_steps))
+                elif _board_full(cells2):
+                    vals.append(float(max_steps))
+                elif remaining_rounds <= 1:
+                    vals.append(leaf_cache[tuple(_encode_board(cells2))])
+                else:
+                    vals.append(score_after_opponent_move(cells2, remaining_rounds - 1, leaf_cache))
+            return max(vals)
+
+        def score_after_opponent_move(cells2, remaining_rounds, leaf_cache):
+            """cells2: real board after the opponent's move, our turn again.
+            Returns our best achievable worst-case score from here."""
+            check_deadline()
+            our_legal = _narrow_to_center(_legal_columns(cells2, width, height), width, search_max_branching)
+            if not our_legal:
+                return float(max_steps)
+            best = None
+            for a in our_legal:
+                cells3 = _apply_move(cells2, a, AGENT, width, height)
+                if _wins_for(cells3, AGENT, width, height, win_len):
+                    return -float(max_steps)  # a forced win exists deeper -- short-circuit
+                s = score_after_our_move(cells3, remaining_rounds, leaf_cache)
+                if best is None or s < best:
+                    best = s
+            return best
+
+        leaf_cache = {}
+        for a in surviving_actions:
+            collect_leaves(action_cells1[a], search_rounds, leaf_cache)
+        check_deadline()  # right before the one batched NN call -- don't start it with no budget left
+        if leaf_cache:
+            leaf_cache.update(leaf_batch_values(list(leaf_cache.keys())))
+        check_deadline()  # and once more right after -- don't walk the tree on a stale/over-budget result
+
+        best_a, best_score = None, None
+        for a in surviving_actions:
+            s = score_after_our_move(action_cells1[a], search_rounds, leaf_cache)
+            if best_score is None or s < best_score:
+                best_a, best_score = a, s
+        return best_a
+
+    root_legal = _legal_columns(cells_now, width, height)
+    if not root_legal:
+        return None
+    # Center-out root ordering: NOT a pruning change (every legal column is
+    # still considered), only fixes which column wins a TIE. Left-to-right
+    # order otherwise defaults ties to the LEFTMOST column, an arbitrary
+    # bias with no game-theoretic basis -- center columns are the real
+    # stronger choice under a tie.
+    center = (width - 1) / 2
+    root_legal = sorted(root_legal, key=lambda c: abs(c - center))
+
+    surviving_actions, action_cells1 = [], {}
+    for a in root_legal:
+        cells1 = _apply_move(cells_now, a, AGENT, width, height)
+        if _wins_for(cells1, AGENT, width, height, win_len):
+            return a  # immediate win -- take it, no need to consider anything else
+        surviving_actions.append(a)
+        action_cells1[a] = cells1
+
+    base_a = run_search(rounds, max_branching, deadline=None)  # always computed -- the guaranteed-safe fallback
+
+    if deeper_rounds is not None:
+        try:
+            return run_search(deeper_rounds, deeper_max_branching, deadline=time.time() + deeper_time_budget)
+        except _RoundSearchTimeout:
+            pass  # deeper attempt didn't finish in time -- fall back to base_a exactly as if deeper_rounds=None
+
+    return base_a

connectx/env.py

@@ -0,0 +1,417 @@
+"""
+ConnectX environment. Two board sizes, same code:
+- Small (`ConnectXEnv()`, 4x4, win_len=3): scaled down so an exact BFS
+  oracle stays tractable -- used for the domain's own self-test below.
+- Real (`ConnectXEnv(width=7, height=6, win_len=4)`, Kaggle's actual board):
+  `bfs_solve` returns None unconditionally -- the full game tree isn't
+  exhaustively searchable at this scale. Trained via on-policy Monte Carlo
+  value learning instead (see verifier.py / train.py), not oracle regression.
+
+State: the WIDTH*HEIGHT board, each cell one-hot over {EMPTY, AGENT,
+OPPONENT}. Row 0 = top; dropping into a column fills the lowest (highest
+row index) empty cell, standard Connect-4 gravity.
+
+Actions: DROP(column) for each column, plus one always-legal PASS action
+(guarantees `always_legal_actions` is a real, non-empty, unconditional
+subset). num_actions = WIDTH + 1.
+
+The fixed training opponent (deterministic unless the epsilon knobs below
+are set): after the agent's move, if the agent didn't already win or fill
+the board, the opponent (1) takes an immediate win if one exists, (2) else
+blocks the agent's immediate win if one exists, (3) else plays the leftmost
+legal column. This is this domain's one honest, named limitation for real
+Kaggle play: the model is trained against THIS specific opponent shape (plus
+diversification, see below), not whatever real opponent Kaggle's matchmaking
+actually pairs it against.
+
+Reward: -1 per agent action. "Solved" = the AGENT has win_len in a row after
+its own move. A loss (opponent wins) or a draw is terminal (`done=True`) but
+NOT solved -- this distinction matters: see search.py's comment on why
+trusting `done` alone as "solved" is a real bug for an adversarial domain.
+"""
+import random
+from collections import deque
+
+from .environment import Environment
+
+DEFAULT_WIDTH = 4
+DEFAULT_HEIGHT = 4
+DEFAULT_WIN_LEN = 3
+
+EMPTY, AGENT, OPPONENT = 0, 1, 2
+CELL_WIDTH = 3
+
+# Above this many cells, exhaustive BFS is not attempted. The small-board
+# default (16 cells) stays well under this; the real board (42 cells) is
+# always above it.
+BFS_MAX_CELLS = 20
+
+
+def _onehot(idx, n):
+    v = [0] * n
+    v[idx] = 1
+    return tuple(v)
+
+
+def _onehot_index(bits):
+    """Robust to a search-time DECODED state whose slot isn't cleanly
+    one-hot (real states, built via `_encode_board`, never hit the
+    fallback)."""
+    return bits.index(1) if 1 in bits else 0
+
+
+def _encode_board(cells):
+    return tuple(b for c in cells for b in _onehot(c, CELL_WIDTH))
+
+
+def _decode_board(state):
+    cells = []
+    off = 0
+    while off < len(state):
+        cells.append(_onehot_index(state[off:off + CELL_WIDTH]))
+        off += CELL_WIDTH
+    return cells
+
+
+def _rc(row, col, width):
+    return row * width + col
+
+
+def _lowest_empty_row(cells, col, width, height):
+    """Gravity: the row closest to the bottom that's still empty in this
+    column, or None if the column is full."""
+    for row in range(height - 1, -1, -1):
+        if cells[_rc(row, col, width)] == EMPTY:
+            return row
+    return None
+
+
+def _legal_columns(cells, width, height):
+    return [c for c in range(width) if _lowest_empty_row(cells, c, width, height) is not None]
+
+
+def _wins_for(cells, value, width, height, win_len):
+    """Whether `value` (AGENT or OPPONENT) has win_len in a row anywhere --
+    horizontal, vertical, or either diagonal."""
+    for row in range(height):
+        for col in range(width):
+            if cells[_rc(row, col, width)] != value:
+                continue
+            for dr, dc in ((0, 1), (1, 0), (1, 1), (1, -1)):
+                end_row = row + dr * (win_len - 1)
+                end_col = col + dc * (win_len - 1)
+                if not (0 <= end_row < height and 0 <= end_col < width):
+                    continue
+                if all(cells[_rc(row + dr * k, col + dc * k, width)] == value for k in range(win_len)):
+                    return True
+    return False
+
+
+def _board_full(cells):
+    return all(c != EMPTY for c in cells)
+
+
+def _stronger_opponent_move(cells, width, height, win_len):
+    """A second, deliberately stronger deterministic opponent -- same
+    win-now/block-immediate-win base as `_fixed_opponent_move`, plus one
+    more ply: among moves that survive those two checks, avoid any that
+    would hand the AGENT an immediate winning reply next turn, if a safer
+    alternative exists. Mixing this into TRAINING (see
+    `opponent_strong_epsilon` below) gives the value head real exposure to
+    a harder-to-punish opponent, not just noise around the weak one."""
+    legal = _legal_columns(cells, width, height)
+    for col in legal:
+        row = _lowest_empty_row(cells, col, width, height)
+        trial = list(cells)
+        trial[_rc(row, col, width)] = OPPONENT
+        if _wins_for(trial, OPPONENT, width, height, win_len):
+            return col
+    for col in legal:
+        row = _lowest_empty_row(cells, col, width, height)
+        trial = list(cells)
+        trial[_rc(row, col, width)] = AGENT
+        if _wins_for(trial, AGENT, width, height, win_len):
+            return col
+    safe = []
+    for col in legal:
+        row = _lowest_empty_row(cells, col, width, height)
+        nxt = list(cells)
+        nxt[_rc(row, col, width)] = OPPONENT
+        if _board_full(nxt):
+            safe.append(col)
+            continue
+        agent_can_win = False
+        for col2 in _legal_columns(nxt, width, height):
+            row2 = _lowest_empty_row(nxt, col2, width, height)
+            trial2 = list(nxt)
+            trial2[_rc(row2, col2, width)] = AGENT
+            if _wins_for(trial2, AGENT, width, height, win_len):
+                agent_can_win = True
+                break
+        if not agent_can_win:
+            safe.append(col)
+    return random.choice(safe) if safe else legal[0]
+
+
+def _fixed_opponent_move(cells, width, height, win_len, opponent_epsilon=0.0, opponent_strong_epsilon=0.0,
+                          opponent_selfplay_epsilon=0.0, opponent_policy_fn=None):
+    """Deterministic base heuristic: win now if possible, else block the
+    agent's immediate win, else leftmost legal column.
+
+    `opponent_epsilon`: with this probability, ignore the heuristic and
+    play a uniformly random legal column instead -- diversifies training
+    trajectories (a fully deterministic opponent means every training walk
+    from a matching starting side is the SAME exact game).
+
+    `opponent_strong_epsilon`: with this probability (checked after the
+    roll above), delegate the whole move to `_stronger_opponent_move`
+    instead -- direct training exposure to a harder opponent, not just
+    noise around the weak one.
+
+    `opponent_selfplay_epsilon` / `opponent_policy_fn`: with this
+    probability (checked last), delegate to an arbitrary caller-supplied
+    move function -- in practice, a frozen snapshot of this same model's
+    own move choice, viewed from the opponent's side (see
+    `train.make_selfplay_pool_opponent_fn`). This is the actual
+    "self-play" mechanism: every opponent above is a fixed, non-learning
+    heuristic the trained policy eventually plateaus against; self-play
+    is what lets it face something that keeps getting better."""
+    legal = _legal_columns(cells, width, height)
+    if opponent_epsilon > 0.0 and random.random() < opponent_epsilon:
+        return random.choice(legal)
+    if opponent_strong_epsilon > 0.0 and random.random() < opponent_strong_epsilon:
+        return _stronger_opponent_move(cells, width, height, win_len)
+    if opponent_selfplay_epsilon > 0.0 and opponent_policy_fn is not None \
+            and random.random() < opponent_selfplay_epsilon:
+        col = opponent_policy_fn(cells)
+        if col in legal:
+            return col
+    for col in legal:
+        row = _lowest_empty_row(cells, col, width, height)
+        trial = list(cells)
+        trial[_rc(row, col, width)] = OPPONENT
+        if _wins_for(trial, OPPONENT, width, height, win_len):
+            return col
+    for col in legal:
+        row = _lowest_empty_row(cells, col, width, height)
+        trial = list(cells)
+        trial[_rc(row, col, width)] = AGENT
+        if _wins_for(trial, AGENT, width, height, win_len):
+            return col
+    return legal[0]
+
+
+def is_solved(state, width, height, win_len):
+    return _wins_for(_decode_board(state), AGENT, width, height, win_len)
+
+
+def is_legal(state, action_idx, width, height):
+    pass_action = width
+    if action_idx == pass_action:
+        return True
+    if not (0 <= action_idx < width):
+        return False
+    return _lowest_empty_row(_decode_board(state), action_idx, width, height) is not None
+
+
+def step(state, action_idx, width, height, win_len, opponent_epsilon=0.0, opponent_strong_epsilon=0.0,
+         opponent_selfplay_epsilon=0.0, opponent_policy_fn=None):
+    pass_action = width
+    cells = list(_decode_board(state))
+    reward = -1.0
+
+    if action_idx != pass_action:
+        row = _lowest_empty_row(cells, action_idx, width, height)
+        cells[_rc(row, action_idx, width)] = AGENT
+
+    if _wins_for(cells, AGENT, width, height, win_len):
+        return _encode_board(cells), reward, True
+    if _board_full(cells):
+        return _encode_board(cells), reward, True  # draw -- terminal, not solved
+
+    opp_col = _fixed_opponent_move(cells, width, height, win_len, opponent_epsilon=opponent_epsilon,
+                                    opponent_strong_epsilon=opponent_strong_epsilon,
+                                    opponent_selfplay_epsilon=opponent_selfplay_epsilon,
+                                    opponent_policy_fn=opponent_policy_fn)
+    opp_row = _lowest_empty_row(cells, opp_col, width, height)
+    cells[_rc(opp_row, opp_col, width)] = OPPONENT
+
+    if _wins_for(cells, OPPONENT, width, height, win_len):
+        return _encode_board(cells), reward, True  # loss -- terminal, not solved
+    done = _board_full(cells)  # draw after opponent's move
+    return _encode_board(cells), reward, done
+
+
+def random_problem(rng, width, height):
+    """Every game starts from an empty board."""
+    return _encode_board([EMPTY] * (width * height)), None
+
+
+def bfs_solve(state, width, height, win_len, max_depth=8):
+    """Exact BFS for a forced win against the fixed opponent baked into
+    `step` -- not a general Connect-4 solver. Returns None above
+    BFS_MAX_CELLS (the real 7x6 board is never attempted)."""
+    if width * height > BFS_MAX_CELLS:
+        return None
+    if is_solved(state, width, height, win_len):
+        return []
+    frontier = deque([state])
+    parent = {state: None}
+    action_taken = {}
+    depth = {state: 0}
+    num_actions = width + 1
+    while frontier:
+        cur = frontier.popleft()
+        if depth[cur] >= max_depth:
+            continue
+        for a_idx in range(num_actions):
+            if not is_legal(cur, a_idx, width, height):
+                continue
+            nxt, _reward, done = step(cur, a_idx, width, height, win_len)
+            if nxt in parent:
+                continue
+            parent[nxt] = cur
+            action_taken[nxt] = a_idx
+            depth[nxt] = depth[cur] + 1
+            if is_solved(nxt, width, height, win_len):
+                path = []
+                node = nxt
+                while parent[node] is not None:
+                    path.append(action_taken[node])
+                    node = parent[node]
+                path.reverse()
+                return path
+            if not done:  # loss/draw states are terminal dead ends, don't expand
+                frontier.append(nxt)
+    return None
+
+
+_CELL_CHAR = {EMPTY: ".", AGENT: "A", OPPONENT: "O"}
+
+
+def format_state(state, width, height):
+    cells = _decode_board(state)
+    rows = [" ".join(_CELL_CHAR[cells[_rc(row, col, width)]] for col in range(width)) for row in range(height)]
+    return "\n" + "\n".join(rows)
+
+
+def format_action(action_idx, width):
+    return "PASS" if action_idx == width else f"DROP(col={action_idx})"
+
+
+class ConnectXEnv(Environment):
+    """width/height/win_len fixed per instance -- state_dim/num_actions
+    depend on them. Default (4x4, win_len=3) is the small, BFS-checkable
+    board this module's own self-test uses; pass width=7, height=6,
+    win_len=4 for the real Kaggle board."""
+
+    def __init__(self, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, win_len=DEFAULT_WIN_LEN,
+                 opponent_epsilon=0.0, opponent_strong_epsilon=0.0,
+                 opponent_selfplay_epsilon=0.0, opponent_policy_fn=None):
+        self.width = width
+        self.height = height
+        self.win_len = win_len
+        # Default 0.0 keeps this instance's `step` fully deterministic
+        # (required by the small-board self-test's BFS oracle). The real-
+        # board TRAINING env sets these > 0; its EVAL env keeps them at
+        # the default so it's still graded against the originally-defined
+        # fixed opponent.
+        self.opponent_epsilon = opponent_epsilon
+        self.opponent_strong_epsilon = opponent_strong_epsilon
+        self.opponent_selfplay_epsilon = opponent_selfplay_epsilon
+        self.opponent_policy_fn = opponent_policy_fn
+
+    @property
+    def state_dim(self):
+        return self.width * self.height * CELL_WIDTH
+
+    @property
+    def num_actions(self):
+        return self.width + 1
+
+    @property
+    def always_legal_actions(self):
+        return [self.width]  # PASS
+
+    def is_solved(self, state):
+        return is_solved(state, self.width, self.height, self.win_len)
+
+    def is_legal(self, state, action_idx):
+        return is_legal(state, action_idx, self.width, self.height)
+
+    def step(self, state, action_idx):
+        return step(state, action_idx, self.width, self.height, self.win_len,
+                     opponent_epsilon=self.opponent_epsilon,
+                     opponent_strong_epsilon=self.opponent_strong_epsilon,
+                     opponent_selfplay_epsilon=self.opponent_selfplay_epsilon,
+                     opponent_policy_fn=self.opponent_policy_fn)
+
+    def random_problem(self, rng, **kwargs):
+        return random_problem(rng, self.width, self.height)
+
+    def bfs_solve(self, state, max_depth=8):
+        return bfs_solve(state, self.width, self.height, self.win_len, max_depth=max_depth)
+
+    def format_state(self, state):
+        return format_state(state, self.width, self.height)
+
+    def format_action(self, action_idx):
+        return format_action(action_idx, self.width)
+
+
+if __name__ == "__main__":
+    env = ConnectXEnv()
+    print(f"state_dim={env.state_dim}  num_actions={env.num_actions}  "
+          f"board={env.width}x{env.height}  win_len={env.win_len}\n")
+
+    start, _ = env.random_problem(random.Random(0))
+    print(f"Empty board: {format_state(start, env.width, env.height)}")
+    path = env.bfs_solve(start)
+    assert path is not None, "no forced win found against the fixed opponent from an empty board"
+    print(f"Oracle's forced-win path: {[env.format_action(a) for a in path]}  (len={len(path)})")
+    cur = start
+    for a in path:
+        cur, r, done = env.step(cur, a)
+        print(f"  after {env.format_action(a)} (reward={r:.0f}, done={done}): {env.format_state(cur)}")
+    assert env.is_solved(cur), "oracle path did not reach a solved (agent-won) state"
+    print("\nCONFIRMED: the exact oracle finds a genuine forced win against the fixed opponent.\n")
+
+    print("=== Random-legal-play smoke test (30 games, no crashes, always terminates) ===")
+    rng = random.Random(1)
+    solved_count, loss_count, draw_count = 0, 0, 0
+    for _i in range(30):
+        state, _ = env.random_problem(rng)
+        for _ in range(env.width * env.height + 1):
+            legal = [a for a in range(env.num_actions) if env.is_legal(state, a)]
+            assert legal, "always_legal_actions guarantee violated -- PASS should always be legal"
+            a = rng.choice([a for a in legal if a != env.width] or legal)
+            state, _r, done = env.step(state, a)
+            if done:
+                break
+        else:
+            raise AssertionError("game did not terminate within the move cap")
+        if env.is_solved(state):
+            solved_count += 1
+        elif _board_full(_decode_board(state)):
+            draw_count += 1
+        else:
+            loss_count += 1
+    print(f"agent wins={solved_count}  losses={loss_count}  draws={draw_count}  (out of 30, random legal play)")
+    print("\nAll games terminated cleanly, always_legal_actions held in every state, no crashes.")
+
+    print("\n=== Real-board smoke test (7x6, win_len=4, no BFS oracle at this scale) ===")
+    real_env = ConnectXEnv(width=7, height=6, win_len=4)
+    print(f"state_dim={real_env.state_dim}  num_actions={real_env.num_actions}")
+    assert real_env.bfs_solve(real_env.random_problem(random.Random(0))[0]) is None, \
+        "bfs_solve should return None at real-board scale (no oracle by design)"
+    rng = random.Random(2)
+    state, _ = real_env.random_problem(rng)
+    for _ in range(real_env.width * real_env.height + 1):
+        legal = [a for a in range(real_env.num_actions) if real_env.is_legal(state, a)]
+        assert legal
+        a = rng.choice([a for a in legal if a != real_env.width] or legal)
+        state, _r, done = real_env.step(state, a)
+        if done:
+            break
+    print(real_env.format_state(state))
+    print("Real-board game ran to completion with no crashes; bfs_solve correctly returns None.")

connectx/environment.py

@@ -0,0 +1,70 @@
+"""
+Minimal common interface for a "domain" this codebase can train on. A state
+is any fixed-length tuple of numbers; how those numbers are interpreted is
+entirely up to the implementation below (`connectx_env.py`). Nothing in
+`model.py` / `train_utils.py` / `search.py` needs to change to support a new
+domain that implements this interface.
+"""
+from abc import ABC, abstractmethod
+
+
+class Environment(ABC):
+    # Every state built by this codebase is a discrete, exactly-hashable
+    # tuple (a one-hot-encoded Connect-4 board) -- kept as a class attribute
+    # rather than hardcoded into `search.py`'s cycle-detection logic so a
+    # future continuous-state domain could override it there without
+    # touching this interface.
+    discrete_state = True
+
+    def observe(self, state):
+        """What the model is allowed to see, as a function of the true
+        state. Default: full observability (identity). ConnectX is fully
+        observable, so this is never overridden -- kept as an explicit
+        extension point rather than removed, since every place that feeds
+        a state to the model calls this first, not the raw state."""
+        return state
+
+    @property
+    @abstractmethod
+    def state_dim(self):
+        """Length of the fixed-size numeric tuple representing a state."""
+
+    @property
+    @abstractmethod
+    def num_actions(self):
+        """Size of the fixed, discrete action space."""
+
+    @property
+    @abstractmethod
+    def always_legal_actions(self):
+        """Action indices legal from EVERY state, unconditionally -- used
+        by search's neurosymbolic decode-gate to fall back on safely."""
+
+    @abstractmethod
+    def is_solved(self, state):
+        ...
+
+    @abstractmethod
+    def is_legal(self, state, action_idx):
+        ...
+
+    @abstractmethod
+    def step(self, state, action_idx):
+        """Returns (next_state, reward, done). Assumes legality."""
+
+    @abstractmethod
+    def random_problem(self, rng, **kwargs):
+        """Returns (state, answer) -- answer is domain-specific (unused for
+        ConnectX, every game starts from the same empty board)."""
+
+    @abstractmethod
+    def bfs_solve(self, state, max_depth=8):
+        """Exact oracle: shortest forced win, or None if not found within
+        max_depth (or if the state space is too large to search -- see
+        connectx_env.py's BFS_MAX_CELLS)."""
+
+    def format_state(self, state):
+        return str(state)
+
+    def format_action(self, action_idx):
+        return str(action_idx)

connectx/episodic_memory.py

@@ -0,0 +1,136 @@
+"""
+A k-NN lookup table over real (encoded state, outcome) pairs from actually-
+played self-play games, consulted alongside the learned value head at
+decision time (see `search._evaluate_with_memory` / `adversarial_search.py`).
+Distinct from the trained model itself -- nothing here is learned, it's a
+cache of real experience the model can fall back on when its own value
+estimate might be shaky.
+
+Precedented by episodic control (Blundell et al., "Model-Free Episodic
+Control"; Pritzel et al., "Neural Episodic Control") and case-based
+reasoning, not a novel mechanism.
+"""
+import torch
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class EpisodicMemory:
+    """Stores (z, outcome) pairs -- z is a REAL encoded state's latent
+    (never a predicted/imagined one, so memory never compounds its own
+    errors), outcome is that state's real remaining-steps-to-win (for a
+    won game) or a fixed penalty (for a lost/drawn one) -- same label scale
+    as the value head, so blending stays consistent."""
+
+    def __init__(self):
+        self._zs = []
+        self._outcomes = []
+        self._state_keys = []
+        self._seen = {}  # state_key -> index, for exact-match dedup
+        self._trust_scale = None
+
+    def __len__(self):
+        return len(self._zs)
+
+    def add(self, z, outcome, state_key=None):
+        """`state_key`: an optional hashable identity for the real state
+        this (z, outcome) came from. When given and already stored, this
+        is a revisit of the exact same state -- keep whichever copy has
+        the better (smaller) outcome instead of appending a duplicate."""
+        if state_key is not None and state_key in self._seen:
+            idx = self._seen[state_key]
+            if outcome < self._outcomes[idx]:
+                self._zs[idx] = z.detach().to("cpu")
+                self._outcomes[idx] = float(outcome)
+                self._trust_scale = None
+            return
+        self._zs.append(z.detach().to("cpu"))
+        self._outcomes.append(float(outcome))
+        self._state_keys.append(state_key)
+        if state_key is not None:
+            self._seen[state_key] = len(self._zs) - 1
+        self._trust_scale = None
+
+    def _stacked(self):
+        return torch.stack(self._zs).to(DEVICE), torch.tensor(self._outcomes, device=DEVICE)
+
+    @torch.no_grad()
+    def trust_scale(self, sample_size=300):
+        """A self-calibrating distance scale for this memory's own latent
+        space: the median nearest-OTHER-neighbor distance among a random
+        subsample of stored points. A query distance much smaller than
+        this means "genuinely close match found"; much larger means
+        "nothing like this was ever stored." Self-calibrating per
+        checkpoint/latent-dim rather than a hand-picked constant."""
+        if self._trust_scale is not None:
+            return self._trust_scale
+        n = len(self._zs)
+        if n < 2:
+            self._trust_scale = 1.0
+            return self._trust_scale
+        Z, _outcomes = self._stacked()
+        if n > sample_size:
+            idx = torch.randperm(n, device=DEVICE)[:sample_size]
+            sample = Z[idx]
+        else:
+            sample = Z
+        dists = torch.cdist(sample, Z)
+        dists = torch.where(dists > 1e-6, dists, torch.full_like(dists, float("inf")))
+        nn_dist = dists.min(dim=1).values
+        nn_dist = nn_dist[torch.isfinite(nn_dist)]
+        self._trust_scale = nn_dist.median().item() if len(nn_dist) > 0 else 1.0
+        return self._trust_scale
+
+    @torch.no_grad()
+    def query_batch(self, zs, k=5):
+        """zs: [B, latent_dim]. Returns (blended_estimates [B], trust [B]).
+        Each row's blend weights its k nearest stored neighbors by inverse
+        distance. `trust` is `exp(-mean_distance / trust_scale)` -- a 0..1
+        confidence already normalized against this memory's own typical
+        spacing, so callers can scale their blend weight by it directly."""
+        Z, outcomes = self._stacked()
+        zq = zs.detach().to(DEVICE)
+        dists = torch.cdist(zq, Z)
+        k = min(k, len(self._zs))
+        topk_dists, topk_idx = torch.topk(dists, k, largest=False, dim=1)
+        topk_outcomes = outcomes[topk_idx]
+        weights = 1.0 / (topk_dists + 1e-2)
+        weights = weights / weights.sum(dim=1, keepdim=True)
+        blended = (weights * topk_outcomes).sum(dim=1)
+        mean_dist = topk_dists.mean(dim=1)
+        trust = torch.exp(-mean_dist / self.trust_scale())
+        return blended, trust
+
+
+@torch.no_grad()
+def add_trajectory_from_real_path(model, normalizer, memory, path_states, env=None):
+    """Adds every state of an actually-played, WON trajectory as a positive
+    (attraction) example -- outcome = real distance to the end of this
+    path."""
+    from .train_utils import states_to_tensor
+
+    dedup = env is not None and env.discrete_state
+    T = len(path_states) - 1
+    for t, s in enumerate(path_states):
+        observed = env.observe(s) if env is not None else s
+        z = normalizer.normalize(states_to_tensor([observed]).to(DEVICE))
+        z = model.encode(z)[0]
+        memory.add(z, T - t, state_key=s if dedup else None)
+
+
+@torch.no_grad()
+def add_negative_trajectory_from_real_path(model, normalizer, memory, path_states, penalty, env=None):
+    """Adds every state of a LOST/drawn trajectory as a negative
+    (repulsion) example -- every state gets the SAME fixed penalty label,
+    deliberately uniform across the whole walk. No new retrieval mechanism
+    needed: `EpisodicMemory.query_batch`'s existing k-NN blend already
+    treats a nearby HIGH-outcome entry as repulsion by construction, the
+    exact mirror of how a low one acts as attraction."""
+    from .train_utils import states_to_tensor
+
+    dedup = env is not None and env.discrete_state
+    for s in path_states:
+        observed = env.observe(s) if env is not None else s
+        z = normalizer.normalize(states_to_tensor([observed]).to(DEVICE))
+        z = model.encode(z)[0]
+        memory.add(z, float(penalty), state_key=s if dedup else None)

connectx/lora.py

@@ -0,0 +1,63 @@
+"""
+Generic LoRA (Low-Rank Adaptation, Hu et al. 2021) wrapper -- module-agnostic,
+works on any nn.Module built from nn.Linear layers. Used here to fine-tune
+the value head under self-play with a capacity-constrained update instead of
+a full-parameter one, which empirically fine-tunes more reliably on a small,
+self-generated data distribution without regressing.
+"""
+import torch
+from torch import nn
+
+
+class LoRALinear(nn.Module):
+    """Wraps an existing nn.Linear, freezing its original weight/bias and
+    adding a trainable low-rank delta:
+        output = frozen_linear(x) + scaling * (x @ A^T) @ B^T
+    `B` is initialized to ZERO, so the wrapped layer's output is byte-
+    identical to the original before any training happens."""
+
+    def __init__(self, linear, rank=4, alpha=1.0):
+        super().__init__()
+        assert isinstance(linear, nn.Linear), f"LoRALinear only wraps nn.Linear, got {type(linear)}"
+        self.linear = linear
+        for p in self.linear.parameters():
+            p.requires_grad = False
+        self.rank = rank
+        self.scaling = alpha / rank
+        device = linear.weight.device
+        dtype = linear.weight.dtype
+        self.lora_A = nn.Parameter(torch.randn(rank, linear.in_features, device=device, dtype=dtype) * 0.01)
+        self.lora_B = nn.Parameter(torch.zeros(linear.out_features, rank, device=device, dtype=dtype))
+
+    def forward(self, x):
+        base = self.linear(x)
+        delta = (x @ self.lora_A.t()) @ self.lora_B.t()
+        return base + self.scaling * delta
+
+    def lora_parameters(self):
+        return [self.lora_A, self.lora_B]
+
+    def merge_into_base(self):
+        """Fold the current LoRA delta into the frozen base weight and
+        reset B to zero -- used to "commit" a trained adapter back into a
+        plain nn.Linear-equivalent state so the saved checkpoint is an
+        ordinary state_dict, loadable with zero LoRA-awareness downstream."""
+        with torch.no_grad():
+            delta_w = self.scaling * (self.lora_B @ self.lora_A)
+            self.linear.weight.add_(delta_w)
+            self.lora_B.zero_()
+
+
+def apply_lora(module, rank=4, alpha=1.0):
+    """Recursively replace every nn.Linear submodule of `module` with a
+    LoRALinear wrapper. Returns the new LoRA parameters for the caller's
+    optimizer. Modifies `module` in place."""
+    lora_params = []
+    for name, child in list(module.named_children()):
+        if isinstance(child, nn.Linear):
+            wrapped = LoRALinear(child, rank=rank, alpha=alpha)
+            setattr(module, name, wrapped)
+            lora_params.extend(wrapped.lora_parameters())
+        else:
+            lora_params.extend(apply_lora(child, rank=rank, alpha=alpha))
+    return lora_params

connectx/memory_build.py

@@ -0,0 +1,55 @@
+"""
+Builds an EpisodicMemory offline, at packaging time, by playing self-play
+games against a MIXED opponent (weak heuristic + random + the stronger
+1-ply-deeper heuristic) using the model's own real adversarial search.
+WON games are stored as positive (attraction) examples; LOST/drawn games
+are ALSO stored, as negative (repulsion) examples -- deliberately mixed
+opposition, not just the single fixed weak training opponent, so a loss
+has to have actually lost to a real mix of opposition before it gets
+stored as "this is bad," rather than encoding one narrow opponent's
+particular blind spots as universal truth.
+"""
+import torch
+
+from .env import ConnectXEnv
+from .adversarial_search import real_adversarial_plan_action
+from .episodic_memory import EpisodicMemory, add_trajectory_from_real_path, add_negative_trajectory_from_real_path
+
+
+@torch.no_grad()
+def build_episodic_memory(env, model, normalizer, rng, n_games=500, opponent_epsilon=0.2,
+                           opponent_strong_epsilon=0.3, adversarial_rounds=2, penalty=None):
+    """`penalty` (default 2x max_steps, deliberately bigger than the value
+    head's own unsolved_penalty): these are discrete stored memory points,
+    not a training-loss target, so being a bit more emphatic buys sharper
+    repulsion without the overfitting risk more gradient steps would
+    carry."""
+    memory = EpisodicMemory()
+    diverse_env = ConnectXEnv(width=env.width, height=env.height, win_len=env.win_len,
+                               opponent_epsilon=opponent_epsilon, opponent_strong_epsilon=opponent_strong_epsilon)
+    max_steps = (env.width * env.height) // 2 + 2
+    if penalty is None:
+        penalty = 2 * max_steps
+    wins, losses = 0, 0
+    for _ in range(n_games):
+        state, _ = diverse_env.random_problem(rng)
+        path_states = [state]
+        for _ in range(max_steps):
+            if diverse_env.is_solved(state):
+                break
+            a = real_adversarial_plan_action(diverse_env, model, normalizer, state, rounds=adversarial_rounds)
+            if a is None:
+                break
+            state, _r, done = diverse_env.step(state, a)
+            path_states.append(state)
+            if done:
+                break
+        if diverse_env.is_solved(state):
+            wins += 1
+            add_trajectory_from_real_path(model, normalizer, memory, path_states, env=diverse_env)
+        else:
+            losses += 1
+            add_negative_trajectory_from_real_path(model, normalizer, memory, path_states, penalty, env=diverse_env)
+    print(f"  built memory from {wins} won + {losses} lost self-play games "
+          f"(mixed opponent) -> {len(memory)} stored states")
+    return memory

connectx/model.py

@@ -0,0 +1,104 @@
+"""
+The four learned pieces: an encoder (real state -> latent), a dynamics model
+(latent + action -> predicted next latent + reward), a value head (latent ->
+estimated cost-to-go), and a diagnostic decoder (latent -> reconstructed
+state, used only to sanity-check the latent isn't collapsing -- never
+consulted for a real decision at inference time).
+
+Config-driven: nothing here is hardcoded to Connect-4 beyond `state_dim`/
+`num_actions`, so this file would work for any domain with a fixed-length
+tuple state and a fixed discrete action space.
+"""
+import torch
+import torch.nn as nn
+
+
+def mlp(dims, out_activation=None):
+    layers = []
+    for i in range(len(dims) - 1):
+        layers.append(nn.Linear(dims[i], dims[i + 1]))
+        is_last = i == len(dims) - 2
+        if not is_last:
+            layers.append(nn.ReLU())
+        elif out_activation is not None:
+            layers.append(out_activation)
+    return nn.Sequential(*layers)
+
+
+class Encoder(nn.Module):
+    def __init__(self, state_dim, latent_dim, hidden_dim=128):
+        super().__init__()
+        self.net = mlp([state_dim, hidden_dim, hidden_dim, latent_dim])
+
+    def forward(self, state):
+        return self.net(state)
+
+
+class DynamicsModel(nn.Module):
+    def __init__(self, latent_dim, num_actions, hidden_dim=128):
+        super().__init__()
+        self.num_actions = num_actions
+        self.trunk = mlp([latent_dim + num_actions, hidden_dim, hidden_dim])
+        self.next_latent_head = nn.Linear(hidden_dim, latent_dim)
+        self.reward_head = nn.Linear(hidden_dim, 1)
+
+    def forward(self, z, action_idx):
+        action_onehot = nn.functional.one_hot(action_idx, self.num_actions).float()
+        h = self.trunk(torch.cat([z, action_onehot], dim=-1))
+        next_z = self.next_latent_head(h)
+        reward = self.reward_head(h).squeeze(-1)
+        return next_z, reward
+
+
+class ValueHead(nn.Module):
+    def __init__(self, latent_dim, hidden_dim=128):
+        super().__init__()
+        self.net = mlp([latent_dim, hidden_dim, hidden_dim, 1])
+
+    def forward(self, z):
+        return self.net(z).squeeze(-1)
+
+
+class Decoder(nn.Module):
+    def __init__(self, latent_dim, state_dim, hidden_dim=128):
+        super().__init__()
+        self.net = mlp([latent_dim, hidden_dim, hidden_dim, state_dim])
+
+    def forward(self, z):
+        return self.net(z)
+
+
+class WorldModel(nn.Module):
+    def __init__(self, state_dim, num_actions, latent_dim=64, hidden_dim=128):
+        super().__init__()
+        self.state_dim = state_dim
+        self.num_actions = num_actions
+        self.latent_dim = latent_dim
+        self.encoder = Encoder(state_dim, latent_dim, hidden_dim)
+        self.dynamics = DynamicsModel(latent_dim, num_actions, hidden_dim)
+        self.value = ValueHead(latent_dim, hidden_dim)
+        self.decoder = Decoder(latent_dim, state_dim, hidden_dim)
+        # Value-target normalization stats: NOT learned, set once by
+        # whichever value-training pass runs (see verifier.py's
+        # train_mc_value_onpolicy) from the actual label distribution it
+        # sees. Buffers (not plain attributes) so they save/load with the
+        # checkpoint automatically.
+        self.register_buffer("value_target_mean", torch.tensor(0.0))
+        self.register_buffer("value_target_std", torch.tensor(1.0))
+
+    def encode(self, state):
+        return self.encoder(state)
+
+    def imagine_step(self, z, action_idx):
+        return self.dynamics(z, action_idx)
+
+    def evaluate(self, z):
+        """Always returns real-scale value estimates (remaining cost, same
+        units `imagine_step`'s predicted reward uses) -- the head internally
+        predicts a normalized target, denormalized here so no caller needs
+        to know normalization is happening."""
+        raw = self.value(z)
+        return raw * self.value_target_std + self.value_target_mean
+
+    def reconstruct(self, z):
+        return self.decoder(z)

connectx/search.py

@@ -0,0 +1,179 @@
+"""
+Latent-space lookahead solver ("Baseline A" at depth=1, real branching
+search at depth>1) -- this is the ORIGINAL search approach this codebase
+started with, kept here as the comparison baseline `adversarial_search.py`
+(the real approach actually deployed) is measured against. See the
+whitepaper's results table for why real board-space search decisively beats
+this for an adversarial domain.
+
+The neurosymbolic decode-gate (`caution` below) exists because an earlier
+version that let the dynamics model imagine arbitrarily deep with no
+legality check at all let it extrapolate to state/action combinations it
+never saw during training, corrupting even the very first move's score.
+With probability `caution`, each beam entry's latent is decoded back to an
+estimated real state (via the trained decoder) and the REAL legal-action
+mask is computed from that -- decoding is used only to filter which moves
+are allowed, never to make the value judgement itself, which stays fully
+latent.
+"""
+import random
+
+import torch
+
+from .model import WorldModel
+from .train_utils import StateNormalizer, states_to_tensor
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def load_checkpoint(path):
+    ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
+    model = WorldModel(
+        state_dim=ckpt["state_dim"],
+        num_actions=ckpt["num_actions"],
+        latent_dim=ckpt["latent_dim"],
+        hidden_dim=ckpt.get("hidden_dim", 256),
+    ).to(DEVICE)
+    model.load_state_dict(ckpt["model_state"], strict=False)
+    model.eval()
+
+    normalizer = StateNormalizer.__new__(StateNormalizer)
+    normalizer.mean = ckpt["norm_mean"].to(DEVICE)
+    normalizer.std = ckpt["norm_std"].to(DEVICE)
+    return model, normalizer
+
+
+class VisitedStates:
+    """Plain hash-set cycle detector -- ConnectX is fully discrete, so this
+    is always an O(1) membership check."""
+
+    def __init__(self, initial_state):
+        self._set = {initial_state}
+
+    def __contains__(self, state):
+        return state in self._set
+
+    def add(self, state):
+        self._set.add(state)
+
+
+def _evaluate_with_memory(model, cand_z, memory, memory_weight, memory_k):
+    """model.evaluate(cand_z), optionally blended with an EpisodicMemory's
+    k-NN lookup at the same latents. `memory=None` is the exact original
+    behavior with zero overhead. Blend weight is scaled by the memory's own
+    self-calibrated `trust` (~1 for a genuinely close match, ~0 for nothing
+    similar ever stored), so a distant, irrelevant neighbor doesn't get
+    blended in at the same weight as a close one."""
+    values = model.evaluate(cand_z)
+    if memory is not None and len(memory) > 0 and memory_weight > 0:
+        blended, trust = memory.query_batch(cand_z, k=memory_k)
+        w = memory_weight * trust
+        values = (1 - w) * values + w * blended
+    return values
+
+
+@torch.no_grad()
+def plan_action(env, model, normalizer, real_state, depth=3, beam_width=8, caution=1.0, rng=None,
+                 memory=None, memory_weight=0.25, memory_k=5):
+    """Best first action for real_state, chosen by beam search purely in
+    latent space. depth=1 reduces exactly to Baseline A (no lookahead
+    beyond the immediate predicted next state)."""
+    rng = rng or random
+
+    root_legal = [a for a in range(env.num_actions) if env.is_legal(real_state, a)]
+    if not root_legal:
+        return None
+
+    z0 = normalizer.normalize(states_to_tensor([env.observe(real_state)]).to(DEVICE))
+    z0 = model.encode(z0)[0]
+
+    # beam entries: (predicted_z, action_seq, cumulative_predicted_reward)
+    beam = [(z0, [], 0.0)]
+
+    for step in range(depth):
+        if step == 0:
+            allowed_per_entry = [root_legal]
+        elif rng.random() < caution:
+            zs = torch.stack([z for z, _seq, _r in beam])
+            decoded = normalizer.denormalize(model.reconstruct(zs)).round()
+            allowed_per_entry = []
+            for row in decoded.tolist():
+                decoded_state = tuple(int(x) for x in row)
+                legal = [a for a in range(env.num_actions) if env.is_legal(decoded_state, a)]
+                allowed_per_entry.append(legal if legal else env.always_legal_actions)
+        else:
+            allowed_per_entry = [env.always_legal_actions] * len(beam)
+
+        candidates = []
+        for (z, seq, cum_r), allowed in zip(beam, allowed_per_entry):
+            z_batch = z.unsqueeze(0).repeat(len(allowed), 1)
+            a_batch = torch.tensor(allowed, dtype=torch.long, device=DEVICE)
+            next_z_batch, reward_batch = model.imagine_step(z_batch, a_batch)
+            for i, a in enumerate(allowed):
+                candidates.append((next_z_batch[i], seq + [a], cum_r + reward_batch[i].item()))
+
+        if not candidates:
+            break
+
+        cand_z = torch.stack([c[0] for c in candidates])
+        values = _evaluate_with_memory(model, cand_z, memory, memory_weight, memory_k)
+        cum_rewards = torch.tensor([c[2] for c in candidates], dtype=torch.float32, device=DEVICE)
+        scores = -cum_rewards + values  # value is a cost estimate; combine with accumulated reward
+
+        k = min(beam_width, len(candidates))
+        top_idx = torch.topk(scores, k, largest=False).indices.tolist()
+        beam = [(candidates[i][0], candidates[i][1], candidates[i][2]) for i in top_idx]
+
+    final_z = torch.stack([b[0] for b in beam])
+    final_values = _evaluate_with_memory(model, final_z, memory, memory_weight, memory_k)
+    final_cum_rewards = torch.tensor([b[2] for b in beam], dtype=torch.float32, device=DEVICE)
+    final_scores = -final_cum_rewards + final_values
+    best_idx = torch.argmin(final_scores).item()
+    return beam[best_idx][1][0]
+
+
+def solve_with_search_counted(env, model, normalizer, state, depth, beam_width, caution=1.0, max_total_steps=12,
+                               memory=None, memory_weight=0.25, memory_k=5):
+    """Runs `plan_action` step by step against the real environment,
+    tracking visited states to fail fast on a cycle.
+
+    The `is_solved` check on a `done` step is NOT a redundant safety net --
+    it's a real, previously-fixed bug class: `done=True` fires on a LOSS or
+    a DRAW in this domain, not just a win (unlike every single-agent puzzle
+    domain, where a dead end is simply never marked done, so `done` and
+    `is_solved` always agreed). Trusting `done` alone here silently counted
+    real losses as wins."""
+    cur = state
+    visited = VisitedStates(state)
+    for i in range(max_total_steps):
+        if env.is_solved(cur):
+            return True, i
+        a = plan_action(env, model, normalizer, cur, depth=depth, beam_width=beam_width, caution=caution,
+                         memory=memory, memory_weight=memory_weight, memory_k=memory_k)
+        if a is None:
+            return False, None
+        next_state, _, done = env.step(cur, a)
+        if next_state in visited:
+            return False, None
+        visited.add(next_state)
+        cur = next_state
+        if done:
+            return env.is_solved(cur), (i + 1 if env.is_solved(cur) else None)
+    return env.is_solved(cur), (max_total_steps if env.is_solved(cur) else None)
+
+
+def evaluate(env, model, normalizer, problems, depth, beam_width, caution=1.0, max_total_steps=12, label="",
+             memory=None, memory_weight=0.25, memory_k=5):
+    solved = 0
+    total_steps = 0
+    for state, _answer in problems:
+        ok, steps = solve_with_search_counted(env, model, normalizer, state, depth, beam_width, caution,
+                                               max_total_steps, memory=memory, memory_weight=memory_weight,
+                                               memory_k=memory_k)
+        if ok:
+            solved += 1
+            total_steps += steps
+    n = len(problems)
+    avg_steps = total_steps / solved if solved else float("nan")
+    print(f"{label:30s} solve_rate={solved/n:.3f} ({solved}/{n})  avg_steps_when_solved={avg_steps:.2f}")
+    return solved / n, avg_steps

connectx/train_utils.py

@@ -0,0 +1,205 @@
+"""
+Stage 1: self-supervised training of encoder + dynamics + decoder on random
+(state, action, next_state, reward) transitions and multi-step rollouts --
+no labels, no oracle, works for any domain implementing `environment.py`'s
+`Environment` interface.
+
+(Value-head training -- "stage 2" -- lives in `verifier.py`, since ConnectX
+uses on-policy Monte Carlo value learning rather than an oracle-labeled
+regression: there's no tractable exact solver at the real 7x6 board scale.)
+"""
+import torch
+import torch.nn as nn
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def sample_legal_action(env, rng, state):
+    legal = [i for i in range(env.num_actions) if env.is_legal(state, i)]
+    return rng.choice(legal)
+
+
+def generate_transitions(env, rng, n_problems=4000, walk_len=6, **problem_kwargs):
+    """Random-walk from random problem starts, collecting (s, a, s', r)
+    along the way -- covers realistic reachable states, not just problem
+    starts."""
+    transitions = []
+    for _ in range(n_problems):
+        state, _ = env.random_problem(rng, **problem_kwargs)
+        for _ in range(walk_len):
+            if env.is_solved(state):
+                break
+            a_idx = sample_legal_action(env, rng, state)
+            next_state, reward, done = env.step(state, a_idx)
+            transitions.append((state, a_idx, next_state, reward))
+            state = next_state
+            if done:
+                break
+    return transitions
+
+
+def generate_rollout_sequences(env, rng, n_problems=4000, k=3, **problem_kwargs):
+    """K-step (states, actions) sequences for the unrolled/open-loop
+    training objective below."""
+    sequences = []
+    for _ in range(n_problems):
+        state, _ = env.random_problem(rng, **problem_kwargs)
+        states = [state]
+        actions = []
+        cur = state
+        for _ in range(k):
+            a_idx = sample_legal_action(env, rng, cur)
+            cur, _, _ = env.step(cur, a_idx)
+            actions.append(a_idx)
+            states.append(cur)
+        sequences.append((states, actions))
+    return sequences
+
+
+class StateNormalizer:
+    def __init__(self, states):
+        t = torch.tensor(states, dtype=torch.float32)
+        self.mean = t.mean(dim=0)
+        self.std = t.std(dim=0).clamp_min(1e-3)
+
+    def to(self, device):
+        self.mean = self.mean.to(device)
+        self.std = self.std.to(device)
+        return self
+
+    def normalize(self, state_tensor):
+        return (state_tensor - self.mean) / self.std
+
+    def denormalize(self, norm_tensor):
+        return norm_tensor * self.std + self.mean
+
+
+def states_to_tensor(states):
+    return torch.tensor([list(s) for s in states], dtype=torch.float32)
+
+
+def train_stage1(model, normalizer, transitions, sequences=None, k=3,
+                  epochs=400, batch_size=256, lr=1e-3):
+    """`sequences` (see generate_rollout_sequences) adds a k-step UNROLLED
+    loss: the dynamics model is chained k times, feeding its own predicted
+    latent back in as input at each step (never re-encoding the real
+    intermediate state). This matches how the model is actually used at
+    inference (chained multi-step search) -- training on 1-step transitions
+    alone leaves compounding rollout error unaddressed."""
+    state_dim = normalizer.mean.shape[0]
+    states = states_to_tensor([t[0] for t in transitions]).to(DEVICE)
+    actions = torch.tensor([t[1] for t in transitions], dtype=torch.long).to(DEVICE)
+    next_states = states_to_tensor([t[2] for t in transitions]).to(DEVICE)
+    rewards = torch.tensor([t[3] for t in transitions], dtype=torch.float32).to(DEVICE)
+
+    norm_states = normalizer.normalize(states)
+    norm_next_states = normalizer.normalize(next_states)
+
+    if sequences is not None:
+        seq_states_raw = torch.stack([states_to_tensor(s) for s, _ in sequences]).to(DEVICE)  # [N, k+1, D]
+        seq_actions = torch.tensor([a for _, a in sequences], dtype=torch.long).to(DEVICE)  # [N, k]
+        norm_seq_states = normalizer.normalize(seq_states_raw.view(-1, state_dim)).view(seq_states_raw.shape)
+
+    n = states.shape[0]
+    opt = torch.optim.Adam(
+        list(model.encoder.parameters()) + list(model.dynamics.parameters()) + list(model.decoder.parameters()),
+        lr=lr,
+    )
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
+    mse = nn.MSELoss()
+
+    for epoch in range(epochs):
+        perm = torch.randperm(n, device=DEVICE)
+        total_loss = 0.0
+        for i in range(0, n, batch_size):
+            idx = perm[i:i + batch_size]
+            s, a, s_next, r = norm_states[idx], actions[idx], norm_next_states[idx], rewards[idx]
+
+            z = model.encode(s)
+            z_next_target = model.encode(s_next)
+            pred_next_z, pred_reward = model.imagine_step(z, a)
+
+            recon = model.reconstruct(z)
+            recon_next_from_dynamics = model.reconstruct(pred_next_z)
+
+            loss = (
+                mse(recon, s)
+                + mse(pred_next_z, z_next_target)
+                + mse(recon_next_from_dynamics, s_next)
+                + mse(pred_reward, r)
+            )
+
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            total_loss += loss.item() * idx.shape[0]
+
+        if sequences is not None:
+            n_seq = norm_seq_states.shape[0]
+            seq_perm = torch.randperm(n_seq, device=DEVICE)
+            unrolled_total = 0.0
+            for i in range(0, n_seq, batch_size):
+                sidx = seq_perm[i:i + batch_size]
+                seq_s = norm_seq_states[sidx]  # [B, k+1, D]
+                seq_a = seq_actions[sidx]  # [B, k]
+
+                z = model.encode(seq_s[:, 0, :])
+                loss_unrolled = 0.0
+                for step in range(k):
+                    z, _pred_r = model.imagine_step(z, seq_a[:, step])
+                    target = seq_s[:, step + 1, :]
+                    target_z = model.encode(target)
+                    loss_unrolled = loss_unrolled + mse(z, target_z) + mse(model.reconstruct(z), target)
+                loss_unrolled = loss_unrolled / k
+
+                opt.zero_grad()
+                loss_unrolled.backward()
+                opt.step()
+                unrolled_total += loss_unrolled.item() * sidx.shape[0]
+
+        sched.step()
+        if (epoch + 1) % 20 == 0 or epoch == 0:
+            msg = f"  [stage1] epoch {epoch+1:3d}/{epochs}  loss={total_loss/n:.4f}"
+            if sequences is not None:
+                msg += f"  unrolled_loss={unrolled_total/n_seq:.4f}"
+            print(msg)
+
+
+@torch.no_grad()
+def eval_stage1(model, normalizer, transitions):
+    """Decoder reconstruction / dynamics-rollout exact-match accuracy (after
+    rounding to the nearest integer) -- the diagnostic that checks the
+    latent isn't collapsing to something the decoder can't read back out."""
+    states = states_to_tensor([t[0] for t in transitions]).to(DEVICE)
+    actions = torch.tensor([t[1] for t in transitions], dtype=torch.long).to(DEVICE)
+    next_states = states_to_tensor([t[2] for t in transitions]).to(DEVICE)
+
+    norm_states = normalizer.normalize(states)
+    z = model.encode(norm_states)
+    recon = normalizer.denormalize(model.reconstruct(z))
+    pred_next_z, _ = model.imagine_step(z, actions)
+    recon_next = normalizer.denormalize(model.reconstruct(pred_next_z))
+
+    recon_acc = (recon.round() == states).all(dim=1).float().mean().item()
+    dyn_acc = (recon_next.round() == next_states).all(dim=1).float().mean().item()
+    return {"decoder_recon_exact_acc": recon_acc, "dynamics_rollout_exact_acc": dyn_acc}
+
+
+@torch.no_grad()
+def eval_multistep_rollout(model, normalizer, sequences, k):
+    """Chained (open-loop) rollout exact-match accuracy at each step 1..k --
+    exposes compounding error the way a single-step eval can't."""
+    state_dim = normalizer.mean.shape[0]
+    seq_states_raw = torch.stack([states_to_tensor(s) for s, _ in sequences]).to(DEVICE)
+    seq_actions = torch.tensor([a for _, a in sequences], dtype=torch.long).to(DEVICE)
+    norm_seq_states = normalizer.normalize(seq_states_raw.view(-1, state_dim)).view(seq_states_raw.shape)
+
+    z = model.encode(norm_seq_states[:, 0, :])
+    results = {}
+    for step in range(k):
+        z, _ = model.imagine_step(z, seq_actions[:, step])
+        recon = normalizer.denormalize(model.reconstruct(z))
+        real = seq_states_raw[:, step + 1, :]
+        acc = (recon.round() == real).all(dim=1).float().mean().item()
+        results[f"k={step+1}"] = acc
+    return results

connectx/verifier.py

@@ -0,0 +1,136 @@
+"""
+On-policy Monte Carlo value-head training -- no oracle, no bootstrapping,
+no target network. Walk the environment with the current value head's own
+epsilon-greedy policy, and label every visited state along a walk that
+actually reached a solved state with its REALIZED return (what happened,
+not a network's own possibly-wrong estimate of what happens next). This is
+the only value-training method used for ConnectX, since the real 7x6 board
+has no tractable exact solver to regress against.
+
+`unsolved_penalty` is what makes this work for an ADVERSARIAL domain
+specifically: `env.step` can return `done=True` on a LOSS (the opponent
+won) or a draw, not just our own win. Discarding every one of those walks
+(the natural default for a single-agent puzzle, where "unsolved" just means
+"further away") would mean the value head never sees a single labeled
+example of "this leads to losing" -- exactly the signal an adversarial
+domain needs to learn to avoid bad moves.
+"""
+import collections
+
+import torch
+import torch.nn as nn
+
+from .train_utils import states_to_tensor, DEVICE
+
+
+def _observed_states_to_tensor(env, states):
+    return states_to_tensor([env.observe(s) for s in states])
+
+
+def _greedy_walk_action(env, model, normalizer, state, rng, epsilon):
+    """Epsilon-greedy action choice using the value head's own 1-step
+    lookahead, scored the same way `search.plan_action`'s depth=1 does
+    (`-reward + value`, using the REAL reward/next-state from `env.step`,
+    not an imagined one) -- so training optimizes for the exact decision
+    rule actually used at inference."""
+    legal = [a for a in range(env.num_actions) if env.is_legal(state, a)]
+    if rng.random() < epsilon:
+        return rng.choice(legal)
+    next_states, rewards = [], []
+    for a in legal:
+        ns, r, _done = env.step(state, a)
+        next_states.append(ns)
+        rewards.append(r)
+    z = model.encode(normalizer.normalize(_observed_states_to_tensor(env, next_states).to(DEVICE)))
+    values = model.evaluate(z)
+    rewards_t = torch.tensor(rewards, dtype=torch.float32, device=DEVICE)
+    scores = -rewards_t + values
+    return legal[torch.argmin(scores).item()]
+
+
+def generate_mc_walks(env, model, normalizer, rng, n_problems, max_steps, epsilon,
+                       unsolved_penalty=None, **problem_kwargs):
+    """Complete walks (up to max_steps) using the current value head's
+    epsilon-greedy policy. A walk that reaches `is_solved` labels every
+    visited state with its real steps-remaining. A walk that doesn't
+    (a loss, a draw, or simply running out of steps) is discarded UNLESS
+    `unsolved_penalty` is set, in which case every state along it gets that
+    fixed, uniform label instead -- deliberately not scaled by how early or
+    late the walk went wrong; a state that's actually fine keeps
+    reappearing in OTHER (solved) walks too, so its label self-corrects
+    over many rounds rather than staying pinned to one bad walk's worst
+    case."""
+    labeled = []
+    for _ in range(n_problems):
+        state, _ = env.random_problem(rng, **problem_kwargs)
+        path_states = [state]
+        for _ in range(max_steps):
+            if env.is_solved(state):
+                break
+            a = _greedy_walk_action(env, model, normalizer, state, rng, epsilon)
+            state, _reward, done = env.step(state, a)
+            path_states.append(state)
+            if done:
+                break
+        if env.is_solved(path_states[-1]):
+            T = len(path_states) - 1
+            for t, s in enumerate(path_states):
+                labeled.append((s, float(T - t)))
+        elif unsolved_penalty is not None:
+            for s in path_states:
+                labeled.append((s, float(unsolved_penalty)))
+    return labeled
+
+
+def train_mc_value_onpolicy(env, model, normalizer, rng, n_rounds=15, n_problems_per_round=400,
+                             max_steps=10, epochs_per_round=40, lr=1e-3,
+                             epsilon_start=1.0, epsilon_end=0.05, warmup_rounds=3,
+                             replay_capacity=20000, min_replay_before_train=30,
+                             unsolved_penalty=None, verbose_every=5, **problem_kwargs):
+    """`epsilon_start=1.0` + `warmup_rounds`: a freshly-initialized value
+    head's "greedy" choice is pure noise, which can be WORSE than uniform
+    random at stumbling into a solved state by chance. Holding epsilon at
+    1.0 (pure random walk) for the first `warmup_rounds` guarantees a
+    baseline solve rate to bootstrap training from, before annealing toward
+    exploitation."""
+    opt = torch.optim.Adam(model.value.parameters(), lr=lr)
+    buffer = collections.deque(maxlen=replay_capacity)
+
+    for round_idx in range(n_rounds):
+        if round_idx < warmup_rounds:
+            epsilon = 1.0
+        else:
+            progress = (round_idx - warmup_rounds) / max(1, n_rounds - 1 - warmup_rounds)
+            epsilon = epsilon_start + (epsilon_end - epsilon_start) * progress
+        labeled = generate_mc_walks(env, model, normalizer, rng, n_problems_per_round, max_steps, epsilon,
+                                     unsolved_penalty=unsolved_penalty, **problem_kwargs)
+        buffer.extend(labeled)
+        if len(buffer) < min_replay_before_train:
+            if verbose_every:
+                print(f"  [mc-onpolicy] round {round_idx+1}/{n_rounds}  epsilon={epsilon:.2f}  "
+                      f"only {len(buffer)} labeled states so far (need {min_replay_before_train}) -- skipping fit")
+            continue
+
+        all_data = list(buffer)
+        states_t = _observed_states_to_tensor(env, [s for s, _ in all_data]).to(DEVICE)
+        returns_t = torch.tensor([r for _, r in all_data], dtype=torch.float32, device=DEVICE)
+
+        model.value_target_mean.copy_(returns_t.mean())
+        model.value_target_std.copy_(returns_t.std().clamp(min=1e-3))
+        returns_norm = (returns_t - model.value_target_mean) / model.value_target_std
+
+        with torch.no_grad():
+            z_states = model.encode(normalizer.normalize(states_t))
+
+        n = len(all_data)
+        for _epoch in range(epochs_per_round):
+            perm = torch.randperm(n, device=DEVICE)
+            pred_norm = model.value(z_states[perm])
+            loss = nn.functional.mse_loss(pred_norm, returns_norm[perm])
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+
+        if verbose_every and (round_idx + 1) % verbose_every == 0:
+            print(f"  [mc-onpolicy] round {round_idx+1}/{n_rounds}  epsilon={epsilon:.2f}  "
+                  f"buffer_size={n}  loss={loss.item():.4f}  return_mean={model.value_target_mean.item():.2f}")

docs/build_whitepaper.py

@@ -0,0 +1,336 @@
+# One-off build script for WHITEPAPER.pdf -- not part of the codebase
+# itself, run once to (re)generate the PDF from this content. Requires
+# `pip install fpdf2`. Run from the docs/ directory: `python build_whitepaper.py`.
+from fpdf import FPDF
+
+TITLE = "Latent-Space Simulation Meets a Real Adversary: A World-Model Approach to Kaggle ConnectX"
+AUTHOR = "Alexandros Titonis"
+
+pdf = FPDF(format="A4")
+pdf.set_auto_page_break(auto=True, margin=20)
+pdf.set_margins(20, 20, 20)
+pdf.add_page()
+pdf.set_font("Helvetica", "B", 18)
+pdf.multi_cell(0, 9, TITLE)
+pdf.ln(2)
+pdf.set_font("Helvetica", "", 11)
+pdf.set_text_color(90, 90, 90)
+pdf.multi_cell(0, 6, "A reinforcement-learning world model, real adversarial search, and an exact "
+                      "endgame solver applied to Kaggle's ConnectX competition.")
+pdf.ln(1)
+pdf.multi_cell(0, 6, f"{AUTHOR}  |  Competition: https://kaggle.com/competitions/connectx")
+pdf.set_text_color(0, 0, 0)
+pdf.ln(4)
+
+
+def h1(text):
+    pdf.set_font("Helvetica", "B", 14)
+    pdf.ln(4)
+    pdf.multi_cell(0, 8, text)
+    pdf.ln(1)
+
+
+def h2(text):
+    pdf.set_font("Helvetica", "B", 12)
+    pdf.ln(2)
+    pdf.multi_cell(0, 7, text)
+    pdf.ln(1)
+
+
+def p(text):
+    pdf.set_font("Helvetica", "", 10.5)
+    pdf.multi_cell(0, 5.6, text)
+    pdf.ln(1.5)
+
+
+def bullet(text):
+    pdf.set_font("Helvetica", "", 10.5)
+    pdf.set_x(pdf.l_margin + 4)
+    pdf.multi_cell(0, 5.6, "-  " + text)
+    pdf.ln(0.5)
+
+
+def table(headers, rows, widths, body_size=9):
+    """Proper multi-line table: wraps each cell's text to its column width
+    FIRST (via a dry-run multi_cell), takes the max line count in the row,
+    then draws every cell's border/text at that shared row height -- avoids
+    fpdf2's plain `cell()` silently misaligning columns when text embeds a
+    newline or wraps."""
+    line_h = 5.0
+    pdf.set_font("Helvetica", "B", 9.5)
+    header_h = 7.0
+    x_start, y_start = pdf.get_x(), pdf.get_y()
+    x = x_start
+    for htext, w in zip(headers, widths):
+        pdf.set_xy(x, y_start)
+        pdf.cell(w, header_h, htext, border=1)
+        x += w
+    pdf.set_xy(x_start, y_start + header_h)
+
+    pdf.set_font("Helvetica", "", body_size)
+    for row in rows:
+        cell_lines = []
+        for text, w in zip(row, widths):
+            lines = pdf.multi_cell(w - 2, line_h, text, dry_run=True, output="LINES")
+            cell_lines.append(lines)
+        n_lines = max(len(lns) for lns in cell_lines)
+        row_h = n_lines * line_h
+        x_start, y_start = pdf.get_x(), pdf.get_y()
+        if y_start + row_h > pdf.page_break_trigger:
+            pdf.add_page()
+            x_start, y_start = pdf.get_x(), pdf.get_y()
+        x = x_start
+        for lines, w in zip(cell_lines, widths):
+            pdf.rect(x, y_start, w, row_h)
+            for i, line in enumerate(lines):
+                pdf.set_xy(x + 1, y_start + i * line_h)
+                pdf.cell(w - 2, line_h, line)
+            x += w
+        pdf.set_xy(x_start, y_start + row_h)
+    pdf.ln(4)
+
+
+# ---------------------------------------------------------------------------
+h1("1. What this is, and what it isn't")
+p("The core idea: instead of generating a move token-by-token or via a hand-written heuristic, "
+  "a small learned model (encoder + latent dynamics + value head + a diagnostic decoder) rolls "
+  "forward candidate moves in latent space, and a search procedure picks the move whose imagined "
+  "future looks best -- MuZero-style planning [1].")
+p("This is one instance of a GENERAL latent-space-simulation architecture, not a Connect-4-specific "
+  "system: the same encoder/dynamics/value/decoder + search pattern used here has also been applied "
+  "to algebraic equation solving, constraint-satisfaction-style logic puzzles, and route-planning "
+  "problems, with zero domain-specific changes to the core training/search code -- swapping domains "
+  "means implementing one small interface (see `environment.py`), not rewriting the architecture. "
+  "ConnectX is the instance documented in this paper because it is the first genuinely ADVERSARIAL "
+  "two-player domain the architecture was applied to: a single-agent puzzle solver only ever has to "
+  "be right, while a Connect-4 agent has to be right against an opponent actively trying to make it "
+  "wrong. That one change breaks assumptions the rest of the pipeline had never been asked to "
+  "question, and is the actual subject of this write-up.")
+p("This is NOT a hand-written Connect-4 bot. The board simulator and rules are exact (no need to "
+  "make a neural network learn something 40 lines of Python computes for free), but every DECISION "
+  "-- what a position is worth, which past games are relevant, when the value head should update "
+  "online -- comes from a learned model trained without a Connect-4-specific oracle.")
+
+h2("1.1 How this compares to typical RL game-playing systems")
+p("Systems in the AlphaZero/MuZero lineage [1][2] are normally trained at large scale: thousands of "
+  "self-play games generated in parallel across many machines (often TPU or GPU clusters), running "
+  "for days. This project's result was trained on a single consumer laptop -- see Section 1.2 for "
+  "the exact specs -- with a full training-pipeline run (stage 1 + stage 2 + self-play) confirmed at "
+  "roughly one hour of wall-clock time, and on the order of tens of thousands of self-play games, not "
+  "millions. This is presented honestly as evidence that the ARCHITECTURE (real search over exactly-"
+  "known rules, substituting for an expensive learned-world-model rollout) is a genuinely efficient "
+  "way to reach competent play in a domain with known rules -- not a claim that this matches "
+  "AlphaZero-level strength; Section 6 states plainly where this still falls short of a fully-solved "
+  "or superhuman opponent.")
+
+h2("1.2 Training hardware")
+p("Every result in this paper was produced on a single machine: an RTX 4060 laptop GPU, a Ryzen AI "
+  "9 HX 370 CPU, and 32GB RAM. No cluster, no cloud training run, no multi-GPU parallelism -- "
+  "everything in this repository, including a full from-scratch training run, was reproduced on that "
+  "one laptop.")
+
+h1("2. Architecture")
+
+h2("2.1 The easy path: collapse the adversary into an ordinary MDP")
+p("The cheapest way to test whether this works at all: bake a fixed, non-learning opponent policy "
+  "directly into the environment's step() function. From the agent's perspective this is then "
+  "indistinguishable from a single-agent puzzle -- it picks an action, the environment returns the "
+  "resulting state (already reflecting the opponent's reply). Zero changes needed to the rest of "
+  "the pipeline. At the real 7x6 board scale there is no exact oracle (the game tree is too large "
+  "for brute-force search), so the value head is trained via on-policy Monte Carlo returns instead "
+  "of regression against oracle labels.")
+
+h2("2.2 Why the dynamics model can't imagine \"after my move, before theirs\"")
+p("step() bundles the agent's move and the fixed opponent's reply into ONE transition, so the "
+  "trained dynamics model only ever saw full round-trips as single training examples -- it "
+  "structurally cannot represent \"the board immediately after my move, before their reply\" as a "
+  "state, because it never saw that state shape. Asking it to imagine that intermediate state would "
+  "be extrapolating outside its training distribution.")
+
+h2("2.3 The fix: real search over the real board, learned value only at the leaf")
+p("Since Connect-4's rules are exactly known, there's no need to make the network imagine anything "
+  "mid-ply. The deployed search does the adversarial ply in REAL board space: enumerate the agent's "
+  "real legal moves; for each, enumerate the opponent's real legal replies and assume they pick "
+  "whichever hurts the agent most (a genuine minimax, not a single guess). The learned value head is "
+  "used only as the LEAF evaluator, on a real, never-imagined state. This is structurally closer to "
+  "AlphaZero's design (real search + a learned value net, viable because the rules are exactly "
+  "known) [2] than MuZero's (search over a LEARNED model, needed specifically when true dynamics "
+  "are not available) [1].")
+p("Result of this one architectural change: 95-100% win rate vs. random play, 100% vs. the fixed "
+  "weak training heuristic, up to 66.7% vs. a stronger 1-ply-deeper heuristic never seen during "
+  "training -- beating every latent-search configuration tried on every metric simultaneously (the "
+  "best latent-search result had been 75% / 50% / 11.7%).")
+
+h2("2.4 The exact endgame solver")
+p("Real search over real board states opens a door latent-space search never could: once a position "
+  "narrows down to a handful of legal columns -- which happens naturally as the board fills -- the "
+  "remaining game tree is small regardless of how many plies are left, and can be SOLVED EXACTLY, "
+  "no learned value head, no guessing. The solver is a memoized alpha-beta minimax with center-out "
+  "move ordering, triggered whenever 5 or fewer columns remain legal (calibrated empirically against "
+  "real Kaggle replay data: 5 columns / 24 empty cells solves in under half a second; 6 columns did "
+  "not finish inside a 5-second cap with this plain-Python implementation). It is bounded by a hard "
+  "wall-clock deadline as an independent safety net -- a position outside the calibrated safe zone "
+  "falls through to the round-based search rather than risking the real move-time budget.")
+p("The insight that made this cheap: branching factor is controlled by how many columns are LEGAL, "
+  "not by how many cells are EMPTY. A position can have twenty or more empty cells and still be "
+  "trivial to solve exactly if only a few columns remain open -- exactly the shape a real Connect-4 "
+  "endgame takes once most of the board has filled.")
+
+h2("2.5 Episodic memory, LoRA fine-tuning, and online learning")
+p("Three more general-purpose mechanisms were applied on top of the search:")
+bullet("Episodic memory: a k-NN lookup over real self-play trajectories, both won and lost (a "
+       "negative/repulsion signal, not just positive examples), blended into the value estimate at "
+       "decision time [3][4]. Small but real positive effect vs. random play; no measurable effect "
+       "vs. either fixed heuristic opponent.")
+bullet("LoRA self-play fine-tuning: capacity-constrained low-rank deltas applied to the value head "
+       "during self-play [5], rather than an unconstrained further gradient update -- 75.0% to "
+       "81.7% win rate vs. the stronger heuristic.")
+bullet("Curriculum self-play: fixing a real gap where the stronger heuristic was never actually "
+       "mixed into self-play training at all -- 81.7% to 85.0%, the best-confirmed result for this "
+       "domain.")
+bullet("Best-effort online learning (value head only, updated live as real games are played): "
+       "isolated via a clean A/B test (the same submission loaded twice, one copy's update call "
+       "disabled) -- currently measures exactly 0.0% effect against the synthetic test-opponent set "
+       "at evaluation scale. Not evidence it is useless in general -- Kaggle's real opponent pool is "
+       "far wider than three fixed synthetic opponents -- just an honest, now-measured answer.")
+
+h1("3. Every bug found and fixed")
+p("Seven real, generic-class bugs were found and fixed during development. Each is recorded here "
+  "because several are the kind of mistake that is easy to make again in a different adversarial "
+  "domain.")
+table(
+    ["#", "Bug", "Found via"],
+    [
+        ["1", "Internal eval reported a false 100% win rate (done was assumed to imply a win)",
+         "Tracing one \"solved\" game through an independent alternating-turn harness"],
+        ["2", "The value head never saw a losing example (unsolved walks were silently discarded)",
+         "Root-cause tracing after bug 1"],
+        ["3", "Every prediction in the packaged submission was silently corrupted (a spurious extra ReLU)",
+         "Diffing the packaged output directly against the real model, node by node"],
+        ["4", "\"Same seed\" training runs were not actually reproducible (unseeded global RNG)",
+         "Two same-seed runs produced different results"],
+        ["5", "A whole session's A/B comparisons were confounded (one shared RNG across a batch)",
+         "An opening-move-only change swung the overall win rate by double digits"],
+        ["6", "A left-column tie-break bias with no game-theoretic basis",
+         "Direct user-observed pattern in real play"],
+        ["7", "The online learner taught the value head that losing and drawing are equally bad",
+         "Mining real Kaggle replay data via the Kaggle API for missed blocks"],
+    ],
+    [10, 88, 72],
+)
+p("The generalizable lesson: bugs 2 and 7 are the same mistake, made twice, in two different places. "
+  "\"Treat a loss and a draw identically\" is a natural default everywhere ELSE in a single-agent "
+  "pipeline (every prior puzzle domain has exactly one failure mode -- unsolved -- with no "
+  "distinction worth making), and it takes deliberate effort to remember that an adversarial domain "
+  "has two qualitatively different bad outcomes.")
+
+h1("4. Case study: diagnosing a real loss (zugzwang), and a first attempt at a fix")
+p("A live, observed loss -- an opponent slowly building what looked, in hindsight, like an "
+  "obviously winning diagonal, with the agent apparently ignoring it -- was root-caused not by "
+  "guesswork but by pulling the actual Kaggle replay and solving the real endgame exhaustively, "
+  "working backward through the game.")
+p("Finding: the position was already a forced loss with 23 empty cells still on the board -- at "
+  "least 15 real plies before the deployed 2-round search could possibly have seen it coming. The "
+  "mechanism is a genuine, textbook Connect-4 ZUGZWANG / PARITY TRAP (classic odd/even threat "
+  "theory [6]): once only two columns remained legal for an extended stretch, which player is "
+  "forced to place the fatal piece is decided purely by the parity of total remaining cells across "
+  "those two columns -- invisible to any bounded-depth positional search, and not addressed by the "
+  "exact endgame solver above (its own calibrated safe zone starts well inside the window where "
+  "this trap was already unavoidable).")
+
+h2("4.1 A parity-heuristic attempt -- tried, measured, and honestly a mixed result")
+p("A column-parity feature was built and tested directly, rather than left purely theoretical: for "
+  "each still-open column with r empty cells, under naive same-column-only alternation the player "
+  "who would place the TOP piece is the mover if r is odd, the other player if r is even. This gives "
+  "a cheap, real-valued \"parity cost\" per board position -- more even-parity open columns scored as "
+  "worse for the agent -- blended into the leaf value estimate at a tunable weight.")
+p("This is explicitly NOT full Claimeven: a genuine claimeven strategy requires REACTIVE move-"
+  "pairing enforced across an entire game, not a one-shot column count at a single position. It was "
+  "built and measured as a testable nudge, not assumed to work because the theory motivates it.")
+p("Measured against the trusted test harness at several weights: the feature produces a REAL, "
+  "repeatable effect, but not a clean win. At weight 0.5, win rate vs. the stronger heuristic rose "
+  "from 70.0% to 77.5% -- but win rate vs. random-legal play fell from 100.0% to 90.0% in the same "
+  "run. Smaller weights (0.1-0.3) showed the same pattern at reduced magnitude in both directions. "
+  "Net honest verdict: this specific heuristic trades performance against weaker opponents for a "
+  "partial, not fully consistent gain against a stronger one -- a real, measured signal in the "
+  "direction the theory predicts, but not yet a confirmed fix. It ships in this repository as an "
+  "opt-in, OFF-by-default parameter (`parity_weight=0.0` in `adversarial_search.py`) rather than a "
+  "new default, exactly so a reader can reproduce this exact finding rather than take it on faith.")
+p("Two angles remain genuinely untried: a value head deliberately trained on zugzwang-rich self-play "
+  "positions (so the pattern becomes an implicit learned feature rather than a hand-tuned scalar "
+  "nudge), and a full, correctly reactive Claimeven implementation rather than a single-position "
+  "heuristic proxy.")
+
+h1("5. Results")
+p("All win rates below are against an independent, trusted test harness (an alternating-turn engine "
+  "that does NOT reuse the environment's own bundled step() -- exactly the shortcut a submission "
+  "validator needs to avoid). Three opponents throughout: random-legal play, the fixed weak "
+  "heuristic the model trained against, and a 1-ply-deeper \"stronger\" heuristic never seen during "
+  "training.")
+table(
+    ["Configuration", "vs random", "vs weak heur.", "vs stronger heur."],
+    [
+        ["Latent beam search, depth=1 (original)", "63.3%", "0.0%", "6.7%"],
+        ["Latent beam search, depth=3", "55.0%", "0.0%", "1.7%"],
+        ["Real adversarial search (1 round)", "96.7%", "100.0%", "66.7%"],
+        ["+ memory + online learner (bug 7 present)", "98.3%", "78.3%*", "50.0%"],
+        ["+ bug 7 (loss-penalty) fixed", "100.0%", "100.0%", "80.0%"],
+        ["+ LoRA self-play fine-tune", "-", "-", "81.7%"],
+        ["+ curriculum self-play (best-ever)", "-", "-", "85.0%"],
+        ["+ exact endgame solver (current, deployed)", "100.0%", "100.0%", "83.3%"],
+    ],
+    [95, 30, 32, 33],
+)
+p("* The weak-heuristic dip in that one row is explained, not a mystery: that harness ran all "
+  "opponent blocks sequentially in one process, so the online learner had already drifted from 60 "
+  "preceding random-opponent games by the time it reached this block.")
+p("On Kaggle's own rating (a TrueSkill-style Gaussian score that starts uncertain and converges "
+  "over dozens of real games against other real submissions): a freshly-uploaded submission's "
+  "rating is not comparable to one that has had a day to settle. Confirmed directly during "
+  "development -- a previously-deployed submission itself started near 300 right after upload and "
+  "climbed past 450 only after 27 real games. 7x6 Connect-4 is a mathematically SOLVED game (the "
+  "first player wins with perfect play) [6], so the real competitive pool likely includes near-"
+  "perfect solvers -- the results above demonstrate the ARCHITECTURE works on this domain; they "
+  "should not be read as a leaderboard-rating prediction.")
+
+h1("6. Honest limitations")
+bullet("The zugzwang-avoidance problem (Section 4) is diagnosed and a first fix attempt is measured, "
+       "but not solved -- see Section 4.1's own honest verdict.")
+bullet("Online learning's real-world value is unconfirmed -- it measures zero effect against three "
+       "fixed synthetic opponents, which is not the same as zero effect against Kaggle's actual, "
+       "much wider real pool.")
+bullet("Deeper search beyond 2 rounds (without the endgame solver) was tried and found WORSE, not "
+       "just unexplored -- unpruned 3-round search blows the time budget, and the only pruning "
+       "width that stayed safe tested worse than plain 2-round search.")
+bullet("The real Kaggle competitive pool likely includes near-perfect solvers, since 7x6 Connect-4 "
+       "is a solved game -- see Section 5's closing note.")
+
+h1("References")
+refs = [
+    "[1] Schrittwieser, J. et al. \"Mastering Atari, Go, Chess and Shogi by Planning with a Learned "
+    "Model.\" Nature, 2020.",
+    "[2] Silver, D. et al. \"A general reinforcement learning algorithm that masters chess, shogi, "
+    "and Go through self-play.\" Science, 2018.",
+    "[3] Blundell, C. et al. \"Model-Free Episodic Control.\" arXiv:1606.04460, 2016.",
+    "[4] Pritzel, A. et al. \"Neural Episodic Control.\" ICML, 2017.",
+    "[5] Hu, E. J. et al. \"LoRA: Low-Rank Adaptation of Large Language Models.\" "
+    "arXiv:2106.09685, 2021.",
+    "[6] Allis, L. V. \"A Knowledge-Based Approach of Connect-Four.\" M.Sc. thesis, Vrije "
+    "Universiteit Amsterdam, 1988. (First complete game-theoretic solution of Connect-4; the "
+    "classical source for odd/even threat and zugzwang theory referenced in Section 4.)",
+    "[7] Mnih, V. et al. \"Human-level control through deep reinforcement learning.\" Nature, 2015. "
+    "(Experience-replay precedent used by the on-policy Monte Carlo value-training loop.)",
+    "[8] Adam, Addison Howard, and Bovard Doerschuk-Tiberi. \"Connect X.\" "
+    "https://kaggle.com/competitions/connectx, 2020. Kaggle.",
+    "[9] Kaggle. \"kaggle_environments.\" https://github.com/Kaggle/kaggle-environments. "
+    "(Board/config schema and fixed-opponent-baked-into-step convention are ported from this "
+    "package's own design, not from any published agent's code.)",
+]
+pdf.set_font("Helvetica", "", 9.5)
+for r in refs:
+    pdf.set_x(pdf.l_margin)
+    pdf.multi_cell(0, 5.2, r)
+    pdf.ln(1.5)
+
+pdf.output("WHITEPAPER.pdf")
+print("Wrote WHITEPAPER.pdf")

scripts/build_submission.py

@@ -0,0 +1,843 @@
+"""
+Packages connectx_checkpoint.pt into a SINGLE self-contained submission.py
+Kaggle can actually run -- REWRITTEN 2026-08-10, per explicit user
+direction ("let's combine all these in the submission file with
+memories, the simulations, the weak learner etc") to fold in THREE
+confirmed pieces from this session's follow-up work, replacing the
+previous latent-beam-search-only version entirely:
+
+1. **Real adversarial search** (was: latent beam search). Per
+   connectx_adversarial_search.py's confirmed result (95-100%/48.3% vs
+   random/weak/stronger, beating latent search on every metric by a wide
+   margin, confirmed on 2 checkpoints): `env.step()` bundles agent+
+   opponent-reply into one transition, so the trained dynamics model was
+   never shown "the board right after my move, before their reply" --
+   it structurally can't imagine that state. Since ConnectX's rules ARE
+   exactly known, this ply is done in REAL board space instead (exact
+   enumeration of our moves, exact enumeration of the opponent's
+   worst-case real reply), with the learned value head used ONLY as the
+   leaf evaluator. This means `dynamics`/`decoder` are NO LONGER NEEDED
+   at all (the old latent search's neurosymbolic decode-gate is
+   structurally unnecessary once every ply is real, not imagined) --
+   only `encoder`+`value` weights are embedded now, a smaller submission.
+2. **Episodic memory (positive + negative)**, built OFFLINE (this
+   script, at build time) from self-play games against a MIXED opponent
+   (weak heuristic + random + the stronger 1-ply-deeper heuristic, per
+   explicit user caution -- "so if the opponent is weak it doesn't learn
+   the bad ways too" -- see memory_build.py). Won games
+   stored as positive (low remaining-steps) examples, lost/drawn games
+   as negative (high, fixed-penalty) examples -- one EpisodicMemory,
+   blended into every leaf evaluation via the exact same k-NN
+   inverse-distance/trust-scaled formula as episodic_memory.py's
+   `query_batch`, replicated here in plain torch (no project import,
+   this file must stay standalone).
+3. **Best-effort online learning ("the weak learner")** -- value-head-
+   ONLY updates (matching this session's own confirmed finding: decoder
+   updates regressed structured-opponent performance at this data scale,
+   so the decoder is excluded entirely here, consistent with "the
+   working side only"), applied incrementally as real games are played,
+   mirroring continuous_learner.py's confirmed-safe recipe (small
+   replay buffer, EMA-updated value_target_mean/std, a few Adam steps
+   per update, lr=1e-5) -- reimplemented here in plain torch since this
+   file can't import continuous_learner.py.
+
+   **Honest, load-bearing caveat, stated plainly rather than oversold**:
+   Kaggle's `agent(observation, configuration)` interface gives no
+   direct "episode ended, here's the result" callback -- this file
+   infers a completed episode two ways, both using ONLY information
+   actually available across calls: (a) our own move immediately wins
+   or draws (directly observable -- we know the board we just produced),
+   or (b) the NEXT call arrives with a completely empty board while a
+   previous episode's trajectory is still buffered -- inferred as a LOSS
+   (we didn't win/draw it ourselves, so it must have ended on the
+   opponent's move). This whole mechanism is a NO-OP, gracefully, unless
+   Kaggle's real evaluation infrastructure reuses the same process across
+   multiple episodes for this submission over time (its own rules page,
+   read earlier this session, doesn't confirm or deny this -- see
+   [[project_connectx_kaggle]]) -- if each episode gets a fresh process,
+   this buffer simply starts empty every time and nothing is lost, no
+   crash, no wasted budget beyond one negligible check.
+"""
+import base64
+import io
+
+import torch
+
+CKPT_PATH = "checkpoints/connectx_checkpoint.pt"
+OUT_PATH = "submission.py"
+
+
+def _encode_tensor_blob(ck, memory_zs, memory_outcomes):
+    """encoder+value weights only (see module docstring -- dynamics/
+    decoder are no longer needed by the real adversarial search), plus
+    value_target_mean/std (top-level buffers, not nested under a
+    submodule prefix) and the offline-built episodic memory's raw
+    (z, remaining_steps) pairs."""
+    keep = {k: v for k, v in ck["model_state"].items()
+            if k.startswith("encoder.") or k.startswith("value.")
+            or k in ("value_target_mean", "value_target_std")}
+    payload = {
+        "weights": keep,
+        "norm_mean": ck["norm_mean"],
+        "norm_std": ck["norm_std"],
+        "state_dim": ck["state_dim"],
+        "num_actions": ck["num_actions"],
+        "latent_dim": ck["latent_dim"],
+        "hidden_dim": ck["hidden_dim"],
+        "board_width": ck["board_width"],
+        "board_height": ck["board_height"],
+        "win_len": ck["win_len"],
+        "memory_zs": torch.stack(memory_zs) if memory_zs else torch.zeros(0, ck["latent_dim"]),
+        "memory_outcomes": torch.tensor(memory_outcomes, dtype=torch.float32),
+    }
+    buf = io.BytesIO()
+    torch.save(payload, buf)
+    return base64.b64encode(buf.getvalue()).decode("ascii")
+
+
+SUBMISSION_TEMPLATE = '''\
+"""
+Auto-generated by build_submission.py -- DO NOT hand-edit
+(regenerate instead). Self-contained Kaggle ConnectX submission: no
+imports beyond torch/base64/io, so it runs standalone in Kaggle's
+evaluation sandbox.
+
+Policy: ONE ROUND of REAL adversarial search (exact enumeration of our
+legal moves, exact enumeration of the opponent's real legal replies,
+worst-case-for-us selected -- a genuine minimax over EXACTLY KNOWN board
+dynamics, not an imagined latent transition) -- the learned value head
+is used ONLY as the leaf evaluator on a real, never-imagined state,
+optionally blended with an offline-built episodic memory (won AND lost
+self-play games, see module docstring). A best-effort online value-head
+update also runs across real games as they're played -- see module
+docstring's honest caveat about when this can/can't actually do
+anything, given Kaggle's evaluation interface.
+
+**Honest, named limitation** (see connectx_env.py / [[project_connectx_kaggle]]):
+the base checkpoint was trained via self-play against a small set of
+fixed/self-generated opponents, not against Kaggle's real matchmaking
+pool -- see that project's memory entry for the full picture, including
+this session's confirmed numbers against synthetic test opponents.
+"""
+import base64
+import collections
+import io
+import time
+
+import torch
+
+_MEMORY_WEIGHT = {memory_weight}
+_MEMORY_K = {memory_k}
+_ONLINE_LR = {online_lr}
+_ONLINE_UPDATES_PER_EPISODE = {online_updates_per_episode}
+_ONLINE_BATCH_SIZE = {online_batch_size}
+_UNSOLVED_PENALTY_MULT = {unsolved_penalty_mult}  # x max_steps, matches this session's convention
+_ADV_ROUNDS = {adv_rounds}  # real adversarial search rounds -- see _adversarial_plan_action's docstring for timing
+# `_ENDGAME_MAX_COLS`/`_ENDGAME_TIME_BUDGET` (added 2026-08-11): below
+# this many legal columns, `_exact_endgame_solve` (a real, no-NN,
+# alpha-beta minimax to the true end of the game) is tried FIRST and used
+# directly if it finishes in time -- see that function's own docstring
+# for the calibration and the exact failure mode (a zugzwang/parity trap
+# invisible to any bounded-depth search) this targets. `_ENDGAME_MAX_COLS
+# = 0` disables this path entirely.
+_ENDGAME_MAX_COLS = {endgame_max_cols}
+_ENDGAME_TIME_BUDGET = {endgame_time_budget}
+# `_DEEPER_ROUNDS`/`_DEEPER_MAX_BRANCHING`/`_DEEPER_TIME_BUDGET`: real,
+# mined-from-real-games evidence showed `_ADV_ROUNDS` sometimes sees ZERO
+# danger on a position (every column looks equally safe) 2-4 plies before
+# a trap that one round DEEPER already narrows down to exactly one safe
+# column -- `_ADV_ROUNDS` isn't wrong about what it can see, it just can't
+# see far enough to avoid a fork the opponent is setting up. A deeper
+# search is provably too slow to run on EVERY move (measured 8-13s at a
+# 6-7-legal-column branching factor) -- so this is a SAFE, opportunistic
+# escalation, not a blanket depth increase: `_DEEPER_ROUNDS = None`
+# disables it entirely, reproducing the original `_ADV_ROUNDS`-only
+# behavior byte-for-byte. When enabled, AFTER computing the normal-
+# `_ADV_ROUNDS` answer (always -- the guaranteed-safe fallback), a
+# `_DEEPER_ROUNDS`-round search is attempted under a hard
+# `_DEEPER_TIME_BUDGET` deadline; if it finishes in time its answer is
+# used instead (strictly more information, never less), if it times out
+# the original answer is returned completely unchanged. Calibrated via a
+# 180-game regression suite (random/weak/stronger opponents): zero
+# win-rate regression, max observed single-move time 1.641s --
+# comfortably under Kaggle's 2s budget.
+_DEEPER_ROUNDS = {deeper_rounds}
+_DEEPER_MAX_BRANCHING = {deeper_max_branching}
+_DEEPER_TIME_BUDGET = {deeper_time_budget}
+# `_ONLINE_ENABLED` (added 2026-08-10, right before submitting -- explicit
+# user decision after reading the competition's own rule "An Agent's sole
+# purpose is to generate an action. Activities/code which do not directly
+# contribute to this will be considered malicious...": the online "weak
+# learner"'s gradient updates are arguably in service of generating BETTER
+# actions, not unrelated activity, but it's a genuine judgment call with
+# real (if likely small) risk, not a zero-risk one -- played safe rather
+# than assume it's fine. False disables it CLEANLY (no buffer/episode-
+# tracking side-state at all when off, not just a no-op update call) so
+# a disabled submission's `agent()` genuinely does nothing but generate
+# an action, matching the rule as literally as possible.
+_ONLINE_ENABLED = {online_enabled}
+
+_BLOB_B64 = (
+{blob_literal}
+)
+
+
+def _load():
+    payload = torch.load(io.BytesIO(base64.b64decode(_BLOB_B64)), map_location="cpu")
+    return payload
+
+
+_P = _load()
+_W = _P["weights"]
+_NORM_MEAN = _P["norm_mean"]
+_NORM_STD = _P["norm_std"]
+_LATENT_DIM = _P["latent_dim"]
+_NUM_ACTIONS = _P["num_actions"]  # includes the training-time PASS action (index WIDTH)
+_WIDTH = _P["board_width"]
+_HEIGHT = _P["board_height"]
+_WIN_LEN = _P["win_len"]
+_PASS_ACTION = _WIDTH
+_CELL_WIDTH = 3
+_EMPTY, _AGENT, _OPPONENT = 0, 1, 2
+_MAX_STEPS = (_WIDTH * _HEIGHT) // 2 + 2
+# `_UNSOLVED_PENALTY` (used ONLY by the online learner's episode-ending
+# label, matching continuous_learner.py's own 1x-max_steps convention)
+# and `_LOSS_PENALTY` (used ONLY by the adversarial search's "opponent
+# wins" terminal case) are DELIBERATELY SEPARATE constants -- a real bug
+# found and fixed 2026-08-10, right after this build was already live:
+# an earlier version used _UNSOLVED_PENALTY (1x max_steps) for BOTH,
+# which meant the search scored "the opponent wins outright" EXACTLY
+# THE SAME as "it's a mere draw" -- losing must be unambiguously worse
+# than a draw for the search to reliably prioritize blocking a real
+# threat over a merely-mediocre move, matching connectx_adversarial_search.py's
+# original, correct 2x convention. Confirmed as the direct, mechanistic
+# cause of a real observed failure: the deployed agent missed blocking
+# an opponent's obvious 3-in-a-column vertical threat, scoring the
+# blocking move WORSE (23.463) than a non-blocking move that let the
+# opponent win outright (23.000, since the loss was scored at only
+# max_steps=23, indistinguishable from ordinary mediocre play).
+_UNSOLVED_PENALTY = _UNSOLVED_PENALTY_MULT * _MAX_STEPS
+_LOSS_PENALTY = 2 * _MAX_STEPS
+
+# Memory tensors (offline-built, see module docstring) -- fixed, never
+# grow at runtime (only the ONLINE value-head buffer below does).
+_MEMORY_Z = _P["memory_zs"]
+_MEMORY_OUTCOMES = _P["memory_outcomes"]
+if _MEMORY_Z.shape[0] >= 2:
+    _d = torch.cdist(_MEMORY_Z, _MEMORY_Z)
+    _d = torch.where(_d > 1e-6, _d, torch.full_like(_d, float("inf")))
+    _nn = _d.min(dim=1).values
+    _nn = _nn[torch.isfinite(_nn)]
+    _MEMORY_TRUST_SCALE = _nn.median().item() if len(_nn) > 0 else 1.0
+else:
+    _MEMORY_TRUST_SCALE = 1.0
+
+# --- Value head params made trainable for the online "weak learner"
+# (see module docstring's honest caveat) -- encoder stays FROZEN
+# (never in this optimizer), matching continuous_learner.py's confirmed
+# recipe: only the value head updates online. When `_ONLINE_ENABLED` is
+# False, NONE of this setup happens at all (no optimizer, no
+# requires_grad, no buffers) -- `agent()` genuinely does nothing but
+# generate an action in that case, not just a disabled-but-present
+# mechanism. ---
+if _ONLINE_ENABLED:
+    _VALUE_PARAM_KEYS = [k for k in _W if k.startswith("value.")]
+    for _k in _VALUE_PARAM_KEYS:
+        _W[_k].requires_grad_(True)
+    # Buffers, not trained parameters (EMA-updated in-place under
+    # no_grad, matching continuous_learner.py's own convention) -- never
+    # added to the optimizer below.
+    _VALUE_TARGET_MEAN = _W.get("value_target_mean", torch.tensor(0.0)).clone()
+    _VALUE_TARGET_STD = _W.get("value_target_std", torch.tensor(1.0)).clone()
+    _ONLINE_OPT = torch.optim.Adam([_W[k] for k in _VALUE_PARAM_KEYS], lr=_ONLINE_LR)
+    _REPLAY_BUFFER = collections.deque(maxlen=2000)  # (state_vec: list[float], label: float)
+    _EPISODE_STATES = []  # real one-hot state vectors seen/produced so far THIS episode
+    _EPISODE_LAST_PIECES = None  # total board piece count as of our last recorded state THIS episode
+else:
+    _VALUE_TARGET_MEAN = _W.get("value_target_mean", torch.tensor(0.0))
+    _VALUE_TARGET_STD = _W.get("value_target_std", torch.tensor(1.0))
+
+
+def _linear(x, w_key, b_key):
+    return torch.nn.functional.linear(x, _W[w_key], _W[b_key])
+
+
+def _mlp3(x, prefix):
+    """Replicates model.py's `mlp([in, hidden, hidden, out])`: Linear ->
+    ReLU -> Linear -> ReLU -> Linear (params at Sequential indices
+    0/2/4, confirmed against the actual saved state_dict keys)."""
+    h = torch.relu(_linear(x, f"{{prefix}}.net.0.weight", f"{{prefix}}.net.0.bias"))
+    h = torch.relu(_linear(h, f"{{prefix}}.net.2.weight", f"{{prefix}}.net.2.bias"))
+    return _linear(h, f"{{prefix}}.net.4.weight", f"{{prefix}}.net.4.bias")
+
+
+def _encode(state_vec):
+    return _mlp3(state_vec, "encoder")
+
+
+def _value_raw(z):
+    return _mlp3(z, "value").squeeze(-1)
+
+
+def _value(z):
+    """Real-scale value estimate (remaining steps), see model.py's
+    WorldModel.evaluate -- denormalizes the network's raw prediction."""
+    return _value_raw(z) * _VALUE_TARGET_STD + _VALUE_TARGET_MEAN
+
+
+def _memory_blend(z_batch, raw_values):
+    """Same k-NN inverse-distance/trust-scaled blend as
+    episodic_memory.py's EpisodicMemory.query_batch -- replicated here
+    in plain torch (this file can't import that module)."""
+    if _MEMORY_Z.shape[0] == 0 or _MEMORY_WEIGHT <= 0:
+        return raw_values
+    dists = torch.cdist(z_batch, _MEMORY_Z)  # [B, N]
+    k = min(_MEMORY_K, _MEMORY_Z.shape[0])
+    topk_dists, topk_idx = torch.topk(dists, k, largest=False, dim=1)
+    topk_outcomes = _MEMORY_OUTCOMES[topk_idx]
+    weights = 1.0 / (topk_dists + 1e-2)
+    weights = weights / weights.sum(dim=1, keepdim=True)
+    blended = (weights * topk_outcomes).sum(dim=1)
+    mean_dist = topk_dists.mean(dim=1)
+    trust = torch.exp(-mean_dist / _MEMORY_TRUST_SCALE)
+    w = _MEMORY_WEIGHT * trust
+    return (1 - w) * raw_values + w * blended
+
+
+# --- Plain-Python board helpers (no torch) -- mirrors connectx_env.py's
+# free functions exactly, duplicated here (not imported) since this file
+# must be standalone. ---
+
+def _onehot(idx, n):
+    v = [0] * n
+    v[idx] = 1
+    return v
+
+
+def _rc(row, col):
+    return row * _WIDTH + col
+
+
+def _encode_board(cells):
+    out = []
+    for c in cells:
+        out.extend(_onehot(c, _CELL_WIDTH))
+    return out
+
+
+def _lowest_empty_row(cells, col):
+    for row in range(_HEIGHT - 1, -1, -1):
+        if cells[_rc(row, col)] == _EMPTY:
+            return row
+    return None
+
+
+def _legal_columns(cells):
+    return [c for c in range(_WIDTH) if _lowest_empty_row(cells, c) is not None]
+
+
+def _wins_for(cells, mark):
+    for row in range(_HEIGHT):
+        for col in range(_WIDTH):
+            if cells[_rc(row, col)] != mark:
+                continue
+            for dr, dc in ((0, 1), (1, 0), (1, 1), (1, -1)):
+                er, ec = row + dr * (_WIN_LEN - 1), col + dc * (_WIN_LEN - 1)
+                if not (0 <= er < _HEIGHT and 0 <= ec < _WIDTH):
+                    continue
+                if all(cells[_rc(row + dr * k, col + dc * k)] == mark for k in range(_WIN_LEN)):
+                    return True
+    return False
+
+
+def _board_full(cells):
+    return all(c != _EMPTY for c in cells)
+
+
+def _apply_move(cells, col, mark):
+    row = _lowest_empty_row(cells, col)
+    new_cells = list(cells)
+    new_cells[_rc(row, col)] = mark
+    return new_cells
+
+
+def _kaggle_board_to_cells(board, mark):
+    """Kaggle's board: flat list, row-major, 0=empty/1=P1/2=P2, row 0 =
+    top -- SAME convention connectx_env.py already uses, confirmed
+    against kaggle_environments' own connectx.json. `mark` tells us
+    which of Kaggle's 1/2 is US."""
+    opponent_mark = 2 if mark == 1 else 1
+    cells = []
+    for v in board:
+        if v == 0:
+            cells.append(_EMPTY)
+        elif v == mark:
+            cells.append(_AGENT)
+        else:
+            assert v == opponent_mark
+            cells.append(_OPPONENT)
+    return cells
+
+
+def _leaf_batch_values(states):
+    if not states:
+        return {{}}
+    state_t = torch.tensor(states, dtype=torch.float32)
+    norm_t = (state_t - _NORM_MEAN) / _NORM_STD
+    z = _encode(norm_t)
+    vals = _memory_blend(z, _value(z))
+    return dict(zip(states, vals.tolist()))
+
+
+def _narrow_to_center(legal_cols, max_branching):
+    """Prunes a legal-column list down to `max_branching` columns closest
+    to the board's center -- free, real Connect-4 domain knowledge (a
+    center column touches more potential 4-in-a-row lines than an edge
+    one, same theory as the empty-board opening hint). `max_branching=
+    None` is a no-op -- exact, unpruned enumeration. Only ever applied to
+    OUR OWN follow-up move choices at the deeper-escalation's round 2+
+    (see `_DEEPER_ROUNDS`'s docstring) -- never to `_ADV_ROUNDS`'s own
+    (always-unpruned) path, and never to the opponent's reply enumeration
+    at ANY round (that's what makes this a genuine worst-case
+    guarantee -- narrowing it would mean silently ignoring some of the
+    opponent's real threats)."""
+    if max_branching is None or len(legal_cols) <= max_branching:
+        return legal_cols
+    center = (_WIDTH - 1) / 2
+    return sorted(legal_cols, key=lambda c: abs(c - center))[:max_branching]
+
+
+class _RoundSearchTimeout(Exception):
+    pass
+
+
+def _check_deadline(deadline):
+    if deadline is not None and time.time() > deadline:
+        raise _RoundSearchTimeout()
+
+
+def _collect_leaves(cells1, remaining_rounds, leaf_cache, max_branching=None, deadline=None):
+    _check_deadline(deadline)
+    if _board_full(cells1):
+        return
+    for opp_col in _legal_columns(cells1):
+        cells2 = _apply_move(cells1, opp_col, _OPPONENT)
+        if _wins_for(cells2, _OPPONENT) or _board_full(cells2):
+            continue
+        if remaining_rounds <= 1:
+            leaf_cache[tuple(_encode_board(cells2))] = None
+        else:
+            for a2 in _narrow_to_center(_legal_columns(cells2), max_branching):
+                cells3 = _apply_move(cells2, a2, _AGENT)
+                if _wins_for(cells3, _AGENT):
+                    continue
+                _collect_leaves(cells3, remaining_rounds - 1, leaf_cache, max_branching, deadline)
+
+
+def _score_after_our_move(cells1, remaining_rounds, leaf_cache, max_branching=None, deadline=None):
+    """cells1: real board right after OUR move (caller already ruled out
+    an immediate win here). Returns our worst-case score -- opponent
+    picks whichever real reply hurts us most. Reads leaf values from
+    `leaf_cache` (already populated by ONE upfront batched call over the
+    WHOLE tree -- see _adversarial_plan_action) instead of calling the
+    value head again at every node."""
+    if _board_full(cells1):
+        return float(_MAX_STEPS)
+    vals = []
+    for opp_col in _legal_columns(cells1):
+        cells2 = _apply_move(cells1, opp_col, _OPPONENT)
+        if _wins_for(cells2, _OPPONENT):
+            vals.append(float(_LOSS_PENALTY))  # opponent wins -- worse than a mere draw, see _LOSS_PENALTY's comment
+        elif _board_full(cells2):
+            vals.append(float(_MAX_STEPS))
+        elif remaining_rounds <= 1:
+            vals.append(leaf_cache[tuple(_encode_board(cells2))])
+        else:
+            vals.append(_score_after_opponent_move(cells2, remaining_rounds - 1, leaf_cache, max_branching, deadline))
+    return max(vals)
+
+
+def _score_after_opponent_move(cells2, remaining_rounds, leaf_cache, max_branching=None, deadline=None):
+    """cells2: real board after the opponent's move, our turn again.
+    Returns OUR best achievable worst-case score from here."""
+    _check_deadline(deadline)
+    our_legal = _narrow_to_center(_legal_columns(cells2), max_branching)
+    if not our_legal:
+        return float(_MAX_STEPS)
+    best = None
+    for a in our_legal:
+        cells3 = _apply_move(cells2, a, _AGENT)
+        if _wins_for(cells3, _AGENT):
+            return -float(_MAX_STEPS)  # a forced win exists deeper -- short-circuit
+        s = _score_after_our_move(cells3, remaining_rounds, leaf_cache, max_branching, deadline)
+        if best is None or s < best:
+            best = s
+    return best
+
+
+class _EndgameTimeout(Exception):
+    pass
+
+
+def _exact_endgame_solve(cells0, mover, deadline):
+    """Exact (no NN) alpha-beta minimax to the true end of the game --
+    see adversarial_search.py's identical function for
+    the full docstring/calibration; this is a plain-torch-free, standalone
+    port (same convention as every other function in this file) so the
+    packaged submission never imports the project. Returns
+    `(best_action, value)` (value from `mover`'s own perspective, +1/-1/0)
+    or `(None, None)` if `deadline` was hit first."""
+    memo = {{}}
+    center = (_WIDTH - 1) / 2
+
+    def solve(cells, to_move, alpha, beta):
+        if time.time() > deadline:
+            raise _EndgameTimeout()
+        key = (tuple(cells), to_move)
+        cached = memo.get(key)
+        if cached is not None:
+            return cached
+        other = _OPPONENT if to_move == _AGENT else _AGENT
+        legal = sorted(_legal_columns(cells), key=lambda c: abs(c - center))
+        if not legal:
+            memo[key] = 0.0
+            return 0.0
+        if to_move == _AGENT:
+            best = -2.0
+            for c in legal:
+                nxt = _apply_move(cells, c, to_move)
+                if _wins_for(nxt, to_move):
+                    val = 1.0
+                elif _board_full(nxt):
+                    val = 0.0
+                else:
+                    val = solve(nxt, other, alpha, beta)
+                best = max(best, val)
+                alpha = max(alpha, best)
+                if alpha >= beta:
+                    break
+        else:
+            best = 2.0
+            for c in legal:
+                nxt = _apply_move(cells, c, to_move)
+                if _wins_for(nxt, to_move):
+                    val = -1.0
+                elif _board_full(nxt):
+                    val = 0.0
+                else:
+                    val = solve(nxt, other, alpha, beta)
+                best = min(best, val)
+                beta = min(beta, best)
+                if alpha >= beta:
+                    break
+        memo[key] = best
+        return best
+
+    root_legal = _legal_columns(cells0)
+    if not root_legal:
+        return None, None
+    root_legal = sorted(root_legal, key=lambda c: abs(c - center))
+    other = _OPPONENT if mover == _AGENT else _AGENT
+    try:
+        best_a, best_val = None, None
+        for c in root_legal:
+            nxt = _apply_move(cells0, c, mover)
+            if _wins_for(nxt, mover):
+                val = 1.0 if mover == _AGENT else -1.0
+            elif _board_full(nxt):
+                val = 0.0
+            else:
+                val = solve(nxt, other, -1.0, 1.0)
+            if best_val is None or (mover == _AGENT and val > best_val) or (mover == _OPPONENT and val < best_val):
+                best_a, best_val = c, val
+            if (mover == _AGENT and best_val == 1.0) or (mover == _OPPONENT and best_val == -1.0):
+                break
+        return best_a, best_val
+    except _EndgameTimeout:
+        return None, None
+
+
+def _run_search(surviving_actions, action_cells1, search_rounds, max_branching=None, deadline=None):
+    """One full leaf-collect + batched-eval + minimax pass at a given
+    (rounds, max_branching) setting -- factored out so it can be called
+    at two different depths, see `_DEEPER_ROUNDS`'s docstring above.
+    `deadline`: propagated into `_collect_leaves`/`_score_after_opponent_
+    move` (checked at both exponential-blowup recursion points) AND
+    checked again here, immediately around the ONE batched NN forward
+    pass -- that call is otherwise UNGUARDED/uninterruptible once
+    started, so bailing out right before it (rather than only inside the
+    pure-Python recursion) avoids ever starting an expensive tensor op
+    with no time budget left for it."""
+    leaf_cache = {{}}
+    for a in surviving_actions:
+        _collect_leaves(action_cells1[a], search_rounds, leaf_cache, max_branching, deadline)
+    _check_deadline(deadline)
+    if leaf_cache:
+        leaf_cache.update(_leaf_batch_values(list(leaf_cache.keys())))
+    _check_deadline(deadline)  # don't walk the tree on a stale/over-budget result either
+
+    best_a, best_score = None, None
+    for a in surviving_actions:
+        s = _score_after_our_move(action_cells1[a], search_rounds, leaf_cache, max_branching, deadline)
+        if best_score is None or s < best_score:
+            best_a, best_score = a, s
+    return best_a
+
+
+@torch.no_grad()
+def _adversarial_plan_action(cells0):
+    """`_ADV_ROUNDS` real adversarial rounds (our move, then the
+    opponent's worst-case real reply, repeated) before falling back to
+    the learned value head + memory blend as the leaf evaluator -- every
+    transition at every round is EXACT (real board simulation, never
+    imagined). Root action never returns PASS.
+
+    **Two-phase, GLOBALLY batched leaf evaluation** (fixed 2026-08-10,
+    same day, right before submitting -- a real timing bug caught just
+    in time, see connectx_adversarial_search.py's identical fix for the
+    full story): calling the leaf evaluator separately at every node in
+    the tree (the first version of `rounds>1`) measured up to 2.3s/move
+    against the offline-built ~2600-state memory -- OVER Kaggle's 2s
+    budget. Fixed by walking the tree TWICE (pure Python, cheap): once
+    to collect every non-terminal leaf across the WHOLE tree into one
+    deduplicated set (transpositions collapse for free), then ONE single
+    batched value+memory call, then a second walk doing the actual
+    minimax from the precomputed lookup. Re-measured after the fix
+    across 60 diverse positions (including the maximal-branching empty-
+    board case): rounds=1 max 0.427s, rounds=2 max 0.375s -- comfortably
+    (~5x) under budget again."""
+    root_legal = _legal_columns(cells0)
+    if not root_legal:
+        return None
+
+    if _ENDGAME_MAX_COLS and len(root_legal) <= _ENDGAME_MAX_COLS:
+        exact_a, _exact_val = _exact_endgame_solve(cells0, _AGENT, deadline=time.time() + _ENDGAME_TIME_BUDGET)
+        if exact_a is not None:
+            return exact_a
+        # else: timed out -- fall through to the round-based search below
+        # exactly as if this check had never happened.
+
+    # Center-out root ordering -- NOT a pruning change (every legal column
+    # is still considered, nothing narrowed), only fixes which column wins
+    # a TIE. The scoring loop below uses strict `<`, so the first action
+    # seen at a given score silently wins ties; left-to-right order made
+    # that default to the LEFTMOST column, an arbitrary, exploitable bias
+    # with no game-theoretic basis (unlike the player-1 opening hint,
+    # which deliberately picks center for a real reason). Center columns
+    # are the real stronger choice under a tie (more potential 4-in-a-row
+    # lines pass through them, same fact `_narrow_to_center` already uses
+    # for pruning) -- found from a direct user-observed pattern in real
+    # play ("when we are second we put in left going to right").
+    _center = (_WIDTH - 1) / 2
+    root_legal = sorted(root_legal, key=lambda c: abs(c - _center))
+
+    surviving_actions, action_cells1 = [], {{}}
+    for a in root_legal:
+        cells1 = _apply_move(cells0, a, _AGENT)
+        if _wins_for(cells1, _AGENT):
+            return a  # immediate win -- take it, no need to consider anything else
+        surviving_actions.append(a)
+        action_cells1[a] = cells1
+
+    base_a = _run_search(surviving_actions, action_cells1, _ADV_ROUNDS)  # always computed -- guaranteed-safe fallback
+
+    if _DEEPER_ROUNDS is not None:
+        try:
+            return _run_search(surviving_actions, action_cells1, _DEEPER_ROUNDS,
+                                max_branching=_DEEPER_MAX_BRANCHING,
+                                deadline=time.time() + _DEEPER_TIME_BUDGET)
+        except _RoundSearchTimeout:
+            pass  # didn't finish in time -- fall back to base_a exactly as if _DEEPER_ROUNDS were None
+
+    return base_a
+
+
+def _online_update(path_states, label):
+    """A FEW Adam steps on a mixed old+new batch from the persisted
+    replay buffer -- value head ONLY (encoder frozen), mirrors
+    continuous_learner.py's confirmed-safe recipe exactly (small
+    updates, EMA-scaled value targets, never a full retrain on just the
+    latest episode). `label`: either "steps" (a real win -- each state
+    labeled with its real remaining-step count) or a fixed penalty
+    (loss/draw -- every state in the walk labeled uniformly bad, same
+    convention as this session's `unsolved_penalty`). Only ever called
+    from `agent()`'s `_ONLINE_ENABLED`-guarded blocks, but a defensive
+    no-op guard here too -- never trust a single call site alone for
+    something this load-bearing."""
+    global _VALUE_TARGET_MEAN, _VALUE_TARGET_STD
+    if not _ONLINE_ENABLED:
+        return
+    if label == "steps":
+        T = len(path_states) - 1
+        for t, s in enumerate(path_states):
+            _REPLAY_BUFFER.append((list(s), float(T - t)))
+    else:
+        for s in path_states:
+            _REPLAY_BUFFER.append((list(s), float(label)))
+
+    if len(_REPLAY_BUFFER) < 8:
+        return
+    pool = list(_REPLAY_BUFFER)
+    states_t = torch.tensor([s for s, _r in pool], dtype=torch.float32)
+    returns_t = torch.tensor([r for _s, r in pool], dtype=torch.float32)
+
+    momentum = 0.98
+    new_mean, new_std = returns_t.mean(), returns_t.std().clamp(min=1e-3)
+    with torch.no_grad():
+        _VALUE_TARGET_MEAN.mul_(momentum).add_(new_mean, alpha=1 - momentum)
+        _VALUE_TARGET_STD.mul_(momentum).add_(new_std, alpha=1 - momentum)
+    returns_norm = (returns_t - _VALUE_TARGET_MEAN) / _VALUE_TARGET_STD
+
+    norm_states_t = (states_t - _NORM_MEAN) / _NORM_STD
+    with torch.no_grad():
+        z_all = _encode(norm_states_t)
+
+    n = len(pool)
+    bs = min(_ONLINE_BATCH_SIZE, n)
+    for _ in range(_ONLINE_UPDATES_PER_EPISODE):
+        idx = torch.randperm(n)[:bs]
+        pred = _value_raw(z_all[idx])
+        loss = torch.nn.functional.mse_loss(pred, returns_norm[idx])
+        _ONLINE_OPT.zero_grad()
+        loss.backward()
+        _ONLINE_OPT.step()
+
+
+def agent(observation, configuration):
+    global _EPISODE_STATES, _EPISODE_LAST_PIECES
+    board = list(observation.board)
+    mark = observation.mark
+    cells = _kaggle_board_to_cells(board, mark)
+
+    # See _ONLINE_ENABLED's own comment above -- when False, NONE of the
+    # episode-tracking/online-update machinery below runs at all, not
+    # just a no-op call: `agent()` genuinely does nothing but pick a
+    # move in that case.
+    if _ONLINE_ENABLED:
+        cur_pieces = sum(1 for v in board if v != 0)
+        # See module docstring's honest caveat -- detecting "a previous
+        # episode ended without us ever winning/drawing it ourselves"
+        # needs care: checking for an ALL-EMPTY board only works when we
+        # happen to be the FIRST mover in the new episode -- as the
+        # second mover, the very first board we see already has the
+        # opponent's first piece on it, so that check would silently
+        # miss the boundary and keep appending to a STALE trajectory
+        # from the already-ended previous episode (a real bug, caught
+        # before submission: our own test harness alternates which side
+        # we play, exactly the condition that triggers it). Robust fix:
+        # within one genuinely continuing episode, the board's total
+        # piece count increases by EXACTLY 1 between our own consecutive
+        # calls (one opponent move happened since we last acted) -- any
+        # other delta means a new episode has started, whichever side we
+        # were on. Infer a LOSS (the only remaining possibility -- our
+        # own win/draw is caught below, right after our own move).
+        #
+        # `_LOSS_PENALTY`, NOT `_UNSOLVED_PENALTY` (fixed 2026-08-10,
+        # follow-up session -- found from a direct user-observed real-game
+        # pattern, "one move before losing, ours plays leftmost"): this is
+        # the exact same mistake as the already-fixed "attacks but never
+        # defends" search bug, just unfixed in a SECOND place. The two
+        # penalties were introduced specifically so the SEARCH treats an
+        # opponent win as worse than a mere draw -- but the online
+        # learner's own training label here used `_UNSOLVED_PENALTY` (the
+        # DRAW value) for a genuine LOSS too, teaching the value head that
+        # losing and drawing are equally bad. Confirmed via real losses
+        # mined from actual Kaggle replays: the fresh (never-online-
+        # updated) search correctly blocks in all 3 traced cases, but the
+        # live, online-drifted process played the losing move instead --
+        # this conflated label is the direct mechanism.
+        if _EPISODE_STATES and cur_pieces != _EPISODE_LAST_PIECES + 1:
+            _online_update(_EPISODE_STATES, float(_LOSS_PENALTY))
+            _EPISODE_STATES = []
+        if not _EPISODE_STATES:
+            _EPISODE_STATES.append(tuple(_encode_board(cells)))
+
+    legal_cols = _legal_columns(cells)
+    if not legal_cols:
+        return 0  # should never happen -- Kaggle only calls us on a non-terminal state
+
+    # Free, EXACT domain knowledge (same "neurosymbolic gate" philosophy
+    # as every other domain's hand-given hint in this project): on a
+    # completely empty board, the center column is the known-best
+    # Connect-4 opening. Costs nothing, never worse than guessing.
+    if all(c == _EMPTY for c in cells):
+        best_action = _WIDTH // 2
+    else:
+        best_action = _adversarial_plan_action(cells)
+        if best_action is None:
+            return legal_cols[0]
+
+    if not _ONLINE_ENABLED:
+        return int(best_action)
+
+    post_cells = _apply_move(cells, best_action, _AGENT)
+    _EPISODE_STATES.append(tuple(_encode_board(post_cells)))
+    _EPISODE_LAST_PIECES = sum(1 for v in board if v != 0) + 1
+
+    if _wins_for(post_cells, _AGENT):
+        _online_update(_EPISODE_STATES, "steps")
+        _EPISODE_STATES = []
+    elif _board_full(post_cells):
+        _online_update(_EPISODE_STATES, float(_UNSOLVED_PENALTY))
+        _EPISODE_STATES = []
+
+    return int(best_action)
+'''
+
+
+def main(ckpt_path=CKPT_PATH, memory_ckpt_path=None, n_memory_games=500,
+         memory_opponent_epsilon=0.2, memory_opponent_strong_epsilon=0.3,
+         memory_weight=0.25, memory_k=5, online_lr=1e-5, online_updates_per_episode=4,
+         online_batch_size=256, unsolved_penalty_mult=1.0, adv_rounds=2, seed=0,
+         online_enabled=False, endgame_max_cols=5, endgame_time_budget=1.2,
+         deeper_rounds=None, deeper_max_branching=4, deeper_time_budget=0.6):
+    import random
+    from connectx.env import ConnectXEnv
+    from connectx.memory_build import build_episodic_memory
+    from connectx.search import load_checkpoint
+
+    ck = torch.load(ckpt_path, map_location="cpu")
+
+    # Memory is built using the SAME real adversarial search (rounds=
+    # adv_rounds) the deployed submission actually plays with, so the
+    # stored trajectories are representative of the real deployed agent's
+    # own play, not a different/weaker search's games.
+    print(f"Building offline episodic memory ({n_memory_games} self-play games, mixed opponent, "
+          f"real adversarial search rounds={adv_rounds})...")
+    mem_ckpt = memory_ckpt_path or ckpt_path
+    model, normalizer = load_checkpoint(mem_ckpt)
+    env = ConnectXEnv(width=ck["board_width"], height=ck["board_height"], win_len=ck["win_len"])
+    rng = random.Random(seed)
+    # env.py's opponent_epsilon/opponent_strong_epsilon rolls read Python's
+    # GLOBAL random module directly, not this `rng` object -- without this,
+    # "same seed" memory-building runs are silently NOT reproducible.
+    random.seed(seed)
+    memory = build_episodic_memory(env, model, normalizer, rng, n_games=n_memory_games,
+                                    opponent_epsilon=memory_opponent_epsilon,
+                                    opponent_strong_epsilon=memory_opponent_strong_epsilon,
+                                    adversarial_rounds=adv_rounds)
+    memory_zs = [z.detach().cpu() for z in memory._zs]
+    memory_outcomes = list(memory._outcomes)
+
+    blob = _encode_tensor_blob(ck, memory_zs, memory_outcomes)
+    width = 100
+    chunks = [blob[i:i + width] for i in range(0, len(blob), width)]
+    blob_literal = "\n".join(f'    "{c}"' for c in chunks)
+
+    out = SUBMISSION_TEMPLATE.format(
+        blob_literal=blob_literal, memory_weight=memory_weight, memory_k=memory_k,
+        online_lr=online_lr, online_updates_per_episode=online_updates_per_episode,
+        online_batch_size=online_batch_size, unsolved_penalty_mult=unsolved_penalty_mult,
+        adv_rounds=adv_rounds, online_enabled=online_enabled,
+        endgame_max_cols=endgame_max_cols, endgame_time_budget=endgame_time_budget,
+        deeper_rounds=deeper_rounds, deeper_max_branching=deeper_max_branching,
+        deeper_time_budget=deeper_time_budget,
+    )
+    with open(OUT_PATH, "w") as f:
+        f.write(out)
+    size_kb = len(out.encode("utf-8")) / 1024
+    print(f"Wrote {OUT_PATH} ({len(memory_zs)} memory states, {size_kb:.1f} KB)")
+
+
+if __name__ == "__main__":
+    main()

scripts/lora_selfplay_finetune.py

@@ -0,0 +1,122 @@
+"""
+LoRA self-play fine-tune of an ALREADY-TRAINED checkpoint's value head --
+this is the step that actually produced the deployed checkpoint, not a
+fresh training run. Starts from the real, working checkpoint (encoder/
+dynamics/decoder frozen, matching `train.py`'s own self-play stage, which
+never touches them either) and fine-tunes the value head under self-play
+through a LoRA-constrained delta rather than an unconstrained further Adam
+update -- directly tests whether constraining HOW MUCH the value head can
+move (not just how much data/how many rounds) helps it absorb self-play
+signal without regressing calibration.
+
+`opponent_strong_epsilon` (the curriculum variant): mixes the stronger
+1-ply-deeper heuristic into training ALONGSIDE the self-play snapshots and
+the original weak heuristic -- without this, self-play rounds only ever
+train against [self-play snapshots, weak heuristic], and never the harder
+fixed opponent at all. This is what produced this project's best-confirmed
+result (see the whitepaper's results table).
+
+rank=4/alpha=4.0: found to work well for a value head elsewhere in this
+project's development on this same architecture; not independently
+re-tuned for ConnectX specifically.
+"""
+import copy
+import random
+
+import torch
+
+from connectx.env import ConnectXEnv
+from scripts.train import make_selfplay_pool_opponent_fn, random_baseline_win_rate
+from connectx.verifier import train_mc_value_onpolicy
+from connectx.lora import LoRALinear, apply_lora
+from connectx.search import load_checkpoint, evaluate
+
+
+def _merge_and_unwrap(module):
+    """Merge every LoRALinear's delta into its frozen base weight, then
+    replace the wrapper with the plain (now-merged) nn.Linear -- so the
+    saved checkpoint is an ORDINARY WorldModel state_dict, loadable by
+    every existing caller with zero LoRA-awareness needed downstream."""
+    for name, child in list(module.named_children()):
+        if isinstance(child, LoRALinear):
+            child.merge_into_base()
+            setattr(module, name, child.linear)
+        else:
+            _merge_and_unwrap(child)
+
+
+def main(ckpt_path="checkpoints/connectx_checkpoint.pt", seed=0, rank=4, alpha=4.0,
+         selfplay_rounds=5, selfplay_epsilon=0.4, selfplay_mc_rounds_per_iter=15,
+         mc_problems_per_round=400, selfplay_pool_size=5,
+         opponent_epsilon=0.15, opponent_strong_epsilon=0.0,
+         save_path="checkpoints/connectx_checkpoint_lora_selfplay.pt"):
+    raw_ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    model, normalizer = load_checkpoint(ckpt_path)
+
+    rng = random.Random(seed)
+    torch.manual_seed(seed)
+    random.seed(seed)  # env.py's opponent stochasticity reads the GLOBAL
+    # random module, not this local rng -- see train.py's own note
+
+    env = ConnectXEnv(width=raw_ckpt["board_width"], height=raw_ckpt["board_height"],
+                       win_len=raw_ckpt["win_len"])
+    max_steps = (env.width * env.height) // 2 + 2
+
+    lora_params = apply_lora(model.value, rank=rank, alpha=alpha)
+    n_lora = sum(p.numel() for p in lora_params)
+    n_frozen = sum(p.numel() for p in model.value.parameters()) - n_lora
+    print(f"LoRA rank={rank} alpha={alpha} on model.value: "
+          f"{n_lora} trainable params, {n_frozen} frozen (base value head)")
+
+    frozen_pool = []
+    for sp_round in range(selfplay_rounds):
+        frozen_model = copy.deepcopy(model).eval()
+        for p in frozen_model.parameters():
+            p.requires_grad_(False)
+        frozen_pool.append(frozen_model)
+        if len(frozen_pool) > selfplay_pool_size:
+            frozen_pool.pop(0)
+        print(f"\nSelf-play LoRA fine-tune round {sp_round + 1}/{selfplay_rounds} "
+              f"(opponent_selfplay_epsilon={selfplay_epsilon}, pool_size={len(frozen_pool)})...")
+        opponent_env = ConnectXEnv(width=env.width, height=env.height, win_len=env.win_len)
+        selfplay_fn = make_selfplay_pool_opponent_fn(frozen_pool, normalizer, opponent_env)
+        selfplay_train_env = ConnectXEnv(width=env.width, height=env.height, win_len=env.win_len,
+                                          opponent_epsilon=opponent_epsilon,
+                                          opponent_strong_epsilon=opponent_strong_epsilon,
+                                          opponent_selfplay_epsilon=selfplay_epsilon,
+                                          opponent_policy_fn=selfplay_fn)
+        # Adam(model.value.parameters()) inside train_mc_value_onpolicy
+        # naturally trains ONLY the LoRA deltas here: the wrapped base
+        # linears have requires_grad=False (set by apply_lora), so their
+        # .grad stays None and Adam's step() skips them.
+        train_mc_value_onpolicy(selfplay_train_env, model, normalizer, rng,
+                                 n_rounds=selfplay_mc_rounds_per_iter,
+                                 n_problems_per_round=mc_problems_per_round,
+                                 max_steps=max_steps, unsolved_penalty=max_steps)
+
+    _merge_and_unwrap(model.value)
+
+    torch.save({
+        "model_state": model.state_dict(),
+        "norm_mean": normalizer.mean.cpu(),
+        "norm_std": normalizer.std.cpu(),
+        "state_dim": raw_ckpt["state_dim"],
+        "num_actions": raw_ckpt["num_actions"],
+        "latent_dim": raw_ckpt["latent_dim"],
+        "hidden_dim": raw_ckpt["hidden_dim"],
+        "board_width": env.width,
+        "board_height": env.height,
+        "win_len": env.win_len,
+    }, save_path)
+    print(f"\nSaved LoRA-self-play-fine-tuned checkpoint to {save_path}")
+
+    print("\n" + "=" * 20 + " EVALUATION (no oracle -- vs. random-legal-play baseline only) " + "=" * 20)
+    eval_rng = random.Random(999)
+    problems = [env.random_problem(eval_rng) for _ in range(150)]
+    random_baseline_win_rate(env, problems, max_steps, random.Random(1000))
+    evaluate(env, model, normalizer, problems, depth=1, beam_width=8, max_total_steps=max_steps,
+             label="Baseline A (model, depth=1)")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,212 @@
+"""
+Trains a WorldModel checkpoint on the real Kaggle ConnectX board (7x6,
+win_len=4). No oracle exists at this scale (see env.py's BFS_MAX_CELLS), so
+this is a genuinely no-oracle recipe throughout:
+
+1. Stage 1 (dynamics + decoder): self-supervised, `train_utils.train_stage1`
+   on random transitions/rollouts -- never depends on an oracle for any
+   domain.
+2. Stage 2 (value head): `verifier.train_mc_value_onpolicy`, on-policy
+   Monte Carlo, `unsolved_penalty=max_steps` -- this is an adversarial
+   domain, where a training walk can end in a LOSS, not just run out of
+   steps; without this penalty every losing walk is silently discarded and
+   the value head never learns to avoid losing moves at all (confirmed as
+   the direct cause of a real 0%, worse-than-random win rate before this
+   was added).
+3. Self-play fine-tune: each round freezes the current model as an
+   opponent (viewed from the opponent's side, via `_swap_agent_opponent`)
+   and trains further against a small, capped POOL of past frozen
+   snapshots (not just the latest one) mixed with the original fixed
+   heuristic -- every opponent used above is fixed and non-learning, which
+   is the actual ceiling on how strong a non-self-play policy can get.
+
+Evaluation has no oracle-ceiling comparison to report against (none exists
+at this scale) -- instead: win rate against the ORIGINAL fixed opponent
+(deterministic, `opponent_epsilon=0`) the model is ultimately graded
+against, plus a random-legal-play win rate as the required floor.
+"""
+import copy
+import random
+
+import torch
+
+from connectx.env import ConnectXEnv, AGENT, OPPONENT, EMPTY, _encode_board
+from connectx.model import WorldModel
+from connectx.train_utils import DEVICE, StateNormalizer, generate_transitions, generate_rollout_sequences, \
+    train_stage1, eval_stage1, eval_multistep_rollout
+from connectx.verifier import train_mc_value_onpolicy
+from connectx.search import evaluate
+from connectx.adversarial_search import real_adversarial_plan_action
+
+
+def _swap_agent_opponent(cells):
+    """AGENT<->OPPONENT relabeling, EMPTY unchanged -- lets a model that
+    was only ever trained to play as AGENT evaluate a position from the
+    OTHER side's perspective, by pretending that side is AGENT instead."""
+    return [c if c == EMPTY else (OPPONENT if c == AGENT else AGENT) for c in cells]
+
+
+def make_selfplay_pool_opponent_fn(frozen_pool, frozen_normalizer, opponent_env):
+    """Builds an `opponent_policy_fn(cells) -> column` (see env.py's
+    `opponent_policy_fn` extension point) that plays using a frozen
+    snapshot of this same architecture's own trained judgment, not a
+    hand-written heuristic. Each call picks a snapshot uniformly at random
+    from `frozen_pool` (a small population of past snapshots, not just the
+    latest one -- a coarse approximation of real population-based self-
+    play/fictitious play, so the live policy can't narrowly overfit to
+    counter-play against whatever the single latest snapshot happens to
+    do) and uses the real adversarial search (rounds=1 -- this function is
+    called on the order of 100K+ times across a training run, so a slower
+    multi-round search would balloon total training time)."""
+
+    def opponent_policy_fn(cells):
+        frozen_model = random.choice(frozen_pool)
+        swapped_state = tuple(_encode_board(_swap_agent_opponent(cells)))
+        with torch.no_grad():
+            a = real_adversarial_plan_action(opponent_env, frozen_model, frozen_normalizer, swapped_state, rounds=1)
+        return a
+
+    return opponent_policy_fn
+
+
+def random_baseline_win_rate(env, problems, max_steps, rng):
+    """Required control: an agent picking uniformly among its legal
+    non-PASS columns. The floor the trained model needs to beat -- not a
+    certified ceiling (no oracle exists at this scale), just the honest
+    "did training do anything at all" check."""
+    wins = 0
+    for state, _ in problems:
+        cur = state
+        for _ in range(max_steps):
+            if env.is_solved(cur):
+                break
+            legal = [a for a in range(env.num_actions) if env.is_legal(cur, a)]
+            non_pass = [a for a in legal if a != env.width]
+            a = rng.choice(non_pass or legal)
+            cur, _r, done = env.step(cur, a)
+            if done:
+                break
+        if env.is_solved(cur):
+            wins += 1
+    n = len(problems)
+    print(f"{'Random-legal-play baseline':30s} win_rate={wins/n:.3f} ({wins}/{n})")
+    return wins / n
+
+
+def main(seed=0, ckpt_path="checkpoints/connectx_checkpoint.pt", latent_dim=96, hidden_dim=256,
+         n_problems=2000, walk_len=8, mc_rounds=25, mc_problems_per_round=400,
+         opponent_epsilon=0.15, opponent_strong_epsilon=0.0,
+         selfplay_rounds=5, selfplay_epsilon=0.4, selfplay_mc_rounds_per_iter=15,
+         selfplay_pool_size=5):
+    # Two env instances, same board, different opponent determinism: `env`
+    # (opponent_epsilon=0, the pure deterministic opponent) is what
+    # evaluation is graded against. `train_env` (opponent_epsilon>0) is
+    # used ONLY for generating training data -- a perfectly deterministic
+    # opponent means every training walk from a matching starting side is
+    # the SAME exact game, a real, diagnosed weakness (the model
+    # reproducibly lost as first player against this exact opponent while
+    # winning as second player).
+    env = ConnectXEnv(width=7, height=6, win_len=4)
+    train_env = ConnectXEnv(width=7, height=6, win_len=4, opponent_epsilon=opponent_epsilon,
+                             opponent_strong_epsilon=opponent_strong_epsilon)
+    rng = random.Random(seed)
+    torch.manual_seed(seed)
+    # Also seed the GLOBAL random module: env.py's opponent_epsilon/
+    # opponent_strong_epsilon rolls read `random.random()`/`random.choice()`
+    # directly, not this function's own seeded `rng` -- without this,
+    # "same seed" runs are silently not reproducible whenever opponent
+    # stochasticity is enabled.
+    random.seed(seed)
+    max_steps = (env.width * env.height) // 2 + 2
+
+    print(f"Domain: ConnectX (real board), {env.width}x{env.height}, win_len={env.win_len}, "
+          f"num_actions={env.num_actions}, state_dim={env.state_dim}, "
+          f"train opponent_epsilon={opponent_epsilon}, opponent_strong_epsilon={opponent_strong_epsilon}\n")
+    assert env.bfs_solve(env.random_problem(rng)[0]) is None, \
+        "expected no oracle at real-board scale -- see env.py's BFS_MAX_CELLS"
+
+    print("Generating stage-1 (dynamics) data...")
+    train_transitions = generate_transitions(train_env, rng, n_problems=n_problems, walk_len=walk_len)
+    val_transitions = generate_transitions(env, rng, n_problems=300, walk_len=walk_len)
+    print(f"  {len(train_transitions)} train transitions, {len(val_transitions)} val transitions")
+
+    unroll_k = 4
+    train_sequences = generate_rollout_sequences(train_env, rng, n_problems=n_problems, k=unroll_k)
+    val_sequences = generate_rollout_sequences(env, rng, n_problems=300, k=unroll_k)
+    print(f"  {len(train_sequences)} train sequences, {len(val_sequences)} val sequences (k={unroll_k})")
+
+    all_states_for_norm = [t[0] for t in train_transitions] + [t[2] for t in train_transitions]
+    normalizer = StateNormalizer(all_states_for_norm).to(DEVICE)
+
+    model = WorldModel(env.state_dim, env.num_actions, latent_dim=latent_dim, hidden_dim=hidden_dim).to(DEVICE)
+
+    print("\nStage 1: training encoder + dynamics + decoder...")
+    train_stage1(model, normalizer, train_transitions, sequences=train_sequences, k=unroll_k)
+
+    print("\nStage-1 val metrics:")
+    print(" ", eval_stage1(model, normalizer, val_transitions))
+    print(" ", eval_multistep_rollout(model, normalizer, val_sequences, k=unroll_k))
+
+    print("\nStage 2: no-oracle value training (on-policy Monte Carlo, bfs_solve never called)...")
+    train_mc_value_onpolicy(train_env, model, normalizer, rng, n_rounds=mc_rounds,
+                             n_problems_per_round=mc_problems_per_round, max_steps=max_steps,
+                             unsolved_penalty=max_steps)
+
+    # Self-play fine-tune: only the value head trains during
+    # train_mc_value_onpolicy (encoder/dynamics/decoder stay fixed from
+    # stage 1), so each frozen snapshot's encoder/dynamics are identical
+    # to the live model's -- only the value judgment (and therefore the
+    # self-play opponent's move choices) differs round to round.
+    # `selfplay_pool_size` keeps a capped, small population of past
+    # snapshots (drops the oldest once full) rather than only the single
+    # latest one, or an unbounded pool that would let early, still-weak
+    # snapshots dominate forever.
+    frozen_pool = []
+    for sp_round in range(selfplay_rounds):
+        frozen_model = copy.deepcopy(model).eval()
+        for p in frozen_model.parameters():
+            p.requires_grad_(False)
+        frozen_pool.append(frozen_model)
+        if len(frozen_pool) > selfplay_pool_size:
+            frozen_pool.pop(0)
+        print(f"\nSelf-play fine-tune round {sp_round + 1}/{selfplay_rounds} "
+              f"(opponent_selfplay_epsilon={selfplay_epsilon}, pool_size={len(frozen_pool)})...")
+        opponent_env = ConnectXEnv(width=7, height=6, win_len=4)  # plain -- used only for is_legal/num_actions
+        selfplay_fn = make_selfplay_pool_opponent_fn(frozen_pool, normalizer, opponent_env)
+        selfplay_train_env = ConnectXEnv(width=7, height=6, win_len=4,
+                                          opponent_epsilon=opponent_epsilon,
+                                          opponent_selfplay_epsilon=selfplay_epsilon,
+                                          opponent_policy_fn=selfplay_fn)
+        train_mc_value_onpolicy(selfplay_train_env, model, normalizer, rng, n_rounds=selfplay_mc_rounds_per_iter,
+                                 n_problems_per_round=mc_problems_per_round, max_steps=max_steps,
+                                 unsolved_penalty=max_steps)
+
+    if ckpt_path:
+        torch.save({
+            "model_state": model.state_dict(),
+            "norm_mean": normalizer.mean.cpu(),
+            "norm_std": normalizer.std.cpu(),
+            "state_dim": env.state_dim,
+            "num_actions": env.num_actions,
+            "latent_dim": latent_dim,
+            "hidden_dim": hidden_dim,
+            "board_width": env.width,
+            "board_height": env.height,
+            "win_len": env.win_len,
+        }, ckpt_path)
+        print(f"\nSaved checkpoint to {ckpt_path}")
+
+    print("\n" + "=" * 20 + " EVALUATION (no oracle -- vs. random-legal-play baseline only) " + "=" * 20)
+    eval_rng = random.Random(999)
+    problems = [env.random_problem(eval_rng) for _ in range(150)]
+    random_baseline_win_rate(env, problems, max_steps, random.Random(1000))
+    evaluate(env, model, normalizer, problems, depth=1, beam_width=8, max_total_steps=max_steps,
+             label="Baseline A (model, depth=1)")
+    evaluate(env, model, normalizer, problems, depth=3, beam_width=8, max_total_steps=max_steps,
+             label="Search (model, depth=3)")
+
+    return model, normalizer
+
+
+if __name__ == "__main__":
+    main()

tests/submission_test.py

@@ -0,0 +1,183 @@
+"""
+Validates connectx_submission.py against a REAL, independent, single-ply
+Connect-4 game engine -- deliberately NOT reusing connectx_env.py's
+`step` (which bundles our move and its own fixed opponent's reply into
+one call, exactly the shortcut this test needs to NOT rely on, since the
+whole point is to check the submission's `agent(observation,
+configuration)` interface the way Kaggle's real engine actually drives
+it: one call per ply, alternating sides, real `board`/`mark`/
+`configuration` objects).
+
+Three opponents, since a single opponent can't tell us much:
+1. Random-legal-column opponent (a floor -- if this loses to random, the
+   translation from Kaggle's board format is definitely broken).
+2. The SAME fixed heuristic connectx_env.py trains against (win-now /
+   block / leftmost) -- a sanity check that connectx_train.py's own
+   100% internal solve-rate result actually reproduces when the
+   submission is driven through a REAL alternating-turn engine instead
+   of the training env's bundled step().
+3. A slightly stronger heuristic (look-two-ahead: also block if leaving
+   the opponent a forced win one ply later) -- a first, cheap check of
+   what this project's own docs already name as the honest limitation
+   (trained against one fixed weak opponent, no guarantee against a
+   different/stronger one).
+"""
+import importlib.util
+import random
+import sys
+from types import SimpleNamespace
+
+WIDTH, HEIGHT, WIN_LEN = 7, 6, 4
+
+
+def _rc(row, col):
+    return row * WIDTH + col
+
+
+def _lowest_empty_row(board, col):
+    for row in range(HEIGHT - 1, -1, -1):
+        if board[_rc(row, col)] == 0:
+            return row
+    return None
+
+
+def legal_columns(board):
+    return [c for c in range(WIDTH) if _lowest_empty_row(board, c) is not None]
+
+
+def wins_for(board, mark):
+    for row in range(HEIGHT):
+        for col in range(WIDTH):
+            if board[_rc(row, col)] != mark:
+                continue
+            for dr, dc in ((0, 1), (1, 0), (1, 1), (1, -1)):
+                er, ec = row + dr * (WIN_LEN - 1), col + dc * (WIN_LEN - 1)
+                if not (0 <= er < HEIGHT and 0 <= ec < WIDTH):
+                    continue
+                if all(board[_rc(row + dr * k, col + dc * k)] == mark for k in range(WIN_LEN)):
+                    return True
+    return False
+
+
+def board_full(board):
+    return all(c != 0 for c in board)
+
+
+def drop(board, col, mark):
+    row = _lowest_empty_row(board, col)
+    board = list(board)
+    board[_rc(row, col)] = mark
+    return board
+
+
+def random_opponent(board, mark, rng):
+    return rng.choice(legal_columns(board))
+
+
+def weak_heuristic_opponent(board, mark, rng):
+    """The exact policy connectx_env.py bakes into training (win-now /
+    block / leftmost) -- reimplemented independently here (not imported)
+    so this test doesn't share a bug with the thing it's checking."""
+    opp = 2 if mark == 1 else 1
+    legal = legal_columns(board)
+    for col in legal:
+        if wins_for(drop(board, col, mark), mark):
+            return col
+    for col in legal:
+        if wins_for(drop(board, col, opp), opp):
+            return col
+    return legal[0]
+
+
+def stronger_heuristic_opponent(board, mark, rng):
+    """weak_heuristic_opponent plus a 1-ply-deeper check: among moves that
+    pass the first two checks, avoid any move that hands the OPPONENT an
+    immediate winning reply next turn, if a safer alternative exists."""
+    opp = 2 if mark == 1 else 1
+    legal = legal_columns(board)
+    for col in legal:
+        if wins_for(drop(board, col, mark), mark):
+            return col
+    for col in legal:
+        if wins_for(drop(board, col, opp), opp):
+            return col
+    safe = []
+    for col in legal:
+        nxt = drop(board, col, mark)
+        if board_full(nxt):
+            safe.append(col)
+            continue
+        opp_can_win = any(wins_for(drop(nxt, c2, opp), opp) for c2 in legal_columns(nxt))
+        if not opp_can_win:
+            safe.append(col)
+    return rng.choice(safe) if safe else legal[0]
+
+
+def _load_agent(path):
+    spec = importlib.util.spec_from_file_location("connectx_submission", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod.agent
+
+
+def play_game(our_agent, opponent_fn, our_mark, rng, max_plies=WIDTH * HEIGHT):
+    board = [0] * (WIDTH * HEIGHT)
+    opp_mark = 2 if our_mark == 1 else 1
+    config = SimpleNamespace(columns=WIDTH, rows=HEIGHT, inarow=WIN_LEN)
+    current = 1  # P1 always moves first
+    for _ply in range(max_plies):
+        if current == our_mark:
+            obs = SimpleNamespace(board=list(board), mark=our_mark)
+            col = our_agent(obs, config)
+            if col not in legal_columns(board):  # safety net -- must never trigger
+                col = legal_columns(board)[0]
+            board = drop(board, col, our_mark)
+            if wins_for(board, our_mark):
+                return "win"
+        else:
+            col = opponent_fn(board, opp_mark, rng)
+            board = drop(board, col, opp_mark)
+            if wins_for(board, opp_mark):
+                return "loss"
+        if board_full(board):
+            return "draw"
+        current = opp_mark if current == our_mark else our_mark
+    return "draw"
+
+
+def run_match(our_agent, opponent_fn, opponent_name, n_games, seed):
+    """Each game gets its OWN independently-seeded rng (`seed`, `i`) --
+    NOT one shared rng consumed sequentially across all `n_games` (the
+    original version of this function). That original design meant a
+    different agent taking even one different early action would shift
+    every subsequent random draw for the REST of the batch, silently
+    turning "compare agent A vs agent B on the same 60 games" into two
+    barely-related sequences of games -- confirmed as a real confound,
+    not a hypothetical one: testing an opening-move change that only
+    ever affects the very first ply still swung the OVERALL 60-game
+    win rate by double digits, which a change isolated to ply 0 has no
+    honest mechanism to cause on its own against a fixed opponent
+    POLICY (only via this exact RNG-cascade artifact). Per-game seeding
+    makes different agents' results directly, fairly comparable game-by-
+    game."""
+    results = {"win": 0, "loss": 0, "draw": 0}
+    for i in range(n_games):
+        rng = random.Random(seed * 100_003 + i)  # large odd multiplier, cheap decorrelation across games
+        our_mark = 1 if i % 2 == 0 else 2  # alternate who moves first
+        outcome = play_game(our_agent, opponent_fn, our_mark, rng)
+        results[outcome] += 1
+    n = n_games
+    print(f"vs {opponent_name:28s} win={results['win']}/{n} ({results['win']/n:.1%})  "
+          f"loss={results['loss']}/{n}  draw={results['draw']}/{n}")
+    return results
+
+
+if __name__ == "__main__":
+    path = sys.argv[1] if len(sys.argv) > 1 else "submission.py"
+    our_agent = _load_agent(path)
+    print(f"Loaded agent from {path}\n")
+
+    n_games = 60
+    run_match(our_agent, random_opponent, "random-legal-play", n_games, seed=0)
+    run_match(our_agent, weak_heuristic_opponent, "weak heuristic (=training opp)", n_games, seed=1)
+    run_match(our_agent, stronger_heuristic_opponent, "stronger heuristic (1-ply deeper)", n_games, seed=2)