Efficient Monte Carlo Tree Search via On-the-Fly State-Conditioned Action Abstraction

This repository contains the official implementation of the publication: Efficient Monte Carlo Tree Search via On-the-Fly State-Conditioned Action Abstraction, Yunhyeok Kwak, Inwoo Hwang, Dooyoung Kim, Sanghack Lee, Byoung-Tak Zhang, The 40th Conference on Uncertainty in Artificial Intelligence, 2024.

📦 Installation

Docker image

The recommended setup is the published Docker image:

docker pull yunkwak/efficient-mcts:1.0

For an exactly pinned image, use the digest:

docker pull yunkwak/efficient-mcts@sha256:b50c57d2d842b406affeee73413d9e926ed827c4e1ca4d699a1cfd658457a256

The image contains the Python, JAX, CUDA, and Python package dependencies. It does not bundle a checkout of this repository. Clone the repository on the host, mount it into the container, then install the local pine package in editable mode:

git clone https://github.com/yun-kwak/efficient-mcts.git
cd efficient-mcts

docker run --rm --gpus all -it \
  -v "$PWD":/workspace/efficient-mcts \
  -w /workspace/efficient-mcts \
  -e WANDB_MODE=disabled \
  yunkwak/efficient-mcts:1.0 \
  bash

pip install -e pine
python -m pytest -q pine/tests

🚀 Quick Start

Run the unit tests:

python -m pytest -q pine/tests

Run one paper preset:

WANDB_MODE=disabled python scripts/run_paper_preset.py \
  --task doorkey_easy \
  --method ours \
  --seed 1 \
  --save_dir ./outputs

Run the default Sokoban experiment:

python scripts/run_experiment.py \
  --env_id Sokoban-PushAndPull-7x7-B1-C3 \
  --exp_name sokoban-push-pull \
  --run_name seed1

✒️ Citation

If you use this code in your research, please cite the following paper:

@inproceedings{
    kwak2024efficient,
    title={Efficient Monte Carlo Tree Search via On-the-Fly State-Conditioned Action Abstraction},
    author={Yunhyeok Kwak and Inwoo Hwang and Dooyoung Kim and Sanghack Lee and Byoung-Tak Zhang},
    booktitle={The 40th Conference on Uncertainty in Artificial Intelligence},
    year={2024},
    url={https://openreview.net/forum?id=UvDsWevxUI}
}

📖 Credits

This repository is based on the following repositories:

• jax_muzero
• Haiku
• Mctx
• Flax
• Monte Carlo Tree Search With Iteratively Refining State Abstractions, NeurIPS 2021, Sokota et al.