Tree Search for LLM Agent Reinforcement Learning

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

ICLR 2026; tree-search RL training directly improves long-horizon LLM agent policies

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Tree-GRPO replaces chain-based independent rollout sampling in LLM agent RL with an agent-step-level tree search, enabling ~1.5× more rollouts under the same token/tool-call budget. The tree structure naturally yields process-level supervision signals from outcome rewards alone, equivalent in gradient structure to step-level DPO. Across 11 benchmarks, Tree-GRPO consistently outperforms chain-based GRPO and GSPO, with relative gains up to 69% on small models.

Problem

Multi-turn agentic RL suffers from two compounding problems: (1) heavy rollout cost — independent chain-based sampling redundantly re-runs entire trajectories, making rollout the dominant training cost; (2) sparse supervision — outcome rewards assign identical credit to every step in a long trajectory, making it hard to attribute success or failure to specific actions, leading to training instability and collapse.

Method

Tree-GRPO builds a tree where each node is a complete Thought-Action-Observation step (not token/sentence level). Training proceeds in three phases repeated \(L\) times:

1. Initialization: Sample \(M\) independent chain trajectories from policy \(\pi_\theta\) in parallel.
2. Sampling: Randomly select \(N\) non-leaf nodes from each tree \(T_i\) for expansion.
3. Expansion: From each selected node with context \(H_{<t}\), continue rollout with \(N\) new branches, appending them to the tree.

Total rollouts per prompt: \(G = M \times (L \times N + 1)\). Expected token/tool-call budget: $\(\mathbb{E}[B_{\text{tree}}] = M \cdot B + L \cdot N \cdot B/2\)$ versus \(G \cdot B\) for chain-based, yielding ~\(2\times\) savings in budget for the same group size.

Advantage estimation operates at two levels: - Intra-tree: siblings at each branch point share a prefix; their outcome reward differences provide step-level process signals. Advantage is normalized within-tree: \(\hat{A}^{\text{Intra}} = (R(H_i) - \mu_{\text{tree}}) / \sigma_{\text{tree}}\). - Inter-tree: standard group relative normalization across all \(M\) trees for a stable baseline. - Final advantage: \(\hat{A}^{\text{tree}}(H_i) = \hat{A}^{\text{Intra-tree}}(H_i) + \hat{A}^{\text{Inter-tree}}(H_i)\).

Theoretical result: Under a binary preference assumption, the gradient of intra-tree GRPO is structurally equivalent to step-level DPO — both optimize \(w \cdot [\nabla_\theta \log p_\theta(H^{\text{win}}_{\geq t}) - \nabla_\theta \log p_\theta(H^{\text{loss}}_{\geq t})]\) — differing only in the weight term \(w\).

Key Contributions

• Agent-step-level tree search rollout strategy replacing independent chain sampling, with explicit token and tool-call budget control.
• Intra-tree + inter-tree group relative advantage estimation that extracts step-level process signals from outcome rewards without additional annotation.
• Formal proof that intra-tree GRPO is structurally equivalent to step-level DPO (Proposition 3.1).
• Empirical demonstration that Tree-GRPO can train a base model into a multi-turn ReAct agent without SFT, even under extremely limited rollout budgets.

Results

• Multi-Hop QA (EM), small models: Tree-GRPO vs. chain GRPO — Qwen2.5-1.5b: +9.5% relative on multi-hop avg. (47.5 vs. 43.4) and +69% relative on single-hop avg. (19.1 vs. 11.3); Llama3.2-3b: +2.7% multi-hop, +38% single-hop.
• Budget ablation (Qwen2.5-3b): At budget≈2 (quarter of standard), Tree-GRPO achieves 36.8 avg. multi-hop EM vs. chain-based 31.8 (+16% rel.) while GRPO at budget≈4 scores 36.8 — matching chain-based at 2× the cost.
• Web-Agent QA (F1): Qwen2.5-14b Tree-GRPO: 21.0 WebWalkerQA avg. vs. GRPO 16.4 (+28%); 12.8 vs. 12.4 overall avg.; GAIA improvements drive most gains.
• Large models (Qwen2.5-14b, 7b): gains narrow (+1.1–8.4% relative) but remain consistent.
• Tree-GRPO achieves comparable or superior performance to chain-based GRPO at one-quarter the rollout budget (Qwen2.5-3b, budget≈2 vs. budget≈8).

Limitations

• Gains diminish for large models (14b+), suggesting the process signal benefit is most significant when model capacity is a bottleneck.
• Single-hop tasks with shallow tree depth (typically 2 turns) yield limited benefit from process signals.
• Web-agent performance is constrained by scarcity and quality mismatch of available training data; RL improvements on very hard benchmarks (BrowseComp) are marginal.
• Tree search requires a sequential expand-then-parallelize scheme that is less naturally suited to fully parallelized LLM inference engines than chain-based sampling.
• The tree configuration hyperparameters (\(M\), \(N\), \(L\)) affect exploration-exploitation trade-offs and require tuning.

Relevance to Agentic AI / LLM Agents

Tree-GRPO directly addresses one of the core bottlenecks in scaling agentic RL: the extreme token and API-call cost of rollout collection combined with the sparsity of reward signal in long-horizon multi-turn trajectories. By demonstrating that process-level supervision can be recovered from outcome rewards via tree structure alone — without process reward models or human annotation — it offers a plug-and-play upgrade to any GRPO-family agentic RL pipeline. The theoretical connection to step-level DPO grounds the empirical gains and links online RL with offline preference learning, which is an active convergence point in agent training methodology. The ability to train small base models (1.5b–3b) into functional ReAct agents under severe budget constraints makes this especially relevant for resource-constrained agentic deployments.