wtedw /
nanoalphazero
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://nanoalphaze…
wtedw /
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://nanoalphaze…
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.
Discuss this product