TRACE团队 投稿
量子位 | 公众号 QbitAI
大模型做强化学习后训练,最容易被低估的成本是什么?
在普通数学题里,成本可能主要来自一遍遍生成长链推理答案。
但到了Agentic场景,事情会更麻烦。
模型不只是写答案,而是要一步步思考、调用工具、接收环境反馈,再继续决定下一步做什么。每一次完整交互轨迹,都是一次rollout。
如果任务很长,工具调用很多,环境反馈又复杂,那么rollout预算很快就会变成训练里的硬瓶颈。
更麻烦的是:不是所有rollout都同样有用。
在RLVR(Reinforcement Learning with Verifiable Rewards,可验证奖励强化学习)里,很多训练流程只看最终答案是否正确。成功记为1,失败记为0。这种outcome-only reward很方便,也更容易规模化。
但它有一个老问题:如果一组rollout全对,或者全错,模型其实很难从这组样本里学到“该偏好哪条路径”。
到了多轮Agent任务里,这个问题又被放大了一层。
因为一次rollout里有很多中间步骤:搜索、推理、调用API、读取观察结果。最终奖励只落在整条轨迹末端,前面的每一步都被同一个终局奖励覆盖。这样一来,训练信号就很稀疏,信用分配也很粗。
针对这个问题,来自清华大学和腾讯混元团队的研究者提出了TRACE(Tree Rollout Allocation for Contrastive Exploration),一个面向Agentic RLVR的统一rollout预算分配框架。
一句话概括:TRACE不再把rollout预算平均撒到每个prompt上,也不只在prompt层面挑题,而是把Agent轨迹看成一棵树,主动把预算分配给那些最可能产生“成功/失败对比”的prompt根节点和中间前缀节点。
实验显示,在数学推理、多跳问答、函数调用三类多轮Agentic任务上,TRACE在相同rollout预算下稳定超过GRPO、PCL和随机树分支TreePO。比如在Qwen3-14B的Multi-Hop QA上,TRACE将平均准确率从GRPO的51.2提升到54.0,提升2.8个点;在Qwen3-8B的DeepScaler数学任务上,TRACE将有效样本比例从GRPO的26.8%提升到60.6%。

RLVR的问题,不只是“样本少”,而是“对比少”
先从RLVR里最常见的一类训练信号说起。
一个prompt给模型,模型生成多条回答。答案能自动验证,就给每条回答一个最终奖励。
如果有的回答对、有的回答错,那么这组样本就很有价值。因为模型可以比较:在同一个问题下面,哪些轨迹更值得强化,哪些轨迹应该被压低。
但如果一组回答全部正确,或者全部错误,事情就不一样了。
从奖励分布看,这组样本没有方差。对GRPO这类组内相对优化方法来说,它很难提供有效的偏好信号。
这也是很多RLVR论文反复遇到的问题:太简单的题,模型全答对;太难的题,模型全答错。两种情况都会消耗rollout,但训练信号并不强。
过去很多工作主要在prompt层面处理这个问题。
比如筛掉太简单或太难的题,只保留成功率处在中间区域的prompt;或者给不同prompt分配不同数量的rollout。
这个方向当然有用,但在Agentic任务里还不够。
因为Agent不是一次性吐出完整答案。它会经历多轮thought-action-observation:先想一步,做一个动作,拿到环境反馈,然后再想下一步。
论文把第t轮交互写成:

其中τt是思考,at是动作,ot是环境观察。
到第t步为止,模型看到的历史是:

这里的x就是原始prompt。
也就是说,prompt只是树的根节点;每一个中间历史Ht,都可以看成一个可以继续分叉的前缀节点。
这就引出了TRACE的核心问题:
既然中间前缀也会决定后续成败,那么rollout预算为什么只能在prompt层面分配?

△TRACE的核心直觉:在固定rollout预算下,平均采样会产生较弱的outcome contrast;把预算转向更可能出现成功/失败分歧的根节点和前缀节点,可以形成更强的训练信号
TRACE想做的事:把outcome reward变成更多局部对比
在outcome-only RLVR里,最终奖励只有一个。
但如果两条轨迹共享同一个前缀,只是在某一步之后走向不同分支,最后一个成功、一个失败,那么它们就天然构成了一个局部对比。
这比单独看一条完整轨迹有用得多。
因为它告诉模型:在同一个历史状态下,后续某些决策更可能导向成功,另一些决策更可能导向失败。
TRACE的想法就是围绕这个“对比”展开。
它会估计一个历史前缀Ht后续成功的概率:

直观理解,这就是:给定当前走到的状态,从这里继续往下走,最终成功的概率是多少。
如果
接近0,说明从这里往下几乎都会失败。如果
接近1,说明从这里往下几乎都会成功。
这两种情况都不太能制造对比。
真正值得继续探索的,是那些成功概率处在中间区域的节点。因为它们更可能在后续分支中同时出现成功和失败。
论文把这种思路概括为mixed-reward contrast allocation(混合奖励对比分配):rollout预算应该优先给那些后代轨迹更可能出现混合终局奖励的锚点。
这里的锚点有两类。
一类是prompt根节点。
另一类是rollout过程中已经访问过的中间前缀。

△论文的对比分配诊断:很多根节点和前缀节点的经验成功率接近0或1,真正有对比潜力的节点只占一部分;按对比价值排序后,少量高价值锚点能覆盖更多有效pair contrast
Stage 1:先在prompt根节点上分配预算
TRACE的第一步,是在一个候选prompt batch里决定:哪些prompt跳过,哪些prompt多采样。
设第i个prompt的预测成功概率为vi,如果给它分配m条根rollout,那么这m条rollout同时包含成功和失败的概率可以写成:

这个式子很好理解。
表示m条rollout全部成功;
表示m条rollout全部失败。
两者都不是我们最想要的情况。
所以用
,就能表示“这组rollout里既有成功也有失败”的概率。
在TRACE里,每个prompt的分配数量属于:

这里有个细节:TRACE不给单个prompt分配1条rollout。
原因也很直接。只有一条轨迹时,本身没法形成组内对比。要么跳过这个prompt,要么至少给它1条rollout,让它有机会出现成功/失败差异。
所以Stage 1的作用,可以理解为:在prompt层面先把预算推向更可能产生有效对比的问题。
Stage 2:再到中间前缀上继续分叉
只在prompt层面做预算分配,还是没有完全用上Agent轨迹里的结构。
比如一个Agent已经搜索到了某些证据,或者已经调用过一个函数,接下来可能有几种不同走法。
这时候,继续从原始prompt重新采样不一定最划算。更直接的做法,是从当前前缀继续分叉,看看不同continuation会不会走向不同结果。
这就是TRACE的第二步:local prefix expansion(局部前缀扩展)。
假设某条原始rollout的最终奖励是
,某个中间前缀是
。如果从这里额外采k条continuation,TRACE使用的前缀价值是:

这个公式看着长,但意思不复杂。
如果原来的事实分支是成功的,TRACE希望额外分支里至少出现一次失败。
如果原来的事实分支是失败的,TRACE希望额外分支里至少出现一次成功。
也就是说,Stage 2不是为了随便多生成几条continuation,而是为了在同一个前缀下面制造“反事实对比”。
这样,最终奖励虽然仍然只有0/1,但树结构会让奖励变得更像局部偏好信号。

△TRACE框架示意图:共享的prefix value predictor先给prompt根节点和中间前缀打分,再通过全局根分配和局部前缀扩展构造rollouttree,最后用树结构数据更新预测器和策略模型
关键不是换优化器,而是在优化器前面换一种采样方式
TRACE容易被误解成一种新的RL优化器。
但论文里强调得很清楚:TRACE更像是一个上游的rolloutacquisition layer。
也就是说,它决定的是:在训练前,应该去哪里采样、从哪里分叉、哪些节点值得继续花预算。
至于采完以后怎么把这棵树变成梯度,可以接不同的tree-aware policy optimizer。
这点很重要。
GRPO是平铺式采样。PCL会在prompt层面做预测式筛选。TreePO引入树结构,但分支位置还是随机的。
TRACE的差异在于,它把“树上的哪个位置值得分叉”也变成了一个可学习、可分配的问题。
论文里对几类方法做了一个很直观的对比:

换句话说,TRACE不是简单地“多采样”,而是在固定预算下,把采样变成了一个结构化分配问题。
实验设置:数学推理、多跳问答、函数调用
论文把TRACE放在三类多轮任务上测试。
第一类是Mathematical Reasoning(数学推理)。
训练数据来自DeepScaler corpus。评测包括AIME24、AMC23、MATH500、MinervaMath、OlympiadBench,以及分布外的MMLU-Pro、ARC-c、GPQA-diamond。
第二类是Multi-Hop QA(多跳问答)。
训练使用HotpotQA,评测覆盖HotpotQA、2WikiMultiHopQA、MuSiQue和Bamboogle。这个设置里还接入了本地E5检索服务器和Wikipedia dump,因此更接近典型的检索型Agent任务。
第三类是Function Calling(函数调用)。
训练使用BFCL v4 multi-turn split的80%,评测使用剩余20%,并覆盖base、long-context、missing-function、missing-parameter等子集。
主实验使用Qwen3-8B和Qwen3-14B作为策略模型。附录中还额外测试了Llama-3.2-3B-Instruct。
对比方法主要包括:
ReAct:不做RL微调,只在相同Agent工具框架下评测基础模型; GRPO:prompt层面随机采样的基础RLVR方法; PCL:基于prompt难度预测的主动筛选方法; TreePO:使用树结构rollout和tree-aware update,但分支位置随机; TRACE:同时学习根节点分配和前缀分支分配。
论文强调,所有方法使用相同rollout预算口径,所以性能差异主要来自预算分配,而不是额外多拿了样本。
主结果:同样预算下,TRACE的曲线整体更高
先看训练过程中的准确率曲线。

△Qwen3-8B与Qwen3-14B在DeepScaler、HotpotQA、BFCL三类任务上的训练曲线。紫色曲线为TRACE,在相同rollout预算下整体高于GRPO、PCL和TreePO
从曲线看,TRACE在三类任务和两种模型规模上都更稳。
数学推理上,Qwen3-8B的分布内平均分从GRPO的70.0提升到TRACE的71.1;Qwen3-14B从73.5提升到74.9。
多跳问答上,Qwen3-8B的平均分从GRPO的48.5提升到TRACE的50.6;Qwen3-14B从51.2提升到54.0。
函数调用上,Qwen3-8B的平均成功率从GRPO的43.5提升到TRACE的46.2;Qwen3-14B从46.1提升到48.0。
更值得注意的是TreePO对比。
TreePO和TRACE都使用树结构rollout,也都能做tree-aware update。区别在于,TreePO的树分支是随机的,TRACE的树分支是预测式分配的。
结果上,TRACE仍然超过TreePO。这说明增益不只是来自“有一棵树”,更来自“知道树上哪里值得继续分叉”。
详细结果:三个任务上都不是单点提升
论文在附录里给出了完整benchmark表格。
数学推理部分,TRACE在分布内和分布外平均分上都优于其他方法。

在Multi-Hop QA和Function Calling上,TRACE也保持了平均分优势。

这些结果里有一个细节:TRACE并不是每个单项都第一。
比如Qwen3-8B数学任务中,AIME24和AMC23的最高值并不都来自TRACE。
但TRACE在平均分上更稳,说明它提升的不是某个benchmark的偶然波动,而是整体训练效率和采样质量。
为什么能提升:有效样本比例明显更高
只看准确率,还不能完全说明TRACE在做什么。
更关键的是effective ratio,也就是训练中真正形成成功/失败对比的样本比例。

