minimind/course/notes/26_agent_rollout_batch_notes.md

第26节：Agent Rollout 单条轨迹与批量采样笔记

我的问题笔记 | 聚焦 rollout_single 为什么是单条轨迹，以及能不能多个样本一起采样

---

目录

1. Q1: rollout_single 是只跑一条 message 吗？
2. Q2: 当前 rollout_batch 真的是批量生成吗？
3. Q3: 底层 rollout_engine.rollout 支持 batch 吗？
4. Q4: 为什么 Agent rollout 没有直接写成 batch？
5. Q5: 如果要改成批量采样，主循环应该长什么样？
6. Q6: old_logps 和 completion_mask 为什么要右移对齐？
7. 一句话总结

---

Q1: rollout_single 是只跑一条 message 吗？

不是。

rollout_single 一次跑的是 一个样本的一条 trajectory，不是一条单独的 message。

代码 trainer/train_agent.py:98-119

def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"):
    ...
    for turn in range(max_turns):
        context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools, open_thinking=open_thinking)
        inputs = tokenizer(context, return_tensors="pt", add_special_tokens=False).to(device)
        ...
        rollout_result = rollout_engine.rollout(
            prompt_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            num_generations=1,
            max_new_tokens=max_new_tokens,
            temperature=0.8,
        )

这里的 messages 是消息列表，可能长这样：

[
    {"role": "system", "content": "..."},
    {"role": "user", "content": "..."}
]

也可能已经包含历史 assistant/tool 消息。

所以它的边界是：

单个样本
-> 单条采样轨迹
-> 轨迹内部最多 max_turns 轮 assistant/tool 交互

因为 context 是单个字符串：

inputs = tokenizer(context, return_tensors="pt", add_special_tokens=False)

所以这里的 tensor 形状通常是：

inputs["input_ids"].shape      = [1, prompt_len]
inputs["attention_mask"].shape = [1, prompt_len]

这个 1 表示 batch size 是 1。

---

Q2: 当前 rollout_batch 真的是批量生成吗？

不是严格意义上的批量生成。

当前 rollout_batch 只是外层名字叫 batch，内部还是循环调用 rollout_single。

代码 trainer/train_agent.py:159-180

def rollout_batch(rollout_engine, tokenizer, messages_batch, tools_batch, num_gen, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"):
    ...
    for messages, tools in zip(messages_batch, tools_batch):
        for _ in range(num_gen):
            msgs_copy = [dict(m) for m in messages]
            completion, context, prompt_ids, response_ids, response_mask, response_old_logps, turn_outputs, unfinished = rollout_single(
                rollout_engine, tokenizer, msgs_copy, tools, max_turns, max_new_tokens, thinking_ratio, device
            )
            all_completions.append(completion)
            ...

这段实际执行顺序是：

batch 里的第 1 个样本
  -> 第 1 条 trajectory
  -> 第 2 条 trajectory
  -> ...

batch 里的第 2 个样本
  -> 第 1 条 trajectory
  -> 第 2 条 trajectory
  -> ...

也就是：

messages_batch 提供多个样本；
num_gen 给每个样本采多条轨迹；
但每条轨迹仍然是单独调用 rollout_single。

所以当前写法简单，但生成效率不高。

---

Q3: 底层 rollout_engine.rollout 支持 batch 吗？

支持。

至少 TorchRolloutEngine.rollout 本身是按 batched tensor 写的。

代码 trainer/rollout_engine.py:75-92

output_ids = model.generate(
    input_ids=prompt_ids.repeat_interleave(num_generations, dim=0),
    attention_mask=attention_mask.repeat_interleave(num_generations, dim=0),
    max_new_tokens=max_new_tokens,
    do_sample=True,
    temperature=temperature,
    num_return_sequences=1,
    pad_token_id=self.tokenizer.pad_token_id,
    eos_token_id=self.tokenizer.eos_token_id,
).clone()  # [B, P+R]
prompt_len = prompt_ids.size(1)
completion_ids = output_ids[:, prompt_len:]  # [B, R]
...
completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)

如果输入是：

prompt_ids.shape      = [B, P]
attention_mask.shape  = [B, P]
num_generations       = G

那么这两行：

prompt_ids.repeat_interleave(num_generations, dim=0)
attention_mask.repeat_interleave(num_generations, dim=0)

会把 batch 展开成：

