跳到正文
AWS Machine Learning Blog· Daniel Quang·· 12 小时前AI 评分59

AWS 实战教程:在 SageMaker AI 上做多轮强化学习训练 gpt-oss-20b 玩 Wordle

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

AI 导读

AWS 机器学习博客发布在 Amazon SageMaker AI 上进行多轮强化学习(MTRL)的实战教程,agent 部署在 Amazon Bedrock AgentCore,用 GRPO 更新 LoRA 适配器,以 Wordle 为示例任务。

正文 · AI 翻译

大多数微调教模型对单个提示生成更好的答案。智能体任务则不同。提交报销单或对支持工单进行分类的智能体需要执行一系列动作,而每一步都取决于之前的步骤。很少有唯一正确的动作序列可供模仿,而唯一可靠的反馈是整个序列是否成功。监督微调没有针对这种情况的标签,而单轮强化学习无法为在结果出现前三步所做的决策分配功劳。

Amazon SageMaker AI 上的多轮强化学习(MTRL)正是基于这种信号来训练模型。SageMaker AI 通过 AWS Lambda 转发器,将提示发送给你托管在 Amazon Bedrock AgentCore 上的智能体,或托管在你自己的基础设施上的智能体,例如 Amazon Elastic Kubernetes Service(Amazon EKS)、Amazon Elastic Compute Cloud(Amazon EC2)或 AWS Fargate。本文示例使用 AgentCore。智能体调用正在训练的策略模型,在你的环境中执行任务所需的任意多轮操作,并在回合结束时报告一个奖励值。你无需指定正确的轨迹,只需指定如何对结果评分。

本文以 Wordle 为例,这是一种玩家有六次机会猜测一个隐藏的五个字母单词的游戏。每次猜测后,如果字母在正确位置则变为绿色,如果在单词中但位置不同则为黄色,如果不在单词中则为灰色。这个游戏很适合 MTRL,因为规则很少,每次猜测都取决于上一次的反馈,而且结果可以客观评分。Wordle 只是一个教学示例,前沿模型已经能玩得很好。在实践中,你会使用 MTRL 让一个小型、低成本的模型擅长时间特定于你业务的多步骤任务,例如公共模型从未训练过的内部工作流程。同样的流水线也适用于那种场景。

在我们的运行中,MTRL 以约 339 美元的训练代币成本,让 gpt-oss-20b 在其默认推理努力下对保留词的解决率从 61% 提升到 84%,在其训练时所处的低推理努力下从 15% 提升到 69%。在本文中,我们将端到端地介绍整个流水线:将智能体和奖励部署到 AgentCore、准备数据集、训练、在 MLflow 中监控,以及在服务设置下进行评估。我们还会分享我们遇到的奖励设计错误和训练失败。你的数据集、智能体和训练后的模型都保留在你的 AWS 账户中,并且 不会被用于训练基础模型。所有代码都在配套的 aws-samples 仓库中。

解决方案概览

该系统分为两部分。SageMaker AI 负责训练作业和策略模型。AgentCore 托管智能体、Wordle 环境和奖励函数。下图展示了它们如何连接。

SageMaker AI training job connected to the AgentCore-hosted agent, Wordle environment, and reward function

图 1:SageMaker AI 上的 MTRL,训练作业在每次 rollout 中调用一次 AgentCore 托管的智能体,以对策略模型进行采样、玩游戏并报告奖励

每次发布任务时,MTRL 服务会使用一条训练数据调用 Amazon Bedrock AgentCore 的一项功能——AgentCore runtime,每条数据包含一个秘密单词。runtime 中的 Strands 智能体通过 SageMaker Job Runtime 的 OpenAI 兼容端点调用策略模型。模型通过一个工具 guess_word 进行最多六次猜测,每次猜测后,环境会返回逐字母的反馈以及已知信息的摘要。游戏结束时,智能体返回一个奖励。组相对策略优化(GRPO)利用同一个单词多次发布任务所产生的奖励分布,来更新策略模型上的低秩适配(LoRA)适配器。

