Giters

mohitd404 / LightZero

[NeurIPS 2023 Spotlight] LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

LightZero

Twitter PyPI PyPI - Python Version Loc Comments

Code Test Badge Creation Package Release

GitHub Org's stars GitHub stars GitHub forks GitHub commit activity GitHub issues GitHub pulls Contributors GitHub license

Updated on 2023.09.21 LightZero-v0.0.2

LightZero is a lightweight, efficient, and easy-to-understand open-source algorithm toolkit that combines Monte Carlo Tree Search (MCTS) and Deep Reinforcement Learning (RL).

English | 简体中文(Simplified Chinese) | Paper

Background

The integration of Monte Carlo Tree Search and Deep Reinforcement Learning, exemplified by AlphaZero and MuZero, has achieved unprecedented performance levels in various games, including Go and Atari. This advanced methodology has also made significant strides in scientific domains like protein structure prediction and the search for matrix multiplication algorithms. The following is an overview of the historical evolution of the Monte Carlo Tree Search algorithm series: pipeline

Overview

LightZero is an open-source algorithm toolkit that combines MCTS and RL for PyTorch. It provides support for a range of MCTS-based RL algorithms and applications with the following advantages:

  • Lightweight.
  • Efficient.
  • Easy-to-understand.

For further details, please refer to Features, Framework Structure and Integrated Algorithms.

LightZero aims to promote the standardization of the MCTS+RL algorithm family to accelerate related research and applications. A performance comparison of all implemented algorithms under a unified framework is presented in the Benchmark.

Outline

Features

Lightweight: LightZero integrates multiple MCTS algorithm families and can solve decision-making problems with various attributes in a lightweight framework. The algorithms and environments LightZero implemented can be found here.

Efficient: LightZero uses mixed heterogeneous computing programming to improve computational efficiency for the most time-consuming part of MCTS algorithms.

Easy-to-understand: LightZero provides detailed documentation and algorithm framework diagrams for all integrated algorithms to help users understand the algorithm's core and compare the differences and similarities between algorithms under the same paradigm. LightZero also provides function call graphs and network structure diagrams for algorithm code implementation, making it easier for users to locate critical code. All the documentation can be found here.

Framework Structure

Image Description 2

The above picture is the framework pipeline of LightZero. We briefly introduce the three core modules below:

Model: Model is used to define the network structure, including the __init__ function for initializing the network structure and the forward function for computing the network's forward propagation.

Policy: Policy defines the way the network is updated and interacts with the environment, including three processes: the learning process, the collecting process, and the evaluation process.

MCTS: MCTS defines the structure of the Monte Carlo search tree and the way it interacts with the Policy. The implementation of MCTS includes two languages: Python and C++, implemented in ptree and ctree, respectively.

For the file structure of LightZero, please refer to lightzero_file_structure.

Integrated Algorithms

LightZero is a library with a PyTorch implementation of MCTS algorithms (sometimes combined with cython and cpp), including:

The environments and algorithms currently supported by LightZero are shown in the table below:

Env./Algo. AlphaZero MuZero EfficientZero Sampled EfficientZero Gumbel MuZero Stochastic MuZero
TicTacToe ✔ ✔ 🔒 🔒 ✔ 🔒
Gomoku ✔ ✔ 🔒 🔒 ✔ 🔒
Connect4 ✔ ✔ 🔒 🔒 🔒 🔒
2048 ✔ ✔ 🔒 🔒 🔒 ✔
Chess 🔒 🔒 🔒 🔒 🔒 🔒
Go 🔒 🔒 🔒 🔒 🔒 🔒
CartPole --- ✔ ✔ ✔ ✔ ✔
Pendulum --- ✔ ✔ ✔ ✔ 🔒
LunarLander --- ✔ ✔ ✔ ✔ ✔
BipedalWalker --- ✔ ✔ ✔ ✔ 🔒
Atari --- ✔ ✔ ✔ ✔ ✔
MuJoCo --- ✔ ✔ ✔ 🔒 🔒
MiniGrid --- 🔒 🔒 🔒 🔒 🔒

(1): "✔" means that the corresponding item is finished and well-tested.

(2): "🔒" means that the corresponding item is in the waiting-list (Work In Progress).

(3): "---" means that this algorithm doesn't support this environment.

Installation

You can install the latest LightZero in development from the GitHub source codes with the following command:

git clone https://github.com/opendilab/LightZero.git
cd LightZero
pip3 install -e .

Kindly note that LightZero currently supports compilation only on Linux and macOS platforms. We are actively working towards extending this support to the Windows platform. Your patience during this transition is greatly appreciated.

