TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning
基本信息
- 标题: TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning
- 第一作者: Heming Zou (Tsinghua University, Tencent LLM Department)
- 研究团队: tsinghua_university, Tencent LLM Department
- 会议/期刊: arXiv:2606.11119v1 [cs.LG]
- 代码: (未在论文中明确提供公开代码链接)
- PDF 文件: [TRACE Paper](file:///C:/Users/admin/.openclaw/workspace/attachment/papers/20260611_trace_unified_rollout_budget_allocation_agentic_rl.pdf)
研究摘要
在大型语言模型(Large Language Models, LLMs)向具备自主决策与长程推理能力的智能体(Agentic Systems)演进的进程中,强化学习(Reinforcement Learning, RL)已成为塑造模型推理行为与工具调用能力的关键技术路径。具体而言,基于可验证奖励的强化学习(Reinforcement Learning with Verifiable Rewards, RLVR)通过为模型生成的推理链(Chain-of-Thought, CoT)与多轮环境交互提供明确的终端反馈,引导模型在广阔的解空间中探索更优策略。然而,这一范式的成功背后隐藏着一个根本性的计算瓶颈:rollout——即让模型完整执行一条从初始提示到最终答案的推理轨迹——在复杂的多轮任务中极为昂贵。当模型需要在数学推理、多跳问答或函数调用等场景中反复与环境交互、调用工具、整合观察结果时,每一次rollout都意味着大量token的生成与推理计算。在有限的计算预算下,如何决定"在哪里投入rollout资源",而非简单地"增加rollout数量",成为了RLVR从理论优雅走向实际高效的核心设计问题。
TRACE这篇论文正是在这一背景下提出了一个深刻而统一的视角。作者们敏锐地指出,当前RLVR的采样低效性并非源于rollout总量不足,而是源于rollout的"分布错配":在提示(prompt)层面,过于简单或过于困难的实例往往产生低方差的同质反馈——所有rollout要么全部成功、要么全部失败,无法为策略优化提供有效的对比信号;在轨迹(trajectory)层面,将单一的终端奖励赋给整条多轮rollout,使得中间各个决策步骤的局部贡献被淹没在稀疏的奖励信号中,信用分配(credit assignment)问题严重恶化。现有工作虽然已通过提示筛选(prompt selection)和根级rollout分配(rollout-count allocation)缓解了提示层面的低效性,但它们始终将rollout视为从根到叶的独立原子轨迹,忽视了多轮交互中内部前缀(prefix)的语义价值——每一个已完成的ReAct式思考-行动-观察回合,实际上都构成了一个可被重新审视、分支探索的决策节点。
TRACE的核心洞见在于,将rollout预算分配重新框架化为**混合奖励对比构建(mixed-reward contrast construction)**问题。作者们将每一次rollout建模为一棵树的生长过程:提示根节点(prompt root)是树的深度零点,而每个已访问的内部前缀则是非根的分支点(branching anchor)。在这个统一的树视角下,提示筛选(零预算跳过)、根级rollout数量分配(正预算激活)与前缀级续展(continuation branching)都被纳入同一个预算分配原则:将有限的rollout资源导向那些后代集合最有可能同时包含成功与失败结果的锚点。这种锚点上的成功与失败共存,恰好构成了策略优化所需的对比信号——共享前缀的兄弟姐妹rollout之间的差异,隐式地定义了关于"从此前缀出发,何种续展更优"的偏好对(preference pair),从而将稀疏的终端奖励转化为密集的内部信用信号。
在技术实现上,TRACE提出了一个预测器引导的、预算约束的树rollout框架。它包含两个互补阶段:全局根分配(global root allocation)在大量候选提示池中决定哪些提示值得投入预算以及投入多少;局部前缀扩展(local prefix expansion)则在已激活提示的已访问前缀上,根据预测的成功率决定哪些前缀应该获得额外的分支续展。两个阶段共享一个可泛化的条件成功概率预测器(conditional success probability predictor),该预测器从前缀历史中学习估计每个锚点的成功概率,并据此指导预算分配。实验结果表明,在数学推理、多跳问答和函数调用三个代表性的多轮智能体任务上,TRACE在相同的rollout预算下,显著优于均匀采样基线(GRPO)、提示难度筛选基线(PCL)以及随机树分配基线(TreePO),不仅提升了最终策略的准确性,更大幅提高了训练过程中有效对比样本的比例。
这一工作的理论意义在于,它首次将提示层面的预算分配与回合层面的分支探索统一在了一个基于树结构的数学框架下,并通过严格的命题证明(Prefix Information Proposition、Prefix Uncertainty Proposition、Activation Allocation Proposition)确立了前缀信息在降低群体难度预测误差、量化剩余对比潜力以及增强策略梯度信号方面的根本价值。对于实践者而言,TRACE提供了一个即插即用的训练框架,只需在现有RLVR流水线中替换采样策略,即可在不增加计算预算的前提下获得更优的策略收敛速度与最终性能。
理论框架
TRACE的理论根基建立在三个相互关联的洞察之上,它们共同构成了从"独立轨迹采样"到"树形对比构建"的范式转换。
智能体RLVR的前缀历史形式化构成了理论大厦的第一块基石。在ReAct交互框架中,智能体在每个回合
在此形式化之上,作者定义了条件成功概率(conditional success probability)
命题1(Prefix Information Improves Group Difficulty Prediction) 建立了前缀信息在预测上的根本优势。该命题考虑一个预测器
命题2(Prefix Uncertainty as Remaining Contrast Potential) 将条件成功概率与剩余对比潜力联系起来。作者引入鞅(martingale)视角,将条件成功预报的二次变差定义为
这个等式的深刻之处在于,它将前缀处的"剩余不确定性"量化为伯努利方差的形式。
命题3(Activation Allocation Dominates Uniform) 则将对比构建与策略梯度信号直接联系起来。作者定义了聚合局部策略梯度贡献
且两者的组合也满足:
这一命题的核心洞见是:在二元奖励设置下,局部梯度贡献仅在锚点的后代同时观察到成功与失败时才非零。因此,期望平方梯度范数恰好等于"激活概率"乘以"条件梯度规模"。TRACE通过将预算导向对比丰富的锚点,提高了激活概率,从而直接增强了策略更新的信号强度。这完成了从"预算分配"到"梯度信号"的理论闭环:TRACE不是在rollout之后被动地利用对比,而是在rollout之前主动地选择那些最可能产生对比的锚点,从而在采样阶段就为优化器准备了更丰富的信息。
技术架构
TRACE的技术实现是一个从全局到局部、从预测到优化的完整系统,其核心架构可以概括为"预测器引导的树形预算分配 + 树感知策略优化"的两阶段流水线。
全局根分配(Global Root Allocation) 是系统的第一阶段。在每个训练步骤
其中
这一优化问题通过动态规划高效求解,其计算开销相对于rollout生成而言微不足道。在实现上,作者特意排除了
局部前缀扩展(Local Prefix Expansion) 是第二阶段,也是TRACE区别于现有工作的核心创新。在提示
其中
这个公式的直观含义是:如果前缀
值得注意的是,TRACE的前缀扩展是提示级局部化的:每个提示的扩展预算
共享前缀价值预测器(Shared Prefix Value Predictor) 是贯穿两个阶段的信息中枢。预测器
预测器的训练目标是通过递归树支撑值(recursive tree-backed targets)实现的。对于rollout树中的任意节点
其中监督集合
树感知策略优化(Tree-Aware Policy Optimization) 接收完成的rollout树,并利用任何支持前缀级信用分配的树感知优化器进行策略更新。TRACE本身不绑定特定的优化器,而是将树结构作为通用接口提供给下游。在实验中,作者使用了基于组相对策略优化(Group Relative Policy Optimization, GRPO)的变体,但将对比样本从提示级扩展到前缀级——同一前缀下的不同续展被视作一个对比组,其终端奖励的差异隐式定义了关于续展选择的偏好对。这种设计将树rollout的对比潜力转化为可直接用于策略梯度计算的偏好信号,实现了从采样到优化的端到端信息流动。
实验评估
TRACE的实验设计覆盖了三个代表性的多轮智能体任务,构成了对方法泛化能力的严格检验。这三个任务分别对应不同的认知需求:数学推理要求形式化推演与符号计算,多跳问答要求信息检索与证据整合,函数调用要求结构化输出与工具编排。这种任务多样性确保了实验结论的广度,而非局限于单一领域的过拟合。
在数学推理任务中,TRACE在DeepScaler语料库上进行训练,这是一个大规模的数学问题集合,涵盖从高中竞赛到奥林匹克级别的难度。评估在多个in-distribution和out-of-distribution基准上进行:AIME24(美国数学邀请赛)、AMC23(美国数学竞赛)、MATH500(标准数学推理测试)、MinervaMath(科学数学问题)、OlympiadBench(奥林匹克级别双语问题),以及三个分布外测试:MMLU-Pro(多学科大学水平理解)、ARC-c(AI2推理挑战)和GPQA-diamond(研究生级别科学问答)。这种评估矩阵的设计意图是检验TRACE训练的策略是否仅记住了训练分布的解题模式,还是真正提升了模型的推理能力本身。
多跳问答任务在HotpotQA训练集上训练,评估集包括HotpotQA验证集、2WikiMultiHopQA(维基百科多跳问答)、MuSiQue(多跳问题合成)和Bamboogle。系统配备了一个基于E5嵌入模型的本地检索服务器,构建在维基百科dump之上,为智能体提供实时信息检索能力。这一设置更接近真实世界中的知识工作场景:智能体需要自主决定何时搜索、如何整合多源信息、以及何时给出最终答案。
函数调用任务使用BFCL v4(Berkeley Function Calling Leaderboard)的多轮分割,训练集占80%,测试集为剩余的20%。测试进一步细分为base(基础函数调用)、long-context(长上下文)、missing-function(缺失函数)和missing-parameter(缺失参数)四个子集,检验智能体在不同函数规范完整性条件下的鲁棒性。
主要实验使用Qwen3-8B和Qwen3-14B作为策略骨干模型,并在附录中补充了Llama-3.2-3B-Instruct的结果。基线选择经过精心考虑:ReAct评估了基础模型的零样本多轮能力;GRPO代表了当前标准的提示级均匀采样RLVR;PCL(Prompt Curriculum Learning)是提示难度筛选的先进方法,主动过滤简单或困难提示;TreePO在GRPO基础上增加了树形rollout和树感知更新,但前缀分配是随机的。这一基线矩阵的设计使得TRACE的增益可以被精确归因:与GRPO对比衡量了树形采样的价值,与PCL对比衡量了前缀级分配超越提示级筛选的价值,与TreePO对比衡量了智能分配超越随机分配的价值。
实验结果呈现出一致且令人信服的图景。在相同rollout预算下,TRACE在所有三个任务和两个模型规模上均优于所有基线。以数学推理(DeepScaler)为例,TRACE将Qwen3-8B的平均in-distribution准确率从GRPO的70.0提升至71.1,Qwen3-14B从73.5提升至74.9。虽然这些数字的绝对提升看似 modest,但它们在高度竞争的标准化基准上具有统计意义,且 TRACE 的增益在所有任务中一致出现,这远胜于单一任务的偶然波动。
更 revealing 的指标是有效比率(effective ratio)——训练批次中包含成功与失败终端叶子的提示比例。这一指标直接衡量了rollout预算被转化为对比信号的效率。在数学推理任务上,TRACE将Qwen3-8B的有效比率从GRPO的26.8%提升至60.6%,Qwen3-14B从34.7%提升至59.7%。这意味着,在相同的计算开销下,TRACE产生的策略更新信息密度是传统方法的两倍以上。这一提升的源头正是TRACE的对比驱动分配:通过将预算从已饱和的锚点(成功率接近0或1)重新导向不确定区(成功率接近0.5),TRACE确保了每一单位计算资源都投入在最可能产生学习信号的位置。
消融实验进一步验证了设计的必要性。表1展示了在Qwen3-8B HotpotQA上,同时激活根分配和前缀分配(Active-Active)取得了50.6%的准确率和52.3%的有效比率,优于仅激活根分配(49.8%)或仅激活前缀分配(50.0%)的配置,表明两个阶段具有互补性:根分配筛选出适合rollout的提示,而前缀分配在已激活提示内部挖掘更细粒度的对比。预算兼容性实验(表2)显示,TRACE在不同预算形状(M=512/1024根预算,N=2/6扩展因子)下均优于随机TreePO,且预算形状的影响显著——在总预算2048的条件下,更宽的根覆盖(1024根×2扩展)优于更深的根探索(512根×6扩展),表明锚点覆盖的广度比单个锚点的深度更重要,因为找到正确的对比锚点本身就是价值所在。
| M | N | 方法 | 平均准确率 | 平均有效比率 |
|---|---|---|---|---|
| 512 | 2 | TreePO | 48.8 | 32.2 |
| 512 | 2 | TRACE | 49.7 | 42.4 |
| 512 | 6 | TreePO | 49.4 | 37.7 |
| 512 | 6 | TRACE | 50.3 | 47.8 |
| 1024 | 2 | TreePO | 49.5 | 42.8 |
| 1024 | 2 | TRACE | 50.6 | 52.3 |
预测器诊断实验(图6和图7)展示了轻量级预测器在提示和前缀两级都学到了可用的难度排名。Spearman相关系数表明,即使训练主要由提示级监督主导,预测器仍能可靠地将信号迁移到内部前缀——这验证了命题1的理论预期:前缀信息确实改善了群体难度预测,且这种改善是可学习的。
案例研究
TRACE的框架在数学推理任务中展现出最为直观的运作逻辑。考虑一个典型的AIME级别问题:智能体需要调用Python解释器进行符号计算或数值验证。在ReAct框架下,智能体首先生成一个思考过程,决定使用Python工具,然后编写代码,接收执行结果,再基于结果进行下一步推理。在传统的GRPO训练中,这一过程被压缩为一条从问题到最终答案的线性轨迹,终端奖励仅告知"正确"或"错误",但不解释"在哪里出错"。
TRACE的树形视角彻底改变了这一图景。假设在第一条裸rollout中,智能体在第二个回合决定使用Python计算一个中间表达式,但代码中存在边界错误。在GRPO中,这条失败的轨迹只会贡献一个稀疏的负信号,而成功轨迹的正信号无法直接告诉策略"第二个回合的代码应该修正"。但在TRACE中,这个失败的第二个回合前缀被记录为一个已访问的节点。预测器评估该前缀的成功率——如果它处于中间区域(比如0.4),TRACE会为此前缀分配额外的续展分支。这些续展中,有些智能体可能修正了代码逻辑并成功,有些则可能犯了不同错误。共享前缀的对比(相同的前两个回合,不同的第三个回合及之后)隐式地告诉策略:"在前缀
在多跳问答任务中,TRACE的价值体现在信息检索的决策点上。HotpotQA问题通常需要跨多个文档的推理,智能体必须决定何时搜索、使用什么关键词、以及何时停止搜索并综合答案。在裸rollout阶段,如果智能体在第三个回合使用了一个过于宽泛的搜索查询,导致返回了大量无关文档,后续回合可能陷入信息过载。在GRPO中,这条轨迹的失败仅贡献一个终端负信号。但在TRACE中,"第三个回合的宽泛查询"这一前缀被识别为一个高价值锚点——预测器可能判断该前缀仍有成功潜力(因为正确的信息可能存在于搜索结果中),从而分配续展分支。在这些续展中,有些智能体可能选择更精确的筛选策略,有些可能重新搜索。前缀级对比再次提供了局部学习信号:"在已经执行了宽泛搜索的条件下,什么样的后续策略更可能成功?"
函数调用任务中的案例则揭示了TRACE在结构化输出上的应用。BFCL v4要求智能体根据用户请求和可用工具描述生成精确的函数调用序列。在多轮场景中,智能体可能需要先调用一个函数获取中间结果,再基于该结果调用第二个函数。如果第一个函数调用返回了错误类型的数据(例如字符串而非数字),第二个函数调用可能因参数类型不匹配而失败。在GRPO中,这一失败仅贡献一个终端负信号。TRACE则能将第一个函数调用后、第二个调用前的状态识别为一个关键前缀,并分配续展探索不同的参数解析或转换策略。这种能力对于构建能够在复杂API生态中自主导航的智能体至关重要。
这些案例共同揭示了一个深层原则:TRACE通过将rollout从线性轨迹转化为分支树,将稀疏的终端奖励转化为密集的内部对比。每一个成功与失败的兄弟姐妹对,都是关于"在此前缀条件下,何为优、何为劣"的一条隐式偏好数据。这些偏好数据的积累,使得策略模型能够在不依赖人工设计的过程奖励或庞大的rollout预算的情况下,逐步学会在多轮决策的每个关键节点做出更优选择。
综合价值与局限
TRACE在理论层面提供了一个统一而优雅的框架,将此前分散的提示筛选、rollout数量分配和过程监督尝试整合在一个基于树形结构的数学视角下。其最核心的理论贡献是确立了"混合奖励对比"作为RLVR采样效率的根本原则:预算的价值不在于rollout的数量,而在于rollout的对比丰富度。这一原则不仅解释了为什么均匀采样是低效的,也为未来更复杂的采样策略提供了设计准则。三个命题的严格证明为这一直觉提供了坚实的数学基础,特别是命题2将条件成功概率的伯努利方差与鞅的二次变差联系起来,展现了作者深厚的概率论功底。
在实践层面,TRACE的最大优势是其即插即用性。现有的RLVR流水线通常已经具备了提示生成、rollout执行、奖励计算和策略更新的基础设施。TRACE的引入仅需替换采样策略模块,无需改变优化器或奖励机制。全局根分配和局部前缀扩展的计算开销相对于LLM推理本身而言微不足道,因此不会引入显著的系统延迟。此外,TRACE的提示级局部化设计(前缀扩展在单个提示内完成,无需跨提示同步)使其天然兼容现有的分布式训练框架,这是从研究原型走向生产系统的关键工程考量。
然而,TRACE也存在若干值得正视的局限。首先,框架目前主要针对二元终端奖励(成功/失败)进行设计,其理论分析和价值函数都基于伯努利变量的假设。虽然这一假设在数学推理、问答和函数调用等具有明确正确答案的任务中成立,但在更开放、奖励更细粒度的任务中(如创意写作、开放式对话评估),
其次,预测器的性能直接决定了TRACE的分配质量。尽管实验表明轻量级的Qwen3-0.6B critic能够学到可用的难度排名,但预测器的准确率存在上限。在理论上,如果预测器是完美的,TRACE的分配也是最优的;但在实践中,预测误差可能导致预算被导向次优锚点。当前实现采用简单的在线MSE更新,作者指出未来可以通过持续学习技术和参数高效微调(如LoRA)来缓解预测器面临的稳定性-可塑性困境,同时降低更新开销。
第三,实验范围虽然覆盖了三个代表性任务,但均属于相对结构化的智能体场景——数学问题有明确答案、问答有标准知识库、函数调用有规范API。在更复杂、非平稳的开放环境中(如实时网页浏览、多用户协作、物理世界交互),前缀的语义稳定性可能下降,预测器的泛化能力可能面临更大挑战。这些更复杂的场景有待未来探索。
从更宏观的视角来看,TRACE代表了一个重要的方法论转向:从"如何生成更多rollout"到"如何在rollout树中做出更优的分支决策"。这与计算机科学中从暴力搜索到启发式搜索的演进一脉相承。在AlphaGo中,Monte Carlo Tree Search(MCTS)通过选择性地扩展最有价值的节点,在巨大的博弈树中实现了超越人类顶尖棋手的能力。TRACE将类似的树搜索思想引入了LLM策略学习,但其独特之处在于:TRACE不是在一个固定的游戏规则树上搜索,而是在一个由语言模型自身生成的、动态变化的语义树上进行预算分配。这种"自生成树"的特性使TRACE的理论分析比传统MCTS更为复杂,但也使其适用范围远超固定规则的游戏场景。
延伸阅读与思考
TRACE的理论根基可以追溯到多个相互交织的研究传统。在强化学习领域,树搜索方法如MCTS和AlphaZero系列(Feng et al., 2023; Koh et al., 2024)已经在离散决策空间中展示了选择性分支的价值。在LLM推理领域,Chain-of-Thought prompting(Wei et al., 2022)和ReAct框架(Yao et al., 2022)为多轮交互提供了基础架构。在RLVR领域,GRPO(Shao et al., 2024)建立了组相对策略优化的标准范式,而PCL(Gao et al., 2025)和Bots(Shen et al., 2025)等提示筛选方法则从另一角度探索了采样效率。TRACE的独特贡献在于将这些线索整合在一个统一的数学框架中,并通过严格的理论分析确立了前缀级信息的价值。
与TRACE最直接相关的竞争方法是TreePO(Ji et al., 2025)和TreeRL(Hou et al., 2025),它们同样探索了树形rollout在LLM RL中的应用。然而,这些工作主要关注树结构的构建和树感知更新的设计,而未将预算分配本身作为优化问题来求解。TRACE填补了这一空白:它证明即使给定相同的树形基础设施,智能的分配策略(TRACE)仍能显著优于随机分配(TreePO),这意味着树结构的潜力尚未被充分挖掘。
从未来研究的角度看,TRACE开辟了多个富有前景的方向。首先是预测器的持续进化。当前预测器采用固定架构的在线更新,但正如作者所指出,前缀历史分布随着策略改进而不断漂移,这本质上是一个非平稳的连续学习问题。将弹性权重巩固(EWC)、无遗忘学习(LwF)等持续学习技术引入预测器训练,或采用动态架构调整(如神经架构搜索),可能显著提升预测器的长期稳定性与适应性。此外,将预测器从标量成功概率扩展为更丰富的价值分布(如预测成功概率的后验分布),可以支持更风险敏感的预算分配策略。
第二个方向是与测试时计算扩展的协同。Snell et al.(2024)等人的研究表明,在推理阶段增加测试时计算(如通过多次采样和多数投票)可以显著提升LLM性能。TRACE的训练阶段树形rollout实际上可以被视为一种"训练时计算扩展",它与测试时计算扩展之间存在潜在的协同效应:在训练阶段学会了高效分支的策略,在测试阶段可能更擅长利用额外的推理时间预算。探索这种训练-测试计算的联合优化是一个激动人心的方向。
第三个方向是跨任务迁移与元学习。TRACE的预测器在不同任务上表现出可迁移的预测能力(从提示级到前缀级的信号迁移)。如果这种迁移能力可以进一步扩展为跨任务迁移——即在一个任务上训练的前缀预测器,能够快速适应新任务的前缀分布——那么TRACE可能成为一种元学习框架,显著降低新任务上的训练成本。
TRACE还引发了一个更深层的方法论反思:在LLM时代,我们是否应该重新思考"计算"与"数据"的边界?传统上,计算是指模型的前向/反向传播,数据是指静态的预训练语料或人工标注。但TRACE展示了一种新的可能性:计算本身可以生成数据——通过策略性的rollout分支,系统生成了高质量的隐式偏好对,这些偏好对构成了策略学习的燃料。在这种视角下,计算不再是数据的消费者,而是数据的创造者。TRACE的预算分配问题,本质上是在问:如何以最高效的计算投入,创造最有价值的训练数据?这一问题可能在未来的AI系统中变得越来越核心,因为当模型足够强大时,人工标注数据的边际价值递减,而模型自身生成的结构化对比数据可能成为持续提升的关键。
最后,一个令人深思的观察是:TRACE的对比驱动原则与人类学习中的"最近发展区"(Zone of Proximal Development, Vygotsky)概念存在深刻的平行。教育心理学研究表明,人类在挑战既不过于简单(已掌握)也不过于困难(超出能力)的任务时学习最有效。TRACE的
笔记创建时间: 2026-06-11
阅读方式: L2 深度阅读
Topics:
- "reinforce_learning"
- "reasoning"
- "agent_architecture"
- "test_time_scaling"
- "llm"
References: - "tsinghua_university"
- "hotpotqa"
- "bfcl_v3"
- "monte_carlo_tree_search"