05-ReTool 源码精读#

目标：沿一条数学题轨迹，追踪它如何从原始数据变成 prompt、模型动作、工具 observation、reward、advantage、PPO 训练张量，最后进入评测与分析。

阅读原则：每章都回答四个问题——当前对象是什么、边界在哪里、状态如何变化、下一模块接收什么。

阅读进度#

• 第一章：prepare_data.py——原始数据如何变成训练 JSONL
• 第二章：data.py——MathExample 与循环 batch
• 第三章：protocol.py——消息、文本与 token 边界
• 第四章：sandbox.py——代码调用如何变成 observation
• 第五章：reward.py——最终答案与稀疏奖励
• 第六章：rollout.py——多轮轨迹状态机
• 第七章：train.py——轨迹到 PPO 张量
• 第八章：eval.py——Base/checkpoint 统一评测
• 第九章：analysis.py——汇总指标与作图

第一章：prepare_data.py——原始数据如何变成训练 JSONL#

源码位置：https://github.com/KMnO4-zx/agentic-rl-lab/blob/main/05-retool/prepare_data.py

1.1 先抓住本章主线#

这一章只做数据结构清洗，还没有模型采样，也没有 token 和张量。

Hugging Face DAPO-Math-17k
        │
        │ 每行是一个嵌套 dict
        ▼
normalize_row()
        │
        │ 提取 question / answer，清理题目模板
        ▼
标准 Python dict
{id, question, answer, data_source}
        │
        ├── 固定 seed 打乱
        ├── 前 50 条 → dev.jsonl
        └── 其余条目 → train.jsonl

本章最重要的对象变化是：

原始 Dataset row
    → 标准化 dict
    → list[dict]
    → 每行一个 JSON 对象的 JSONL 文件

这里的 question 和 answer 始终是 Python 字符串。此时没有调用 tokenizer，不存在 token 边界，也不存在 loss mask。 token 化要等到 protocol.py 构造模型 prompt 时才发生。

1.2 为什么选择 DAPO-Math-17k#

常量：

DATASET_ID = "BytedTsinghua-SIA/DAPO-Math-17k"

ReTool 的 reward 是 outcome-only reward：模型最终给出答案后，系统需要把它和参考答案做数学等价判断。因此训练样本至少要提供：

• question：给模型的问题；
• answer：可验证的参考答案。

DAPO-Math-17k 已经包含这两类信息，只是原始字段是 verl 风格的嵌套结构，需要先压平。

一个简化后的原始样本可以理解为：

row = {
    "prompt": [
        {"role": "user", "content": "外层答题模板 + 真正的数学题 + 外层提醒"}
    ],
    "reward_model": {"ground_truth": "3"},
    "extra_info": {"index": "sample-001"},
    "data_source": "math_dapo",
}

脚本希望把它变成：

{"id":"sample-001","question":"真正的数学题","answer":"3","data_source":"math_dapo"}

注意：这里没有保存原始推理过程。训练时模型会自己生成轨迹，脚本只需要问题和可验证的最终答案。

1.3 第一处清洗：去掉互相冲突的答题模板#

DAPO 原始问题外面包着固定提示，要求最终使用：

Answer: $Answer

但 ReTool 的 system prompt 和 reward 约定最终答案必须写成：

\boxed{答案}

如果不清洗，同一条 prompt 会同时要求两种最终格式。模型可能遵从 DAPO 的旧模板输出 Answer: 3，而 reward.py 只寻找 \boxed{3}，于是数学答案明明正确，仍会因为格式不合法得到 -1。

脚本把固定模板拆成两个常量：

PROMPT_PREFIX = ...
PROMPT_SUFFIX = ...

strip_dapo_template() 的状态变化为：

原字符串
  → 如果开头完全匹配 PROMPT_PREFIX，就删除前缀
  → 如果结尾完全匹配 PROMPT_SUFFIX，就删除后缀
  → strip() 删除题目首尾空白
  → 返回题目正文

重点：这是精确匹配，不是模糊匹配#

startswith() 和 endswith() 要求字符完全一致。如果数据里的模板发生了空格、换行或措辞变化，对应部分就不会被删除。

函数允许只命中一边：前缀匹配就删前缀，后缀匹配就删后缀。源码注释里的“模板不完整时原样保留”更准确地说，应理解为“没有匹配到的部分原样保留”。

1.4 核心函数：normalize_row()#

函数签名：

def normalize_row(index: int, row: dict[str, Any]) -> dict[str, Any] | None:

它有两种返回结果：

• 返回标准化字典：样本有效；
• 返回 None：样本缺少必要字段，后面会被过滤。

第一步：检查并读取 prompt#