对于后文的内容,有三个设计要点值得注意:

  • 智能体和奖励位于同一处。部署到 AgentCore 的单个 Python 文件同时包含环境、工具和评分函数。
  • 数据集仅包含秘密单词。智能体自带的单词列表由开源的 SCOWL 词表构建,仅用于评分。
  • 奖励以结果为主导。成功解谜根据速度获得 1.05–1.55 的奖励,失败为 -0.5,从未做出有效猜测则为 -1.5。一个小的塑形项会奖励每个难度层级内的良好策略。

支持的模型和 AWS 区域

MTRL 是无服务器的:你只需选择基础模型,服务便会预置训练和采样基础设施。在撰写本文时,它支持以下基础模型,每个模型都通过 LoRA 适配器进行训练:

基础模型 AWS 区域
Amazon Nova 2 Lite 美国东部(弗吉尼亚北部)(us-east-1)、美国西部(俄勒冈)(us-west-2)
GPT-OSS-20B 美国东部(弗吉尼亚北部)(us-east-1)、美国西部(俄勒冈)(us-west-2)
Gemma 4 31B(指令微调版) 美国西部(俄勒冈)(us-west-2)
Qwen 3.6 27B 美国西部(俄勒冈)(us-west-2)

我们选择 GPT-OSS-20B 是因为它是四者中每 token 价格最低的,且其推理力度可以按请求进行控制。有关当前列表,请参阅 MTRL 文档。

定价

你无需选择实例或规划集群规模。你按三个维度的每百万 token 付费,价格由基础模型决定:

  • 预填充(Prefill):策略模型处理的输入 token。在多轮任务中,这些 token 会不断累积,因为每一轮都会重新读取历史记录。
  • 采样(Sample):策略模型生成的 token,包括推理 token。
  • 训练(Train):在更新 LoRA 适配器的反向传播过程中处理的 token。

这次 100 步的运行在 12 小时 40 分钟内,通过 105,486 次发布任务,共消耗了 4.916 亿个预填充 token、1.597 亿个采样 token 和 6.433 亿个训练 token。按照撰写本文时美国东部(弗吉尼亚北部)区域 (us-east-1) 的 GPT-OSS-20B 公开定价,这大约是 339 美元:

维度 Token 数 每百万价格 费用
预填充 491,616,806 $0.12 $58.99
采样 159,737,098 $0.30 $47.92
训练 643,308,554 $0.36 $231.59
总计 1,294,662,458 $338.50

训练维度占主导,并随轨迹长度而扩展,因此更短的游戏和更简洁的推理可以减少费用。AgentCore runtime 会单独对智能体使用的 vCPU 和内存进行计费。对于这次运行,根据 runtime 的 Amazon CloudWatch 使用指标按 us-east-1 公开价格计算,费用约为 79 美元。DescribeJob 在 BillableTokenUsage 下报告 token 数量。当前费率请参阅 Amazon SageMaker AI 定价。

先决条件

要完成本演练,你需要满足以下条件。另请参阅 MTRL 先决条件。

  • 需要 Python 3.10 或更高版本、uv、Node 20 或更高版本,以及 AgentCore CLI(npm install -g @aws/agentcore)。安装方法及 CLI 所需的 IAM 权限请参阅 Amazon Bedrock AgentCore 入门。
  • 在启动任务的机器上需要 botocore 1.43.37 或更高版本。这是首个包含 AgentRFT CreateJob 操作的版本。
  • 两个 AWS Identity and Access Management (IAM) 角色,在步骤 1 中说明。
  • 一个与智能体位于同一区域的 Amazon Simple Storage Service (Amazon S3) 存储桶。训练任务也必须在该区域运行,因为 CreateJob 会拒绝跨区域的运行时 ARN。

步骤 1:创建 IAM 角色

MTRL 使用两个角色,其中一个角色的信任策略是 SageMaker AI 通用文档中没有提到的。

智能体运行时角色。agentcore deploy 会自动为你创建此角色。它信任 bedrock-agentcore.amazonaws.com,并需要 AmazonSageMakerJobRuntimeAccess 托管策略,该策略允许智能体调用策略模型并报告奖励。示例代码库中的 CDK 堆栈会自动附加此策略。

