在 Amazon SageMaker AI 上通过多轮强化学习微调搜索 Agent

由大语言模型（LLM）驱动的搜索 Agent 正在改变企业检索信息的方式。与要求用户构造完美查询不同，搜索 Agent 会自主决定搜索什么、采用何种检索策略，以及何时停止搜索。它会跨越多轮交互执行这些操作，并根据已经检索到的内容不断调整方法。

然而，要让这种多步骤行为稳定有效地运行并不容易。没有任何基础模型天生就了解你的工具或环境。向小模型提供提示词，你很难获得可靠的多轮行为；向前沿模型提供提示词，它通常能够工作，但你需要为这种能力付出延迟和成本。微调提供了第三条路径：直接教会小模型使用你的工具并适应你的环境。最终，你可以同时获得小模型的速度和成本优势，以及原本需要前沿模型才能达到的可靠性。

尽管微调是顺理成章的下一步，传统方法各自都存在不足。监督微调（SFT）依赖专家提供理想多轮轨迹的示范，而这些示范收集成本高昂，并且通常并不存在于你的具体场景中。单轮强化学习（RL），例如基于可验证奖励的强化学习（RLVR），一次只为一个响应评分。但搜索 Agent 会在多个回合中作出相互依赖的决策，每一步都建立在此前上下文之上。孤立地优化单一步骤，完全无法捕捉这些依赖关系。

你需要的是一种能够针对完整多轮轨迹优化 Agent 的训练方法。奖励信号只需反映最终结果是否良好即可。这正是多轮强化学习（MTRL）所提供的能力。它训练 Agent 在完整步骤序列中作出良好决策，固化与你环境相关的行为，并让你能够控制输出质量。由于运行的是更小、更专用的模型，你也能获得更快、更低成本的推理。

本文介绍了如何使用 Amazon SageMaker AI 多轮强化学习（MTRL） 微调一个搜索 Agent，并分享我们在检索质量和可靠性方面观察到的结果。首先，我们说明 Amazon SageMaker AI MTRL 是什么，以及它如何工作。

什么是 Amazon SageMaker AI MTRL？

借助 Amazon SageMaker AI MTRL，你可以在多轮交互场景中使用强化学习微调 LLM。它将 Agent 任务定义为一连串决策，使用多轮 rollout 生成训练数据，并利用策略梯度算法优化模型。

Amazon SageMaker AI MTRL 提供：

• 模块化 Agent-环境接口：你可以保持低代码集成。你可以定义自定义奖励、自定义工具循环和多轮对话形态。
• **无服务器执行：**你可以按 token 付费获得生产规模的 Agent 强化学习能力，无需预置或管理 GPU 集群。
• 异步 rollout 和轨迹收集：你可以并行执行生成和梯度更新，并通过有界的离策略陈旧性控制，使训练保持快速且不会偏离当前策略过远。
• **原生算法库：**你可以从近端策略优化（PPO）、裁剪重要性采样策略优化（CISPO）和重要性采样（IS）损失中选择，并搭配基于组的优势估计器（GRPO、GRPO pass@k、RLOO 等）。
• 可恢复训练：你可以将长时间训练拆分到多个任务中，从而突破单个任务的时间限制。
• 轨迹和奖励可观测性：你可以在由 Amazon SageMaker AI 管理的 MLflow 中，逐回合、跨训练步骤检查 Agent 的行为。
• 评估任务：在部署到 Amazon SageMaker AI endpoint 或 Amazon Bedrock 之前，你可以报告奖励、pass@k 和轨迹指标。

这使 MTRL 非常适合搜索 Agent：你拥有清晰的奖励信号（检索质量）、多轮交互循环（Agent 发出查询并接收结果），以及定义明确的环境（搜索工具）。

环境与设置

我们的设置依赖以下条件：

• 拥有可在美国西部（俄勒冈）AWS 区域（us-west-2）访问 Amazon SageMaker AI 的 AWS 账户。
• 已将训练和验证数据集按要求格式上传至 Amazon Simple Storage Service（Amazon S3）（详情参见 MTRL 文档）。
• 已部署一个 Agent endpoint，为 MTRL 环境在 rollout 期间调用提供 BM25 和向量搜索工具。
• 熟悉 Amazon SageMaker AI 和 Python。

方案概述

本文使用 Amazon SageMaker AI MTRL，为搜索 Agent 微调 Qwen3.6-27B 模型（该模型在美国西部［俄勒冈］区域 us-west-2 受支持）。搜索 Agent 是一种由 LLM 驱动的系统，能够自主使用搜索工具查找、收集和综合信息，以回答问题或完成任务。

