feat: add MemAgent multi-turn training example

[TabsPhasers]

概要

本 PR 为 ReLax 增加一套可运行的 MemAgent 训练与评测流程。

MemAgent 先按 token 将长文档切成多个 chunk。每轮模型根据 question + previous memory + current chunk 生成新的 memory，并用新结果覆盖旧 memory。处理完全部 chunk 后，模型根据 question + final memory 生成最终答案。训练阶段会保存每次 memory 更新和最终回答，并将它们逐轮展开成独立训练数据。

主要改动包括：

• 在 examples/mem_agent/ 实现多轮 memory rollout、最终答案 reward 和轨迹展开；
• 为每个 memory turn 和 final turn 保存 tokens、response_length、loss_mask、rollout_log_probs；
• 在 trajectory 维度完成 GRPO reward normalization，再将最终 advantage 分配到各个 turn；
• 增加 HotpotQA 数据准备、训练、checkpoint 转换、HotpotQA/RULER-HQA 评测和结果比较脚本；
• 支持 custom converter 动态产生训练行，并让 Actor 按本轮实际行数完整消费；
• 补充自定义训练文档、单元测试、CPU 回归测试和固定实验配方。

关联 #236 和 Contributor Program Task 36。

背景

ReLax 已支持替换 rollout 生成逻辑。MemAgent 还需要处理动态的 1:N 训练数据展开：一条 trajectory 会根据实际 chunk 数生成多条训练行。

TransferQueue 的容量和 Actor 的读取数量如果仍按 trajectory 数计算，可能丢失部分 memory turn、少读训练行，或因等待错误行数而卡住。GRPO normalization 的执行顺序也很重要；先展开 turn 会让 chunk 较多的问题在组内统计中占据更高权重。

本实现明确了两项约束：

1. 在每条 trajectory 仍对应一条记录时完成 reward 和 advantage 计算；
2. turn 展开后的保存行数、传输行数和消费行数完全一致。

实现

1. 多轮 memory rollout

question + long context
        |
        v
按 token 切分 context
        |
        v
question + old memory + chunk_i
        |
        v
生成 new memory_i
（覆盖旧 memory，并执行 token 上限）
        |
        v
重复处理配置范围内的全部 chunk
        |
        v
question + final memory
        |
        v
生成最终 \boxed{} 答案

每个 chunk 对应一次独立请求。下一轮重新构造 prompt，只传入截断后的 memory 文本。空的 memory 更新同样会覆盖旧值，确保每轮更新语义一致。

最新 Task 36 验收要求不限定模型规模。本次正式复现固定使用
Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca。ReLax 与 VIME
先按已经完成验证的短配方进行控制变量训练，再用相同的长文档配置评测
RULER-HQA，检查从训练长度到更长文档的泛化效果。

控制变量训练配方如下：

配置项	值
模型	Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca
数据	冻结的 1,000 条 HotpotQA 训练数据及其顺序
Chunk 上限	512 tokens
Memory 上限	128 tokens
最终答案上限	64 tokens
最大 chunk 数	4
GRPO group size	8
信用分配	split
学习率	1e-6
KL 系数	0.001
对照 checkpoint	ReLax 与 VIME 使用相同训练步数和样本暴露量

长文档评测配方如下：

配置项	值
数据	BytedTsinghua-SIA/hotpotqa@27275ff4fee67ac0acb6478e405e7ac07efbdc1a 的 eval_50/200/800
Chunk 上限	2048 tokens
Memory 上限	1024 tokens
最终答案上限	256 tokens
最大 chunk 数	64
对比模型	基础模型、VIME checkpoint、ReLax checkpoint

最终 prompt 的输入范围固定为 question、最后一轮 memory 和答案格式说明。原始 context、当前 chunk 和此前各轮 memory 均在该阶段结束前释放。