prompt = row.get("prompt")
if not isinstance(prompt, list) or not prompt:
    return None

这里规定 prompt 必须是非空列表。随后只读取第一条消息：

question = strip_dapo_template(str(prompt[0].get("content") or ""))

数据契约是：

row["prompt"]
    └── 第 0 个消息 dict
          └── "content" → 问题字符串

or "" 用来把缺失或空的 content 变为空字符串；后面的非空校验会丢弃它。

第二步：读取参考答案#

reward_model = row.get("reward_model") or {}
ground_truth = (
    reward_model.get("ground_truth")
    if isinstance(reward_model, dict)
    else None
)

这一步先确认 reward_model 是字典，再读取 ground_truth，避免直接索引不存在的键。

有些数据可能把答案包成列表，因此源码做了兼容：

if isinstance(ground_truth, list):
    ground_truth = ground_truth[0] if ground_truth else None

也就是说，多答案列表只保留第一个答案；空列表视为缺失答案。

最后统一转换为去除首尾空白的字符串：

answer = str(ground_truth or "").strip()

第三步：过滤无效样本#

if not question or not answer:
    return None

问题或参考答案任意一个为空，这条数据都无法用于当前训练：

• 没有问题，模型没有输入；
• 没有答案，outcome reward 无法判定对错。

第四步：压平为四字段字典#

return {
    "id": str(row.get("extra_info", {}).get("index") or index),
    "question": question,
    "answer": answer,
    "data_source": str(row.get("data_source") or "dapo_math"),
}

四个字段的职责如下：

这一步叫“标准化”而不是“token 化”：它只是统一字段和字符串格式。

1.5 一行较难的 Python：海象运算符过滤#

main() 用一段列表推导完成“逐行标准化 + 丢弃无效行”：

records = [
    record
    for index, row in enumerate(dataset)
    if (record := normalize_row(index, row)) is not None
]

初学者可以把它展开理解成：

records = []
for index, row in enumerate(dataset):
    record = normalize_row(index, row)
    if record is not None:
        records.append(record)

:= 是赋值表达式，常称“海象运算符”。这里让 normalize_row() 只调用一次，同时把返回结果用于判断和追加。

完成后：

dataset: Hugging Face Dataset
records: list[dict[str, Any]]

如果一条有效数据也没有，源码立即报错，而不是生成两个看似正常但实际为空的文件：

if not records:
    raise ValueError(...)

1.6 两条数据加载路径#

命令行参数 --raw 的默认值是：

05-retool/datasets/raw/dapo-math-17k.parquet

加载逻辑：

本地 raw parquet 存在
    → load_dataset("parquet", data_files=..., split="train")

本地 raw parquet 不存在
    → load_dataset(DATASET_ID, split="train")

两条路径最终都返回可遍历的 Hugging Face Dataset，所以下游清洗逻辑不需要区分数据来自本地还是 Hub。

易混点#

从 Hub 加载时，datasets 库会使用自己的缓存，但这段脚本不会主动把原始 parquet 写入默认的 datasets/raw/dapo-math-17k.parquet。因此 README 展示的 raw/ 文件更像是可选的人工下载/缓存路线，而不是这段脚本必然生成的产物。

1.7 固定打乱、切分 train/dev#

清洗完成后执行：

random.Random(args.seed).shuffle(records)
dev_records = records[: args.dev_size]
train_records = records[args.dev_size :]

为什么要先打乱#

如果原始数据按来源、难度或题型排序，直接取前 50 条可能得到有偏的开发集。先打乱再切分，可以减少这种顺序偏差。

为什么使用 random.Random(args.seed)#

它创建一个局部随机数生成器：

• 相同输入顺序 + 相同 seed，得到相同切分；
• 不污染其他代码使用的全局随机状态；
• 默认 seed 为 42，便于重复实验。

dev 集做什么#

默认前 50 条写入 dev.jsonl，用于 smoke test 和快速验证，明确不参与训练。其余数据写入 train.jsonl。

这里的 dev 不是最终论文评测集。最终模型评测使用的是 AIME 2025；两者职责不同：

dev.jsonl  → 开发阶段快速检查训练/rollout 链路
AIME 2025  → 比较 Base 与各 checkpoint 的正式评测

1.8 write_jsonl()：为什么是一行一个 JSON#

with path.open("w", encoding="utf-8") as file:
    for record in records:
        file.write(json.dumps(record, ensure_ascii=False) + "\n")

JSONL 的特点是每一行都是一个独立 JSON 对象：

{"id":"1","question":"...","answer":"3","data_source":"math_dapo"}
{"id":"2","question":"...","answer":"7","data_source":"math_dapo"}