我们聚焦于企业搜索场景，其中 Agent 可使用两种工具：

• **词法搜索（BM25）：**通过统计词频查找精确关键词匹配。最适合包含特定术语或标识符的查询。
• **向量搜索：**将查询和文档转换为嵌入向量并计算相似度。建议用于语义型或概念型查询。

我们还限制了回合数（一个回合是一次用户与助手交互），以防止 Agent 生成过长的响应，并鼓励高效的搜索行为。

训练设置

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

数据集

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

数据集	训练/测试	公开链接	说明
FRAMES	训练	HuggingFace	需要综合多篇 Wikipedia 文章信息的多跳事实问答。
BRIGHT	训练	HuggingFace	覆盖 12 个领域、需要大量推理的检索任务；相关性需要推理判断，而非关键词重合。
Enterprise RAG	训练	GitHub	一个基准数据集，包含超过 50 万篇合成企业文档和 500 个问题，用于在企业内部知识上训练和评估 RAG 系统。
ESCI	训练	GitHub	一个多语言产品搜索数据集，将查询-产品对标记为 Exact、Substitute、Complement 或 Irrelevant。
Musique	训练	GitHub	通过组合单跳问题构建而成的问答数据集，形成具有挑战性的多跳推理问题。
MLQA	训练	HuggingFace	一个平行抽取式问答基准，覆盖七种语言，用于评估跨语言理解能力。
FreshStack	测试	HuggingFace	覆盖五个主题的近期开发者技术问答数据，来源为 Stack Overflow 和文档。
WixQA	测试	HuggingFace	基于 Wix 知识库的帮助中心支持问答。
BrowseComp-Plus	测试	GitHub	针对约 10 万份经人工验证的开放网络文档的高难度深度研究查询。
Wands	测试	HuggingFace	一个 Wayfair 产品搜索数据集，包含超过 42,000 条查询-产品相关性判断，标记为 Exact、Partial 或 Irrelevant。

我们根据文档，将数据集预处理为 MTRL 服务能够识别的格式。我们还从每个数据集的训练实例中额外预留 5% 作为验证实例。

奖励函数

我们的主要指标是 nDCG@10（排名第 10 位的归一化折损累计增益）。nDCG@10 是一种标准信息检索指标，用于衡量检索结果中排名前 10 的文档与理想排序的匹配程度。它会奖励将高度相关文档排在列表靠前位置的系统；得分为 1.0 表示完美排序，0.0 则表示没有检索到相关文档。我们直接将 nDCG@10 用作 MTRL 的奖励函数。这是一种轨迹级奖励：Agent 完成完整的多轮搜索后，奖励反映最终检索文档与真实标注的匹配程度。

当 Agent 达到最大回合数，或在单个回合中达到最大采样 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 设置每个训练步骤中的 prompt 数量。
• rollout_max_concurrency: 32 控制在轨迹收集期间并行运行的 rollout 数量。

这就是全部设置。通常需要强化学习专业知识才能确定的选择——算法、优势估计器、离策略陈旧性边界——均使用默认设置。从基础模型到下一节的结果，只需要前述配置。

结果

MTRL 微调在四个留出基准中的三个上提升了搜索 Agent 的表现，并且在所有基准上显著提高了可靠性。最大的提升来自 BrowseComp-Plus（nDCG@10 提升 23.7%）和 WixQA（提升 18.4%）；Wands 上有较小提升，FreshStack 上则略有回退。可靠性提升是更显著的结果。在 BrowseComp-Plus 上，失败率从 22.89% 降至 0.68%。这意味着 Agent 不仅学会了更好地搜索，也学会了在回合和 token 预算内完成任务。接下来的部分将介绍训练曲线和各基准的具体数据。

训练进度

下图展示了训练过程中训练实例和验证实例上的奖励（nDCG@10）。x 轴表示训练步骤，y 轴表示 nDCG@10 得分。两条线分别跟踪训练奖励和验证奖励，二者均稳步上升后趋于平台期。MTRL 训练可能持续多天，而 MTRL 服务的默认时间限制为 24 小时，可通过 CreateJob JSON schema 进行调整。如果任务因超时或基础设施错误而停止或失败，你可以借助恢复支持，从先前 checkpoint 继续训练。

图 1：训练期间训练集和验证集上的 nDCG@10 奖励

该图展示了训练集和验证集的 nDCG@10 与训练步骤之间的关系。随着微调推进，两条曲线均稳步上升，并在训练结束时达到饱和，表明进一步训练可能不会带来收益。

