TreePS-RAG: Tree-based Process Supervision for Reinforcement Learning in Agentic RAG

TL;DR

TreePS-RAG introduces an online, tree-structured reinforcement learning framework for agentic retrieval-augmented generation that enables dense, step-wise credit assignment without intermediate annotations or auxiliary rewards. By modeling agentic rollouts as a rollout tree and estimating node values via Monte Carlo averages over descendant leaves, it delivers fine-grained process advantages while maintaining an outcome-based objective. An efficient online tree construction strategy with depth-bound expansion and similarity-based pruning preserves exploration under a fixed compute budget. Across seven QA benchmarks and multiple backbones, TreePS-RAG consistently outperforms both outcome-only and existing process-supervised RL baselines, demonstrating the practical value of online, tree-based process supervision for complex reasoning-enabled retrieval tasks.

Abstract

Agentic retrieval-augmented generation (RAG) formulates question answering as a multi-step interaction between reasoning and information retrieval, and has recently been advanced by reinforcement learning (RL) with outcome-based supervision. While effective, relying solely on sparse final rewards limits step-wise credit assignment and provides weak guidance for intermediate reasoning and actions. Recent efforts explore process-level supervision, but typically depend on offline constructed training data, which risks distribution shift, or require costly intermediate annotations. We present TreePS-RAG, an online, tree-based RL framework for agentic RAG that enables step-wise credit assignment while retaining standard outcome-only rewards. Our key insight is to model agentic RAG reasoning as a rollout tree, where each reasoning step naturally maps to a node. This tree structure allows step utility to be estimated via Monte Carlo estimation over its descendant outcomes, yielding fine-grained process advantages without requiring intermediate labels. To make this paradigm practical, we introduce an efficient online tree construction strategy that preserves exploration diversity under a constrained computational budget. With a rollout cost comparable to strong baselines like Search-R1, experiments on seven multi-hop and general QA benchmarks across multiple model scales show that TreePS-RAG consistently and significantly outperforms both outcome-supervised and leading process-supervised RL methods.

Figures (5)

• Figure 1: Overview of TreePS-RAG. Left: Online tree construction and process supervision with $N=8$, $D=4$, $N_{\text{retain}}=2$. Retained parent nodes are expanded with $B_d$ continuations. Similarity-based clustering is applied among sibling search children to prune redundant nodes (crossed out) and retain $\hat{N}_{\text{retain}}$ diverse candidates (solid circles). Node values ($V$) are estimated via Monte Carlo averaging over descendant leaf outcomes to compute process advantages. Right: Tree-to-trajectory transformation for policy optimization. Root-to-leaf paths are collected into a set (one example highlighted in green). Process advantages are uniformly assigned to all model-generated tokens within a step $(r_i, a_i)$, while observation tokens $(o_i)$ are masked out. For fair comparison, only $N$ trajectories is randomly sampled during training to maintain comparable to standard outcome-based RL methods.
• Figure 2: Analysis of continuation-based reasoning.
• Figure 3: Template for TreePS-RAG. {input_question} will be replaced with specific question during training and inference. We follow jin2025searchr1trainingllmsreason and only replace <think> with <reason> to adapt for Qwen3 models.
• Figure 4: Case study 1 where TreePS-RAG correctly answers the question while Search-R1 fails, given the same reasoning prefix. In the retrieved documents, evidence supporting the correct answer is highlighted with red bold underlines, whereas the information related to Search-R1's erroneous answer is marked with black bold underlines.
• Figure 5: Case study 2 where TreePS-RAG correctly answers the question while Search-R1 fails to get needed information given the same reasoning prefix. In the retrieved documents, evidence supporting the correct answer is highlighted with black bold underlines.