Muax
An open-source implementation of MuZero with JAX and TensorFlow.
Muax is an open-source library that provides implementations of MuZero variants using JAX and TensorFlow. It supports distributed reinforcement learning and training.
The project integrates with DeepMind’s Acme framework and allows for comparing Monte Carlo Tree Search (MCTS) policies on various multi-agent games from OpenSpiel, such as Go.
You can find the repository on GitHub: bwfbowen/muax.