TRACE: Uma Estrutura Unificada de Alocação de Orçamento de Rolagem para Aprendizado por Reforço Agêntico Eficiente

Resumo

Aprendizado por reforço com recompensas verificáveis (RLVR) é uma abordagem promissora para aprimorar o raciocínio e o comportamento agentivo em modelos de linguagem de grande escala. No entanto, a otimização de políticas intensiva em rollouts é frequentemente limitada por contraste insuficiente de recompensa, que surge quando prompts excessivamente simples ou complexos geram feedback de baixa variância e quando recompensas baseadas apenas no resultado atribuem a mesma avaliação terminal a cada decisão em um rollout de múltiplos turnos. Esforços anteriores concentraram-se em alocar os recursos de rollout disponíveis para prompts promissores, mas exploram apenas a informatividade da amostra no nível do prompt, negligenciando a variação na informatividade no nível de prefixo entre os turnos de um mesmo rollout. Este trabalho tem como alvo o RL agentivo de múltiplos turnos, modelando cada turno de pensamento-ação-observação no estilo ReAct como um nó semanticamente distinto, permitindo que a alocação de orçamento se estenda das raízes dos prompts aos prefixos no nível de turno com continuações adicionais, o que naturalmente forma rollouts estruturados em árvore. Apresentamos o Tree Rollout Allocation for Contrastive Exploration (TRACE), uma estrutura unificada de alocação de rollout que aprimora o contraste de recompensa dentro de um orçamento de amostragem fixo. Tecnicamente, o TRACE aloca o orçamento de rollout tanto para raízes de prompts quanto para prefixos intermediários com maior probabilidade de produzir recompensas terminais mistas. Um preditor generalizável compartilhado estima a probabilidade de sucesso condicional nesses pontos de ancoragem a partir dos históricos de prefixo, orientando essa alocação. A estrutura arbórea adaptativa resultante enriquece o feedback baseado apenas no resultado e amplifica o sinal de atualização da política. Empiricamente, o TRACE alcança desempenho competitivo e ganhos de eficiência em benchmarks agentivos típicos, por exemplo, melhorando a precisão média do Qwen3-14B em Multi-Hop QA em 2,8 pontos em relação às linhas de base concorrentes, com custo de amostragem equivalente.

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.