Hello HN,I built a complete, game-agnostic implementation of AlphaZero in JAX.repo: https://github.com/wtedw/nanoAlphaZerodemo (NN + MCTS run locally in your browser): https://nanoalphazero.wtedw.comIt uses no human data, can train grandmaster-level chess models, and supports a variety of games: Chess, Go 3x3 - 9x9, Hex 4x4 - 9x9, Connect FourYou can also use this repo to train AlphaZero on any custom 2-player, perfect-information game (this needs more documentation).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 func...