测试表现

下表将微调后的模型与原始 Qwen3.6-27B 在四个留出测试数据集上的表现进行比较。表中所有指标均来自我们在上一节所列数据集上执行的评估。

列定义：

• Size 是数据集中的问题数量。
• nDCG@10 是标准的归一化折损累计增益。失败任务的 nDCG@10 记为 0。
• Failure rate 是 Agent 发生错误的问题占比，例如超出回合上限或 token 预算。
• Turns (avg) 是每个问题的平均回合数。

基准	规模	模型	nDCG@10	失败率	回合数
WixQA	400	Qwen3.6-27B	0.5725	0.67%	4.3
WixQA	400	Qwen3.6-27B-finetune	0.6781	0.17%	4.5
Wands	147	Qwen3.6-27B	0.5762	0.00%	2.2
Wands	147	Qwen3.6-27B-finetune	0.6112	0.00%	2.9
FreshStack	672	Qwen3.6-27B	0.4112	0.20%	3.1
FreshStack	672	Qwen3.6-27B-finetune	0.4089	0.05%	2.8
BrowseComp-Plus	830	Qwen3.6-27B	0.5136	22.89%	7.0
BrowseComp-Plus	830	Qwen3.6-27B-finetune	0.6354	0.68%	6.3

结果表明，RL 微调带来了有意义的差异：

• **更好的搜索

Search agents powered by large language models (LLMs) are transforming how enterprises retrieve information. Rather than requiring users to craft the perfect query, a search agent autonomously decides what to search for, which retrieval strategy to use, and when to stop searching. It does this across multiple rounds of interaction, refining its approach based on what it has already retrieved.

However, getting this multi-step behavior to work well is hard. No base model arrives knowing your tools or your environment. Prompt a small model and you rarely get dependable multi-turn behavior. Prompt a frontier model and it often works, but you pay for that capability in latency and cost. Fine-tuning offers a third path: you teach a small model your tools and environment directly. The result is a small model’s speed and cost with the reliability that would otherwise require a frontier model.

Even though fine-tuning is the natural next step, the traditional approaches each fall short. Supervised fine-tuning (SFT) depends on expert demonstrations of ideal multi-turn trajectories, which are costly to collect and usually don’t exist for your setup. Single-turn reinforcement learning (RL), such as RL with verifiable rewards (RLVR), scores one response at a time. But a search agent makes interdependent decisions across many turns, each one building on the context before it. Optimizing a single step in isolation misses those dependencies entirely.

What you need is a training approach that optimizes the agent across the full multi-turn trajectory. The reward signal only needs to reflect whether the final outcome was good. That is exactly what multi-turn reinforcement learning (MTRL) provides. It trains the agent to make good decisions across a full sequence of steps, bakes in your environment-specific behavior, and gives you control over output quality. Because you run a smaller, specialized model, you also get faster, cheaper inference.

In this post, we describe how we fine-tuned a search agent using Amazon SageMaker AI multi-turn reinforcement learning (MTRL) and share the results we observed in retrieval quality and reliability. We begin by explaining what Amazon SageMaker AI MTRL is and how it works.

What is Amazon SageMaker AI MTRL?

With Amazon SageMaker AI MTRL, you can fine-tune LLMs using reinforcement learning in multi-turn interaction settings. It frames an agentic task as a sequence of decisions, uses multi-turn rollouts to generate training data, and optimizes the model with policy gradient algorithms.

Amazon SageMaker AI MTRL offers:

• Modular agent-environment interface : You keep integration low-code. You define custom rewards, custom tool loops, and multi-turn conversation shapes.
• Serverless execution: You get production-scale agentic RL at per-token pricing without provisioning or managing GPU clusters.
• Asynchronous rollout and trajectory collection : You run generation and gradient updates in parallel with bounded off-policy staleness, so training stays fast without drifting too far from the current policy.
• A native algorithm library: You choose from Proximal Policy Optimization (PPO), Clipped Importance Sampling Policy Optimization (CISPO), and importance-sampling (IS) losses, paired with group-based advantage estimators (GRPO, GRPO pass@k, RLOO, and more).
• Resumable training : You can split long training runs across multiple jobs to work beyond single-job time limits.
• Trajectory and reward observability : You can inspect what your agent did turn by turn and across training steps in MLflow managed by Amazon SageMaker AI.
• Evaluation jobs : You can report reward, pass@k, and trajectory metrics before deploying to an Amazon SageMaker AI endpoint or Amazon Bedrock.

