第二章：`data.py`——`MathExample` 与循环 batch#

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

2.1 先抓住本章主线#

上一章把原始数据写成了 JSONL：

{"id":"dapo-0027","question":"If 2x + 1 = 7, what is x?","answer":"3","data_source":"math_dapo"}

json

data.py 的工作是把这些磁盘文本重新变成训练代码方便使用的 Python 对象，并根据训练 step 提供固定大小的问题批次：

train.jsonl
    │
    │ load_examples()
    ▼
list[MathExample]
    │
    │ shuffled_examples(seed)
    ▼
固定顺序打乱后的 list[MathExample]
    │
    │ take_batch(start, batch_size)
    ▼
当前 step 的问题 batch
    │
    │ rollout_batch(..., group_size=G)
    ▼
每道题扩展成 G 条 Trajectory

text

本章仍处在数据对象层：没有 tokenizer，没有 prompt token，没有 logprob、reward、advantage 或张量。

2.2 `MathExample`：一道题的标准对象#

源码定义：

@dataclass(frozen=True)
class MathExample:
    id: str
    question: str
    answer: str
    data_source: str

python

它把一道题固定为四个具名字段。上一章的字典：

{
    "id": "dapo-0027",
    "question": "If 2x + 1 = 7, what is x?",
    "answer": "3",
    "data_source": "math_dapo",
}

python

读取后变成：

MathExample(
    id="dapo-0027",
    question="If 2x + 1 = 7, what is x?",
    answer="3",
    data_source="math_dapo",
)

python

为什么不用普通字典一直传下去#

具名对象带来三个好处：

结构明确：看到 MathExample 就知道必需的四个字段；
访问清晰：使用 example.question，而不是到处写 row["question"]；
类型提示有效：编辑器和静态检查工具知道字段名称及类型。

如果误写：

example.queston

python

运行时会明确报属性不存在；使用任意字典时，这类错误更容易藏在动态键名里。

2.3 `@dataclass` 帮我们生成了什么#

@dataclass 会根据字段声明自动生成常用方法，例如初始化方法和便于阅读的字符串表示。

因此源码不需要手写：

class MathExample:
    def __init__(self, id, question, answer, data_source):
        self.id = id
        self.question = question
        self.answer = answer
        self.data_source = data_source

python

创建对象时仍然可以使用清楚的关键字参数：

example = MathExample(
    id="dapo-0027",
    question="If 2x + 1 = 7, what is x?",
    answer="3",
    data_source="math_dapo",
)

python

2.4 `frozen=True`：题目对象不可重新赋值#

源码不是普通的 @dataclass，而是：

@dataclass(frozen=True)

python

这意味着对象创建后不能再给字段赋新值。例如：

example.answer = "4"

python

会抛出 FrozenInstanceError。

这很适合原始训练题：rollout 可以读取问题和参考答案，但不应该在生成轨迹时意外修改它们。后续每条 Trajectory 都会引用一个 MathExample，题目对象保持稳定，轨迹自己的动态状态则放在 Trajectory 内。

可以先记住这个对象分工：

MathExample：静态、不可变
├── 问题
├── 参考答案
└── 数据标识

Trajectory：动态、会变化
├── 多轮消息
├── assistant turns
├── code calls
├── done
├── reward
└── advantage

text

严格来说，frozen=True 禁止的是字段重新绑定，不会递归冻结字段内部的可变对象。不过 MathExample 的四个字段都是字符串，字符串本身也不可变，因此当前对象可以视为稳定的值对象。

2.5 `load_examples()`：JSONL 行到对象#

函数签名：

def load_examples(path: str | Path) -> list[MathExample]:

python

输入路径既可以是字符串，也可以是 Path：

load_examples("datasets/train.jsonl")
load_examples(Path("datasets/train.jsonl"))

python

返回值是：

list[MathExample]

python

也就是把整个文件一次性读进内存。

第一步：创建空列表#

examples: list[MathExample] = []

python

此时：

examples = []

text

类型标注告诉读者，这个列表只应装 MathExample。

第二步：打开文件#

with Path(path).open(encoding="utf-8") as file:

python

这里做了两件事：

Path(path) 统一字符串路径和 Path 对象；
with 代码块结束时自动关闭文件，即使读取中途抛出异常也会关闭。

第三步：逐行解析 JSON#

for line in file:
    row = json.loads(line)

python

假设当前文本行是：

{"id":"dapo-0027","question":"If 2x + 1 = 7, what is x?","answer":"3","data_source":"math_dapo"}

text

json.loads() 后得到普通字典：

row = {
    "id": "dapo-0027",
    "question": "If 2x + 1 = 7, what is x?",
    "answer": "3",
    "data_source": "math_dapo",
}

python

这里的 loads 可以理解为“load string”：它解析的是已经读进内存的字符串。文件本身由外层 for line in file 负责读取。

第四步：构造并追加 `MathExample`#

examples.append(
    MathExample(
        id=str(row["id"]),
        question=row["question"],
        answer=row["answer"],
        data_source=row["data_source"],
    )
)

python

对象变化为：

JSONL 文本行
    ↓ json.loads()
dict
    ↓ MathExample(...)
MathExample
    ↓ append()
list[MathExample]

text

只有 id 在这里再次显式执行了 str()，确保编号统一为字符串。其他三个字段直接信任 prepare_data.py 的输出契约。

文件读完后：

return examples

python

2.6 为什么读取阶段比较“严格”#

load_examples() 没有再次容错清洗。例如：

空白行会让 json.loads() 报错；
JSON 语法错误会抛 JSONDecodeError；
缺少 question 等字段会抛 KeyError；
question、answer 不会再次执行 strip()；
重复的 id 不会被检测。

这不是遗漏，而是模块边界的选择：

prepare_data.py：负责清洗和建立数据契约
data.py：信任契约，快速加载训练数据

text

如果 data.py 再次默默修复坏数据，反而可能掩盖数据准备阶段的问题。当前实现选择“尽早明显失败”。

2.7 `shuffled_examples()`：训练开始时固定打乱顺序#

源码：

def shuffled_examples(path: str | Path, seed: int) -> list[MathExample]:
    examples = load_examples(path)
    random.Random(seed).shuffle(examples)
    return examples

python

状态变化为：

train.jsonl 文件顺序
    ↓ load_examples()
[A, B, C, D, E]
    ↓ Random(seed).shuffle()，原地改变列表顺序
[C, A, E, B, D]  （仅为示意）

text

shuffle() 改变的是列表中对象的排列顺序，不会修改任何 MathExample 的内容。

为什么 `prepare_data.py` 已经 shuffle，这里还要再 shuffle#

两次打乱目的不同：

prepare_data.py 的 shuffle
    → 决定哪些题进入 dev，哪些题进入 train

data.py 的 shuffle
    → 决定 train 中题目被训练 step 消费的顺序

text

训练脚本可以使用自己的 --seed 控制训练顺序，而不需要重新生成数据文件。

这两次即使都默认使用 seed 42，也不能理解为“第二次会把第一次抵消”。第一次打乱后已经切走 dev，第二次面对的是一个长度和内容都不同的 train 列表。

固定 seed 的边界#

相同的数据文件顺序和相同 seed 会得到相同的训练顺序。若 JSONL 内容或行顺序变化，seed 相同也不能保证结果相同。

2.8 `take_batch()`：把列表当成首尾相接的环#

函数签名：

def take_batch(
    examples: list[MathExample],
    start: int,
    batch_size: int,
) -> list[MathExample]:

python

源码核心只有一行：

return [
    examples[(start + offset) % len(examples)]
    for offset in range(batch_size)
]

python

关键是取模运算 % len(examples)。它让超过列表末尾的位置绕回开头。

假设训练顺序为：

索引：  0  1  2  3  4
题目： [A, B, C, D, E]

text

现在设置：

batch_size = 2

python

不同 step 会得到：

step	`start = step × 2`	访问索引	batch
0	0	`0, 1`	`[A, B]`
1	2	`2, 3`	`[C, D]`
2	4	`4, 0`	`[E, A]`
3	6	`1, 2`	`[B, C]`

以 step 2 为例：

offset = 0 → (4 + 0) % 5 = 4 → E
offset = 1 → (4 + 1) % 5 = 0 → A

text

所以 batch 可以跨越列表尾部，不会因为“剩余题目不足一个完整 batch”而丢弃 E。

可以把数据列表想象成一个环：

A → B → C → D → E
↑               ↓
└───────────────┘

text

2.9 `start` 在训练循环中如何计算#

train.py 中的调用是：

batch = take_batch(
    examples,
    step * args.questions_per_batch,
    args.questions_per_batch,
)

python

因此：

start = step × questions_per_batch

text

这相当于把训练问题看成一条无限重复的数据流：

A, B, C, D, E, A, B, C, D, E, A, ...

text

每个 step 从上一个 step 结束的位置继续取，不会在每个 step 都重新随机抽题。随机性来自训练启动时的一次固定 shuffle。

当前训练没有显式的 epoch 循环，主要进度单位是 step。当 step 足够多时，题目自然会循环出现。

2.10 空数据为什么单独处理#

源码先判断：

if not examples:
    return []

python

若不提前返回，后面的：

... % len(examples)

python

就会变成对 0 取模并抛出异常。

不过训练入口不会静默接受空数据。train.py 随后还有：

if not examples:
    raise ValueError("训练数据为空，请先运行 prepare_data.py")

python

职责分工是：

take_batch() 对空列表安全返回；
train.py 在业务入口给出更清楚的错误信息。

2.11 问题 batch 不等于轨迹 batch#

这是本章最重要的数量关系。

假设：

questions_per_batch = 2
group_size = 8

text

take_batch() 只会返回两道题：

batch = [question_A, question_B]

python

随后 rollout_batch() 会为每道题采样 8 条轨迹：

问题 A
├── trajectory A0
├── trajectory A1
├── ...
└── trajectory A7

问题 B
├── trajectory B0
├── trajectory B1
├── ...
└── trajectory B7

text

因此总轨迹数是：

问题数量 × group_size
= 2 × 8
= 16 条轨迹

text

对象数量变化为：

2 个 MathExample
    ↓ rollout 首轮分叉
16 个 Trajectory

text

同一道题的 8 条轨迹构成一个 reward group，稍后会在组内计算相对 advantage。这里的 data.py 只负责选出问题，不负责复制轨迹或计算 advantage。

2.12 `MathExample` 在 rollout 中放在哪里#

rollout.py 的 Trajectory 包含：

example: MathExample

python

因此每条轨迹都能随时访问：

trajectory.example.question
trajectory.example.answer

python

两者用途不同：

example.question
    → 构造 system + user 初始消息
    → 模型可以看到

example.answer
    → 轨迹结束后交给 reward.py
    → 模型看不到

text

这是一个重要的信息边界：参考答案跟着轨迹对象流转，是为了评分；它不会进入模型上下文。

2.13 训练与评测如何复用 `MathExample`#

训练路径：

train.jsonl
  → load_examples()
  → MathExample
  → rollout_batch()

text

评测路径略有不同。eval.py 读取 AIME 2025 后直接构造：

MathExample(
    id=...,
    question=str(row["problem"]).strip(),
    answer=str(row["answer"]),
    data_source="aime_2025",
)

python

也就是说，训练集与评测集的磁盘格式可以不同，但进入 rollout 前都被统一成 MathExample：

DAPO JSONL ── load_examples() ──┐
                               ├─→ MathExample → rollout
AIME Dataset ── 直接构造 ───────┘

text

这就是标准数据对象的价值：rollout 不必关心上游数据是 JSONL、Hugging Face Dataset，还是其他格式。

2.14 `max_train_samples` 如何影响数据流#

训练入口在 shuffle 之后还可能执行：

if args.max_train_samples > 0:
    examples = examples[: args.max_train_samples]

python

因此它选择的是打乱后列表的前 N 条：

完整 train.jsonl
    ↓ 固定 seed shuffle
训练顺序列表
    ↓ 截取前 N 条
本次实验实际使用的题目池
    ↓ take_batch() 循环取题

text

这对低成本 smoke test 很有用：可以固定一个小题目池反复验证训练闭环，而不必改动原始数据文件。

但也要意识到，如果训练 step 很多，小题目池会被反复使用，不代表模型看到了大量不同题目。

2.15 重点与疑难点#

重点 1：数据顺序是固定随机，不是每步随机采样#

训练启动时 shuffle 一次，之后 take_batch() 顺序循环。这样便于复现每个 step 使用了哪些问题。

重点 2：静态题目和动态轨迹分离#

MathExample 保存不可变的题目事实；模型生成的消息、代码调用和训练信号属于 Trajectory。不要把“题目样本”和“模型轨迹”混为同一个对象。

重点 3：参考答案随对象携带，但不会泄露给模型#

answer 会进入 Trajectory.example，供最终评分使用；构造模型初始消息时只取 question。

疑难点 1：加载器一次把全部数据读进内存#

load_examples() 返回完整列表，不是流式迭代器。DAPO-Math-17k 的规模可以接受这种实现；若扩展到超大数据集，可能需要流式读取或数据索引。

疑难点 2：`batch_size` 大于数据量时，单个 batch 内会重复题目#

例如：

take_batch([A, B], start=0, batch_size=5)

python

结果是：

[A, B, A, B, A]

python

