Amazon SageMaker AI的多轮强化学习：实用指南

中文版正在随流水线自动整理，稍后刷新即可；可先阅读下方英文原文。

本文由 AI宝库自动抓取并翻译，版权归原作者所有；如需查看原文请点下方链接。

Artificial Intelligence

Multi-Turn Reinforcement Learning on Amazon SageMaker AI: A hands-on guide

Most fine-tuning teaches a model to produce a better answer to a single prompt. Agentic tasks are different. An agent that files an expense report or triages a support ticket takes a sequence of actions, and each step depends on what came before. There’s rarely one correct sequence to imitate, and the only reliable feedback is whether the whole sequence succeeded. Supervised fine-tuning has no labels for that, and single-turn reinforcement learning cannot credit a decision made three steps before the outcome.

Multi-Turn Reinforcement Learning (MTRL) on Amazon SageMaker AI trains models on exactly that signal. SageMaker AI sends prompts to an agent you host 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 here uses AgentCore. The agent calls the policy model being trained, acts in your environment over as many turns as the task needs, and reports one reward when the episode ends. You don’t specify the right trajectory, only how to score the result.

This post uses Wordle, a game in which a player has six tries to guess a hidden five-letter word, as the example. After each guess, every letter turns green if it is in the right position, yellow if it’s in the word but in a different position, and gray if it isn’t in the word. The game suits MTRL because the rules are few, each guess depends on the feedback from the last, and the outcome is scored objectively. Wordle is a teaching example, and frontier models already play it well. In practice, you would use MTRL to make a small, low-cost model good at a multi-step task specific to your business, such as an internal workflow that public models were never trained on. The same pipeline applies to that setting.

In our run, MTRL took gpt-oss-20b from a 61 to an 84 percent solve rate on held-out words at its default reasoning effort, and from 15 to 69 percent at the low effort it trained at, for about $339 in training tokens. In this post, we walk through the pipeline end-to-end: deploying the agent and reward to AgentCore, preparing the dataset, training, monitoring in MLflow, and evaluating at serving settings. We also share the reward-design mistakes and training failures we hit. Your datasets, agent, and trained model stay in your AWS account and are not used to train the base models. All code is in the accompanying aws-samples repository.

Solution overview

The system has two halves. SageMaker AI owns the training job and the policy model. AgentCore hosts the agent, the Wordle environment, and the reward function. The following diagram shows how they connect.

For every rollout, the MTRL service invokes the AgentCore runtime, a capability of Amazon Bedrock AgentCore, with one training row, which 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, and after each guess the environment returns per-letter feedback and a summary of what is known. When the game ends, the agent returns one reward. Group Relative Policy Optimization (GRPO) uses the spread of rewards across several rollouts of the same word to update a Low-Rank Adaptation (LoRA) adapter on the policy model.

Three design points matter for what follows:

- The agent and reward live in one place. A single Python file deployed to AgentCore holds the environment, the tool, and the scoring function.

- The dataset carries only the secret word. The word lists bundled with the agent 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–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.

Supported models and AWS Regions

MTRL is serverless: you pick a base model and the service provisions the training and sampling infrastructure. At the time of writing, it supports the following base models, each trained with a LoRA adapter:

We chose GPT-OSS-20B because it has the lowest per-token price of the four and its reasoning effort is controllable per request. See the MTRL documentation for the current list.

Pricing

You don’t choose instances or size a cluster. You pay per million tokens in three dimensions, priced by base model:

- Prefill: input tokens the policy model processes. These accumulate in a multi-turn task, because every turn re-reads the history.

- Sample: tokens the policy model generates, including reasoning tokens.

- Train: 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 the US East (N. Virginia) Region (us-east-1) at the time of writing, that is about $339:

The train dimension dominates and scales with trajectory length, so shorter games and less verbose reasoning cut the bill. The AgentCore runtime bills separately for the vCPU and memory the agent uses. For this run, that was about $79, from the runtime’s Amazon CloudWatch usage metrics at us-east-1 list prices. DescribeJob reports the token counts under BillableTokenUsage. See Amazon SageMaker AI pricing for current rates.

Prerequisites

To follow this walkthrough, you need the following. See also the MTRL prerequisites.

- Python 3.10 or later, uv, Node 20 or later, and the AgentCore CLI (npm install -g @aws/agentcore). See Get started with Amazon Bedrock AgentCore for installation and the IAM permissions the CLI needs.

- botocore 1.43.37 or later, on the machine that launches the job. This is the first version that includes the AgentRFT CreateJob operation.

- Two AWS Identity and Access Management (IAM) roles, described in Step 1.

- An Amazon Simple Storage Service (Amazon S3) bucket in the same Region as the agent. The training job must also run in that Region, because CreateJob rejects a cross-Region runtime ARN.

Step 1: Create the IAM roles

