一、业务背景：reward 在抖，问题往往不在奖励函数
我在训一个会自己查天气、自己调用计算器的多轮 tool-use agent。数据是人工标过的多轮对话：每一轮模型要决定要不要调工具、调哪个、参数怎么填；工具返回一段 JSON 之后，再把这段观测拼回上下文，继续生成，直到给出最终答案。听起来像标准的「ReAct / tool calling + RL」故事，真正上手之后最烦人的不是接口格式，而是训练信号对不齐。
跑了几个 epoch，reward 曲线抖得很厉害，策略几乎学不动。排查顺序其实很普通：先看 JSONL 有没有角色标错、有没有缺 tool 轮次；再看规则奖励——有没有合法 <tool_call>、答案里有没有 ground truth、长度是否离谱——这些都「看起来正常」。奖励函数再怎么加减项，曲线也不稳。最后才去翻训练代码：每个 token 的 loss 到底乘了什么 mask。事后回想，这个顺序本身也值得记一笔——Agent RL 出问题时，直觉总是先怪奖励，其实更该先怀疑「你到底在对哪些 token 求梯度」。
问题本质可以压成一句话。多轮轨迹里混着两类 token：一类是模型自己生成的（assistant 话术、<tool_call> 片段）；一类是环境塞回来的（工具返回、以后还有检索块）。模型对环境观测没有控制权——它不能「少生成一点天气 JSON」来提高 reward。如果 loss 把这两类都当成 next-token 优化目标，等于在策略梯度里灌噪声。你以为在优化「会不会调用工具」，实际有一半梯度在强迫模型去模仿一段它根本不该负责的环境文本。
所以真正要解决的，不是再发明一个更炫的奖励项，而是把训练目标说清楚：环境 token 必须留在上下文里给模型看，但不能进策略梯度。 下面写的 env_mask，就是为这件事单独开的一维标记。
补一句定位：这里只把 Agent RL 里最容易讲错、也最容易写错的那一层——环境观测与策略目标的边界——讲清楚、讲落地、可核对。
二、技术方案：env_mask 是独立于 attention 的一维
2.1 语义约定
在 agent/base_loop.py 文件头，约定写得很短：
env_mask:
1 = 模型生成 token（参与策略梯度与 KL）
0 = 环境观测 token（tool 结果、检索块等）
很多人第一反应是：「这不就是 attention mask 吗？」不是。Attention mask 回答的是「谁能看见谁」：工具返回必须被看见，否则下一轮生成没有依据。env_mask 回答的是「谁可以当训练目标」。环境观测对前向是上下文，对反向是禁区。两件事正交，硬塞进同一个 mask，要么看不见工具结果，要么继续污染梯度——都会错。把这个区分讲给别人听时，我通常会补一句：如果只能记住一个设计点，就记住「看见 ≠ 学习」。
2.2 打标发生在 rollout，而不是事后猜
实际打标在 agent/tool_loop.py 的 ToolAgentLoop.rollout_single。模型对当前上下文 generate 完一段，对应 token 全部标 1，同时记下 old_logps；若解析出 tool call，执行 mock/真实工具，再用 chat template 算出「相对上一轮 messages 新增的观测 token」，经 agent/context_utils.py 的 append_env_tokens 追加进 response_ids，并标 0：
# tool_loop：模型段
response_ids.extend(new_ids)
env_mask.extend([1] * len(new_ids))
old_logps.extend(new_logps)
# append_env_tokens：环境段
response_ids = response_ids + obs
env_mask = env_mask + [0] * len(obs)
old_logps = old_logps + [0.0] * len(obs)
环境段的 old_logps 目前补 0：最后乘 completion_mask 时这些位置不进梯度，但若训练循环先算 ratio = exp(new_logp − old_logp)、再乘 mask，old_logp=0 会在中间冒出 exp(new_logp) 这种量级怪异的值，存在数值风险。更稳妥的做法是环境段直接补成与当前 new_logp 相同（令 ratio=1），再靠 completion_mask 挡梯度。关键点仍是：mask 与轨迹同源生成，而不是训练前用启发式从文本里反推「哪一段像工具输出」——后者一改 template 就碎。
2.3 Packing：保 prompt，裁 response 尾
多轮轨迹进 GRPO/CISPO 之前，走 pack_rollouts_for_training：把每条样本的 prompt_ids + response_ids 拼成 batch，并把 env_mask、old_logps 对齐到同一长度。上下文一长就超 max_total_len。截断策略写在 _fit_rollout_to_max_len：
- 默认：保留完整 prompt 前缀，只裁 response 尾部；
- 仅当 prompt 本身已超过上限：对 prompt 做左截（保留靠近回复的一侧），并丢掉整个 response。
prompt_len 直接用「保留下来的 prompt 长度」，不从 env_mask 里重新推断边界。这听起来像实现细节，实际决定了后面 advantage 打在哪些 token 上——第三节会展开旧右裁怎么把 credit 弄脏。
2.4 训练：loss 乘 completion_mask
在 trainer/train_agent.py，先用 completion_mask_from_env_mask 把全序列 env_mask 右移一位，对齐「预测下一个 token」的位置，并且只在 mask=1 的区间里找 EOS（避免环境 JSON 里偶然出现的 eos id 提前掐断）。然后对 per-token 的 GRPO/CISPO loss 做：
policy_loss = (
((per_token_loss * completion_mask).sum(dim=1)[valid_rows]
/ token_counts[valid_rows].clamp(min=1)).mean()
if valid_rows.any() else per_token_loss.sum() * 0.0
)
实现上还有一个容易忽略的细节：prompt 段在 env_mask 里也是 0。拼装时是「prompt 全 0 + response 段按生成/环境打标」。右移成 completion_mask 之后，prompt 内部的 next-token 位置同样不会进策略损失；真正开始有 1 的地方，对应「预测第一个模型生成 token」及之后的可控段。这样不需要再维护一套单独的 prompt mask，语义保持单一。
2.5 示意图
图 1：序列上的 env_mask（1=模型可控进 loss，0=环境观测只当上下文），以及用于 next-token 的 completion_mask = env_mask[:, 1:]。
图 2：训练路径——打标 → packing → completion_mask → loss × mask。
2.6 Tool-Use 与 Search 共用语义
同一套 1/0 语义用在 search_loop：检索回来的 ... 标 0，模型写的 等标 1。train_search_agent.py 复用同一条 packing / align_rewards_to_packing 路径。这样不是「为 Search 再写一套框架」，而是换一种环境观测源。对读者而言，这也降低了认知负担：先搞懂 Tool-Use 上的 mask，Search 只是把「工具 JSON」换成「检索块」，训练侧几乎不用换脑子。
tensor 侧的 padding / info_mask 拼接，参考了 RL-RAG的做法，完整版权链写在 agent/tensor_helper.py 与 README Acknowledgements。Agent loop、packing、训练脚本是本项目原创。
三、踩坑记录：两个会把训练信号弄脏的细节
踩坑这一节比公式值钱，因为它们都是「看起来能跑、实际在学错东西」。
3.1 Chat template delta：改模板后的静默错位
多轮里每追加一段工具结果，都要知道「新增了哪些 token」。实现上用 encode_chat_template_delta：对 messages 变化前后各做一次 apply_chat_template + tokenize，再取后缀差。
有一阵子一切正常，直到我改了 chat template，给 assistant 侧加了 thinking 相关结构。改完之后 reward 突然随机抖，有时还会掉到不合理的负向。数据没动，奖励函数没动，唯一变量是 template。加日志后看到：前后两次 tokenize 的公共前缀已经对不上了，但旧代码完全不报警，直接按长度差截后段。截出来的「增量」并不是真正的工具观测，却被标成 env_mask=0 拼进轨迹——等于用错误的边界在训练。
现在的行为：前缀不一致时打 warning，并提供 strict=True 直接抛错（单测 / 排查用）。length-based fallback 仍保留，因为 template 稳定时它够用、也快。但前提必须说破——增量 tokenize 默认假设 chat template 不会改写历史 span。Agent 场景里这个假设经常不成立：thinking 开关、tools 字段、add_generation_prompt 都会让「同一条 messages」的 token 前缀漂移。后续更稳的方向是 LCS 对齐，或者放弃两次全量 tokenize 求差，改为显式构造观测块再 tokenize。
3.2 Packing 右裁：截断落在回复中段时的 credit 错位
旧 packing 大致是：对整段做右截到 max_total_len，再试图从 mask 反推 prompt 边界。一旦截断点落在 assistant 轨迹中间，prompt_len 和真实 prompt 错位，会出现两类灾难：
- 梯度打在「当前 batch 里根本没作为目标的位置」，或边界 token 语义被拧歪；
- reward 仍按截断前的完整 completion 算，advantage 给没进 loss 的尾部记功——策略更新与评价对象不是同一条轨迹。
现在 _fit_rollout_to_max_len 改成保 prompt、裁 response 尾；align_rewards_to_packing 用 packing 后模型真正看见的 response 重解码，作为奖励函数输入，截断样本把 unfinished 置为真。这样 loss 看见的轨迹 与 reward 评价的轨迹 对齐。这是整条 Agent RL 链路上最有含金量的正确性修复之一：它不新增算法名词，但直接决定你「在优化什么」。
把两个坑放在一起看，共性是：训练信号的「时间轴」和「文本轴」一旦错位，曲线会假装你在学习。 Delta 错位污染的是 mask 边界；packing 错位污染的是 prompt/response 边界与 reward 对象。修法也都偏「工程正确性」：报警、严格模式、改截断策略、强制 reward 跟 packed 文本走——而不是换一个更大的基座模型指望问题消失。
四、数据指标：只写测过的
机器：NVIDIA GeForce RTX 3090；环境 PyTorch 2.5.1+cu121。数字来自 Full 包 smoke（同目录 LiteAgentRL_博客实测数据.md；本稿引用未改数）。
| 场景 | 墙钟 | 峰值显存 | 模型配置 |
|---|---|---|---|
| Phase2：generate_smoke_datasets + Pretrain/SFT/DPO/GRPO | 31.624 s | 396 MiB | tiny 0.533 M（hidden=64，2 层） |
| Phase4：Tool-Use Agent RL（train_agent，走 env_mask） | 7.985 s | 418 MiB | 同上 |
说明：表格里的 0.533 M 是 smoke 用的 tiny 配置（hidden=64，2 层）；正文写的博客实测默认配置参数量是 63.912 M（LiteConfig，hidden=768，8 层）。两者不是同一套模型，勿混读。
默认 LiteConfig()（hidden=768，8 层，vocab=6400，GQA 8/4 heads）参数量按 sum(p.numel()) 实测为 63.912 M。仓库里仍有「26M-class」的习惯称呼，与当前 numel 合计不一致；写进博客时以实测为准，避免把营销口径当成物理参数量。
必须划清边界：这些数字证明的是——在极小模型上，整条 env_mask 训练链路可以在几十秒、约四百兆显存内跑通；它们不证明正式语料上的收敛速度或最终 reward。本轮也没有做「关闭 env_mask vs 打开」的对照实验：仓库没有一键开关，我也没有伪造前后曲线。缺数据就标明未测，比编一个漂亮对比更负责。Phase5–7 全链路总时长、768×8 正式训练峰值显存，同样未在本稿引用范围内；
五、AI 协作、开源边界，以及你怎么复现
开发过程中用了 AI 辅助：读跨文件调用、起草单测、帮我盯边界条件。举一个具体例子：completion_mask_from_env_mask 截断 EOS 时，必须只在模型生成段（mask=1）里找，否则环境返回的文本里一旦出现与 eos 相同的 id，轨迹会被提前掐断，后面的 tool/answer 全部失踪。这个点是在把需求讲清楚时被提醒到的；落地实现、补测试、看 phase3 用例是否覆盖，是自己改到绿的。另一个协作痕迹是 packing 单测：用极短假序列断言「超长时 prompt 前缀仍在、response 被裁尾、env_mask 与 ids 等长」——AI 可以生成第一版断言，但我改过输入样例，确保测的是「保 prompt」而不是「随便裁一段还能 shape 对齐」。对 AI 的健康用法大概就是这样——加速搜索与草稿，不外包判断。
开源策略和这篇博客对齐：对外 Core 包放的是机制与可跑脚本，不是本机全部资产。包含 agent/（env_mask、packing、Tool/Search loop、tensor_helper 版权声明）、trainer/、model/、services/、smoke/数据生成脚本与测试；故意不放 *.pth、样本 .jsonl、.faiss。
python scripts/generate_smoke_datasets.py
python scripts/run_phase2_smoke.py
python scripts/run_phase4_smoke.py
就可以在本地复现文中的链路验证。要跑更长的生产数据与权重，需要自己准备语料。
GitHub：【待补】
仓库 README 与本文发布后会互相挂链接，方便从文章跳进代码、从代码跳回设计动机。想对照实现，优先读：
-
agent/base_loop.py— 语义、packing、align_rewards_to_packing -
agent/tool_loop.py / agent/context_utils.py— 打标与 delta -
trainer/train_agent.py— completion_mask 进 loss
若你只想快速验证「我这台机器也能跑通」，用上面三条 smoke 命令即可；若你想改语义（例如把某一类特殊 token 也标成环境），从 append_env_tokens 与 completion_mask_from_env_mask 两处下手通常最短。
小结
多轮 Tool Agent「训不动」时，先问一句：loss 有没有在优化模型根本控制不了的 token？env_mask 把「可见」和「可学」拆开；packing 保证截断之后轨迹与 reward 仍指向同一段文本；delta tokenize 则提醒你——模板一变，静默错位比显式报错更危险。3090 上 Phase2 31.624 s / 396 MiB、Phase4 7.985 s / 418 MiB，只说明机制可复现、smoke 够轻；真正的业务效果，仍然取决于你的数据、奖励与更长的训练曲线。先把目标函数写对，再谈算法名词，会省掉很多假忙。若这篇对你排查 Agent RL 有一点帮助，欢迎对照仓库里的三个入口文件自己改一版 mask 语义，跑一遍本地的 phase3 与 phase4 相关测试即可——那才是这篇博客想留下的东西。
