TSR: Trajectory-Search Rollouts for Multi-Turn RL of LLM Agents

TL;DR

TSR introduces a tree-style search method during training to generate higher-quality multi-turn trajectories for LLM agents, improving stability and performance in reinforcement learning tasks.

Contribution

It proposes TSR, a novel training-time search approach that enhances trajectory quality and learning stability for multi-turn RL of LLM agents, compatible with standard policy optimizers.

Findings

Achieved up to 15% performance improvements on benchmark tasks.

Enhanced training stability and trajectory quality.

Compatible with multiple search strategies and policy gradient methods.

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 state-based feedback. This improves rollout quality and stabilizes learning while remaining compatible with standard policy gradient optimizers, 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 modest, one-time increase in training compute. By moving search from inference time to the rollout stage of training, TSR provides a modular and general mechanism for stronger multi-turn agent learning, complementary to existing frameworks and rejection-sampling-style selection methods.