Installation with Docker

We also provide a Dockerfile that sets up an environment with all dependencies needed to run the LightZero library. This Docker image is based on Ubuntu 20.04 and installs Python 3.8, along with other necessary tools and libraries. Here's how to use our Dockerfile to build a Docker image, run a container from this image, and execute LightZero code inside the container.

  1. Download the Dockerfile: The Dockerfile is located in the root directory of the LightZero repository. Download this file to your local machine.
  2. Prepare the build context: Create a new empty directory on your local machine, move the Dockerfile into this directory, and navigate into this directory. This step helps to avoid sending unnecessary files to the Docker daemon during the build process. 
    mkdir lightzero-docker
mv Dockerfile lightzero-docker/
cd lightzero-docker/
  3. Build the Docker image: Use the following command to build the Docker image. This command should be run from inside the directory that contains the Dockerfile. 
    docker build -t ubuntu-py38-lz:latest -f ./Dockerfile .
  4. Run a container from the image: Use the following command to start a container from the image in interactive mode with a Bash shell. 
    docker run -dit --rm ubuntu-py38-lz:latest /bin/bash
  5. Execute LightZero code inside the container: Once you're inside the container, you can run the example Python script with the following command: 
    python ./LightZero/zoo/classic_control/cartpole/config/cartpole_muzero_config.py

Quick Start

Train a MuZero agent to play CartPole:

cd LightZero
python3 -u zoo/classic_control/cartpole/config/cartpole_muzero_config.py

Train a MuZero agent to play Pong:

cd LightZero
python3 -u zoo/atari/config/atari_muzero_config.py

Train a MuZero agent to play TicTacToe:

cd LightZero
python3 -u zoo/board_games/tictactoe/config/tictactoe_muzero_bot_mode_config.py

Benchmark

Click to collapse

tictactoe_bot-mode_main connect4_bot-mode_main gomoku_bot-mode_main

pong_main qbert_main mspacman_main mspacman_sez_K

"Factored Policy" indicates that the agent learns a policy network that outputs a categorical distribution. After manual discretization, the dimensions of the action space for the five environments are 11, 49 (7^2), 256 (4^4), 64 (4^3), and 4096 (4^6), respectively. On the other hand, "Gaussian Policy" refers to the agent learning a policy network that directly outputs parameters (mu and sigma) for a Gaussian distribution.

pendulum_main pendulum_sez_K lunarlander_main

bipedalwalker_main hopper_main walker2d_main

pong_gmz_ns mspacman_gmz_ns gomoku_bot-mode_gmz_ns lunarlander_gmz_ns

2048_stochasticmz_mz mspacman_gmz_ns

keycorridors3r3_exploration fourrooms_exploration

Awesome-MCTS Notes

Paper Notes

The following are the detailed paper notes (in Chinese) of the above algorithms:

Click to collapse

Algo. Overview

The following are the overview MCTS principle diagrams of the above algorithms:

Click to expand

Awesome-MCTS Papers

Here is a collection of research papers about Monte Carlo Tree Search. This Section will be continuously updated to track the frontier of MCTS.

Key Papers

Click to expand

LightZero Implemented series

AlphaGo series

MuZero series

MCTS Analysis

MCTS Application

Other Papers

Click to expand

ICML

ICLR

NeurIPS

Other Conference or Journal

Feedback and Contribution

  • File an issue on Github

  • Contact our email (opendilab@pjlab.org.cn)

  • We appreciate all the feedback and contributions to improve LightZero, both algorithms and system designs.

Citation

@misc{lightzero,
      title={LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios},
      author={Yazhe Niu and Yuan Pu and Zhenjie Yang and Xueyan Li and Tong Zhou and Jiyuan Ren and Shuai Hu and Hongsheng Li and Yu Liu},
      year={2023},
      eprint={2310.08348},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Acknowledgments

This repo is partially based on the following repo, many thanks to their pioneering work:

Special thanks to @PaParaZz1, @karroyan, @nighood, @jayyoung0802, @timothijoe, @TuTuHuss, @HarryXuancy, @puyuan1996, @HansBug for their contributions and support.

License

All code within this repository is under Apache License 2.0.

Thanks to our contributors

(Back to top)

About

[NeurIPS 2023 Spotlight] LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios

License:Apache License 2.0

Languages

Language:Python 88.3%Language:C++ 9.0%Language:Cython 2.3%Language:Shell 0.2%Language:Dockerfile 0.1%Language:Makefile 0.1%