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

來源原標題

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

作者

Huibin Shen

原始標籤

Advanced (300) · Amazon SageMaker AI · Technical How-to

原始連結

https://aws.amazon.com/blogs/machine-learning/fine-tune-a-search-agent-with-multi-turn-rl-on-amazon-sagemaker-ai/

摘要原文

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. 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: 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). Our setup relied on the following: 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: 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. This section walks through the three key components of our training setup: the datasets, the reward function, and the MTRL job configuration. We use the following datasets for training and testing: 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. 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. 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: The key hyperparameters we change are: 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. 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. 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. 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. The results show that RL fine-tuning made a meaningful difference: If you do explore this approach, remember to clean up the resources afterward to avoid ongoing charges: 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: If you want to apply this approach to your own agents, here are some useful starting points: