feat: add MemAgent multi-turn training example

docs/en/guide/customize-training.md

@@ -211,6 +211,12 @@ generate.manages_inference_permit = True
 
 Specify via launch script (`--custom-generate-function-path examples.deepeyes.rollout.generate`), or per eval dataset via `custom_generate_function_path` in eval config.
 
+If you replace the higher-level function selected by `--rollout-function-path`, return
+`RolloutFnTrainOutput`. Its `metrics` member is for observability. A rollout that enables
+`custom_train_expanded_batch` and transfers a data-dependent number of rows must also set
+`train_row_count` to the exact post-conversion row count placed in the current training partition.
+Ordinary 1:1 rollout functions should leave `train_row_count` as `None`.
+
 ### Per-request concurrency scheduling for multi-turn rollout
 
 By default `generate_and_rm` holds one session-level concurrency permit (`GenerateState.semaphore`) for the entire custom `generate` call — a multi-turn rollout keeps the slot even while running env/tool steps, which hurts engine utilization.

docs/zh/guide/customize-training.md

@@ -207,6 +207,11 @@ generate.manages_inference_permit = True
 
 通过启动脚本指定（`--custom-generate-function-path examples.deepeyes.rollout.generate`），或在评估数据集配置中通过 `custom_generate_function_path` 按数据集设置。
 
+如果替换 `--rollout-function-path` 指向的更高层 rollout 函数，应返回
+`RolloutFnTrainOutput`。其中 `metrics` 只用于可观测性；若 rollout 开启
+`custom_train_expanded_batch`，且实际传输行数由 converter 动态展开决定，还必须把当前训练分区中
+实际写入的转换后行数填入 `train_row_count`。普通 1:1 rollout 应保持 `train_row_count=None`。
+
 ### 多轮 Rollout 的请求级并发调度
 
 默认情况下，`generate_and_rm` 会为整个自定义 `generate` 调用持有一把会话级并发锁（`GenerateState.semaphore`）——多轮 rollout 在环境/工具执行期间也一直占用名额，降低推理引擎利用率。

examples/mem_agent/README.md

