AWS shows multi-turn reinforcement learning on SageMaker AI with a Wordle agent

Why multi-turn tasks need a different training approach

AWS argues that most fine-tuning aims at a better answer to a single prompt, while agentic tasks such as filing an expense report or triaging a support ticket involve a sequence of actions where each step depends on the previous ones. There is rarely one correct sequence to imitate, and the only reliable feedback is whether the whole sequence succeeded.

According to the post, supervised fine-tuning has no labels for that situation, and single-turn reinforcement learning cannot credit a decision made three steps before the outcome. MTRL on SageMaker AI is built to train on that end-of-episode signal. You do not specify the right trajectory, only how to score the result.

SageMaker AI sends prompts to an agent hosted on Amazon Bedrock AgentCore, or on your own infrastructure such as Amazon EKS, Amazon EC2 or AWS Fargate, through an AWS Lambda forwarder. The agent calls the policy model being trained, acts in the environment over as many turns as needed, and reports one reward when the episode ends.

How the Wordle example works

AWS chose Wordle because the rules are few, each guess depends on the feedback from the last, and the outcome can be scored objectively. The company stresses that Wordle is a teaching example and that frontier models already play it well. In practice, it says, MTRL is meant to make a small, low-cost model good at a multi-step task specific to a business, such as an internal workflow that public models were never trained on.

The system has two halves. SageMaker AI owns the training job and the policy model, while AgentCore hosts the agent, the Wordle environment and the reward function. For each rollout, the MTRL service invokes the AgentCore runtime with one training row holding one secret word. A Strands agent calls the policy model and plays up to six guesses through a single tool, guess_word. Group Relative Policy Optimization (GRPO) then uses the spread of rewards across several rollouts of the same word to update a Low-Rank Adaptation (LoRA) adapter on the policy model.

The reward is outcome-dominant. A solve pays 1.05 to 1.55 depending on speed, a loss pays -0.5, and never making a valid guess pays -1.5. A small shaping term rewards good strategy within each tier. AWS also says datasets, the agent and the trained model stay in the customer's AWS account and are not used to train the base models.

Results, supported models and cost

In AWS's run, MTRL took gpt-oss-20b from 61 to 84 percent on held-out words at its default reasoning effort, and from 15 to 69 percent at the low effort it was trained at. AWS chose GPT-OSS-20B because it has the lowest per-token price of the four supported models and its reasoning effort is controllable per request.

MTRL is serverless: users pick a base model and the service provisions the infrastructure. At the time of writing, the supported models, each trained with a LoRA adapter, are:

• Amazon Nova 2 Lite: US East (N. Virginia) and US West (Oregon)
• GPT-OSS-20B: US East (N. Virginia) and US West (Oregon)
• Gemma 4 31B (instruction-tuned): US West (Oregon)
• Qwen 3.6 27B: US West (Oregon)

Pricing is per million tokens across three dimensions: prefill, sample and train. The 100-step run used 491.6 million prefill, 159.7 million sample and 643.3 million train tokens across 105,486 rollouts in 12 hours 40 minutes. At GPT-OSS-20B list prices in us-east-1, the total came to $338.50, of which train tokens accounted for $231.59. The AgentCore runtime is billed separately and cost about $79 for this run, according to AWS.

Setup requirements and what it means for teams

The walkthrough requires Python 3.10 or later, uv, Node 20 or later and the AgentCore CLI. It also needs botocore 1.43.37 or later, described as the first version that includes the AgentRFT CreateJob operation, plus an Amazon S3 bucket in the same Region as the agent. Cross-Region runtime ARNs are rejected by CreateJob.

Two IAM roles are involved. The agent runtime role is created by agentcore deploy. The training job role must trust job.sagemaker.amazonaws.com with both sts:AssumeRole and sts:TagSession. AWS warns that using the generic sagemaker.amazonaws.com principal makes CreateJob fail with "Could not assume role."

For teams building agents, the post points to a way to tune a small model on a business-specific workflow using only a scoring function rather than labeled trajectories. AWS notes that the train dimension dominates cost and scales with trajectory length, so shorter tasks and less verbose reasoning reduce the bill. It also says it shares the reward-design mistakes and training failures it encountered, and that all code is in an aws-samples repository.