Merge 138c367ddb into d65ef2c00e

README.md

@@ -911,15 +911,17 @@ response = client.chat.completions.create(
 
 在介绍实现具体算法之前，我先以个人理解的极简视角，阐述所有Policy Optimization (PO)算法的统一共性。
 
-所有RL算法的本质都只是在优化一个期望：
+说到底，这里讨论的 PO 算法都只是在优化一个期望：
 
-$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q), o \sim \pi(O|q)} \left[ \underbrace{f(r_t)}_{\text{策略项}} \cdot \underbrace{g(A_t)}_{\text{优势项}} - \underbrace{h(\text{KL}_t)}_{\text{正则项}} \right]$$
+$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q),\, o \sim \pi_\theta(\cdot \mid q)} \left[ \underbrace{\Phi(r_t, A_t)}_{\text{策略目标}} - \underbrace{h(\text{KL}_t)}_{\text{正则项}} \right]$$
 
-训练时，只需**最小化负目标函数**，即: $\mathcal{L}_{PO} = -\mathcal{J}_{PO}$
+训练时，只需**最小化负目标函数**，即：
+
+$$\mathcal{L}_{PO} = -\mathcal{J}_{PO}$$
 
 这个框架只包含三个核心组件：
-* **策略项** $f(r_t)$: 如何使用概率比 $r_t$? 即告诉模型新旧策略偏差有多大，是否探索到了更好的token
+* **策略项** $\Phi(r_t, A_t)$: 如何结合概率比 $r_t$ 和优势 $A_t$ 更新策略
-* **优势项** $g(A_t)$: 如何计算优势 $A_t$, 这很重要！大模型算对定积分也不足为奇，小模型回答对加减法优势通常都是正的
+* **优势项** $A_t$: 如何计算优势，这很重要！大模型算对定积分也不足为奇，小模型回答对加减法优势通常都是正的
 * **正则项** $h(\text{KL}_t)$: 如何约束变化幅度 $\text{KL}_t$, 既防止跑偏又防止管的太死
 
 <details>
@@ -929,7 +931,7 @@ $$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q), o \sim \pi(O|q)} \left[ \underbrac
 |------|------|------|------|
 | $q$ | 问题/提示词 | 从数据集 $P(Q)$ 中采样 | - |
 | $o$ | 模型输出序列 | 由策略 $\pi$ 生成 | - |
-| $r_t$ | 概率比 | $r_t = \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{ref}(o_t \mid q, o_{<t})}$ | $(0, +\infty)$ |
+| $r_t$ | 概率比 | $r_t = \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{\mathrm{old}}(o_t \mid q, o_{<t})}$ | $(0, +\infty)$ |
 | $A_t$ | 优势函数 | 衡量某个动作相比基线有多好 | $(-\infty, +\infty)$ |
 | $\text{KL}_t$ | KL散度 | 防止策略偏离参考模型太远 | $[0, +\infty)$ |
 
