Fine-tune a search agent with multi-turn RL on Amazon SageMaker AI

在 Amazon SageMaker AI 上使用多轮强化学习微调搜索代理 | 人工智能

在 Amazon SageMaker AI 上使用多轮强化学习微调搜索代理

由大型语言模型（LLMs）驱动的搜索代理正在改变企业检索信息的方式。用户无需精心构造完美查询，搜索代理可自主决定搜索内容、采用的检索策略以及何时停止搜索。它通过多轮交互逐步优化策略，根据已检索到的信息不断调整方法。

然而，要让这种多步骤行为有效运行颇具挑战。没有基础模型能预先了解你的工具或环境。提示小型模型时，通常难以获得可靠的多轮交互行为；提示前沿模型时虽然通常有效，但需要付出延迟和成本的代价。微调提供了一种折中方案：你可以直接向小型模型传授工具和环境知识。最终实现小型模型的高效低成本，同时获得前沿模型级别的可靠性。

尽管微调是自然的下一步，但传统方法各有局限。监督微调（SFT）依赖专家演示的理想多轮交互轨迹，但收集这些数据成本高昂，且通常无法适用于你的场景。单轮强化学习（RL），如基于可验证奖励的强化学习（RLVR），每次仅评估单个响应。但搜索代理在多轮交互中做出相互依赖的决策，每一步都基于之前的上下文。单独优化某一步骤会完全忽略这些依赖关系。

你真正需要的是能针对完整多轮交互轨迹进行优化的训练方法。奖励信号只需反映最终结果是否理想。这正是多轮强化学习（MTRL）所提供的。它训练代理在完整步骤序列中做出优质决策，融入环境特定行为，并让你掌控输出质量。由于使用小型专用模型，还能实现更快速、更低成本的推理。

在本文中，我们将介绍如何使用 Amazon SageMaker AI 多轮强化学习（MTRL）微调搜索代理，并分享我们在检索质量和可靠性方面观察到的成果。首先我们将解释 Amazon SageMaker AI MTRL 是什么以及其工作原理。

什么是 Amazon SageMaker AI MTRL？

通过 Amazon SageMaker AI MTRL，你可以在多轮交互场景中使用强化学习对大型语言模型（LLMs）进行微调。它将智能体任务框架视为一系列决策，通过多轮交互生成训练数据，并使用策略梯度算法优化模型。

Amazon SageMaker AI MTRL 提供以下特性：

• 模块化智能体-环境接口：你可使用低代码方式实现集成。可自定义奖励函数、工具循环和多轮对话结构。

• 无服务器执行：无需配置或管理 GPU 集群，即可按 token 计费实现生产级智能体强化学习。

• 异步交互与轨迹收集：生成过程与梯度更新并行执行，通过有界离线策略陈旧度控制，确保训练速度且不偏离当前策略太远。

• 原生算法库：您可以选择近端策略优化（PPO）、截断重要性采样策略优化（CISPO）和重要性采样（IS）损失函数，配合基于分组的优势估计器（GRPO、GRPO pass@k、RLOO 等）。

• 可恢复训练：您可以将长时间训练任务拆分到多个作业中，突破单个作业的时间限制。

• 轨迹和奖励可观测性：您可以通过 Amazon SageMaker AI 管理的 MLflow 检查代理在每一步操作和训练过程中的行为轨迹。

• 评估任务：在将模型部署到 Amazon SageMaker AI 终端节点或 Amazon Bedrock 之前，您可以报告奖励值、pass@k 和轨迹指标。

这使得 MTRL 成为搜索代理的天然选择：您拥有明确的奖励信号（检索质量）、多轮交互循环（代理发出查询并接收结果）以及明确定义的环境（搜索工具）。

环境与配置

我们的配置依赖以下内容：

• 具有访问美国西部（俄勒冈）AWS 区域（us-west-2）中 Amazon SageMaker AI 权限的 AWS 账户。

• 已上传到 Amazon Simple Storage Service（Amazon S3）的训练和验证数据集，格式符合要求（详见 MTRL 文档）。

• 已部署的代理终端节点，为 MTRL 环境在 rollout 过程中提供 BM25 和向量搜索工具。

• 熟悉 Amazon SageMaker AI 和 Python。

解决方案概述

