GRPO 算法解析：DeepSeek 的强化学习新方案

对每个 prompt p:
1. 采样 G 个回复 {r_1, r_2, ..., r_G}
2. 计算每个回复的奖励 {R_1, R_2, ..., R_G}
3. 组内归一化：advantage_i = (R_i - mean(R)) / std(R)
4. 用归一化后的 advantage 更新策略

# PPO 的 Advantage 计算
def ppo_advantage(rewards, values, gamma=0.99, lam=0.95):
    """需要 Critic 网络预测 values"""
    advantages = []
    returns = []
    gae = 0
    for t in reversed(range(len(rewards))):
        delta = rewards[t] + gamma * values[t + 1] - values[t]
        gae = delta + gamma * lam * gae
        advantages.insert(0, gae)
    return advantages

# GRPO 的 Advantage 计算
def grpo_advantage(rewards_per_group: list):
    """不需要 Critic，直接用组内统计"""
    group_mean = np.mean(rewards_per_group)
    group_std = np.std(rewards_per_group)
    
    advantages = [(r - group_mean) / (group_std + 1e-8) 
                  for r in rewards_per_group]
    return advantages

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

class GRPOTrainer:
    def __init__(self, model, ref_model, reward_model, tokenizer, config):
        self.model = model          # 策略模型（训练）
        self.ref_model = ref_model  # 参考模型（冻结）
        self.reward_model = reward_model  # 奖励模型
        self.tokenizer = tokenizer
        self.config = config
    
    def train_step(self, prompts: list):
        batch_size = len(prompts)
        group_size = self.config.group_size  # G=8 或 16
        
        all_advantages = []
        all_old_logps = []
        all_responses = []
        
        # 1. 对每个 prompt 采样 G 个回复
        for prompt in prompts:
            responses = []
            rewards = []
            
            for _ in range(group_size):
                # 生成回复
                response = self._generate(prompt, temperature=0.8)
                responses.append(response)
                
                # 计算奖励
                reward = self.reward_model.score(prompt, response)
                rewards.append(reward)
            
            # 2. GRPO Advantage 计算
            advantages = self._compute_grpo_advantage(rewards)
            all_advantages.extend(advantages)
            all_responses.extend(responses)
        
        # 3. 计算旧策略的 log prob
        old_logps = self._compute_logps(self.model, prompts_expanded, all_responses)
        all_old_logps = old_logps.detach()
        
        # 4. PPO-style 更新（多轮 epoch）
        for epoch in range(self.config.ppo_epochs):
            new_logps = self._compute_logps(self.model, prompts_expanded, all_responses)
            ref_logps = self._compute_logps(self.ref_model, prompts_expanded, all_responses)
            
            # 5. 计算 Loss
            loss = self._grpo_loss(
                new_logps=new_logps,
                old_logps=all_old_logps,
                ref_logps=ref_logps,
                advantages=torch.tensor(all_advantages),
                beta=self.config.kl_coef  # KL 散度惩罚
            )
            
            loss.backward()
            self.optimizer.step()
            self.optimizer.zero_grad()
        
        return loss.item()
    
    def _grpo_loss(self, new_logps, old_logps, ref_logps, advantages, beta=0.04):
        # Ratio
        ratio = torch.exp(new_logps - old_logps)
        
        # Clipped ratio (PPO clip)
        clipped_ratio = torch.clamp(ratio, 1 - self.config.clip_range, 
                                     1 + self.config.clip_range)
        
        # Policy loss
        policy_loss = -torch.min(ratio * advantages, clipped_ratio * advantages).mean()
        
        # KL penalty (GRPO 使用 k3 估计器)
        kl = (torch.exp(ref_logps - new_logps) - (ref_logps - new_logps) - 1).mean()
        
        return policy_loss + beta * kl
    
    def _compute_grpo_advantage(self, rewards: list):
        """GRPO 核心创新：组内相对优势"""
        mean_r = np.mean(rewards)
        std_r = np.std(rewards)
        advantages = [(r - mean_r) / (std_r + 1e-8) for r in rewards]
        return advantages

class RewardComposer:
    def __init__(self):
        self.reward_functions = {
            "correctness": self._correctness_reward,
            "format": self._format_reward,
            "reasoning": self._reasoning_reward,
            "safety": self._safety_reward,
        }
    
    def compute_reward(self, prompt: str, response: str) -> float:
        total = 0
        for name, func in self.reward_functions.items():
            score = func(prompt, response)
            total += score * self.weights.get(name, 1.0)
        return total
    
    def _correctness_reward(self, prompt: str, response: str) -> float:
        """正确性奖励：答案是否正确"""
        # 对于数学题：验证最终答案
        if has_ground_truth(prompt):
            gt = get_ground_truth(prompt)
            pred = extract_answer(response)
            return 1.0 if pred == gt else 0.0
        
        # 对于开放题：用 Judge 模型
        return self.judge.score(prompt, response)
    
    def _format_reward(self, prompt: str, response: str) -> float:
        """格式奖励：是否包含 <think> 标签等"""
        score = 0.0" in response:
            score += 0.5
        if "<answer>" in response and "</answer>" in response:
            score += 0.5
        return score
    
    def _reasoning_reward(self, prompt: str, response: str) -> float:
        """推理质量奖励"""
        # 检查推理步骤的完整性
        reasoning_steps = extract_reasoning_steps(response)
        if len(reasoning_steps) < 2:
            return 0.0
        
        # 检查逻辑连贯性
        coherence = self._check_coherence(reasoning_steps)
        return coherence

# DeepSeek R1 的训练流程
"""
Stage 1: 冷启动 SFT
  - 用高质量推理数据做 SFT
  - 让模型学会基本的推理格式

Stage 2: GRPO 训练（核心）
  - 奖励函数：正确性 + 格式 + 推理质量
  - Group Size: 8-16
  - 多轮迭代：每轮基于上一轮模型生成新数据

Stage 3: 拒绝采样 + SFT
  - 从 GRPO 模型采样多个回复
  - 保留正确的回复做 SFT

Stage 4: 再次 GRPO
  - 进一步对齐
"""