基于树搜索的智能体轮次策略优化 AT^2PO: Agentic Turn-based Policy Optimization via Tree Search
提出AT2PO框架,通过熵引导树扩展、轮次级信用分配和轮次策略优化解决多轮智能体强化学习三大核心挑战
前置知识
智能体强化学习
智能体强化学习是将强化学习应用于大语言模型智能体的方法,通过优化策略让智能体在多轮交互中学习如何更好地推理和调用工具。智能体通过与环境(包括搜索引擎、数据库等工具)进行多轮交互,根据最终获得的奖励信号调整自己的行为策略。与传统的RLHF不同,智能体RL需要处理工具调用、中间状态转换等复杂的多轮决策过程。
本文的核心贡献就是针对多轮智能体强化学习的特殊挑战设计优化框架,理解智能体如何与环境交互是理解AT2PO方法的基础。
ReAct范式
ReAct(Reasoning + Acting)是一种让大语言模型进行推理和行动的范式,模型交替输出思维链和工具调用指令。在每个轮次中,智能体首先进行内部推理,然后决定是调用外部工具还是返回最终答案。工具的返回结果会被加入到上下文中,成为下一轮推理的输入,形成一个结构化的多轮交互过程。
本文采用ReAct范式建模智能体与环境交互,轮次级的策略优化正是基于这种交替推理-行动的自然结构。
重要性采样比率
重要性采样比率$\frac{\pi_\theta(a|s)}{\pi_{\theta_{old}}(a|s)}$用于衡量新策略与旧策略在相同状态下采取相同动作的概率变化程度。在策略优化算法中,这个比率被用于计算策略梯度,通常配合clipping机制防止策略更新过大导致训练不稳定。比率接近1表示策略变化小,偏离1表示策略变化大。
ATPO的核心创新就是将重要性采样从token级和序列级改为轮次级,理解这个概念是理解为什么轮次级优化更稳定的关键。
熵引导树搜索
熵引导树搜索是一种优先扩展高不确定性节点的树搜索策略。策略熵$H_\pi(y|x) = -\sum p(y|x)\log p(y|x)$衡量了模型在给定状态下输出分布的不确定性,熵越高表示模型越不确定,越值得进一步探索。通过选择熵最高的节点进行分支,可以在有限的计算预算内最大化探索效率。
这是AT2PO在rollout阶段的核心技术,解决了传统方法探索多样性不足的问题,是理解AT2PO如何提升轨迹质量的关键。
信用分配问题
信用分配问题是指在长序列决策中,当只有最终奖励信号时,如何将奖励合理地分配给序列中的每个决策步骤。在多轮智能体任务中,奖励只在任务结束时给出(如答案是否正确),但需要评估每个轮次的推理或工具调用对最终结果的贡献程度,这是一个典型的稀疏奖励下的信用分配难题。
这是多轮智能体RL的核心挑战之一,AT2PO的Turn-wise Credit Assignment就是为了解决这个问题而设计。
研究动机
现有智能体强化学习方法面临三大核心挑战。首先是探索多样性有限:在有限的rollout预算下,传统方法依赖随机或启发式的树扩展策略,无法优先探索高不确定性或高潜力的轮次,导致轨迹质量和多样性不足。例如,在HotpotQA等需要多跳推理的任务中,平均需要3.71个轮次,但如果探索不充分,智能体可能错过关键的推理路径。其次是稀疏奖励分配问题:多轮轨迹通常只在结束时给出奖励(如答案是否正确),这导致很难将成功或失败的责任归因到具体某个轮次的推理或工具调用,学习信号稀疏且延迟。最关键的是策略优化错位:现有方法将智能体输出视为平坦的token序列,要么在token级别应用重要性采样(如GRPO),要么在序列级别应用(如GSPO),但智能体任务是轮次结构的——推理、工具调用、结果接收交替进行,这种结构性错位导致梯度不稳定,学习效率低下。
本文的目标是本文的目标是设计一个统一的多轮智能体强化学习框架,能够同时解决探索多样性有限、稀疏奖励分配和策略优化错位这三大挑战。具体来说,要在有限的计算预算内生成更多样化和高质量的轨迹,从稀疏的最终奖励中提取精细的轮次级监督信号,并设计与多轮决策自然粒度对齐的策略优化算法,从而提升智能体在复杂多轮任务上的性能。
与已有工作不同的是,本文的独特切入点在于充分认识到智能体任务的轮次结构本质,将整个强化学习流程(rollout、rewarding、training)都围绕轮次级设计。与Tree-GRPO等现有方法不同,AT2PO不仅采用树结构组织rollout,更重要的是在树扩展时主动选择高熵节点(而不是随机或启发式),在信用分配时利用树拓扑进行轮次级价值传播(而不是只用最终奖励),在策略优化时设计轮次级重要性采样和clipping(而不是token级或序列级)。这种端到端的轮次级设计是本文的核心创新点,也是区别于所有现有方法的本质特征。
核心方法
AT2PO框架包含三个协同工作的组件。在rollout阶段,采用熵引导树扩展策略,通过计算每个节点的策略熵,优先从不确定性最高的轮次进行分支,在有限的计算预算内最大化探索效率。在rewarding阶段,利用树结构进行轮次级信用分配,通过蒙特卡洛bootstrap将叶节点的最终奖励沿树向上传播,计算每个中间节点的价值和优势估计,从而将稀疏的轨迹级奖励转化为精细的轮次级监督信号。在training阶段,设计智能体轮次策略优化(ATPO),将重要性采样和clipping操作从token级或序列级改为轮次级,使优化目标与多轮决策的自然粒度对齐。这三个组件相互配合:树结构提供了多样性轨迹和拓扑信息,信用分配利用拓扑提取精细监督,ATPO确保优化过程充分利用这些监督信号。
AT2PO的核心创新在于端到端的轮次级设计。首先,熵引导树扩展不仅让rollout更有效率,更重要的是树结构为后续的信用分配提供了拓扑基础。其次,轮次级信用分配利用树形拓扑,通过$V_n = r_n$(叶子)或$V_n = \sum_{c \in C(n)} w_c V_c$(中间)的递归公式,将稀疏的最终奖励传播到每个轮次,其中权重$w_c = \frac{H_\pi(c)}{\sum_{c' \in C(n)} H_\pi(c')$考虑了子节点的熵,高熵子节点获得更大权重。最后,ATPO定义轮次级重要性采样比率$\rho^{turn}_{i,t}(\theta) = \frac{\pi_\theta(y^{k(t)_i}|x, y^{<k(t)_i})}{\pi_{\theta_{old}}(y^{k(t)_i}|x, y^{<k(t)_i})}$,其中$y^{k(t)_i}$表示第$k(t)$个轮次的所有token,同一个轮次内的所有token共享这个比率。这种设计在token级细粒度控制和序列级稳定性之间取得了平衡。
方法步骤详情
AT2PO的完整工作流程分为三个阶段。第一阶段是树结构rollout:对于每个训练batch中的提示$x_i$,首先用当前策略$\pi_\theta$生成M条独立的链式轨迹,初始化树$T_i$。然后进行L次迭代:计算每个非叶子节点的策略熵$H_\pi(n)$,选择熵最高的K个节点,从这些节点的状态重新生成轨迹分支作为新分支添加到树中。分支惩罚系数$\alpha$防止某个节点过度扩展。最终每个提示生成$M + LK$个叶子节点。第二阶段是轮次级信用分配:对每个叶子节点分配归一化的最终奖励$\hat{r}_n = \frac{r_n - \text{mean}\{r_n\}}{\text{std}\{r_n\}}$,然后通过递归公式$V_n = \sum_{c \in C(n)} w_c V_c$(叶子节点$V_n = \hat{r}_n$)向上传播计算每个节点的价值,其中权重$w_c$基于子节点熵归一化。优势估计采用$A_n = V_n$。第三阶段是ATPO训练:对每个样本的每个轮次计算轮次级重要性采样比率,应用clipping策略:$J_{ATPO}(\theta) = \mathbb{E}_{x \sim \mathcal{D}, \{y_i\}_{i=1}^G \sim \pi_{\theta_{old}}(\cdot|x)}[\frac{1}{G}\sum_{i=1}^G \frac{1}{|y_i|}\sum_{t=1}^{|y_i|} M_{i,t} \min(\rho^{turn}_{i,t}(\theta)\hat{A}_{i,t}, \text{clip}(\rho^{turn}_{i,t}(\theta), 1-\epsilon_l, 1+\epsilon_r)\hat{A}_{i,t})]$,其中$M_{i,t}$只包含推理引擎生成的token,排除工具响应token。
技术新颖性
AT2PO的技术新颖性体现在三个方面。首先,熵引导树扩展将探索决策从随机/启发式改为数据驱动,主动选择高不确定性轮次进行分支,这比Tree-GRPO的随机扩展更高效。其次,轮次级信用分配首次利用树拓扑进行价值传播,无需额外的监督信号就能从稀疏奖励中提取精细的轮次级学习信号,这比只用最终奖励的方法更精细。最后,ATPO的轮次级优化在token级和序列级之间找到了最佳平衡:token级方法(如GRPO)由于单个token的采样比率方差大,容易导致大量token被clipped,学习效率低;序列级方法(如GSPO)虽然稳定但损失了细粒度控制;ATPO让同一轮次的token共享比率,减少了clipped token数量,同时保留了足够的细粒度。实验显示,AT2PO在三个backbone模型上平均比最强baseline提升最多1.84个百分点,特别是在需要更多轮次的多跳任务上提升更明显。
实验结果
AT2PO在七个基准数据集上的一致性提升验证了其有效性。在Qwen3-4B模型上,AT2PO在多跳数据集(HotpotQA、2WikiMultihopQA、MuSiQue、Bamboogle)上平均达到48.81%,比最强baseline AEPO的46.95%提升了1.86个百分点;在单跳数据集(NQ、TriviaQA、PopQA)上平均达到56.44%,比DAPO的56.33%略有提升。在Qwen3-8B模型上,多跳任务平均50.15%,比AEPO的48.62%提升1.53个百分点;单跳任务平均58.82%,与DAPO持平。在Qwen2.5-7B模型上,多跳任务平均45.83%,比AEPO的44.51%提升1.32个百分点;单跳任务平均56.34%,与DAPO持平。值得注意的是,AT2PO在多跳任务上的提升始终大于单跳任务,这正好验证了轮次级设计的价值:多跳任务平均需要3.71个轮次,而单跳任务43.5%的样本只需要1个轮次,轮次越多的场景越能发挥AT2PO的优势。训练动态分析显示,AT2PO保持了最稳定的熵轨迹,避免了GRPO等token级方法早期熵崩塌的问题,这直接对应了更好的探索能力。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| HotpotQA (Qwen3-4B) | Exact Match | 49.44% | AEPO 46.36% | +3.08pp |
| 2WikiMultihopQA (Qwen3-4B) | Exact Match | 52.99% | AEPO 51.78% | +1.21pp |
| MuSiQue (Qwen3-4B) | Exact Match | 24.80% | AEPO 23.47% | +1.33pp |
| Bamboogle (Qwen3-4B) | Exact Match | 56.80% | DAPO 51.20% | +5.60pp |
| Natural Questions (Qwen3-4B) | Exact Match | 47.90% | DAPO 47.50% | +0.40pp |
| TriviaQA (Qwen3-4B) | Exact Match | 65.32% | DAPO 65.84% | -0.52pp |
| PopQA (Qwen3-4B) | Exact Match | 51.81% | DAPO 51.03% | +0.78pp |
| Multi-Hop Avg (Qwen3-8B) | Exact Match | 50.15% | AEPO 48.62% | +1.53pp |
| Single-Hop Avg (Qwen3-8B) | Exact Match | 58.82% | DAPO 58.53% | +0.29pp |
局限与改进
作者承认的局限性包括:树扩展比线性rollout引入了额外的计算开销,特别是在有充足计算资源时,扩展过程的多次迭代限制了并行化效率;另一个局限是本文只评估了搜索型智能体环境,未来需要在更多样的环境中评估泛化能力。我自己观察到的额外局限是:AT2PO在单跳任务上的提升有限,有时甚至被使用3倍rollout预算的DAPO超过,这表明在轮次较少的场景下轮次级设计的边际收益递减;另外,实验只评估了Exact Match这种二元奖励,没有探索更细粒度的奖励信号(如F1-score)是否会影响不同方法的表现;最后,消融实验显示不同的信用分配公式(如TD误差$L_n = V_n - V_{n_{parent}}$)表现不佳,但作者没有深入分析为什么在智能体RL中直接使用$V_n$作为优势更有效,这是理论理解上的一个空白。
独立分析的弱点
AT2PO在单轮或少轮任务上的收益有限是一个明显的弱点。在NQ、TriviaQA、PopQA等单跳数据集上,43.5%的样本只需要一个工具调用,这种场景下轮次级设计的优势无法充分发挥。改进方向可以是设计自适应机制,根据任务的轮次复杂度动态调整优化粒度——对于简单任务自动降级到序列级优化,对于复杂任务使用轮次级优化。另一个弱点是树扩展的迭代过程限制了并行化,导致计算开销增加。改进方向可以探索更高效的并行树搜索算法,如异步扩展策略,让不同节点的扩展并行进行而不是串行迭代。此外,当前方法只考虑了策略熵作为扩展标准,没有考虑历史探索效果,改进方向可以引入multi-armed bandit等机制,同时平衡exploration和exploitation。
未来方向
作者提出的未来工作方向包括探索更高效的rollout策略以提升并行化效率,以及在更多样的智能体环境中评估AT2PO的泛化能力。基于实验结果,我认为有价值的延伸方向包括:1)将AT2PO扩展到更复杂的多模态智能体场景,如视觉-语言导航任务,其中轮次结构更加复杂;2)研究不同奖励信号(连续奖励、部分奖励、人类反馈)对轮次级信用分配的影响,当前只用二元EM奖励限制了信用的精细度;3)探索ATPO与其他稳定性技术的结合,如自适应clipping、梯度归一化等,进一步提升训练稳定性;4)理论分析为什么在智能体RL中直接使用$V_n$作为优势优于TD误差,这可能揭示智能体RL与传统RLVR的本质区别;5)研究AT2PO在更大规模模型(如70B参数)上的表现,验证轮次级设计的可扩展性。
复现评估
论文的复现性良好。代码已在GitHub开源(https://github.com/zzfoutofspace/ATPO),使用VeRL框架实现,这是一个成熟的RL基础设施。实验使用公开的QA基准(HotpotQA、2WikiMultihopQA、MuSiQue、Bamboogle、NQ、TriviaQA、PopQA)和本地搭建的Wikipedia搜索引擎(2100万条目,e5-base-v2检索器),无需付费API。硬件要求是8张NVIDIA H20 GPU,这在学术环境中可及。超参数设置明确:batch size 64、mini-batch size 8、最大响应长度6192、最大工具调用6次、AT2PO配置$M=10$、$L=2$、$K=6$、clipping阈值$3 \times 10^{-3}$和$4 \times 10^{-3}$。训练240步。论文还详细描述了不同baseline的配置,包括GRPO、DAPO、GSPO、AEPO、Tree-GRPO的超参数。唯一需要注意的是,作者提到Tree-GRPO在Qwen3模型上训练不稳定,可能的原因是retokenization drift,这是复现时需要特别注意的技术细节。总体来说,复现难度中等,主要挑战是硬件资源和环境配置。
论文图表
上图展示了一个完整的搜索智能体rollout序列的token级熵分布,横轴是token位置,纵轴是熵值,不同颜色代表不同轮次。下图是轮次级熵的箱线图分析,展示了不同轮次之间熵分布的显著差异。数据显示不同轮次的token熵值存在实质性差异,这验证了轮次级设计的必要性。
这张图对理解论文至关重要,因为它直观地展示了多轮智能体轨迹中不同轮次的不确定性差异,这是AT2PO设计熵引导树扩展和轮次级优化的实证基础。如果不同轮次的熵分布相同,那么轮次级设计就没有意义,但图显示差异显著,因此证明了问题的存在和解决方向的有效性。