TRACE: Un Marco Unificado de Asignación de Presupuesto de Despliegue para el Aprendizaje por Refuerzo Agéntico Eficiente

Resumen

El aprendizaje por refuerzo con recompensas verificables (RLVR) es un enfoque prometedor para mejorar el razonamiento y el comportamiento agéntico en modelos de lenguaje grandes. Sin embargo, la optimización de políticas intensiva en despliegues a menudo se ve limitada por un contraste de recompensa insuficiente, que surge cuando instrucciones excesivamente simples o complejas generan retroalimentación de baja varianza, y cuando las recompensas basadas únicamente en el resultado asignan la misma evaluación terminal a cada decisión en un despliegue de múltiples turnos. Trabajos anteriores se han centrado en asignar los recursos de despliegue disponibles a instrucciones prometedoras, pero solo aprovechan la informatividad de la muestra a nivel de instrucción y descuidan la variación en la informatividad a nivel de prefijo entre los distintos turnos dentro de un mismo despliegue. Este trabajo aborda el RL agéntico de múltiples turnos modelando cada turno de pensamiento-acción-observación al estilo ReAct como un nodo semánticamente distinto, lo que permite que la asignación de presupuesto se extienda desde las raíces de las instrucciones hasta los prefijos a nivel de turno con continuaciones adicionales, formando naturalmente despliegues en estructura de árbol. Introducimos la Asignación de Despliegues en Árbol para Exploración Contrastiva (TRACE, por sus siglas en inglés), un marco unificado de asignación de despliegues que mejora el contraste de recompensa dentro de un presupuesto de muestreo fijo. Técnicamente, TRACE asigna el presupuesto de despliegue tanto a raíces de instrucciones como a prefijos intermedios que tienen mayor probabilidad de generar recompensas terminales mixtas. Un predictor generalizable compartido estima la probabilidad condicional de éxito en estos anclajes a partir de los historiales de prefijos para guiar esta asignación. La estructura arbórea adaptativa resultante enriquece la retroalimentación basada únicamente en el resultado y amplifica la señal de actualización de la política. Empíricamente, TRACE logra un rendimiento competitivo y ganancias en eficiencia en benchmarks agénticos típicos; por ejemplo, mejora la precisión promedio de Qwen3-14B en Preguntas y Respuestas de Múltiples Saltos en 2.8 puntos en comparación con líneas base competitivas, con igual costo de muestreo.

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.