Capturing token IDs during agentic interactions for better reinforcement learning
TL;DR · AI 摘要
亚马逊开源工具Turnstile通过精确记录token ID提升强化学习训练效果,验证显示两类智能体训练效率提升。
核心要点
- Turnstile工具可精确记录生成时的token级历史数据
- 文本微小格式变化会导致token ID差异影响训练效果
- 开源项目已验证文本和多模态智能体训练效果提升
结构提纲
按章节快速跳转。
思维导图
用一张图看清主题之间的关系。
查看大纲文本(无障碍 / 无 JS 友好)
- 强化学习训练优化
- 核心挑战
- token ID敏感性
- 解决方案
- Turnstile工具
- 精确记录机制
- 验证结果
- 文本智能体提升
- 多模态智能体提升
金句 / Highlights
值得收藏与分享的关键句。
文本微小格式变化会导致token ID差异,影响训练效果。
Turnstile在生成时刻记录无歧义的token级历史数据。
验证显示文本和多模态智能体训练效率均提升。
强化学习(RL)是我们提升语言模型执行复杂多步骤任务(如编写代码、网站导航或科研工作流)能力的重要技术。在这些场景中,模型并非独立运行,而是通过我们称为 _适配器_ 的软件组件进行封装,该组件允许模型调用工具、观察使用结果,并决定下一步操作。通过 RL 优化此类模型时,我们让模型在适配器中尝试大量任务,对每次尝试的表现进行评分,并利用评分引导模型参数向成功的选择靠拢。
难点在于数据记录的精确性。要将评分结果转化为参数更新,训练器需要完整记录模型输出的原始内容——而非摘要,也非看似完整但丢失关键信息的对话转录。
在模型内部,文本被视作称为 _token_ 的编号单元序列;一个英文句子可能包含 10-20 个 token,每个 token 都由称为 _分词器_ 的软件分配唯一整数 ID。即使两个字符串在转录中看起来完全相同,格式的微小变化可能导致不同的 token ID。这种细微差异足以使训练器基于与模型实际经历略有不同的历史进行优化。
今天我们发布 Turnstile,这是一个用 Rust 编程语言编写的轻量级代理程序,位于任何智能体适配器与运行模型的后端系统之间。Turnstile 在生成时刻(唯一能保证历史记录无歧义的时刻)精确记录每个请求的 token 级别历史,随后导出一个与框架无关的通用轨迹数据,可无缝对接您现有的 RL 训练系统。
我们使用 Turnstile 实现真实的 RL 训练运行。在本次报告的验证中,两个不同智能体——一个纯文本编码智能体和一个多模态计算机使用智能体——在 RL 训练过程中持续改进。在两种情况下,智能体适配器均保持不变,Turnstile 记录的数据直接流入训练系统,实现了端到端的预期学习信号。
Token、执行轨迹与智能体转录的潜在偏差
以下三个术语将在后续内容中频繁出现,现进行明确界定:
_分词器_ 是一个确定性函数,用于在文本与整数 token ID 之间进行双向转换。每个模型都绑定特定分词器,不可替换。分词器具有严格性:一个多余的空格、JSON 格式工具调用的写法差异,或 _聊天模板_(服务系统用于将角色和消息封装为模型可读单文本的格式字符串)的微小变化,都可能改变 token ID,即使渲染后的文本对人类观察者看起来完全相同。我们将这种因重新运行分词器导致的不匹配称为 _重分词漂移_,而将因聊天模板格式变化导致的不匹配称为 _模板漂移_。
一个 _执行过程_ 记录了完成任务的一次尝试:包括提示内容、所有工具调用、工具反馈、模型的响应以及最终结果。生成该执行过程的模型版本被称为 _行为策略_。策略梯度强化学习的数学原理只有在训练者针对行为策略实际看到的上下文优化模型行为时才能有效运作。如果我们重新渲染提示内容并得到略微不同的标记序列,此时模型将面对行为策略不熟悉的上下文。训练信号会退化,有时甚至难以察觉,因为模型看起来仍然在学习。
这就是为什么代理框架反而会加剧问题而非改善问题。框架并非静态提示;在单次执行过程中,它可能会压缩旧消息以节省上下文空间、重试格式错误的工具调用、分支到子代理、合并结果后返回或总结历史记录。所有这些都是正常且有用的代理行为。但每次重写都可能让下一次请求的标记序列进一步偏离模型上一轮实际生成的内容。框架生成的转录记录忠实记录了 _对话_,但通常无法忠实记录 _标记_,而训练者需要的是这些标记。
在代理边界捕获标记
Turnstile 的核心设计选择是:在执行过程结束后停止尝试从文本中重建标记级状态。我们选择在生成时立即捕获这些信息,此时状态已经准确,且无需修改框架。
代理使用的是现代代理框架已经普遍支持的相同 HTTP API:OpenAI Chat Completions API,该接口已成为"发送消息列表、获取响应"的事实标准。框架通过 Turnstile 创建执行组时,只需将 Chat Completions 客户端指向 Turnstile 的地址而非真实后端即可,其余部分无需改动。在后台,所有请求都会通过 Turnstile 路由到推理后端(目前是 SGLang,未来计划支持 vLLM)。Turnstile 会记录模型采样的精确标记 ID、每个标记的 _对数概率_(模型对每个标记的置信度值,以对数形式表示;训练者需要这些值来计算更新)以及一个 _损失掩码_,该掩码标记了哪些标记是由模型生成且应参与训练,哪些来自用户、工具或系统提示且不应参与训练。
from openai import OpenAI import turnstile
with turnstile.Proxy( backend_url="http://localhost:30000", model="Qwen/Qwen3-1.7B", ) as proxy: group_id = proxy.create_group(max_trajectory_tokens=4096)
client = OpenAI( base_url=f"http://{proxy.addr}/group/{group_id}/v1", api_key="-", )
client.chat.completions.create( model="Qwen/Qwen3-1.7B", messages=[{"role": "user", "content": "Solve this Python task."}], )
sequences = proxy.get_training_sequences(group_id)
当部署完成时,适配器会向 Turnstile 请求记录的轨迹。每条轨迹是一个 _TrainingSequence_ 对象,包含标记 ID、日志概率、完整序列的损失掩码,以及记录序列中哪些时间段使用了模型权重的哪个版本。(训练器需要这些 _权重版本边界_ 来判断模型参数是否因异步更新而在部署中途发生了变化。)将这些数据转换为训练器所需的特定批次形状是适配器的常规工作:附加奖励值、扩展掩码,然后传递给训练器。
Turnstile 位于代理适配器和推理后端之间,随着请求的流动记录原生标记的部署状态。适配器继续使用普通的聊天补全接口进行交互;训练器接收到标记 ID、日志概率、掩码和权重版本边界。
现有适配器可保持黑盒状态
目前存在许多现成的代理适配器(如 OpenHands、Codex 和 Terminus),它们本身已具备实用性但并非专为训练运行设计。在没有 Turnstile 的情况下,使用这些适配器驱动强化学习训练需要其记录标记 ID、日志概率、掩码和路由追踪;实际上,这类工作通常由训练系统内部构建的适配器形状组件完成。无论哪种方式,适配器最终都充当了标记级的强化学习数据管道,这属于错误的抽象层级。适配器知道它原本打算发送给模型的信息,但通常并不了解模型实际使用的精确标记序列、缓存状态、路由追踪或处理后的多模态输入。这些信息都存储在后端。
通过 Turnstile,生产环境的适配器无需记录训练数据。它只需将聊天补全客户端指向 Turnstile 而非推理后端,其余部分保持不变。Turnstile 记录面向模型的部署状态:它无需理解适配器的私有控制流或上下文变化的语义原因。如果下一个请求是对早期请求的忠实标记级扩展,Turnstile 会将其合并到同一条轨迹中。如果适配器压缩了内存、重写了历史、合并了子代理结果或以无法证明等效的方式修改了前缀,Turnstile 会启动新序列并保持可训练后缀的真实性。所有这些特性都适用于严格的黑盒场景:即使是一个对训练系统封闭的专有适配器,也能在完全无需源码级集成的情况下驱动强化学习训练。
以下是两个使用开源适配器以黑盒方式使用 Turnstile 的示例。
文本模式。使用 OpenHands 作为适配器,在 Mostly Basic Python Problems (MBPP) 数据集上训练 Qwen3-1.7B。图表显示了纯文本编码运行的训练奖励 _(左)_ 和保留评估奖励 _(右)_。
多模态。使用 OSWorld 的标准 PromptAgent 适配器,在屏幕截图驱动的桌面操作上训练 Qwen3-VL-8B 的 OSWorld 计算机使用任务。图表显示了训练分布奖励。每个提示的平均奖励从约 0.2 提升至约 0.71,经过约 165 个部署步骤。
多轮代理与前缀感知轨迹
一种简单粗暴的存储方式是将每个请求视为独立的训练样本。这种方法效率低下,因为每个新请求都包含到目前为止的完整对话,会导致相同 token 字符串被反复重复存储。这种方法还存在细微的错误,因为它会丢失对话轮次之间的关联性。
相反,Turnstile 将多轮对话展开过程存储为一条持续增长的 token 路径——只要该路径忠实反映模型实际看到的内容。当后续请求只是在先前请求的基础上添加少量新 token(如用户新消息、工具结果或 LLM 回应)时,Turnstile 会在 token 层面识别重叠部分——不是通过比较渲染后的字符串,而是通过验证先前记录的 token ID 是否确实出现在新请求的起始位置。如果确认无误,这两个轮次就会合并为一条连续的可训练序列,损失掩码能准确识别哪些部分是模型的输出。
简单存储方式会在每轮对话中重复存储当前对话历史(即“前缀”)。而 Turnstile 在 token 层面前缀不变时会将多轮合并为单一序列,当下次请求的前缀发生变化时则会分叉为新序列。无论哪种情况,可训练的后缀部分始终准确无误。
当后续请求无法安全地从前序请求扩展时(例如由于工具修改了早期消息或 token 因某些原因发生偏移),Turnstile 不会强行伪装成可以扩展。它会启动新的训练序列。我们称这种操作为“轨迹爆炸”。虽然这会消耗比乐观方案更多的训练 token,但能确保用于 RL 训练的每个 token 字符串都是行为策略实际观察到的内容。关键不在于追求紧凑性,而在于绝不向训练器提供虚假的历史信息。
专家混合路由引入了隐藏维度
一些现代模型采用_专家混合_(MoE)架构,其中只有模型参数中的一小部分——称为专家——会针对任何给定输入 token 被激活。专家的选择本身是计算过程的一部分,由每一层的小型_路由网络_决定。路由决策依赖于每一层的激活状态,先前 token 的处理方式出现细微差异可能导致不同专家被选中。这对 RL 至关重要,因为即使两个请求具有相同的 token ID,不同运行时这些 token 可能被路由到不同专家。
去年秋天,北京大学的研究人员及其同事分析了这种差异,并提出应记录 MoE 模型的路由决策,以便训练器能够重放这些决策。我们采用了相同的原则。当启用 MoE 记录时,Turnstile 会向推理后端请求路由轨迹并将其与 token 一并记录。每次扩展 token 路径时,它都会检查共享前缀的路由是否与上一轮记录一致。如果不一致(例如由于键值缓存缺失迫使后端重新计算前缀并选择不同专家),Turnstile 会拆分轨迹而非在错误路由下进行训练。
在专家混合模型中,相同 token ID 可能被路由到不同专家。在此示例中,完整的缓存命中确保了请求一和请求二的路由一致。但在冷启动重新计算后,相同 token 却被路由到不同专家。
多模态展开
当模型同时以图像作为输入时——即视觉-语言模型——rollout需要额外维护一组状态信息。模型无法直接看到上传图像的原始字节,而是接收图像处理器的输出结果。该处理器是一个固定流程,会调整图像尺寸、裁剪图像、将其转换为像素值的张量,并在文本中插入占位符标记以标识图像位置。由于不同版本的处理器、不同的缩放策略或不同的补丁几何形状,相同图像字节可能生成不同的张量,且占位符数量会随着输入维度变化。如果强化学习训练器需要从头重新处理图像,其训练时看到的视觉前缀可能与行为策略实际接收到的视觉前缀不一致。
Turnstile将图像处理视为rollout的一部分。当请求包含图像时,Turnstile会解码图像、哈希并存储原始字节以备审计、运行模型配置的处理器,并将处理后的像素特征按占位符出现的顺序,与token ID一同记录在相同轨迹中。导出的序列同时包含原始和处理后的视觉数据,因此训练器可根据需要选择使用任意一种。
class TrainingSequence: tokens: list[int] logprobs: list[float] segment_info: list[tuple[bool, int]] weight_versions: list[tuple[int, str]] routed_experts: str | None images: list[TrainingImage] processed_images: list[TrainingProcessedImage]
应用方向
Turnstile目前处于早期阶段。当前实现包含Rust核心、SGLang后端、用于进程内训练脚本的Python绑定、前缀感知的多轮捕获、可选的MoE路由捕获以及多模态支持。近期工作将扩展至更多领域:vLLM后端、更多训练框架适配器以及更全面的多模态模型覆盖。长期目标与最初设计保持一致:智能体框架不应被迫成为强化学习数据管道,训练器不应需要通过渲染文本猜测发生了什么。模型生成了token,我们将其完整记录。
参考资料
- OSWorld 及其 开源智能体框架
- OpenHands
- THUDM/slime
- Prime Intellect渲染器博客
- Prime Intellect渲染器GitHub仓库
- prime-rl算法(扩展属性、多轮轨迹合并)
- prime-rl推理(路由重放)
- Polar:在任意框架上规模化智能体强化学习
- NVIDIA-NeMo/ProRL-Agent-Server
- 通过对齐训练与推理路由稳定MoE强化学习
- rLLM
- Agent Lightning
- Strands Agents SGLang提供者
致谢
特别感谢Keagan Long、Daisy Lin、Changlong Yu和Yifei Wang对本工作的贡献。