训练和评测使用不同的文档长度是本次预先固定的泛化测试设计。VIME 对照训练必须采用与 ReLax 相同的条件，基础模型、VIME 和 ReLax 也必须使用相同的评测条件。现有 ReLax iter_0000074 使用短配方训练，因此不能将其描述为 VIME 默认长配方的训练结果。

2. 训练数据展开

rollout 保存全部 memory turns 和 final turn。converter 先校验整条 trajectory，然后调用 ReLax 现有的 GRPO reward 后处理，最后逐 turn 展开。

每条轨迹必须满足以下条件：

• turn index 从 0 开始连续递增；
• 若干 memory turn 后只能有一个 final turn；
• response_length 为正，且处于对应 turn 的 token 上限内；
• response_length <= len(tokens)；
• len(loss_mask) == response_length；
• len(rollout_log_probs) == response_length；
• 同一个 prompt group 内的 trajectory 具有相同 turn 数；
• trajectory 状态完整，可用于训练。

默认信用分配方式为 split：一条 trajectory 的 advantage 除以 turn 数，每个 turn 获得其中一份。代码也实现并测试了 share，该模式为每个 turn 分配完整 advantage。

3. Reward 计算

中间 memory turn 不单独计分。reward function 从最终回答的 \boxed{} 中抽取答案，再按照 HotpotQA normalized exact match 计算一次 0/1 reward。缺少合法 box、答案不匹配和生成失败的结果均记为 0。

输出格式属于训练目标。训练与评测共用同一套答案抽取和规范化逻辑。

4. ReLax 接入方式

MemAgent 使用现有扩展点接入：

• --custom-generate-function-path：多轮 memory rollout；
• --custom-convert-samples-to-train-data-path：trajectory converter；
• --custom-rm-path：最终答案 reward；
• --custom-config-path：chunk、memory、展开和信用分配配置。

SGLang 核心保持原样，Controller 只处理通用的训练数据行数和队列分组信息。

框架侧增加了动态 1:N 展开需要的通用能力：

• 在线传输和调试回放共用 custom converter；
• custom_train_sample_expansion_factor 为 TransferQueue 预留最大容量；
• custom_train_data_group_size 指定转换后数据在 TransferQueue 中的分组大小；
• RolloutFnTrainOutput.train_row_count 返回当前训练分区实际写入的行数；
• 展开训练模式按 rollout_id 等待行数，生成或转换失败时同步结束等待；
• Actor 按实际行数读取数据，并检查总行数可被 data-parallel size 整除。

train_row_count 专门承载训练行数，metrics 继续用于监控指标。普通 1:1 rollout、Agentic 和 Forge 的非展开路径使用默认值 train_row_count=None。启用 custom_train_expanded_batch 的高层 custom rollout 必须返回精确行数，缺失时会直接报错。

主要文件

模块	主要文件	职责
Rollout	examples/mem_agent/rollout.py, prompts.py	token 切块、memory 覆盖、final prompt 输入隔离
训练数据	examples/mem_agent/convert.py	校验 trajectory，先归一化再展开 turn
Reward	examples/mem_agent/reward.py, metrics.py	boxed exact match 和诊断指标
数据	prepare_data.py, prepare_data.sh	HotpotQA/RULER-HQA 下载、转换和 SHA-256 manifest
配方	config*.yaml, run-*.sh	0.6B 控制变量训练与长文档评测配方
评测	eval_ruler_hqa.py, compare_results.py	多轮 memory 评测、逐样本结果和 3 个百分点比较
框架	rollout 返回值、TransferQueue、Controller、Actor	动态展开行数的产生、传递和消费
测试与文档	tests/examples/mem_agent/ 及相关回归	行为、异常路径、默认路径和配方检查

运行方式

准备固定版本的数据：

DATA_DIR=/data/mem-agent \
bash examples/mem_agent/prepare-data.sh

运行 Qwen3-0.6B 配方的 2-step 冒烟测试：

