Hello HN,

I built a complete, game-agnostic implementation of AlphaZero in JAX.

repo: https://github.com/wtedw/nanoAlphaZero

demo (NN + MCTS run locally in your browser): https://nanoalphazero.wtedw.com

It uses no human data, can train grandmaster-level chess models, and supports a variety of games: Chess, Go 3x3 - 9x9, Hex 4x4 - 9x9, Connect Four

You can also use this repo to train AlphaZero on any custom 2-player, perfect-information game.

How does it work?

At a high level, the entire AlphaZero algorithm gets compiled into a single jitted run_fn that repeatedly performs self-play and model updates:

  state = make_alphazero()

  def run_fn(state):
      games = selfplay(state)  # using Gumbel MuZero

      # Move active games into the self-play buffer
      # Move completed games into the replay buffer

      state = train(state, replay_buffer.sample())

      return state

  while True:
      state = run_fn(state)

There are no threads, queues, or distributed workers to manage. It is just one large JAX function.

The repo is primarily focused on making large-scale AlphaZero experimentation fast and easy to run. Training strong models is secondary and mostly serves as a sanity check that the underlying logic is sound.

Documentation on training custom / complex games is sparse, so if you have any questions feel free to message me.