TRACE: Единая структура распределения бюджета развертывания для эффективного агентного обучения с подкреплением

Аннотация

Обучение с подкреплением на основе верифицируемых вознаграждений (RLVR) является перспективным подходом для улучшения рассуждений и агентного поведения больших языковых моделей. Однако оптимизация политики с интенсивным развёртыванием часто ограничена недостаточной контрастностью вознаграждений, возникающей, когда чрезмерно простые или сложные промпты генерируют обратную связь с низкой дисперсией, а также когда вознаграждения, основанные только на конечных результатах, присваивают одинаковую терминальную оценку каждому решению в многошаговом развёртывании. Предыдущие усилия были сосредоточены на распределении доступных ресурсов развёртывания по перспективным промптам, однако они учитывают информативность выборки только на уровне промпта и игнорируют вариативность информативности на уровне префиксов между шагами внутри одного и того же развёртывания. Данная работа направлена на многошаговое агентное RL путём моделирования каждого шага "мысль-действие-наблюдение" в стиле ReAct как семантически отдельного узла, что позволяет распространить распределение бюджета с корневых промптов на префиксы на уровне шагов с последующими продолжениями, естественным образом формируя древовидные развёртывания. Мы представляем TRACE (Tree Rollout Allocation for Contrastive Exploration) — единую структуру распределения развёртываний, которая улучшает контрастность вознаграждений в рамках фиксированного бюджета выборки. Технически TRACE выделяет бюджет развёртывания как корневым промптам, так и промежуточным префиксам, которые с наибольшей вероятностью приведут к смешанным терминальным вознаграждениям. Общий обобщаемый предиктор оценивает условную вероятность успеха в этих опорных точках на основе истории префиксов, направляя распределение. Полученная адаптивная древовидная структура обогащает обратную связь, основанную только на конечных результатах, и усиливает сигнал обновления политики. Эмпирически TRACE достигает конкурентоспособной производительности и выигрыша в эффективности на типичных агентных бенчмарках, например, повышая среднюю точность Qwen3-14B на Multi-Hop QA на 2,8 пункта по сравнению с сильными базовыми линиями при равной стоимости выборки.

English

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.