MODEL_PATH=/data/models/Qwen3-0.6B \
DATA_DIR=/data/mem-agent-pilot \
RUN_ROOT=/data/runs/mem-agent-relax \
NUM_ROLLOUT=2 \
bash examples/mem_agent/run-qwen3-0.6B-train.sh

同一入口可通过 NUM_ROLLOUT、LOAD_PATH 和 START_ROLLOUT_ID 运行完整训练或从 checkpoint 恢复。训练结束后使用 convert-to-hf.sh 转换选定 checkpoint，再通过统一评测入口运行 HotpotQA/RULER-HQA。训练日志、逐 rollout 结果、TensorBoard event、GPU 监控和 reward 汇总均保留在 RUN_ROOT。

在基础模型、同配方 VIME checkpoint 和 ReLax checkpoint 之间进行长文档同条件比较：

BASE_MODEL_PATH=/data/models/Qwen3-0.6B \
VIME_MODEL_PATH=/data/checkpoints/vime-hf \
RELAX_MODEL_PATH=/data/checkpoints/mem-agent-relax-hf \
TOKENIZER_PATH=/data/models/Qwen3-0.6B \
DATA_DIR=/data/mem-agent \
RESULTS_DIR=/data/results/vime-vs-relax \
LENGTHS="50 200 800" \
bash examples/mem_agent/run-paired-eval.sh

评测产物包括逐样本 JSONL 和汇总 JSON。比较脚本会检查结果是否为空、请求是否成功，以及输入 SHA、评测格式和控制变量是否一致。

测试结果

本文对应的本地提交为 1f66ac8632d3b379116cfd73dcd34e3c60f7bc22，父提交为 merge commit 467cfa1，其中已合入 origin/main@682d474。本地工作树干净。

