Multi‑Turn Reinforcement Learning on SageMaker AI: What the Wordle Test Shows

Multi‑Turn Reinforcement Learning on SageMaker AI: What the Wordle Test Shows

According to Artificial Intelligence, Amazon SageMaker AI now offers Multi‑Turn Reinforcement Learning (MTRL), a way to train agents that must act over several steps instead of answering a single prompt. The blog walks through a Wordle‑style demo, reports a jump from 61 % to 84 % solve rate on a 20‑billion‑parameter model, and breaks down a $339 token bill.

How MTRL Works Under the Hood

MTRL treats each episode—an entire game or workflow—as a single training example. The environment (here a Wordle engine) lives in Amazon Bedrock AgentCore. For every rollout the training job sends a secret word to the agent runtime, which calls the policy model (the model being tuned) through an OpenAI‑compatible endpoint. The model can make up to six guesses via a guess_word tool; after each guess the environment returns per‑letter feedback. When the episode ends the agent reports one scalar reward. The reward‑shaping logic is simple: a win earns 1.05–1.55 (higher for fewer guesses), a loss –0.5, and an invalid guess –1.5, plus a tiny bonus for good intermediate strategy. The service then applies Group Relative Policy Optimization (GRPO) to the spread of rewards and updates a Low‑Rank Adaptation (LoRA) adapter on the base model.

Cost Breakdown of the Wordle Run

The authors ran a 100‑step training job on the open‑source GPT‑OSS‑20B model, which has the lowest per‑token price of the supported models. Token consumption and Amazon‑listed rates gave the following bill:

Dimension Tokens Price per 1 M Cost
Prefill (input) 491,616,806 $0.12 $58.99
Sample (generated) 159,737,098 $0.30 $47.92
Train (back‑prop) 643,308,554 $0.36 $231.59
Total 1,294,662,458 — $338.50

The training dimension dominates because every turn re‑reads the full conversation history, inflating the backward‑pass token count. Shorter games or lower‑effort reasoning (the REASONING_EFFORT=low flag) cut the sample tokens by 63–83 % and reduce the bill. The AgentCore runtime itself added about $79 for vCPU and memory usage, measured via CloudWatch.

Performance Gains and Limitations

Before MTRL, GPT‑OSS‑20B solved Wordle on held‑out words 61 % of the time when using its default reasoning setting. After 105,486 rollouts over 12 h 40 min, the solve rate rose to 84 %. At the low‑effort setting the model went from a dismal 15 % to 69 % solve. Those numbers prove that a single scalar reward can steer a large language model toward better multi‑step decision making.

However, the test also exposed a few practical snags:

  • The reward function is the only supervision signal. If the reward is mis‑specified, the model can learn pathological strategies (the authors note early “shaping term” mistakes).
  • The training loop assumes a deterministic environment. Real‑world workflows often involve external APIs that can fail or return variable latency, which would break the tight rollout contract.
  • The token‑based pricing means that any increase in episode length—more turns, richer tool output, or verbose reasoning—directly inflates cost.

What Actually Changes for Your Business

The trade‑off is clear: you trade cheap single‑turn fine‑tuning for a higher‑cost, higher‑variance loop that can teach a model to chain actions. In practice this means you can take a cheap base model, add a small LoRA adapter (a few megabytes), and end up with an agent that can file an expense report or triage a ticket without needing a giant, expensive foundation model. The upside is lower inference cost because the fine‑tuned model is smaller; the downside is the upfront token bill and the need to write a reliable environment and reward.

Who benefits? Small‑to‑medium teams that have a well‑defined, rule‑based workflow and can simulate it locally. Large enterprises with already‑built data pipelines may find the engineering overhead too high compared to using a managed agent service. Who loses? Vendors that sell single‑turn fine‑tuning as a one‑size‑fits‑all solution; their offering cannot cover the multi‑step credit problem that MTRL solves.

What to watch next: the upcoming support for additional base models (e.g., Gemma 4 31B) and the rollout of the new “low‑effort” pricing tier that could halve sample token costs. Also keep an eye on AWS’s announced “reward‑function library,” which aims to standardise shaping terms and reduce the trial‑and‑error phase the authors describe.

Quick Start Checklist for Today

  1. Clone the AWS sample repo and install the AgentCore CLI (npm i -g @aws/agentcore).
  2. Create the two IAM roles exactly as shown (note the trust principal job.sagemaker.amazonaws.com).
  3. Deploy the Wordle agent with agentcore deploy and note the runtime ARN.
  4. Run the smoke‑test payload (aws bedrock-agentcore invoke-agent-runtime …) to verify the chain works.
  5. Replace the Wordle logic with a simple internal workflow (e.g., “approve a leave request”) and adjust the reward function to give +1 for a successful approval and –1 for a rejection.
  6. Kick off a short training run (aws sagemaker create-training-job …) with a low‑effort setting to keep costs under $50.

If the test run produces a positive reward on a few dozen rollouts, you already have a working multi‑turn agent you can start experimenting with in production.

Sources

Read next

We count page views without cookies — no identifier, nothing stored on your device. Accept to allow cookies for analytics.