KoMaster

Building a Go AI to beat me—a 6-dan player—using AlphaGo-style neural nets, MCTS, and reinforcement learning against KataGo on PSU's GPU cluster.

Comment

Building a Go AI to Beat a 6-Dan Player

Inspiration

I've played Go competitively for years, reaching 6-dan — a level where most players have internalized thousands of patterns, joseki, and strategic principles built up over a lifetime of study. When DeepMind's AlphaGo defeated Lee Sedol in 2016, I watched with a mix of awe and skepticism. Could a machine truly understand Go, or was it just brute-forcing search trees? I wanted to find out firsthand — not by reading papers, but by building one myself.

The goal was simple and personal: train a Go AI from scratch that can beat me.

How It Works

The system is an AlphaGo-style pipeline built in PyTorch, running on a university GPU cluster. It has three main components:

1. Neural Network Architecture

The model uses a dual-headed residual network — the same family of architectures used in AlphaGo Zero. A shared convolutional trunk processes the board state, feeding into two separate output heads:

  • Policy head — outputs a probability distribution $\pi(a \mid s)$ over all legal moves, answering "what move should I play?"
  • Value head — outputs a scalar $v(s) \in [-1, 1]$, answering "who is winning?"

The network takes the board position $s$ as a multi-channel tensor (encoding current stones, ko state, and move history) and produces:

$$f_\theta(s) = (\mathbf{p}, v), \quad \mathbf{p} \in \Delta^{361}, \quad v \in [-1, 1]$$

2. Monte Carlo Tree Search (MCTS)

During game play, the AI doesn't just follow the raw policy — it runs MCTS to search ahead. Each simulation selects moves by maximizing an upper confidence bound:

$$\text{UCB}(s, a) = Q(s, a) + c_{\text{puct}} \cdot \pi(a \mid s) \cdot \frac{\sqrt{N(s)}}{1 + N(s, a)}$$

where $Q(s,a)$ is the average value of subtrees visited via action $a$, $N(s)$ is the visit count of the parent node, and $c_{\text{puct}}$ controls the exploration-exploitation tradeoff. With 400 simulations per move, the AI considers a wide tree of possibilities before committing to a decision.

3. Training Pipeline

Training happens in two phases:

Phase 1 — Supervised Learning: The model is pretrained on foxq_elite_7d_plus.pkl, a dataset of games played by 7-dan and above players on the Fox Go Server. The loss is:

$$\mathcal{L}{\text{SL}} = -\sum_a \pi^*(a) \log p\theta(a \mid s) + \lambda \cdot (v_\theta(s) - z)^2$$

where $\pi^*$ is the expert move distribution and $z \in {-1, +1}$ is the game outcome. After 25 epochs, the model learns reasonable opening patterns and basic strategic intuitions.

Phase 2 — Reinforcement Learning: The model plays games against KataGo, a superhuman Go engine, and updates on the resulting data. Rather than pure self-play (which is noisy and slow to converge), training against KataGo provides high-quality, adversarial feedback every single game. The RL objective mirrors the supervised loss but uses self-generated game data:

$$\mathcal{L}{\text{RL}} = -\sum_t \pi{\text{MCTS}}(a_t) \log p_\theta(a_t \mid s_t) + \lambda \cdot (v_\theta(s_t) - z_t)^2$$

with a mixed replay buffer combining KataGo games, self-play, and a fraction of expert data to prevent forgetting.

What I Built

  • A custom BatchedMCTS engine that runs parallel game generation on GPU
  • Asynchronous KataGo integration via background threads, so GPU training never waits on engine I/O
  • A training loop (train_gpu.py) designed for the university cluster, with checkpoint recovery and TensorBoard logging
  • A 100-iteration RL curriculum with configurable expert data mixing ratios

Challenges

The Value Head Collapse

The hardest problem I encountered was catastrophic forgetting in the value head. Around iteration 22 of RL training, the value loss collapsed from $\approx 0.38$ to $\approx 0.02$ — suspiciously low. The model had become extremely confident in its position evaluations, but completely wrong. Resignation rates dropped from 95% to 0%, and games that should end in 150 moves were dragging out to the 300-move cap.

This created a death spiral:

  1. Value head predicts garbage → resignation never triggers
  2. Games run to maximum length → each iteration takes 3× longer
  3. Longer games generate noisier value targets → value head gets worse

The root cause was that self-play outcomes from a weak model are too noisy to supervise a value head reliably. With an 80/20 expert/self-play mix, the expert data masked the problem in the loss metrics — but the damage was real.

The fix: Restart RL from the pretrained checkpoint with resignation disabled, letting the policy head continue improving while the value head issue is deferred. The long-term solution is KataGo value targets — using KataGo's position evaluations as ground-truth $v$ labels instead of binary game outcomes. This is closer to what AlphaGo Zero actually does in production.

Debugging on a Remote Cluster

Running on a shared university machine (e5-cse-135-19) added friction at every step: conda environments that vanish between sessions, module systems that need manual loading, jobs interrupted by cluster policy, and no easy way to interactively debug a GPU process. I learned to make training scripts self-recovering — checkpointing every 3 iterations and resuming cleanly from any interruption.

What I Learned

  • The value head is the hardest part. Policy learning is relatively forgiving; value learning is brittle and collapses easily under noisy targets.
  • Monitoring the right metrics matters. Loss curves looked fine while resignation rates were silently dying — a lesson in the limits of scalar metrics.
  • KataGo as a curriculum is powerful. Training against a superhuman opponent compresses what would take hundreds of self-play iterations into a fraction of the compute.
  • Preserving supervised learning progress is crucial. The 25 epochs of pretraining represent real, hard-won knowledge. Overwriting it carelessly during RL is a trap.

What's Next

  • Implement KataGo value targets to fix the value head properly
  • Re-enable resignation once value estimates are trustworthy
  • Adaptive expert mixing: transition $70\% \to 60\% \to 50\%$ expert data as the model matures
  • Curriculum learning: automatically increase KataGo's visit count as win rate improves
  • The ultimate test: sit down, open a board, and play.

Built With

Share this project:

Updates