←新闻详情

Amazon Science更新于 10月12日 19:19

Amazon Science 发布 Rust 代理 Turnstile,为智能体强化学习捕获 token ID

Capturing token IDs during agentic interactions for better reinforcement learning - Amazon Science

AI 导读

Amazon Science 发布用 Rust 编写的代理 Turnstile,位于智能体 harness 与模型后端之间,在生成时刻记录模型采样的精确 token ID、逐 token 对数概率和 loss mask,并导出与框架无关的轨迹供现有 RL 训练栈使用。Turnstile 兼容 OpenAI Chat Completions API,harness 无需改动,后端目前为 SGLang,计划支持 vLLM。在验证中,一个纯文本编码智能体和一个多模态计算机使用智能体在 RL 训练中持续提升。

正文

当前显示原文,中文全文译文尚未就绪。

A new Rust proxy called Turnstile sits between the model backend and the agent harness to capture information lost in mere text transcripts.

Reinforcement learning (RL) is one of the techniques we use to make language models better at sustained, multistep tasks like writing code, navigating a website, or carrying out a research workflow. The model doesn't act alone in those settings; it's wrapped in a piece of software we call a harness, which lets it call tools, observe the results of using them, and decide what to do next. To improve such a model with RL, we let it attempt many tasks inside the harness, score how well each attempt went, and use the score to nudge the model's parameters toward the choices that worked.

The hard part turns out to be the bookkeeping. To turn a scored attempt into a parameter update, the trainer needs an exact record of what the model produced — not a summary, and not a transcript that looks as though it captures a complete exchange but drops vital information.

Internally, models see text as a sequence of numbered units called tokens; an English sentence might have 10 or 20 tokens, each assigned an integer ID by a piece of software called a tokenizer. Two strings that look identical in a transcript can map to different token IDs after a small change of formatting, and that gap, however small, is enough to make the trainer optimize against a slightly different past than the one the model actually experienced.

Today we're releasing Turnstile, a small proxy written in the Rust programming language that sits between any agent harness and the backend system that runs the model. Turnstile records the exact token-level history of every request as it happens, at the only point where that history is unambiguously correct: the moment of generation. It then exports a generic, framework-neutral trajectory that can feed into whichever RL training stack you already use.

We use Turnstile to drive real RL training runs. In the validations we report here, two different agents — a text-only coding agent and a multimodal computer-use agent — improved steadily over the course of their RL runs. In both cases the agent harness was left unchanged, and the data Turnstile records flowed directly into the training stack and produced the expected learning signal end to end.

Tokens, rollouts, and why agent transcripts can lie

Three pieces of vocabulary do most of the work in the rest of this post, so it's worth grounding them up front.

A tokenizer is a deterministic function that turns text into a list of integer token IDs and token IDs back into text. Each model is paired with a specific tokenizer; you can't substitute one for another. The tokenizer is unforgiving: a stray space, a different way of writing a tool call as JSON, or a slightly different chat template (the format string a serving system uses to wrap roles and messages into a single text input the model can read) can change the token IDs even when the rendered text looks the same to a human. We'll call this kind of mismatch retokenization drift when it comes from rerunning a tokenizer over text we’ve already seen and chat template drift when it comes from the surrounding format changing under us.

A rollout is one recorded attempt at a task: the prompt, every tool call, tool feedback, the model’s responses, and the final outcome. The version of the model that produced the rollout is called the behavior policy. The mathematics of policy-gradient RL works cleanly only when the trainer optimizes the model's behavior against the context the behavior policy actually saw. If we rerender the prompt and end up with a slightly different token sequence, we're now training the model against a context unfamiliar to the behavior policy. The training signal degrades, sometimes invisibly, since the model still appears to be learning.

That is why agent harnesses make the problem worse rather than better. A harness is not a static prompt; during a single rollout it may compact older messages to save context, retry a malformed tool call, branch into subagents, merge their results back, or summarize history. All of that is normal, useful agent behavior. But each rewrite is another chance for the next request's token sequence to drift away from what the model actually generated last turn. The transcript a harness produces is a faithful record of the conversation; it is not, in general, a faithful record of the tokens, and it is the tokens the trainer needs.

Capturing tokens at the proxy boundary

Turnstile's central design choice is to stop trying to reconstruct token-level state from text after the rollout is over. We capture it at the moment of generation, where it is already correct, and we do that without changing the harness.

The proxy speaks the same HTTP API every modern agent harness already speaks: the OpenAI Chat Completions API, which has become the de facto standard for "send a list of messages, get back a response." The harness creates a rollout group with Turnstile, points its Chat Completions client at Turnstile's address instead of the real backend, and runs unchanged. Behind the scenes, every request flows through Turnstile to the inference backend (today SGLang, with vLLM planned). Turnstile records the exact token IDs the model sampled, the per-token log probabilities (the model's own confidence values for each token, expressed as logarithms; the trainer needs these to compute its update), and a loss mask that marks which tokens were generated by the model and should contribute to training versus which came from the user, tools, or the system prompt and should not.

来源

Amazon Science