TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning
在固定 rollout 预算下，通过树结构分配——而非平铺均匀采样——将 outcome-only 奖励转化为更密集的对比信号，从而显著提升多轮 Agentic RL 的训练效率。
研究动机
1. 问题难在哪
RLVR（Reinforcement Learning with Verifiable Rewards）已成为训练 LLM 推理和 Agentic 能力的核心范式。然而在多轮 ReAct 风格的 Agent 任务中，rollout 代价极高（需要长链思考 + 环境交互），而学习信号却极度稀疏：
- Prompt 级别：过简单或过难的 prompt 会产生接近 0 或 1 的成功率，奖励方差极低，对模型更新几乎无益。
- 轨迹级别：outcome-only 奖励将整条轨迹的所有 token 分配同一个终端奖励，无法区分哪一步决策对结果起了关键作用（credit assignment 稀疏）。
2. 现有方法的不足
| 方法类型 | 策略 | 局限 |
|---|---|---|
| GRPO | 随机均匀采样 | 大量 rollout 来自"无对比度"的 prompt，浪费预算 |
| PCL（Prompt Curriculum Learning） | 按难度过滤 prompt | 只在 prompt root 层操作，忽略轨迹内部 prefix 的信息差异 |
| TreePO | 树结构 rollout + 随机 branching | 引入了树结构但分支点选取随机，未利用 prefix 级别的不确定性 |
3. 核心 Insight
将 rollout 采样从"平铺"转化为树结构分配问题：prompt root 和轨迹中每个 turn 结束后的 prefix 都是候选"anchor"，将预算分配给最可能产生混合终端奖励（既有成功又有失败）的 anchor，从而最大化奖励对比信号的密度。
核心方法
1. 形式化：多轮 RLVR 的 Prefix History
每个 ReAct turn 被建模为节点 \(n_t := \langle \tau_t, a_t, o_t \rangle\)（thought、action、observation），prefix history 为 \(H_t := (x, n_1, \ldots, n_t)\)。针对任意 prefix \(H_t\) 定义条件成功概率：
\[ V_t^\pi := \mathbb{E}^\pi\bigl[r(H_T) \mid H_t\bigr] \]
这是 TRACE 中所有分配决策的核心评分量。
2. 混合奖励对比——统一分配原则
对于任意 anchor（root 或 prefix），分配 \(m\) 个 continuation 后产生混合奖励的概率为：
\[ \text{Contrast}(V, m) = 1 - V^m - (1-V)^m \]
该函数在 \(V=0.5\) 处最大，在 \(V \to 0\) 或 \(V \to 1\) 时趋近 0——只有中间难度的 anchor 才值得分配预算。
命题 2 进一步证明，prefix 处的剩余对比潜力等于 Bernoulli 方差：\(\mathbb{E}^\pi[[Z]_{t:T} \mid H_t] = V_t^\pi(1 - V_t^\pi)\)，从而为上述启发式提供了理论依据。
3. 阶段一：全局 Root 分配
给定候选 prompt 池 \(\{x_1,\ldots,x_B\}\)，预测每个 prompt 的成功概率 \(\tilde{V}_\psi(x_i)\)，求解如下整数规划：
\[ \max_{m_1,\ldots,m_B} \sum_{i=1}^{B} \bigl[1 - \tilde{v}_i^{m_i} - (1-\tilde{v}_i)^{m_i}\bigr] \quad \text{s.t.} \sum m_i = M,\ m_i \in \{0\} \cup \{2,\ldots,M\} \]
\(m_i=0\) 表示跳过该 prompt；\(m_i \geq 2\) 确保组内至少有两条 rollout 以支持 GRPO 式的组对比更新。动态规划求解，计算开销可忽略。
4. 阶段二：局部 Prefix 扩展
每个活跃 prompt \(x_i\) 完成 \(m_i\) 条 bare rollout 后，对每个已访问的 non-terminal prefix \(H_{i,j,t}\) 计算分配价值：
\[ V_{\text{pref}}(i,j,t,k) := 1 - \bigl[r_{i,j}\,\tilde{V}_\psi(H_{i,j,t}) + (1-r_{i,j})(1-\tilde{V}_\psi(H_{i,j,t}))\bigr]^k \]
即"再采样 \(k\) 条 continuation 后，至少有一条翻转已观察奖励"的概率。然后在局部预算 \(m_i N\) 下求解分配，不等待其他 prompt（prompt-local，系统友好）。
5. 共享成功率预测器
在线训练单个轻量 \(\tilde{V}_\psi(H_t)\)，同时服务 root 评分和 prefix 评分。训练目标是树上的递归经验均值：
\[ \hat{V}(y) = \frac{1}{n_y} \sum_{c \in C(y)} n_c \hat{V}(c), \quad \mathcal{L}_{\text{value}} = \frac{1}{|\mathcal{S}|}\sum_{y\in\mathcal{S}} (\tilde{V}_\psi(y) - \hat{V}(y))^2 \]
root 样本更多，anchor prefix 样本较少，但实验表明预测器能将 prefix 难度迁移到从未见过的中间步骤上。
6. 系统整体流程
7. 方法对比一览
| 方法 | Prompt 选择 | Root 预算 | Prefix 扩展 | 树结构更新 | 学习式分配 |
|---|---|---|---|---|---|
| GRPO | 随机 | 均匀 | × | × | × |
| PCL | 预测式 | 固定 | × | × | 仅 root |
| TreePO | 随机 | 均匀 | 随机 | ✓ | × |
| TRACE | 预测式 | 自适应 | 预测式 | ✓ | Root + Prefix |
主要实验结果
实验设置
- 模型：Qwen3-8B、Qwen3-14B（附录另报 Llama-3.2-3B）
- 数学推理：训练集 DeepScaler；测试 AIME24、AMC23、MATH500、MinervaMath、OlympiadBench（in-dist）以及 MMLU-Pro、ARC-c、GPQA（out-of-dist）
- Multi-Hop QA：训练 HotpotQA；测试 HotpotQA、2WikiMultiHopQA、MuSiQue、Bamboogle，配合本地 E5 检索服务器
- Function Calling：训练 BFCL v4 多轮分割 80%；测试剩余 20%（Base、Long、Miss-Func、Miss-Param）
- 预算控制：所有方法使用完全相同的 rollout 预算，增益来自分配而非额外采样
准确率提升（vs GRPO 基线）
有效比率（Effective Ratio）提升
有效比率 = batch 中奖励组包含"成功 + 失败"的 prompt 比例。越高表明每次更新获得的对比信号越多。
详细数值对比
数学推理完整表格（Qwen3-8B & 14B）
| 模型 | 方法 | AIME24 | AMC23 | MATH500 | Minerva | Olympiad | Avg↑ | MMLU-Pro | ARC-c | GPQA | OOD Avg↑ |
|---|---|---|---|---|---|---|---|---|---|---|---|
| 8B | ReAct | 52.3 | 87.0 | 89.8 | 37.8 | 58.6 | 65.1 | 68.5 | 95.3 | 52.4 | 72.1 |
| GRPO | 63.6 | 91.0 | 92.3 | 40.2 | 62.8 | 70.0 | 72.7 | 95.3 | 55.9 | 74.6 | |
| PCL | 64.0 | 91.5 | 92.8 | 40.7 | 63.0 | 70.4 | 72.9 | 95.4 | 56.2 | 74.8 | |
| TreePO | 64.3 | 91.3 | 92.6 | 40.9 | 63.1 | 70.4 | 73.0 | 95.2 | 56.4 | 74.9 | |
| TRACE | 63.9 | 91.2 | 93.4 | 41.2 | 65.8 | 71.1 | 73.4 | 95.7 | 56.9 | 75.3 | |
| 14B | ReAct | 61.5 | 89.4 | 91.8 | 40.0 | 62.3 | 69.0 | 75.1 | 96.2 | 59.2 | 76.8 |
| GRPO | 65.6 | 94.3 | 94.2 | 45.2 | 68.4 | 73.5 | 75.9 | 96.1 | 59.4 | 77.1 | |
| PCL | 66.2 | 95.1 | 95.0 | 45.1 | 67.6 | 73.9 | 76.1 | 96.2 | 59.7 | 77.3 | |
| TreePO | 66.4 | 95.0 | 94.8 | 45.4 | 67.8 | 74.0 | 76.0 | 96.4 | 59.8 | 77.4 | |
| TRACE | 66.1 | 94.7 | 94.9 | 47.3 | 71.5 | 74.9 | 76.5 | 96.5 | 60.4 | 77.8 |
Multi-Hop QA & Function Calling 完整表格
| 模型 | 方法 | Hotpot | 2Wiki | Musiq | Bamb | QA Avg↑ | Base | Miss-Func | Miss-Param | Long | FC Avg↑ |
|---|---|---|---|---|---|---|---|---|---|---|---|
| 8B | ReAct | 41.9 | 37.6 | 18.4 | 44.7 | 35.7 | 36.7 | 21.5 | 31.1 | 22.3 | 28.0 |
| GRPO | 53.2 | 48.7 | 27.4 | 64.7 | 48.5 | 58.5 | 31.0 | 45.8 | 38.7 | 43.5 | |
| PCL | 53.6 | 49.2 | 27.9 | 65.0 | 48.9 | 59.0 | 32.1 | 46.8 | 39.3 | 44.3 | |
| TreePO | 54.0 | 49.7 | 28.4 | 65.9 | 49.5 | 58.9 | 32.0 | 46.5 | 39.4 | 44.2 | |
| TRACE | 55.7 | 50.9 | 29.5 | 66.3 | 50.6 | 61.2 | 34.4 | 48.8 | 40.4 | 46.2 | |
| 14B | ReAct | 46.4 | 44.2 | 22.8 | 55.3 | 42.2 | 51.1 | 23.6 | 35.6 | 22.8 | 33.6 |
| GRPO | 55.3 | 52.7 | 30.1 | 66.8 | 51.2 | 70.6 | 32.6 | 40.9 | 36.4 | 46.1 | |
| PCL | 55.6 | 52.9 | 30.2 | 67.2 | 51.5 | 70.0 | 31.8 | 41.4 | 40.4 | 45.9 | |
| TreePO | 55.1 | 55.3 | 31.0 | 70.6 | 53.0 | 70.8 | 33.0 | 43.2 | 39.4 | 46.6 | |
| TRACE | 57.8 | 53.1 | 33.7 | 71.3 | 54.0 | 72.4 | 34.8 | 44.6 | 40.2 | 48.0 |
消融实验
两阶段各自的贡献（Qwen3-8B Multi-Hop QA）
| 阶段一（Root 分配） | 阶段二（Prefix 扩展） | Avg Acc | Avg Eff. Ratio |
|---|---|---|---|
| 均匀 | 均匀 | 49.5 | 42.8% |
| Active | 均匀 | 49.8 | 49.1% |
| 均匀 | Active | 50.0 | 47.3% |
| Active | Active | 50.6 | 52.3% |
两阶段增益可叠加：root 分配筛出有对比潜力的 prompt，prefix 扩展进一步在轨迹内部发现对比点。
预算形状的影响（M × N，总预算 = M(1 + N/2)）
| M | N | 总预算 | TreePO Acc | TRACE Acc | TreePO Eff% | TRACE Eff% |
|---|---|---|---|---|---|---|
| 512 | 2 | 1024 | 48.8 | 49.7 | 32.2 | 42.4 |
| 512 | 6 | 2048 | 49.4 | 50.3 | 37.7 | 47.8 |
| 1024 | 2 | 2048 | 49.5 | 50.6 | 42.8 | 52.3 |
相同总预算下，更宽的 root 覆盖（M=1024, N=2）优于更深的 prefix 扩展（M=512, N=6）。瓶颈不只是 rollout 数量，而是预算能否覆盖奖励尚未确定的状态。
局限与展望
- 依赖可验证终端奖励：TRACE 的混合奖励对比目标假设存在清晰的 binary terminal reward。对于没有明确验证信号的开放式生成任务（写作、对话），需要重新设计对比目标。
- 预测器质量上限：分配效果受 \(\tilde{V}_\psi\) 精度约束。本文使用了"vanilla"实现，更强的预测器（如大模型 value head、更多 prefix 监督）有望进一步提升预算利用率。
- 评估范围：主实验仅覆盖三类任务和两个 Qwen3 尺寸，更复杂的非平稳 agentic 场景（持续学习、长期环境）留待未来探索。
- 研究趋势：TRACE 代表了从"采多少 rollout"转向"在树的哪里采 rollout"的设计范式转变。随着 LLM Agent 向更长链路、更复杂工具交互发展，基于过程级别信号的采样策略将成为提升训练效率的核心手段，与 PRM、MCTS、off-policy replay 等方向形成互补。