Tutorials

A hands-on guide demonstrates how to implement multi-turn reinforcement learning on Amazon SageMaker AI.

SageMaker AI uses multi-turn reinforcement learning to train agents from final rewards; a Wordle demo lifted GPT-OSS-20B solve rates to 84% for $339.

A hands-on guide demonstrates how to implement multi-turn reinforcement learning on Amazon SageMaker AI.

Key Takeaways

  • Amazon SageMaker Multi-Turn Reinforcement Learning trains agents from final episode outcomes instead of single labeled responses, making it useful for workflows where success depends on a sequence of choices.
  • The AWS Wordle example lifted gpt-oss-20b’s solve rate from 61% to 84% on held-out words at default reasoning effort, and from 15% to 69% during training, for about $339 in token costs.
  • The pipeline splits responsibilities: SageMaker AI owns the policy model and training job, while AgentCore hosts the agent, environment, and reward function that reports one outcome-based reward per episode.
  • Training uses Group Relative Policy Optimization to update a LoRA adapter based on reward differences across multiple rollouts, and costs are dominated by train tokens as trajectory length increases.
  • MTRL is serverless, supports a limited set of base models and regions, and requires careful setup around IAM roles, runtime ARNs, SDK versions, and same-region deployment.

Meta description: Amazon SageMaker Multi-Turn Reinforcement Learning trains agents on success-only feedback. A Wordle example lifted solve rate to 84 percent for $339.

A hands-on guide demonstrates how to implement multi-turn reinforcement learning on Amazon SageMaker AI.

Tags: Amazon SageMaker AI, agent training, reinforcement learning, Wordle, AWS pricing

Featured image: abstract friendly illustration of a multi-turn reinforcement learning loop between a model, an agent, and a scoring function

For developers building agents, Amazon's Multi-Turn Reinforcement Learning on Amazon SageMaker AI addresses a training gap: agentic work rarely has one labeled answer, because success depends on a sequence of choices. The AWS walkthrough uses Wordle as a small, objective task and reports that gpt-oss-20b improved from a 61 percent solve rate to 84 percent on held-out words at its default reasoning effort. At the low reasoning effort used during training, the solve rate rose from 15 percent to 69 percent, with about $339 in training-token costs.

Why sequence tasks need a different training signal

Supervised fine-tuning usually maps one prompt to one preferred response. An agent that files an expense report or triages a support ticket does something closer to a game: each action changes the state, and the final outcome is the clearest signal. Multi-Turn Reinforcement Learning trains on that outcome. SageMaker AI sends prompts to an agent hosted on Amazon Bedrock AgentCore, or on your own infrastructure such as Amazon Elastic Kubernetes Service (Amazon EKS), Amazon Elastic Compute Cloud (Amazon EC2), or AWS Fargate, through an AWS Lambda forwarder. The example uses AgentCore. The agent calls the policy model being trained, acts over multiple turns, and reports one reward when the episode ends. You supply a scoring rule, not a perfect path.

A hands-on guide demonstrates how to implement multi-turn reinforcement learning on Amazon SageMaker AI.

Wordle fits because the rules are small: six guesses, one hidden five-letter word, and per-letter feedback after each guess. Green means the letter is correct and in the right place, yellow means it is in the word but misplaced, and gray means it is absent. The game is only a teaching example; strong frontier models already play it well. The practical target is a small, lower-cost model that must become good at a specific multi-step business workflow that public models were never trained on.

What the Wordle run produced

The pipeline has two sides. SageMaker AI owns the training job and the policy model. AgentCore hosts the agent, the Wordle environment, and the reward function. For every rollout, the MTRL service invokes the AgentCore runtime with one training row that holds one secret word. A Strands agent in the runtime calls the policy model through the SageMaker Job Runtime's OpenAI-compatible endpoint. The model plays up to six guesses through one tool, guess_word. After each guess, the environment returns letter feedback and a summary of what is known. When the game ends, the agent returns one reward.

Group Relative Policy Optimization uses the spread of rewards across several rollouts of the same word to update a Low-Rank Adaptation adapter on the policy model. Three design choices shape the run. The agent and reward live in one Python file deployed to AgentCore. The dataset carries only the secret word; the bundled word lists are built from the open-source SCOWL word lists and are used only for scoring. 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.

Models, regions, and the $339 token bill

MTRL is serverless: you choose a base model, and the service provisions the training and sampling infrastructure. At the time of writing, the supported base models, each trained with a LoRA adapter, are Amazon Nova 2 Lite in US East (N. Virginia) and US West (Oregon), GPT-OSS-20B in the same two regions, Gemma 4 31B instruction-tuned in US West (Oregon), and Qwen 3.6 27B in US West (Oregon). AWS chose GPT-OSS-20B because it had the lowest per-token price of the four and its reasoning effort is controllable per request.

You pay per million tokens in three dimensions: prefill, sample, and train. Prefill is input tokens the policy model processes, which accumulate because each turn re-reads the history. Sample is tokens the policy model generates, including reasoning tokens. Train is tokens processed in the backward pass that updates the LoRA adapter. The 100-step run consumed 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 (N. Virginia) at the time of writing, the token bill was about $339.