This makes MTRL a natural fit for search agents: you have a clear reward signal (retrieval quality), a multi-turn interaction loop (the agent issuing queries and receiving results), and a well-defined environment (the search tools).

Environment and setup

Our setup relied on the following:

• An AWS account with access to Amazon SageMaker AI in the US West (Oregon) AWS Region (us-west-2).
• Training and validation datasets uploaded to Amazon Simple Storage Service (Amazon S3) in the required format (see the MTRL documentation for details).
• A deployed agent endpoint that exposes BM25 and vector search tools for the MTRL environment to call during rollouts.
• Familiarity with Amazon SageMaker AI and Python.

Solution overview

In this post, we use Amazon SageMaker AI MTRL to fine-tune a Qwen3.6-27B model (supported in the US West (Oregon) Region (us-west-2)) for a search agent. The search agent is an LLM-powered system that autonomously uses search tools to find, gather, and synthesize information to answer a question or complete a task.

We focus on an enterprise search setting where the agent has two tools available:

• Lexical search (BM25): Finds exact keyword matches by counting word frequencies. Best for queries with specific terms or identifiers.
• Vector search: Converts queries and documents into embedding vectors and computes similarities. This is recommended for semantic or conceptual queries.

We also limit the number of turns (one turn is one round of user-assistant interaction), preventing the agent from generating excessively long responses and encouraging efficient search behavior.

Training setup

This section walks through the three key components of our training setup: the datasets, the reward function, and the MTRL job configuration.

Datasets

We use the following datasets for training and testing:

Dataset	Train/Test	Public Link	Description
FRAMES	Train	HuggingFace	Multi-hop factoid QA requiring synthesis across multiple Wikipedia articles.
BRIGHT	Train	HuggingFace	Reasoning-intensive retrieval across 12 domains, where relevance needs reasoning, not keyword overlap.
Enterprise RAG	Train	GitHub	A benchmark of over 500,000 synthetic enterprise documents and 500 questions for training and evaluating RAG systems on internal company knowledge
ESCI	Train	GitHub	A multilingual product-search dataset labeling query–product pairs as Exact, Substitute, Complement, or Irrelevant
Musique	Train	GitHub	A question-answering dataset built by composing single-hop questions into challenging multi-hop reasoning problems.
MLQA	Train	HuggingFace	A parallel extractive question-answering benchmark covering seven languages for evaluating cross-lingual comprehension.
FreshStack	Test	HuggingFace	Recent developer-tech Q&A over Stack Overflow + docs for five topics.
WixQA	Test	HuggingFace	Help-center support QA over the Wix knowledge base.
BrowseComp-Plus	Test	GitHub	Hard deep-research queries over ~100K human-verified open-web documents.
Wands	Test	HuggingFace	A Wayfair product-search dataset containing over 42,000 query–product relevance judgments labeled as Exact, Partial, or Irrelevant.

We preprocess the datasets into the format recognized by the MTRL service according to the documentation. We further reserve 5 percent of the training instances within each dataset as validation instances.

Reward function

Our main metric is nDCG@10 (Normalized Discounted Cumulative Gain at rank 10). nDCG@10 is a standard information retrieval metric that measures how well the top 10 retrieved documents match the ideal ranking. It rewards systems that place highly relevant documents near the top of the list, with a score of 1.0 meaning perfect ranking and 0.0 meaning no relevant documents retrieved. We use nDCG@10 directly as the reward function in MTRL. This is a trajectory-level reward where the agent completes its full multi-turn search, and the reward reflects how well the final retrieved documents match the ground truth.

When the agent reaches the maximum number of turns or the maximum sampling tokens in a single turn, we assign a reward of -1 to explicitly teach the model to avoid these failure modes. This kind of penalty-based reward design is effective at steering the model away from undesirable behavior without needing to hand-craft complex intermediate rewards.

MTRL job configuration

One of the appeals of Amazon SageMaker AI MTRL is how little you need to configure to get going. We change three hyperparameters and leave everything else at the default values. The following code snippet shows how to configure and launch the training job using the 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()

The key hyperparameters we change are:

• max_epochs: 1 controls how many full passes over the training data.
• global_batch_size: 128 sets the number of prompts per training step.
• rollout_max_concurrency: 32 controls how many rollouts run in parallel during trajectory collection.

That’s the whole setup. The choices that usually demand RL expertise, the algorithm, the advantage estimator, the off-policy staleness bounds, all run on defaults. Getting from the base model to the results in the next section required nothing more than the preceding configuration.

Results

