Training Agents with Multi-Turn Reinforcement Learning on SageMaker AI

According to a technical guide published by Daniel Quang on the AWS Blog, Multi-Turn Reinforcement Learning (MTRL) training in Amazon SageMaker AI makes it possible to train models for AI agent tasks (agentic tasks). In such tasks, an agent takes a sequence of actions where each step depends on the actions that preceded it, there is rarely a single correct sequence to imitate, and the only reliable feedback is whether the entire sequence succeeded. Supervised fine-tuning does not include labels for such situations, and single-turn reinforcement learning cannot reward a decision made three steps before the final outcome. The MTRL service in SageMaker AI trains models based on this signal, where the user does not specify the desired trajectory of actions, but only how to score the outcome.

The game Wordle serves as an accompanying example, in which a player has six tries to guess a hidden five-letter word, and after each guess receives feedback where a letter turns green if it is in the right position, yellow if it is in the word but in a different position, and gray if it is not in the word. In the experiment described in the guide, MTRL training increased the solve rate of the gpt-oss-20b model from 61 percent to 84 percent on held-out words at its default reasoning effort, and from 15 percent to 69 percent at low reasoning effort, for a token cost of approximately $339.

Solution Architecture, Supported Models, and Pricing Model

The system consists of two primary parts: SageMaker AI manages the training job and the policy model, while Amazon Bedrock AgentCore hosts the agent, the game environment, and the reward function. For every rollout, the MTRL service invokes the AgentCore runtime with one training row holding a secret word. A Strands agent in the runtime calls the policy model through an OpenAI-compatible endpoint of the SageMaker Job Runtime. The model makes up to six guesses through a dedicated tool called guess_word, and after each guess the environment returns per-letter feedback and a summary of known information. When the game ends, the agent returns a single reward. The Group Relative Policy Optimization (GRPO) algorithm uses the spread of rewards across several rollouts of the same word to update a LoRA adapter on top of the policy model.

At the time of writing the guide, the service operates on a serverless model and supports four base models, each trained with a LoRA adapter: the Amazon Nova 2 Lite and GPT-OSS-20B models are supported in the US East (N. Virginia) and US West (Oregon) Regions, while Gemma 4 31B (instruction-tuned) and Qwen 3.6 27B are supported in the US West (Oregon) Region. In the guide, GPT-OSS-20B was selected because it has the lowest per-token price of the four and its reasoning effort can be controlled per request.

Pricing is calculated per million tokens across three dimensions: Prefill (input tokens re-read at each turn, accumulating in multi-step tasks), Sample (tokens generated by the model, including reasoning tokens), and Train (tokens processed in the backward pass that updates the LoRA adapter). The 100-step training run consumed 491.6 million prefill tokens (costing $58.99 at $0.12 per million), 159.7 million sample tokens (costing $47.92 at $0.30 per million), and 643.3 million train tokens (costing $231.59 at $0.36 per million), totaling $338.50 across 105,486 rollouts over 12 hours and 40 minutes. vCPU and memory consumption in the AgentCore runtime was billed separately at approximately $79.

Permissions Setup, Agent Deployment, and Runtime Verification

The configuration process begins by creating two IAM roles:

1. Agent runtime role: Created using the agentcore deploy command, trusts bedrock-agentcore.amazonaws.com, and requires the AmazonSageMakerJobRuntimeAccess managed policy.
2. Training job role: Configured by the user, required to trust job.sagemaker.amazonaws.com with sts:AssumeRole and sts:TagSession actions, and requires the AmazonSageMakerJobFullAccess managed policy. The guide notes that the principal is job.sagemaker.amazonaws.com rather than generic sagemaker.amazonaws.com, which leads to a "Could not assume role" error.

The entire agent is defined in a single Python file (main.py) and deployed using the AgentCore CLI (which requires Python 3.10 or later, uv, and Node 20 or later, with botocore 1.43.37 or later required on the machine launching the job). The runtime reads two environment variables: REASONING_EFFORT, set to low to reduce token usage by 63 to 83 percent compared to medium, and REASONING_PROMPT_HINT, set to true to add a hint to the system prompt. The entry point function wraps the sagemaker_rft_handler decorator within BedrockAgentCoreApp and @app.entrypoint, where the entry point decorator must be the outermost to start an HTTP server and prevent a 30-second timeout error.

Before committing training resources, the authors recommend running a smoke test using the aws bedrock-agentcore invoke-agent-runtime command with an updated account ID. The test passes a fake job ARN and verifies that CloudWatch logs show a 400 error reading "Could not find job," confirming that the parsing, authentication, and reporting stages are functioning properly and returning a reward of -1.5.

