[论文解读] TSR: Trajectory-Search Rollouts for Multi-Turn RL of LLM Agents
TSR 引入训练时轨迹搜索,以为多-turn RL 的大语言模型代理产生更高质量的逐回合 rollout,在 Sokoban、FrozenLake 和 WebShop 的性能与稳定性方面实现与优化器无关的集成。
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 在稀疏或延迟的逐回合奖励下的脆弱性提供动机与解决方案。
- 提出一种训练时 rollout 生成器,在不改变优化目标的前提下提升轨迹质量。
- 显示轻量级逐回合树搜索可以与优化器无关且兼容 PPO/GRPO。
- 在固定的训练时计算增加下,展示在多样化环境中的性能与稳定性提升。
提出的方法
- 将多轮 RL 正规化为部分可观测马尔可夫决策过程(POMDP),并聚焦于 rollout 生成质量。
- 提出 TSR,一种基于树搜索的 rollout 生成器,在每一回合扩展候选动作(M),并用任务特定函数 S 对它们进行评分。
- 以 Best-of-N、束搜索(beam search)和浅层前瞻(shallow lookahead)策略实例化 TSR,以为策略更新选择高质量轨迹。
- 与实例级过滤相结合,保持任务多样性并缓解模式塌陷。
- 在 Sokoban、FrozenLake、WebShop 上使用固定的一次性训练计算开销评估 TSR 与 PPO/GRPO 的表现。
- 使用稀疏/延迟奖励的代理分数来引导在挑战性环境中的 rollout 评分。
实验结果
研究问题
- RQ1在训练时为逐回合轨迹搜索增加适度的计算预算,是否能在不改变目标函数的前提下提升多轮 RL 的性能?
- RQ2不同的 TSR 搜索策略(Best-of-N、束搜索、前瞻)在性能、稳定性和推理效率方面的比较如何?
- RQ3将 TSR 与实例级过滤结合,是否比简单的 rollout 采样获得更好的多样性与学习信号?
- RQ4 TSR 训练的小模型是否能够匹配或超越在多轮任务上的更大型通用模型?
- RQ5在存在回声陷阱(Echo Trap)的情况下, TSR 对训练稳定性和梯度行为有何影响?
主要发现
- TSR 在所有评估任务和模型规模上始终优于实例过滤基线。
- 束搜索通常在 TSR 变体中提供最强的性能提升和更快的收敛。
- 在 0.5B 模型上,TSR 在 WebShop 上实现了最高达 15% 的绝对性能提升。
- 使用 TSR 训练可以通过缩短平均响应长度和互动回合数提升推理效率。
- 在 TSR 下梯度范数保持稳定,表明回声陷阱风险降低。
- 一个经过 TSR 训练的 0.5B 代理可以在 Sokoban 上超越如 GPT-4o 等更大模型,在 Sokoban 和 FrozenLake 上也优于 Qwen-72B。
- 对搜索预算的扩展显示边际效应递减,支持在主要实验中采用中等预算(如 B=2, M=4)。
更好的研究,从现在开始
从论文设计到论文写作,大幅缩短您的研究时间。
无需绑定信用卡
本解读由 AI 生成,并经人工编辑审核。