@@ -0,0 +1,111 @@
+# MemAgent on ReLax
+
+This example trains Qwen3-4B to update a bounded textual memory while reading a long document chunk by chunk. Every memory-update turn and the final-answer turn is saved as an independent training row. Only the final boxed answer receives a rule-based reward; GRPO normalization happens before the trajectory is expanded.
+
+The reward summary reads the trajectory-level 0/1 outcome from `rollout_result/train/<rollout_id>.jsonl`, before memory turns are expanded. This avoids weighting long trajectories more heavily and is robust to Ray log de-duplication. `run-pipeline.sh` then runs `summarize_reward.py` and writes `training-reward.summary.json`, containing every raw point, the first/last-window means, their delta, and the peak. It rejects a run whose rollout ids are incomplete instead of producing a partial trend.
+
+The reproducibility contract is frozen to:
+
+- model: `Qwen/Qwen3-4B@1cfa9a7208912126459214e8b04321603b3df60c`;
+- dataset: `BytedTsinghua-SIA/hotpotqa@27275ff4fee67ac0acb6478e405e7ac07efbdc1a`;
+- chunk/memory/final limits: 2048/1024/256 tokens, at most 64 chunks;
+- GRPO group size 8, split credit, LR `1e-6`, KL coefficient `0.001`;
+- 100 rollout steps with checkpoints every 50 steps.
+
+## Prepare model and data
+
+Download the exact model revision to a local directory with your preferred Hugging Face client. Then prepare all frozen train/eval files and their SHA-256 manifest:
+
+```bash
+DATA_DIR=/data/mem-agent bash examples/mem_agent/prepare-data.sh
+```
+
+`prepare-data.sh` downloads `hotpotqa_train_32k.parquet`, `hotpotqa_dev.parquet`, and `eval_50/200/800.json` at the pinned dataset revision. It writes converted JSONL files plus `artifact_manifest.json`.
+
+## Train
+
+```bash
+MODEL_PATH=/data/models/Qwen3-4B \
+DATA_DIR=/data/mem-agent \
+SAVE_DIR=/data/checkpoints/mem-agent-relax \
+bash examples/mem_agent/run-qwen3-4B-train.sh
+```
+
+For the required two-step correctness smoke, add `NUM_ROLLOUT=2` and use the same command. A smoke run validates the pipeline but is not an effects result.
+
+The train-side SGLang context envelope is 8192 tokens because each request is an independent chunk turn (2K chunk + 1K memory + response), not the concatenated trajectory. This retains the frozen 9216-token per-GPU packing budget and sample-mean loss used for split credit.
+
+## Convert and evaluate
+
+```bash
+MODEL_PATH=/data/models/Qwen3-4B \
+CHECKPOINT_DIR=/data/checkpoints/mem-agent-relax \
+CHECKPOINT_TAG=iter_0000099 \
+bash examples/mem_agent/convert-to-hf.sh
+
+MODEL_PATH=/data/checkpoints/mem-agent-relax-HF/iter_0000099 \
+TOKENIZER_PATH=/data/models/Qwen3-4B \
+DATA_DIR=/data/mem-agent \
+RESULTS_DIR=/data/results/mem-agent-relax \
+bash examples/mem_agent/run-eval.sh
+```
+
+The evaluator writes raw per-sample JSONL and a summary JSON for HotpotQA dev and RULER-HQA 50/200/800. Its 64-chunk limit is the effective value of fixed VIME's official `run-eval.sh`: that script sources `_common.sh`, which exports `MEM_MAX_CHUNKS=64`, even though the Python evaluator alone has a 512 fallback. Failed requests keep their ground truth in the raw file and remain in the denominator with score zero. Formal comparison additionally rejects empty runs and any run with request errors. Each summary records the normalized input file SHA-256 and evaluator schema version, so equal paths with different bytes or incompatible evaluator revisions cannot be compared. `boxed_em_pct` is the HotpotQA reward-compatible accuracy and `sub_em_pct` is the primary VIME-compatible RULER-HQA metric. Set `MODE=base` to run the optional single-context diagnostic; its context truncation always preserves the question and answer instruction.
+
+`TOKENIZER_PATH` should point to the frozen base snapshot. `run-pipeline.sh` preserves it automatically before switching `MODEL_PATH` to the converted checkpoint. When `NUM_ROLLOUT=2` is used, the pipeline also selects `iter_0000001` automatically instead of the 100-step default `iter_0000099`.
+
+For the VIME reproduction tolerance, evaluate an official VIME checkpoint when one is available; otherwise use a checkpoint produced once from the fixed VIME recipe. The acceptance runner holds the tokenizer, data, prompts, sampling parameters, recurrent inference mode, and evaluator constant across frozen base, VIME, and ReLax. It evaluates RULER-HQA 50/200/800 by default; `LENGTHS` can freeze a smaller pre-agreed subset before any result is observed. It exits non-zero unless every selected VIME/ReLax `sub_em_pct` gap is at most 3 percentage points, ReLax beats frozen base on every selected RULER-HQA `sub_em_pct`, and ReLax beats frozen base on HotpotQA `boxed_em_pct`. Raw per-sample files are retained for review:
+
+```bash
+BASE_MODEL_PATH=/data/models/Qwen3-4B \
+VIME_MODEL_PATH=/data/checkpoints/vime-hf \
+RELAX_MODEL_PATH=/data/checkpoints/mem-agent-relax-hf \
+TOKENIZER_PATH=/data/models/Qwen3-4B \
+DATA_DIR=/data/mem-agent \
+RESULTS_DIR=/data/results/vime-vs-relax \
+LENGTHS="50 200 800" \
+bash examples/mem_agent/run-paired-eval.sh
+```
+
+## One-command chain
+
+With the environment paths set, `run-pipeline.sh` executes data preparation, training, checkpoint conversion, and evaluation in order:
+
+```bash
+MODEL_PATH=/data/models/Qwen3-4B \
+DATA_DIR=/data/mem-agent \
+SAVE_DIR=/data/checkpoints/mem-agent-relax \
+RESULTS_DIR=/data/results/mem-agent-relax \
+bash examples/mem_agent/run-pipeline.sh
+```
+
+GPU execution is intentionally not started by the CPU test suite. The caller remains responsible for starting the ReLax/Ray environment described by the repository deployment guide.
+
+## Qwen3-0.6B single-4090 pilot
+
+The 0.6B recipe is a low-cost pipeline and learnability diagnostic; it does not replace the frozen Qwen3-4B VIME/ReLax acceptance run. It fixes `Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca`, disables Qwen3 thinking, uses 512-token chunks, a 128-token memory, a 64-token final answer, and at most four chunks.
+
+Prepare 24 immutable 2--4-chunk candidates before allocating a GPU:
+
+```bash
+TOKENIZER_PATH=/data/models/Qwen3-0.6B \
+DATA_DIR=/data/task36-pilot \
+bash examples/mem_agent/prepare-pilot-candidates.sh
+```
+
+The baseline samples every candidate eight times. Selection fails unless at least 12 prompts have both a success and a failure; 2--6 successes out of 8 are preferred. Eight prompts become the training split and four disjoint prompts become the held-out pilot split. This Pass@N-screened set deliberately supplies GRPO reward variance and must not be reported as an unbiased HotpotQA metric.
+
+```bash
+MODEL_PATH=/data/models/Qwen3-0.6B \
+DATA_DIR=/data/task36-pilot \
+RESULTS_DIR=/data/task36-runs/baseline \
+bash examples/mem_agent/run-qwen3-0.6B-baseline.sh
+
+MODEL_PATH=/data/models/Qwen3-0.6B \
+DATA_DIR=/data/task36-pilot \
+RUN_ROOT=/data/task36-runs/train \
+NUM_ROLLOUT=2 \
+bash examples/mem_agent/run-qwen3-0.6B-train.sh
+```
+
+The train script is TP=1 and produces the complete Ray job log, TensorBoard events, `training-reward.summary.json`, exact `training-reward.csv` points, and `training-reward.svg`. Run the two-step smoke first. Only after its generated/transferred/consumed row counts agree should a longer run start. Converted checkpoints are evaluated against the frozen pilot split with the same seed, sampling parameters, prompt path, and tokenizer via `run-qwen3-0.6B-eval.sh`.

