DHDev0/Muzero

Pytorch Implementation of MuZero for gym environment. It support any Discrete , Box and Box2D configuration for the action space and observation space.

Muzero

Pytorch Implementation of MuZero for Gymnasium environment. It should support any Discrete ,Box or Box2D configuration for the observation space and action space.

I documented the code as closely as possible next to the Muzero paper. ( You will find two types of comments: self explain variables and comments within the code)

P.S. Gymnasium class Discrete, Box, and Box2D refer to different types of observation or action spaces in a simulation. Discrete refers to a 1D array of integers, Box refers to a 1D array of floating point numbers, and Box2D refers to a multi-dimensional array of integers or floating point numbers.

MuZero -> MuZero Unplugged -> Stochastic MuZero

Table of contents

• Getting started
  • Local Installation
  • Docker
  • Dependency
• Usage
  • Jupyter Notebook
  • CLI
• Features
• How to make your own custom gym environment?
• Authors
• Subjects
• License

Getting started

Local Installation

PIP dependency : requirement.txt

git clone https://github.com/DHDev0/Muzero.git

cd Muzero

pip install -r requirements.txt

If you experience some difficulty refer to the first cell Tutorial or use the dockerfile.

Docker

Build image: (building time: 22 min , memory consumption: 8.75 GB)

docker build -t muzero .

(do not forget the ending dot)

Start container:

docker run --cpus 2 --gpus 1 -p 8888:8888 muzero
#or
docker run --cpus 2 --gpus 1 --memory 2000M -p 8888:8888 muzero
#or
docker run --cpus 2 --gpus 1 --memory 2000M -p 8888:8888 --storage-opt size=15g muzero

The docker run will start a jupyter lab on https://localhost:8888//lab?token=token (you need the token) with all the necessary dependency for cpu and gpu(Nvidia) compute.

Option meaning:
--cpus 2 -> Number of allocated (2) cpu core
--gpus 1 -> Number of allocated (1) gpu
--storage-opt size=15gb -> Allocated storage capacity 15gb (not working with windows WSL)
--memory 2000M -> Allocated RAM capacity of 2GB
-p 8888:8888 -> open port 8888 for jupyter lab (default port of the Dockerfile)

Stop the container:

docker stop $(docker ps -q --filter ancestor=muzero)

Delete the container:

docker rmi -f muzero

Dependency

Language :

• Python 3.8 to 3.10 (bound by the retro compatibility of Ray and Pytorch)

Library :

• torch 1.13.0
• torchvision 0.14.0
• ray 2.0.1
• gymnasium 0.27.0
• matplotlib >=3.0
• numpy 1.21.5

More details at: requirement.txt

Usage

Jupyter Notebook

For practical example, you can use the Tutorial.

CLI

Set your config file (example): https://github.com/DHDev0/Muzero/blob/main/config/

First and foremost cd to the project folder:

cd Muzero

Training :

python muzero_cli.py train config/experiment_133_config.json

Training with report

python muzero_cli.py train report config/experiment_133_config.json

Inference (play game with specific model) :

python muzero_cli.py train play config/experiment_133_config.json

Training and Inference :

python muzero_cli.py train play config/experiment_133_config.json

Benchmark model :

python muzero_cli.py benchmark config/experiment_133_config.json

Training + Report + Inference + Benchmark :

python muzero_cli.py train report play benchmark play config/experiment_133_config.json

Features

[x] implemented , [ ] unimplemented

• Work for any Gymnasium environments/games. (any combination of continous or/and discrete action and observation space)
• MLP network for game state observation. (Multilayer perceptron)
• LSTM network for game state observation. (LSTM)
• Transformer decoder for game state observation. (Transformer)
• Residual network for RGB observation using render. (Resnet-v2 + MLP)
• Residual LSTM network for RGB observation using render. (Resnet-v2 + LSTM)
• MCTS with 0 simulation (use of prior) or any number of simulation.
• Model weights automatically saved at best selfplay average reward.
• Priority or Uniform for sampling in replay buffer.
• Manage illegal move with negative reward.
• Scale the loss using the importance sampling ratio.
• Custom "Loss function" class to apply transformation and loss on label/prediction.
• Load your pretrained model from tag number.
• Single player mode.
• Training / Inference report. (not live, end of training)
• Single/Multi GPU or Single/Multi CPU for inference, training and self-play.
• Support mix precision for training and inference.(torch_type: bfloat16,float16,float32,float64)
• Pytorch gradient scaler for mix precision in training.
• Tutorial with jupyter notebook.
• Pretrained weights for cartpole. (you will find weight, report and config file)
• Commented with link/page to the paper.
• Support : Windows , Linux , MacOS.
• Fix pytorch linear layer initialization. (refer to : https://tinyurl.com/ykrmcnce)
• Support of Gymnasium 0.27.0
• Two player mode and more.
• Hyperparameter search. (pseudo-code available in self_play.py)
• Training and deploy on cloud cluster using Kubeflow, Airflow or Ray for aws,gcp and azure.

How to make your own custom gym environment?

Refer to the Gym documentation

You will be able to call your custom gym environment in muzero after you register it in gym.

Authors

• Daniel Derycke

Subjects

Deep reinforcement learning

License

GPL-3.0 license