Tree Search for LLM Agent Reinforcement Learning
Paper: arXiv:2509.21240 Code: AMAP-ML/Tree-GRPO Code reference:
main@19bf3fa(2026-01-27)
1. 论文一句话总结
这篇论文提出 Tree-GRPO：把 LLM agent RL 中原本彼此独立的 chain rollout 改成以 ReAct agent step 为节点的 tree-search rollout，通过共享前缀降低 token / tool-call 预算，并把 outcome-only reward 反推出 intra-tree 与 inter-tree 两级 group relative advantage，从而在不额外标注 step reward 的前提下得到类似 step-level DPO 的过程监督信号。
核心新意不是简单地“用树搜索增强采样”，而是把树结构放进在线 RL 的 credit assignment：同一棵树里的不同分支天然共享历史前缀，分支之后的叶子 outcome reward 差异可以视作对这个分叉点后续行为的偏好信号；跨树分组又提供更稳定的全局 baseline。论文把这两个信号相加，仍然落在 GRPO/PPO-style clipped objective 里，因此可以复用现有 veRL / Search-R1 类训练栈。
2. 研究问题：为什么 chain-based agent RL 不够
LLM agent 的多轮任务通常不是一次性生成答案，而是按 thought -> action/tool call -> observation 的循环推进。论文把一个完整 agent trajectory 写作 (\mathcal{H}={(\tau_0,\alpha_0,o_0),\ldots,(\tau_T,\alpha_T,o_T)})，其中 (\tau) 是模型思考文本，(\alpha) 是动作或工具调用，(o) 是环境观察。训练时最常见的 reward 仍是最终答案是否正确、F1/EM 是否高、web QA 是否命中，也就是只有 trajectory 末端才有 outcome reward。
这种设置带来两个耦合问题。第一是 rollout 成本高：chain-based GRPO 为同一个 prompt 采样多条完整 trajectory，每条都要重新走完整工具调用和上下文生成，公共前缀无法复用。第二是 credit assignment 粗糙：同一条 trajectory 内所有 step 被赋予同一个 advantage；如果最终错了，模型无法知道是检索 query 错、工具调用错、阅读错，还是最后答案组织错。对于长 horizon agent，这种稀疏监督会让小模型尤其难以稳定学会多轮工具使用。
Tree-GRPO 的定位是：不引入外部过程标注、不训练额外 verifier，也不把 test-time MCTS 直接搬到训练中；它只改变 rollout 组织方式和 advantage 分组方式，使相同 token/tool 预算下生成更多可比较的分支，并从分支间 outcome 差异中提取过程信号。
3. 方法总览
3.1 Tree-search rollout：以完整 agent step 为节点
论文没有把 token 或 sentence 当作树节点，而是把一个完整 agent 交互 step ((\tau,\alpha,o)) 当作节点。这个设计很关键：LLM agent 的可解释决策边界通常不是 token，而是一次“想法 + 动作 + 环境反馈”。在这个粒度上扩展，既能保留语义完整性，也更容易把工具调用、检索观察和后续生成对齐到同一个 step。
采样流程分三步：
- Initialization：对每个 prompt (x_i)，先并行生成 (M) 条独立 chain trajectory，作为 (M) 棵树的初始骨架。
- Sampling：从每棵树中随机采样 (N) 个非叶节点作为 expansion node。非叶节点意味着它已经包含从 root 到当前 step 的上下文，但还可以从这里继续生成新的后续轨迹。
- Expansion：把 root 到该节点的完整上下文 (\mathcal{H}{<t}) 与原 prompt 输入给当前 policy (\pi\theta)，继续生成新的 suffix (\mathcal{H}_{\ge t})，并把它插回原树形成新分支。重复 (L) 轮后，每个 prompt 约得到 (M\times(L\times N+1)) 条 rollout。
论文给出的期望预算直觉是：如果一条完整 trajectory 的 token/tool-call 成本为 (B)，初始化 (M) 条链要 (M B)；一次随机 expansion 平均从中间节点继续，约花 (B/2)；因此 tree sampling 的期望预算为
[ \mathbb{E}[B_{tree}] = M\cdot B + L\cdot N\cdot B/2. ]
这解释了为什么树搜索能在近似相同预算下拿到更多 rollout：新分支复用了已生成的公共前缀。论文默认设置中 chain-based RL group size 为 4，而 Tree-GRPO 常用 ((M=2,N=2,L=1))，最终每个 prompt 可以形成 2 棵树、每棵树约 3 个候选叶子，对应总共 6 个可训练 rollout，但预算表仍按约 4 条完整 trajectory 的成本比较。
3.2 Tree-based group relative advantage
chain-based GRPO 对同一 prompt 下的完整 trajectories 做 group normalization。问题是 trajectory 级 advantage 会被广播到所有 token/step：
[ A(\mathcal{H})=A({(\tau_0,\alpha_0,o_0),\ldots,(\tau_T,\alpha_T,o_T)}), ]
也就是每个中间行为都拿同一个信用。这对长任务很不友好。
Tree-GRPO 把分组分成两级：
- Intra-tree group：在同一棵树内比较共享公共前缀的不同分支。由于这些分支在某个节点之后才分叉，叶子 outcome reward 的差异可以更接近“这一步之后的选择好坏”。
- Inter-tree group：在同一 prompt 的所有树/rollout 间比较，提供更稳定的全局 baseline，缓解单棵树内样本少导致的 normalization 不可靠。
论文将两者写成统一形式：
[ \hat{A}{Intra/Inter-tree}(\mathcal{H}^i)= \frac{R(\mathcal{H}^i)-\mathrm{mean}({R(\mathcal{H}^j)}{j}^{G_{Intra/Inter-tree}(\mathcal{T}i)})} {\mathrm{std}({R(\mathcal{H}^j)}{j}^{G_{Intra/Inter-tree}(\mathcal{T}_i)})}. ]
最终 advantage 是加和：
[ \hat{A}{tree}(\mathcal{H}^i)=\hat{A}{Intra-tree}(\mathcal{H}^i)+\hat{A}_{Inter-tree}(\mathcal{H}^i). ]
直觉上，intra-tree 告诉模型“在同一个上下文分叉点之后，哪个 suffix 更好”；inter-tree 告诉模型“相对同 prompt 的其他独立探索，这条 trajectory 是否整体更好”。前者偏过程 credit，后者偏稳定 baseline。论文的 ablation 也显示只用 intra-tree 在预算 4 的 Qwen2.5-3B multi-hop 设置下会 collapse，只用 inter-tree 有提升但不如两者相加。
3.3 Policy objective
Tree-GRPO 保持 PPO/GRPO 风格的 clipped objective。每条 tree-search rollout (\mathcal{H}^i) 的每个 token/step 使用同一个 (\hat{A}{tree}(\mathcal{H}^i))，并计算 ratio (r{i,t}(\theta))、clip、KL penalty。方法层面的变化在 rollout 生成和 advantage 估计，而不是重新发明 optimizer。
这个选择有工程意义：如果已有系统能跑 GRPO，只要 rollout worker 能返回 tree_uid、advantage estimator 能读到 intra/inter 分组，就能把 Tree-GRPO 接进训练。论文开源实现也基本走这个路线：训练脚本设置 algorithm.adv_estimator=tree，rollout 侧打开 actor_rollout_ref.rollout.tree_search=true，trainer 侧把 uid 和 tree_uid 两个 key 分别用于 inter-tree 和 intra-tree normalization。
3.4 理论解释：intra-tree GRPO 近似 step-level preference learning
论文给了一个二元偏好设定：在某个中间节点 ((x,\mathcal{H}_{<t})) 后，后续 trajectory 分成 winning suffix 和 losing suffix，reward 分别为 1 和 0。step-level DPO 会最大化 winning vs losing 的 Bradley-Terry likelihood，梯度包含
[ \nabla_\theta \log p_\theta(\mathcal{H}^{win}{\ge t})- \nabla\theta \log p_\theta(\mathcal{H}^{loss}_{\ge t}). ]
论文的 Proposition 说明，在这个 setting 下，intra-tree GRPO 和 step-level DPO 可以写成同一种 gradient structure：
[ \nabla_\theta J_{unified}(\theta)=w\cdot \left(\nabla_\theta \log p_\theta(\mathcal{H}^{win}{\ge t})- \nabla\theta \log p_\theta(\mathcal{H}^{loss}_{\ge t})\right), ]
差别只在 weight term (w)。因此 Tree-GRPO 的 intra-tree 部分可以被理解为：在在线采样过程中，利用同一 tree branching point 下的 outcome 差异隐式构造 step-level preference objective。这个结论不是说 Tree-GRPO 完全等价于显式 DPO 数据训练，而是说明它的优化方向含有“偏好更好 suffix、压低更差 suffix”的过程监督成分。
4. 算法伪代码
4.1 Tree-search rollout
# input: prompt x, policy pi, tree params M, N, L, final leaves K
roots = []
for m in range(M):
root = TreeNode(prompt=x, depth=0)
chain = rollout_full_agent_trajectory(pi, x) # thought/action/obs steps
root.insert_chain(chain)
roots.append(root)
for _ in range(L):
expansion_nodes = []
for root in roots:
expansion_nodes += root.sample_non_leaf_nodes(N, mode="random")
for node in expansion_nodes:
prefix = node.path_from_root() # shared context H_<t
suffix = rollout_suffix(pi, x, prefix) # new H_>=t
node.insert_branch(suffix)
final_leaves = []
for root in roots:
final_leaves += root.sample_leaf(K)
return final_leaves, tree_uid_for_each_leaf4.2 Advantage 计算
# rewards: outcome reward for final leaves under the same prompt
for leaf in leaves:
inter_group = leaves_with_same_prompt_uid(leaf.uid)
intra_group = leaves_with_same_tree_uid(leaf.tree_uid)
A_inter = normalize(leaf.reward, inter_group.rewards)
A_intra = normalize(leaf.reward, intra_group.rewards)
leaf.advantage = A_inter + A_intra这里的关键不是 leaf 数量本身，而是 tree_uid 保留了分支来源，使训练端能区分“全 prompt 内比较”和“同树内比较”。如果没有 tree_uid，tree sampling 只会退化成更多样本的 GRPO，无法得到论文强调的过程偏好信号。
5. 实验设置与结果
5.1 数据、任务和指标
论文在 11 个 benchmark、3 类 QA 任务上评估：
- Single-Hop QA：NQ、TriviaQA、PopQA，指标为 EM。
- Multi-Hop QA：HotpotQA、2WikiMultiHopQA、MuSiQue、Bamboogle，指标为 EM。
- Web-Agent QA：SimpleQA、GAIA、WebWalkerQA、BrowseComp，指标为 F1。
训练集按任务类型分开。single-hop 使用 NQ train split；multi-hop 混合 HotpotQA、2WikiMultiHopQA、MuSiQue；web-agent 训练集由 ASearcher-35K 采样 2k 条过滤样本加 WebDancer 释放的 200 条 hard web QA 构成。模型包括 Qwen2.5 1.5B/3B/7B/14B 与 Llama3.2-3B；默认每个 prompt 的 rollout budget 约为 4。baseline 包括 Direct Inference、Search-o1、ReAct、GRPO 与 GSPO。
5.2 主结果：小模型和 multi-hop 提升最明显
| 模型 | 方法 | Single-Hop Avg | Multi-Hop Avg | 主要变化 |
|---|---|---|---|---|
| Qwen2.5-1.5B | GRPO | 43.4 | 11.3 | baseline |
| Qwen2.5-1.5B | Tree-GRPO | 47.5 | 19.1 | single-hop +9.5%，multi-hop +69% |
| Qwen2.5-3B | GRPO | 48.1 | 31.8 | baseline |
| Qwen2.5-3B | Tree-GRPO | 50.0 | 36.8 | single-hop +4.0%，multi-hop +16% |
| Qwen2.5-7B | GRPO | 50.5 | 36.4 | baseline |
| Qwen2.5-7B | Tree-GRPO | 52.2 | 37.8 | single-hop +3.4%，multi-hop +3.9% |
| Qwen2.5-14B | GRPO | 55.1 | 41.8 | baseline |
| Qwen2.5-14B | Tree-GRPO | 55.7 | 45.3 | single-hop +1.1%，multi-hop +8.4% |
Qwen2.5 系列上，Tree-GRPO 相比 GRPO 在所有模型规模上都提升，但收益分布不均：multi-hop 和小模型最明显。Qwen2.5-1.5B 的 multi-hop average 从 11.3 到 19.1，说明 tree-based process signal 对“需要多轮检索/推理”的 setting 最有价值；14B 上 single-hop 只涨 1.1%，说明大模型和短任务里 chain-based RL 的稀疏监督问题相对没那么严重。
Llama3.2-3B 上也有类似趋势：GRPO single-hop avg 48.7、multi-hop avg 26.7；Tree-GRPO 达到 50.0 和 36.8，multi-hop 相对提升 38%。这支持论文的 claim：Tree-GRPO 不是 Qwen-specific trick，而是对 agent RL 的采样/信用分配问题做了更通用的改造。
5.3 Web-agent QA 结果
Web-agent QA 的表格按 SimpleQA、GAIA Lv.1/2/3、WebWalkerQA Easy/Med/Hard、BrowseComp 报 F1。Qwen2.5-7B 上，GRPO 在 SimpleQA/GAIA avg/WebWalkerQA avg/BrowseComp 分别为 61.5 / 14.7 / 10.9 / 2.3，Tree-GRPO 为 62.4 / 16.8 / 11.2 / 2.7。Qwen2.5-14B 上，GRPO 为 65.4 / 16.4 / 12.4 / 2.4，Tree-GRPO 为 67.8 / 21.0 / 12.8 / 2.6。提升最明显的是 GAIA 平均分，尤其 Qwen2.5-14B 从 16.4 到 21.0，符合论文说法：Tree-GRPO 对需要 web browsing、工具使用和多步信息整合的问题更有帮助。
5.4 rollout budget 对比
| 预算 | Chain-based Multi-Hop Avg | Tree-based 设置 | Tree-based Multi-Hop Avg | 结论 |
|---|---|---|---|---|
| ≈2/prompt | 14.9 | M=1,N=2,L=1 | 31.6 | 树搜索在很低预算下复用前缀，收益最大 |
| ≈4/prompt | 31.8 | M=2,N=2,L=1 | 36.8 | 默认设置，multi-hop +16% |
| ≈8/prompt | 31.4 | M=4,N=2,L=1 | 35.1 | 仍优于 chain，但边际收益下降 |
| ≈16/prompt | 33.9 | M=4,N=5,L=1 | 37.3 | 最好树结构仍有 +10% |
论文最强的效率 claim 是：Tree-GRPO 在约 4/prompt 的预算下已经超过 chain-based 约 16/prompt 的结果。例如 Qwen2.5-3B multi-hop 上，Tree-GRPO ((M=2,N=2,L=1)) 得到 36.8，而 chain-based 预算 16 只有 33.9。这就是 README 中“only a quarter of the rollout budget”的来源。
5.5 ablation：intra + inter 必须一起用
在 Qwen2.5-3B multi-hop 上，预算 ≈4/prompt 时只用 (A_{Intra-tree}) 发生 collapse；只用 (A_{Inter-tree}) 得到 33.8；两者相加得到 36.8。预算 ≈8/prompt 时，只用 intra 得到 34.6，只用 inter 得到 33.4，两者相加为 36.4。这个结果说明：intra-tree 的过程偏好很重要，但样本少时 baseline 方差大；inter-tree 提供稳定性；二者组合才是 Tree-GRPO 的关键。
另一个重要 ablation 是树节点粒度。论文比较 token/sentence-level tree search 与 agent-step-level tree search，发现 token/sentence 粒度在 single-hop 和 multi-hop 都更差，甚至会低于 vanilla GRPO。原因很直观：agent 的可行动决策边界是完整工具调用和环境反馈，token 层分叉会破坏语义完整性，也更难把 outcome reward 映射到有意义的过程 credit。
6. 开源代码搜索与实现映射
代码搜索结果：论文 arXiv 页面和 GitHub 搜索均指向公开仓库 AMAP-ML/Tree-GRPO。本笔记基于 main@19bf3fa（commit date 2026-01-27）检查代码。仓库基于 Search-R1 与 veRL，核心改动集中在 tree-search rollout manager、TreeNode、Tree-GRPO advantage estimator 和训练脚本。
| 论文概念 | 代码位置 | 证据与解释 |
|---|---|---|
| Tree-GRPO 训练入口 | train_multihopqa_tree_search.sh, train_singlehopqa_tree_search.sh, train_webagent_tree_search.sh | 训练脚本设置 tree_search_m=2, tree_search_n=2, tree_search_l=1, tree_search_k=3；设置 algorithm.adv_estimator=tree 与 actor_rollout_ref.rollout.tree_search=true。 |
| chain-based baseline | train_multihopqa_grpo.sh 等 | GRPO baseline 设置 algorithm.adv_estimator=grpo，rollout 不开 tree_search，并设置 actor_rollout_ref.rollout.n_agent=4，对应 group size 4。 |
| tree rollout 参数 | search_r1/llm_agent/generation_ts.py | GenerationConfig 中 m 表示树数量，n 表示每轮扩展节点数，l 表示 expansion 轮数，k 表示每棵树最终采样叶子数，expand_mode='random'。 |
| 初始化与扩展 | generation_ts.py::run_llm_loop_tree_search | 先为每个 prompt 创建 root node，再调用 gen_action_chain 得到初始链；每轮从每棵树调用 root.get_expand_node(n)，把节点上下文重新 pad 成 batch 后继续生成新分支。 |
| 叶子采样与 reward | generation_ts.py, tree_node.py | rollout 结束后每棵树 sample_leaf(k)；reward_fn(final_output) 计算叶子原始 outcome score；reward_mode=base 时 token-level score 只放在有效 response 的最后一个 token。 |
| intra/inter advantage | verl/trainer/ppo/ray_trainer_ts.py | 当 adv_estimator == 'tree'，代码分别用 uid 计算 inter-tree advantages，用 tree_uid 计算 inner/intra-tree advantages，并把二者相加写入 data.batch['advantages']。 |
| outcome normalization | verl/trainer/ppo/core_algos.py | compute_grpo_outcome_advantage 对同组 outcome rewards 做 mean/std normalization；compute_tree_grpo_outcome_advantage 还提供两次 normalization 的 tree_2norm 变体，但默认训练脚本用的是 tree。 |
6.1 训练配置细节
以 train_multihopqa_tree_search.sh 为例，实际训练配置不是从 base config 推测出来的，而是脚本显式覆盖：8 GPU，train_batch_size=8*64=512，val_batch_size=8*16=128，actor_ppo_mini_batch_size=8*8=64，actor_ppo_micro_batch_size=8*4=32；prompt/response/observation 长度分别设置为 max_prompt_length=4096、max_response_length=500、max_start_length=2048、max_obs_length=500；actor LR 由脚本变量控制，KL loss coef 为 0.001，rollout temperature 为 1，state_masking=true，rollout.reward_mode=base，rollout.expand_mode=random。
这点很重要：论文方法中“树结构产生过程监督”的理论解释容易让人误以为代码会把每个中间节点都显式打 step reward；实际默认训练脚本采用 reward_mode=base，reward 仍然来自叶子 outcome。过程信号主要通过 tree_uid 分组后的 advantage normalization 体现，而不是通过额外 reward model。
6.2 代码级伪代码：当前仓库默认路径
# train_*_tree_search.sh
algorithm.adv_estimator = "tree"
rollout.tree_search = True
rollout.ts_m = 2
rollout.ts_n = 2
rollout.ts_l = 1
rollout.ts_k = 3
rollout.reward_mode = "base"
rollout.expand_mode = "random"
# ray_trainer_ts.py
if adv_estimator == "tree":
inter_A = compute_grpo_outcome_advantage(rewards, index=uid)
intra_A = compute_grpo_outcome_advantage(rewards, index=tree_uid)
advantages = inter_A + intra_A这里 uid 是 prompt-level group id，tree_uid 是树内 group id。这样实现与论文公式完全对应：inter-tree 是同 prompt 下所有 rollout 的 group relative advantage，intra-tree 是同一棵树内的 group relative advantage。工程上，这比在 rollout 阶段维护复杂的 branch reward backprop 简洁得多，也更容易接到现有 GRPO trainer。
7. 图表解读
tree_chain_comp 展示了 chain-based 与 tree-based rollout 的本质区别：chain-based 每条轨迹从头生成，只有 outcome reward；tree-based 共享 prefix，在每个 branching point 都能比较后续 suffix 的好坏。这个图是理解论文的关键。
tree_search_comp_plot 对应“树搜索粒度”实验，说明 agent-step node 优于 token/sentence node。它支持论文的一个隐含设计原则：agent RL 中的 credit assignment 应该对齐 action/observation 边界，而不是盲目追求更细的 token 粒度。
8. 局限与复现注意
第一，Tree-GRPO 的收益依赖任务是否真的有可复用前缀和可比较分支。对于单轮短答或本身不需要工具使用的任务，tree rollout 的额外实现复杂度未必划算；主表也显示 14B single-hop 上只有小幅提升。
第二，intra-tree signal 可能高方差。论文 ablation 中预算 4 时只用 intra-tree 会 collapse，说明“同树内过程偏好”本身不一定稳定，必须配合 inter-tree baseline。实际复现时如果改变 (M,N,L,K)，不能只看最终 rollout 数量，还要看每棵树内可用于 normalization 的叶子数是否足够。
第三，树扩展策略当前开源实现默认是 random。论文没有系统探索 learned expansion policy、uncertainty-guided node selection 或 value-guided search；因此 Tree-GRPO 更像是一个可并行的 tree sampling / credit assignment 框架，而不是完整的 MCTS-style planning algorithm。
第四，代码中默认 reward_mode=base，也就是 outcome reward 仍主要落在最终有效 token。论文中的“step-level process supervision”应理解为由 tree grouping 产生的 implicit process signal，而不是显式逐步 reward 标注。如果下游任务有密集过程 reward，Tree-GRPO 可能还可以叠加，但这不是本文实验的主要设置。
第五，web-agent 训练数据尚有开放性。GitHub README 写到 web-agent QA 的训练和测试集会继续释放，复现 web-agent 表格可能比 single-hop/multi-hop 更依赖作者后续数据发布、搜索 API、缓存和环境配置。
9. 我对这篇论文的定位
这篇论文适合放在 Agent/RL，因为它的主要贡献不是通用 LLM post-training，也不是 test-time tree search，而是面向多轮工具使用 agent 的在线 RL 采样和信用分配。它和普通 GRPO/GSPO 的差别在于：优化器外壳相似，但 rollout topology 和 advantage grouping 改变了监督粒度。
如果把它和 Agent Lightning、ARTIST、Planner-R1 等 agent RL 工作比较，Tree-GRPO 的优点是非常“基础设施友好”：不要求把 agent runtime 拆成复杂 MDP 训练接口，也不要求额外 process label；只要 rollout worker 能保存树结构，trainer 能读取 tree_uid，就可以在现有 GRPO pipeline 上实验。代价是 tree rollout manager 明显更复杂，且随机扩展策略是否最优仍未解决。
10. 可迁移启发
对自己的 agent RL 系统，这篇论文最值得借鉴的是三点：
- 把采样组织成可复用前缀的结构：即便不做完整 Tree-GRPO，只要任务有多轮工具调用，复用 prefix 就能显著降低 rollout 成本。
- 把 advantage 分组从 prompt 维度扩展到 trajectory topology 维度：
uid之外再加tree_uid、branch_uid或prefix_uid，可以在 outcome-only reward 下构造更细的相对比较。 - 优先对齐语义 step，而不是 token：对于工具调用 agent，step-level 节点比 token-level 节点更符合环境交互边界，也更利于解释错误来源。
如果要继续推进，可以尝试把 random expansion 改成 uncertainty/value guided expansion，把失败分支的中间 observation 用于 error taxonomy，或者在已有 dense reward 场景中比较 reward_mode=base 与 reward_mode=tree_diff。但这些都应该建立在本文的基本 observation 上：tree structure 本身已经能从 outcome reward 中提取一部分过程监督。