Tsinghua and Tencent Propose TRACE Framework to Optimize Rollout Budget in LLM Post-Training

In Agentic RLVR, models require multiple rounds of thinking, tool calls, and environment feedback. Each complete interaction trajectory is a rollout. When tasks are long and tool calls are numerous, rollout budget becomes a training bottleneck. Worse, not all rollouts are equally useful: if a group of rollouts are all correct or all wrong, the model struggles to learn preference signals. Under outcome-only rewards, training signals are sparse and credit assignment is coarse.

Previous work mainly addressed this at the prompt level, e.g., filtering problems of moderate difficulty or allocating different numbers of rollouts. But in agent tasks, the prompt is just the root node of a tree; intermediate prefixes (historical states) also determine subsequent success or failure. TRACE's core insight is that rollout budget should not be allocated only at the prompt level but should extend to intermediate prefix nodes.

TRACE consists of two stages. Stage 1 (global root allocation): in a prompt batch, based on predicted success probability, skip or increase sampling so that the rollout group for the same prompt is more likely to contain both successes and failures. Stage 2 (local prefix expansion): fork from already sampled intermediate prefixes to create counterfactual contrasts—if the original branch succeeded, the extra branch expects at least one failure, and vice versa.

TRACE relies on a shared prefix value predictor to estimate the subsequent success probability of root nodes and prefix nodes. The predictor shows usable ranking correlation at both prompt-level and prefix-level, indicating that historical prefixes contain learnable local uncertainty signals.

TRACE is not a new RL optimizer but an upstream rollout acquisition layer. It decides sampling locations and fork points, after which it can interface with tree-aware policy optimizers such as GRPO, PCL, or TreePO.

The paper tests on three types of multi-round agentic tasks: mathematical reasoning (DeepScaler), multi-hop QA (HotpotQA), and function calling (BFCL), using Qwen3-8B and Qwen3-14B as policy models. All methods use the same rollout budget metric.

Main results: TRACE yields overall accuracy curves above GRPO, PCL, and TreePO under the same budget. For example, on Qwen3-14B, math reasoning average score rises from 73.5 (GRPO) to 74.9; multi-hop QA from 51.2 to 54.0; function calling from 46.1 to 48.0. TRACE does not rank first on every individual benchmark, but its average is more stable, indicating improved overall training efficiency.

The effective sample ratio (proportion of samples forming success/failure contrasts) improves significantly: on Qwen3-8B DeepScaler, GRPO achieves 26.8% while TRACE reaches 60.6%; on Qwen3-14B, it rises from 34.7% to 59.7%. This confirms that TRACE's core benefit comes from contrast construction.

Ablation experiments are conducted on Qwen3-8B Multi-Hop QA. Using only root allocation (Stage 1) raises the effective sample ratio from 42.8 to 49.1; using only prefix allocation (Stage 2) raises accuracy to 50.0 and effective sample ratio to 47.3; combining both yields the highest accuracy and effective sample ratio. This shows that budget allocation in Agentic RLVR is not a single-layer problem.

Budget shape experiments show that with the same total budget of 2048, the (1024,2) setting (wider root coverage) outperforms (512,6) (deeper branching). This indicates that rollout efficiency depends not only on total count but also on budget shape. For multi-hop QA, wider root coverage may be more valuable.

TRACE is primarily designed for outcome-only RLVR. If tasks lack clear verifiable terminal rewards, the mixed-outcome contrast objective may need redesign. In the current implementation, the predictor scoring path incurs additional computational overhead, but the predictor update itself accounts for a small portion. The paper argues that when rollouts themselves are very expensive, spending some computation to judge where rollouts are more worthwhile is cost-effective.

TRACE's significance lies in shifting the sampling problem in Agentic RLVR from 'how many trajectories to sample' to 'where to fork.' As agent tasks become longer and more open-ended, flat multi-sampling becomes prohibitively expensive. TRACE offers a way to obtain effective training signals more frequently under a fixed budget.