TreeRL: LLM Reinforcement Learning with On-Policy Tree Search

Paper · arXiv 2506.11902 · Published June 13, 2025
Reinforcement LearningReward ModelsReasoning Model ArchitecturesDeep Research Agents

Reinforcement learning (RL) with tree search has demonstrated superior performance in traditional reasoning tasks. Compared to conventional independent chain sampling strategies with outcome supervision, tree search enables better exploration of the reasoning space and provides dense, on-policy process rewards during RL training but remains under-explored in On-Policy LLM RL. We propose TreeRL, a reinforcement learning framework that directly incorporates on-policy tree search for RL training. Our approach includes intermediate supervision and eliminates the need for separate reward model training. Existing approaches typically train a separate process reward model, which can suffer from distribution mismatch and reward hacking. We also introduce a cost-effective tree search approach that achieves higher search efficiency under the same generation token budget by strategically branching from high-uncertainty intermediate steps rather than using random branching. Experiments on challenging math and code reasoning benchmarks demonstrate that TreeRL achieves superior performance compared to traditional ChainRL, highlighting the potential of tree search for LLM. TreeRL is open-sourced at https://github.com/THUDM/TreeRL.

Introduction. Large language models (LLMs) have demonstrated remarkable capabilities across diverse complex reasoning tasks (Achiam et al., 2023; Team et al., 2023; Dubey et al., 2024), including mathematics (Shao et al., 2024b), programming (Lozhkov et al., 2024; Zhu et al., 2024), and autonomous agents (Zhou et al., 2024). Reinforcement learning (RL) has emerged as a powerful approach to significantly improve the reasoning abilities of LLMs by optimizing the policy through reward feedback (OpenAI, 2024; Guo et al., 2025; Hou et al., 2025; Shao et al., 2024b). Current RL methods for LLM training generally independently sample multiple trajectories (Shao et al., 2024b; Wang et al., 2024b; Touvron et al., 2023) and obtain reward signals based on the final answers. However, tree search, which has demonstrated significant success in other domains like AlphaZero (Silver et al., 2017), remains underdeveloped in reinforcement learning for LLM reasoning.

Discussion / Conclusion. This work presents TreeRL, an RL approach that combines tree search with process supervision to enhance LLM reasoning. EPTree improves response diversity and performance over traditional methods like MCTS and i.i.d multi-chain sampling. Then, we conduct reinforcement learning with EP- Tree and the derived process supervision from tree search. Experiments on math reasoning tasks show that the TreeRL outperforms existing techniques, highlighting the potential of RL with tree search to advance LLM in complex reasoning tasks. In this work, we propose to improve reinforcement learning with on-policy tree search. While this approach demonstrates promising performance, it does come with several limitations. First, current LLM inference engines do not offer special optimizations for tree search, meaning the proposed EPTree still requires 2+ iterations, resulting in a performance that is approximately 2× slower than multi-chain sampling.