在本文中，我们使用 Amazon SageMaker AI MTRL 对 Qwen3.6-27B 模型（支持美国西部（俄勒冈）区域（us-west-2））进行微调，用于搜索代理。搜索代理是一个基于大语言模型（LLM）的系统，能够自主使用搜索工具查找、收集并综合信息以回答问题或完成任务。

我们聚焦于企业搜索场景，其中代理拥有以下两种工具：

• 词法搜索（BM25）：通过统计词频寻找精确关键词匹配。适用于包含特定术语或标识符的查询。

• 向量搜索：将查询和文档转换为嵌入向量并计算相似度。推荐用于语义或概念性查询。

我们还限制了交互轮数（一轮为用户与代理的一次交互），防止代理生成过长响应并鼓励高效搜索行为。

训练配置

本节将介绍训练配置的三个关键组成部分：数据集、奖励函数和 MTRL 作业配置。

数据集

我们使用以下数据集进行训练和测试：

| 数据集 | 训练/测试 | 公共链接 | 描述 | |-------------|-----------|------------------|----------------------------------------------------------------------| | FRAMES | 训练 | HuggingFace | 需要跨多篇维基百科文章综合推理的多跳事实问答数据集。 | | BRIGHT | 训练 | N/A | 跨 12 个领域的推理密集型检索任务，相关性需要推理而非关键词重叠。 | | 企业 RAG | 训练 | GitHub | 包含 50 万+ 个合成企业文档和 500 个问题的企业知识 RAG 系统基准测试。| | ESCI | 训练 | N/A | 多语言产品搜索数据集，标注查询-产品对为精确匹配、替代品、互补品或无关。| | Musique | 训练 | N/A | 通过将单跳问题组合成具有挑战性的多跳推理问题构建的问答数据集。 | | MLQA | 训练 | N/A | 覆盖七种语言的并行抽取式问答基准测试，用于评估跨语言理解能力。 | | FreshStack | 测试 | N/A | 最近的开发者技术问答数据集，涵盖 Stack Overflow 和五个主题的文档。 | | WixQA | 测试 | N/A | N/A（未提供详细信息） |

奖励函数

我们的主要评估指标是nDCG@10（Normalized Discounted Cumulative Gain at rank 10）。nDCG@10是信息检索领域的标准评估指标，用于衡量前10个检索结果与理想排序的匹配程度。该指标对将高度相关文档排在前列的系统给予奖励，得分1.0表示完美排序，0.0表示未检索到任何相关文档。我们在MTRL中直接使用nDCG@10作为奖励函数。这是一个轨迹级奖励，要求智能体完成完整的多轮搜索，奖励值反映最终检索结果与真实标签的匹配程度。

当智能体达到最大轮数限制或单轮最大采样token数时，我们会分配-1的奖励值，以明确教导模型避免这些失败模式。这种基于惩罚的奖励设计能有效引导模型远离不良行为，无需手动设计复杂的中间奖励。

MTRL任务配置

Amazon SageMaker AI MTRL的一大优势在于其配置简便性。我们仅修改了三个超参数，其余参数均使用默认值。以下代码片段展示了如何使用MultiTurnRLTrainer SDK配置并启动训练任务：

from sagemaker.modules.train import MultiTurnRLTrainer

trainer = MultiTurnRLTrainer(
    model_id="qwen3.6-27b",
    agent_endpoint="<your-agent-endpoint>",
    training_dataset_s3_uri="s3://amzn-s3-demo-bucket/datasets/train/",
    validation_dataset_s3_uri="s3://amzn-s3-demo-bucket/datasets/validation/",
    hyperparameters={
        "max_epochs": 1,
        "global_batch_size": 128,
        "rollout_max_concurrency": 32,
    },
    output_s3_uri="s3://amzn-s3-demo-bucket/output/",
)

trainer.train()

我们修改的关键超参数包括：

• max_epochs: 1 控制对训练数据的完整遍历次数
• global_batch_size: 128 设置每个训练步骤的提示数量
• rollout_max_concurrency: 32 控制轨迹收集期间并行执行的rollout数量

这就是全部配置。通常需要RL专业知识的决策（如算法选择、优势估计器、离线策略陈旧度限制等）都使用默认设置。从基础模型到下一章节的实验结果，仅需上述配置即可实现。

