06 - 最新研究进展 (2024-2025)¶

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

class DecisionTransformer(nn.Module):  # 继承nn.Module定义网络层
    """Decision Transformer模型"""

    def __init__(self, state_dim, action_dim, hidden_dim=128,
                 n_layers=3, n_heads=4, max_len=100):
        super(DecisionTransformer, self).__init__()

        self.state_dim = state_dim
        self.action_dim = action_dim
        self.hidden_dim = hidden_dim
        self.max_len = max_len

        # 嵌入层
        self.embed_return = nn.Linear(1, hidden_dim)
        self.embed_state = nn.Linear(state_dim, hidden_dim)
        self.embed_action = nn.Linear(action_dim, hidden_dim)

        # 位置编码
        self.pos_embed = nn.Parameter(torch.randn(1, max_len, hidden_dim))

        # Transformer
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=n_heads,
            dim_feedforward=hidden_dim * 4,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)

        # 预测头
        self.predict_action = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Tanh()
        )

        self.predict_return = nn.Linear(hidden_dim, 1)
        self.predict_state = nn.Linear(hidden_dim, state_dim)

    def forward(self, returns, states, actions, timesteps):
        """
        前向传播

        参数:
            returns: (batch, seq_len, 1)
            states: (batch, seq_len, state_dim)
            actions: (batch, seq_len, action_dim)
            timesteps: (batch, seq_len)

        返回:
            预测的动作、状态、回报
        """
        batch_size, seq_len = states.shape[0], states.shape[1]

        # 嵌入
        return_embeddings = self.embed_return(returns)
        state_embeddings = self.embed_state(states)
        action_embeddings = self.embed_action(actions)

        # 交错排列: [r_1, s_1, a_1, r_2, s_2, a_2, ...]
        stacked_inputs = torch.stack(  # torch.stack沿新维度拼接张量
            [return_embeddings, state_embeddings, action_embeddings], dim=2
        ).reshape(batch_size, 3 * seq_len, self.hidden_dim)  # 重塑张量形状

        # 添加位置编码
        pos_embeddings = self.pos_embed[:, :3 * seq_len, :]
        stacked_inputs = stacked_inputs + pos_embeddings

        # Transformer
        transformer_out = self.transformer(stacked_inputs)

        # 提取动作预测（每隔3个token）
        action_preds = self.predict_action(transformer_out[:, 2::3])

        return action_preds

    def get_action(self, returns, states, actions, timesteps):
        """获取动作预测"""
        action_preds = self.forward(returns, states, actions, timesteps)
        return action_preds[:, -1, :]  # 返回最后一个动作

class DTTrainer:
    """Decision Transformer训练器"""

    def __init__(self, state_dim, action_dim, lr=1e-4):
        self.model = DecisionTransformer(state_dim, action_dim)
        self.optimizer = optim.Adam(self.model.parameters(), lr=lr)
        self.loss_fn = nn.MSELoss()

    def train_step(self, trajectories):
        """
        训练一步

        参数:
            trajectories: 轨迹批次
                - returns: (batch, seq_len, 1)
                - states: (batch, seq_len, state_dim)
                - actions: (batch, seq_len, action_dim)
        """
        returns = trajectories['returns']
        states = trajectories['states']
        actions = trajectories['actions']

        # 前向传播
        action_preds = self.model(returns, states, actions, None)

        # 计算损失（只计算动作预测）
        loss = self.loss_fn(action_preds, actions)

        # 反向传播
        self.optimizer.zero_grad()  # 清零梯度
        loss.backward()  # 反向传播计算梯度
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 0.25)
        self.optimizer.step()  # 更新参数

        return loss.item()  # 将单元素张量转为Python数值

class OnlineDecisionTransformer:
    """Online Decision Transformer"""

    def __init__(self, state_dim, action_dim):
        self.dt_model = DecisionTransformer(state_dim, action_dim)
        self.exploration_bonus = 0.1  # 探索奖励系数

    def select_action(self, state, target_return):
        """选择动作（带探索）"""
        # 使用DT模型预测动作
        with torch.no_grad():  # 禁用梯度计算，节省内存
            action = self.dt_model.get_action(
                target_return, state, prev_actions, timesteps
            )

        # 添加探索噪声
        action = action + np.random.normal(0, self.exploration_bonus, size=action.shape)

        return action

    def online_update(self, new_trajectory):
        """在线更新"""
        # 将新轨迹加入缓冲区
        self.replay_buffer.add(new_trajectory)

        # 从缓冲区采样训练
        batch = self.replay_buffer.sample()
        loss = self.dt_model.train_step(batch)

        return loss