DimensionTokensPrice per 1MCost
Prefill491,616,806$0.12$58.99
Sample159,737,098$0.30$47.92
Train643,308,554$0.36$231.59
Total1,294,662,458$338.50

The train dimension dominates and scales with trajectory length, so shorter games and less verbose reasoning reduce cost. The AgentCore runtime bills separately for vCPU and memory used by the agent; for this run, that was about $79, based on the runtime's Amazon CloudWatch usage metrics at us-east-1 list prices. DescribeJob reports token counts under BillableTokenUsage.

Setup steps that decide whether training starts

  1. Prepare the roles and Region. You need Python 3.10 or later, uv, Node 20 or later, and the AgentCore CLI, installed with npm install -g @aws/agentcore. The machine launching the job needs botocore 1.43.37 or later, the first version that includes the AgentRFT CreateJob operation. You also need two Identity and Access Management roles and an Amazon Simple Storage Service bucket in the same Region as the agent. The training job must run in that Region because CreateJob rejects a cross-Region runtime ARN.

The agent runtime role is created by agentcore deploy. It trusts bedrock-agentcore.amazonaws.com and needs the AmazonSageMakerJobRuntimeAccess managed policy, which lets the agent call the policy model and report rewards. The training job role is created manually. Its trust policy must name job.sagemaker.amazonaws.com and grant both sts:AssumeRole and sts:TagSession. It also needs AmazonSageMakerJobFullAccess. A common trap is using sagemaker.amazonaws.com instead; CreateJob then fails with Could not assume role, and the message names the correct principal only after retries.

  1. Deploy the agent. Clone the aws-samples/sample-multi-turn-rl-wordle repository, set the account and Region in Wordle/agentcore/aws-targets.json, then run agentcore validate and agentcore deploy. The CLI builds dependencies for the runtime's Linux ARM64 platform, creates the execution role with AmazonSageMakerJobRuntimeAccess attached, and deploys through AWS CloudFormation. AWS recommends this over a console code upload because the roughly 700 MB of compiled dependencies or the policy attachment are easy to get wrong. Note the runtime ARN for the training job.

The runtime reads REASONING_EFFORT=low and REASONING_PROMPT_HINT=true from agentcore.json. The first sets reasoning_effort on every request to the policy model; in the evaluation jobs, low effort used 63 to 83 percent fewer sample tokens per game than medium. The second adds Reasoning: low to the system prompt, which is how gpt-oss reads the setting in its Harmony chat format. Neither variable gives the model any hint about the secret word. The agent is one file, main.py. The sagemaker_rft_handler decorator does not start an HTTP server by itself; it must be wrapped with BedrockAgentCoreApp and @app.entrypoint, with the entrypoint decorator outermost. Without that, every invocation times out after 30 seconds.

  1. Smoke-test the chain before GPU time. The repository includes Wordle/payload.json with a training-shaped prompt and a sample job ARN. Replace the sample account ID with your own, then invoke the AgentCore runtime. A failure in Amazon CloudWatch logs reading Could not find job is the expected proof that parsing, authentication, and reward reporting work. The invocation returns reward -1.5 because the agent never reaches the policy model. A 403 reading not authorized to access this resource usually means the job ARN still names another account. A 30-second timeout usually means the HTTP server is not starting.
  1. Prepare a task-only dataset. MTRL datasets do not include target responses, transcripts, or labels. Each row is one task for the agent to attempt. In the Wordle example, that row carries the secret word. The walkthrough also covers training, monitoring in MLflow, and evaluating at serving settings. Your datasets, agent, and trained model stay in your AWS account and are not used to train the base models.

Common mistakes

  • Trusting the generic sagemaker.amazonaws.com principal instead of job.sagemaker.amazonaws.com for the training job role.
  • Forgetting sts:TagSession in the training job role trust policy.
  • Deploying through the console when the AgentCore CLI dependency build and policy attachment are already handled.
  • Omitting BedrockAgentCoreApp and the outer @app.entrypoint, causing every runtime invocation to time out after 30 seconds.
  • Leaving a job ARN from another account in the smoke-test payload, which creates a 403 that looks like a broken IAM setup.
  • Running the training job in a different Region from the runtime or dataset bucket, which CreateJob rejects.
  • Ignoring trajectory length and reasoning verbosity, because train tokens scale with the multi-turn sequence and dominate the bill.

What readers should watch

The Wordle example is not meant to show that agents need training to play a common game. It shows how to turn a success-only signal into a policy update for a small model. Readers evaluating this pipeline should watch three things: whether the supported model and region list expands beyond the current set, how much of the cost is driven by long trajectories, and whether outcome-only scoring can transfer from objective games to internal workflows where the reward is harder to define. The sample code and the documented failure modes are the practical part of the announcement, because they show where a multi-turn training job is likely to break before it reaches the first successful rollout.

Comments (0)

  1. No comments yet. Be the first to share what worked for you.

Leave a comment

Comments are reviewed before they appear. Your email address is not published.