
AWS以Wordle演示SageMaker AI多轮强化学习,详解奖励、评估与成本
AWS的一份实践指南以Wordle为例,介绍Amazon SageMaker AI上的多轮强化学习,即MTRL。模型根据多步骤任务产生的奖励接受训练。作者报告,GPT-OSS-20B在预留测试词上的解题率有所提高,同时记录了奖励设计、运行故障和成本方面的经验。这些数字来自指南中的具体实验,不能据此确定其他任务、模型或部署设置的表现。
模型、代理与游戏如何协作
Wordle要求代理在六次尝试内猜出一个隐藏的五字母单词。在该示例中,SageMaker AI负责训练任务和策略模型,Amazon Bedrock AgentCore则承载Strands代理、游戏环境及奖励函数。代理调用SageMaker Job Runtime提供的OpenAI兼容端点,再通过guess_word工具获得反馈。因此,一局游戏可以包含多次模型响应和工具调用,而奖励用于评价这些步骤构成的完整过程。
训练使用GRPO比较同一个目标词的多次运行所获奖励,并据此更新LoRA适配器。示例在每局结束后返回一个奖励;MTRL也允许通过UpdateReward提交逐轮奖励。整局奖励是这份示例的选择,并非指南所述接口唯一接受的奖励方式。
代理还可以运行在客户的EKS、EC2或Fargate等基础设施上,通过AWS Lambda转发函数接入。指南将MTRL描述为无服务器服务,由服务配置训练和采样基础设施。因此,案例把模型训练与代理运行分别交给SageMaker AI和AgentCore,同时保留了代理运行在其他客户基础设施上的接入方式。
数据提供任务,不提供示范答案
MTRL数据行包含任务,不需要目标回答、对话记录或标签。服务会把prompt列,或第一列的内容,原样传给代理。示例把指令、隐藏答案及追踪标识编码在任务中,但策略模型始终看不到隐藏答案。这里需要区分代理接收的任务数据与模型实际可见的信息:答案用于环境处理,并未作为解题提示展示给模型。
示例准备了640个训练词,以及128个与训练集不重叠的验证词。这样的划分使指南能够分别报告训练和预留词上的评估结果。AWS还表示,数据集、代理和训练后的模型都留在客户自己的AWS账户中,不会用于训练基础模型。这是指南对数据处理的说明,并非独立审计结论。
启动程序使用MultiTurnRLTrainer,并要求具备MLflow应用。配置将采样上限设为8192个token,最多训练20个epoch、100步。每批包含128个任务,每个任务运行八次,因此每个epoch有五步,每步执行1024局游戏。这些参数说明了该示例如何从任务数量组织出每步的游戏数量。
指南特别提醒,训练期间重新部署代理会重启运行时,并可能导致整个训练任务失败。代理部署因此也是实验运行条件的一部分;即使训练由服务管理,正在承接游戏请求的代理运行时仍会影响任务能否完成。
奖励不能让失败超过成功
作者最初采用累加每次猜测奖励的设计,但这种方式出现了与游戏目标不一致的排序:一局失败仍可能获得约+0.49,而第三次才猜中的得分可能高于第一次就猜中。这个问题涉及的不只是奖励数值大小,而是奖励是否正确表达更快成功、较慢成功和失败之间的优先关系。
修订后的函数以最终结果项为主,再加入有界的辅助奖励。辅助部分被限制在−1到1之间,成功时的权重为0.04,失败时为0.15。按照作者的说明,这样可以保持快速解出、较慢解出、失败、没有有效猜测的顺序,避免辅助奖励改变主要结果之间的高低关系。
最低分的成功也能获得+1.01,高于服务设定的1.0成功阈值。因而,奖励函数既要与游戏的实际结果一致,也要与服务用于识别成功的阈值相容。指南中的修改针对的正是这两个具体要求,而不是单纯增加猜测过程中的奖励。
评估设置决定如何理解提升
作者报告,在中等推理投入、温度1.0的条件下,预留词解题率从61%升至84%。其中84%是两次运行的平均值,两次分别达到80%和88%。在低推理投入下,解题率从15%升至69%。这些结果分别对应所述推理设置,不能把不同条件下的数字合并成一个不带条件的模型能力指标。
训练过程中的验证采用温度0,其解题率最高为48%,结束时为37%。这与温度1.0下报告的评估数字采用不同设置,需要分开理解。作者没有测试继续增加训练步数,是否能改善实际服务设置下的结果,因此指南也没有证明更多训练必然带来进一步提升。
推理投入还影响采样消耗。作者称,与中等投入相比,低投入每局评估使用的采样token减少了63%至83%。结合两种投入下分别报告的解题率,读者可以看到该实验同时记录了效果和消耗;这些比例仍限于所报告的评估游戏,不能直接视为其他任务的节省幅度。
MLflow记录奖励、token和轮次统计;CloudWatch中的OpenTelemetry追踪则记录模型轮次与工具调用。两类记录分别呈现训练统计和游戏交互过程,为查看奖励变化及具体调用提供依据。指南中的运行经验也说明,仅有最终解题率不足以描述整个实验过程。
运行故障与格式问题
在32个并发运行下,作者遇到了凭证提供程序无法及时提供凭证的问题;将令牌缓存一小时又造成身份验证失败。其修复方案采用可自行刷新的会话、较短的缓存时间,并在收到403响应后重新签名。指南还指出,若运行根本没有到达策略模型,可能导致整个训练任务失败。
已经执行采样却没有有效猜测的游戏则获得−1.5奖励,100步实验中共出现51局。这与未到达模型的运行故障属于不同处理情形。作者还发现,过高的学习率会使工具调用变成普通文本,因此加入猜测恢复,并对格式问题施加惩罚。案例由此同时记录了基础设施故障和模型输出格式偏离两方面的问题。
费用数字与结束后的资源
这次100步训练运行了12小时40分钟,共完成105486次运行,消耗4.916亿个prefill token、1.597亿个采样token和6.433亿个训练token。作者按照当时us-east-1区域的标价,报告训练token费用约339美元,AgentCore运行时使用费用约79美元。这些是所述项目的历史费用,不能作为当前报价,也不能直接推算其他实验的总成本。
评估系列的token费用约为1.20美元,总用时约一小时;每次包含128个词的评估耗时四至七分钟。训练与评估的费用和耗时分别对应各自的运行,阅读时应保留这一范围,不能把评估系列的低费用误认为完整训练的费用。
训练和评估任务完成后会停止计费,但AgentCore运行时、其IAM角色、MLflow应用和S3对象仍然保留。指南表示,训练角色及模型包本身不计费。因此,任务完成与相关资源是否仍存在是两个需要分别查看的状态。
上述结果、技术经验和费用均来自这份AWS实践指南,没有独立验证来源。该案例展示了作者如何组织一次Wordle多轮强化学习实验,以及如何解释其奖励、评估和运行记录;它不保证其他任务能够取得相同提升或相同成本。