推理模型从零构建 第7章：改进 GRPO 训练

def compute_advantage_stats(rewards_list):
    rewards = torch.tensor(rewards_list)
    advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-4)

    adv_avg = advantages.mean().item()  # 应该永远是 0（GRPO 归一化的性质）
    adv_std = advantages.std().item()   # 信息量更大的指标
    return advantages, adv_avg, adv_std

# 正常情况：有对有错 → std ≈ 1（健康的学习信号）
compute_advantage_stats([1., 1., 0., 0.])
# → mean=0.0000, std=0.9998

# 全对或全错 → std = 0（没有学习信号！）
compute_advantage_stats([0., 0., 0., 0.])
# → mean=0.0000, std=0.0000

def sequence_logprob_and_entropy(model, token_ids, prompt_len):
    logits = model(token_ids.unsqueeze(0)).squeeze(0).float()
    logprobs = torch.log_softmax(logits, dim=-1)

    # 选取实际生成的 token 的 logprob（和第6章相同）
    targets = token_ids[1:]
    selected = logprobs[:-1].gather(1, targets.unsqueeze(-1)).squeeze(-1)
    logp_all_steps = torch.sum(selected[prompt_len - 1:])

    # 新增：计算回答部分的平均 entropy
    all_answer_logprobs = logprobs[:-1][prompt_len - 1:]
    all_answer_probs = torch.exp(all_answer_logprobs)       # logprob → prob
    plogp = all_answer_probs * all_answer_logprobs           # p * log(p)
    step_entropy = -torch.sum(plogp, dim=-1)                 # 每个 token 的 entropy
    entropy_all_steps = torch.mean(step_entropy)             # 所有回答 token 的平均

    return logp_all_steps, entropy_all_steps

# old_logps: 更新前模型对 rollout 的 logprob
# new_logps: 更新后模型对同样 rollout 的 logprob
log_ratio = new_logps - old_logps
ratio = torch.exp(log_ratio)

# DeepSeek-R1 用 clip_eps=10（非常宽松的裁剪）
clip_eps = 10.0
clipped_ratio = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps)

# PPO-style clipping：正 advantage 取 min，负 advantage 取 max
adv = advantages.detach()
unclipped = ratio * adv
clipped = clipped_ratio * adv

obj = torch.where(
    adv >= 0,
    torch.minimum(unclipped, clipped),   # 限制正更新的上界
    torch.maximum(unclipped, clipped),   # 限制负更新的下界
)

clipped_pg_loss = -torch.mean(obj)

# 未裁剪: -2.5764
# 裁剪后: -2.3998 → 更新幅度略微减小，防止过激更新

import copy

# 冻结 reference model（原始模型的副本，不参与训练）
ref_model = copy.deepcopy(model).to(device)
ref_model.eval()
for p in ref_model.parameters():
    p.requires_grad = False

# 在 compute_grpo_loss 中新增 KL 计算：
for _ in range(num_rollouts):
    token_ids, prompt_len, text = sample_response(model, ...)

    # 当前模型的 logprob
    logp = sequence_logprob(model, token_ids, prompt_len)

    # reference model 的 logprob（不需要梯度）
    with torch.no_grad():
        ref_logp = sequence_logprob(ref_model, token_ids, prompt_len)

    # ...

# Policy gradient loss + KL 正则化
pg_loss = -(advantages.detach() * logps).mean()
kl_loss = kl_coeff * torch.mean(logps - ref_logps)  # 惩罚偏离 reference
loss = pg_loss + kl_loss  # 总 loss

# Base tokenizer 不认识 
tokenizer_base.encode("")   # → [13708, 766, 29]  拆成3个token

# 添加 special tokens 后
tokenizer_base._tok.add_special_tokens(
    ["", "", "", ""]
)
tokenizer_base.encode("")   # → [151667]  单个token
tokenizer_base.encode("")  # → [151668]  单个token

# 或直接用 reasoning model 的 tokenizer（已原生支持）
tokenizer.encode("")   # → [151667]
tokenizer.encode("")  # → [151668]

def reward_format(token_ids, prompt_len,
                  start_think_id=151667, end_think_id=151668):
    """检查生成部分是否包含正确顺序的  """
    try:
        gen = token_ids[prompt_len:].tolist()
        return float(gen.index(start_think_id) < gen.index(end_think_id))
    except ValueError:
        return 0.0  # 缺少任一标签 → 0

# 在 GRPO loss 中组合两种 reward
rlvr_reward = reward_rlvr(text, example["answer"])       # 正确性
format_reward = reward_format(token_ids, prompt_len)      # 格式
reward = rlvr_reward + format_reward_weight * format_reward  # 总 reward

# 同时修改 prompt template 引导模型使用  格式
def render_prompt_with_think_tokens(prompt):
    return (
        "You are a helpful math assistant.\n"
        "When solving the problem, first write your reasoning "
        "inside  and  tags.\n"
        "Then write the final result on a new line in the exact format:\n"
        "\\boxed{ANSWER}\n\n"
        f"Question:\n{prompt}\n\nAnswer:"
    )