[論文レビュー] TSR: Trajectory-Search Rollouts for Multi-Turn RL of LLM Agents

Aladin Djuhera, Swanand Kadhe+3 more|arXiv (Cornell University)|Feb 12, 2026

Reinforcement Learning in Robotics被引用数 0

ひとこと要約

TSR は訓練時の軌跡探索を導入し、LLMエージェントの多ターン強化学習における各ターンのロールアウト品質を向上させ、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.

研究の動機と目的

• sparse or delayed per-turn rewards. への対処と多ターン RL の脆弱性を動機づけ、解決する。
• 最適化目的を変更せず、軌跡品質を改善する訓練時ロールアウト生成器を提案する。
• 軽量な各ターンの木探索がオプティマイザに依存しないこと、PPO/GRPO と互換であることを示す。
• 固定の訓練時計算増加の下で、多様な環境における性能と安定性の向上を実証する。

提案手法

• 多ターン RL を POMDP として形式化し、ロールアウト生成品質に焦点を当てる。
• TSR（木探索ベースのロールアウト生成器）を導入し、各ターンで候補行動を展開（M）し、タスク固有の関数 S で評価する。
• Best-of-N、ビーム探索、浅い先読み戦略を用いて TSR を具体化し、方針更新のための高品質軌跡を選択する。
• TSR をインスタンスレベルのフィルタリングと組み合わせて、タスク多様性を維持しモード崩壊を抑制する。
• Sokoban、FrozenLake、WebShop に対して、固定の一回限りの訓練計算増加を用いて TSR を PPO/GRPO で評価する。
• 疎な/遅延報酬の代理スコアを用いて、難しい環境におけるロールアウトスコアリングを誘導する。

実験結果

リサーチクエスチョン

• RQ1 各ターンの軌跡探索に対して modest な訓練時計算予算を追加するだけで、目的を変更せずに多-turn RL の性能を改善できるか？
• RQ2 TSR の異なる探索戦略（Best-of-N、ビーム、先読み）は、性能、安定性、推論効率の観点でどう比較されるか？
• RQ3 TSR とインスタンスレベルのフィルタリングを組み合わせると、単純なロールアウトサンプリングより多様性と学習信号が改善されるか？
• RQ4 TSR で訓練された小型モデルは、多ターンタスクでより大きな汎用モデルと同等かそれを上回ることができるか？
• RQ5 Echo Trap の存在下で、 TSR が訓練の安定性と勾配挙動に与える影響はどうなるか？

主な発見

• TSR は評価されたすべてのタスクとモデルサイズで、インスタンスフィルタリングベースラインを一貫して上回る。
• ビーム探索は通常、 TSR の変種の中で最も強い性能向上とより速い収束をもたらす。
• 0.5B モデルで WebShop に対して絶対的に最大15%の性能向上を達成。
• TSR を用いた訓練は、平均応答長と相互作用ターンを削減することにより推論効率を改善。
• 勾配ノルムは TSR 下で安定しており、Echo Trap のリスクを低減。
• 0.5B の TSR 訓練エージェントは Sokoban で GPT-4o のような大規模モデルを凌駕し、Sokoban および FrozenLake では Qwen-72B を超える。
• 探索予算を拡張すると収益が逓減するため、メイン実験には適度な予算（例: B=2, M=4）を正当化する。