检查项	结果
MemAgent、custom converter、展开行数跟踪、Agentic 和结构化返回值定向测试	94 passed
CPU 全量测试（单独处理 Linux /dev/shm 用例）	1001 passed, 284 skipped
tests/test_s3_model_loader.py 其余用例	50 passed, 1 deselected
pre-commit run --all-files --show-diff-on-failure	passed
Ruff check / format check	passed
bash -n examples/mem_agent/*.sh	passed
Python compileall	passed
git diff --check	passed

被排除的断言要求系统提供默认 /dev/shm，本机 macOS 不具备该目录。该文件的其余 50 个用例已经执行。本轮测试全部使用 CPU。

测试覆盖：

• chunk 边界、最大 chunk 数和 context truncation；
• memory 覆盖、空更新和 token 硬上限；
• final prompt 的 question + final memory 输入范围；
• 全部 memory/final turn 的完整展开；
• tokens、mask 和 log-probs 的长度对齐；
• trajectory 维度 reward normalization 及 split/share 信用分配；
• 展开行数的产生、等待、失败唤醒和 Actor 消费；
• custom converter 的在线传输与调试回放；
• 普通 rollout 以及 Agentic/Forge 的非展开路径。

Qwen3-0.6B 已完成实验

Qwen3-0.6B 是本次正式复现模型。以下结果已经覆盖真实训练控制流、reward 趋势、checkpoint 转换和 HotpotQA 留出集评测。剩余效果实验为 VIME 控制变量训练以及三个模型的 RULER-HQA 长文档对照评测。

实验配置

配置项	值
模型	Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca
数据来源	hotpotqa/hotpot_qa@1908d6afbbead072334abe2965f91bd2709910ab，distractor train
精度与并行配置	BF16，TP=1，1×RTX 4090
Chunk / memory / final 上限	512 / 128 / 64 tokens
最大 chunk 数	4
Thinking	关闭
采样参数	temperature 1.0，top_p 1.0
评测 seed	4242

先在 4,000 条 HotpotQA 候选上运行 Pass@8，共完成 32,000 条 trajectory，请求错误数为 0。随后固定 1,000 个唯一训练 ID。另取 500 条与训练集无交集的留出评测数据（held-out），选取过程与基础模型的答题结果无关。

2-step 冒烟测试

单张 RTX 4090 连续完成 2 个 rollout step。每一步的数据数量为：

8 trajectories
  -> 32 saved turns
  -> 32 transferred rows
  -> 32 consumed rows

Ray job 状态为 SUCCEEDED，运行期间没有 hang 或 OOM。第二步记录到非零 loss 和 gradient，并成功保存 iter_0000001 checkpoint。这组结果用于确认训练链路和行数契约可以连续运行。

Reward 趋势

20-step 趋势检查中，前 5 个 rollout 的 reward 均值为 0.365625，后 5 个为 0.734375，绝对增加 0.36875。

之后从 iter_0000019 恢复训练，连续完成 rollout 0..91，Actor 训练指标记录到 step 90。训练 reward 已形成明显提升趋势，归档中最新完整、可转换的 checkpoint 为 iter_0000074，后续统一使用该 checkpoint 评测。

按 trajectory 逐题汇总的 reward 如下：

统计窗口	Reward 均值
前 10 个 rollout	0.478125
后 10 个 rollout	0.725000
绝对变化	+0.246875
最高值	rollout 70：0.953125

训练归档保留 3 个原始 TensorBoard event 文件，共 77 个 scalar tags、6,777 个 scalar points。指标包括 reward、loss、grad norm、KL、entropy、learning rate、log-prob 差异、response length、truncation、throughput 和各阶段耗时。

同条件留出集评测

iter_0000074 已转换为 Hugging Face 格式。它与原始 Qwen3-0.6B 在同一组 500 条留出数据上各评测一次。两组评测使用相同的数据、tokenizer、prompt、chunk/memory 配置、采样参数、停止条件、seed 和 evaluator，均完成 500/500 条请求，错误数为 0。

模型	Boxed EM	sub_EM	F1
基础模型 Qwen3-0.6B	17.0%	27.4%	25.2092%
ReLax iter_0000074	41.6%	47.0%	49.1406%
绝对变化	+24.6pp	+19.6pp	+23.9314pp

同一组 500 题中，boxed EM 有 148 条从错变对、25 条从对变错、60 条两边都对、267 条两边都错。基础模型有 189 条输出缺少可解析的 box，checkpoint 中该数量为 0。checkpoint 仍有 292 条 boxed 答案不匹配，因此当前提升同时包含格式遵循和答题能力两方面变化。

验收进度

要求	当前证据	状态
训练、转换和评测入口可执行，至少连续训练 2 step	0.6B 单卡 2-step、checkpoint 保存、HF 转换和同条件评测	已通过
按 token 切块；memory 覆盖且不超限；final 只用 question + memory	单元测试、CPU 集成测试和真实 rollout	已通过
全部 turn 展开；字段严格对齐；最终 reward 正确分配	单元测试及每步 32 saved/transferred/consumed rows	已通过
训练 reward 有明显提升	20-step 和 92-rollout reward 曲线	已通过
HotpotQA checkpoint 优于基础模型	同一组 500 条留出数据的 0.6B 同条件评测	已通过
RULER-HQA 50/200/800 优于基础模型	基础模型与 ReLax iter_0000074 的长文档评测	待验证
ReLax 与 VIME 的 RULER-HQA sub_EM 绝对差不超过 3pp	VIME 使用相同 0.6B 短配方训练，三者使用相同长文档配置评测	待验证

兼容范围与已知开销

• 普通 1:1 rollout 使用默认的固定 batch 路径，train_row_count 保持默认值 None；
• train_row_count 是 RolloutFnTrainOutput 的可选字段，现有非展开调用方保持原用法；
• 展开训练模式要求高层 custom rollout 返回精确行数；
• trajectory 变长会增加生成请求、TransferQueue 容量、通信量和训练行数；
• trajectory 校验失败时训练会立即报错，以免丢弃尾部 turn 或等待错误行数；
• checkpoint 格式保持不变，已有训练配置和 checkpoint 可以继续用于非 MemAgent 任务。

剩余效果实验

官方验收要求中，流程可运行、行为正确和训练数据正确已有完整证据。效果验收中的 reward 提升和 HotpotQA 基础模型对照也已完成。剩余工作集中在一组控制变量实验：

1. 使用与 ReLax 相同的 1,000 条训练数据及顺序，独立训练一次 VIME。模型初始化、prompt、512/128/4 短配方、GRPO 参数、seed、训练步数和样本暴露量对齐 ReLax iter_0000074 的实际训练记录。
2. 对基础模型、VIME checkpoint 和 ReLax iter_0000074 运行同条件 RULER-HQA 评测。评测固定 chunk 上限 2048 tokens、memory 上限 1024 tokens、最大 chunk 数 64，并覆盖 eval_50/200/800。
3. 报告三个模型的总样本数、成功数、错误数和聚合指标，保存逐样本结果；确认 ReLax 在预选长度上优于基础模型，并检查 abs(sub_EM_ReLax - sub_EM_VIME) <= 3pp。

VIME 训练和正式评测将按现有 ReLax 实验的同一证据口径归档，包括 debug 日志、TensorBoard event、rollout、checkpoint、转换产物、配置、环境信息和 SHA-256 清单。PR 当前公开检查已通过；完整 CPU 回归结果和 macOS 上不适用的 /dev/shm 用例已在上文分别说明。

ReLax iter_0000074 使用 512/128/4 短配方训练；RULER-HQA 长文档配置用于检验其泛化能力。最终效果验收结论取决于 RULER-HQA 基础模型提升和 VIME/ReLax 3pp 对照结果。

审核顺序

1. examples/mem_agent/rollout.py：检查每轮 memory 覆盖和 final prompt 的输入范围；
2. examples/mem_agent/convert.py：检查 reward normalization 的执行顺序和全部 turn 的展开结果；
3. relax/engine/rollout/base_types.py、relax/engine/rollout/sglang_rollout.py 和 relax/distributed/ray/rollout.py：检查实际行数的产生、传递和等待逻辑；
4. relax/utils/utils.py、Controller 和 Actor：检查 converter、队列容量、数据分组和消费数量；
5. tests/examples/mem_agent/ 及展开行数相关测试：检查异常输入和普通训练路径；
6. scripts 和 README：检查固定模型、数据、参数，以及基础模型/VIME/ReLax 的评测条件。

custom_train_data_group_size 表示 TransferQueue 一次交给训练端的行分组大小。默认 converter 继续按原有 group size 传输数据；MemAgent 在 trajectory 维度完成 GRPO normalization 后，将每个展开行作为一个可独立消费的分组。

变更类型

• 缺陷修复
• 新功能
• 破坏性变更
• 文档更新
• 仅重构
• 性能优化
• CI/CD 或构建调整

检查清单

• MemAgent 的主要实现集中在 examples/mem_agent/ 目录
• 代码注释说明了整体流程以及切块、记忆长度和训练数据对齐等关键约束
• 本次实现没有修改 SGLang 核心代码
• CPU 测试、pre-commit、Shell 语法检查和 git diff --check 均已通过
• ReLax 使用 Qwen3-0.6B 在单张 RTX 4090 上连续训练 2 个 step，运行期间没有报错或卡死
• 基于 Qwen3-0.6B 训练得到的 ReLax checkpoint 已完成格式转换，并在相同配置的留出集上完成评测
• 官方验收要求及每项要求的当前完成情况已经逐条记录
• 在相同配置下完成 Qwen3-0.6B 基础模型、VIME checkpoint 和 ReLax checkpoint 的 RULER-HQA 评测
• 对 VIME 与 ReLax 的评测结果进行成对比较，确认所选 RULER-HQA 指标的绝对差不超过 3 个百分点