alextoti1/WorldModel-ConnectX

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 (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.

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

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.