[논문 리뷰] TSR: Trajectory-Search Rollouts for Multi-Turn RL of LLM Agents

Aladin Djuhera, Swanand Kadhe외 3명|arXiv (Cornell University)|2026. 02. 12.

Reinforcement Learning in Robotics인용 수 0

한 줄 요약

TSR은 학습 시점 트레이젝토리 검색을 도입하여 LLM 에이전트의 다중 턴 RL에서 턴당 롤아웃의 품질을 높이고, Optimizer에 종속되지 않는 통합으로 Sokoban, FrozenLake, WebShop 전반에서 성능과 안정성을 향상시킵니다.

ABSTRACT

Advances in large language models (LLMs) are driving a shift toward using reinforcement learning (RL) to train agents from iterative, multi-turn interactions across tasks. However, multi-turn RL remains challenging as rewards are often sparse or delayed, and environments can be stochastic. In this regime, naive trajectory sampling can hinder exploitation and induce mode collapse. We propose TSR (Trajectory-Search Rollouts), a training-time approach that repurposes test-time scaling ideas for improved per-turn rollout generation. TSR performs lightweight tree-style search to construct high-quality trajectories by selecting high-scoring actions at each turn using task-specific feedback. This improves rollout quality and stabilizes learning while leaving the underlying optimization objective unchanged, making TSR optimizer-agnostic. We instantiate TSR with best-of-N, beam, and shallow lookahead search, and pair it with PPO and GRPO, achieving up to 15% performance gains and more stable learning on Sokoban, FrozenLake, and WebShop tasks at a one-time increase in training compute. By moving search from inference time to the rollout stage of training, TSR provides a simple and general mechanism for stronger multi-turn agent learning, complementary to existing frameworks and rejection-sampling-style selection methods.

연구 동기 및 목표

• 희소하거나 지연된 각 턴 보상으로 인한 다중 턴 RL의 취약성과 불안정성을 동기부여하고 해결합니다.
• 최적화 목표를 바꾸지 않으면서 트레이닝 시점의 롤아웃 생성기를 제안하여 궤적 품질을 향상시킵니다.
• 가벼운 턴당 트리탐색이 옵티마이저에 독립적이고 PPO/GRPO와 호환될 수 있음을 보여줍니다.
• 고정된 학습 시간 계산 증가 하에서 다양한 환경에서 성능 및 안정성 향상을 입증합니다.

제안 방법

• 다중 턴 RL을 POMDP로 정식화하고 롤아웃 생성 품질에 초점을 맞춥니다.
• 트리 탐색 기반 롤아웃 생성기인 TSR을 도입하여 턴당 후보 행동(M)을 확장하고 작업 특화 함수 S로 점수를 매깁니다.
• Best-of-N, 빔 탐색, 얕은 선헤드(lookahead) 전략으로 TSR을 구체화하여 정책 업데이트를 위한 고품질 궤적을 선택합니다.
• 태스크 다양성을 유지하고 모드 붕괴를 완화하기 위해 TSR과 인스턴스 수준 필터링을 결합합니다.
• 고정된 일회성 학습 계산 증가를 사용하여 Sokoban, FrozenLake, WebShop에서 PPO/GRPO로 TSR을 평가합니다.
• 희소/지연 보상에 대한 프록시 점수를 사용하여 도전적인 환경에서 롤아웃 점수를 안내합니다.

실험 결과

연구 질문

• RQ1 모든 연구 목표를 바꾸지 않고 per-turn 궤적 검색에 대한 보통의 학습 시간 계산 예산을 추가하면 다중 턴 RL 성능이 향상될 수 있는가?
• RQ2 다양한 TSR 검색 전략(Best-of-N, 빔, Lookahead)이 성능, 안정성, 추론 효율성 측면에서 어떻게 비교되는가?
• RQ3 TSR을 인스턴스 수준 필터링과 결합하면 naive 롤아웃 샘플링보다 다양성과 학습 신호가 더 잘 확보되는가?
• RQ4 TSR로 학습된 소형 모델이 다중 턴 과제에서 더 큰 일반 모델과 맞먹거나 능가할 수 있는가?
• RQ5 Echo Trap 존재하에서 TSR이 학습 안정성과 그래디언트 동작에 미치는 영향은 무엇인가?

주요 결과

• TSR은 평가된 모든 작업 및 모델 크기에서 인스턴스 필터링 기준선보다 일관되게 우수한 성능을 보였습니다.
• 빌드-탐색(BEAM) 전략이 일반적으로 TSR 변형들 중에서 가장 강력한 성능 향상과 더 빠른 수렴을 제공합니다.
• TSR은 0.5B 모델에서 WebShop에 대해 최대 15%의 절대 성능 향상을 달성합니다.
• TSR로 학습하면 평균 응답 길이와 상호 작용 턴 수를 줄여 추론 효율이 향상됩니다.
• 그래디언트 노름은 TSR 하에서 안정적으로 유지되어 Echo Trap 위험이 감소합니다.
• 0.5B TSR 학습 에이전트가 Sokoban에서 더 큰 모델인 GPT-4o를 능가하고 Sokoban과 FrozenLake에서 Qwen-72B를 능가합니다.
• 검색 예산을 확대하면 수익이 감소하는 수익 체증이 나타나므로 메인 실험에서는 예산을 보통으로 설정하는 것이 타당합니다(B=2, M=4 등).