README.md

# 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).