AWS shows how to train a search agent with multi-turn RL on SageMaker AI

Why multi-turn training for search agents

According to AWS, a search agent powered by a large language model decides on its own what to search for, which retrieval strategy to use and when to stop, refining its approach over several rounds based on what it has already retrieved. Getting that behavior to work reliably is hard. AWS says no base model arrives knowing a company's tools or environment, a small prompted model rarely gives dependable multi-turn behavior, and a frontier model often works but costs more in latency and money.

Fine-tuning is presented as a third path: teach a small model your tools and environment directly. AWS argues that the usual methods fall short. Supervised fine-tuning depends on expert demonstrations of ideal multi-turn trajectories, which are costly to collect and usually do not exist for a given setup. Single-turn reinforcement learning, such as RL with verifiable rewards (RLVR), scores one response at a time and so misses the dependencies between decisions across turns.

Multi-turn reinforcement learning instead optimizes the agent across the full trajectory, and the reward only needs to reflect whether the final outcome was good. AWS adds that running a smaller, specialized model also brings faster and cheaper inference.

What SageMaker AI MTRL offers

Amazon SageMaker AI MTRL fine-tunes LLMs with reinforcement learning in multi-turn settings. It treats an agentic task as a sequence of decisions, generates training data through multi-turn rollouts and optimizes the model with policy gradient algorithms. AWS lists these features:

• Modular agent-environment interface: low-code integration, with custom rewards, tool loops and conversation shapes.
• Serverless execution: per-token pricing without provisioning or managing GPU clusters.
• Asynchronous rollout and trajectory collection: generation and gradient updates run in parallel with bounded off-policy staleness.
• Native algorithm library: PPO, CISPO and importance-sampling losses, paired with group-based advantage estimators such as GRPO, GRPO pass@k and RLOO.
• Resumable training: long runs can be split across multiple jobs.
• Observability: turn-by-turn trajectories and rewards can be inspected in MLflow managed by Amazon SageMaker AI.
• Evaluation jobs: reward, pass@k and trajectory metrics before deploying to a SageMaker AI endpoint or Amazon Bedrock.

The training setup

The example uses Qwen3.6-27B in the US West (Oregon) Region (us-west-2). The agent has two tools: lexical search with BM25, which finds exact keyword matches, and vector search, which compares embeddings and is recommended for semantic or conceptual queries. The number of turns is capped to discourage excessively long responses and encourage efficient search.

Training data combines six datasets: FRAMES, BRIGHT, Enterprise RAG, ESCI, Musique and MLQA. Testing uses FreshStack, WixQA, BrowseComp-Plus and Wands. AWS reserves 5 percent of the training instances in each dataset for validation.

The reward is nDCG@10, a standard retrieval metric where 1.0 means perfect ranking and 0.0 means no relevant documents retrieved. It is applied at the trajectory level, after the agent finishes its full search. If the agent hits the maximum number of turns or the maximum sampling tokens in a single turn, it receives a reward of -1. AWS says this penalty steers the model away from those failures without hand-crafted intermediate rewards.

For the job itself, AWS says it changed three hyperparameters and left everything else at default values. Its MultiTurnRLTrainer SDK example sets max_epochs to 1, global_batch_size to 128 and rollout_max_concurrency to 32. The source text provided breaks off before the results section.

What it means for enterprise teams

The post positions MTRL as a fit for search agents because they have a clear reward signal in retrieval quality, a multi-turn loop of queries and results, and a well-defined environment of search tools. For teams without expert demonstrations or GPU clusters, AWS's pitch is a lower-configuration route to a small model that behaves reliably in their own environment. The retrieval quality and reliability results that AWS says it observed are not included in the portion of the source available here.