△有效样本比例曲线。TRACE在三类任务和两种模型规模下都更容易构造出非退化奖励组,因此每单位rollout更容易变成有效更新信号
DeepScaler数学任务上,TRACE的提升非常明显。
在Qwen3-8B上,GRPO的平均有效样本比例是26.8%,TRACE提升到60.6%。
在Qwen3-14B上,GRPO是34.7%,TRACE提升到59.7%。
这说明TRACE的核心收益确实来自“对比构造”。
同样数量的rollout,GRPO更可能浪费在全对或全错的组里;TRACE则更倾向于把预算推到尚未饱和、仍有分歧空间的位置。
这也是它在Agentic RLVR里有现实意义的地方。
随着Agent任务越来越长,单纯多采样会越来越贵。比起生成更多轨迹,更重要的可能是:让已有预算更频繁地打到有效信号上。
消融实验:根节点分配和前缀分配都不能少
为了验证TRACE的两个阶段是不是都有用,论文做了消融。
实验对象是Qwen3-8B的Multi-Hop QA(HotpotQA)。研究者把Stage 1或Stage 2分别换成uniform allocation,看准确率和有效样本比例的变化。

结果很清楚。
只做根节点分配,能把有效样本比例从42.8拉到49.1。
只做前缀分配,也能把准确率拉到50.0,有效样本比例到47.3。
但两者叠加之后,准确率和有效样本比例都最高。
这说明Agentic RLVR的预算分配不是一个单层问题。
prompt重要,轨迹中间的前缀也重要。
预算形状也重要:不是rollout总数越多越好
论文还做了一个很有意思的预算敏感性实验。
在TRACE的设置里,根预算是M,扩展因子是N,期望rollout预算可写成:
M(1+N/2)
也就是说,(M,N)=(512,6)和(1024,2)的总预算都是2048,但一个更偏向深挖少量根节点,另一个更偏向覆盖更多根节点。
实验结果如下:

TRACE在每个预算设置下都优于随机TreePO。
但更有意思的是,两个同样总预算为2048的设置里,(1024,2)比(512,6)更强。
这说明rollout效率不只取决于总数,还取决于预算形状。
对于HotpotQA这样的多跳问答任务,更广的根节点覆盖,可能比在少量prompt上更深地继续分叉更有价值。
预测器能不能真的学到“哪里有对比”?
TRACE依赖一个shared predictor来估计根节点和前缀节点的成功概率。
如果这个预测器只是随机猜,那整个预算分配就站不住。
论文因此额外报告了prompt-level和prefix-level的Spearman相关性。

△Prompt-level预测质量:预测的prompt难度与经验成功率之间存在可用的排序相关性

△Prefix-level预测质量:同一个预测器也能迁移到中间前缀,说明历史前缀里确实包含可学习的局部不确定性信号
这部分结果的意义在于:TRACE并不是只学会了“这道题难不难”。
它还在一定程度上学会了“走到这个中间状态之后,后面还有没有可能分出成败”。
对于Agentic场景,这点很关键。
因为真正的决策往往不发生在prompt输入的那一刻,而发生在中间的工具调用、检索查询、部分推理状态里。
如果模型能识别这些中间状态的分歧潜力,就能把outcome-only reward变成更密集的局部训练信号。
换到Llama-3.2-3B,TRACE仍然有效
除了Qwen3-8B和Qwen3-14B,论文还在Llama-3.2-3B-Instruct上做了额外测试。
任务还是Multi-Hop QA。


△Llama-3.2-3B在Multi-Hop QA上的平均准确率曲线,TRACE继续保持领先
这里的提升幅度更明显。
TRACE的平均分达到32.3,相比GRPO的28.5、TreePO的29.2,都有比较明显的优势。
这说明TRACE不完全依赖某个特定Qwen backbone。
至少在论文的额外实验里,它可以迁移到另一个较小规模的开源模型上。
成本问题:预测器更新很小,但打分路径还可以优化
既然TRACE多了一个预测器和分配器,它会不会把节省下来的rollout成本又花回计算开销里?
论文在HotpotQA上报告了训练时间拆解。