MTRL uses two roles, and one of them has a trust policy that the general SageMaker AI documentation doesn’t mention.

Agent runtime role. agentcore deploy creates this role for you. It trusts bedrock-agentcore.amazonaws.com and needs the AmazonSageMakerJobRuntimeAccess managed policy, which lets the agent call the policy model and report rewards. The CDK stack in the sample repository attaches this policy automatically.

Training job role. You create this one. It must trust job.sagemaker.amazonaws.com with both sts:AssumeRole and sts:TagSession, and it needs the AmazonSageMakerJobFullAccess managed policy.

Note: The trust principal is job.sagemaker.amazonaws.com, not sagemaker.amazonaws.com. With the generic principal, CreateJob fails with “Could not assume role,” and the message names the correct principal only after retries.

Step 2: Deploy the agent to Amazon Bedrock AgentCore

Clone the sample repository, set your account and Region in Wordle/agentcore/aws-targets.json, and deploy:

The CLI builds the agent’s dependencies for the runtime’s Linux ARM64 platform, creates the execution role with AmazonSageMakerJobRuntimeAccess attached, and deploys through AWS CloudFormation. We recommend it over uploading a code package on the console, where it’s common to get the roughly 700 MB of compiled dependencies or the policy attachment wrong. Note the runtime ARN in the output for the training job.

The runtime reads two environment variables from agentcore.json. REASONING_EFFORT=low sets reasoning_effort on every request to the policy model. In our evaluation jobs, low effort used 63 to 83 percent fewer sample tokens per game than medium. REASONING_PROMPT_HINT=true also 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. Simplified, its entrypoint looks like the following. For your own task, you change the task parsing, the tools, and the reward:

Important: The sagemaker_rft_handler decorator does not start an HTTP server by itself. Wrap it with BedrockAgentCoreApp and @app.entrypoint, with the entrypoint decorator outermost. Without this, every invocation times out after 30 seconds.

Step 3: Verify the runtime

Before spending GPU time, confirm the whole chain works with a payload shaped like the one the training service sends. The repository includes one as Wordle/payload.json:

The fake job ARN is expected to fail, and the way it fails is the proof. Put your own account ID into it first: for a job ARN in another account, the service returns a 403 before it checks whether the job exists, which looks like a broken IAM setup. From the Wordle directory, run:

In the runtime’s Amazon CloudWatch logs, you see the agent parse the prompt, authenticate to the Job Runtime endpoint, and receive a 400 error reading “Could not find job.” That failure confirms that parsing, authentication, and error reporting work. The command itself prints {“reward”: -1.5}, because the agent never reached the policy model. A 403 reading “not authorized to access this resource” means the job ARN still names another account. A 30-second timeout instead means the HTTP server is not starting (see the Important note in Step 2).

Step 4: Prepare the dataset

The MTRL dataset has no target responses, transcripts, or labels. Each row is one task for the agent to attempt. Here are the first three rows of the training file:

The service reads the prompt column (or the first column) and passes the string to your agent verbatim, without parsing it. We made the string a small JSON object with three fields:

- prompt: the instruction the agent shows the policy model, identical on every row.

- answer: the secret word. Only the environment and the reward function use it. The policy model never sees it.

- id: a row ID for tracing a rollout in the logs back to its source row.

The sample repository’s make_dataset.py draws 640 training words (five full batches of 128) and 128 validation words from the answer list with no overlap, so the validation solve rate measures generalization. The service also accepts Apache Parquet, JSON arrays, and CSV. Upload both files to Amazon S3:

Because the service never inspects the prompt string, a production task can pack a full message list, tool configuration, or reward specification into it. See the MTRL documentation for examples, including how to protect sensitive content in the string.

Step 5: Launch the training job

The launcher uses the MultiTurnRLTrainer class from the sagemaker-train SDK. Point it at the runtime ARN from Step 2, the role from Step 1, the datasets from Step 4, and an MLflow app for metrics. The MLflow app is required, so if you don’t have one, create it first with aws sagemaker create-mlflow-app:

The launcher overrides three hyperparameters and leaves every other one at the gpt-oss-20b default. It sets sampling_max_tokens to 8192, the service maximum, because at the default of 4096 gpt-oss often ran out of tokens while reasoning, before its first tool call. It sets max_epochs to 20 and max_steps to 100: 640 rows at the default global_batch_size of 128 is 5 steps per epoch, and with the default group_size of 8, each step plays 1,024 games. The job moves through Starting, Downloading, and Pending before reaching Training. After it’s training, your AgentCore runtime begins receiving rollout invocations.

Important: Don’t redeploy the agent while a job is running. Redeploying restarts the runtime, in-flight rollouts see no sampling requests, and the service fails the entire job.

Step 6: Monitor training in MLflow

The service logs reward, token, and turn statist

Auto-fetched and translated by AI Gem Box. All rights belong to the original author.