class VLAModel(nn.Module):
    """视觉-语言-动作模型"""

    def __init__(self, vocab_size, image_size, action_dim):
        super(VLAModel, self).__init__()

        # 视觉编码器
        self.vision_encoder = VisionEncoder(image_size)

        # 语言编码器
        self.language_encoder = LanguageEncoder(vocab_size)

        # 多模态融合
        self.fusion = MultiModalFusion()

        # 动作解码器
        self.action_decoder = ActionDecoder(action_dim)

    def forward(self, image, text, previous_actions):
        """
        前向传播

        参数:
            image: 视觉输入
            text: 语言指令
            previous_actions: 历史动作
        """
        # 编码
        visual_features = self.vision_encoder(image)
        language_features = self.language_encoder(text)

        # 融合
        fused_features = self.fusion(visual_features, language_features, previous_actions)

        # 解码动作
        action = self.action_decoder(fused_features)

        return action

import torch
import torch.nn.functional as F

class GRPOTrainer:
    """
    GRPO训练器（简化版）

    核心思想：
    1. 对每个问题，用当前策略采样G个回答
    2. 用规则/奖励模型对每个回答打分
    3. 在组内做标准化得到相对优势
    4. 用PPO-clip目标更新策略
    """

    def __init__(self, policy_model, ref_model, reward_fn,
                 group_size=8, clip_eps=0.2, kl_coeff=0.01, lr=1e-6):
        self.policy = policy_model        # 策略模型 π_θ
        self.ref_model = ref_model          # 参考模型 π_ref（冻结）
        self.reward_fn = reward_fn          # 奖励函数 r(q, o)
        self.group_size = group_size        # 每个问题采样数 G
        self.clip_eps = clip_eps            # PPO裁剪参数 ε
        self.kl_coeff = kl_coeff            # KL散度系数 β
        self.optimizer = torch.optim.Adam(self.policy.parameters(), lr=lr)

    def compute_group_advantages(self, rewards):
        """
        计算组内相对优势（GRPO核心）

        参数:
            rewards: shape (batch_size, group_size) -- 每组G个回答的奖励
        返回:
            advantages: shape (batch_size, group_size) -- 标准化后的优势
        """
        # 组内均值和标准差
        mean_r = rewards.mean(dim=-1, keepdim=True)       # (batch, 1)
        std_r = rewards.std(dim=-1, keepdim=True) + 1e-8  # (batch, 1)

        # 标准化
        advantages = (rewards - mean_r) / std_r            # (batch, G)
        return advantages

    def compute_kl_divergence(self, log_probs_policy, log_probs_ref):
        """计算策略与参考模型的KL散度"""
        kl = (log_probs_policy - log_probs_ref).mean()
        return kl

    def train_step(self, questions):
        """
        GRPO训练一步

        参数:
            questions: 一批问题 [q_1, q_2, ..., q_B]
        """
        all_losses = []

        for question in questions:
            # ===== Step 1: 采样一组回答 =====
            with torch.no_grad():
                outputs = []
                for _ in range(self.group_size):
                    output = self.policy.generate(question, temperature=0.7)
                    outputs.append(output)

            # ===== Step 2: 计算奖励 =====
            rewards = torch.tensor([
                self.reward_fn(question, output) for output in outputs
            ])  # shape: (G,)

            # ===== Step 3: 计算组内相对优势 =====
            advantages = self.compute_group_advantages(rewards.unsqueeze(0)).squeeze(0)  # unsqueeze增加一个维度  # squeeze压缩维度

            # ===== Step 4: 计算策略比率和损失 =====
            for i, output in enumerate(outputs):  # enumerate同时获取索引和元素
                # 当前策略的对数概率
                log_prob = self.policy.log_prob(question, output)
                # 旧策略的对数概率（采样时的）
                with torch.no_grad():
                    log_prob_old = self.policy.log_prob(question, output)
                    log_prob_ref = self.ref_model.log_prob(question, output)

                # 策略比率
                ratio = torch.exp(log_prob - log_prob_old)

                # PPO-Clip损失
                surr1 = ratio * advantages[i]
                surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages[i]
                policy_loss = -torch.min(surr1, surr2)

                # KL惩罚
                kl_penalty = self.kl_coeff * self.compute_kl_divergence(log_prob, log_prob_ref)

                loss = policy_loss + kl_penalty
                all_losses.append(loss)

        # ===== Step 5: 反向传播 =====
        total_loss = torch.stack(all_losses).mean()
        self.optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.policy.parameters(), max_norm=1.0)
        self.optimizer.step()

        return total_loss.item()

def math_reward_function(question, answer):
    """
    数学推理奖励函数示例

    DeepSeek-R1使用的奖励信号：
    - 准确性奖励：答案是否正确（0/1）
    - 格式奖励：是否按要求格式输出（如\\boxed{}）
    """
    # 提取模型输出的最终答案
    predicted = extract_answer(answer)
    # 获取标准答案
    ground_truth = get_ground_truth(question)

    # 准确性奖励
    accuracy_reward = 1.0 if predicted == ground_truth else 0.0

    # 格式奖励（鼓励使用思维链和规范格式）
    format_reward = 0.0
    if "<think>" in answer and "</think>" in answer:
        format_reward += 0.1
    if "\\boxed{" in answer:
        format_reward += 0.1

    return accuracy_reward + format_reward