△HotpotQA训练时间拆解。当前实现中,预测器参数更新本身占比不大;较大的额外开销主要来自尚未充分优化的predictor scoring路径
从图里可以看到,predictor update本身不是主要成本。
真正值得继续优化的是predictor scoring,也就是给大量根节点和前缀节点打分的路径。
论文也没有把TRACE包装成一个完全没有代价的方法。
它更像是在说:当rollout本身非常贵时,花一部分计算去判断哪里更值得rollout,可能是划算的。
尤其在长程Agent任务里,这个判断会越来越重要。
这篇工作的意义:Agentic RLVR需要从“抽多少条”走向“在哪里分叉”
这篇论文最值得关注的地方,不是某一个公式,而是它把Agentic RLVR的采样问题换了一个角度。
过去我们经常问:一个prompt要采几条rollout?
TRACE进一步追问:一条轨迹已经走到中间状态了,接下来应该从哪里继续分叉,才最可能产生有用的训练信号?
在普通RLVR里,prompt是天然的采样单位。
但在多轮Agent任务里,prompt只是起点。真正的信息,往往藏在模型与环境交互之后形成的前缀状态里。
TRACE把这些前缀也纳入预算分配,并用共享预测器估计它们的后续成功概率。
如果一个前缀几乎必然成功,继续分叉价值不大。
如果一个前缀几乎必然失败,继续分叉也很可能浪费。
最值得花预算的,是那些还没有“塌缩”的状态,也就是后续既可能成、也可能败的位置。
从这个角度看,TRACE的目标不是让模型看见更多样本,而是让模型更频繁地看见“有差异的样本”。
这对Agentic RL后训练很现实。
因为未来的Agent任务往往会更长、更开放,也更依赖工具和环境交互。只靠平铺式多采样,成本很容易失控。
当rollout变贵之后,训练效率的关键就不只是“生成多少”,而是“哪些地方值得生成”。
当然,论文也给出了限制。
TRACE主要面向outcome-only RLVR。如果任务没有清晰可验证的终局奖励,那么mixed-outcome contrast这个目标可能需要重新设计。
此外,论文里的预测器还是一个相对基础的实现。更强的预测器、更复杂的非平稳Agent环境,以及更大规模的工具使用场景,都还需要后续验证。
但至少这篇工作已经给出了一个清晰信号:
Agentic RLVR的样本效率优化,正在从“筛prompt”,走向“分配整棵rollout tree”。
论文标题:TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning
论文作者:Heming Zou、Qi Wang、Yun Qu、Yuhang Jiang、Lizhou Cai、Yixiu Mao、Ru Peng、Xin Xu、Weijie Liu、Kai Yang、Saiyong Yang、Xiangyang Ji
主要作者:邹鹤鸣,季向阳教授团队博士,研究范围覆盖大模型高效后训练、持续学习与脑启发智能,致力于提升多轮交互式智能体在RL后训练中的采样–反馈效率与决策信用分配,抑制自演进中的灾难性遗忘,并基于脑启发机制探索更省算力的前沿学习范式。在ICML、NeurIPS、ICLR等国际顶级会议以一作/共一发表多篇论文,工作贯穿算法创新与工程落地全栈优化,目前腾讯混元大模型团队科研实习生。
机构:清华大学自动化系季向阳教授团队、腾讯混元(LLM Departmen,Tencent)
论文地址:https://arxiv.org/pdf/2606.11119v1
一键三连「点赞」「转发」「小心心」
欢迎在评论区留下你的想法!
— 完 —
【学术投稿】请在工作日发送邮件至:ai@qbitai.com,标题注明【投稿】,并告诉我们:你是谁,从哪来,投稿内容附上项目/主页链接,以及联系方式。
🌟 点亮星标 🌟
内容中包含的图片若涉及版权问题,请及时与我们联系删除



评论
沙发等你来抢