训练任务角色。这个角色由你创建。它必须信任 job.sagemaker.amazonaws.com,并同时拥有 sts:AssumeRole 和 sts:TagSession,还需要 AmazonSageMakerJobFullAccess 托管策略。

{
    "Version": "2012-10-17",
    "Statement": [{
        "Effect": "Allow",
        "Principal": {"Service": "job.sagemaker.amazonaws.com"},
        "Action": ["sts:AssumeRole", "sts:TagSession"]
    }]
}

注意:信任主体是 job.sagemaker.amazonaws.com,而不是 sagemaker.amazonaws.com。如果使用通用主体,CreateJob 会失败并报错“Could not assume role”,而且只有在重试之后,错误消息才会指出正确的主体。

步骤 2:将智能体部署到 Amazon Bedrock AgentCore

克隆示例代码库,在 Wordle/agentcore/aws-targets.json 中设置你的账户和区域,然后部署:

git clone https://github.com/aws-samples/sample-multi-turn-rl-wordle
cd sample-multi-turn-rl-wordle/Wordle
agentcore validate
agentcore deploy

CLI 会为运行时的 Linux ARM64 平台构建智能体的依赖项,创建附加了 AmazonSageMakerJobRuntimeAccess 的执行角色,并通过 AWS CloudFormation 进行部署。我们推荐使用 CLI,而不是在控制台上上传代码包——后者很容易在约 700 MB 的编译依赖项或策略附加环节出错。请记下输出中的运行时 ARN,训练任务会用到它。

Amazon Bedrock AgentCore console showing the Wordle agent runtime with status Ready

图 2:Amazon Bedrock AgentCore 控制台中状态为 Ready 的 Wordle 智能体运行时

运行时会从 agentcore.json 读取两个环境变量。REASONING_EFFORT=low 会在每次对策略模型的请求上设置 reasoning_effort。在我们的评估任务中,低推理力度比中等力度每个游戏少用 63% 到 83% 的采样 token。REASONING_PROMPT_HINT=true 还会在系统提示词中加入 “Reasoning: low”,gpt-oss 正是通过其 Harmony 聊天格式中的这一设置来读取该配置的。这两个变量都不会向模型透露任何关于秘密单词的提示。

AgentCore console environment variables panel with REASONING_EFFORT set to low

图 3:运行时环境变量,显示 REASONING_EFFORT 设置为 low

该智能体只有一个文件 main.py。简化后的入口点如下所示。对于你自己的任务,你需要修改任务解析、工具和奖励:

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()}

重要提示:sagemaker_rft_handler 装饰器本身不会启动 HTTP 服务器。需要用 BedrockAgentCoreApp 和 @app.entrypoint 包裹它,且入口点装饰器在最外层。否则,每次调用都会在 30 秒后超时。

步骤 3:验证运行时

在消耗 GPU 时间之前,先用一个形状与训练服务所发送负载相同的负载来确认整条链路是否正常工作。代码库中提供了一个示例,位于 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"
  }
}

这个伪造的任务 ARN 预期会失败,而它的失败方式本身就是验证。请先把你自己的账户 ID 填进去:对于其他账户中的任务 ARN,服务会在检查任务是否存在之前就返回 403,这看起来像是 IAM 配置有问题。在 Wordle 目录下运行:

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

在运行时的 Amazon CloudWatch 日志中,你会看到智能体解析提示词、向 Job Runtime 端点进行身份验证,并收到一条 400 错误,内容为“Could not find job”。这个失败恰恰确认了解析、身份验证和错误报告都在正常工作。命令本身输出 {“reward”: -1.5},因为智能体从未到达策略模型。如果出现内容为“not authorized to access this resource”的 403 错误,说明任务 ARN 仍然指向另一个账户。如果是 30 秒超时,则说明 HTTP 服务器没有启动(参见第 2 步中的重要说明)。

第 4 步:准备数据集

MTRL 数据集没有目标响应、对话记录或标签。每一行都是智能体要尝试完成的一个任务。以下是训练文件的前三行:

