From f82990b0ed7f36709594bd478279f13e8e5649c1 Mon Sep 17 00:00:00 2001
From: HITSZJiahaoLiu <97160673+HITLiuJiahao@users.noreply.github.com>
Date: Tue, 22 Sep 2026 23:17:32 +0800
Subject: [PATCH] feat: add multi-turn tool-use on-policy distillation

---
 README.md                  |   2 +
 README_en.md               |   2 +
 docs/agent_opd.md          | 110 ++++++++++++++
 tests/test_agent_opd.py    | 294 +++++++++++++++++++++++++++++++++++++
 trainer/agent_opd_utils.py | 158 ++++++++++++++++++++
 trainer/train_agent_opd.py | 235 +++++++++++++++++++++++++++++
 6 files changed, 801 insertions(+)
 create mode 100644 docs/agent_opd.md
 create mode 100644 tests/test_agent_opd.py
 create mode 100644 trainer/agent_opd_utils.py
 create mode 100644 trainer/train_agent_opd.py

diff --git a/README.md b/README.md
index dadc82c..85f65e3 100644
--- a/README.md
+++ b/README.md
@@ -750,6 +750,8 @@ python eval_llm.py --weight full_sft
 
 仓库中提供的 `train_distillation.py` 更适合作为理解白盒蒸馏流程的参考实现：它完整展示了教师/学生双模型加载、`CE + KL` 混合损失、温度缩放、MoE 与 dense 组合蒸馏，以及断点续训和分布式训练等关键细节。
 
+多轮工具交互的在线蒸馏参考实现见 [`train_agent_opd.py`](trainer/train_agent_opd.py)：学生执行工具并读取结果，冻结教师沿着学生实际生成的交互轨迹提供 token 分布监督。使用方式、单设备限制与验证范围见 [Agent OPD 指南](docs/agent_opd.md)。
+
 ```bash
 # 方式1
 torchrun --nproc_per_node 1 train_distillation.py
diff --git a/README_en.md b/README_en.md
index 95fb03a..8b91899 100644
--- a/README_en.md
+++ b/README_en.md
@@ -750,6 +750,8 @@ White-box distillation typically fits the teacher distribution in addition to th
 
 The `train_distillation.py` script is intended as a reference implementation for understanding the white-box distillation pipeline: it demonstrates teacher/student dual model loading, `CE + KL` mixed loss, temperature scaling, MoE and dense combination distillation, as well as key details such as checkpoint resume and distributed training.
 
+For on-policy distillation through multi-turn tool interactions, see [`train_agent_opd.py`](trainer/train_agent_opd.py). The student executes tools and reads their results while a frozen teacher supervises token distributions along the student's sampled trajectory. See the [Agent OPD guide](docs/agent_opd.md) for usage, single-device scope and validation limits.
+
 ```bash
 # Method 1
 torchrun --nproc_per_node 1 train_distillation.py