class STaRTrainer:
    """STaR (Self-Taught Reasoner) 训练器"""

    def __init__(self, model, tokenizer, num_iterations=5):
        self.model = model
        self.tokenizer = tokenizer
        self.num_iterations = num_iterations

    def generate_rationale(self, question, hint=None):
        """生成推理过程"""
        if hint:
            prompt = f"Question: {question}\nHint: The answer is {hint}\nLet's think step by step:"
        else:
            prompt = f"Question: {question}\nLet's think step by step:"

        output = self.model.generate(prompt, max_length=512, temperature=0.7)
        rationale, answer = self.parse_output(output)
        return rationale, answer

    def star_iteration(self, dataset):
        """STaR一轮迭代"""
        training_data = []

        for question, ground_truth in dataset:
            # Step 1: 直接推理
            rationale, predicted = self.generate_rationale(question)

            if predicted == ground_truth:
                # Step 2: 正确 → 保留
                training_data.append({
                    'question': question,
                    'rationale': rationale,
                    'answer': predicted
                })
            else:
                # Step 3: 错误 → 合理化（提供答案作为hint）
                rationale_hint, predicted_hint = self.generate_rationale(
                    question, hint=ground_truth
                )
                if predicted_hint == ground_truth:
                    training_data.append({
                        'question': question,
                        'rationale': rationale_hint,
                        'answer': predicted_hint
                    })

        # Step 4: 微调
        self.finetune(training_data)
        return len(training_data)

    def train(self, dataset):
        """完整STaR训练"""
        for i in range(self.num_iterations):
            n_samples = self.star_iteration(dataset)
            accuracy = self.evaluate(dataset)
            print(f"Iteration {i+1}: {n_samples} samples, accuracy={accuracy:.2%}")

class ProcessRewardModel:
    """过程奖励模型（PRM）"""

    def __init__(self, base_model, step_separator="\\n\\n"):
        self.model = base_model
        self.step_separator = step_separator
        self.reward_head = torch.nn.Linear(base_model.hidden_size, 1)

    def score_steps(self, question, reasoning_steps):
        """
        对推理过程的每一步打分

        参数:
            question: 问题
            reasoning_steps: ["Step 1: ...", "Step 2: ...", ...]

        返回:
            step_scores: 每步的得分 [0.9, 0.85, 0.3, ...]
        """
        step_scores = []
        context = f"Question: {question}\n"

        for step in reasoning_steps:
            context += step + self.step_separator
            # 编码到当前步骤的上下文
            hidden = self.model.encode(context)
            # 取最后一个token的隐藏状态
            score = torch.sigmoid(self.reward_head(hidden[:, -1, :]))
            step_scores.append(score.item())

        return step_scores

    def best_of_n_with_prm(self, question, candidates, aggregation="min"):
        """
        使用PRM进行Best-of-N选择

        aggregation:
          - "min": 取所有步骤得分的最小值（保守策略）
          - "last": 取最后一步的得分
          - "prod": 取所有步骤得分的乘积
        """
        best_score = -float('inf')
        best_candidate = None

        for candidate in candidates:
            steps = candidate.split(self.step_separator)
            scores = self.score_steps(question, steps)

            if aggregation == "min":
                final_score = min(scores)
            elif aggregation == "last":
                final_score = scores[-1]  # [-1]负索引取最后元素
            elif aggregation == "prod":
                final_score = 1.0
                for s in scores:
                    final_score *= s

            if final_score > best_score:
                best_score = final_score
                best_candidate = candidate

        return best_candidate, best_score

