Agentic RL中的Token-In Token-Out原则：如何正确实现多轮工具调用训练
在强化学习（Reinforcement Learning, RL）框架下训练大语言模型（Large Language Model, LLM）已成为构建高级智能体的核心路径。当任务局限于单轮交互时，训练过程通常稳定高效：损失曲线平滑下降，奖励信号合理，模型快速收敛。然而，一旦将场景扩展至支持工具调用的多轮代理（agentic）设置——即模型可在推理过程中动态调用外部函数（如计算器、代码解释器或API）并基于返回结果继续生成——许多开发者会遭遇意想不到的训练崩溃：损失值无故飙升、梯度爆炸，甚至出现张量形状不匹配的错误。这些现象的背后，往往隐藏着一个被广泛忽视但至关重要的原则性错误：违反了Token-In, Token-Out（TITO）不变量。
TITO原则的本质在于：强化学习的梯度更新必须严格作用于模型实际采样生成的令牌序列之上，而非任何经过重新编码的近似版本。这一要求在单轮场景中天然满足，但在多轮工具调用循环中极易被破坏。典型的错误做法是：在每一轮模型生成后，先将输出令牌解码为文本，解析其中是否包含工具调用；若存在，则执行工具并将结果以新消息形式追加到对话历史中；随后，在计算最终损失前，将整个更新后的对话历史重新通过聊天模板（chat template）和分词器（tokenizer）进行编码。这种看似自然的流程，实则埋下了致命隐患。
问题的根源在于分词过程的非可逆性。现代LLM普遍采用字节对编码（Byte Pair Encoding, BPE）或其变体作为分词策略。BPE虽能为给定字符串提供一种贪心的、规范的分词方式，但同一字符串可能存在多种合法的令牌序列表示。例如，单词“hello”可能被编码为单个令牌[hello]，也可能被拆分为[he][llo]，两者解码后均得到相同文本。更复杂的情况出现在结构化输出中：JSON格式的工具调用参数可能因空格、换行、键顺序或布尔值大小写（true vs True）的不同而产生不同的令牌序列。因此，当你对模型生成的令牌序列进行解码再重新编码时，极有可能得到一个与原始采样序列不同的整数ID列表。此时，若将梯度反向传播至这个“新”序列，就等同于在优化一个模型从未见过的分布，导致训练信号完全失真。
这种错误实现方式可称为“Message-In, Token-Out”（MITO）循环。它不仅引发令牌漂移，还带来第二个技术难题：损失掩码（loss mask）的重建困难。在RL训练中，我们仅希望对模型自身生成的令牌计算损失，而忽略用户输入或工具返回的内容。MITO循环在最后一步才对完整对话进行一次性编码，丢失了各轮次的边界信息。为了区分哪些令牌属于模型输出，不得不编写复杂的解析逻辑，根据聊天模板中的角色标记（如<|im_start|>assistant）反向推断每个令牌的归属。这不仅增加工程复杂度，且不同模型家族的模板差异巨大，需为每个模型定制解析器，维护成本高昂。
要彻底解决上述问题，关键在于重构训练循环，严格遵循TITO Done Right的核心准则：永不重新编码已解码的令牌。具体实现上，应维护一个全局的令牌缓冲区（token buffer），该缓冲区直接累积模型在每一步采样得到的原始令牌ID。对话历史（messages list）仅作为辅助元数据用于逻辑判断，绝不参与最终的令牌序列构建。当模型生成一段输出后，我们将其令牌ID追加至缓冲区，同时对其进行解码以检测是否存在工具调用。若检测到有效调用，则执行工具获取结果；但关键在于，工具结果不会以“消息”形式重新渲染整个对话，而是通过一种增量方式转换为令牌并追加到缓冲区末尾。
这种增量转换依赖于聊天模板的一个关键属性：前缀保持性（prefix preservation）。具体而言，对于任意以助手工具调用结尾的对话历史，当追加一个工具返回消息后，新对话的令牌化结果必须以前一状态的令牌序列作为严格前缀。数学表达为：render(messages + [tool_response])[:len(render(messages))] == render(messages)。若此条件成立，则可通过“差分法”精确计算出工具响应所对应的令牌增量（delta）：分别对包含与不包含工具响应的两个虚拟对话进行令牌化，取后者减去前者的后缀部分，即为需追加的令牌序列。
令人欣慰的是，对当前主流开源模型家族的实证测试表明，前缀保持性在工具消息场景下几乎是一种默认行为。包括Qwen2.5/3.x系列、Llama 3.1/3.2/4、Gemma 4、Function Gemma、GLM-4.5/5、DeepSeek-V3.1/R1等在内的十八个模型家族均天然满足该属性。唯一的例外是早期版本的Qwen3，其模板在推理时会根据是否为最后一轮动态插入<think>标签，导致前缀不一致。但修复方案极为简单：修改Jinja模板中的一行条件判断，强制始终渲染思考块即可。这一发现具有重大实践意义——它意味着开发者无需为每个模型手动实现复杂的渲染器（renderer），只需验证其前缀保持性，即可安全采用TITO循环。
当然，现实场景中仍存在挑战TITO原则的边缘情况。首先是**历史重写（history rewriting）**机制，如Z.ai的clear_thinking模式会移除除最后一轮外的所有思考块，或长程编码代理在接近上下文上限时对历史进行摘要压缩。这类操作本质上修改了模型后续步骤所依赖的上下文，使其与原始采样轨迹不一致，直接破坏了RL目标函数的定义基础。对此，合理的应对策略是：将最后一次历史重写点之前的所有内容视为固定提示（prompt），对其施加零损失掩码；仅对重写点之后由模型真实采样的令牌进行梯度更新。虽然这会缩短有效训练序列的长度，但保证了剩余部分的数学正确性。
其次是**序列截断（truncation）**问题。当生成过程因达到最大序列长度而提前终止时，输出可能缺少闭合标记（如<|im_end|>）。对于依赖完整结构进行桥接的渲染器方案，这会导致无法安全扩展序列，被迫回退到危险的重新渲染。但在TITO框架下，截断并非问题：缓冲区忠实地保存了模型实际生成的不完整令牌序列，损失掩码仅覆盖这些真实存在的令牌，无需关心其是否构成语法完整的结构。解析器在遇到不完整工具调用时自然会跳过执行，符合资源耗尽的实际情况。
综上所述，构建稳健的Agentic RL训练系统，并不需要为每个模型家族重复造轮子式地实现专用渲染器。只要坚持“永不重新编码已解码令牌”这一黄金法则，并确保所用聊天模板满足针对工具消息的前缀保持性（该条件在实践中极易满足），即可通过简洁的增量令牌追加机制，实现数学上严谨、工程上高效的多轮工具调用训练。这种方法不仅规避了令牌漂移和损失掩码错位两大陷阱，还将模板相关的复杂性降至最低，为大规模智能体强化学习提供了可靠的基础架构。