训练默认配置不会遇到这个问题，但极小数据 smoke test 需要注意。同一题若在一个 batch 中重复出现，它们拥有不同的 question_index，rollout 会把它们当作两个独立 group。

疑难点 3：函数没有校验负数和零 batch#

batch_size=0 或负数时，range(batch_size) 为空，函数返回空列表。正常训练应保证 questions_per_batch > 0；当前辅助函数本身没有强制检查。

疑难点 4：返回列表复用原对象#

take_batch() 创建了一个新的列表，但列表里的元素仍是原来的 MathExample 引用，并没有深拷贝。由于 MathExample 是 frozen 且字段都是字符串，这种复用既安全又节省内存。

2.16 本章对象账本#

阶段	对象	典型类型	是否含 token	是否含训练信号
磁盘记录	`line`	`str`	否	只有参考答案
JSON 解析结果	`row`	`dict`	否	只有参考答案
单道标准题	`example`	`MathExample`	否	只有参考答案
全部训练题	`examples`	`list[MathExample]`	否	只有参考答案
当前问题批次	`batch`	`list[MathExample]`	否	只有参考答案
下一模块产物	`trajectories`	`list[Trajectory]`	后续会有	后续会有

当前对象边界仍是“一道题一个 MathExample”。下一章进入 protocol.py 后，question 将第一次被放进结构化消息，并通过 tokenizer 变成模型真正接收的 token ID 序列。

2.17 本章小结#

用一句话概括 data.py：

它把 JSONL 中的四字段记录加载为不可变的 MathExample，按固定 seed 决定训练顺序，再通过取模把有限题目列表变成可供任意多个训练 step 连续消费的问题流。

读完本章应能回答：

MathExample 为什么使用 @dataclass(frozen=True)？
load_examples() 中 line → row → MathExample 分别是什么类型？
数据准备阶段和训练加载阶段为什么各 shuffle 一次？
take_batch() 如何使用取模跨越列表尾部？
2 道问题、group_size=8 最终会产生多少条轨迹？
为什么 answer 在轨迹对象里，却不会泄露到模型 prompt？

ReTool 源码精读02

第二章：data.py——MathExample 与循环 batch#

2.1 先抓住本章主线#

2.2 MathExample：一道题的标准对象#

为什么不用普通字典一直传下去#

2.3 @dataclass 帮我们生成了什么#

2.4 frozen=True：题目对象不可重新赋值#

2.5 load_examples()：JSONL 行到对象#

第一步：创建空列表#

第二步：打开文件#

第三步：逐行解析 JSON#

第四步：构造并追加 MathExample#

2.6 为什么读取阶段比较“严格”#

2.7 shuffled_examples()：训练开始时固定打乱顺序#

为什么 prepare_data.py 已经 shuffle，这里还要再 shuffle#

固定 seed 的边界#

2.8 take_batch()：把列表当成首尾相接的环#

2.9 start 在训练循环中如何计算#

2.10 空数据为什么单独处理#

2.11 问题 batch 不等于轨迹 batch#

2.12 MathExample 在 rollout 中放在哪里#

2.13 训练与评测如何复用 MathExample#

2.14 max_train_samples 如何影响数据流#

2.15 重点与疑难点#

重点 1：数据顺序是固定随机，不是每步随机采样#

重点 2：静态题目和动态轨迹分离#

重点 3：参考答案随对象携带，但不会泄露给模型#

疑难点 1：加载器一次把全部数据读进内存#

疑难点 2：batch_size 大于数据量时，单个 batch 内会重复题目#

疑难点 3：函数没有校验负数和零 batch#

疑难点 4：返回列表复用原对象#

2.16 本章对象账本#

2.17 本章小结#

第二章：`data.py`——`MathExample` 与循环 batch#

2.2 `MathExample`：一道题的标准对象#

2.3 `@dataclass` 帮我们生成了什么#

2.4 `frozen=True`：题目对象不可重新赋值#

2.5 `load_examples()`：JSONL 行到对象#

第四步：构造并追加 `MathExample`#

2.7 `shuffled_examples()`：训练开始时固定打乱顺序#

为什么 `prepare_data.py` 已经 shuffle，这里还要再 shuffle#

2.8 `take_batch()`：把列表当成首尾相接的环#

2.9 `start` 在训练循环中如何计算#

2.12 `MathExample` 在 rollout 中放在哪里#

2.13 训练与评测如何复用 `MathExample`#

2.14 `max_train_samples` 如何影响数据流#

疑难点 2：`batch_size` 大于数据量时，单个 batch 内会重复题目#