Building AI for Connect 4

Play the game → · Watch demo →

A deep dive into training neural networks to master a classic game — Connect 4, Longhorn Edition.

---

Demo

Login	Choose Your Battle

Model Report	Transformer Wins

---

Connect 4 is a solved game: with perfect play, the first player can always force a win. But finding that perfect play requires analyzing millions of possible board positions. This project trains neural networks that learn to play at a high level by studying expert moves from Monte Carlo Tree Search (MCTS), without explicitly programming game strategy.

We built two architectures — a Convolutional Neural Network (CNN) and a Transformer — and compared their performance. Both treat the problem as classification over 7 columns (where to drop a chip).

---

The Challenge

• Training data: Moves generated by MCTS (the same family of algorithms used by AlphaGo and AlphaZero). MCTS simulates thousands of random games from each position to find statistically strong moves.
• Pipeline: Generate games → Run MCTS → Record moves → Mirror data → Train NNs.
• Dataset: 703,111 positions from MCTS with 2000 simulations per move. After horizontal mirroring (exploiting the board’s symmetry), we have 1.4M+ training examples.

Data creation: even coverage across game phases

We varied the number of random opening moves from 4 to 14 using a right-tailed probability distribution (mode at 6). That gives more even coverage across early-, mid-, and late-game board states instead of biasing toward the opening.

Board encoding

Each position is encoded as a 6×7×2 tensor:

• Channel 0: 1 where Player 1 (orange) has a chip, 0 elsewhere
• Channel 1: 1 where Player 2 (white) has a chip, 0 elsewhere

This lets the network clearly distinguish both players’ positions.

---

Network architectures

	CNN	Transformer
Style	AlphaZero-style residual CNN	Vision Transformer (ViT)
Idea	2D convolutions scan the 6×7 board; residual blocks detect patterns (lines, threats, blocks)	Attention over all 42 positions; 6 blocks, 8 heads, hidden dim 128
Training data	700K+ positions (mirrored → 1.4M+)	1.4M+ positions (with mirror augmentation)
Outputs	Policy (best column) + value (win probability)	Policy + value

Both output a policy (probability over 7 columns) and a value (scalar from -1 to +1 for expected outcome).

---

Training results

Metric	CNN	Transformer
Policy accuracy (top-1)	~72%	~74% (77% in our final run)
Policy accuracy (top-2)	~88%	~90% (92% in our final run)
Training	~32 epochs, ~15 min (Colab A100)	~15 epochs, ~70 min (Colab T4)

Both models predict the exact MCTS move around ~73% of the time and include the best move in their top-2 about ~90% of the time — with instant inference instead of 2000 simulations per move.

---

Open in Google Colab

Run the notebooks from this repo directly in Colab:

Notebook	Description
Data viewer	Inspect the dataset, mirroring code, and how to train a 2-headed net; run simulated games.
Data generation	Generate MCTS data (random opening 4–14, right-tailed); build the rolling mcts_distribution.npz.
CNN builder	Train the residual CNN (policy + value heads).
Transformer	Train the Vision Transformer for Connect 4.
Play Connect 4	Play against the CNN or Transformer in the notebook.

---

Repo layout

• data_viewer.ipynb — Dataset exploration, mirroring, training discussion, simulated games
• data_generation/data_generation_no_guide.ipynb — MCTS data generation (4–14 random opening, right-tailed)
• network_training/CNN/cnn_builder.ipynb — CNN training
• network_training/Transformer/transformer_colab.ipynb — Transformer training
• Connect4_Demo/play_connect4.ipynb — Interactive play vs CNN or Transformer
• mirroring/mirror.py — Horizontal board mirroring for augmentation
• data_balance/balance.py — Move-depth loss weighting
• aws_docker_backend/ — AWS inference backend for the web app
• anvil_webpage/ — Anvil web app (play in the browser)
• REPORT.md — Detailed technical report (data pipeline, architectures, AWS, Anvil)

---

Credits

Development team

• Alina Hota
• Arturo Juarez
• Zan Merrill
• Rohini Sondole

Technologies: TensorFlow/Keras, Anvil, AWS Lightsail, Docker, Python

Source: github.com/AHMerrill/connect-4

---

Think you can beat our AI? Play at the Anvil app or use the Play Connect 4 notebook.