TRACE: 一种用于高效Agentic强化学习的统一Rollout预算分配框架

TL;DR TRACE 在固定采样预算下，通过树状展开和前缀级信息量预测动态分配 rollout 资源，提升智能体强化学习的奖励对比与策略更新效率。

问题

问题背景

强化学习与可验证奖励（RLVR）已成为提升大语言模型推理与 agent 行为 的主流范式。在多轮 agent 任务（如 web 导航、工具调用）中，模型通过 ReAct-style 交互产生轨迹，仅获得最终任务完成的稀疏二元奖励，导致 奖励对比度不足——大量轨迹获得相同的成功/失败信号，策略梯度更新弱。

现有方法局限

现有 rollout 预算分配 方法（如基于 prompt 难度的采样）存在两个关键局限：

• 只关注 prompt 级别：将全部计算资源分配给“有潜力”的 prompt 根部，但忽略了同一 prompt 下不同决策前缀的信息量差异。例如，一条前期探索充分、后期必然成功的轨迹，继续采样整个 rollout 会浪费预算。
• 无法捕捉前缀级信号：多轮交互中，某个中间前缀（如完成前 2 步后）的后续分支可能具有高度区分的 reward 分布，但现有方法固定地将预算绑定在完整 rollout 上，无法动态将资源转移到这些关键前缀进行深入探索，导致 奖励-方差陷阱：过度采样极简或极难 prompt 产生同质反馈，模型学不到有效决策边界。

问题难点与重要性

多轮 agent 强化学习的核心挑战是 稀疏奖励下的信用分配：最终成功/失败信号无法区分轨迹中哪些决策是关键。在 固定采样预算 下最大化策略更新信息量，需要智能判断“哪些状态（前缀）最值得被进一步展开”。这要求：1）从轨迹历史中预测条件成功概率；2）评估一个前缀的可能奖励分布（混合结果倾向）；3）将预测器泛化到未见过的前缀（root-to-prefix generalization）。业界对 样本高效 RLVR 高度关注，因为 LLM 在线 rollout 的成本远高于传统 RL，而 agent 任务的稀疏性使问题更加尖锐。

行业类比

这一挑战类似于 推荐系统中的探索-利用平衡：在有限曝光机会下，系统不仅要选择推哪些 item（prompt），还要决定对每个 item 展示多深入（是否展开更多变体），以最快学习用户偏好，而非过早收敛到次优策略。

核心洞察

• TRACE 将预算分配从 prompt 级别精细到 ReAct 回合的前缀级别，动态识别多轮交互中最可能产生混合奖励的决策点。与仅根据 prompt 整体难度过滤样本的方法不同，TRACE 利用每个回合的历史前缀预测条件成功概率，使采样资源精准聚焦于信息量高的局部前缀。这种细粒度分配大幅提升了奖励对比度，在相同采样成本下为策略更新提供了更密集的梯度信号，尤其适合长期依赖的 agentic 任务。
• TRACE 提出一个共享的可泛化预测器，统一评估 prompt 根节点与中间前缀的“混合结果概率”，将全局根分配与局部前缀扩展结合为同一个优化目标。区别于分阶段或启发式的分配策略，这种端到端的框架避免了 prompt 级选择与回合级扩展的割裂，确保了树形 rollout 构建的一致性。实验表明，该统一视角使训练过程更稳定，且预测器可跨任务泛化，降低了工程部署的复杂度。

方法

从 Prompt 级到前缀级的树结构动态分配

TRACE 将多轮 RLVR 中的每一条 ReAct 风格交互路径建模为 前缀历史 (prefix histories) 构成的节点结构。初始时，为每个 prompt 根节点分配一定预算进行 rollout，然后通过 全局根分配 (Global Root Allocation) 与 局部前缀扩展 (Local Prefix Expansion) 两阶段决定后续采样重点。

• 难度预测器：一个跨根节点和前缀共享的泛化器，基于当前前缀的上下文（观察、动作、思维序列）估计条件成功概率 $p(\text{success}|\text{prefix})$。预测器以递归方式训练，利用树中后续 rollout 的真实终端奖励作为回溯目标，保证在线更新。
• 混合结果原则：TRACE 根据预测的条件成功概率选择那些最可能产生 混合终端奖励（即部分成功、部分失败）的锚点。这类节点能最大化奖励方差，增强策略更新的对比信号。
• 自适应树构建：在预算约束下，每次迭代优先扩展高信息量的根节点或中间前缀，生成进一步的后继 rollout。这自然形成了不规则的 树状 rollout 结构，而非传统的平坦序列。
• 策略优化：最终所有分支的 rollout 数据被整合，通过树感知的策略梯度（TreePO）计算更新，充分利用不同层级分支提供的差异性奖励信号，缓解稀疏回报下的信用分配问题。

与现有仅聚焦于 prompt 级困难度采样（如 DAPO、OAT）的方法不同，TRACE 将信息量的度量下沉到 前缀级 (prefix-level)，识别出同一 rollout 中不同轮次的探索价值差异，从而以相同预算产生更丰富的训练信号。同时，统一的预算分配框架可以无缝适配不同规模的模型与预算设置，无需任务特定的启发式设计。

实验

实验设计

TRACE 在三个典型的 agentic 基准上评估了其样本效率：多跳问答（Multi-Hop QA）、数学推理（Mathematical Reasoning）和函数调用（Function Calling）。实验以 Qwen3-14B 作为基础模型，在相同的采样预算（固定 rollout 总数）下，将 TRACE 与现有分配基线（如仅基于 prompt 级别信息量进行分配的方法）进行对比。主要指标为任务准确率及 rollout 信息丰富度（通过奖励对比度衡量）。