{"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\"}"}

该服务读取 prompt 列(或第一列),并将字符串原样传递给你的智能体,不做解析。我们把字符串做成了包含三个字段的小型 JSON 对象:

  • prompt: 智能体展示给策略模型的指令,每一行都相同。
  • answer: 秘密单词。只有环境和奖励函数会使用它,策略模型永远看不到它。
  • id: 行 ID,用于在日志中将一次 rollout 追溯到其来源行。

示例仓库的 make_dataset.py 从答案列表中无重叠地抽取 640 个训练单词(即五个完整的 128 大小的批次)和 128 个验证单词,因此验证解题率可以衡量泛化能力。该服务还支持 Apache Parquet、JSON 数组和 CSV。将这两个文件上传到 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/

由于该服务从不检查 prompt 字符串的内容,生产环境的任务可以在其中打包完整的消息列表、工具配置或奖励规范。更多示例(包括如何保护字符串中的敏感内容),请参阅 MTRL 文档。

第 5 步:启动训练任务

启动器使用 sagemaker-train SDK 中的 MultiTurnRLTrainer 类。将其指向第 2 步中的运行时 ARN、第 1 步中的角色、第 4 步中的数据集,以及一个用于指标的 MLflow 应用。MLflow 应用是必需的,如果你还没有,请先用 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>

启动器覆盖了三个超参数,其余均保持 gpt-oss-20b 的默认值。它将 sampling_max_tokens 设为 8192(即服务上限),因为在默认值 4096 下,gpt-oss 在进行推理时常常在第一次工具调用之前就耗尽了 token。它将 max_epochs 设为 20,max_steps 设为 100:640 行数据在默认 global_batch_size 为 128 时每个 epoch 为 5 步,而在默认 group_size 为 8 时,每步进行 1,024 局游戏。任务会依次经过 Starting、Downloading 和 Pending 状态,然后进入 Training。开始训练后,你的 AgentCore 运行时将开始接收 rollout 调用。

重要提示:在任务运行期间,不要重新部署智能体。重新部署会重启运行时,进行中的 rollout 将收不到采样请求,服务会使整个任务失败。

第 6 步:在 MLflow 中监控训练

该服务在每一步都会记录奖励、token 和轮次统计信息,并每隔 val_every 步运行一次验证。需要关注的关键指标是 val/reward/pass_at_1,即留出单词的解题比例。每次解题得分至少为 1.01,因此能超过服务的 success_threshold 阈值 1.0,并被计入 pass@1。下图基于本次运行的 MLflow 历史数据构建:

Two line charts tracking mean reward per game and validation solve rate over 100 training steps

图 4:100 步运行的训练曲线,左侧为每局平均奖励,右侧为验证解题率

训练奖励在 15 步内攀升至约 +0.57 并保持在该水平。验证奖励较低,主要原因是温度 1.0 与 0 的差异,而非过拟合(见第 7 步)。验证解决率在前 10 步从 2% 跃升至 31%,随后在 35% 到 48% 之间波动。在 128 个单词的规模下,这些波动大多只是噪声。我们没有测试在服务设置下更多训练步数是否有所帮助。

运行时的 OpenTelemetry 跟踪记录在 CloudWatch 中,记录了每一轮模型交互和每次工具调用。下面是训练后的模型在一次评估作业(中等推理力度,温度 1.0)中玩的一局游戏,推理内容已删减。秘密单词是 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!

来自 AgentCore 运行时跟踪记录的一次训练后模型推演。模型把绿色字母保留在原位,从不重复使用灰色字母,并利用每个黄色字母排除位置——尽管训练数据中没有任何内容教它这样做。

第 7 步:在服务设置下评估训练后的模型

作业完成后,它在 SageMaker AI Studio 中的页面会显示状态以及保存训练后模型的模型包组。DescribeJob 在 JobCategory 设置为 AgentRFT 时,会返回输出模型包的 ARN 以及计费 token 数量。

SageMaker AI Studio Jobs page showing the completed multi-turn RL training job

图 5:SageMaker AI Studio 中 Jobs、Training、Multi-turn RL 标签页下已完成的训练作业

