Implementation for ACL'25 paper TreeRL: LLM Reinforcement Learning with On-Policy Tree Search. The implementation is based on OpenRLHF

TreeRL is a reinforcement learning framework that directly incorporates on-policy tree search for training, eliminating the need for separate reward model training while providing better exploration of reasoning space through strategic branching from high-uncertainty steps. Experiments on math and code reasoning benchmarks demonstrate that TreeRL can achieve consistent better performance than ChainRL under the same expeeriment setting.

Currently, we provide the following code of PRIME, you can find more details in each directory.

• scripts provides the script to start the training. And the RL implemenration is in openrlhf.
• EPTree: the tree sampling implementation.
• The training data can be found in datasets(datasets). We use the data from T1 for SFT.

If you find this work is helpful to your research, please consider citing our paper:

@inproceedings{treerl,
  title={TreeRL: LLM Reinforcement Learning with On-Policy Tree Search},
  author={Hou, Zhenyu* and Hu, Ziniu* and Li, Yujiang* and Lu, Rui* and Tang, Jie and Dong, Yuxiao},
  booktitle={ACL},
  year={2025}
}