MTRL微调在四个保留基准测试中的三个上提升了搜索代理的表现，并在所有基准测试中显著提高了其可靠性。最大的提升出现在BrowseComp-Plus（nDCG@10提升23.7%）和WixQA（提升18.4%），在Wands上提升较小，在FreshStack上略有下降。可靠性提升是更显著的结果。在BrowseComp-Plus上，失败率从22.89%降至0.68%。这意味着代理不仅学会了更好地搜索，还学会了在单次交互和令牌预算内完成任务。接下来的章节将展示训练曲线和各基准测试的具体数据。

训练进展

下图显示了训练和验证实例在整个训练过程中的奖励值（nDCG@10）。x轴表示训练步数，y轴显示nDCG@10得分。两条曲线分别跟踪训练和验证奖励值，均稳步上升后趋于平稳。MTRL训练可能持续多天，MTRL服务默认时间限制为24小时，可通过CreateJob JSON schema进行调整。如果任务因超时或基础设施错误被停止或失败，可通过恢复功能从之前的检查点继续训练。

图1：训练和验证集上nDCG@10奖励值随训练过程的变化

该图将nDCG@10与训练步数对应，展示了训练和验证集的表现。随着微调的进行，两条曲线均稳步上升并在训练结束时趋于饱和，表明进一步训练不会带来显著提升。

测试性能

下表比较了微调模型与原始Qwen3.6-27B在四个保留测试数据集上的表现。表中所有指标均来自前文所述数据集的评估运行结果。

列定义：

• Size：数据集中的问题数量。
• nDCG@10：标准归一化折扣累积增益。失败任务的nDCG@10设为0。
• Failure rate：代理遇到错误（例如超出回合限制或令牌预算）的问题占比。
• Turns (avg)：每个问题的平均回合数。

benchmark | size | model | ndcg@10 | fail_rate | turns ---|---|---|---|---|--- 400 | Qwen3.6-27B | 0.5725 | 0.67% | 4.3 Qwen3.6-27B-finetune | 0.6781 | 0.17% | 4.5 147 | 0.5762 | 0.00% | 2.2 0.6112 | 2.9 672 | 0.4112 | 0.20% | 3.1 0.4089 | 0.05% | 2.8 830 | 0.5136 | 22.89% | 7.0 0.6354 | 0.68% | 6.3

结果表明强化学习微调产生了显著影响：

• 更好的搜索质量：如上表所示，nDCG@10在WixQA上提升18.4%（从0.5725到0.6781），在Wands上提升6%（从0.5762到0.6112），在BrowseComp-Plus上提升23.7%（从0.5136到0.6354）。FreshStack上得分略有下降。
• 更少的失败：为失败案例设置-1奖励非常有效。在BrowseComp-Plus上，失败率从22.89%降至0.68%。

清理

如果您尝试这种方案，请记得后续清理资源以避免持续收费：

• 如果MTRL训练作业仍在运行，请通过Amazon SageMaker AI控制台停止或删除该作业。
• 如果不再需要Amazon S3中的模型文件，请删除这些文件。
• 删除为评估创建的任何已部署端点。

结论

在本文中，我们分享了使用 Amazon SageMaker AI MTRL 通过强化学习对基于大语言模型的搜索代理进行微调的实践经验。通过简单的配置，微调后的模型在检索质量、失败率和回合效率（每问题更少的交互回合）方面均取得了显著提升。

我们使用 Amazon SageMaker AI MTRL 的关键经验总结：

• 低准入门槛：默认配置即可良好运行，无需深度强化学习专业知识即可上手。

• 无服务器架构：无需预置 GPU 或管理分布式训练集群。按 token 计费。

• 可恢复训练：长时间训练任务可拆分到多个作业中执行。

• 直接优化：直接针对实际任务指标进行训练，而非使用代理损失函数。

• 完整可观测性：可通过 Amazon SageMaker AI 管理的 MLflow 逐回合检查轨迹，理解代理的学习过程。

下一步

如果你想将此方法应用到自己的代理系统中，以下是一些有用的起点：

• 查阅 Amazon SageMaker AI MTRL 文档。

• 按要求格式准备数据集。

• 准备你的代理并根据任务特定指标定义奖励函数。

• 从默认配置开始，我们在实践中发现其效果良好。

• 进行评估并持续迭代。

作者信息

'\"`