/LightZero

LightZero: A lightweight and efficient MCTS/AlphaZero/MuZero algorithm toolkit.

Primary LanguagePythonApache License 2.0Apache-2.0

LightZero

Twitter PyPI PyPI - Python Version Loc Comments

Docs Deploy Code Test Badge Creation Package Release codecov

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

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

English | 简体中文

Background

The method of combining Monte Carlo Tree Search and Deep Reinforcement Learning represented by AlphaZero and MuZero has achieved superhuman level in various games such as Go and Atari, has also made gratifying progress in scientific fields such as protein structure prediction, matrix multiplication algorithm search, etc. 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 toolkits that combines MCTS and RL for PyTorch. It provides support for a range of MCTS-based RL algorithms and applications with 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 documentations 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./Alg. AlphaZero MuZero EfficientZero Sampled EfficientZero
Atari --- ✔ ✔ ✔
tictactoe ✔ ✔ 🔒 🔒
gomoku ✔ ✔ 🔒 🔒
chess 🔒 🔒 🔒 🔒
go 🔒 🔒 🔒 🔒
lunarlander --- ✔ ✔ ✔
bipedalwalker --- ✔ ✔ ✔
cartpole --- ✔ ✔ ✔
pendulum --- ✔ ✔ ✔

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

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

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

Installation

You can install 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 .

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

Below are the benchmark results of MuZero, MuZero w/ SSL , EfficientZero and Sampled EfficientZero on three discrete action space games in Atari.

Image Description 1 Image Description 2 Image Description 3 Image Description 4

Below are the benchmark results of Sampled EfficientZero with Factored/Gaussian policy representation on three continuous action space games: Pendulum-v1, LunarLanderContinuous-v2, BipedalWalker-v3.

Where Factored Policy indicates that the agent learns a policy network that outputs a categorical distribution, the dimensions of the action space for the three environments are 11, 49 (7^2), and 256 (4^4), respectively, after manual discretization. On the other hand, Gaussian Policy indicates that the agent learns a policy network that outputs parameters (mu and sigma) for a Gaussian distribution.

Image Description 1 Image Description 2 Image Description 3 Image Description 3

Below are the benchmark results of AlphaZero and MuZero on two board_games: TicTacToe, Gomoku.

Image Description 1 Image Description 2

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

mcts

alphazero

muzero

efficientzero

sampled muzero

gumbel muzero

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 feedbacks and contributions to improve LightZero, both algorithms and system designs.

Citation

@misc{lightzero,
    title={{LightZero: OpenDILab} A lightweight and efficient MCTS/AlphaZero/MuZero algorithm toolkits.},
    author={LightZero Contributors},
    publisher = {GitHub},
    howpublished = {\url{https://github.com/opendilab/LightZero}},
    year={2023},
}

Acknowledgments

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

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

License

All code within this repository is under Apache License 2.0.