# 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 .`, 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=)` — 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.`) 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).