TRACE: 효율적인 에이전트 기반 강화 학습을 위한 통합된 배포 예산 할당 프레임워크

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

검증 가능한 보상을 활용한 강화 학습(RLVR)은 대규모 언어 모델에서 추론 능력과 에이전트 행동을 향상시키는 유망한 접근 방식입니다. 그러나 정책 최적화 과정에서 많은 롤아웃(rollout, 시뮬레이션 실행)이 필요하며, 이는 종종 불충분한 보상 대비로 인해 제한됩니다. 과도하게 단순하거나 복잡한 프롬프트가 낮은 변동성을 가진 피드백을 생성하고, 결과 정보만을 사용하는 보상이 다중 턴 롤아웃에서 모든 결정에 동일한 최종 평가를 부여하는 경우 이러한 문제가 발생합니다. 기존 연구는 사용 가능한 롤아웃 자원을 유망한 프롬프트에 할당하는 데 초점을 맞추었지만, 이는 프롬프트 수준의 샘플 정보만 활용하고 동일한 롤아웃 내에서 턴(turn) 간의 접두사(prefix) 수준 정보 변동성을 고려하지 못합니다. 본 연구에서는 ReAct 스타일의 사고-행동-관찰 턴을 의미적으로 구별되는 노드로 모델링하여 다중 턴 에이전트 기반 강화 학습을 목표로 합니다. 이를 통해 예산 할당은 프롬프트 루트에서 시작하여 턴 수준의 접두사로 확장되며, 추가적인 연관 관계를 형성하여 트리 구조의 롤아웃을 자연스럽게 생성합니다. 우리는 Tree Rollout Allocation for Contrastive Exploration (TRACE)이라는 통합된 롤아웃 할당 프레임워크를 제안하며, 이는 고정된 샘플링 예산 내에서 보상 대비를 향상시킵니다. TRACE는 기술적으로 다양한 최종 보상을 얻을 가능성이 높은 프롬프트 루트와 중간 접두사에 롤아웃 예산을 할당합니다. 공유되고 일반화 가능한 예측 모델은 이러한 지점에서의 조건부 성공 확률을 추정하여, 이 할당 과정을 안내합니다. 결과적으로 생성된 적응형 트리 구조는 결과 정보만을 사용하는 피드백을 풍부하게 만들고 정책 업데이트 신호를 증폭시킵니다. 실험적으로 TRACE는 일반적인 에이전트 기반 벤치마크에서 경쟁력 있는 성능과 효율성 향상을 달성했습니다. 예를 들어, Qwen3-14B Multi-Hop QA의 평균 정확도를 경쟁 모델 대비 2.8 포인트 개선하는 결과를 얻었습니다 (동일한 샘플링 비용 기준).

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.