@@ -1085,7 +1087,7 @@ $$\mathcal{L}_{PPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \text{clip}(r_t, 1-\v
 
 其中：
 - **策略项**: $f(r_t) = \min(r_t, \text{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (裁剪概率比防止更新过激)
-- **优势项**: $g(A_t) = R - V(s)$ (通过Critic网络估计价值函数)
+- **优势项**: $A_t$ 通常由Critic网络估计，也可以使用GAE进行计算
 - **正则项**: $h(\text{KL}_t) = \beta \cdot \mathbb{E}[\text{KL}]$ (全局KL散度约束)
 
 对比DPO而言，
@@ -1126,8 +1128,8 @@ python train_ppo.py
 $$\mathcal{L}_{GRPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t) - \beta \cdot \text{KL}_t\right]$$
 
 其中：
-- **策略项**: $f(r_t) = \min(r_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (使用概率比的对称 clip 裁剪)
+- **策略项**: $f(r_t, A_t) = \min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ (对概率比和优势项一起做裁剪，防止更新过激)
-- **优势项**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (组内归一化，消除Critic网络)
+- **优势项**: $g(A_{i,j}) = \frac{R_{i,j} - \mu_i}{\sigma_i + \epsilon}$ (组内归一化，消除Critic网络)
 - **正则项**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token级KL散度约束)
 
 对于同一个问题，模型生成 N 个回答并计算各自奖励，再用组内平均奖励作为 baseline。高于 baseline 的回答被鼓励，低于 baseline 的回答被抑制，因此无需额外训练 critic 网络。
@@ -1160,11 +1162,11 @@ CISPO 的关注点并不是重新设计 group baseline，而是用非常小的 l
 
 **CISPO损失**：
 
-$$\mathcal{L}_{CISPO} = -\mathbb{E}\left[\min(r_t, \varepsilon_{max}) \cdot A_t \cdot \log \pi_\theta(a_t|s) - \beta \cdot \text{KL}_t\right]$$
+$$\mathcal{L}_{CISPO} = -\mathbb{E}\left[\min(r_t, \varepsilon_{\mathrm{high}}) \cdot A_t \cdot \log \pi_\theta(a_t|s) - \beta \cdot \text{KL}_t\right]$$
 
 其中：
-- **策略项**: $f(r_t) = \min(r_t, \varepsilon_{max}) \cdot \log \pi_\theta(a_t|s)$ (ratio 只作为裁剪后的权重)
+- **策略项**: $f(r_t) = \min(r_t, \varepsilon_{\mathrm{high}}) \cdot \log \pi_\theta(a_t|s)$ (ratio 只作为裁剪后的权重)
-- **优势项**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (可直接沿用 GRPO 的组内相对优势)
+- **优势项**: $g(A_{i,j}) = \frac{R_{i,j} - \mu_i}{\sigma_i + \epsilon}$ (可直接沿用 GRPO 的组内相对优势)
 - **正则项**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token级KL散度约束)
 
 CISPO在GRPO基础上，把原本容易被clip成常数的策略项改写成“裁剪权重 × log 概率”的形式。这样ratio即使被截断，也不会把梯度路径一起截断。因此可以直接把CISPO视作GRPO的loss变体来实现，而不是单独维护一套独立脚本。这里不再单列实验。只需在 `train_grpo.py` 把 `loss_type` 配置为 `cispo`，其余训练流程仍沿用 GRPO 的分组采样、奖励计算与优势构造逻辑即可。
@@ -1269,9 +1271,9 @@ python eval_toolcall.py --weight agent
 | 算法 | 策略项 $f(r_t)$ | 优势项 $g(A_t)$ | 正则项 $h(\text{KL}_t)$ | 训练模型数 |
 |------|----------------|----------------|----------------------|----------|
 | **DPO** | $\log r_w - \log r_l$ | 无显式优势项 | 隐含在 $\beta$ 中 | 1 (前向参与 2) | 
-| **PPO** | $\min(r, \text{clip}(r))$ | $R - V(s)$ | $\beta \cdot \mathbb{E}[\text{KL}]$ | 2 | 
+| **PPO** | $\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ | $A_t$（通常由 Critic 估计，也可以使用 GAE） | $\beta \cdot \mathbb{E}[\text{KL}]$ | 2 |
-| **GRPO** | $\min(r, \text{clip}(r))$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 |
+| **GRPO** | $\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ | $A_{i,j}=\frac{R_{i,j}-\mu_i}{\sigma_i+\epsilon}$ | $\beta \cdot \text{KL}_t$ | 1 |
-| **CISPO** | $\mathrm{clip}(r, 0, \varepsilon_{max}) \cdot A_t \cdot \log \pi_\theta$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 | 
+| **CISPO** | $\mathrm{clip}(r, 0, \varepsilon_{\mathrm{high}}) \cdot A_t \cdot \log \pi_\theta$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 | 
 
 **说白了，这些 RL 算法不是割裂独立的，而是在统一优化视角下，对同一目标函数进行不同设计权衡后形成的自然变体，呈现为一种优美自洽的统一。**
 

README_en.md

@@ -911,15 +911,17 @@ However, in specific practice, their differences go beyond just the feedback sou
 
 Before introducing the implementation of specific algorithms, let me first describe the unified commonality of all Policy Optimization (PO) algorithms from my own minimalist perspective.
 
-The essence of all RL algorithms is just optimizing an expectation:
+Put simply, the PO algorithms discussed here are all optimizing an expectation:
 
-$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q), o \sim \pi(O|q)} \left[ \underbrace{f(r_t)}_{\text{policy term}} \cdot \underbrace{g(A_t)}_{\text{advantage term}} - \underbrace{h(\text{KL}_t)}_{\text{regularization term}} \right]$$
+$$\mathcal{J}_{PO} = \mathbb{E}_{q \sim P(Q),\, o \sim \pi_\theta(\cdot \mid q)} \left[ \underbrace{\Phi(r_t, A_t)}_{\text{policy objective}} - \underbrace{h(\text{KL}_t)}_{\text{regularization term}} \right]$$
 
