Implementing AlphaZero for Gomoku

Introduction#

AlphaZero combines self-play, Monte Carlo tree search (MCTS), and a deep neural network in a closed training loop. It reached superhuman performance in Go, chess, and shogi. This note uses Gomoku as a smaller setting and presents a minimal implementation together with engineering suggestions.

1. Algorithm overview#

Policy-value network fθ(s) → (p, v)
- p: prior probabilities for the legal actions.
- v: the estimated outcome for the current player in state s, in [-1, 1].
MCTS: combines the network’s priors and value estimate to search the game tree and produce a stronger policy π.
Self-play: uses MCTS as the behavior policy, with temperature-controlled sampling, to collect (s, π, z) examples.
Training: minimizes L = (z - v)^2 - π·log p + c||θ||^2.

2. Environment and board representation#

Board: an N × N grid (9×9 or 11×11 by default); each move passes the turn to the other player.
Terminal conditions: five consecutive stones horizontally, vertically, or diagonally, or a full-board draw.
State encoding:
- Two planes: stones belonging to the current player and stones belonging to the opponent.
- Optional planes: the most recent move and an indicator for whose turn it is.

3. Policy-value network (PyTorch)#

1
import torch, torch.nn as nn, torch.nn.functional as F
2

3
class PolicyValueNet(nn.Module):
4
    def __init__(self, board_size=9, channels=64):
5
        super().__init__()
6
        self.board_size = board_size
7
        self.conv = nn.Sequential(
8
            nn.Conv2d(2, channels, 3, padding=1), nn.ReLU(),
9
            nn.Conv2d(channels, channels, 3, padding=1), nn.ReLU(),
10
            nn.Conv2d(channels, channels, 3, padding=1), nn.ReLU(),
11
        )
12
        # policy head
13
        self.policy_head = nn.Sequential(
14
            nn.Conv2d(channels, 2, 1), nn.ReLU(),
15
            nn.Flatten(),
16
            nn.Linear(2*board_size*board_size, board_size*board_size)
17
        )
18
        # value head
19
        self.value_head = nn.Sequential(
20
            nn.Conv2d(channels, 1, 1), nn.ReLU(),
21
            nn.Flatten(),
22
            nn.Linear(board_size*board_size, 64), nn.ReLU(),
23
            nn.Linear(64, 1), nn.Tanh()
24
        )
25
    def forward(self, x):
26
        h = self.conv(x)
27
        p = self.policy_head(h)
28
        v = self.value_head(h)
29
        return p, v

4. MCTS implementation (PUCT)#

1
import math, numpy as np
2

3
class Node:
4
    def __init__(self, prior):
5
        self.P = prior   # prior probability
6
        self.N = 0       # visit count
7
        self.W = 0.0     # accumulated value
8
        self.Q = 0.0     # mean value
9
        self.children = {}  # action -> Node
10

11
class MCTS:
12
    def __init__(self, net, board_size, c_puct=1.5):
13
        self.net = net
14
        self.board_size = board_size
15
        self.c_puct = c_puct
16
        self.root = Node(1.0)
17

18
    def search(self, state, legal_moves, to_tensor):
19
        path = []
20
        node = self.root
21
        # 1) Selection
22
        while node.children:
23
            # Select the action with the highest UCB score
24
            best_a, best_child, best_score = None, None, -1e9
25
            for a, ch in node.children.items():
26
                u = self.c_puct * ch.P * math.sqrt(node.N + 1) / (1 + ch.N)
27
                score = ch.Q + u
28
                if score > best_score:
29
                    best_a, best_child, best_score = a, ch, score
30
            path.append((node, best_a))
31
            node = best_child
32
            # The external environment advances the state. This example
33
            # simplifies the flow and expands only the first leaf.
34
            if not node.children:
35
                break
36
        # 2) Expansion & evaluation
37
        x = to_tensor(state)  # [1, 2, N, N]
38
        with torch.no_grad():
39
            logits, v = self.net(x)
40
            p = torch.softmax(logits, dim=1).cpu().numpy()[0]
41
            v = float(v.item())
42
        # Keep priors for legal actions only
43
        priors = {a: p[a] for a in legal_moves}
44
        norm = sum(priors.values()) + 1e-8
45
        for a in priors:
46
            priors[a] /= norm
47
        leaf = Node(1.0)
48
        leaf.children = {a: Node(priors[a]) for a in legal_moves}
49
        node.children = leaf.children
50
        # 3) Backup
51
        for n, _ in path:
52
            n.N += 1
53
            n.W += v
