Bots for Connect4
by Mihai Nica
Explore Monte Carlo Tree Search step by step and play Connect 4 against an MCTS-powered bot.
About this notebook
This notebook teaches Monte Carlo Tree Search for Connect 4 and other k-in- a-line games. It is one part of a series about the AlphaZero algorithm. There is no external dataset. The notebook generates each board position from the game rules. The code uses marimo, JAX, NumPy, einops, tqdm, and matplotlib.
You set the first player to Human, AI, or AI vs AI. You click a column button to place a piece. The AI then runs many random rollouts and plays its move. A bar chart shows the simulation count for each action and the average score. A line chart shows the root value across turns. A timeline slider lets you return to earlier turns. A separate sandbox has four buttons for the Select, Expand, Simulate, and Backup steps. Sliders and checkboxes change the board size, the number in a line needed to win, gravity, and the diagonal rule.
JAX compiles the game rules and runs batched random rollouts with jax.lax.scan. The tree stores its nodes in NumPy arrays for fast in-place updates. UCB scores guide the search at each level. The backup step records each result along the path to the root. A Mermaid diagram draws the search tree. One slider sets the rollout budget for the AI.