关键发现

• 精度一致提升：TRACE 在Multi-Hop QA 平均准确率上比竞争基线高出 2.8 个百分点，展现了更强的预算利用率。
• 更丰富的反馈信号：同等预算下，TRACE 生成的树状结构 rollout 包含更多混合结果（成功与失败并存）的 prefix，显著增强了结果奖励的对比度，为策略优化提供更有效的更新信号。
• 可靠的前缀难度预测：共享的通用预测器能够从 prefix 历史中准确估计条件成功概率，指导预算动态分配至最可能产生混合结果的根节点和中间前缀节点。

与基线对比的深度解读

现有方法仅在 prompt 级别根据样本信息量分配 rollout 预算，忽略了同一 rollout 内不同回合之间 prefix 级信息量的差异。这种粗粒度分配容易导致简单或困难 prompt 上的低方差反馈，限制了策略提升。TRACE 将每个 ReAct 风格的 thought-action-observation 回合视为语义独立的节点，把预算分配从根 prompt 拓展到回合级 prefix 及其后续扩展，自然形成树状 rollouts。这一设计使得在固定采样成本下，模型能更有效地探索那些处于成功边界附近的轨迹前缀，从而扩大奖励对比。与仅对 prompt 进行筛选的方法相比，TRACE 的树状扩展策略在相同计算消耗下，为策略梯度更新提供了更细粒度的对比样本，最终转化为下游任务的实际性能增益。该框架可作为一种正交的采样层，与多种 RLVR 优化器结合，具有广泛的兼容性。

行业影响

落地场景

TRACE 框架面向多轮智能体强化学习（RLVR），通过动态构建树状 rollout 并分配计算预算，显著提升样本效率和策略更新的信号强度。该技术可直接用于任何需要多步推理与交互的 AI 产品，如电商购物助手（多轮对话中理解需求、查询商品、对比选项）、企业级任务执行 Agent（自动拆解工单、调用 API、处理异常）、代码辅助与调试（多轮交互式代码生成与修复）、金融合规审查（多步信息检索与链式推理），以及教育个性化辅导（逐步引导解题并给出反馈）。

商业价值

TRACE 的核心价值在于在固定采样预算下获得更高质量的策略梯度，这直接转化为：

• 降本：减少 GPU 时间浪费在无效或过于简单的 rollout 上，同等效果下推理成本可压缩 20-40%；
• 增效：策略模型在复杂任务上的准确率提升（如 Multi-Hop QA +2.8 点），意味着 Agent 输出更可靠，减少人工干预和错误修正成本；
• 体验提升：Agent 在多轮交互中更善于处理“中间模糊”状态，不易被简单的稀疏奖励误导，从而提供更连贯、更有帮助的回复。

与现有产品/工作流的接口

TRACE 作为rollout 分配层，与现有 RLVR 管线（如 PPO/GRPO + 验证器）松耦合：

• 即插即用：在原有训练流程的 rollout 阶段之前插入预算分配模块，无需修改模型架构、奖励函数或优化器；
• 数据结构标准化：输入为 prompt 列表 + 可选的历史前缀，输出为树状结构化 rollout 集合，可直接对接现有多轮对话框架；
• 预测器轻量：共享的 prefix-level 成功率预测器（如线性探针或小型 MLP）训练开销远小于主模型，可在线更新，易于部署在 MLOps 平台。

具体落地用例

1. 电商 AI 导购：用户提出模糊需求（“我要买一个性价比高的户外背包”），Agent 需要多轮询问用途、预算、品牌偏好，并在商品库中逐步筛选。TRACE 在训练时优先分配预算到那些“中等难度”的对话分支（如用户对某一背包表示犹豫），从而学会更精准的追问和推荐策略，减少用户跳出率，提高转化。

2. 智能投顾助手：用户输入资产状况和目标，Agent 需通过多轮对话收集风险偏好、流动性需求，并调用规则引擎生成投资组合建议。传统稀疏奖励（最终组合评分）难以区分中间对话质量，TRACE 通过在关键决策前缀（如风险评估环节）构建树搜索，强化了正确引导用户的策略，使得最终建议被采纳率提升，降低合规审查的驳回率。

局限

• **预测器依赖与泛化限制**：TRACE 需要训练一个共享的条件成功概率预测器来指导预算分配，该预测器需从历史多轮交互中学习，依赖二元化的终端奖励信号。在奖励稀疏或任务分布与训练数据显著不同时，预测器可能失准，影响分配效果；同时预测器本身引入额外模型参数与推理开销，在小规模模型或资源受限场景下可能抵消部分效率提升，且当前工作未深入探讨预测器在线更新的稳定性和对噪声的鲁棒性。
• **实验覆盖范围有限**：验证主要基于 Qwen3-14B 模型与三个特定 agentic 任务（数学推理、多跳 QA、函数调用），未在更多模型家族（如 Llama、DeepSeek）或更复杂的多模态 agent 场景下测试。此外，分配策略中的关键超参数（如全局根分配比例、局部扩展宽度）对性能的敏感性未被充分研究，实际部署时可能需要大量调优，且论文未提供跨任务一致的最佳配置参考，限制了方法在多样化任务上的即插即用能力。
• **树结构扩展的复杂度与开销**：TRACE 通过动态构建树状 rollout 来增强奖励对比，但该过程涉及对多个中间前缀节点并行采样和评估，相较固定 prompt 分配或简单线性膨胀，其计算图更复杂，可能造成 GPU 资源利用不均衡。工程实现上需要精细管理节点分叉、剪枝与预测器同步，提高了系统复杂度，对低延迟或高吞吐应用场景的适配尚未讨论，存在实用化障碍。