[B * G, P]

生成结果大致是：

output_ids.shape       = [B * G, P + R]
completion_ids.shape   = [B * G, R]
per_token_logps.shape  = [B * G, R]
completions            = 长度 B * G 的字符串列表

所以限制不在底层 rollout engine，而在 agent 多轮工具循环的组织方式。

---

Q4: 为什么 Agent rollout 没有直接写成 batch？

因为 Agent rollout 不是“一次 prompt -> 一次 answer”的普通生成，而是带环境反馈的动态循环。

代码 trainer/train_agent.py:132-152

calls = parse_tool_calls(new_text)
if not calls:
    break
unfinished = turn == max_turns - 1
messages.append({"role": "assistant", "content": new_text})
for call in calls:
    name, raw = call.get("name", ""), call.get("arguments", {})
    ...
    result = execute_tool(name, raw)
    result_str = (json.dumps(result, ensure_ascii=False) if result else '{"error": "tool not found"}')[:2048]
    messages.append({"role": "tool", "content": result_str})

observe_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
observe_ids = tokenizer(observe_context, return_tensors="pt", add_special_tokens=False)["input_ids"][0].tolist()
current_len = len(prompt_ids) + len(response_ids)
obs_delta = observe_ids[current_len:]
response_ids.extend(obs_delta)
response_mask.extend([0] * len(obs_delta))
response_old_logps.extend([0.0] * len(obs_delta))

这段说明了几个 batch 难点：

难点	为什么麻烦
prompt 长度不同	多个 context 要 padding 成同一个 tensor
是否继续不同	有的样本没有 tool call，要结束；有的样本要继续下一轮
tool 数量不同	一条 response 里可能 0 个、1 个或多个 tool call
messages 会变	每条 trajectory 都要追加自己的 assistant/tool 消息
observation 不是 policy token	tool observation 要放进 response_ids，但 response_mask 要填 0
logprob 长度不同	policy 生成 token 有 old logprob，环境 observation token 没有

所以不能简单地把：

tokenizer(context, return_tensors="pt")

改成：

tokenizer(contexts, return_tensors="pt", padding=True)

然后就结束了。

真正的问题是：每一轮之后，batch 里的 trajectory 会分化。

---

Q5: 如果要改成批量采样，主循环应该长什么样？

正确思路是做 动态 active batch。

也就是先把每个样本展开成多条 trajectory，然后每一轮只把还没结束的 trajectory 收集起来批量生成。

伪代码：

trajectories = []

for messages, tools in zip(messages_batch, tools_batch):
    for _ in range(num_gen):
        trajectories.append({
            "messages": [dict(m) for m in messages],
            "tools": tools,
            "active": True,
            "open_thinking": random.random() < thinking_ratio,
            "prompt_ids": None,
            "response_ids": [],
            "response_mask": [],
            "response_old_logps": [],
            "turn_outputs": [],
            "unfinished": False,
        })

for turn in range(max_turns):
    active_trajs = [t for t in trajectories if t["active"]]
    if not active_trajs:
        break

    contexts = [
        tokenizer.apply_chat_template(
            t["messages"],
            tokenize=False,
            add_generation_prompt=True,
            tools=t["tools"],
            open_thinking=t["open_thinking"],
        )
        for t in active_trajs
    ]

    inputs = tokenizer(
        contexts,
        return_tensors="pt",
        padding=True,
        add_special_tokens=False,
    ).to(device)

    rollout_result = rollout_engine.rollout(
        prompt_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        num_generations=1,
        max_new_tokens=max_new_tokens,
        temperature=0.8,
    )

    for row, traj in enumerate(active_trajs):
        new_text = rollout_result.completions[row]
        calls = parse_tool_calls(new_text)

        if not calls:
            traj["active"] = False
            continue

        traj["messages"].append({"role": "assistant", "content": new_text})
        for call in calls:
            result = execute_tool(call["name"], call["arguments"])
            traj["messages"].append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)})

这里最重要的是：

batch 维度不是固定的原始样本 batch；
batch 维度是当前 turn 仍然 active 的 trajectory 集合。

实现时还要额外小心：

点	要保证什么
prompt_ids	每条 trajectory 只记录最初 prompt 的有效 token
padding	批量 tokenizer 后要用 attention_mask 找真实长度
response_mask	policy 生成 token 填 1，tool observation delta 填 0
response_old_logps	policy token 填 rollout logprob，observation token 填 0.0
final_context	每条 trajectory 各自维护，不要被 batch 内其它样本污染
num_generations	展开成 trajectory 后，rollout 每轮通常用 num_generations=1 更好管理

