Heming Zou, Qi Wang, Yun Qu, Yuhang Jiang, Lizhou Cai, Yixiu Mao, Ru Peng, Xin Xu, Weijie Liu, Kai Yang, Saiyong Yang, Xiangyang Ji
5 min
Abstract
Reinforcement learning with verifiable rewards (RLVR) is a promising approach for enhancing reasoning and agentic behavior in large language models. However, rollout-intensive policy optimization is often limited by insufficient reward contrast, arising when overly simple or complex prompts generate low-variance feedback and when outcome-only rewards assign the same terminal assessment to every decision in a multi-turn rollout. Past efforts have focused on allocating available rollout resources to promising prompts, yet they only leverage sample informativeness at the prompt level and neglect variation in prefix-level informativeness across turns within the same rollout. This work targets multi-turn agentic RL by modeling each ReAct-style thought-action-observation turn as a semantically distinct node, allowing budget allocation to extend from prompt roots to turn-level prefixes with further continuations, which naturally forms tree-structured rollouts. We introduce Tree Rollout Allocation for Contrastive Exploration (TRACE), a unified rollout allocation framework that enhances reward contrast within a fixed sampling budget. Technically, TRACE allocates rollout budget to both prompt roots and intermediate prefixes that are most likely to yield mixed terminal rewards. A shared generalizable predictor estimates conditional success probability at these anchors from prefix histories to guide this allocation. The resulting adaptive tree structure enriches outcome-only feedback and amplifies the policy-update signal. Empirically, TRACE achieves competitive performance and efficiency gains on typical agentic benchmarks, e.g., improving Qwen3-14B Multi-Hop QA average accuracy by 2.8 points over competitive baselines at equal sampling cost.
Sam: Does it actually perform better in practice?
Alex: The study suggests it does. On a benchmark involving multi-step reasoning questions — the kind where you have to chain several pieces of information together to reach an answer — TRACE improved accuracy by nearly three percentage points compared to methods that spread their budget evenly across all attempts. That's a meaningful gain without spending any additional compute.
Sam: That's worth paying attention to, because it means the improvement comes from being smarter about the same resources, not just throwing more power at the problem.
Alex: Exactly. And there's a practical engineering challenge the paper also addresses. When you're running a tree of decision paths, different branches finish at very different speeds. Some paths are short; others take much longer. In a standard setup, the whole batch has to wait for the slowest one to finish before anything gets processed — what the paper calls the "long-tail" problem.
Sam: Like a group project where everyone has to wait for the one person who hasn't finished yet.
Alex: A good comparison. TRACE handles this by keeping all the rollouts for a given task on the same computing worker, so a slow branch in one task doesn't hold up progress on everything else. Each worker manages its own workload independently.
Sam: And the predictor itself — does it get better over time, or is it fixed from the start?
Alex: It learns continuously. It looks at which paths actually succeeded and which failed, and updates its estimates accordingly. So early in training, its guesses about where the fifty-fifty points are might be rough. But as the agent accumulates experience, the predictor gets sharper at identifying exactly where the uncertainty lies.
Sam: Which means the whole system feeds itself. Better predictions lead to better-targeted training, which produces better data for the predictor to learn from.
Alex: That's the core loop. It's a self-reinforcing cycle — identify uncertainty, concentrate resources there, learn from the contrast, and repeat. What makes TRACE notable isn't any single component, but the way these pieces work together to turn a relatively unfocused process into a deliberate one.
Sam: It's a bit like the difference between a student who just re-reads their notes hoping something sticks, versus one who actively hunts for the specific gaps in their understanding.
Alex: That's a fair way to put it. And the paper's broader argument is that as AI tasks get longer and more complex, this kind of targeted allocation will matter more, not less — because the gap between useful and useless training data only grows as tasks become harder.
Sam: A genuinely interesting direction. Thanks for walking through it.
Alex: Thanks for listening to ResearchPod.