54
            n.Q = n.W / n.N
55
        return v
56

57
    def get_policy(self, temperature=1.0):
58
        # Derive a visit-count distribution from the root
59
        visits = np.zeros(self.board_size * self.board_size, dtype=np.float32)
60
        for a, ch in self.root.children.items():
61
            visits[a] = ch.N
62
        if temperature == 0:
63
            pi = np.zeros_like(visits)
64
            pi[np.argmax(visits)] = 1.0
65
            return pi
66
        visits = visits ** (1.0 / max(1e-6, temperature))
67
        if visits.sum() == 0:
68
            visits += 1.0
69
        return visits / visits.sum()

The example leaves out the mechanics for advancing the environment state during search. In a full implementation, the environment must be copied and stepped through; this sketch focuses on the PUCT selection rule and how it combines values with priors.

5. Minimal self-play and training loop#

1
from collections import deque
2

3
def encode_state(board, current_player):
4
    # board: shape (N, N), values in {0, 1, -1}; 1 moves first, -1 second
5
    cur = (board == current_player).astype(np.float32)
6
    opp = (board == -current_player).astype(np.float32)
7
    x = np.stack([cur, opp], axis=0)  # [2, N, N]
8
    return torch.from_numpy(x).unsqueeze(0)  # [1, 2, N, N]
9

10
class ReplayBuffer:
11
    def __init__(self, capacity=50000):
12
        self.buf = deque(maxlen=capacity)
13
    def push(self, s, pi, z):
14
        self.buf.append((s, pi, z))
15
    def sample(self, batch):
16
        idx = np.random.choice(len(self.buf), batch, replace=False)
17
        s, pi, z = zip(*[self.buf[i] for i in idx])
18
        return torch.cat(s), torch.tensor(np.array(pi), dtype=torch.float32), torch.tensor(z, dtype=torch.float32).unsqueeze(1)
19

20
def loss_fn(outputs, targets_pi, targets_v, l2=1e-4):
21
    logits, v = outputs
22
    loss_p = F.cross_entropy(logits, targets_pi.argmax(dim=1))
23
    loss_v = F.mse_loss(v, targets_v)
24
    l2_penalty = sum((p**2).sum() for p in net.parameters()) * l2
25
    return loss_p + loss_v + l2_penalty

Training-loop outline:

Run MCTS for a number of simulations in each game (for example, 100–800) and use the resulting π as the training target.
At the end of a game, assign z ∈ {+1, -1, 0} for a win, loss, or draw.
Train the network for several steps after collecting a batch of games. A learning-rate schedule and data augmentation (rotations and reflections) can help.

6. Suggested parameters and engineering notes#

Board size: 9×9 converges faster; 11×11 is more challenging.
MCTS: 200–800 simulations; c_puct ≈ 1.0–2.5; root Dirichlet noise α ≈ 0.3 encourages exploration.
Temperature: use T = 1 for the first 10 moves, then gradually move toward greedy selection with T → 0.
Data augmentation: the eight board symmetries (rotations and reflections) can improve sample efficiency.
Mixed precision and gradient accumulation: can speed up training and reduce memory pressure.

7. Playing and visualization#

Use a web Canvas/HTML page to display the board and call a backend for move inference.
Add a Gomoku widget to an “algorithm visualization lab” and display MCTS visit heatmaps in real time.

8. Summary#

For a closer look at the search loop, return to the MCTS overview or step through the MCTS playground. Compare its random rollouts with the policy priors and value estimates used here: where does each method get its feedback?

This note outlines a minimal AlphaZero-style Gomoku system: a policy-value network, PUCT MCTS, self-play, and training. Start by getting the loop running on a 9×9 board, then try a more capable network and more simulations to improve play strength.

References:

Silver et al., Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm (AlphaZero).
Mastering the Game of Go without Human Knowledge (AlphaGo Zero), Nature, 2017.

Implementing AlphaZero for Gomoku

🌿 What is a Digital Garden?

Introduction#

1. Algorithm overview#

2. Environment and board representation#

3. Policy-value network (PyTorch)#

4. MCTS implementation (PUCT)#

5. Minimal self-play and training loop#

6. Suggested parameters and engineering notes#

7. Playing and visualization#

8. Summary#

Article concept network

Local Knowledge Neighborhood

Implementing CFR for Kuhn Poker

Interpreting RL Results: Sampling Gains, Reward Loopholes, and Generalization

How GRPO Compares Responses: Group Advantages, Zero Variation, and Length Weights

Monte Carlo Tree Search

Causal Effects & Selection Bias