MTRL fine-tuning improved the search agent on three of four held-out benchmarks and made it substantially more reliable across all of them. The largest gains came on BrowseComp-Plus (+23.7 percent nDCG@10) and WixQA (+18.4 percent), with a smaller gain on Wands and a slight regression on FreshStack. The reliability improvement was the more striking result. On BrowseComp-Plus, the failure rate fell from 22.89 percent to 0.68 percent. This means the agent learned not only to search better but also to finish the task within its turn and token budget. The sections that follow walk through the training curves and the per-benchmark numbers.

Training progress

The following figure shows the reward (nDCG@10) on both training and validation instances over the course of training. The x-axis represents training steps and the y-axis shows the nDCG@10 score. Two lines track training and validation reward, both rising steadily before plateauing. MTRL training can span multiple days and the MTRL service has a default time limit of 24 hours, which is adjustable through the CreateJob JSON schema. If a job is stopped or has failed because of timeout or infrastructure error, you can continue training from a previous checkpoint with resume support.

Figure 1: nDCG@10 reward on the training and validation sets over the course of training

The figure plots nDCG@10 against training steps for both training and validation sets. Both curves increase steadily as fine-tuning progresses and saturate by the end of training, suggesting further training would not be beneficial.

Test performance

The following table compares the fine-tuned model against the original Qwen3.6-27B on four held-out test datasets. All metrics in the table are computed from our evaluation runs on the datasets listed in the preceding section.

Column definitions:

• Size is the number of questions in the dataset.
• nDCG@10 is the standard Normalized Discounted Cumulative Gain. The failed tasks are assigned nDCG@10 of 0.
• Failure rate is the percentage of questions where the agent hits an error (for example, exceeding the turn limit or token budget).
• Turns (avg) is the average number of turns per question.

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

The results show that RL fine-tuning made a meaningful difference:

• Better search quality: As shown in the preceding table, nDCG@10 improved on WixQA (from 0.5725 to 0.6781, a +18.4 percent gain), on Wands (from 0.5762 to 0.6112, a +6 percent), and on BrowseComp-Plus (from 0.5136 to 0.6354, a +23.7 percent gain). The score decreased slightly on FreshStack.
• Fewer failures: Assigning a reward of -1 for failure cases was highly effective. On BrowseComp-Plus, the failure rate dropped from 22.89 percent to only 0.68 percent.

Clean up

If you do explore this approach, remember to clean up the resources afterward to avoid ongoing charges:

• Stop or delete the MTRL training job in the Amazon SageMaker AI console if it’s still running.
• Delete the model artifacts in Amazon S3 if they’re no longer needed.
• Delete any deployed endpoints that were created for evaluation.

Conclusion

In this post, we described our experience using Amazon SageMaker AI MTRL to fine-tune an LLM-powered search agent with reinforcement learning. With minimal configuration, the fine-tuned model showed substantial improvements in retrieval quality, failure rate, and turn efficiency (fewer turns per question).

Key takeaways from our experience with Amazon SageMaker AI MTRL:

• Low barrier to entry: Default configurations work well, and no deep RL expertise is required to get started.
• Serverless infrastructure: No need to provision GPUs or manage distributed training clusters. You pay per-token.
• Resumable training: Long training runs can be split across multiple jobs.
• Direct optimization: You train against your actual task metric rather than proxy losses.
• Full observability: You can inspect trajectories turn by turn in MLflow managed by Amazon SageMaker AI to understand what the agent learned.

Next steps

If you want to apply this approach to your own agents, here are some useful starting points:

1. Consult the Amazon SageMaker AI MTRL documentation.
2. Prepare your dataset in the required format.
3. Prepare your agent and define your reward function based on your task-specific metric.
4. Start with default configurations, which worked well in our experience.
5. Evaluate and iterate.

---

About the authors

Huibin Shen

Huibin is an applied scientist in AWS, currently working on agentic search. His main interest is applying advanced reinforcement learning and retrieval machine learning techniques to solve real-world problems.

Pooja Karadgi

Pooja is a Senior Product Manager at AWS on the SageMaker AI team, where she works on developer experiences and model customization, including reinforcement learning. Her focus is making advanced fine-tuning techniques practical and accessible for builders.

Sifei Li

Sifei is a Senior Software Engineer in Amazon SageMaker where she’s working on building Amazon AI Platforms and was part of the launch team for Amazon SageMaker Model Customization and Multi-turn RL Platform.

Lorenzo Stella

Lorenzo is a Senior Applied Scientist at AWS, where he works on agentic search. His interest spans training of AI models and reinforcement learning, with a focus on decision-making systems and time series forecasting.