/muzero-general

MuZero

Primary LanguagePythonMIT LicenseMIT

Code style: black

MuZero General

A commented and documented implementation of MuZero based on the Google DeepMind paper and the associated pseudocode. It is designed to be easily adaptable for every games or reinforcement learning environments (like gym). You only need to edit the game file with the parameters and the game class. Please refer to the documentation and the example.

MuZero is a model based reinforcement learning algorithm, successor of AlphaZero. It learns to master games without knowing the rules. It only knows actions and then learn to play and master the game. It is at least more efficient than similar algorithms like AlphaZero, SimPLe and World Models. See How it works.

Features

  • Residual Network and Fully connected network in PyTorch
  • Multi-Threaded/Asynchronous/Cluster with Ray
  • CPU/GPU support
  • Multi GPU support (To Do)
  • TensorBoard real-time monitoring
  • Model weights automatically saved at checkpoints
  • Single and two player mode
  • Commented and documented
  • Easily adaptable for new games
  • Examples of board games, Gym and Atari games (See list of implemented games)
  • Pretrained weights available
  • Windows support (Workaround: Use the notebook in Google Colab)

Further improvements

These improvements are active research, they are personal ideas and go beyond MuZero paper. We are open to contributions and other ideas.

Demo

All performances are tracked and displayed in real time in TensorBoard :

cartpole training summary

Testing Lunar Lander :

lunarlander training preview

Games already implemented

  • Cartpole (Tested with the fully connected network)
  • Lunar Lander (Tested in deterministic mode with the fully connected network)
  • Gridworld (Tested with the fully connected network)
  • Tic-tac-toe (Tested with the fully connected network and the residual network)
  • Connect4 (Slightly tested with the residual network)
  • Gomoku
  • Twenty-One / Blackjack (Tested with the residual network)
  • Atari Breakout

Tests are done on Ubuntu with 16 GB RAM / Intel i7 / GTX 1050Ti Max-Q. We make sure to obtain a progression and a level which ensures that it has learned. But we do not systematically reach a human level. For certain environments, we notice a regression after a certain time. The proposed configurations are certainly not optimal and we do not focus for now on the optimization of hyperparameters. Any help is welcome.

Code structure

code structure

See also: MuZero network summary

Getting started

Installation

git clone https://github.com/werner-duvaud/muzero-general.git
cd muzero-general

pip install -r requirements.txt

Run

python muzero.py

To visualize the training results, run in a new terminal:

tensorboard --logdir ./results

Authors

Please use this bibtex if you want to cite this repository (master branch) in your publications:

@misc{muzero-general,
  author       = {Duvaud, Hainaut},
  title        = {MuZero General: Open Reimplementation of MuZero},
  year         = {2019},
  publisher    = {GitHub},
  journal      = {GitHub repository},
  howpublished = {\url{https://github.com/werner-duvaud/muzero-general}},
}

Getting involved