Dataset Preparation and Launching Training

The MTRL dataset does not contain target responses, transcripts, or labels, but only task rows. Each row includes a JSON-formatted string with three fields: prompt presented to the model, answer (the secret word) used only by the environment and reward function and never exposed to the model, and id (row identifier). In the example, 640 training words (five batches of 128) and 128 validation words were generated with no overlap between them, and the files were uploaded to Amazon S3. The service also supports Apache Parquet, JSON arrays, and CSV formats.

The training job is launched using the MultiTurnRLTrainer class from the sagemaker-train package. In the experimental configuration, three hyperparameter overrides were set: sampling_max_tokens was set to 8192 (the service maximum) instead of the 4096 default to prevent the model from exhausting tokens during reasoning before its first tool call; max_epochs was set to 20; and max_steps was set to 100 (640 rows with a default global batch size of 128 produce 5 steps per epoch, and with the default group_size of 8, each step runs 1,024 games). In addition, an MLflow application is required to receive metrics. The guide explicitly warns against redeploying the agent while the training job is running, as restarting the runtime interrupts sampling requests and causes the entire job to fail.

Monitoring, Serving Settings Evaluation, and Resource Cleanup

During training, the service logs reward, token, and turn statistics at every step and runs validation every designated number of steps. The primary metric to monitor is val/reward/pass_at_1, representing the fraction of held-out words solved. Training reward rose to about +0.57 within 15 steps and stabilized there. The validation solve rate jumped from 2 percent to 31 percent in the first 10 steps, and subsequently hovered between 35 and 48 percent.

An evaluation run (JobCategory set to AgentRFTEvaluation) allows evaluating the trained model or base model under serving settings where temperature can be selected, unlike training validation which is fixed at temperature 0:

• At temperature 0, both models frequently end without solving across both reasoning effort levels.
• With sampling enabled at temperature 1.0 and medium reasoning effort, the trained model solved 84 percent of held-out words, compared to 61 percent for the base model.
• The LoRA adapter added between 14 and 54 percentage points to the solve rate across all four decoding settings evaluated. The authors recommend evaluating at actual intended serving settings, since the validation curve at temperature 0 showed a final solve rate of only 37 percent, whereas under serving settings with temperature 1.0 and medium effort it reached 84 percent. These evaluation runs took approximately one hour and cost about $1.20 in tokens. At the conclusion of the process, training and evaluation jobs stop billing when they complete, but the AgentCore runtime, its IAM role, the MLflow app, and the S3 objects persist, and the guide provides commands to delete them.

Lessons from Reward Design and Training Failures

Reward function design required revision after the initial design summed per-guess scoring across the entire game, producing two distortions: a lost game yielded a positive reward of about +0.49 due to accumulated bonuses, and a first-guess solve received a lower score (+0.99) than a third-guess solve (+1.61). The correction separated a dominant outcome term from a bounded shaping term: The applied formula is reward = outcome + w * clip(shaping, -1, 1), where weight w was set to 0.04 for a solve and 0.15 for a loss. The outcome component awards between 1.05 and 1.55 depending on solve speed (1.55 on guess 1 down to 1.05 on guess 6), losing after six guesses awards -0.5, and failing to make a valid guess awards -1.5. The shaping component includes penalties for contradicting known letters and bonuses for testing new letters and shrinking the search space. The low solve weight (0.04) ensures that a faster solve always outscores a slower one, and that any solve ranks above every loss.

In addition, the authors detailed four failures that caused training runs to fail:

1. Bearer token exhaustion under concurrency: Minting a token with a fresh session per call starved the credential provider at 32 concurrent rollouts, while caching for an hour yielded 403 errors during credential rotation. The fix utilized a self-refreshing session, a short cache, and re-signing on 403.
2. Policy model disconnection: When model calls failed authentication, the service cancelled the entire job with the error "No sampling requests were received." The agent must always return a reward (such as -1.5) rather than an error status.
3. Tool calls drifting into plain text: A high learning rate caused the model to produce tool calls as free text rather than structured calls, creating a constant floor reward and causing the gradient to vanish. The agent was updated to parse guesses from raw text with a slight shaping penalty.
4. Early plateau: Training reward stopped rising around step 15 and validation stalled around step 30. The authors noted that future exploration should evaluate using a curriculum seeding games with prior guesses or increasing the number of rollouts per word beyond the default of 8.