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

Resumo Técnico: TRACE - Um Framework Unificado de Alocação de Orçamento de Rollout para Aprendizado por Reforço Agêntico Eficiente

1. Declaração do Problema

O Aprendizado por Reforço com Recompensas Verificáveis (RLVR) emergiu como um mecanismo crítico para aprimorar o raciocínio e os comportamentos agênticos em Grandes Modelos de Linguagem (LLMs). No entanto, os pipelines padrão de RLVR enfrentam ineficiências computacionais significativas devido ao alto custo de gerar trajetórias longas de Cadeia de Pensamento (CoT) e interações intercaladas com o ambiente.

O desafio central reside no contraste de recompensa. Em configurações de recompensa baseadas apenas no resultado (onde uma única recompensa binária é atribuída ao final de um rollout), os sinais de aprendizado são frequentemente esparsos e desiguais:

• Nível de Prompt: Prompts excessivamente simples ou excessivamente complexos geram feedback de baixa variância (ex: todos sucessos ou todos falhas), oferecendo pouco sinal para atualizações de política.
• Nível de Trajetória: Atribuir uma única recompensa terminal a um rollout longo e de múltiplos turnos obscurece a atribuição de crédito local. Sem feedback intermediário, o modelo não consegue distinguir quais turnos ou decisões específicas contribuíram para o resultado final.

Métodos existentes tentam mitigar isso alocando orçamentos de rollout para prompts promissores (seleção de prompt) ou aumentando o número de rollouts por prompt (alocação de contagem de rollout). No entanto, essas abordagens tratam os rollouts como trajetórias atômicas. Uma vez que um prompt é selecionado, todo o rollout é gerado linearmente, negligenciando o potencial de ramificação em turnos intermediários (prefixos) onde o resultado é incerto. Isso resulta em uma coleção "plana" de rollouts que falha em maximizar a densidade de informação dentro de um orçamento de amostragem fixo.

2. Metodologia: TRACE

Os autores propõem o TRACE (Tree Rollout Allocation for Contrastive Exploration), um framework unificado que reimagina a alocação de rollout como um processo estruturado em árvore. Em vez de gerar trajetórias lineares independentes, o TRACE aloca dinamicamente o orçamento tanto para raízes de prompts quanto para prefixos intermediários para maximizar a probabilidade de observar recompensas terminais mistas (sucessos e falhas) dentro do mesmo contexto de condicionamento.

Componentes Técnicos Principais:

• Rollouts Estruturados em Árvore:

  • O framework modela interações no estilo ReAct (Pensamento-Ação-Observação) como nós em uma árvore.
  • Raízes: O prompt inicial ($x$).
  • Prefixos: Históricos intermediários ($H_t$) formados após $t$ turnos.
  • Ramificação: Em qualquer prefixo visitado, o sistema pode alocar orçamento adicional para gerar continuações "contrafatuais" (novos sufixos) que divergem da ramificação "factual".

• Objetivo de Alocação Unificado (Contraste de Recompensa Mista):

  • O objetivo é alocar orçamento para âncoras (raízes ou prefixos) onde a probabilidade de sucesso condicional ($V^\pi_t$) não é nem 0 nem 1, mas intermediária.
  • Base Teórica: O artigo estabelece que a variação quadrática esperada da previsão de sucesso condicional ao longo de um sufixo de rollout é igual a $V^\pi_t(1 - V^\pi_t)$. Este valor representa o "potencial de contraste". Alocar orçamento para âncoras com alto potencial de contraste maximiza a probabilidade de observar ramos irmãos com resultados terminais opostos, criando sinais de preferência pareados implícitos para o otimizador.

• Processo de Alocação de Dois Estágios:

  1. Alocação Global de Raiz: Antes da amostragem, um preditor estima a probabilidade de sucesso para um conjunto de prompts candidatos. O orçamento é atribuído a prompts que provavelmente gerarão resultados mistos (excluindo aqueles com sucesso ou falha determinísticos). Prompts com uma probabilidade prevista próxima de 0 ou 1 recebem orçamento zero.
  2. Expansão de Prefixo Local: Para prompts ativos, o sistema identifica prefixos não terminais visitados. Ele aloca um orçamento local para gerar continuações adicionais a partir desses prefixos. A alocação prioriza prefixos onde o resultado factual observado difere da probabilidade prevista (ou seja, onde uma ramificação contrafatual é provável de inverter o resultado).