---

Q6: old_logps 和 completion_mask 为什么要右移对齐？

核心原因是：Causal LM 的 loss/logprob 不是直接对 input_ids[:, i] 计分，而是用第 i 个位置的 logits 预测第 i + 1 个 token。

代码 trainer/train_agent.py:255-285

for p, r, m, old_lp in zip(prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch):
    ids = p + r
    mask = [0] * len(p) + m
    old_logps = [0.0] * max(len(p) - 1, 0) + old_lp
    ...

full_response_masks = torch.tensor(...)
old_per_token_logps = torch.tensor(...)

res = model_unwrapped(input_ids, attention_mask=full_mask)
logits = res.logits[:, :-1, :]
per_token_logps = F.log_softmax(logits, dim=-1).gather(
    2,
    input_ids[:, 1:].unsqueeze(-1)
).squeeze(-1)

completion_mask = full_response_masks[:, 1:]

这里有三组长度：

张量/列表	长度	含义
ids = p + r	L	完整 token 序列
mask = [0] * len(p) + m	L	按原始 token 位置标记哪些 response token 参与训练
per_token_logps	L - 1	对 input_ids[:, 1:] 这些目标 token 的 logprob

所以 per_token_logps[:, i] 对应的不是 input_ids[:, i]，而是：

per_token_logps[:, i]
= log P(input_ids[:, i + 1] | input_ids[:, :i + 1])

举例：

p = [A, B, C]
r = [D, E]

input_ids = [A, B, C, D, E]

模型实际计分的位置是：

per_token_logps:
预测 B, 预测 C, 预测 D, 预测 E

为什么 old_logps 前面补 len(p) - 1 个 0？

old_lp 只来自 rollout 期间生成的 response token，不包含 prompt 内部 token 的 old logprob。

但是训练时的 per_token_logps 包含 prompt 内部的目标位置：

预测 B
预测 C
预测 D
预测 E

其中 预测 B、预测 C 这两个位置数量正好是：

len(p) - 1

所以要写成：

old_logps = [0.0] * max(len(p) - 1, 0) + old_lp

对上面的例子就是：

old_lp    = [old_logp(D), old_logp(E)]
old_logps = [0.0, 0.0, old_logp(D), old_logp(E)]

前面的 0.0 不是有意义的 old policy 概率，只是占位。它们之后会被 completion_mask 盖掉，不参与 policy loss。

注意这里不是补 len(p) 个 0。因为 per_token_logps 已经丢掉了第一个 token 的目标位置，只剩 input_ids[:, 1:]，prompt 内部还需要对齐的目标 token 是 B, C，数量是 len(p) - 1。

为什么 completion_mask = full_response_masks[:, 1:]？

full_response_masks 是按原始 input_ids 位置标的：

input_ids:            [A, B, C, D, E]
full_response_masks:  [0, 0, 0, 1, 1]

但 per_token_logps 是按目标 token 位置标的：

per_token_logps:      [预测 B, 预测 C, 预测 D, 预测 E]

所以 mask 也必须右移成：

completion_mask = full_response_masks[:, 1:]

对应结果是：

completion_mask:      [0, 0, 1, 1]

这样才表示：

预测 B: 不训练
预测 C: 不训练
预测 D: 训练
预测 E: 训练

如果错用 full_response_masks[:, :-1]，会得到：

[0, 0, 0, 1]

这样第一个 response token D 的 logprob 会被 mask 掉，loss 位置整体错一格。

最后要把 tool observation 也放进这个逻辑里理解：m 里 policy 生成 token 是 1，tool observation delta 是 0；old_lp 里 observation token 也用 0.0 占位。真正决定哪些 token 参与 loss 的是右移后的 completion_mask。

---

一句话总结

rollout_single 当前一次只跑 一条 trajectory，所以 tensor 是 [1, seq_len]。

底层 rollout_engine.rollout 支持 [B, seq_len] 批量生成，但 Agent rollout 有 tool call、tool observation、提前结束、多轮上下文变长这些动态分支。

因此可以批量采样，但要写成 按 active trajectory 动态组 batch，不能只把外层 for 循环删掉。