Hexbot Part 3: Using Neural Networks

This is the third post in a series of posts about my foray into building a bot to play Hex. The focus of this post is on how to augment MCTS with a neural network, and in particular the AlphaZero-style strategy for doing so.

Hexbot Part 1: The 101 Unit
Hexbot Part 2: The Board and Our First Bot
Hexbot Part 3: Using Neural Networks
Hexbot Part 4: Did It Work?

Improving MCTS

In part 2 we went over MCTS, which serves as a tunable baseline for how to choose moves in a game of Hex: more rollouts means stronger play at the cost of thinking time, where “strength” grows logarithmically with rollout count.

We can, of course, do a lot better than basic MCTS, though many paths forward incorporate domain knowledge about Hex strategy. In part 1 we already mentioned alpha-beta search, a different search algorithm entirely which relies on a value function, often a hand-crafted one.

MCTS itself can also be improved in a number of ways. The expand step, for example, can be modified to use heavy rollouts where a heuristic is used instead of choosing moves completely at random. The search step can be augmented by calculating a prior over the possible moves, which influences the weight assigned to moves with low visit count.

The AlphaZero architecture, which we will be applying to Hex, takes essentially both of these approaches, using a neural network to augment an MCTS in two ways:

• A value network is used instead of a pure rollout, giving each new leaf node a number indicating how favorable the network considers the board state for either player, and these are aggregated up the tree in the same way as rollout outcomes.

• A policy network initializes each node with a prior distribution over child nodes, i.e. available moves, which biases the search towards moves the network deems more likely.

The code I used for this type of search is on GitHub, but it’s very similar to the basic MCTS we discussed before just with a different function to maximize when descending the tree that accounts for the values from the policy network.

Improving the neural network

The question of course is how to train these networks. One strategy is to use expert games as a training set, and this was how AlphaGo was initially trained. However a key insight of theirs was that the network-augmented MCTS is an algorithm for taking policy calculations from the model and aggregating them into a better policy. By having the model play games against itself and training the model to predict the improved policy returned by MCTS, the model can improve through self-play. The discovery of AlphaGo Zero is that this idea of improvement through self-play works even when starting from a randomly-initialized network.¹

It may seem a little mysterious that this works at all, and indeed there are a lot of ways for it to go wrong, but the basic idea makes sense. The model is essentially learning to predict what an MCTS search will do, and the self-play is structured to ensure that a diversity of positions are explored and that the problem is computationally feasible by gradually “compressing” recent MCTS effort into the model. The model improves the MCTS, and the MCTS improves the model.

The model architecture

The model is fundamentally a convolutional neural network, taking an “image”² of the board and applying a convolution and several residual convolutional layers to it before taking the final image and flattening it to be mapped into a policy and value output.

The model size has three configuration parameters:

• conv_channels, the number of channels in the intermediate board images.
• conv_layers, the number of residual convolutional layers to apply.
• value_hidden, the number of hidden neurons in the value calculation.

In pseudo-Python, the model calculation looks like this:

def evaluate(self, board):
    im_input = self.board_to_image(board)

    im = self.conv_input(im_input)
    for (ker0, ker1) in self.residual_conv_layers:
        im = self.residual_conv_layer(im, ker0, ker1)

    policy = self.policy_head(im)
    value = self.value_head(im)

    return policy, value

def board_to_image(self, board):
    # an 11x11 image with 2 channels, one for red pieces
    # and one for blue pieces where a "pixel" is 1.0 if
    # there is a piece and 0.0 otherwise
    return ...

def conv_input(self, im):
    # conv2d with conv_channels 3x3 filters
    return im
        .hex_conv2d(self.conv_input_kernels)
        .batch_norm()
        .leaky_relu()

def residual_conv_layer(self, im, ker0, ker1):
    # conv2d with conv_channels 3x3 filters
    return im
        .hex_conv2d(ker0)
        .batch_norm()
        .leaky_relu()
        .hex_conv2d(ker1)
        .batch_norm()
        .add(im) # skip connection
        .leaky_relu()

def policy_head(self, im):
    # conv2d with two 1x1 filters
    return im
        .hex_conv2d(self.policy_kernels)
        .batch_norm()
        .leaky_relu()
        .flatten()
        .matmul(self.policy_linear) # 242 -> 121
        .log_softmax()

def value_head(self, im):
    # conv2d with one 1x1 filter
    return im
        .hex_conv2d(self.value_kernel)
        .batch_norm()
        .leaky_relu()
        .flatten()
        .matmul(self.value_linear0) # 121 -> value_hidden
        .leaky_relu()
        .matmul(self.value_linear1) # value_hidden -> 1

The actual model code is on GitHub, and is built on Burn, a tensor and autodiff library for Rust. hex_conv2d here refers to a normal convolution but with some of the kernel values set to zero, making e.g. a 3x3 kernel only able to incorporate information from cells that are up to 1 cell away. (I was sort of just winging it when I decided to do this but I was worried information being able to travel more quickly in one direction might result in weird behavior.)

This architecture is lifted more or less directly from the AlphaGo Zero paper, with the aforementioned “hex convolution” change.

Training it

The training procedure at this point is straightforward in concept but full of many little optimization³ problems⁴ and fiddly implementation⁵ details⁶. I’ll go over the main points quickly here but will leave discussion of smaller points for another time.

I wrote some code to set up distributed training with components that interact over HTTP:

• A controller, which serves model checkpoints and the self-play log. This was just an HTTP server with a light enough workload to run on my laptop.

• A self-play daemon, which periodically downloads the latest checkpoint and generates self-play data. I ran three instances of this in a Google Cloud Run worker pool with NVIDIA L4 GPUs.

• An optimizer daemon, which follows the self-play log and runs gradient descent on the latest model checkpoint, periodically uploading a new checkpoint to the controller. I ran this component on my laptop while at my computer and moved it to Cloud Run while away (because running my M4 unattended at high loads makes me nervous.)

I trained the model for about 24 hours with the above setup and the following model configuration:

{
  "conv_layers": 16,
  "conv_channels": 32,
  "value_hidden": 256
}

Up next

In the next post, I’ll explain how to decide if the work we’ve done is any good.