Skip to content

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

License

Notifications You must be signed in to change notification settings

opendilab/LightZero

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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 discord badge

Updated on 2024.11.23 LightZero-v0.1.0

English | 简体中文(Simplified Chinese) | Documentation | LightZero Paper | 🔥UniZero Paper | 🔥ReZero Paper

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). For any questions about LightZero, you can consult the RAG-based Q&A assistant: ZeroPal.

🔍 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 Monte Carlo Tree Search (MCTS) and Reinforcement Learning (RL) for PyTorch. It supports a range of MCTS-based RL algorithms and applications, offering several key 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 Sampled MuZero EfficientZero Sampled EfficientZero Gumbel MuZero Stochastic MuZero UniZero Sampled UniZero ReZero
TicTacToe 🔒 🔒 🔒 🔒 🔒 🔒
Gomoku 🔒 🔒 🔒 🔒 🔒
Connect4 🔒 🔒 🔒 🔒 🔒 🔒
2048 --- 🔒 🔒 🔒 🔒 🔒 🔒
Chess 🔒 🔒 🔒 🔒 🔒 🔒 🔒 🔒 🔒 🔒
Go 🔒 🔒 🔒 🔒 🔒 🔒 🔒 🔒 🔒 🔒
CartPole --- 🔒 🔒
Pendulum --- 🔒 🔒
LunarLander --- 🔒
BipedalWalker --- 🔒 🔒 🔒
Atari --- 🔒 🔒
DeepMind Control --- --- --- 🔒 🔒 🔒 🔒
MuJoCo --- 🔒 🔒 🔒 🔒 🔒 🔒
MiniGrid --- 🔒 🔒 🔒 🔒 🔒
Bsuite --- 🔒 🔒 🔒 🔒 🔒
Memory --- 🔒 🔒 🔒 🔒 🔒
SumToThree (billiards) --- 🔒 🔒 🔒 🔒 🔒 🔒 🔒 🔒
MetaDrive --- 🔒 🔒 🔒 🔒 🔒 🔒 🔒 🔒

(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

Train a UniZero agent to play Pong:

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

📚 Documentation

The LightZero documentation can be found here. It contains tutorials and the API reference.

For those interested in customizing environments and algorithms, we provide relevant guides:

Should you have any questions, feel free to contact us for support.

📊 Benchmark

Click to expand

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

You can also refer to the relevant Zhihu column (in Chinese): In-depth Analysis of MCTS+RL Frontier Theories and Applications.

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

🌏 Citation

@article{niu2024lightzero,
  title={LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios},
  author={Niu, Yazhe and Pu, Yuan and Yang, Zhenjie and Li, Xueyan and Zhou, Tong and Ren, Jiyuan and Hu, Shuai and Li, Hongsheng and Liu, Yu},
  journal={Advances in Neural Information Processing Systems},
  volume={36},
  year={2024}
}

@article{pu2024unizero,
  title={UniZero: Generalized and Efficient Planning with Scalable Latent World Models},
  author={Pu, Yuan and Niu, Yazhe and Ren, Jiyuan and Yang, Zhenjie and Li, Hongsheng and Liu, Yu},
  journal={arXiv preprint arXiv:2406.10667},
  year={2024}
}

@article{xuan2024rezero,
  title={ReZero: Boosting MCTS-based Algorithms by Backward-view and Entire-buffer Reanalyze},
  author={Xuan, Chunyu and Niu, Yazhe and Pu, Yuan and Hu, Shuai and Liu, Yu and Yang, Jing},
  journal={arXiv preprint arXiv:2404.16364},
  year={2024}
}

💓 Acknowledgments

This project has been developed partially based on the following pioneering works on GitHub repositories. We express our profound gratitude for these foundational resources:

We would like to extend our special thanks to the following contributors @PaParaZz1, @karroyan, @nighood, @jayyoung0802, @timothijoe, @TuTuHuss, @HarryXuancy, @puyuan1996, @HansBug for their valuable contributions and support to this algorithm library.

Thanks to all who contributed to this project:

🏷️ License

All code within this repository is under Apache License 2.0.

(Back to top)