它适合训练数据，因为可以逐行读取，不必一次解析一个巨大的 JSON 数组。

关键参数：

• 文件模式 "w"：覆盖写入；重复执行脚本会重建输出文件；
• encoding="utf-8"：明确使用 UTF-8；
• ensure_ascii=False：中文等非 ASCII 字符直接保存，而不是写成 \uXXXX；
• 每条记录末尾追加 \n：保证一行一条样本。

输出目录会先由下面的代码创建：

args.output_dir.mkdir(parents=True, exist_ok=True)

parents=True 允许创建缺失的父目录，exist_ok=True 表示目录已经存在也不报错。

1.9 脚本结尾的检查信息#

写完文件后，脚本输出：

• train 样本数；
• dev 样本数；
• answer 字符长度的最小值、中位位置值、最大值；
• 第一条训练样本的前 300 个字符。

特别注意：

len(record["answer"])

统计的是 Python 字符串长度，不是 tokenizer token 数。它只能帮助发现空答案或异常长答案，不能预测模型训练显存或上下文长度。

1.10 main() 的完整执行顺序#

把所有函数串起来，执行入口是：

if __name__ == "__main__"
    │
    ▼
main()
    │
    ├─ parse_args()
    ├─ 创建输出目录
    ├─ 从本地 parquet 或 Hub 加载 Dataset
    ├─ enumerate(dataset)
    │     └─ normalize_row()
    │           ├─ strip_dapo_template()
    │           └─ 无效样本返回 None
    ├─ 固定 seed 原地打乱 records
    ├─ 切分 dev_records / train_records
    ├─ write_jsonl(train.jsonl)
    ├─ write_jsonl(dev.jsonl)
    └─ 打印样本数、答案长度和样例

脚本的输出契约是：只要成功完成，下游就可以假设 JSONL 每条记录都有四个非空字段。

1.11 重点与疑难点#

重点 1：格式模板会间接改变 reward#

数据清洗看似离 RL 很远，但旧的 Answer: 模板如果残留，会诱导模型输出 reward 不接受的格式。数据层的字符串清洗最终会影响奖励分布和训练信号。

重点 2：参考答案不是监督学习标签#

当前实现没有让模型直接模仿 answer，也没有把参考答案拼进 prompt。它只在轨迹结束后供 reward.py 判定结果，因此这是 outcome supervision，不是传统的答案 token 交叉熵监督。

重点 3：此处还没有 token 边界#

本章对象边界是 JSON 字段边界。真正需要严格追踪的 prompt token、assistant token 和 observation token，要从 protocol.py 与 rollout.py 开始。

疑难点 1：源码信任原始 schema#

源码检查了 prompt 是非空列表，但默认 prompt[0] 一定是字典。如果上游 schema 变化为字符串或其他类型，.get("content") 会报错。这是对固定数据集契约的合理简化，不是通用数据清洗器。

疑难点 2：or 会把某些假值视为缺失#

例如：

ground_truth or ""
extra_info.index or index

如果值是数字 0，Python 会把它视为假值。对当前数据集，答案和 id 通常以字符串保存，所以一般不会出问题；若扩展到其他数据源，应区分“值为 0”和“字段缺失”。

疑难点 3：--dev-size 没有边界校验#

若 dev_size 大于等于有效样本总数，train_records 会为空，而结尾访问 train_records[0] 时会报错。默认值 50 对 DAPO-Math-17k 是安全的，但做本地小样本测试时需要留意。

疑难点 4：切分依赖上游迭代顺序#

固定 seed 只能保证：输入记录顺序相同时，打乱结果相同。如果远端数据集版本或上游顺序变化，即使 seed 相同，最终 train/dev 切分也可能变化。严格复现时还需要固定数据集版本或保留本地原始文件。

1.12 本章对象账本#

这里还没有 reward 数值、advantage、logprobs 或任何张量。下一章 data.py 会把每行 JSON 转成不可变的 MathExample 对象，并说明训练循环如何按 step 取题。

1.13 本章小结#

用一句话概括 prepare_data.py：

它把 DAPO 的嵌套 verl 数据压平成 ReTool 需要的四字段 JSONL，并提前消除 Answer: 与 \boxed{} 的协议冲突，再用固定 seed 划出开发集和训练集。

读完本章应能回答：

1. question 和 answer 分别来自原始数据的哪个字段？
2. 为什么必须删除 DAPO 的外层答题模板？
3. normalize_row() 为什么可能返回 None？
4. 固定 seed 能保证什么，不能保证什么？
5. 为什么这里统计的 answer length 不是 token 长度？
6. dev.jsonl 与 AIME 2025 的职责有什么不同？