Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning

This repository contains code and analysis for the paper: Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning. Below is the framework of our proposed method.

Environment Setup

1. [Option A] From .yaml File

• Create environment

  CONDA_OVERRIDE_CUDA=12.4 conda env create --file conda_recipe.yaml
  
  conda activate mcts-rl
  

• Install training packages

  pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu124
  
  pip install --no-cache-dir "vllm==0.8.5.post1" "torch==2.6.0" "torchvision==0.21.0" "torchaudio==2.6.0" "tensordict==0.6.2" torchdata "fsspec[http]<=2025.3.0,>=2023.1.0"
  
  DS_BUILD_FUSED_ADAM=1 DS_BUILD_CPU_ADAM=1 pip install --no-cache-dir deepspeed
  
  pip install -r requirements.txt
  

2. [Option B] Manual SetUp

• Create environment

  conda create -n mcts-rl python==3.10
  
  conda activate mcts-rl
  

• Install training packages

  pip install --no-cache-dir "vllm==0.8.5.post1" "torch==2.6.0" "torchvision==0.21.0" "torchaudio==2.6.0" "tensordict==0.6.2" torchdata
  
  DS_BUILD_FUSED_ADAM=1 DS_BUILD_CPU_ADAM=1 pip install --no-cache-dir deepspeed
  
  pip install "transformers[hf_xet]>=4.51.0" accelerate datasets peft hf-transfer "numpy<2.0.0" "pyarrow>=15.0.0" pandas ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler pytest py-spy tensorboard
  pip install "nvidia-ml-py>=12.560.30" "fastapi[standard]>=0.115.0" "optree>=0.13.0" "pydantic>=2.9" "grpcio>=1.62.1"
  
  pip install -r requirements.txt
  

3. For tree visualization

conda install --channel conda-forge pygraphviz

4. Huggingface Login

git config --global credential.helper store

from huggingface_hub import login
login()

Dataset Download

• Arithmo: akjindal53244/Arithmo-Data

• GSM8K: openai/grade-school-math

• MATH: hendrycks/math

• ARC: AI2 Reasoning Challenge

• AI2S: AI2 Science Questions

• OBQA: Openbook QA

• CSQA: tau/commonsense_qa

• SciQ: SciQ Dataset

Run MCTS-DPO

Our main code include ./mcts_rl/algorithms/mcts and ./mcts_rl/trainers/tsrl_trainer.py

Example script: run.sh

Citation

@article{xie2024monte,
  title={Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning},
  author={Xie, Yuxi and Goyal, Anirudh and Zheng, Wenyue and Kan, Min-Yen and Lillicrap, Timothy P and Kawaguchi, Kenji and Shieh, Michael},
  journal={arXiv preprint arXiv:2405.00451},
  year={2024}
}

---

_{^{This repository is adapted from the code of the works Safe-RLHF.}}