diff --git a/docs/agent_opd.md b/docs/agent_opd.md
new file mode 100644
index 0000000..731339c
--- /dev/null
+++ b/docs/agent_opd.md
@@ -0,0 +1,110 @@
+# Multi-turn tool-use on-policy distillation
+
+`trainer/train_agent_opd.py` is a single-device reference trainer for distilling a
+frozen MiniMind teacher along the student's own tool interactions. It uses the
+local tool implementations in `train_agent.py`; weather, exchange rates and other
+mock tool results are **not live services**. It does not load a reward model.
+
+The student generates a turn, its tool calls are executed, and the returned
+observations condition the next student turn. The teacher sees exactly these
+student-generated token prefixes and observations. No teacher answer replaces a
+student rollout. Prompt text, inserted role markers, tool observations and padding
+are context only. Every generated assistant token, including each turn's first
+EOS, is a distillation target. Later assistant turns remain targets after earlier
+EOS tokens.
+
+## Objective
+
+Following the on-policy data collection idea in
+[GKD (Agarwal et al., ICLR 2024)](https://arxiv.org/abs/2306.13649), the default loss is
+the mean full-vocabulary `T² KL(teacher_T || student_T)` over assistant tokens on
+the sampled interaction. `--kl_direction reverse_kl` instead uses
+`T² KL(student_T || teacher_T)` at the same prefixes. This is a GKD-style
+semi-gradient: sampling and environment execution are detached. It is not a
+REINFORCE estimator of trajectory-level reverse KL.
+
+Sampling defaults to temperature 1, top-p 1 and top-k 0. `--temperature` controls
+student sampling; `--distill_temperature` controls only the loss distributions.
+Both checkpoints must be native MiniMind weights using the **same tokenizer and
+token-to-ID mapping**. A teacher with an unrelated vocabulary is not supported.
+
+## Run
+
+Run from the repository root, after installing the repository dependencies:
+
+```bash
+python trainer/train_agent_opd.py \
+  --student_checkpoint out/full_sft_768.pth \
+  --teacher_checkpoint out/agent_768.pth \
+  --data_path dataset/agent_rl.jsonl \
+  --device cuda:0 --dtype bfloat16 \
+  --max_turns 3 --max_gen_len 256 --max_total_len 2048 \
+  --batch_size 1 --accumulation_steps 4 --max_steps 100
+```
+
+Use an independently trained, stronger teacher. The example filenames refer to
+weights you already have; the trainer does not download them. Architecture flags
+(`--student_hidden_size`, `--teacher_num_layers`, `--teacher_use_moe`, etc.) must
+match their checkpoints. Weights are loaded strictly. For a CPU correctness run,
+use `--device cpu --dtype float32` and small checkpoints. This initial trainer
+supports one process/device; it does not provide DDP or SGLang training.
+
+Data may use `conversations` (as in `agent_rl.jsonl`) or `messages`. Tool schemas
+can appear in the system message's `tools` field or the record's top-level
+`tools` field. A trailing fixed assistant answer is ignored. `gt` is not required:
+
+```json
+{"messages":[{"role":"user","content":"请用工具计算 2+2。"}],"tools":[{"type":"function","function":{"name":"calculate_math","description":"计算数学表达式","parameters":{"type":"object","properties":{"expression":{"type":"string"}},"required":["expression"]}}}]}
+```
+
+Only tools declared for the sample are available. Unknown tools and invalid
+arguments produce an error observation. Context and turn limits stop further
+interaction while retaining already sampled targets. Context is never silently
+left-truncated into a different teacher prefix. Oversized initial prompts raise
+an error; adjust the data or `--max_total_len`.
+
+The output includes an inference weight file, `out/agent_opd/agent_opd_768.pth`,
+and a full-precision `agent_opd_768_resume.pth` containing optimizer/scaler state,
+the next data cursor and random states. Checkpoints are saved after complete
+optimizer updates, including a final partial accumulation window. Resume with
+the original arguments plus:
+
+```bash
+--resume out/agent_opd/agent_opd_768_resume.pth --max_steps 200
+```
+
+`--max_steps` counts total optimizer steps, including resumed steps. Keep the
+teacher, data, architecture, optimizer and sampling arguments unchanged. The
+logged `loss`, `mean_turns`, `tool_calls`, `assistant_tokens` and
+`truncated_fraction` describe training, not task success rates.
+
+## Validation and limits
+
+```bash
+OMP_NUM_THREADS=1 HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false \
+  python -m unittest discover -s tests -p test_agent_opd.py -v
+```
+
+Tests use the real repository tokenizer, local tools and tiny CPU MiniMind
+models. Scripted generation covers deterministic multi-turn boundaries, including
+tool errors, per-turn EOS and multiple calls; other tests exercise actual model
+generation, optimization, teacher freezing and exact CPU checkpoint continuation.
+A controlled two-turn optimization test verifies that teacher divergence falls.
+No pretrained task-quality improvement or CUDA/MPS performance result is claimed.
+
+Before claiming better tool use, evaluate the starting student, teacher and OPD
+checkpoint on held-out prompts with identical decoding and tool limits. Compare
+against fixed-data distillation as well, and report final-answer accuracy,
+successful tool execution, truncation, wall time and peak memory. A falling
+distillation loss alone is not evidence of better answers.
+
+## 中文说明
+
+这是面向多轮工具交互的 OPD 参考实现：学生自己生成工具调用，读取工具结果后继续生成，
+教师沿着学生实际经历的上下文提供分布监督。工具结果参与上下文，但不作为预测标签；每轮
+学生生成的 EOS 都参与损失，后续轮次不会因第一轮 EOS 而被屏蔽。默认使用前向 KL，支持
+反向 KL 的 GKD 式半梯度。教师被冻结，数据不需要 `gt` 或标准答案。
+
+首版仅支持单进程、共享分词器的 MiniMind 教师与学生，复用仓库的本地模拟工具。
+请根据实际权重指定教师/学生结构。测试验证训练链路和数值正确性；正式模型上的能力提升
+仍需独立留出集实验。断点文件保存完整精度和随机状态，恢复时沿用原始参数并加上 `--resume`。
diff --git a/tests/test_agent_opd.py b/tests/test_agent_opd.py
new file mode 100644
index 0000000..11aefc1
--- /dev/null
+++ b/tests/test_agent_opd.py
@@ -0,0 +1,294 @@
+"""CPU regressions with the real tokenizer, tools and tiny MiniMind models."""
+
+import contextlib
+import io
+import json
+import os
+from pathlib import Path
+import tempfile
+import unittest
+from unittest.mock import patch
+
+os.environ.setdefault("HF_HUB_OFFLINE", "1")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+import torch
+from transformers import AutoTokenizer
+
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from trainer.agent_opd_utils import (
+    AgentTrajectory, collect_trajectory, distillation_loss, pack_trajectories,
+)
+from trainer.train_agent import TOOLS
+from trainer import train_agent_opd as trainer
+
+ROOT = Path(__file__).resolve().parents[1]
+MATH_TOOLS = [TOOLS[0]]
+MESSAGES = [{"role": "user", "content": "请用工具计算 2+2。"}]
+CALL = '<tool_call>\n{"name":"calculate_math","arguments":{"expression":"2+2"}}\n</tool_call>'
+
+
+class ScriptedStudent(torch.nn.Module):
+    """Deterministic generation at the model boundary, keeping real tool/token IO."""
+
+    def __init__(self, tokenizer, responses):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.zeros(()))
+        self.responses = responses
+        self.tokenizer = tokenizer
+        self.inputs, self.options = [], []
+
+    def generate(self, input_ids, **kwargs):
+        self.inputs.append(input_ids[0].tolist())
+        self.options.append(kwargs)
+        response = self.responses[len(self.inputs) - 1]
+        ids = self.tokenizer(response, add_special_tokens=False)["input_ids"] if isinstance(response, str) else response
+        ids = ids[:kwargs["max_new_tokens"]]
+        return torch.cat((input_ids, input_ids.new_tensor([ids])), dim=1)
+
+
+class TrajectoryTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        torch.set_num_threads(1)
+        cls.tokenizer = AutoTokenizer.from_pretrained(ROOT / "model")
+        cls.eos = cls.tokenizer.eos_token
+
+    def sample(self, responses, **kwargs):
+        student = ScriptedStudent(self.tokenizer, responses)
+        trajectory = collect_trajectory(student, self.tokenizer, MESSAGES, MATH_TOOLS, **kwargs)
+        return student, trajectory
+
+    def test_tool_result_conditions_next_turn_but_is_not_a_target(self):
+        student, trajectory = self.sample([CALL + self.eos, "结果是 4。" + self.eos])
+        self.assertEqual(trajectory.turns, 2)
+        self.assertEqual(trajectory.tool_calls, 1)
+        self.assertFalse(trajectory.truncated)
+        self.assertIn('"result": "4"', self.tokenizer.decode(student.inputs[1]))
+        assistant_ids = [t for t, keep in zip(trajectory.input_ids, trajectory.assistant_mask) if keep]
+        expected = self.tokenizer(CALL + self.eos, add_special_tokens=False)["input_ids"]
+        expected += self.tokenizer("结果是 4。" + self.eos, add_special_tokens=False)["input_ids"]
+        self.assertEqual(assistant_ids, expected)
+        self.assertEqual(assistant_ids.count(self.tokenizer.eos_token_id), 2)
+        self.assertTrue(student.training)
+        for options in student.options:
+            self.assertEqual((options["top_p"], options["top_k"], options["temperature"]), (1.0, 0, 1.0))
+        self.assertEqual(MESSAGES, [{"role": "user", "content": "请用工具计算 2+2。"}])
+
+    def test_preserves_exact_sampled_ids_and_first_eos_only(self):
+        sampled = self.tokenizer(CALL, add_special_tokens=False)["input_ids"] + [self.tokenizer.eos_token_id] * 3
+        student, trajectory = self.sample([sampled, [self.tokenizer.eos_token_id]])
+        start = len(student.inputs[0])
+        self.assertEqual(trajectory.input_ids[start:start + len(sampled) - 2], sampled[:-2])
+        self.assertEqual(sum(trajectory.assistant_mask), len(sampled) - 1)
+
+    def test_multiple_tool_calls_return_separate_observations(self):
+        second = CALL.replace("2+2", "3*3")
+        student, trajectory = self.sample([CALL + "\n" + second + self.eos, "4 和 9。" + self.eos])
+        context = self.tokenizer.decode(student.inputs[1])
+        self.assertIn('"result": "4"', context)
+        self.assertIn('"result": "9"', context)
+        self.assertEqual(trajectory.tool_calls, 2)
+
+    def test_unknown_tool_and_malformed_arguments_become_observations(self):
+        for call in [CALL.replace("calculate_math", "get_current_weather"),
+                     '<tool_call>{"name":[],"arguments":null}</tool_call>',
+                     '<tool_call>{"name":"calculate_math","arguments":"invalid"}</tool_call>']:
+            with self.subTest(call=call):
+                student, trajectory = self.sample([call + self.eos, "无法计算。" + self.eos])
+                self.assertIn("unknown tool or invalid arguments", self.tokenizer.decode(student.inputs[1]))
+                self.assertEqual(trajectory.turns, 2)
+
+    def test_non_dictionary_tool_json_does_not_crash(self):
+        _, trajectory = self.sample(['<tool_call>[]</tool_call>' + self.eos])
+        self.assertEqual(trajectory.turns, 1)
+
+    def test_turn_limit_keeps_generated_tokens_without_executing_more_tools(self):
+        with patch("trainer.agent_opd_utils.execute_tool") as execute:
+            _, trajectory = self.sample([CALL + self.eos], max_turns=1)
+        self.assertTrue(trajectory.truncated)
+        execute.assert_not_called()
+        self.assertTrue(trajectory.assistant_mask[-1])
+
+    def test_context_limit_never_slices_an_observation_into_a_training_target(self):
+        probe = ScriptedStudent(self.tokenizer, [CALL + self.eos])
+        prompt = self.tokenizer.apply_chat_template(MESSAGES, tools=MATH_TOOLS, tokenize=False,
+                                                    add_generation_prompt=True, open_thinking=False)
+        limit = len(self.tokenizer(prompt, add_special_tokens=False)["input_ids"]) + 8
+        trajectory = collect_trajectory(probe, self.tokenizer, MESSAGES, MATH_TOOLS, max_total_len=limit)
+        self.assertEqual(len(trajectory.input_ids), limit)
+        self.assertTrue(trajectory.truncated)
+        with self.assertRaisesRegex(ValueError, "Prompt leaves no generation space"):
+            collect_trajectory(probe, self.tokenizer, MESSAGES, MATH_TOOLS, max_total_len=1)
+
+    def test_empty_final_answer_still_trains_eos(self):
+        _, trajectory = self.sample([[self.tokenizer.eos_token_id]])
+        self.assertEqual(sum(trajectory.assistant_mask), 1)
+        self.assertFalse(trajectory.truncated)
+
+    def test_real_model_can_generate_an_on_policy_trajectory(self):
+        torch.manual_seed(3)
+        student = MiniMindForCausalLM(MiniMindConfig(hidden_size=32, num_hidden_layers=1))
+        trajectory = collect_trajectory(student, self.tokenizer, MESSAGES, max_new_tokens=2)
+        self.assertGreaterEqual(sum(trajectory.assistant_mask), 1)
+        self.assertLessEqual(sum(trajectory.assistant_mask), 2)
+
+
+class LossTests(unittest.TestCase):
+    def test_shifted_masks_include_both_turns_and_exclude_observations_and_padding(self):
+        trajectories = [AgentTrajectory([1, 10, 2, 30, 31, 11, 2], [0, 1, 1, 0, 0, 1, 1], 2, 1, False),
+                        AgentTrajectory([1, 12, 2], [0, 1, 1], 1, 0, False)]
+        ids, attention, targets = pack_trajectories(trajectories, 2, "cpu")
+        self.assertEqual(targets.tolist(), [[True, True, False, False, True, True],
+                                          [True, True, False, False, False, False]])
+        self.assertEqual(attention[1].tolist(), [1, 1, 1, 0, 0, 0, 0])
+        self.assertEqual(ids[1, -1].item(), 2)
+
+    def test_kl_directions_match_independent_probability_calculation(self):
+        student = torch.tensor([[[0.2, -0.4, 1.0]]], requires_grad=True)
+        teacher = torch.tensor([[[0.8, -0.2, 0.1]]], requires_grad=True)
+        for direction in ("forward_kl", "reverse_kl"):
+            with self.subTest(direction=direction):
+                temperature = 1.7
+                ps, pt = (student.detach() / temperature).softmax(-1), (teacher.detach() / temperature).softmax(-1)
+                p, q = (pt, ps) if direction == "forward_kl" else (ps, pt)
+                expected = (p * (p.log() - q.log())).sum() * temperature ** 2
+                loss = distillation_loss(student, teacher, torch.ones(1, 1, dtype=torch.bool), temperature, direction)
+                torch.testing.assert_close(loss, expected)
+                loss.backward()
+                self.assertIsNone(teacher.grad)
+                self.assertGreater(student.grad.abs().sum().item(), 0)
+                student.grad = None
+
+    def test_masked_tokens_have_no_direct_gradient_even_with_invalid_logits(self):
+        student = torch.randn(1, 5, 4, requires_grad=True)
+        teacher = torch.randn(1, 5, 4, requires_grad=True)
+        mask = torch.tensor([[True, False, False, True, False]])
+        with torch.no_grad():
+            student[:, 1] = float("nan")
+            teacher[:, 1] = float("nan")
+        loss = distillation_loss(student, teacher, mask)
+        self.assertTrue(torch.isfinite(loss))
+        loss.backward()
+        self.assertEqual(student.grad[~mask].abs().sum().item(), 0)
+        self.assertIsNone(teacher.grad)
+
+    def test_empty_mask_is_finite_differentiable_zero(self):
+        student = torch.randn(2, 3, 5, requires_grad=True)
+        loss = distillation_loss(student, student.detach(), torch.zeros(2, 3, dtype=torch.bool))
+        loss.backward()
+        self.assertEqual(loss.item(), 0)
+        self.assertEqual(student.grad.abs().sum().item(), 0)
+
+
+class TrainingTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        torch.set_num_threads(1)
+
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.tmp.cleanup)
+        self.directory = Path(self.tmp.name)
+        self.student_path, self.teacher_path = self.directory / "student.pth", self.directory / "teacher.pth"
+        config = MiniMindConfig(hidden_size=32, num_hidden_layers=1)
+        torch.manual_seed(1)
+        torch.save(MiniMindForCausalLM(config).state_dict(), self.student_path)
+        torch.manual_seed(2)
+        torch.save(MiniMindForCausalLM(config).state_dict(), self.teacher_path)
+        self.data = self.directory / "prompts.jsonl"
+        self.data.write_text("\n".join(json.dumps({"messages": MESSAGES}) for _ in range(3)) + "\n")
+
+    def args(self, name="run", extra=()):
+        return trainer.parse_args([
+            "--student_checkpoint", str(self.student_path), "--teacher_checkpoint", str(self.teacher_path),
+            "--data_path", str(self.data), "--output_dir", str(self.directory / name),
+            "--student_hidden_size", "32", "--student_num_layers", "1",
+            "--teacher_hidden_size", "32", "--teacher_num_layers", "1",
+            "--device", "cpu", "--max_gen_len", "2", "--learning_rate", "0.001",
+            *extra,
+        ])
+
+    def run_train(self, args):
+        with contextlib.redirect_stdout(io.StringIO()):
+            return trainer.train(args)
+
+    def checkpoint(self, name):
+        return torch.load(self.directory / name / "agent_opd_32_resume.pth", weights_only=True)
+
+    def test_actual_training_updates_student_only_and_saves_residual_window(self):
+        loaded = []
+        original = trainer.load_model
+
+        def track(*args):
+            model = original(*args)
+            loaded.append(model)
+            return model
+
+        args = self.args(extra=["--accumulation_steps", "2"])
+        with patch.object(trainer, "load_model", side_effect=track):
+            metrics = self.run_train(args)
+        self.assertEqual(len(metrics), 2)  # 3 microbatches, including the final partial window.
+        student, teacher = loaded
+        initial = torch.load(self.student_path, weights_only=True)
+        self.assertTrue(any(not torch.equal(initial[k], v) for k, v in student.state_dict().items()))
+        teacher_weights = torch.load(self.teacher_path, weights_only=True)
+        for key, value in teacher.state_dict().items():
+            torch.testing.assert_close(value, teacher_weights[key], rtol=0, atol=0)
+        self.assertTrue(all(not p.requires_grad and p.grad is None for p in teacher.parameters()))
+        saved = self.checkpoint("run")
+        self.assertEqual((saved["epoch"], saved["batch"], saved["step"]), (1, 0, 2))
+        for key, value in student.state_dict().items():
+            torch.testing.assert_close(saved["model"][key], value, rtol=0, atol=0)
+
+    def test_resumed_online_sampling_matches_uninterrupted_training(self):
+        self.run_train(self.args("full", ["--max_steps", "2"]))
+        self.run_train(self.args("resume", ["--max_steps", "1"]))
+        path = self.directory / "resume/agent_opd_32_resume.pth"
+        self.run_train(self.args("resume", ["--max_steps", "2", "--resume", str(path)]))
+        full, resumed = self.checkpoint("full"), self.checkpoint("resume")
+        self.assertEqual((full["epoch"], full["batch"], full["step"]),
+                         (resumed["epoch"], resumed["batch"], resumed["step"]))
+        for key in full["model"]:
+            torch.testing.assert_close(full["model"][key], resumed["model"][key], rtol=0, atol=0)
+
+    def test_resume_rejects_a_different_teacher_or_sampling_setup(self):
+        self.run_train(self.args(extra=["--max_steps", "1"]))
+        resume = str(self.directory / "run/agent_opd_32_resume.pth")
+        with self.assertRaisesRegex(ValueError, "same data, teacher, sampling"):
+            self.run_train(self.args(extra=["--resume", resume, "--temperature", "0.8"]))
+
+    def test_agent_dataset_needs_no_ground_truth_answer(self):
+        self.data.write_text(json.dumps({"conversations": [
+            {"role": "system", "content": "Use tools", "tools": json.dumps(MATH_TOOLS)},
+            *MESSAGES, {"role": "assistant", "content": "ignored fixed answer"},
+        ]}) + "\n")
+        messages, tools = trainer.load_prompts(self.data)[0]
+        self.assertEqual(messages[-1]["role"], "user")
+        self.assertEqual(tools, MATH_TOOLS)
+
+    def test_multi_turn_training_reduces_teacher_divergence(self):
+        # Control generation only; both models, KL, backprop and optimizer are real.
+        tokenizer = AutoTokenizer.from_pretrained(ROOT / "model")
+        self.data.write_text(json.dumps({"messages": MESSAGES, "tools": MATH_TOOLS}) + "\n")
+
+        def generate(model, input_ids, **kwargs):
+            has_result = '"result": "4"' in tokenizer.decode(input_ids[0])
+            text = ("结果是 4。" if has_result else CALL) + tokenizer.eos_token
+            new_ids = tokenizer(text, add_special_tokens=False)["input_ids"]
+            return torch.cat((input_ids, input_ids.new_tensor([new_ids])), dim=1)
+
+        args = self.args(extra=["--epochs", "12", "--max_gen_len", "128"])
+        with patch.object(MiniMindForCausalLM, "generate", new=generate):
+            metrics = self.run_train(args)
+        self.assertTrue(all(record["mean_turns"] == 2 and record["tool_calls"] == 1 for record in metrics))
+        self.assertLess(metrics[-1]["loss"], metrics[0]["loss"] * 0.5)
+
+    def test_invalid_training_parameters_fail_before_loading_models(self):
+        for flag, value in [("--grad_clip", "0"), ("--learning_rate", "nan"), ("--max_steps", "-1")]:
+            with self.subTest(flag=flag), self.assertRaises(ValueError):
+                self.run_train(self.args(extra=[flag, value]))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/trainer/agent_opd_utils.py b/trainer/agent_opd_utils.py
new file mode 100644
index 0000000..ef3b755
--- /dev/null
+++ b/trainer/agent_opd_utils.py
@@ -0,0 +1,158 @@
+"""Student-owned tool trajectories and token-level on-policy distillation."""
+
+import json
+from contextlib import nullcontext
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from trainer.train_agent import execute_tool, parse_tool_calls
+
+
+@dataclass
+class AgentTrajectory:
+    input_ids: list
+    assistant_mask: list
+    turns: int
+    tool_calls: int
+    truncated: bool
+
+
+def tool_observation(tokenizer, calls, tools, open_thinking=False):
+    """Render only the environment response and the next assistant prefix."""
+    marker = "<|agent_opd_observation_boundary|>"
+    messages = [{"role": "assistant", "content": marker}]
+    allowed = {tool["function"]["name"] for tool in (tools or [])}
+    for call in calls:
+        name, arguments = call.get("name", ""), call.get("arguments", {})
+        if isinstance(arguments, str):
+            try:
+                arguments = json.loads(arguments)
+            except json.JSONDecodeError:
+                arguments = None
+        result = None
+        if isinstance(name, str) and name in allowed and isinstance(arguments, dict):
+            result = execute_tool(name, arguments)
+        if result is None:
+            result = {"error": "unknown tool or invalid arguments"}
+        messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)[:2048]})
+    rendered = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True,
+        tools=tools, open_thinking=open_thinking,
+    )
+    _, found, observation = rendered.partition(marker)
+    if not found:
+        raise ValueError("The chat template must preserve assistant content")
+    return tokenizer(observation, add_special_tokens=False)["input_ids"]
+
+
+@torch.no_grad()
+def collect_trajectory(student, tokenizer, messages, tools=None, *, max_turns=3,
+                       max_new_tokens=256, max_total_len=2048, temperature=1.0,
+                       open_thinking=False, autocast_ctx=None):
+    """Sample each assistant turn, execute local tools, and preserve sampled IDs.
+
+    Only student-generated tokens (including each turn's first EOS) are targets.
+    Environment observations remain visible context with no direct loss.
+    """
+    if min(max_turns, max_new_tokens, max_total_len) < 1 or temperature <= 0:
+        raise ValueError("Generation limits and temperature must be positive")
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True,
+        tools=tools, open_thinking=open_thinking,
+    )
+    ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
+    if not ids or len(ids) >= max_total_len:
+        raise ValueError("Prompt leaves no generation space; increase --max_total_len")
+    mask = [False] * len(ids)
+    model = getattr(student, "module", student)
+    device = next(model.parameters()).device
+    was_training = model.training
+    model.eval()
+    turns, call_count, truncated = 0, 0, False
+    try:
+        for turn in range(max_turns):
+            inputs = torch.tensor([ids], device=device)
+            with autocast_ctx if autocast_ctx is not None else nullcontext():
+                output = model.generate(
+                    input_ids=inputs, attention_mask=torch.ones_like(inputs),
+                    max_new_tokens=min(max_new_tokens, max_total_len - len(ids)),
+                    do_sample=True, temperature=temperature, top_p=1.0, top_k=0,
+                    repetition_penalty=1.0, num_beams=1, use_cache=True,
+                    eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id,
+                )
+            new_ids = output[0, len(ids):].tolist()
+            if tokenizer.eos_token_id in new_ids:
+                new_ids = new_ids[:new_ids.index(tokenizer.eos_token_id) + 1]
+            if not new_ids:
+                truncated = True
+                break
+            ids.extend(new_ids)
+            mask.extend([True] * len(new_ids))
+            turns += 1
+            if new_ids[-1] != tokenizer.eos_token_id:
+                truncated = True
+                break
+            text = tokenizer.decode(new_ids, skip_special_tokens=True)
+            calls = [call for call in parse_tool_calls(text) if isinstance(call, dict)]
+            call_count += len(calls)
+            if not calls:
+                break
+            if turn + 1 == max_turns:
+                truncated = True
+                break
+            observation = tool_observation(tokenizer, calls, tools, open_thinking)
+            if observation[:1] == [tokenizer.eos_token_id]:
+                observation = observation[1:]  # EOS is already in the sampled IDs.
+            if len(ids) + len(observation) >= max_total_len:
+                truncated = True
+                break
+            ids.extend(observation)
+            mask.extend([False] * len(observation))
+    finally:
+        model.train(was_training)
+    return AgentTrajectory(ids, mask, turns, call_count, truncated)
+
+
+def pack_trajectories(trajectories, pad_token_id, device):
+    if not trajectories:
+        raise ValueError("At least one trajectory is required")
+    if any(len(t.input_ids) != len(t.assistant_mask) for t in trajectories):
+        raise ValueError("Trajectory tokens and masks must have the same length")
+    length = max(len(t.input_ids) for t in trajectories)
+    ids, attention, targets = [], [], []
+    for trajectory in trajectories:
+        size = len(trajectory.input_ids)
+        ids.append(trajectory.input_ids + [pad_token_id] * (length - size))
+        attention.append([1] * size + [0] * (length - size))
+        targets.append(trajectory.assistant_mask + [False] * (length - size))
+    return (
+        torch.tensor(ids, dtype=torch.long, device=device),
+        torch.tensor(attention, dtype=torch.long, device=device),
+        torch.tensor(targets, dtype=torch.bool, device=device)[:, 1:],
+    )
+
+
+def distillation_loss(student_logits, teacher_logits, target_mask,
+                      temperature=1.0, direction="forward_kl"):
+    """Full-vocabulary KL averaged over assistant tokens on sampled prefixes.
+
+    Sampling is detached (the GKD semi-gradient); this is not a REINFORCE
+    estimator of trajectory-level reverse KL. Teacher gradients are blocked.
+    """
+    if student_logits.shape != teacher_logits.shape or student_logits.ndim != 3:
+        raise ValueError("Student and teacher logits must share [batch, time, vocab]")
+    if target_mask.shape != student_logits.shape[:2]:
+        raise ValueError("Target mask must align with the shifted logits")
+    if temperature <= 0 or direction not in {"forward_kl", "reverse_kl"}:
+        raise ValueError("Invalid distillation temperature or KL direction")
+    if not target_mask.any():
+        return student_logits[:, :0, :].sum()
+    student_logps = F.log_softmax(student_logits[target_mask].float() / temperature, dim=-1)
+    teacher_logps = F.log_softmax(teacher_logits.detach()[target_mask].float() / temperature, dim=-1)
+    if direction == "forward_kl":
+        loss = F.kl_div(student_logps, teacher_logps, log_target=True, reduction="batchmean")
+    else:
+        loss = F.kl_div(teacher_logps, student_logps, log_target=True, reduction="batchmean")
+    return loss * temperature ** 2
diff --git a/trainer/train_agent_opd.py b/trainer/train_agent_opd.py
new file mode 100644
index 0000000..6063931
--- /dev/null
+++ b/trainer/train_agent_opd.py
@@ -0,0 +1,235 @@
+"""Single-device on-policy distillation through MiniMind's local tool environment."""
+
+import os
+import sys
+
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+import datasets  # noqa: F401  # Keep pyarrow before torch on Windows (issue #771).
+import argparse
+import json
+import math
+import random
+from contextlib import nullcontext
+from pathlib import Path
+
+import torch
+from transformers import AutoTokenizer
+
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from trainer.agent_opd_utils import collect_trajectory, pack_trajectories, distillation_loss
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_prompts(path):
+    samples = []
+    with open(path, encoding="utf-8") as stream:
+        for line in stream:
+            if not line.strip():
+                continue
+            sample = json.loads(line)
+            messages = sample.get("conversations", sample.get("messages", []))
+            if messages and messages[-1].get("role") == "assistant":
+                messages = messages[:-1]
+            if not messages or messages[-1].get("role") not in {"user", "tool"}:
+                raise ValueError("Each sample must end in a user/tool prompt or a removable assistant answer")
+            tools = sample.get("tools")
+            for message in messages:
+                if message.get("role") == "system" and message.get("tools"):
+                    tools = message["tools"]
+            if isinstance(tools, str):
+                tools = json.loads(tools)
+            samples.append((messages, tools))
+    if not samples:
+        raise ValueError("The prompt dataset is empty")
+    return samples
+
+
+def load_model(path, config, device):
+    weights = torch.load(path, map_location="cpu", weights_only=True)
+    weights = weights.get("model", weights)
+    model = MiniMindForCausalLM(config)
+    model.load_state_dict(weights, strict=True)
+    return model.to(device)
+
+
+def training_signature(args):
+    # A resumed run must retain its teacher, data order, and optimization setup.
+    adjustable = {"resume", "epochs", "max_steps", "save_interval", "output_dir"}
+    return {key: value for key, value in vars(args).items() if key not in adjustable}
+
+
+def save_checkpoint(args, model, optimizer, scaler, epoch, batch, step):
+    directory = Path(args.output_dir)
+    directory.mkdir(parents=True, exist_ok=True)
+    suffix = "_moe" if args.student_use_moe else ""
+    prefix = directory / f"agent_opd_{args.student_hidden_size}{suffix}"
+    weights = {name: tensor.detach().cpu() for name, tensor in model.state_dict().items()}
+    state = {
+        "model": weights, "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict(),
+        "epoch": epoch, "batch": batch, "step": step, "signature": training_signature(args),
+        "torch_rng": torch.get_rng_state(), "python_rng": random.getstate(),
+        "cuda_rng": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else [],
+        "mps_rng": torch.mps.get_rng_state() if next(model.parameters()).device.type == "mps" else None,
+    }
+    for filename, content in [(f"{prefix}_resume.pth", state),
+                              (f"{prefix}.pth", {key: value.half() for key, value in weights.items()})]:
+        torch.save(content, filename + ".tmp")
+        os.replace(filename + ".tmp", filename)
+
+
+def train(args):
+    if int(os.environ.get("WORLD_SIZE", "1")) > 1:
+        raise ValueError("Agent OPD currently supports one training process/device")
+    for key in ("epochs", "batch_size", "accumulation_steps", "max_turns", "max_gen_len",
+                "max_total_len", "save_interval"):
+        if getattr(args, key) < 1:
+            raise ValueError(f"--{key} must be positive")
+    if args.max_steps < 0 or any(not math.isfinite(value) or value <= 0 for value in
+                               (args.learning_rate, args.temperature, args.distill_temperature, args.grad_clip)):
+        raise ValueError("Steps must be nonnegative; learning rate, temperatures and grad clip must be positive")
+    device = torch.device(args.device)
+    if device.type == "cpu" and args.dtype != "float32":
+        raise ValueError("Use --dtype float32 for the CPU reference run")
+    torch.manual_seed(args.seed)
+    random.seed(args.seed)
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
+    if tokenizer.pad_token_id is None or tokenizer.eos_token_id is None:
+        raise ValueError("The shared tokenizer must define PAD and EOS")
+    configs = {
+        who: MiniMindConfig(hidden_size=getattr(args, f"{who}_hidden_size"),
+                            num_hidden_layers=getattr(args, f"{who}_num_layers"),
+                            use_moe=bool(getattr(args, f"{who}_use_moe")))
+        for who in ("student", "teacher")
+    }
+    if len(tokenizer) != configs["student"].vocab_size:
+        raise ValueError("Both MiniMind checkpoints must use the supplied shared tokenizer/vocabulary")
+    if args.max_total_len > min(config.max_position_embeddings for config in configs.values()):
+        raise ValueError("--max_total_len exceeds the model context length")
+    student = load_model(args.student_checkpoint, configs["student"], device)
+    teacher = load_model(args.teacher_checkpoint, configs["teacher"], device).eval().requires_grad_(False)
+    optimizer = torch.optim.AdamW(student.parameters(), lr=args.learning_rate)
+    scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda" and args.dtype == "float16")
+    dtype = getattr(torch, args.dtype)
+    autocast_ctx = nullcontext() if dtype == torch.float32 else torch.autocast(device.type, dtype=dtype)
+    samples = load_prompts(args.data_path)
+    batches_per_epoch = math.ceil(len(samples) / args.batch_size)
+    start_epoch, start_batch, step = 0, 0, 0
+    if args.resume:
+        state = torch.load(args.resume, map_location="cpu", weights_only=True)
+        if state["signature"] != training_signature(args):
+            raise ValueError("Resume requires the same data, teacher, sampling and optimizer arguments")
+        student.load_state_dict(state["model"], strict=True)
+        optimizer.load_state_dict(state["optimizer"])
+        scaler.load_state_dict(state["scaler"])
+        start_epoch, start_batch, step = state["epoch"], state["batch"], state["step"]
+        torch.set_rng_state(state["torch_rng"])
+        random.setstate(state["python_rng"])
+        if state["cuda_rng"]:
+            torch.cuda.set_rng_state_all(state["cuda_rng"])
+        if state["mps_rng"] is not None:
+            torch.mps.set_rng_state(state["mps_rng"])
+    if args.max_steps and step >= args.max_steps:
+        return []
+    student.train()
+    optimizer.zero_grad(set_to_none=True)
+    metrics = []
+    cursor_epoch, cursor_batch = start_epoch, start_batch
+    for epoch in range(start_epoch, args.epochs):
+        order = list(range(len(samples)))
+        random.Random(args.seed + epoch).shuffle(order)
+        first_batch = start_batch if epoch == start_epoch else 0
+        for window in range(first_batch, batches_per_epoch, args.accumulation_steps):
+            count = min(args.accumulation_steps, batches_per_epoch - window)
+            total_loss, turns, calls, tokens, truncated, trajectories_seen = 0.0, 0, 0, 0, 0, 0
+            for batch in range(window, window + count):
+                indices = order[batch * args.batch_size:(batch + 1) * args.batch_size]
+                trajectories = [collect_trajectory(
+                    student, tokenizer, *samples[index], max_turns=args.max_turns,
+                    max_new_tokens=args.max_gen_len, max_total_len=args.max_total_len,
+                    temperature=args.temperature, open_thinking=args.open_thinking,
+                    autocast_ctx=autocast_ctx,
+                ) for index in indices]
+                input_ids, attention_mask, target_mask = pack_trajectories(
+                    trajectories, tokenizer.pad_token_id, device,
+                )
+                if not target_mask.any():
+                    raise ValueError("No assistant tokens were generated for this batch")
+                with torch.no_grad(), autocast_ctx:
+                    teacher_logits = teacher(input_ids, attention_mask=attention_mask).logits[:, :-1]
+                with autocast_ctx:
+                    output = student(input_ids, attention_mask=attention_mask)
+                    kd_loss = distillation_loss(output.logits[:, :-1], teacher_logits, target_mask,
+                                                args.distill_temperature, args.kl_direction)
+                    loss = kd_loss + output.aux_loss if configs["student"].use_moe else kd_loss
+                if not torch.isfinite(loss):
+                    raise FloatingPointError("Non-finite OPD loss")
+                scaler.scale(loss / count).backward()
+                total_loss += loss.item() / count
+                tokens += int(target_mask.sum())
+                turns += sum(t.turns for t in trajectories)
+                calls += sum(t.tool_calls for t in trajectories)
+                truncated += sum(t.truncated for t in trajectories)
+                trajectories_seen += len(trajectories)
+                del teacher_logits, output, loss, kd_loss
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(student.parameters(), args.grad_clip, error_if_nonfinite=True)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+            step += 1
+            cursor_batch = window + count
+            cursor_epoch = epoch
+            if cursor_batch == batches_per_epoch:
+                cursor_epoch, cursor_batch = epoch + 1, 0
+            record = {"step": step, "loss": total_loss, "assistant_tokens": tokens,
+                      "mean_turns": turns / trajectories_seen, "tool_calls": calls,
+                      "truncated_fraction": truncated / trajectories_seen}
+            metrics.append(record)
+            print(json.dumps(record), flush=True)
+            if step % args.save_interval == 0:
+                save_checkpoint(args, student, optimizer, scaler, cursor_epoch, cursor_batch, step)
+            if args.max_steps and step >= args.max_steps:
+                break
+        if args.max_steps and step >= args.max_steps:
+            break
+    save_checkpoint(args, student, optimizer, scaler, cursor_epoch, cursor_batch, step)
+    return metrics
+
+
+def parse_args(argv=None):
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--student_checkpoint", required=True, help="Native MiniMind .pth weights")
+    parser.add_argument("--teacher_checkpoint", required=True, help="Frozen teacher using the same tokenizer")
+    parser.add_argument("--data_path", default=str(ROOT / "dataset/agent_rl.jsonl"))
+    parser.add_argument("--tokenizer_path", default=str(ROOT / "model"))
+    parser.add_argument("--output_dir", default=str(ROOT / "out/agent_opd"))
+    parser.add_argument("--resume", default="", help="Full-precision *_resume.pth checkpoint")
+    for who in ("student", "teacher"):
+        parser.add_argument(f"--{who}_hidden_size", type=int, default=768)
+        parser.add_argument(f"--{who}_num_layers", type=int, default=8)
+        parser.add_argument(f"--{who}_use_moe", type=int, choices=[0, 1], default=0)
+    parser.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--dtype", choices=["float32", "bfloat16", "float16"], default="float32")
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--accumulation_steps", type=int, default=1)
+    parser.add_argument("--learning_rate", type=float, default=1e-6)
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--max_steps", type=int, default=0, help="Optimizer steps; 0 runs all epochs")
+    parser.add_argument("--save_interval", type=int, default=100)
+    parser.add_argument("--max_turns", type=int, default=3)
+    parser.add_argument("--max_gen_len", type=int, default=256)
+    parser.add_argument("--max_total_len", type=int, default=2048)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--distill_temperature", type=float, default=1.0)
+    parser.add_argument("--kl_direction", choices=["forward_kl", "reverse_kl"], default="forward_kl")
+    parser.add_argument("--open_thinking", action="store_true")
+    parser.add_argument("--seed", type=int, default=42)
+    return parser.parse_args(argv)
+
+
+if __name__ == "__main__":
+    train(parse_args())