• Preditor Compartilhado Generalizável:

  • Um preditor leve ($\tilde{V}_\psi$) estima a probabilidade de sucesso condicional a partir de históricos de prefixos.
  • Ele é treinado via regressão de erro quadrático médio em alvos recursivos derivados das árvores de rollout coletadas (médias empíricas bottom-up de recompensas terminais).
  • Este preditor guia tanto a seleção de raiz quanto as decisões de ramificação de prefixo.

• Otimização de Política Consciente de Árvore:

  • As estruturas de árvore resultantes fornecem pares de preferência passo a passo implícitos (ex: uma continuação bem-sida vs. uma continuação falha a partir do mesmo prefixo).
  • Estes são usados para atualizar a política, aproveitando sinais de gradiente de grupo-relativo ou conscientes de árvore que são mais densos do que as atualizações padrão de apenas resultado.

3. Principais Contribuições

1. Formulação Unificada: O trabalho formula a alocação de contraste de recompensa mista como uma visão única de RLVR eficiente em amostras, unificando filtragem de prompt, alocação de contagem de rollout e ramificação de nível de prefixo como decisões estratégicas de orçamento sobre âncoras de árvore de rollout.
2. Derivação Teórica: Os autores derivam utilidades de alocação baseadas em um objetivo de conjunto de descendentes compartilhado, provando que a alocação de nível de prefixo melhora estritamente a previsão de dificuldade de grupo e que alocar orçamento para incerteza intermediária maximiza o potencial de contraste.
3. Framework TRACE: A instanciação desta visão em um framework de rollout em árvore guiado por preditor e com orçamento restrito, que integra alocação global de raiz com expansão de prefixo local, aumentando os sinais de crédito implícitos sem exigir recompensas de processo manufaturadas.

4. Resultados Experimentais

Os autores avaliaram o TRACE em três benchmarks agênticos de múltiplos turnos: Raciocínio Matemático (DeepScaler, AIME, etc.), QA de Múltiplos Saltos (HotpotQA, 2Wiki, etc.) e Chamada de Função (BFCL v4), usando backbones Qwen3-8B e Qwen3-14B.

• Ganhos de Desempenho: Sob um orçamento de amostragem igual, o TRACE superou consistentemente os baselines, incluindo GRPO padrão, métodos de seleção de prompt (PCL) e métodos de árvore aleatórios (TreePO).
  • Exemplo: Em QA de Múltiplos Saltos, o TRACE melhorou a precisão média do Qwen3-14B em 2.8 pontos sobre baselines competitivos.
  • Exemplo: Em Raciocínio Matemático, melhorou a precisão do Qwen3-8B de 70.0% (GRPO) para 71.1%.
• Eficiência e Contraste: O TRACE aumentou significativamente a Razão Efetiva (a fração de prompts/rollouts contendo tanto sucessos quanto falhas).
  • Em Raciocínio Matemático com Qwen3-8B, a razão efetiva subiu de 26.8% (GRPO) para 60.6% (TRACE).
  • Isso indica que o TRACE redireciona com sucesso o orçamento para longe de âncoras saturadas (determinísticas) em direção a regiões ricas em contraste.
• Estudos de Ablação:
  • Tanto a alocação de raiz quanto o estágio de expansão de prefixo contribuíram independentemente para os ganhos de desempenho.
  • O framework mostrou robustez através de diferentes tamanhos de orçamento, onde a "forma" do orçamento (amplitude de cobertura de raiz vs. profundidade de continuação) provou ser crítica.

5. Significância e Alegações

O artigo afirma que o TRACE reformula a amostragem de RLVR agêntico de uma questão de "quantos rollouts independentes desenhar" para "onde ramificar dentro da árvore de rollout".

• Sinais de Aprendizado Mais Densos: Ao converter recompensas de resultado escaláveis em contrastes de recompensa mista mais densos e pares de preferência passo a passo implícitos, o TRACE amplifica o sinal de atualização da política sem aumentar o custo computacional.
• Generalização: A abordagem depende de um preditor compartilhado para raízes e prefixos, demonstrando que históricos de prefixo carregam uma noção aprendível de incerteza local que pode guiar decisões de ramificação.
• Limitações: Os autores observam que o TRACE é projetado primariamente para RLVR de apenas resultado com verificação terminal clara. Tarefas que carecem de recompensas terminais claras podem exigir a revisão do objetivo de contraste. Além disso, a instância atual do preditor é um cenário vanilla, e preditores mais fortes para tarefas agênticas específicas permanecem como uma área para trabalhos futuros.

Em resumo, o TRACE oferece um método sistemático para maximizar o valor informacional de um orçamento de rollout fixo em RL agêntico de múltiplos turnos, construindo dinamicamente rollouts estruturados em árvore que visam a incerteza e maximizam o contraste de recompensa.