class ReasoningSearch:
    """推理搜索策略"""

    @staticmethod  # @staticmethod不需要实例即可调用
    def best_of_n(model, question, n=64, reward_model=None):
        """
        Best-of-N采样
        最简单的方法：生成N个答案，选最好的
        """
        candidates = []
        for _ in range(n):
            response = model.generate(question, temperature=0.7)
            score = reward_model.score(question, response)
            candidates.append((response, score))

        return max(candidates, key=lambda x: x[1])  # lambda匿名函数

    @staticmethod
    def beam_search_with_prm(model, question, beam_width=5, max_steps=10, prm=None):
        """
        带PRM引导的Beam Search
        在推理步骤层面进行搜索
        """
        # 初始beam
        beams = [{"steps": [], "score": 0.0, "context": f"Q: {question}\n"}]

        for step_idx in range(max_steps):
            all_candidates = []

            for beam in beams:
                # 对每个beam，生成多个候选下一步
                for _ in range(beam_width):
                    next_step = model.generate_step(beam["context"], temperature=0.8)

                    # PRM打分
                    step_score = prm.score_step(question, beam["steps"] + [next_step])

                    all_candidates.append({
                        "steps": beam["steps"] + [next_step],
                        "score": beam["score"] + step_score,
                        "context": beam["context"] + next_step + "\n"
                    })

            # 保留top-k
            beams = sorted(all_candidates, key=lambda x: x["score"], reverse=True)[:beam_width]

            # 检查是否有beam已完成
            for beam in beams:
                if is_answer_complete(beam["steps"][-1]):
                    return beam

        return beams[0]

    @staticmethod
    def monte_carlo_tree_search(model, question, num_simulations=100, prm=None):
        """
        蒙特卡洛树搜索（MCTS）用于推理
        在step级别构建搜索树
        """
        root = MCTSNode(state=question, parent=None)

        for _ in range(num_simulations):
            # Selection：选择最有前途的节点
            node = root.select(exploration_weight=1.4)

            # Expansion：生成新的推理步骤
            next_step = model.generate_step(node.state, temperature=1.0)
            child = node.expand(next_step)

            # Simulation：完成剩余推理并评估
            full_solution = model.complete_reasoning(child.state)
            reward = evaluate_solution(question, full_solution)

            # Backpropagation：更新统计
            child.backpropagate(reward)

        # 返回最佳路径
        return root.best_path()

class CodeAgentRewardFunction:
    """代码Agent的奖励函数设计"""

    def __init__(self, test_weight=1.0, quality_weight=0.1, efficiency_weight=0.05):
        self.test_weight = test_weight
        self.quality_weight = quality_weight
        self.efficiency_weight = efficiency_weight

    def compute_reward(self, trajectory):
        """
        计算代码修复轨迹的奖励

        trajectory包含:
          - issue_description: Bug描述
          - actions: Agent执行的操作序列
          - final_patch: 最终提交的代码补丁
          - test_results: 测试运行结果
        """
        # 1. 测试通过奖励（核心）
        test_reward = self.compute_test_reward(trajectory['test_results'])

        # 2. 代码质量奖励
        quality_reward = self.compute_quality_reward(trajectory['final_patch'])

        # 3. 效率惩罚（步骤越少越好）
        efficiency_penalty = len(trajectory['actions']) * self.efficiency_weight

        total_reward = (
            self.test_weight * test_reward +
            self.quality_weight * quality_reward -
            efficiency_penalty
        )

        return total_reward

    def compute_test_reward(self, test_results):
        """基于测试结果的奖励"""
        if test_results['all_passed']:
            return 1.0
        else:
            # 部分奖励：通过的测试比例
            return test_results['passed'] / test_results['total'] * 0.5

    def compute_quality_reward(self, patch):
        """基于代码质量的奖励"""
        score = 0.0
        # 修改行数适当
        if len(patch.split('\n')) < 50:
            score += 0.3
        # 无语法错误
        if not has_syntax_errors(patch):
            score += 0.5
        # 包含注释
        if has_comments(patch):
            score += 0.2
        return score

class CodeExecutionReward:
    """基于代码执行的奖励计算"""

    def __init__(self, timeout=10):
        self.timeout = timeout

    def compute_reward(self, generated_code, test_cases):
        """
        执行代码并计算奖励

        test_cases: [
            {"input": "...", "expected_output": "..."},
            ...
        ]
        """
        reward = 0.0

        # Level 1: 语法检查
        if not self.syntax_check(generated_code):
            return -0.5  # 语法错误

        # Level 2: 编译/解析检查
        if not self.compile_check(generated_code):
            return -0.3  # 编译错误

        # Level 3: 执行测试用例
        passed = 0
        for test in test_cases:
            try:  # try/except捕获异常
                result = self.execute(generated_code, test['input'], timeout=self.timeout)
                if result == test['expected_output']:
                    passed += 1
            except TimeoutError:
                reward -= 0.1  # 超时惩罚
            except RuntimeError:
                reward -= 0.1  # 运行时错误

        # 通过率奖励
        pass_rate = passed / len(test_cases)
        reward += pass_rate  # [0, 1]

        return reward

    def multi_level_reward(self, generated_code, test_cases):
        """
        多层级奖励（更细粒度）

        返回:
            dict: 各维度的得分
        """
        return {
            'syntax': self.syntax_score(generated_code),        # 语法正确性
            'compilable': self.compile_score(generated_code),    # 可编译性
            'functional': self.functional_score(generated_code, test_cases),  # 功能正确
            'efficiency': self.efficiency_score(generated_code, test_cases),  # 运行效率
            'style': self.style_score(generated_code)           # 代码风格
        }