-During training, one only needs to **minimize the negative objective function**, i.e.: $\mathcal{L}_{PO} = -\mathcal{J}_{PO}$
+During training, one only needs to **minimize the negative objective function**, i.e.:
+
+$$\mathcal{L}_{PO} = -\mathcal{J}_{PO}$$
 
 This framework contains only three core components:
-* **Policy term** $f(r_t)$: How to use the probability ratio $r_t$? It tells the model how far the new and old policies have deviated, and whether better tokens have been explored
+* **Policy term** $\Phi(r_t, A_t)$: How to combine the probability ratio $r_t$ and advantage $A_t$ to update the policy
-* **Advantage term** $g(A_t)$: How to compute the advantage $A_t$, this is very important! It's no surprise that large models can solve definite integrals correctly, but for small models, even getting addition and subtraction right usually yields a positive advantage
+* **Advantage term** $A_t$: How to compute the advantage, this is very important! It's no surprise that large models can solve definite integrals correctly, but for small models, even getting addition and subtraction right usually yields a positive advantage
 * **Regularization term** $h(\text{KL}_t)$: How to constrain the magnitude of change $\text{KL}_t$, preventing both drifting too far and constraining too tightly
 
 <details>
@@ -929,7 +931,7 @@ This framework contains only three core components:
 |--------|---------|-------------|-------|
 | $q$ | Question/Prompt | Sampled from dataset $P(Q)$ | - |
 | $o$ | Model output sequence | Generated by policy $\pi$ | - |
-| $r_t$ | Probability ratio | $r_t = \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{ref}(o_t \mid q, o_{<t})}$ | $(0, +\infty)$ |
+| $r_t$ | Probability ratio | $r_t = \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{\mathrm{old}}(o_t \mid q, o_{<t})}$ | $(0, +\infty)$ |
 | $A_t$ | Advantage function | Measures how much better a certain action is compared to the baseline | $(-\infty, +\infty)$ |
 | $\text{KL}_t$ | KL divergence | Prevents the policy from drifting too far from the reference model | $[0, +\infty)$ |
 