一个评估作业(JobCategory 设置为 AgentRFTEvaluation)会通过相同的 AgentCore 运行时和奖励函数运行你的评估提示,可以使用基础模型或训练后的模型包,并将 eval/reward/pass_at_1 记录到 MLflow。与训练验证不同,你可以设置温度。该代码仓库的启动器会连续运行评估,因为默认配额一次只允许一个评估作业(你可以申请提高配额):

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>

Reasoning effort and sampling temperature at each stage: training low and 1.0, validation low and 0 fixed by the service, evaluation jobs at four settings with serving at medium and 1.0.

图 6:各阶段的推理力度和采样温度,其中训练作业将验证固定在温度 0 和部署时的推理力度,而评估作业则允许你自行选择

每次 128 个单词的评估耗时四到七分钟。推理力度来自智能体的环境变量,因此要更改它,你需要编辑 REASONING_EFFORT、重新部署并评估。我们在低和中等推理力度、温度 0 和 1.0 下评估了两个模型:

Bar charts comparing base and trained model solve rate and mean reward across four decoding settings

图 7:四种解码设置下基础模型与训练后模型的对比,左侧为解决率,右侧为每局平均奖励

每个条形的灰色部分是基础模型,蓝色部分是训练带来的提升。低推理力度、温度 0 的那一组是训练作业在第 0 步和第 100 步的验证结果。中等推理力度、温度 1.0 的训练后结果是两次运行的平均值(80% 和 88%)。

在温度 0 下,无论哪种推理力度,两个模型经常在没有解开谜底的情况下结束游戏。开启采样后,中等推理力度为训练后模型带来 15 个百分点的提升,为基础模型带来 46 个百分点。适配器在每种设置下都带来 14 到 54 个百分点的提升,说明训练所教授的内容可以跨设置迁移。

训练作业自身在 MLflow 中的验证曲线却讲述了另一个故事:它在第 90 步达到 48% 的峰值,结束时为 37%,评分温度为 0。超参数参考建议评估时使用温度 0,因为贪心解码能使曲线具有确定性,便于跨运行比较。我们报告的是温度 1.0 的结果,因为我们就是这样部署该模型的。在中等推理力度和温度 1.0 下,训练后的模型能解开 84% 的保留单词,高于基础模型的 61%。

提示:请在你实际部署的设置下进行评估。训练曲线显示这个适配器最终达到 37 percent。而在部署设置下它解决了 84 percent。展示这一点的评估任务花费了约 $1.20 的 token 费用,耗时约一小时。

第 8 步:清理

训练和评估任务完成后即停止计费,但 AgentCore runtime、其 IAM 角色、MLflow 应用以及 S3 对象会继续保留。在项目目录中运行:

# 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

由于 agentcore remove all 会重写 agentcore/ 下的文件,如果你之后想重新部署,请运行 git checkout -- agentcore/。删除 AgentCore-Wordle-default CloudFormation 堆栈会移除相同的资源,但不会动那些文件。训练任务角色和模型包不产生费用。

深入探讨:设计奖励

我们的奖励只做了一次修订,而那个错误很常见。

最初的设计是把 wordle-lora-rl 项目中每次猜测的得分在整个游戏中求和。那种评分在那个项目中可行,因为每次猜测都是独立的训练样本。但对每局游戏求和成一个奖励后,产生了两个倒挂问题,我们只有通过阅读轨迹才发现:

  • 输掉的游戏反而得到正奖励。六轮加成加起来约为 +0.49。
  • 第 1 轮解出的得分低于第 3 轮解出的。第一猜就解出得 +0.99,而第三猜解出得 +1.61,因为额外的轮次收集了更多加成。

修复方案是将一个占主导地位的结果项与一个小的、有界的塑形项分开:

reward = outcome + w * clip(shaping, -1, 1), w = 0.04 if solved, else 0.15

结果项完全由游戏如何结束决定:

结果 数值
在第 n 猜解出 1.05 + 0.5 * (6 – n) / 5,即第 1 猜为 1.55,第 6 猜为 1.05
猜满六次,未解出 -0.5
从未做出有效猜测 -1.5

