Multi-Turn Reinforcement Learning mit Amazon SageMaker AI: Anleitung zur Pipeline
Beim Multi-Turn Reinforcement Learning auf Amazon SageMaker AI lernen Modelle aus dem Ergebnis ganzer mehrstufiger Aufgaben statt aus einzelnen gelabelten Schritten. Eine Anleitung fuehrt anhand von Wordle durch den kompletten Ablauf, vom Agenten-Deployment bis zum Training.
Die meisten Fine-Tuning-Ansätze bringen einem Modell bei, auf einen einzelnen Prompt eine bessere Antwort zu erzeugen. Agentic-Aufgaben sind anders. Ein Agent, der eine Spesenabrechnung einreicht oder ein Support-Ticket einstuft, führt eine Abfolge von Aktionen aus, und jeder Schritt hängt von dem ab, was zuvor geschah. Es gibt selten eine einzige korrekte Abfolge, die man nachahmen könnte, und das einzige verlässliche Feedback ist, ob die gesamte Abfolge erfolgreich war. Supervised Fine-Tuning hat dafür keine Labels, und einstufiges Reinforcement Learning kann eine Entscheidung, die drei Schritte vor dem Ergebnis getroffen wurde, nicht entsprechend bewerten.
Multi-Turn Reinforcement Learning (MTRL) auf Amazon SageMaker AI trainiert Modelle genau auf diesem Signal. SageMaker AI sendet Prompts an einen Agenten, den Sie auf Amazon Bedrock AgentCore hosten oder auf Ihrer eigenen Infrastruktur, etwa Amazon Elastic Kubernetes Service (Amazon EKS), Amazon Elastic Compute Cloud (Amazon EC2) oder AWS Fargate, und zwar über einen AWS Lambda Forwarder. Das Beispiel hier verwendet AgentCore. Der Agent ruft das trainierte Policy-Modell auf, agiert in Ihrer Umgebung über beliebig viele Züge, wie die Aufgabe erfordert, und meldet eine Belohnung, sobald die Episode endet. Sie geben nicht die richtige Trajektorie vor, sondern nur, wie das Ergebnis zu bewerten ist.
Dieser Beitrag verwendet Wordle, ein Spiel, bei dem ein Spieler sechs Versuche hat, ein verstecktes fünf Buchstaben langes Wort zu erraten, als Beispiel. Nach jedem Rateversuch wird jeder Buchstabe grün, wenn er sich an der richtigen Position befindet, gelb, wenn er im Wort enthalten, aber an einer anderen Position ist, und grau, wenn er nicht im Wort vorkommt. Das Spiel eignet sich für MTRL, weil die Regeln wenige sind, jeder Rateversuch vom Feedback des vorherigen abhängt und das Ergebnis objektiv bewertet wird.
Wordle ist ein Lehrbeispiel, und aktuelle Frontier-Modelle spielen es bereits gut. In der Praxis würden Sie MTRL einsetzen, um ein kleines, kostengünstiges Modell in einer mehrstufigen Aufgabe fit zu machen, die spezifisch für Ihr Unternehmen ist, etwa einen internen Workflow, auf dem öffentliche Modelle nie trainiert wurden. Dieselbe Pipeline lässt sich auf dieses Szenario anwenden.
In unserem Lauf brachte MTRL gpt-oss-20b bei zurückgehaltenen Wörtern mit der Standard-Reasoning-Anstrengung von einer Erfolgsquote von 61 auf 84 Prozent, und bei der niedrigen Anstrengung, mit der es trainiert wurde, von 15 auf 69 Prozent, für etwa 339 US-Dollar an Trainings-Tokens. In diesem Beitrag führen wir durch die Pipeline von Anfang bis Ende: das Deployen des Agenten und der Belohnungsfunktion auf AgentCore, die Vorbereitung des Datensatzes, das Training, das Monitoring in MLflow und die Evaluierung unter Serving-Einstellungen. Wir teilen auch die Fehler bei der Belohnungsgestaltung und die Trainingsprobleme, auf die wir stießen. Ihre Datensätze, Ihr Agent und Ihr trainiertes Modell bleiben in Ihrem AWS-Konto und werden nicht zum Training der Basismodelle verwendet.
Das System besteht aus zwei Hälften. SageMaker AI verwaltet den Trainingsjob und das Policy-Modell. AgentCore hostet den Agenten, die Wordle-Umgebung und die Belohnungsfunktion. Das folgende Diagramm zeigt, wie beide verbunden sind.
Abbildung 1: MTRL auf SageMaker AI, wobei der Trainingsjob den auf AgentCore gehosteten Agenten einmal pro Rollout aufruft, um das Policy-Modell zu sampeln, das Spiel zu spielen und eine Belohnung zu melden
Bei jedem Rollout ruft der MTRL-Service die AgentCore Runtime, eine Funktion von Amazon Bedrock AgentCore, mit einer Trainingszeile auf, die ein geheimes Wort enthält. Ein Strands-Agent in der Runtime ruft das Policy-Modell über den OpenAI-kompatiblen Endpunkt des SageMaker Job Runtime auf. Das Modell spielt bis zu sechs Rateversuche über ein einziges Tool namens guess_word, und nach jedem Versuch gibt die Umgebung ein Feedback pro Buchstabe sowie eine Zusammenfassung des bekannten Wissens zurück. Wenn das Spiel endet, liefert der Agent eine Belohnung zurück. Group Relative Policy Optimization (GRPO) nutzt die Streuung der Belohnungen über mehrere Rollouts desselben Worts, um einen Low-Rank Adaptation (LoRA) Adapter am Policy-Modell zu aktualisieren.
Drei Designpunkte sind für das Folgende relevant:
MTRL ist serverless: Sie wählen ein Basismodell, und der Service stellt die Trainings- und Sampling-Infrastruktur bereit. Zum Zeitpunkt der Erstellung dieses Beitrags werden die folgenden Basismodelle unterstützt, die jeweils mit einem LoRA-Adapter trainiert werden:
Wir haben uns für GPT-OSS-20B entschieden, weil es den niedrigsten Preis pro Token der vier Modelle hat und sich seine Reasoning-Anstrengung pro Anfrage steuern lässt. Die aktuelle Liste finden Sie in der MTRL-Dokumentation.
Sie wählen keine Instanzen und dimensionieren keinen Cluster. Sie zahlen pro Million Tokens in drei Dimensionen, bepreist nach Basismodell:
Der Lauf mit 100 Schritten verbrauchte 491,6 Millionen Prefill-, 159,7 Millionen Sample- und 643,3 Millionen Train-Tokens über 105.486 Rollouts in 12 Stunden und 40 Minuten. Zu den Listenpreisen für GPT-OSS-20B in der Region US East (N. Virginia) (us-east-1) zum Zeitpunkt der Erstellung sind das etwa 339 US-Dollar:
Die Train-Dimension dominiert und skaliert mit der Trajektorienlänge, daher senken kürzere Spiele und weniger ausführliches Reasoning die Rechnung. Die AgentCore Runtime stellt separat die vom Agenten genutzten vCPUs und den Speicher in Rechnung. In diesem Lauf waren das etwa 79 US-Dollar, laut den Amazon CloudWatch Nutzungsmetriken der Runtime zu Listenpreisen in us-east-1. DescribeJob meldet die Token-Zahlen unter BillableTokenUsage. Aktuelle Preise finden Sie unter Amazon SageMaker AI Pricing.
Um dieser Anleitung zu folgen, benötigen Sie Folgendes. Sehen Sie sich auch die MTRL-Voraussetzungen an.
MTRL verwendet zwei Rollen, und eine davon hat eine Trust-Policy, die die allgemeine SageMaker AI Dokumentation nicht erwähnt.
Agent-Runtime-Rolle. agentcore deploy erstellt diese Rolle für Sie. Sie vertraut bedrock-agentcore.amazonaws.com und benötigt die verwaltete Policy AmazonSageMakerJobRuntimeAccess, die es dem Agenten erlaubt, das Policy-Modell aufzurufen und Belohnungen zu melden. Der CDK-Stack im Beispiel-Repository hängt diese Policy automatisch an.
Trainingsjob-Rolle. Diese erstellen Sie selbst. Sie muss job.sagemaker.amazonaws.com sowohl mit sts:AssumeRole als auch mit sts:TagSession vertrauen und benötigt die verwaltete Policy AmazonSageMakerJobFullAccess.
{ "Version": "2012-10-17", "Statement": [{ "Effect": "Allow", "Principal": {"Service": "job.sagemaker.amazonaws.com"}, "Action": ["sts:AssumeRole", "sts:TagSession"] }] } Hinweis: Der Trust-Principal ist job.sagemaker.amazonaws.com, nicht sagemaker.amazonaws.com. Mit dem generischen Principal schlägt CreateJob mit der Meldung „Could not assume role“ fehl, und die Fehlermeldung nennt den korrekten Principal erst nach mehreren Wiederholungen.
Klonen Sie das Beispiel-Repository, hinterlegen Sie Ihr Konto und Ihre Region in Wordle/agentcore/aws-targets.json und deployen Sie:
git clone https://github.com/aws-samples/sample-multi-turn-rl-wordle cd sample-multi-turn-rl-wordle/Wordle agentcore validate agentcore deploy Die CLI baut die Abhängigkeiten des Agenten für die Linux ARM64 Plattform der Runtime, erstellt die Ausführungsrolle mit angehängter AmazonSageMakerJobRuntimeAccess und deployt über AWS CloudFormation. Wir empfehlen diesen Weg gegenüber dem Hochladen eines Code-Pakets über die Konsole, bei dem es häufig zu Fehlern bei den rund 700 MB kompilierten Abhängigkeiten oder beim Anhängen der Policy kommt. Notieren Sie sich den Runtime-ARN aus der Ausgabe für den Trainingsjob.
Abbildung 2: Die Wordle-Agent-Runtime in der Amazon Bedrock AgentCore Konsole mit dem Status Ready
Die Runtime liest zwei Umgebungsvariablen aus agentcore.json. REASONING_EFFORT=low setzt reasoning_effort bei jeder Anfrage an das Policy-Modell. In unseren Evaluierungsjobs verbrauchte die niedrige Anstrengung 63 bis 83 Prozent weniger Sample-Tokens pro Spiel als die mittlere. REASONING_PROMPT_HINT=true fügt zudem „Reasoning: low“ zum System-Prompt hinzu, denn so liest gpt-oss die Einstellung in seinem Harmony-Chat-Format. Keine der beiden Variablen gibt dem Modell einen Hinweis auf das geheime Wort.
Abbildung 3: Umgebungsvariablen der Runtime mit REASONING_EFFORT auf low gesetzt
Der Agent besteht aus einer einzigen Datei, main.py. Vereinfacht sieht sein Einstiegspunkt wie folgt aus. Für Ihre eigene Aufgabe passen Sie das Task-Parsing, die Tools und die Belohnungsfunktion an:
app = BedrockAgentCoreApp() @app.entrypoint # äußerste Ebene: startet den HTTP-Server, den AgentCore aufruft @sagemaker_rft_handler # meldet die zurückgegebene Belohnung an SageMaker AI def handle_rollout(payload): prompt, answer = parse_task(payload["prompt"]) game = WordleGame(answer) @tool def guess_word(guess: str) -> str: """Submit a 5-letter Wordle guess and get feedback.""" return game.guess(guess) agent = Agent(model=build_model(), system_prompt=SYSTEM_PROMPT, tools=[guess_word]) agent(prompt) return {"reward": game.final_reward()} Wichtig: Der Decorator sagemaker_rft_handler startet selbst keinen HTTP-Server. Umschließen Sie ihn mit BedrockAgentCoreApp und @app.entrypoint, wobei der Entrypoint-Decorator äußerste Ebene sein muss. Andernfalls läuft bei jedem Aufruf nach 30 Sekunden ein Timeout ab.
Bevor Sie GPU-Zeit investieren, prüfen Sie mit einer Payload, die der vom Trainingsservice gesendeten ähnelt, ob die gesamte Kette funktioniert. Das Repository enthält eine solche Datei als Wordle/payload.json:
{ "prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"beach\", \"id\": \"smoke-test\"}", "metadata": { "jobArn": "arn:aws:sagemaker:us-east-1:111122223333:job/AgentRFT/smoke-test", "rolloutId": "r1" } } Es wird erwartet, dass die gefälschte Job-ARN fehlschlägt, und die Art des Fehlschlagens ist der Beweis. Tragen Sie zuerst Ihre eigene Konto-ID ein: Bei einer Job-ARN in einem anderen Konto gibt der Service einen 403-Fehler zurück, bevor er prüft, ob der Job existiert, was wie eine fehlerhafte IAM-Konfiguration aussieht. Führen Sie aus dem Wordle-Verzeichnis Folgendes aus:
ACCOUNT=$(aws sts get-caller-identity --query Account --output text) sed "s/111122223333/$ACCOUNT/" payload.json > /tmp/payload.json aws bedrock-agentcore invoke-agent-runtime \ --agent-runtime-arn <RUNTIME_ARN> \ --runtime-session-id smoke-test-$(date +%s)000000000000000 \ --payload fileb:///tmp/payload.json /dev/stdout In den Amazon CloudWatch Logs der Runtime sieht man, wie der Agent den Prompt parst, sich am Job Runtime Endpunkt authentifiziert und einen 400 Fehler mit der Meldung „Could not find job“ erhält. Dieser Fehlschlag bestätigt, dass Parsen, Authentifizierung und Fehlerberichterstattung funktionieren. Der Befehl selbst gibt {“reward”: -1.5} aus, weil der Agent das Policy-Modell nie erreicht hat.
Ein 403 mit der Meldung „not authorized to access this resource“ bedeutet, dass der Job ARN noch auf ein anderes Konto verweist.
Ein Timeout von 30 Sekunden bedeutet stattdessen, dass der HTTP Server nicht startet (siehe den Important Hinweis in Schritt 2).
Der MTRL Datensatz enthält keine Zielantworten, Transkripte oder Labels. Jede Zeile ist eine Aufgabe, die der Agent versuchen soll. Hier sind die ersten drei Zeilen der Trainingsdatei:
{"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"riser\", \"id\": \"wordle_train_0000\"}"} {"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"swoon\", \"id\": \"wordle_train_0001\"}"} {"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"dirty\", \"id\": \"wordle_train_0002\"}"} Der Dienst liest die Spalte prompt (oder die erste Spalte) und übergibt die Zeichenfolge unverändert an Ihren Agenten, ohne sie zu parsen. Wir haben die Zeichenfolge zu einem kleinen JSON Objekt mit drei Feldern gemacht:
Das Skript make_dataset.py im Beispiel Repository zieht 640 Trainingswörter (fünf volle Batches mit je 128) und 128 Validierungswörter aus der Antwortliste ohne Überschneidung, sodass die Validierungs-Lösungsrate die Generalisierung misst. Der Dienst akzeptiert außerdem Apache Parquet, JSON Arrays und CSV. Laden Sie beide Dateien nach Amazon S3 hoch:
uv run python make_dataset.py aws s3 cp training-data.jsonl s3://<YOUR_BUCKET>/wordle-mtrl/train/ aws s3 cp validation-data.jsonl s3://<YOUR_BUCKET>/wordle-mtrl/validation/ Da der Dienst die Prompt-Zeichenfolge niemals prüft, kann eine Produktionsaufgabe eine vollständige Nachrichtenliste, eine Tool-Konfiguration oder eine Reward-Spezifikation in sie packen. Beispiele dazu, einschließlich des Schutzes sensibler Inhalte in der Zeichenfolge, finden Sie in der MTRL Dokumentation.
Der Launcher verwendet die Klasse MultiTurnRLTrainer aus dem sagemaker-train SDK. Er zeigt auf die Runtime ARN aus Schritt 2, die Rolle aus Schritt 1, die Datensätze aus Schritt 4 und eine MLflow App für Metriken. Die MLflow App ist erforderlich; falls Sie keine haben, erstellen Sie zuerst eine mit aws sagemaker create-mlflow-app:
uv run python run_mtrl_training.py \ --model openai-reasoning-gpt-oss-20b \ --agent-runtime-arn <RUNTIME_ARN> \ --role-arn arn:aws:iam::<ACCOUNT>:role/<JOB_ROLE> \ --train-dataset s3://<YOUR_BUCKET>/wordle-mtrl/train/training-data.jsonl \ --val-dataset s3://<YOUR_BUCKET>/wordle-mtrl/validation/validation-data.jsonl \ --s3-output-path s3://<YOUR_BUCKET>/wordle-mtrl/output/ \ --mlflow-app-arn arn:aws:sagemaker:<REGION>:<ACCOUNT>:mlflow-app/<APP_ID> Der Launcher überschreibt drei Hyperparameter und lässt alle übrigen auf den gpt-oss-20b Standardwerten. Er setzt sampling_max_tokens auf 8192, das Dienstmaximum, weil gpt-oss beim Standardwert von 4096 beim reasoning oft die Tokens aufbrauchte, bevor der erste Tool-Aufruf erfolgte. Er setzt max_epochs auf 20 und max_steps auf 100: 640 Zeilen bei der Standard-größte global_batch_size von 128 ergeben 5 Schritte pro Epoche, und mit der Standard-group_size von 8 werden bei jedem Schritt 1.024 Spiele gespielt.
Der Job durchläuft die Phasen Starting, Downloading und Pending, bevor er Training erreicht. Sobald er trainiert, beginnt Ihre AgentCore Runtime, Rollout-Aufrufe zu empfangen.
Wichtig: Deployen Sie den Agenten nicht erneut, während ein Job läuft. Eine erneute Bereitstellung startet die Runtime neu, laufende Rollouts erhalten keine Sampling-Anfragen, und der Dienst bricht den gesamten Job fehlgeschlagen ab.
Der Dienst protokolliert bei jedem Schritt Reward-, Token- und Rundenstatistiken und führt alle val_every Schritte eine Validierung durch. Die wichtigste Metrik ist val/reward/pass_at_1, der Anteil der gelösten, zurückgehaltenen Wörter. Jede Lösung erzielt mindestens 1.01 und übersteigt damit den success_threshold des Dienstes von 1.0, sodass pass@1 sie zählt. Das folgende Diagramm basiert auf dem MLflow Verlauf dieses Laufs:
Abbildung 4: Trainingskurven für den 100-Schritte-Lauf, mit mittlerem Reward pro Spiel links und Validierungs-Lösungsrate rechts
Der Trainingsreward steigt innerhalb von 15 Schritten auf etwa +0,57 und bleibt dort. Der Validierungsreward bleibt niedriger, was hauptsächlich an der Temperatur 1.0 gegenüber 0 liegt und nicht an Overfitting (siehe Schritt 7). Die Validierungs-Lösungsrate springt in den ersten 10 Schritten von 2 auf 31 Prozent und bewegt sich danach zwischen 35 und 48 Prozent. Bei 128 Wörtern sind diese Schwankungen meist Rauschen. Wir haben nicht getestet, ob mehr Schritte bei den Serving-Einstellungen helfen.
Die OpenTelemetry Traces der Runtime in CloudWatch erfassen jede Modellrunde und jeden Tool-Aufruf. Hier ist ein Spiel, das das trainierte Modell während eines Evaluierungsjobs gespielt hat (mittlere Anstrengung, Temperatur 1.0), mit gekürztem reasoning. Das geheime Wort war PASTE:
[model] Reasoning: "CRANE" hat 5 unterschiedliche Buchstaben: C,R,A,N,E. Wir raten "CRANE". Tool call: guess_word(guess="CRANE") [tool] Versuch 1/6: CRANE -> C=grau, R=grau, A=gelb, N=grau, E=grün Aktueller Wissensstand: - Korrekte Position (Grün): _ _ _ _ E - Im Wort, falsche Position (Gelb): 'A' (mindestens 1) - Nicht im Wort (Grau): C, N, R - Verbleibende Versuche: 5 [model] Reasoning: "ALIVE": endet auf E, enthält A.
Keine Buchstaben C,N,R. Tool call: guess_word(guess="ALIVE") [tool] Versuch 2/6: ALIVE -> A=gelb, L=grau, I=grau, V=grau, E=grün [model] Reasoning: Also Positionen nicht 1 oder 3 für A. Verbleibende Positionen: 2,4 für A.
Vielleicht "HASTE" raten? Tool call: guess_word(guess="HASTE") [tool] Versuch 3/6: HASTE -> H=grau, A=grün, S=grün, T=grün, E=grün [model] Reasoning: Wir haben das Wort: _ASTE. "PASTE" ist ein Wort.
PASTE passt. Auch "WASTE" passt. Wähle vielleicht "PASTE". Tool call: guess_word(guess="PASTE") [tool] Versuch 4/6: PASTE --- alles grün.
Richtig, gelöst! Ein Modell-Rollout des trainierten Modells aus den Traces der AgentCore Runtime. Das Modell hält grüne Buchstaben an ihrem Platz, verwendet nie einen grauen Buchstaben wieder und nutzt jeden gelben Buchstaben, um Positionen auszuschließen, obwohl die Trainingsdaten ihm das nicht gezeigt haben.
Nach Abschluss des Jobs zeigt seine Seite in SageMaker AI Studio den Status und die Model Package Group, die das trainierte Modell enthält. DescribeJob mit JobCategory auf AgentRFT gibt die Output Model Package ARN und die abrechenbaren Token-Zählungen zurück.
Abbildung 5: Der abgeschlossene Trainingsjob in SageMaker AI Studio auf dem Tab Jobs, Training, Multi-turn RL
Ein Evaluierungsjob ( JobCategory auf AgentRFTEvaluation ) spielt Ihre Evaluierungsprompts durch dieselbe AgentCore Runtime und denselben Reward, auf dem Basismodell oder dem trainierten Paket, und protokolliert eval/reward/pass_at_1 an MLflow. Anders als bei der Trainingsvalidierung können Sie die Temperatur festlegen. Der Launcher des Repositorys führt Evaluierungen hintereinander aus, da das Standard-Kontingent nur einen Evaluierungsjob gleichzeitig erlaubt (eine Erhöhung können Sie anfordern):
uv run python run_mtrl_eval.py --runs base:1.0 trained:1.0 \ --model-package-arn <OUTPUT_MODEL_PACKAGE_ARN> \ --agent-runtime-arn <RUNTIME_ARN> \ --role-arn arn:aws:iam::<ACCOUNT>:role/<JOB_ROLE> \ --dataset s3://<YOUR_BUCKET>/wordle-mtrl/validation/validation-data.jsonl \ --s3-output-path s3://<YOUR_BUCKET>/wordle-mtrl/eval/ \ --mlflow-app-arn arn:aws:sagemaker:<REGION>:<ACCOUNT>:mlflow-app/<APP_ID> Abbildung 6: Reasoning-Aufwand und Sampling-Temperatur in jeder Phase; der Trainingsjob fixiert die Validierung auf Temperatur 0 und der bereitgestellte Aufwand, während Evaluierungsjobs Ihnen die Wahl lassen
Jede Evaluierung mit 128 Wörtern dauerte vier bis sieben Minuten. Der Reasoning-Aufwand stammt aus den Umgebungsvariablen des Agenten; um ihn zu ändern, bearbeiten Sie REASONING_EFFORT, stellen erneut bereit und evaluieren. Wir haben beide Modelle bei niedrigem und mittlerem Aufwand sowie bei Temperatur 0 und 1.0 evaluiert:
Abbildung 7: Basismodell gegenüber trainiertem Modell unter vier Decoding-Einstellungen, mit Lösungsrate links und mittlerem Reward pro Spiel rechts
Das graue Segment jedes Balkens ist das Basismodell, das blaue Segment der Gewinn durch das Training. Das Paar mit niedrigem Aufwand und Temperatur 0 ist die Validierung des Trainingsjobs bei den Schritten 0 und 100. Das Ergebnis des trainierten Modells bei mittlerem Aufwand und Temperatur 1.0 ist der Mittelwert aus zwei Läufen (80 und 88 Prozent).
Bei Temperatur 0 beenden beide Modelle das Spiel oft ohne Lösung, bei beiden Aufwandstufen. Mit aktiviertem Sampling bringt mittlerer Aufwand dem trainierten Modell 15 Punkte und dem Basismodell 46 Punkte. Der Adapter bringt in jeder Einstellung 14 bis 54 Punkte, sodass das, was das Training gelehrt hat, über die Einstellungen hinweg transferiert.
Die eigene Validierungskurve des Trainingsjobs in MLflow erzählte eine andere Geschichte: Sie erreichte bei Schritt 90 mit 48 Prozent ihren Höhepunkt und endete bei 37 Prozent, bewertet bei Temperatur 0. Die Hyperparameter-Referenz empfiehlt Temperatur 0 für Evaluierungen, weil greedy decoding die Kurven deterministisch und über Läufe hinweg vergleichbar macht. Wir berichten Temperatur 1.0, weil wir das Modell so betreiben. Bei mittlerem Aufwand und Temperatur 1.0 löst das trainierte Modell 84 Prozent der zurückgehaltenen Wörter, gegenüber 61 Prozent beim Basismodell.
Tipp: Werten Sie genau unter den Einstellungen aus, unter denen Sie später ausliefern. Die Trainingskurve zeigt, dass dieser Adapter bei 37 Prozent endet. Unter Serving-Einstellungen loest er 84 Prozent. Die Evaluationsjobs, die das zeigen, kosten etwa 1,20 US-Dollar an Tokens und dauern rund eine Stunde.
Trainings- und Evaluationsjobs werden nach Abschluss nicht weiter abgerechnet, aber der AgentCore-Runtime, seine IAM-Rolle, die MLflow-Anwendung und die S3-Objekte bleiben bestehen. Führen Sie vom Projektverzeichnis aus Folgendes aus:
# 1. Agent-Runtime und ihre Ausfuehrungsrolle. remove all setzt die Projektkonfiguration zurueck; # das naechste deploy sieht den leeren Zustand und baut die bereitgestellten Ressourcen ab. cd sample-multi-turn-rl-wordle/Wordle agentcore remove all agentcore deploy # 2. MLflow-Anwendung aws sagemaker delete-mlflow-app --arn arn:aws:sagemaker:<REGION>:<ACCOUNT>:mlflow-app/<APP_ID> # 3.
Datensaetze, Job-Ausgabe und Evaluationsausgabe aws s3 rm s3://<YOUR_BUCKET>/wordle-mtrl/ --recursive Da agentcore remove all die Dateien unter agentcore/ ueberschreibt, fuehren Sie git checkout -- agentcore/ aus, wenn Sie spaeter erneut bereitstellen moechten. Das Loeschen des CloudFormation-Stacks AgentCore-Wordle-default entfernt dieselben Ressourcen und laesst die Dateien unberuehrt. Die Trainingsjob-Rolle und die Modellpakete werden nicht abgerechnet.
Unsere Belohnungsfunktion benoetigte genau eine Ueberarbeitung, und der Fehler ist ein haeufiger.
Der erste Entwurf summierte die Bewertung pro Versuch aus dem Projekt wordle-lora-rl ueber das gesamte Spiel. Dort funktioniert diese Bewertung, weil jeder Versuch ein unabhaengiges Trainingssample ist. Zu einer Belohnung pro Spiel summiert, erzeugte sie zwei Umkehrungen, die wir nur durch das Lesen von Traces entdeckten:
Die Korrektur trennt einen dominanten Ergebnisterm von einem kleinen, begrenzten Shaping-Term:
reward = outcome + w * clip(shaping, -1, 1), w = 0.04 if solved, else 0.15 Der Ergebnisterm wird vollstaendig dadurch bestimmt, wie das Spiel endete:
Der Shaping-Term ist jene Bewertung pro Versuch, summiert ueber das Spiel:
Die Summe wird auf [-1, 1] begrenzt und mit 0.04 bei einem Loesungserfolg und 0.15 bei einer Niederlage skaliert, sodass sie Spiele mit demselben Ergebnis unterscheidet, ein Spiel aber nie in eine andere Stufe verschiebt. Bei typischen Shaping-Werten:
Eine schnellere Loesung schlaegt nun immer eine langsamere, jede Loesung schlaegt jede Niederlage, und Nichtspielen ist das schlechteste Ergebnis. Unsere erste Korrektur verwendete ein einziges Gewicht von 0.15, wodurch eine Loesung im dritten Versuch mit maximalem Shaping (+1.45) eine Loesung im zweiten Versuch mit dem Minimum (+1.25) uebertraf. Benachbarte Versuchszahlen liegen 0.1 auseinander, daher entfernt ein kleineres Loesungsgewicht von 0.04 die Ueberlappung, und die schlechteste Loesung (+1.01) ueberschreitet weiterhin die Erfolgsschwelle von 1.0.
Shaping bleibt relevant, weil GRPO aus der Belohnungsvarianz innerhalb einer Gruppe lernt, sodass zwei verlorene Spiele, die sich in der Hinweiskonsistenz unterscheiden, dem Modell dennoch etwas beibringen. Wie die SageMaker AI-Dokumentation empfiehlt: Beginnen Sie mit einer ergebnisbasierten Belohnung, fuegen Sie Shaping sparsam hinzu und pruefen Sie die Belohnung an echten Trajektorien.
Jeder dieser Fehler kostete mindestens einen Trainingslauf.
MTRL auf SageMaker AI trainiert ein Modell auf dem Ergebnis einer gesamten Episode. Das passt zu agentischen Aufgaben, bei denen die richtige Trajektorie unbekannt ist und nur das Ergebnis bewertet werden kann. Wir verwendeten eine einzige Belohnung auf Episodenebene.
Wenn auch Zwischenschritte bewertet werden koennen, akzeptiert MTRL ueber die UpdateReward-API zudem eine Belohnung pro Zug. In diesem Beitrag erreichte ein Modell mit offenen Gewichten und 20B eine Loesungsquote von 84 Prozent unter Serving-Einstellungen, gegenueber zuvor 61 Prozent, bei etwa 339 US-Dollar an Trainings-Tokens und etwa 79 US-Dollar an AgentCore-Runtime-Nutzung. Die 37 Prozent, die die Trainingskurve meldet, sind ein Artefakt der Art, wie der Trainingsjob die Validierung bewertet; werten Sie daher unter den Einstellungen aus, unter denen Sie ausliefern.
Dieselbe Pipeline laesst sich auf interne Verfahren, Tools und Workflows anwenden, auf die oeffentliche Modelle nie trainiert wurden.
Um zu starten, klonen Sie das Beispiel-Repository, folgen den vorangehenden Schritten und fuehren Sie den Smoke-Test aus Schritt 3 aus, bevor Sie GPU-Zeit investieren. Die zugrunde liegenden APIs finden Sie in der SageMaker AI MTRL-Dokumentation.
Daniel ist Solutions Architect im Team Frontier AI Startups bei AWS, wo er den groessten und strategisch wichtigsten Generative-AI-Startups hilft, ihre Infrastruktur in der Cloud zu skalieren. Er hat einen PhD in Informatik von der University of California, Irvine, mit Forschungsschwerpunkt auf maschinellem Lernen und Bioinformatik.
Hadrien ist Senior Frontier AI Account Manager im AWS Startup-Team und unterstuetzt fuehrende KI-Startups beim Aufbau technischer Integrationen und globaler Go-to-Market-Partnerschaften. Er verfuegt ueber acht Jahre Erfahrung in der Startup-Landschaft.
Vinay ist Head of Solutions Architecture for Frontier AI bei AWS, wo er ein Team von Technologen leitet, das Foundation-Modell-Laboren und KI-nativen Startups beim Aufbau, Training und der Skalierung ihrer Modelle und Anwendungen auf AWS hilft. Er bringt mehr als 20 Jahre Erfahrung im Aufbau technischer Teams bei AWS, Google Cloud und stark wachsenden Startups mit.
Automatisch übersetzt. Das Original auf Englisch ist unten aufklappbar.
Original (Englisch) anzeigen
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.
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.
Figure 1: MTRL on SageMaker AI, where the training job invokes the AgentCore-hosted agent once per rollout to sample the policy model, play the game, and report a reward
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:
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.
You don’t choose instances or size a cluster. You pay per million tokens in three dimensions, priced by base model:
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.
To follow this walkthrough, you need the following. See also the MTRL prerequisites.
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.
{ "Version": "2012-10-17", "Statement": [{ "Effect": "Allow", "Principal": {"Service": "job.sagemaker.amazonaws.com"}, "Action": ["sts:AssumeRole", "sts:TagSession"] }] } 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.
Clone the sample repository, set your account and Region in Wordle/agentcore/ aws-targets.json, and deploy:
git clone https://github.com/aws-samples/sample-multi-turn-rl-wordle cd sample-multi-turn-rl-wordle/Wordle agentcore validate agentcore 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.
Figure 2: The Wordle agent runtime in the Amazon Bedrock AgentCore console with status Ready
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.
Figure 3: Runtime environment variables showing REASONING_EFFORT set to low
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:
app = BedrockAgentCoreApp() @app.entrypoint # outermost: starts the HTTP server AgentCore invokes @sagemaker_rft_handler # reports the returned reward to SageMaker AI def handle_rollout(payload): prompt, answer = parse_task(payload["prompt"]) game = WordleGame(answer) @tool def guess_word(guess: str) -> str: """Submit a 5-letter Wordle guess and get feedback.""" return game.guess(guess) agent = Agent(model=build_model(), system_prompt=SYSTEM_PROMPT, tools=[guess_word]) agent(prompt) return {"reward": game.final_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.
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:
{ "prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"beach\", \"id\": \"smoke-test\"}", "metadata": { "jobArn": "arn:aws:sagemaker:us-east-1:111122223333:job/AgentRFT/smoke-test", "rolloutId": "r1" } } 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:
ACCOUNT=$(aws sts get-caller-identity --query Account --output text) sed "s/111122223333/$ACCOUNT/" payload.json > /tmp/payload.json aws bedrock-agentcore invoke-agent-runtime \ --agent-runtime-arn <RUNTIME_ARN> \ --runtime-session-id smoke-test-$(date +%s)000000000000000 \ --payload fileb:///tmp/payload.json /dev/stdout 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).
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:
{"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"riser\", \"id\": \"wordle_train_0000\"}"} {"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"swoon\", \"id\": \"wordle_train_0001\"}"} {"prompt": "{\"prompt\": \"Guess the 5-letter word\", \"answer\": \"dirty\", \"id\": \"wordle_train_0002\"}"} 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:
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:
uv run python make_dataset.py aws s3 cp training-data.jsonl s3://<YOUR_BUCKET>/wordle-mtrl/train/ aws s3 cp validation-data.jsonl s3://<YOUR_BUCKET>/wordle-mtrl/validation/ 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.
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:
uv run python run_mtrl_training.py \ --model openai-reasoning-gpt-oss-20b \ --agent-runtime-arn <RUNTIME_ARN> \ --role-arn arn:aws:iam::<ACCOUNT>:role/<JOB_ROLE> \ --train-dataset s3://<YOUR_BUCKET>/wordle-mtrl/train/training-data.jsonl \ --val-dataset s3://<YOUR_BUCKET>/wordle-mtrl/validation/validation-data.jsonl \ --s3-output-path s3://<YOUR_BUCKET>/wordle-mtrl/output/ \ --mlflow-app-arn arn:aws:sagemaker:<REGION>:<ACCOUNT>:mlflow-app/<APP_ID> 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.
The service logs reward, token, and turn statistics at every step and runs validation every val_every steps. The metric to watch is val/reward/pass_at_1, the fraction of held-out words solved. Every solve scores at least 1.01, so it clears the service’s success_threshold of 1.0 and pass@1 counts it. The following chart is built from this run’s MLflow history:
Figure 4: Training curves for the 100-step run, with mean reward per game on the left and validation solve rate on the right
Training reward climbs to about +0.57 within 15 steps and holds there. Validation reward stays lower, mostly because of temperature 1.0 versus 0 rather than overfitting (see Step 7). The validation solve rate jumps from 2 to 31 percent in the first 10 steps, then moves between 35 and 48 percent. With 128 words, those swings are mostly noise. We did not test whether more steps help at serving settings.
The runtime’s OpenTelemetry traces in CloudWatch record every model turn and tool call. Here is one game the trained model played during an evaluation job (medium effort, temperature 1.0), with the reasoning trimmed. The secret word was PASTE:
[model] Reasoning: "CRANE" has 5 distinct letters: C,R,A,N,E. We'll guess "CRANE". Tool call: guess_word(guess="CRANE") [tool] Guess 1/6: CRANE -> C=gray, R=gray, A=yellow, N=gray, E=green Current Knowledge: - Correct Position (Green): _ _ _ _ E - In Word, Wrong Position (Yellow): 'A' (at least 1) - Not in Word (Gray): C, N, R - Guesses Remaining: 5 [model] Reasoning: "ALIVE": ends with E, includes A.
Letters not C,N,R. Tool call: guess_word(guess="ALIVE") [tool] Guess 2/6: ALIVE -> A=yellow, L=gray, I=gray, V=gray, E=green [model] Reasoning: So positions not 1 or 3 for A. Positions left: 2,4 for A.
We might guess "HASTE"? Tool call: guess_word(guess="HASTE") [tool] Guess 3/6: HASTE -> H=gray, A=green, S=green, T=green, E=green [model] Reasoning: We have the word: _ASTE. "PASTE" is a word.
PASTE fits. Also "WASTE" also fits Choose "PASTE" maybe. Tool call: guess_word(guess="PASTE") [tool] Guess 4/6: PASTE --- all green.
Correct, you solved it! A trained-model rollout from the AgentCore runtime traces. The model keeps green letters in place, never reuses a gray letter, and uses each yellow letter to rule out positions, though nothing in the training data showed it how.
When the job completes, its page in SageMaker AI Studio shows the status and the model package group that holds the trained model. DescribeJob, with JobCategory set to AgentRFT, returns the output model package ARN and the billable token counts.
Figure 5: The completed training job in SageMaker AI Studio on the Jobs, Training, Multi-turn RL tab
An evaluation job ( JobCategory set to AgentRFTEvaluation ) plays your evaluation prompts through the same AgentCore runtime and reward, on the base model or the trained package, and logs eval/reward/pass_at_1 to MLflow. Unlike training validation, you can set the temperature. The repository’s launcher runs evaluations back to back, because the default quota allows one evaluation job at a time ( you can request an increase ):
uv run python run_mtrl_eval.py --runs base:1.0 trained:1.0 \ --model-package-arn <OUTPUT_MODEL_PACKAGE_ARN> \ --agent-runtime-arn <RUNTIME_ARN> \ --role-arn arn:aws:iam::<ACCOUNT>:role/<JOB_ROLE> \ --dataset s3://<YOUR_BUCKET>/wordle-mtrl/validation/validation-data.jsonl \ --s3-output-path s3://<YOUR_BUCKET>/wordle-mtrl/eval/ \ --mlflow-app-arn arn:aws:sagemaker:<REGION>:<ACCOUNT>:mlflow-app/<APP_ID> Figure 6: Reasoning effort and sampling temperature at each stage, where the training job fixes validation at temperature 0 and the deployed effort while evaluation jobs let you choose
Each 128-word evaluation took four to seven minutes. Reasoning effort comes from the agent’s environment variables, so to change it you edit REASONING_EFFORT, redeploy, and evaluate. We evaluated both models at low and medium effort and at temperature 0 and 1.0:
Figure 7: Base model versus trained model under four decoding settings, with solve rate on the left and mean reward per game on the right
The gray segment of each bar is the base model and the blue segment is the gain from training. The low-effort, temperature-0 pair is the training job’s validation at steps 0 and 100. The medium-effort, temperature-1.0 trained result is the mean of two runs (80 and 88 percent).
At temperature 0, both models often end the game without a solve, at either effort. With sampling on, medium effort adds 15 points for the trained model and 46 for the base model. The adapter adds 14–54 points in every setting, so what training taught transfers across settings.
The training job’s own validation curve in MLflow told a different story: it peaked at 48 percent at step 90 and ended at 37 percent, scored at temperature 0. The hyperparameters reference recommends temperature 0 for evaluation because greedy decoding makes curves deterministic and comparable across runs. We report temperature 1.0 because that’s how we serve the model. At medium effort and temperature 1.0, the trained model solves 84 percent of held-out words, up from 61 percent for the base model.
Tip: Evaluate at the settings you will serve at. The training curve says this adapter ends at 37 percent. At serving settings it solves 84 percent. The evaluation jobs that show this cost about $1.20 in tokens and took about an hour.
Training and evaluation jobs stop billing when they complete, but the AgentCore runtime, its IAM role, the MLflow app, and the S3 objects persist. From the project directory:
# 1. Agent runtime and its execution role. remove all resets the project config; # the next deploy sees the empty state and tears down the deployed resources. cd sample-multi-turn-rl-wordle/Wordle agentcore remove all agentcore deploy # 2. MLflow app aws sagemaker delete-mlflow-app --arn arn:aws:sagemaker:<REGION>:<ACCOUNT>:mlflow-app/<APP_ID> # 3. Datasets, job output, and evaluation output aws s3 rm s3://<YOUR_BUCKET>/wordle-mtrl/ --recursive Because agentcore remove all rewrites the files under agentcore/, run git checkout -- agentcore/ if you want to redeploy later. Deleting the AgentCore-Wordle-default CloudFormation stack removes the same resources and leaves the files alone. The training job role and the model packages are not billed.
Our reward needed one revision, and the mistake is a common one.
The first design summed the per-guess scoring from the wordle-lora-rl project over the whole game. That scoring works there because each guess is an independent training sample. Summed into one reward per game, it produced two inversions we found only by reading traces:
The fix separates a dominant outcome term from a small, bounded shaping term:
reward = outcome + w * clip(shaping, -1, 1), w = 0.04 if solved, else 0.15 The outcome term is set entirely by how the game ended:
The shaping term is that per-guess scoring, summed over the game:
The sum is clipped to [-1, 1] and scaled by 0.04 for a solve and 0.15 for a loss, so it separates games with the same outcome but never moves a game into another tier. With typical shaping values:
A faster solve now always beats a slower one, every solve beats every loss, and never playing is the worst outcome. Our first fix used a single weight of 0.15, which let a guess-3 solve with full shaping (+1.45) outscore a guess-2 solve with the minimum (+1.25). Adjacent guess counts are 0.1 apart, so a smaller solve weight of 0.04 removes the overlap, and the worst solve (+1.01) still clears the 1.0 success threshold. Shaping still matters because GRPO learns from reward variance within a group, so two lost games that differ in clue consistency still teach the model something. As the SageMaker AI documentation advises, start with an outcome-based reward, add shaping sparingly, and check the reward on real trajectories.
Each of these cost at least one training run.
MTRL on SageMaker AI trains a model on the outcome of a whole episode. That fits agentic tasks, where the right trajectory is unknown and only the result can be scored. We used a single episode-level reward. When intermediate steps can be scored too, MTRL also accepts one reward per turn through the UpdateReward API. In this post, an open-weight 20B model reached an 84 percent solve rate at serving settings, up from 61 percent, for about $339 in training tokens and about $79 in AgentCore runtime usage. The 37 percent the training curve reports is an artifact of how the training job scores validation, so evaluate at the settings you will serve at.
The same pipeline applies to internal procedures, tools, and workflows that public models were never trained on.
To get started, clone the sample repository, follow the preceding steps, and run the smoke test in Step 3 before committing GPU time. For the underlying APIs, see the SageMaker AI MTRL documentation.
Daniel is a Solutions Architect on the Frontier AI Startups team at AWS, where he helps the world’s largest and most strategic generative AI startups scale their infrastructure on the cloud. He holds a PhD in Computer Science from the University of California, Irvine, with a research focus on machine learning and bioinformatics.
Hadrien is a Senior Frontier AI Account Manager in the AWS Startup team, helping top AI startups build technical integrations and global go-to-market (GTM) partnerships. He has eight years of experience in the startup landscape.
Vinay is the Head of Solutions Architecture for Frontier AI at AWS, where he leads a team of technologists helping foundation model labs and AI-native startups build, train, and scale their models and applications on AWS. He brings more than 20 years of experience building technical teams across AWS, Google Cloud, and high-growth startups.