Plain Monte-Carlo Tree Search (MCTS) relies on random rollouts (simulations of future game moves) to estimate the value of a game position. AlphaZero’s key insight was to replace random rollouts with a learned neural network that biases the tree search toward promising moves, cutting the number of simulations needed for strong play. In this project you will reproduce that idea end-to-end at toy scale: small enough to actually understand every component, not just call a library.
We’ll pick something with a small enough state space that self-play converges quickly but still nontrivial strategically, e.g. Connect Four (6×7 board) or simplified Othello (6×6 instead of 8×8 board).
You will need to implement:
- Game engine — board representation, legal move generation, win/draw detection, and a simple canonical state encoding (e.g. a small tensor: current player’s pieces, opponent’s pieces, maybe a “whose turn” plane). This needs to be fast (called thousands of times per second during self-play).
- Neural network — a small transformer or CNN that outputs (i) move probabilities over legal actions and (ii) a scalar estimate of current player’s win probability.
- PUCT-based MCTS — Each node stores visit counts, total value, and a prior probability (from the neural network) per child. Selection uses the PUCT formula (UCT variant that incorporates the prior). Instead of random rollouts to a terminal state, leaf evaluation calls the network directly: to give move priors to expand children with and an estimate of the position value (so no rollout to the end of the game needed).
- Self-play data generation loop — play the current network against itself using MCTS-guided move selection (with temperature/exploration noise, for move diversity early in training), recording (state, MCTS visit-count distribution, eventual game outcome) tuples as training data.
- Training loop — sample from the self-play replay buffer, train the network to match the MCTS visit distribution (cross-entropy) and the actual game outcome (MSE), alternating with more self-play generation (creating an iterate-and-improve loop).
You will compare the performance of the trained network+MCTS against a plain-MCTS-with-random-rollouts baseline, in particular: (i) win rate curve across self-play training iterations (to evaluate learning progress) and (ii) how many simulations does the learned version need to match a given plain-MCTS strength.