Tree-OPO: Off-policy Monte Carlo Tree-Guided Advantage Optimization for Multistep Reasoning

2025-09-11

AI Summary

To address advantage saturation and reward collapse in preference-based reinforcement learning (RL) for multi-step reasoning tasks, this paper proposes a policy optimization framework integrating Monte Carlo Tree Search (MCTS) with phased training. The method introduces a tree-structured advantage estimation scheme, leveraging high-quality reasoning trajectories from MCTS to construct prefix-conditioned reward signals for fine-grained assessment of composite reasoning quality. It further combines Group Relative Policy Optimization (GRPO) with an off-policy phased training paradigm to enhance update stability. Empirical results demonstrate that the approach significantly mitigates advantage saturation and reward collapse, improves consistency and generalization in multi-step reasoning, and achieves state-of-the-art performance on mathematical reasoning and code generation benchmarks.

Technology Category

• Knowledge Representation and Reasoning: Preferences
• Reasoning under Uncertainty: Stochastic Optimization
• Multiagent Systems: Multiagent Learning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

Recent advances in reasoning with large language models (LLMs) have shown the effectiveness of Monte Carlo Tree Search (MCTS) for generating high-quality intermediate trajectories, particularly in math and symbolic domains. Inspired by this, we explore how MCTS-derived trajectories, traditionally used for training value or reward models, can be repurposed to improve policy optimization in preference-based reinforcement learning (RL). Specifically, we focus on Group Relative Policy Optimization (GRPO), a recent algorithm that enables preference-consistent policy learning without value networks. We propose a staged GRPO training paradigm where completions are derived from partially revealed MCTS rollouts, introducing a novel tree-structured setting for advantage estimation. This leads to a rich class of prefix-conditioned reward signals, which we analyze theoretically and empirically. Our initial results indicate that while structured advantage estimation can stabilize updates and better reflect compositional reasoning quality, challenges such as advantage saturation and reward signal collapse remain. We propose heuristic and statistical solutions to mitigate these issues and discuss open challenges for learning under staged or tree-like reward structures.

Problem

• Enhance policy optimization using MCTS trajectories
• Address advantage saturation in staged GRPO training
• Mitigate reward signal collapse in tree-structured reasoning

Innovation

• Monte Carlo Tree Search guided advantage optimization
• Staged GRPO training with MCTS rollouts
• Tree-structured advantage estimation for reasoning