
在 Amazon SageMaker AI 上用多轮强化学习训练搜索智能体
一篇 AWS 文章介绍了一种搜索智能体,它会在多轮交互中连续作出选择。智能体自主决定要搜索什么、采用哪一种检索策略,以及何时停止。因此,文中所述训练关注的是完整的决策序列,而不是孤立的一次回答。每个选择都属于一条最终产生检索文档的轨迹,文章据此把智能体训练描述为对整个搜索过程的训练。
文章介绍了 Amazon SageMaker AI MTRL,即多轮强化学习训练方式。按照文章的描述,MTRL 使用多轮 rollout 覆盖智能体的完整决策序列,并采用策略梯度算法。它支持自定义奖励、工具循环和多轮对话结构。所述服务还包括无服务器执行,以及对 rollout 和轨迹的异步收集。文章列出的算法选项包括 PPO、CISPO 和基于重要性采样的损失函数。列出的其他功能还包括可恢复训练、轨迹可观测性和评估任务。
把搜索看作一条决策轨迹
在文章描述的设置中,智能体可以选择执行下一次搜索,也可以决定停止。示例根据查询类型采用不同的检索方法:对于精确词语或标识符,使用 BM25 词法搜索;对于语义或概念性查询,使用向量搜索。对于多步骤智能体来说,这一区分与它的选择有关,因为它既要选择检索策略,也要判断是否需要继续下一步。文章将这些选择视为一个连续过程,而不是彼此独立的搜索动作。
文中训练使用的奖励是轨迹级 nDCG@10,依据最终检索到的文档与真实标注结果的匹配程度计算。这意味着所述奖励反映的是搜索轨迹结束时文档的质量,并将训练目标与该设置下的最终检索质量联系起来。这并不表示每个步骤得到相同的评价,也不表示 BM25 或向量搜索适用于所有查询;文章明确将它们分别用于精确词语或标识符,以及语义或概念性查询。
训练限制与报告结果
文章称,达到最大轮数或每轮采样 token 上限时,奖励为 -1。它还列出了这次训练中调整过的超参数:max_epochs 为 1,global_batch_size 为 128,rollout_max_concurrency 为 32。这些数字描述的是文章所述的配置,并不能说明相同的限制或设置适用于所有智能体和任务。
在四个留出基准上,文章报告 nDCG@10 有三个基准得到提升。BrowseComp-Plus 提升了 23.7%,WixQA 提升了 18.4%。Wands 的提升幅度较小,而 FreshStack 略有回退。对于 BrowseComp-Plus,文章还报告失败率从 22.89% 降至 0.68%。这些结果在不同基准之间并不相同:报告既包含提升,也包含回退。它们是所述评估中的结果,并不能证明所有搜索任务都会呈现同样的变化。
文章建议的实践流程
文章建议先按要求的格式准备数据集,根据任务特定指标定义奖励函数,从默认配置开始,然后进行评估并迭代。这些步骤将训练配置与团队希望衡量的结果联系起来。如果任务依据检索文档的质量来评价,奖励就应体现相关的任务指标。随后,评估和迭代可以根据观察到的结果调整流程,例如留意 BrowseComp-Plus 和 WixQA 的提升,以及 FreshStack 的轻微回退。
这一实践框架将搜索视为一连串动作,并让训练和评估与这条序列相对应。按照文章的描述,MTRL 支持自定义奖励、工具交互和轨迹收集;基准结果则显示,测得的变化可能因评估集而异。文章没有声称每个智能体或每个数据集都会取得相同幅度的提升。它的建议集中在数据准备、与任务指标相连的奖励,以及反复评估。这些内容共同说明,所述多轮训练如何把智能体的搜索决策与最终检索文档的质量联系起来。