TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning

2026-06-09

AI Summary

In multi-turn agent-based reinforcement learning, sparse, outcome-only rewards often hinder policy optimization due to insufficient reward contrast. This work proposes a tree-structured ReAct reasoning framework that extends rollout budget allocation from the prompt level to prefix nodes at the turn level, enabling an adaptive tree-based rollout structure. A shared, generalizable success probability predictor dynamically estimates conditional success likelihood based on prefix history, guiding exploration toward high-information nodes. Under a fixed sampling budget, the approach substantially enhances reward contrast, yielding a 2.8 percentage point average accuracy improvement over strong baselines on benchmarks such as Qwen3-14B multi-hop question answering.

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.

Problem

• reward contrast
• multi-turn agentic reinforcement learning
• rollout budget allocation
• outcome-only rewards
• prefix-level informativeness

Innovation

• rollout allocation
• reward contrast
• tree-structured rollouts
• agentic reinforcement learning
• prefix-level informativeness