@@ -1085,7 +1087,7 @@ $$\mathcal{L}_{PPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \text{clip}(r_t, 1-\v
 
 Where:
 - **Policy term**: $f(r_t) = \min(r_t, \text{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (clips probability ratio to prevent overly aggressive updates)
-- **Advantage term**: $g(A_t) = R - V(s)$ (estimates value function through Critic network)
+- **Advantage term**: $A_t$ is typically estimated by the Critic network, and GAE can also be used
 - **Regularization term**: $h(\text{KL}_t) = \beta \cdot \mathbb{E}[\text{KL}]$ (global KL divergence constraint)
 
 Compared to DPO,
@@ -1126,8 +1128,8 @@ In early 2025, as DeepSeek-R1 exploded in popularity, GRPO from the DeepSeekMath
 $$\mathcal{L}_{GRPO} = -\mathbb{E}\left[\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t) - \beta \cdot \text{KL}_t\right]$$
 
 Where:
-- **Policy term**: $f(r_t) = \min(r_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon))$ (uses symmetric clip on the probability ratio)
+- **Policy term**: $f(r_t, A_t) = \min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ (clips the probability ratio and advantage term together)
-- **Advantage term**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (intra-group normalization, eliminating the Critic network)
+- **Advantage term**: $g(A_{i,j}) = \frac{R_{i,j} - \mu_i}{\sigma_i + \epsilon}$ (intra-group normalization, eliminating the Critic network)
 - **Regularization term**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token-level KL divergence constraint)
 
 For the same question, the model generates N responses and computes their respective rewards, then uses the intra-group average reward as a baseline. Responses above the baseline are encouraged, and responses below the baseline are suppressed, thus no additional critic network needs to be trained.
@@ -1160,11 +1162,11 @@ CISPO does not focus on redesigning the group baseline. Instead, it uses a very
 
 **CISPO Loss**:
 
-$$\mathcal{L}_{CISPO} = -\mathbb{E}\left[\min(r_t, \varepsilon_{max}) \cdot A_t \cdot \log \pi_\theta(a_t|s) - \beta \cdot \text{KL}_t\right]$$
+$$\mathcal{L}_{CISPO} = -\mathbb{E}\left[\min(r_t, \varepsilon_{\mathrm{high}}) \cdot A_t \cdot \log \pi_\theta(a_t|s) - \beta \cdot \text{KL}_t\right]$$
 
 Where:
-- **Policy term**: $f(r_t) = \min(r_t, \varepsilon_{max}) \cdot \log \pi_\theta(a_t|s)$ (ratio serves only as a clipped weight)
+- **Policy term**: $f(r_t) = \min(r_t, \varepsilon_{\mathrm{high}}) \cdot \log \pi_\theta(a_t|s)$ (ratio serves only as a clipped weight)
-- **Advantage term**: $g(A_t) = \frac{R - \mu_{group}}{\sigma_{group}}$ (can directly reuse GRPO's intra-group relative advantage)
+- **Advantage term**: $g(A_{i,j}) = \frac{R_{i,j} - \mu_i}{\sigma_i + \epsilon}$ (can directly reuse GRPO's intra-group relative advantage)
 - **Regularization term**: $h(\text{KL}_t) = \beta \cdot \text{KL}_t$ (token-level KL divergence constraint)
 
 CISPO, building on GRPO, rewrites the policy term that was easily clipped into a constant into the form "clipped weight × log probability". This way, even if the ratio is truncated, the gradient path is not truncated along with it. Therefore, CISPO can be directly viewed as a loss variant of GRPO to implement, rather than maintaining a separate standalone script. No separate experiment is listed here. One only needs to set `loss_type` to `cispo` in `train_grpo.py`; the rest of the training process still follows GRPO's group sampling, reward computation, and advantage construction logic.
@@ -1269,9 +1271,9 @@ Returning to the "**unified framework**", the table below summarizes how differe
 | Algorithm | Policy term $f(r_t)$ | Advantage term $g(A_t)$ | Regularization term $h(\text{KL}_t)$ | Number of training models |
 |-----------|---------------------|------------------------|-------------------------------------|--------------------------|
 | **DPO** | $\log r_w - \log r_l$ | No explicit advantage term | Implicit in $\beta$ | 1 (2 participate in forward) | 
-| **PPO** | $\min(r, \text{clip}(r))$ | $R - V(s)$ | $\beta \cdot \mathbb{E}[\text{KL}]$ | 2 | 
+| **PPO** | $\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ | $A_t$ (typically estimated by the Critic, or computed with GAE) | $\beta \cdot \mathbb{E}[\text{KL}]$ | 2 |
-| **GRPO** | $\min(r, \text{clip}(r))$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 |
+| **GRPO** | $\min(r_t \cdot A_t, \mathrm{clip}(r_t, 1-\varepsilon, 1+\varepsilon) \cdot A_t)$ | $A_{i,j}=\frac{R_{i,j}-\mu_i}{\sigma_i+\epsilon}$ | $\beta \cdot \text{KL}_t$ | 1 |
-| **CISPO** | $\mathrm{clip}(r, 0, \varepsilon_{max}) \cdot A_t \cdot \log \pi_\theta$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 | 
+| **CISPO** | $\mathrm{clip}(r, 0, \varepsilon_{\mathrm{high}}) \cdot A_t \cdot \log \pi_\theta$ | $\frac{R - \mu}{\sigma}$ | $\beta \cdot \text{KL}_t$ | 1 | 
 
 **To put it plainly, these RL algorithms are not isolated from one another. From a unified optimization perspective, they are natural variants formed by making different design trade-offs on the same objective function, presenting a beautifully self-consistent unity.**
 

trainer/train_dapo.py

@@ -0,0 +1,421 @@
+import os
+import sys
+
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import argparse
+import re
+import warnings
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from contextlib import nullcontext
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.nn.utils import clip_grad_norm_
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from transformers import AutoModel
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from dataset.lm_dataset import RLAIFDataset
+from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model
+
+warnings.filterwarnings('ignore')
+
+
+def calculate_rewards(prompts, responses, reward_model, reward_tokenizer, args):
+    """
+    计算复杂奖励，包含格式奖励、模型打分以及 DAPO 特有的长度惩罚。
+    输入的 prompts/responses 已经是展开后的（batch_size * group_size）条目。
+    """
+    def reasoning_model_reward(rewards):
+        # 1. 结构化格式奖励
+        pattern  = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>$"
+        pattern2 = r"^<think>\n.*?\n</think>\n\n<answer>\n.*?\n</answer>$"
+        matches_pattern  = [re.match(pattern,  r, re.S) for r in responses]
+        matches_pattern2 = [re.match(pattern2, r, re.S) for r in responses]
+        format_rewards = [0.5 if (m1 or m2) else 0.0
+                         for m1, m2 in zip(matches_pattern, matches_pattern2)]
+        rewards += torch.tensor(format_rewards, device=args.device)
+
+        # 2. 标签存在奖励
+        def mark_num(text):
+            reward = 0
+            for tag in ["<think>", "</think>", "<answer>", "</answer>"]:
+                if text.count(tag) == 1:
+                    reward += 0.25
+            return reward
+
+        mark_rewards = [mark_num(r) for r in responses]
+        rewards += torch.tensor(mark_rewards, device=args.device)
+        return rewards
+
+    rewards = torch.zeros(len(responses), device=args.device)
+
+    if args.reasoning == 1:
+        rewards = reasoning_model_reward(rewards)
+
+    # 3. 超长惩罚 (Overlong Punishment - DAPO 特性)
+    len_penalty = []
+    for resp in responses:
+        curr_len = len(resp)
+        if curr_len > args.length_threshold:
+            penalty = -args.length_penalty * ((curr_len - args.length_threshold) / 100)
+            len_penalty.append(max(penalty, -1.0))
+        else:
+            len_penalty.append(0.0)
+    rewards += torch.tensor(len_penalty, device=args.device)
+
+    # 4. 奖励模型评分
+    with torch.no_grad():
+        reward_model_scores = []
+        for prompt, response in zip(prompts, responses):
+            pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
+            matches = re.findall(pattern, prompt, re.DOTALL)
+            messages = [{"role": role, "content": content.strip()} for role, content in matches]
+
+            tmp_chat = messages + [{"role": "assistant", "content": response}]
+            score = reward_model.get_score(reward_tokenizer, tmp_chat)
+            scale = 3.0
+            score = max(min(score, scale), -scale)
+
+            if args.reasoning == 1:
+                answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
+                if answer_match:
+                    answer_content = answer_match.group(1).strip()
+                    tmp_chat = messages + [{"role": "assistant", "content": answer_content}]
+                    answer_score = reward_model.get_score(reward_tokenizer, tmp_chat)
+                    answer_score = max(min(answer_score, scale), -scale)
+                    score = score * 0.4 + answer_score * 0.6
+            reward_model_scores.append(score)
+
+        reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
+        rewards += reward_model_scores
+
+    return rewards
+
+
+def dapo_train_epoch(epoch, loader, iters, old_actor_model, ref_model,
+                     actor_scheduler, reward_model, reward_tokenizer,
+                     start_step=0, wandb=None):
+    actor_model.train()
+
+    for step, batch in enumerate(loader, start=start_step + 1):
+        prompts    = batch["prompt"]
+        batch_size = len(prompts)
+        G          = args.group_size
+
+        # 每个 prompt 重复 G 次，展开为 (batch_size * G) 条输入
+        prompts_expanded = [p for p in prompts for _ in range(G)]
+
+        enc = tokenizer(
+            prompts_expanded,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=args.max_seq_len,
+            padding_side="left"
+        ).to(args.device)
+        prompt_length = enc.input_ids.shape[1]
+
+        with torch.no_grad():
+            model_for_gen = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
+            gen_out = model_for_gen.generate(
+                input_ids=enc.input_ids,
+                attention_mask=enc.attention_mask,
+                max_new_tokens=args.max_gen_len,
+                do_sample=True,
+                temperature=0.8,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+
+        # gen_out: (batch_size * G, seq_len)
+        responses_text = [
+            tokenizer.decode(gen_out[i, prompt_length:], skip_special_tokens=True)
+            for i in range(len(prompts_expanded))
+        ]
+        rewards_flat = calculate_rewards(prompts_expanded, responses_text, reward_model, reward_tokenizer, args)
+        # rewards_flat: (batch_size * G,)
+
+        # ------------------------------------------------------------------
+        # advantage = (r - mean_group) / std_group
+        # ------------------------------------------------------------------
+        rewards_grouped = rewards_flat.view(batch_size, G)
+        mean_r          = rewards_grouped.mean(dim=1, keepdim=True)
+        std_r           = rewards_grouped.std(dim=1, keepdim=True).clamp(min=1e-8)
+        advantages      = ((rewards_grouped - mean_r) / std_r).view(-1)  # (batch_size * G,)
+
+        # ------------------------------------------------------------------
+        # Dynamic Sampling
+        # ------------------------------------------------------------------
+        if args.use_dynamic_sampling == 1:
+            if (std_r < 1e-6).all().item():
+                if is_main_process():
+                    Logger(f"Step {step}: All groups have zero reward variance, skipping...")
+                continue
+
+        # ------------------------------------------------------------------
+        # Actor forward
+        # ------------------------------------------------------------------
+        full_mask = (gen_out != tokenizer.pad_token_id).long()
+
+        with autocast_ctx:
+            res      = actor_model(input_ids=gen_out, attention_mask=full_mask)
+            logits   = res.logits
+            aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
+
+        labels      = gen_out[:, 1:].clone()
+        logp_tokens = F.log_softmax(logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1)
+
+        seq_len    = gen_out.size(1) - 1
+        resp_mask  = torch.arange(seq_len, device=gen_out.device).unsqueeze(0) >= prompt_length - 1
+        final_mask = resp_mask & (~labels.eq(tokenizer.pad_token_id))
+
+        # Token-level 平均 log-prob（防止长回复主导梯度）
+        token_count = final_mask.sum(dim=1).clamp(min=1).float()
+        actor_logp  = (logp_tokens * final_mask).sum(dim=1) / token_count
+
+        with torch.no_grad():
+            old_logits      = old_actor_model(input_ids=gen_out, attention_mask=full_mask).logits
+            old_logp_tokens = F.log_softmax(old_logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1)
+            old_logp        = (old_logp_tokens * final_mask).sum(dim=1) / token_count
+
+            ref_logits      = ref_model(input_ids=gen_out, attention_mask=full_mask).logits
+            ref_logp_tokens = F.log_softmax(ref_logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1)
+            ref_logp        = (ref_logp_tokens * final_mask).sum(dim=1) / token_count
+
+        kl     = (actor_logp - old_logp).mean()
+        kl_ref = (actor_logp - ref_logp).mean()
+        ratio  = torch.exp(actor_logp - old_logp)
+
+        # ------------------------------------------------------------------
+        # Decoupled Clip
+        #   Adv > 0：只限上界，允许正向大步更新prompts_expanded = [p for p in prompts for _ in range(G)]
+        #   Adv ≤ 0：只限下界，防止负向更新反弹
+        # ------------------------------------------------------------------
+        surr1 = ratio * advantages
+        clipped_ratio = torch.where(
+            advantages > 0,
+            torch.clamp(ratio, max=1.0 + args.clip_epsilon_high),
+            torch.clamp(ratio, min=1.0 - args.clip_epsilon_low),
+        )
+        surr2       = clipped_ratio * advantages
+        policy_loss = -torch.min(surr1, surr2).mean()
+
+        loss = (policy_loss + args.kl_coef * kl_ref + aux_loss) / args.accumulation_steps
+        loss.backward()
+
+        if (step + 1) % args.accumulation_steps == 0:
+            clip_grad_norm_(actor_model.parameters(), args.grad_clip)
+            actor_optimizer.step()
+            actor_scheduler.step()
+            actor_optimizer.zero_grad()
+
+        if is_main_process():
+            response_ids = gen_out[:, enc.input_ids.shape[1]:]
+            is_eos       = (response_ids == tokenizer.eos_token_id)
+            eos_indices  = torch.argmax(is_eos.int(), dim=1)
+            lengths = torch.where(
+                is_eos.any(dim=1),
+                eos_indices + 1,
+                torch.tensor(response_ids.shape[1], device=is_eos.device)
+            )
+            avg_len    = lengths.float().mean()
+            reward_std = rewards_flat.std().item()
+
+            if wandb is not None:
+                wandb.log({
+                    "actor_loss":       policy_loss.item(),
+                    "reward":           rewards_flat.mean().item(),
+                    "reward_std":       reward_std,
+                    "kl":               kl.item(),
+                    "kl_ref":           kl_ref.item(),
+                    "avg_response_len": avg_len.item(),
+                    "actor_lr":         actor_optimizer.param_groups[0]['lr'],
+                })
+
+            Logger(
+                f"Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), "
+                f"Loss: {policy_loss.item():.4f}, Reward: {rewards_flat.mean().item():.4f}, "
+                f"KL: {kl.item():.4f}, Len: {avg_len.item():.2f}"
+            )
+
+        if (step + 1) % args.update_old_actor_freq == 0:
+            raw_actor = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
+            raw_actor = getattr(raw_actor, '_orig_mod', raw_actor)
+            old_actor_model.load_state_dict({k: v.detach().cpu() for k, v in raw_actor.state_dict().items()})
+            old_actor_model.to(args.device)
+
+        if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
+            actor_model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            raw_actor = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
+            raw_actor = getattr(raw_actor, '_orig_mod', raw_actor)
+            torch.save({k: v.half().cpu() for k, v in raw_actor.state_dict().items()}, ckp)
+            lm_checkpoint(
+                lm_config, weight=args.save_weight, model=actor_model,
+                optimizer=actor_optimizer, epoch=epoch, step=step,
+                wandb=wandb, save_dir='../checkpoints', scheduler=actor_scheduler,
+            )
+            actor_model.train()
+
+        del enc, gen_out, responses_text, rewards_flat, rewards_grouped
+        del advantages, full_mask, logits, labels, logp_tokens, final_mask, token_count
+        del actor_logp, old_logits, old_logp, ref_logits, ref_logp
+        del kl, kl_ref, ratio, surr1, surr2, clipped_ratio, policy_loss, loss
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="MiniMind DAPO (Decoupled Clip and Dynamic Sampling Policy Optimization)"
+    )
+    parser.add_argument("--save_dir",              type=str,   default="../out",                      help="模型保存目录")
+    parser.add_argument('--save_weight',            type=str,   default='dapo',                  help="保存权重的前缀名")
+    parser.add_argument("--epochs",                type=int,   default=1,                             help="训练轮数")
+    parser.add_argument("--batch_size",            type=int,   default=2,                             help="prompt 级别的 batch size（实际前向 = batch_size * group_size）")
+    parser.add_argument("--learning_rate",         type=float, default=8e-8,                          help="Actor 学习率")
+    parser.add_argument("--device",                type=str,   default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype",                 type=str,   default="bfloat16",                    help="混合精度类型")
+    parser.add_argument("--num_workers",           type=int,   default=8,                             help="数据加载线程数")
+    parser.add_argument("--accumulation_steps",    type=int,   default=1,                             help="梯度累积步数")
+    parser.add_argument("--grad_clip",             type=float, default=1.0,                           help="梯度裁剪阈值")
+    parser.add_argument("--log_interval",          type=int,   default=1,                             help="日志打印间隔")
+    parser.add_argument("--save_interval",         type=int,   default=10,                            help="模型保存间隔")
+    parser.add_argument('--hidden_size',            type=int,   default=512,                           help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers',      type=int,   default=8,                             help="隐藏层数量")
+    parser.add_argument('--use_moe',                type=int,   default=0,  choices=[0, 1],            help="是否使用 MoE 架构（0=否，1=是）")
+    parser.add_argument('--max_seq_len',            type=int,   default=66,                            help="Prompt 最大长度")
+    parser.add_argument("--max_gen_len",            type=int,   default=1024,                          help="生成的最大长度")
+    parser.add_argument("--data_path",              type=str,   default="../dataset/rlaif-mini.jsonl", help="RLAIF 数据路径")
+
+    # DAPO 核心参数
+    parser.add_argument("--clip_epsilon_low",      type=float, default=0.1,  help="Decoupled Clip 下界")
+    parser.add_argument("--clip_epsilon_high",     type=float, default=0.2,  help="Decoupled Clip 上界")
+    parser.add_argument("--group_size",            type=int,   default=8,    help="每个 prompt 的采样回复数")
+    parser.add_argument("--use_dynamic_sampling",  type=int,   default=1,   choices=[0, 1], help="是否开启动态采样过滤")
+    parser.add_argument("--length_penalty",        type=float, default=0.1,  help="长度惩罚系数")
+    parser.add_argument("--length_threshold",      type=int,   default=1200, help="触发长度惩罚的字符阈值")
+    parser.add_argument("--kl_coef",               type=float, default=0.02, help="KL 散度惩罚系数（相对 ref_model）")
+    parser.add_argument("--reasoning",             type=int,   default=1,   choices=[0, 1], help="推理模型类型（0=普通模型，1=推理模型）")
+    parser.add_argument("--update_old_actor_freq", type=int,   default=4,   help="更新 old_actor_model 的频率（步数）")
+    parser.add_argument("--reward_model_path",     type=str,   default="../../internlm2-1_8b-reward", help="Reward 模型路径")
+    parser.add_argument('--from_resume',            type=int,   default=0,   choices=[0, 1], help="是否自动检测并续训（0=否，1=是）")
+    parser.add_argument("--use_wandb",             action="store_true",                     help="是否使用 wandb/swanlab")
+    parser.add_argument("--wandb_project",         type=str,   default="MiniMind-DAPO",     help="wandb 项目名")
+    parser.add_argument("--use_compile",           type=int,   default=0,   choices=[0, 1], help="是否使用 torch.compile 加速（0=否，1=是）")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized():
+        args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+
+    # ========== 2. 配置目录、模型参数、检查 ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(
+        hidden_size=args.hidden_size,
+        num_hidden_layers=args.num_hidden_layers,
+        use_moe=bool(args.use_moe)
+    )
+    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') \
+               if args.from_resume == 1 else None
+
+    # ========== 3. 设置混合精度 ==========
+    device_type  = "cuda" if "cuda" in args.device else "cpu"
+    dtype        = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+
+    # ========== 4. 配 wandb (SwanLab) ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume   = 'must' if wandb_id else None
+        wandb.init(
+            project=args.wandb_project,
+            name=f"MiniMind-DAPO-{args.learning_rate}",
+            id=wandb_id,
+            resume=resume,
+        )
+
+    # ========== 5. 初始化模型 ==========
+    base_weight = "reason" if args.reasoning == 1 else "full_sft"
+
+    actor_model, tokenizer = init_model(lm_config, base_weight, device=args.device)
+    if args.use_compile == 1:
+        actor_model = torch.compile(actor_model)
+
+    old_actor_model, _ = init_model(lm_config, base_weight, device=args.device)
+    old_actor_model = old_actor_model.eval().requires_grad_(False)
+
+    ref_model, _ = init_model(lm_config, base_weight, device=args.device)
+    ref_model = ref_model.eval().requires_grad_(False)
+
+    reward_model = AutoModel.from_pretrained(
+        args.reward_model_path,
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+    )
+    reward_model = reward_model.to(args.device).eval().requires_grad_(False)
+    reward_tokenizer = AutoTokenizer.from_pretrained(args.reward_model_path, trust_remote_code=True)
+
+    # ========== 6. 数据、优化器、调度器 ==========
+    train_ds        = RLAIFDataset(args.data_path, tokenizer, max_length=(args.max_seq_len + args.max_gen_len))
+    train_sampler   = DistributedSampler(train_ds) if dist.is_initialized() else None
+    actor_optimizer = optim.AdamW(actor_model.parameters(), lr=args.learning_rate)
+
+    iters           = len(DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler))
+    total_steps     = (iters // args.accumulation_steps) * args.epochs
+    actor_scheduler = CosineAnnealingLR(actor_optimizer, T_max=total_steps, eta_min=args.learning_rate / 10)
+
+    if ckp_data:
+        actor_model.load_state_dict(ckp_data['model'])
+        actor_optimizer.load_state_dict(ckp_data['optimizer'])
+        actor_scheduler.load_state_dict(ckp_data['scheduler'])
+        start_epoch = ckp_data['epoch']
+        start_step  = ckp_data.get('step', 0)
+    else:
+        start_epoch, start_step = 0, 0
+
+    if dist.is_initialized():
+        actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
+        old_actor_model.to(args.device)
+
+    # ========== 7. 打印配置摘要 ==========
+    if is_main_process():
+        Logger("=" * 60)
+        Logger("DAPO Training Config (Critic-free):")
+        Logger(f"  group_size           = {args.group_size}")
+        Logger(f"  clip_epsilon_low     = {args.clip_epsilon_low}")
+        Logger(f"  clip_epsilon_high    = {args.clip_epsilon_high}")
+        Logger(f"  use_dynamic_sampling = {args.use_dynamic_sampling}")
+        Logger(f"  kl_coef              = {args.kl_coef}")
+        Logger(f"  effective_batch      = {args.batch_size} prompts × {args.group_size} = {args.batch_size * args.group_size}")
+        Logger("=" * 60)
+
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch)
+        indices = torch.randperm(len(train_ds)).tolist()
+        skip    = start_step if (epoch == start_epoch and start_step > 0) else 0
+        loader  = DataLoader(
+            train_ds,
+            batch_sampler=SkipBatchSampler(train_sampler or indices, args.batch_size, skip),
+            num_workers=args.num_workers,
+            pin_memory=True,
+        )
+        dapo_train_epoch(
+            epoch, loader, len(loader),
+            old_actor_model, ref_model,
+            actor_scheduler, reward_model, reward_tokenizer,
+            skip, wandb,
+        )
+
+    if dist.is_initialized():
+        dist.destroy_process_group()