塑形项就是那次猜测的评分,在整局游戏中求和:

  • 对与已知的绿色、黄色或灰色字母相矛盾的惩罚,对重复猜测的惩罚,以及对不在词典中的词的惩罚。
  • 对强有力的开局词的熵加成。
  • 对尝试新字母以及缩小剩余可能性集合的加成。

总和被裁剪到 [-1, 1],解出时乘以 0.04,失败时乘以 0.15,因此它能区分结果相同的不同游戏,但绝不会把一局游戏移到另一个档位。在典型的塑形值下:

场景 奖励
第 1 猜解出 +1.550
第 3 猜解出 +1.375
第 6 猜解出 +1.069
六次猜测后失败 -0.427
从未做出有效猜测 -1.500

现在更快的解出总是胜过更慢的解出,每次解出都胜过每次失败,而从未游戏是最差的结果。我们的第一次修复使用了单一的 0.15 权重,导致第 3 猜解出且满塑形(+1.45)的得分超过了第 2 猜解出但塑形最小(+1.25)的得分。相邻猜测次数之间的差距是 0.1,因此把解出权重缩小到 0.04 就消除了重叠,而最差的解出(+1.01)仍高于 1.0 的成功阈值。塑形依然重要,因为 GRPO 是从组内的奖励方差中学习的,所以两局在线索一致性上不同的失败游戏仍然能教会模型一些东西。正如 SageMaker AI 文档所建议的:从基于结果的奖励开始,少量添加塑形,并在真实轨迹上检查奖励。

深入探讨:出了什么问题,按你会遇到的顺序

其中每一项都至少耗费了一次训练。

  • 并发下的 Bearer token。每次调用都用新的 botocore 会话铸造 token,导致在 32 个并发 rollout 时凭证提供者被耗尽,而缓存 token 一小时又会在凭证轮换时返回 403 错误。修复方案是使用一个自刷新的会话、一个短缓存,以及在收到 403 时重新签名。
  • 从未到达策略模型的 rollout 会导致整个任务失败。当一次 rollout 中的所有模型调用都未能通过身份验证时,服务会以“No sampling requests were received.”结束该任务。返回错误状态同样会导致任务失败,所以一定要返回一个奖励。模型被采样但从未做出有效猜测的游戏是没有问题的:它得分为 -1.5,在约 105,000 次 rollout 的 100 步运行中有 51 次这种情况。
  • 工具调用退化为纯文本。当学习率过高时,模型把工具调用写成文本,每次 rollout 都得到相同的保底奖励,梯度随之消失。现在智能体会从原始文本中恢复猜测,并施加一个较小的格式惩罚。
  • 平台期。训练奖励在大约第 15 步后停止上升,验证表现在大约第 30 步后不再改善,不过在 temperature 0 下验证结果低估了模型(第 7 步)。我们接下来会尝试的方案是:用之前的猜测来初始化游戏的课程学习,或者将每个单词的 rollout 数从默认的 group_size 8 增加到更多。

结论

SageMaker AI 上的 MTRL 基于整个回合(episode)的结果来训练模型。这非常适合智能体任务,因为正确的轨迹是未知的,只有结果可以被评分。我们使用了单一的回合级奖励。当中间步骤也可以评分时,MTRL 还可以通过 UpdateReward API 接受每个回合(turn)一个奖励。在这篇文章中,一个开源权重的 20B 模型在服务设置下将解决率从 61% 提升到了 84%,训练 token 花费约 339 美元,AgentCore 运行时使用费约 79 美元。训练曲线报告的 37% 是训练任务对验证进行评分方式的产物,因此请在你将要部署服务的设置下进行评估。

同样的流水线也适用于公共模型从未训练过的内部流程、工具和工作流。

要开始使用,请克隆示例代码库,按照前面的步骤操作,并在投入 GPU 时间之前先运行第 3 步中的冒烟测试。有关底层 API,请参阅 SageMaker AI MTRL 文档。


关于作者

来源: AWS Machine Learning Blog · aws.amazon.com