examples/mem_agent/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+"""MemAgent recurrent-memory training and evaluation example."""

examples/mem_agent/compare_results.py

@@ -0,0 +1,176 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+"""Compare paired VIME and ReLax evaluation summaries."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+
+
+COMPATIBILITY_FIELDS = (
+    "data_file",
+    "data_sha256",
+    "evaluator_schema_version",
+    "mode",
+    "tokenizer",
+    "temperature",
+    "top_p",
+    "sampling_count",
+    "seed",
+    "enable_thinking",
+    "chunk_tokens",
+    "max_memory_tokens",
+    "max_final_tokens",
+    "max_chunks",
+    "max_input_tokens",
+    "server_max_model_len",
+    "total",
+)
+COMPLETENESS_FIELDS = ("successful", "errors")
+
+
+def validate_compatible_summaries(*summaries: dict[str, Any]) -> None:
+    """Reject incomplete runs or mismatched controlled evaluation fields."""
+    if len(summaries) < 2:
+        raise ValueError("At least two summaries are required for compatibility validation.")
+    for summary in summaries:
+        for field in COMPLETENESS_FIELDS:
+            if field not in summary:
+                raise KeyError(f"Completeness field {field!r} must exist in every summary.")
+        total = int(summary.get("total", 0))
+        successful = int(summary["successful"])
+        errors = int(summary["errors"])
+        # A request error is retained as a zero-score row for diagnosis, but
+        # an effects claim must come from a non-empty, fully completed run.
+        if total <= 0 or errors != 0 or successful != total:
+            raise ValueError(
+                f"Evaluation summary is incomplete: total={total}, successful={successful}, errors={errors}."
+            )
+    for field in COMPATIBILITY_FIELDS:
+        if any(field not in summary for summary in summaries):
+            raise KeyError(f"Compatibility field {field!r} must exist in every summary.")
+        values = [summary[field] for summary in summaries]
+        if any(value != values[0] for value in values[1:]):
+            raise ValueError(f"Evaluation summaries differ on controlled field {field!r}: {values}")
+
+
+def compare_pair(
+    label: str,
+    vime_summary: dict[str, Any],
+    relax_summary: dict[str, Any],
+    metric: str = "sub_em_pct",
+    tolerance_pp: float = 3.0,
+) -> dict[str, Any]:
+    """Build one auditable percentage-point comparison.
+
+    Both summaries must come from the same evaluator/data recipe. The function
+    intentionally compares the reported percentage field directly: ``3.0``
+    therefore means three percentage points, not a three-percent relative gap.
+    """
+    if metric not in vime_summary or metric not in relax_summary:
+        raise KeyError(f"Metric {metric!r} must exist in both summaries.")
+    vime_value = float(vime_summary[metric])
+    relax_value = float(relax_summary[metric])
+    gap_pp = abs(relax_value - vime_value)
+    return {
+        "label": label,
+        "metric": metric,
+        "vime": vime_value,
+        "relax": relax_value,
+        "absolute_gap_pp": gap_pp,
+        "tolerance_pp": tolerance_pp,
+        "passed": gap_pp <= tolerance_pp,
+    }
+
+
+def compare_baseline(
+    label: str,
+    base_summary: dict[str, Any],
+    relax_summary: dict[str, Any],
+    metric: str,
+) -> dict[str, Any]:
+    """Require the trained ReLax checkpoint to strictly beat frozen base."""
+    if metric not in base_summary or metric not in relax_summary:
+        raise KeyError(f"Metric {metric!r} must exist in both summaries.")
+    base_value = float(base_summary[metric])
+    relax_value = float(relax_summary[metric])
+    improvement_pp = relax_value - base_value
+    return {
+        "label": label,
+        "metric": metric,
+        "base": base_value,
+        "relax": relax_value,
+        "improvement_pp": improvement_pp,
+        "passed": improvement_pp > 0,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--pair",
+        nargs=3,
+        action="append",
+        metavar=("LABEL", "VIME_SUMMARY", "RELAX_SUMMARY"),
+        default=[],
+        help="Repeat for every RULER-HQA length selected for acceptance.",
+    )
+    parser.add_argument("--metric", default="sub_em_pct")
+    parser.add_argument("--tolerance-pp", type=float, default=3.0)
+    parser.add_argument(
+        "--baseline-pair",
+        nargs=4,
+        action="append",
+        default=[],
+        metavar=("LABEL", "METRIC", "BASE_SUMMARY", "RELAX_SUMMARY"),
+        help="Require ReLax to strictly exceed frozen base; repeat for every required metric/dataset.",
+    )
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    if not args.pair and not args.baseline_pair:
+        parser.error("At least one --pair or --baseline-pair is required.")
+
+    comparisons = []
+    for label, vime_path, relax_path in args.pair:
+        with Path(vime_path).open(encoding="utf-8") as source:
+            vime_summary = json.load(source)
+        with Path(relax_path).open(encoding="utf-8") as source:
+            relax_summary = json.load(source)
+        validate_compatible_summaries(vime_summary, relax_summary)
+        comparison = compare_pair(label, vime_summary, relax_summary, args.metric, args.tolerance_pp)
+        comparison["vime_summary"] = str(vime_path)
+        comparison["relax_summary"] = str(relax_path)
+        comparisons.append(comparison)
+
+    baseline_comparisons = []
+    for label, metric, base_path, relax_path in args.baseline_pair:
+        with Path(base_path).open(encoding="utf-8") as source:
+            base_summary = json.load(source)
+        with Path(relax_path).open(encoding="utf-8") as source:
+            relax_summary = json.load(source)
+        validate_compatible_summaries(base_summary, relax_summary)
+        comparison = compare_baseline(label, base_summary, relax_summary, metric)
+        comparison["base_summary"] = str(base_path)
+        comparison["relax_summary"] = str(relax_path)
+        baseline_comparisons.append(comparison)
+
+    report = {
+        "metric": args.metric,
+        "tolerance_pp": args.tolerance_pp,
+        "passed": all(item["passed"] for item in comparisons + baseline_comparisons),
+        "comparisons": comparisons,
+        "baseline_comparisons": baseline_comparisons,
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open("w", encoding="utf-8") as destination:
+        json.dump(report, destination, ensure_ascii=False, indent=2)
+        destination.write("\n")
+    print(json.dumps(report, ensure_ascii=False, indent=2))
+    if not report["passed"]:
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/mem_agent/config-pilot-qwen3-0.6b.yaml

@@ -0,0 +1,15 @@
+# Qwen3-0.6B/RTX 4090 diagnostic recipe. This does not replace config.yaml's
+# frozen Qwen3-4B/VIME acceptance contract.
+mem_agent_chunk_tokens: 512
+mem_agent_max_memory_tokens: 128
+mem_agent_max_final_tokens: 64
+mem_agent_max_chunks: 4
+mem_agent_enable_thinking: false
+mem_agent_credit_assignment: split
+# Four memory turns plus one final turn is the maximum pilot expansion.
+custom_train_sample_expansion_factor: 5
+custom_train_data_group_size: 1
+custom_train_expanded_batch: true
+mem_agent_train_rows_multiple: 1
+model_id: Qwen/Qwen3-0.6B
+model_revision: c1899de289a04d12100db370d81485cdf75e47ca

examples/mem_agent/config.yaml

@@ -0,0 +1,15 @@
+mem_agent_chunk_tokens: 2048
+mem_agent_max_memory_tokens: 1024
+mem_agent_max_final_tokens: 256
+mem_agent_max_chunks: 64
+mem_agent_credit_assignment: split
+# One trajectory has at most 64 memory turns plus one final-answer turn.
+# These fields reserve enough queue capacity, disable post-expansion GRPO
+# regrouping, and tell the actor to consume every converted row in one step.
+custom_train_sample_expansion_factor: 65
+custom_train_data_group_size: 1
+custom_train_expanded_batch: true
+# TP=2 on eight actor GPUs gives four data-parallel consumers (CP stays 1).
+mem_agent_train_rows_multiple: 4
+model_id: Qwen/Qwen3-4B
+model_revision: 1cfa9a7208912126459214e8b04321603b3df60c

examples/mem_agent/contracts.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+"""Compatibility guards for the MemAgent trajectory contract."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def require_strict_alignment(args: Any) -> None:
+    """Preserve rejection of the retired false-valued compatibility option.
+
+    Alignment validation is always strict and unconditional. The bundled
+    recipes therefore no longer advertise a boolean switch, while older
+    external configs that explicitly requested an unsupported relaxed mode
+    continue to fail instead of silently changing behavior.
+    """
+    if not getattr(args, "mem_agent_strict_alignment", True):
+        raise ValueError("MemAgent training requires mem_agent_strict_alignment=true.")

examples/mem_agent/convert-to-hf.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+RELAX_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)"
+MODEL_PATH="${MODEL_PATH:?Set MODEL_PATH to the frozen base model directory.}"
+CHECKPOINT_DIR="${CHECKPOINT_DIR:?Set CHECKPOINT_DIR to the ReLax checkpoint root.}"
+CHECKPOINT_TAG="${CHECKPOINT_TAG:-iter_0000099}"
+HF_OUTPUT_DIR="${HF_OUTPUT_DIR:-${CHECKPOINT_DIR}-HF/${CHECKPOINT_TAG}}"
+
+python3 "${RELAX_ROOT}/scripts/tools/convert_torch_dist_to_hf_bridge.py" \
+  --input-dir "${CHECKPOINT_DIR}/${CHECKPOINT_TAG}" \
+  --output-dir "${HF_OUTPUT_DIR}" \
+  --origin-hf-dir "${MODEL_PATH}"
+
+echo "Converted checkpoint: ${HF_OUTPUT_DIR}"