第九章 Agentic RL （ Agent 强化学习）¶

# 概念性示意: Agent RL vs Agent SFT
from torch.nn.functional import cross_entropy

# SFT方式训练Agent（模仿学习）
def train_agent_sft(model, expert_demos):
    """
    给定专家的(状态, 动作)对，训练模型模仿。
    问题: 模型只学会了"在这个状态下，专家会做X"，
    但不理解"为什么做X"以及"如果X失败了怎么办"。
    """
    for state, expert_action in expert_demos:
        loss = cross_entropy(model(state), expert_action)
        loss.backward()

# RL方式训练Agent（探索学习）
def train_agent_rl(model, environment):
    """
    让Agent在环境中自主探索，通过奖励信号学习最优策略。
    优点: 模型学会了"什么行为能带来好结果"，包括错误恢复。
    """
    for episode in range(num_episodes):
        state = environment.reset()
        trajectory = []

        while not environment.done:
            # Agent自主选择动作（可能是好的，也可能是坏的）
            action = model.sample_action(state)

            # 环境返回奖励和新状态
            next_state, reward = environment.step(action)
            trajectory.append((state, action, reward))
            state = next_state

        # 根据整个轨迹的奖励更新模型
        update_policy(model, trajectory)

# 三种训练方法的核心区别（概念性代码）
import torch

class PPOTrainer:
    """
    PPO (Proximal Policy Optimization)
    经典RLHF方法，需要独立的奖励模型。
    """
    def train_step(self, prompts, old_model, reward_model):
        # 1. 当前模型和旧模型分别生成回复
        new_responses = self.model.generate(prompts)
        old_responses = old_model.generate(prompts)

        # 2. 奖励模型打分
        rewards = reward_model.score(prompts, new_responses)

        # 3. 计算优势函数
        advantages = self.compute_advantages(rewards)

        # 4. PPO裁剪更新
        ratio = self.model.log_prob(new_responses) / old_model.log_prob(new_responses)
        clipped_ratio = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon)
        loss = -torch.min(ratio * advantages, clipped_ratio * advantages).mean()

        # 需要: 奖励模型 + 旧模型副本 → 显存需求大
        return loss

class DPOTrainer:
    """
    DPO (Direct Preference Optimization)
    直接从偏好对学习，不需要独立的奖励模型。
    """
    def train_step(self, prompts, chosen, rejected, ref_model):
        # 1. 计算当前模型和参考模型的log概率
        pi_chosen = self.model.log_prob(prompts, chosen)
        pi_rejected = self.model.log_prob(prompts, rejected)
        ref_chosen = ref_model.log_prob(prompts, chosen)
        ref_rejected = ref_model.log_prob(prompts, rejected)

        # 2. DPO损失（隐式奖励）
        log_ratio_chosen = pi_chosen - ref_chosen
        log_ratio_rejected = pi_rejected - ref_rejected

        loss = -torch.log(
            torch.sigmoid(self.beta * (log_ratio_chosen - log_ratio_rejected))
        ).mean()

        # 需要: 偏好对数据(chosen/rejected) + 参考模型
        # 不需要: 独立奖励模型
        return loss

class GRPOTrainer:
    """
    GRPO (Group Relative Policy Optimization)
    DeepSeek提出，用组内相对排名代替绝对奖励。
    特别适合Agent场景！
    """
    def train_step(self, prompts, reward_fn):
        """
        GRPO的关键创新：
        1. 对每个prompt生成一组(G个)回复
        2. 用规则奖励函数给每个回复打分
        3. 在组内计算相对优势（不需要奖励模型）
        """
        group_size = 8  # 每个prompt生成8个回复

        all_responses = []
        all_rewards = []

        for prompt in prompts:
            # 1. 采样一组回复
            responses = [self.model.generate(prompt) for _ in range(group_size)]

            # 2. 规则奖励函数打分（不需要训练奖励模型！）
            rewards = [reward_fn(prompt, resp) for resp in responses]

            all_responses.extend(responses)
            all_rewards.extend(rewards)

        # 3. 组内归一化计算优势
        rewards_tensor = torch.tensor(all_rewards)
        # 关键：在每组内做归一化
        for i in range(0, len(rewards_tensor), group_size):
            group = rewards_tensor[i:i+group_size]
            mean = group.mean()
            std = group.std() + 1e-8
            rewards_tensor[i:i+group_size] = (group - mean) / std

        # 4. 策略梯度更新
        loss = 0
        for resp, advantage in zip(all_responses, rewards_tensor):
            log_prob = self.model.log_prob(resp)
            loss -= log_prob * advantage

        return loss / len(all_responses)

# 概念性伪代码：以下辅助函数（has_valid_tool_call_format, extract_tool_calls,
# task_completed_successfully, count_steps, contains_dangerous_action）
# 展示奖励函数的设计思路，实际实现需根据Agent的输出格式定制。

def agent_reward_function(prompt: str, agent_trajectory: str) -> float:
    """
    GRPO的一大优势：可以使用基于规则的奖励函数。
    对于Agent，我们可以精确定义什么是"好的行为"。
    """
    reward = 0.0

    # 1. 格式奖励：Agent的输出是否符合预期格式？
    if has_valid_tool_call_format(agent_trajectory):
        reward += 1.0
    else:
        reward -= 2.0  # 格式错误是严重问题

    # 2. 工具使用奖励：是否正确使用了工具？
    tool_calls = extract_tool_calls(agent_trajectory)
    for call in tool_calls:
        if call.is_valid:
            reward += 0.5
        if call.result_used_in_reasoning:
            reward += 1.0  # 使用了工具结果进行推理

    # 3. 任务完成奖励：最终是否达成目标？
    if task_completed_successfully(prompt, agent_trajectory):
        reward += 5.0

    # 4. 效率奖励：用更少的步骤完成加分
    num_steps = count_steps(agent_trajectory)
    if num_steps <= 3 and task_completed_successfully(prompt, agent_trajectory):
        reward += 2.0  # 高效完成
    elif num_steps > 10:
        reward -= 1.0  # 效率低

    # 5. 安全奖励：避免危险操作
    if contains_dangerous_action(agent_trajectory):
        reward -= 10.0

    return reward

import re
import json
from dataclasses import dataclass
from typing import Callable

@dataclass
class RewardComponent:
    """奖励组件"""
    name: str
    weight: float
    compute: Callable  # (prompt, trajectory) -> float
    description: str

class AgentRewardSystem:
    """
    分层Agent奖励系统
    四个层次: 格式 → 工具使用 → 推理质量 → 任务完成
    """

    def __init__(self):
        self.components = [
            # Level 1: 格式正确性 (基础要求)
            RewardComponent(
                name="format_compliance",
                weight=1.0,
                compute=self._format_reward,
                description="Agent输出是否符合预期的JSON/XML格式"
            ),

            # Level 2: 工具使用质量
            RewardComponent(
                name="tool_usage",
                weight=2.0,
                compute=self._tool_usage_reward,
                description="工具调用是否正确、是否充分利用了工具结果"
            ),

            # Level 3: 推理质量
            RewardComponent(
                name="reasoning_quality",
                weight=1.5,
                compute=self._reasoning_reward,
                description="推理过程是否逻辑清晰、是否有效推进任务"
            ),

            # Level 4: 任务完成度
            RewardComponent(
                name="task_completion",
                weight=3.0,
                compute=self._task_completion_reward,
                description="是否最终完成了用户的任务"
            ),
        ]

    def compute_total_reward(self, prompt: str, trajectory: str) -> dict:
        """计算总奖励及各组件分数"""
        results = {}
        total = 0.0

        for component in self.components:
            score = component.compute(prompt, trajectory)
            weighted = score * component.weight
            results[component.name] = {
                "raw_score": score,
                "weight": component.weight,
                "weighted_score": weighted,
            }
            total += weighted

        results["total"] = total
        return results

    def _format_reward(self, prompt: str, trajectory: str) -> float:
        """
        Level 1: 格式奖励
        检查Agent输出是否符合结构化格式要求
        """
        score = 0.0

        # 检查思考/行动/观察标记
        if "<think>" in trajectory and "</think>" in trajectory:
            score += 0.3
        if "<action>" in trajectory and "</action>" in trajectory:
            score += 0.3

        # 检查JSON工具调用格式
        # 注意：正则提取JSON工具调用较脆弱，可能匹配到非工具调用的JSON片段。
        # 生产环境建议使用结构化解析（如要求LLM输出固定JSON schema）而非正则。
        tool_calls = re.findall(r'\{.*?"name".*?"arguments".*?\}', trajectory)
        if tool_calls:
            for call in tool_calls:
                try:
                    json.loads(call)
                    score += 0.2
                except json.JSONDecodeError:
                    score -= 0.3  # 格式错误扣分

        return min(max(score, -1.0), 1.0)

    def _tool_usage_reward(self, prompt: str, trajectory: str) -> float:
        """
        Level 2: 工具使用奖励
        评估工具使用的合理性和有效性

        注意：当前使用简单的字符串匹配（如 '<action>' 标签）来检测工具调用，
        可能在trajectory包含讨论工具调用格式的文本时产生误报。
        生产环境建议使用结构化的工具调用记录而非从原始文本中提取。
        """
        score = 0.0

        # 是否使用了工具（大多数Agent任务都需要工具）
        tool_count = trajectory.count("<action>")
        if tool_count > 0:
            score += 0.3

        # 是否重复调用相同工具（不好的行为）
        actions = re.findall(r'<action>(.*?)</action>', trajectory)
        if len(actions) > len(set(actions)):
            score -= 0.2  # 重复调用扣分

        # 是否在观察后调整了策略（好的行为）
        observations = re.findall(r'<observation>(.*?)</observation>', trajectory)
        if len(observations) > 1:
            # 检查后续动作是否与前一个不同（说明Agent在学习）
            score += 0.3

        return min(max(score, -1.0), 1.0)

    def _reasoning_reward(self, prompt: str, trajectory: str) -> float:
        """
        Level 3: 推理质量奖励
        这通常需要LLM-as-Judge来评估
        """
        # 简单规则版本
        score = 0.0

        thoughts = re.findall(r'<think>(.*?)</think>', trajectory, re.DOTALL)

        for thought in thoughts:
            # 推理是否有实质内容（不是废话）
            if len(thought.strip()) > 50:
                score += 0.2
            # 是否引用了之前的观察
            if "observation" in thought.lower() or "result" in thought.lower():
                score += 0.1

        return min(max(score, -1.0), 1.0)

    def _task_completion_reward(self, prompt: str, trajectory: str) -> float:
        """
        Level 4: 任务完成奖励
        通常需要外部验证（如运行测试、检查答案）
        """
        # 这里需要根据具体任务类型来实现
        # 示例：代码生成任务
        if "<final_answer>" in trajectory:
            return 0.5  # 至少给出了答案
        return 0.0

# ❌ 常见奖励设计错误

# 错误1: 奖励篇幅而非质量
bad_reward_1 = lambda _, traj: len(traj) / 1000  # lambda匿名函数：_表示忽略第一个参数，只用轨迹长度算奖励
# → Agent会生成冗长无用的内容

# 错误2: 只奖励最终结果，忽略过程
bad_reward_2 = lambda _, traj: 10.0 if "correct" in traj else 0.0  # lambda+三元表达式：包含correct得10分否则0分
# → Agent无法学习中间步骤的好坏

# 错误3: 惩罚过重导致过于保守
bad_reward_3 = lambda _, traj: -100.0 if any_error(traj) else 1.0
# → Agent学会了"什么都不做"以避免惩罚

# ✅ 好的奖励设计原则
"""
1. 过程奖励 > 结果奖励: 每一步都给反馈
2. 密集奖励 > 稀疏奖励: 避免大量0分
3. 相对排名 > 绝对分数: GRPO的核心优势
4. 多维度 > 单一分数: 分开评估格式/工具/推理/结果
5. 可验证 > 主观评价: 尽量用规则而非LLM-as-Judge
"""

"""
GRPO训练工具调用Agent的完整示例

目标: 训练一个Agent学会使用计算器/搜索/代码执行器工具解决数学问题
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.optim import AdamW

# === Step 1: 定义工具环境 ===

class ToolEnvironment:
    """Agent的工具环境"""

    def __init__(self):
        self.tools = {
            "calculator": self._calculator,
            "search": self._search,
            "code_executor": self._code_executor,
        }
        self.tool_descriptions = {
            "calculator": "Evaluate a mathematical expression. Input: expression (str)",
            "search": "Search for factual information. Input: query (str)",
            "code_executor": "Execute Python code and return output. Input: code (str)",
        }

    def _calculator(self, expression: str) -> str:
        try:
            # 安全的数学表达式求值：使用 AST 白名单检查而非字符集过滤
            import ast
            tree = ast.parse(expression, mode='eval')
            for node in ast.walk(tree):
                if not isinstance(node, (ast.Expression, ast.BinOp, ast.UnaryOp,
                                        ast.Constant, ast.Add, ast.Sub, ast.Mult,
                                        ast.Div, ast.Mod, ast.Pow, ast.USub, ast.UAdd)):
                    return "Error: invalid expression"
            return str(eval(compile(tree, '<expr>', 'eval'), {"__builtins__": {}}, {}))
        except Exception as e:
            return f"Error: {e}"

    def _search(self, query: str) -> str:
        # 模拟搜索（实际应用中接入真实搜索API）
        knowledge_base = {
            "pi": "Pi (π) is approximately 3.14159265358979",
            "speed_of_light": "Speed of light is 299,792,458 m/s",
            "earth_radius": "Earth's mean radius is 6,371 km",
        }
        for key, value in knowledge_base.items():
            if key in query.lower():
                return value
        return f"No results found for: {query}"

    def _code_executor(self, code: str) -> str:
        try:
            # 受限的代码执行环境
            local_vars = {}
            # exec动态执行代码字符串；{"__builtins__": {}}清空内置函数以限制权限，local_vars接收执行结果
            exec(code, {"__builtins__": {}}, local_vars)
            return str(local_vars.get("result", "No 'result' variable defined"))
        except Exception as e:
            return f"Execution error: {e}"

    def execute_tool(self, name: str, args: str) -> str:
        if name in self.tools:
            return self.tools[name](args)
        return f"Unknown tool: {name}"

    def get_tool_prompt(self) -> str:
        desc = "\n".join(
            f"- {name}: {desc}"
            for name, desc in self.tool_descriptions.items()
        )
        return f"Available tools:\n{desc}"

# === Step 2: 定义奖励函数 ===

def compute_agent_reward(prompt: str, response: str,
                         environment: ToolEnvironment,
                         expected_answer: str = None) -> float:
    """
    Agent行为的综合奖励函数
    """
    import re
    import json

    reward = 0.0

    # R1: 格式正确性 (±2.0)
    # 检查是否使用了正确的思考-行动格式
    has_thought = bool(re.search(r'<think>.*?</think>', response, re.DOTALL))
    has_action = bool(re.search(r'<action>.*?</action>', response, re.DOTALL))
    has_answer = bool(re.search(r'<answer>.*?</answer>', response, re.DOTALL))

    if has_thought:
        reward += 0.5
    if has_action or has_answer:
        reward += 0.5
    if not (has_thought or has_action or has_answer):
        reward -= 2.0  # 完全不遵循格式

    # R2: 工具调用有效性 (±3.0)
    tool_calls = re.findall(
        r'<action>\s*\{.*?"tool":\s*"(\w+)".*?"input":\s*"(.*?)".*?\}\s*</action>',
        response, re.DOTALL
    )

    valid_calls = 0
    for tool_name, tool_input in tool_calls:
        if tool_name in environment.tools:
            result = environment.execute_tool(tool_name, tool_input)
            if "Error" not in result:
                valid_calls += 1
                reward += 1.0
            else:
                reward -= 0.5
        else:
            reward -= 1.0

    # R3: 答案正确性 (±5.0)
    if expected_answer and has_answer:
        answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)
        if answer_match:
            agent_answer = answer_match.group(1).strip()
            if expected_answer.strip().lower() in agent_answer.lower():
                reward += 5.0
            else:
                reward -= 1.0

    # R4: 效率奖励 (±1.0)
    num_steps = len(tool_calls)
    if num_steps <= 3 and reward > 3:
        reward += 1.0  # 高效完成任务
    elif num_steps > 8:
        reward -= 1.0  # 过多步骤

    return reward

# === Step 3: GRPO训练循环 ===

class GRPOTrainer:
    """
    简化版GRPO训练器
    用于演示核心训练逻辑
    """

    def __init__(self, model, tokenizer, environment,
                 group_size=4, lr=1e-5, kl_coeff=0.1):
        self.model = model
        self.tokenizer = tokenizer
        self.env = environment
        self.group_size = group_size
        self.kl_coeff = kl_coeff
        self.optimizer = AdamW(model.parameters(), lr=lr)

        # 参考模型（用于KL散度约束）
        self.ref_model = AutoModelForCausalLM.from_pretrained(
            model.config.name_or_path
        )
        self.ref_model.eval()

    def train_step(self, batch: list[dict]) -> dict:
        """
        GRPO训练的一步

        batch: [{"prompt": str, "expected_answer": str}, ...]
        """
        all_log_probs = []
        all_ref_log_probs = []
        all_rewards = []
        all_advantages = []

        for item in batch:
            prompt = item["prompt"]
            expected = item.get("expected_answer")

            # 1. 为每个prompt生成group_size个回复
            group_rewards = []
            group_log_probs = []
            group_ref_log_probs = []

            for _ in range(self.group_size):
                # 采样生成
                inputs = self.tokenizer(prompt, return_tensors="pt")
                with torch.no_grad():
                    output = self.model.generate(
                        **inputs,
                        max_new_tokens=512,
                        do_sample=True,
                        temperature=0.7,
                        top_p=0.9,
                    )

                response = self.tokenizer.decode(
                    output[0][inputs.input_ids.shape[1]:],
                    skip_special_tokens=True
                )

                # 计算奖励
                reward = compute_agent_reward(
                    prompt, response, self.env, expected
                )
                group_rewards.append(reward)

                # 计算log概率
                log_prob = self._compute_log_prob(inputs, output)
                group_log_probs.append(log_prob)

                ref_log_prob = self._compute_ref_log_prob(inputs, output)
                group_ref_log_probs.append(ref_log_prob)

            # 2. 组内归一化（GRPO的核心）
            rewards_tensor = torch.tensor(group_rewards)
            mean = rewards_tensor.mean()
            std = rewards_tensor.std() + 1e-8
            advantages = (rewards_tensor - mean) / std

            all_advantages.extend(advantages.tolist())
            all_log_probs.extend(group_log_probs)
            all_ref_log_probs.extend(group_ref_log_probs)
            all_rewards.extend(group_rewards)

        # 3. 策略梯度更新 + KL约束
        policy_loss = 0
        kl_loss = 0

        for log_prob, ref_log_prob, advantage in zip(
            all_log_probs, all_ref_log_probs, all_advantages
        ):
            policy_loss -= log_prob * advantage
            kl_loss += (log_prob - ref_log_prob)  # KL散度近似

        total_loss = policy_loss + self.kl_coeff * kl_loss
        total_loss = total_loss / len(all_log_probs)

        self.optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
        self.optimizer.step()

        return {
            "loss": total_loss.item(),
            "mean_reward": sum(all_rewards) / len(all_rewards),
            "reward_std": torch.tensor(all_rewards).std().item(),
        }

    def _compute_log_prob(self, inputs, output):
        """计算当前模型的log概率（简化近似）"""
        with torch.enable_grad():
            # 注意：output包含完整序列(prompt+生成)，需用它同时作为input和labels
            # 这是粗略近似（包含了prompt部分），精确实现应只计算生成部分的log概率
            # 精确实现应屏蔽prompt token：
            #   prompt_len = inputs.shape[-1]
            #   labels = output.clone()
            #   labels[:, :prompt_len] = -100  # -100为CrossEntropyLoss的忽略索引
            #   outputs = self.model(input_ids=output, labels=labels)
            outputs = self.model(input_ids=output, labels=output)
            return -outputs.loss  # 负的平均交叉熵 ≈ 序列log概率的近似

    def _compute_ref_log_prob(self, inputs, output):
        """计算参考模型的log概率（简化近似）"""
        with torch.no_grad():
            outputs = self.ref_model(input_ids=output, labels=output)
            return -outputs.loss

# === Step 4: 训练数据准备 ===

TRAINING_DATA = [
    {
        "prompt": "What is 15% of 380? Use the calculator tool to compute this.",
        "expected_answer": "57"
    },
    {
        "prompt": "What is the circumference of the Earth in km? Search for Earth's radius and then calculate.",
        "expected_answer": "40030"
    },
    {
        "prompt": "Calculate the compound interest on $1000 at 5% for 3 years.",
        "expected_answer": "1157.625"
    },
]

# === Step 5: 训练主循环 ===

def main():
    """GRPO训练主流程"""
    # 初始化
    model_name = "Qwen/Qwen2.5-1.5B-Instruct"  # 用小模型演示
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    env = ToolEnvironment()

    trainer = GRPOTrainer(
        model=model,
        tokenizer=tokenizer,
        environment=env,
        group_size=4,
        lr=1e-5,
    )

    # 训练循环
    num_epochs = 10
    for epoch in range(num_epochs):
        metrics = trainer.train_step(TRAINING_DATA)
        print(
            f"Epoch {epoch+1}/{num_epochs} | "
            f"Loss: {metrics['loss']:.4f} | "
            f"Mean Reward: {metrics['mean_reward']:.2f} | "
            f"Reward Std: {metrics['reward_std']:.2f}"
        )

    # 保存模型
    model.save_pretrained("./agent-grpo-trained")
    tokenizer.save_pretrained("./agent-grpo-trained")
    print("Training complete! Model saved to ./agent-grpo-trained")

if __name__ == "__main__":
    main()

from abc import ABC, abstractmethod
from typing import Any

class AgentEnvironment(ABC):  # 抽象基类，定义Agent训练环境接口
    """Agent训练环境的基类"""

    @abstractmethod
    def reset(self, task: dict) -> str:
        """重置环境，返回初始观察"""
        pass

    @abstractmethod
    def step(self, action: str) -> tuple[str, float, bool]:
        """
        执行动作
        Returns: (observation, reward, done)
        """
        pass

    @abstractmethod
    def get_available_tools(self) -> list[dict]:
        """返回可用工具列表"""
        pass

class CodeFixEnvironment(AgentEnvironment):
    """
    代码修复环境 - 用于训练代码调试Agent
    """

    def __init__(self, test_cases: list[dict]):
        self.test_cases = test_cases
        self.current_task = None
        self.current_code = ""
        self.step_count = 0
        self.max_steps = 15

    def reset(self, task: dict) -> str:
        self.current_task = task
        self.current_code = task["buggy_code"]
        self.step_count = 0

        return (
            f"Fix the following code:\n```python\n{self.current_code}\n```\n"
            f"Test description: {task['test_description']}"
        )

    def step(self, action: str) -> tuple[str, float, bool]:
        self.step_count += 1

        if self.step_count >= self.max_steps:
            return "Max steps reached", -2.0, True

        # 解析动作
        if action.startswith("edit:"):
            # 编辑代码
            new_code = action[5:].strip()
            self.current_code = new_code
            return f"Code updated:\n{new_code}", 0.5, False

        elif action.startswith("run_tests"):
            # 运行测试
            passed, total, output = self._run_tests()
            reward = (passed / total) * 5.0 - 1.0
            done = passed == total
            return (
                f"Tests: {passed}/{total} passed\n{output}",
                reward,
                done
            )

        elif action.startswith("read:"):
            # 读取文件
            return f"File content:\n{self.current_code}", 0.0, False

        else:
            return "Unknown action", -1.0, False

    def _run_tests(self) -> tuple[int, int, str]:
        """运行测试用例"""
        passed = 0
        total = len(self.current_task.get("tests", []))
        output = []

        for i, test in enumerate(self.current_task.get("tests", [])):
            try:
                exec_globals = {}
                # exec将Agent修改后的代码与测试代码拼接执行，exec_globals为独立命名空间防止变量污染
                exec(self.current_code + "\n" + test["code"], exec_globals)
                passed += 1
                output.append(f"  ✅ Test {i+1}: PASSED")
            except Exception as e:
                output.append(f"  ❌ Test {i+1}: FAILED - {e}")

        return passed, total, "\n".join(output)

    def get_available_tools(self) -> list[dict]:
        return [
            {"name": "edit", "description": "Edit the code. Usage: edit:<new_code>"},
            {"name": "run_tests", "description": "Run the test suite"},
            {"name": "read", "description": "Read the current code"},
        ]

class RLVRTraining:
    """
    RLVR：使用可验证奖励训练推理模型

    核心思想：
    - 数学/编程任务有确定性验证器（单元测试、数学求解器）
    - 不需要人类标注偏好，用验证器提供二进制奖励
    - 可规模化：更多计算→更好推理能力
    """

    def __init__(self):
        self.verifiers = {
            "math": self.math_verifier,
            "code": self.code_verifier,
            "format": self.format_verifier
        }

    def math_verifier(self, problem: str, answer: str) -> float:
        """数学验证器：检查答案是否正确"""
        # 使用 SymPy 或数值方法验证
        expected = self.solve_math(problem)
        return 1.0 if self.numerical_equal(answer, expected) else 0.0

    def code_verifier(self, code: str, test_cases: list) -> float:
        """代码验证器：运行单元测试"""
        passed = 0
        for test in test_cases:
            try:
                exec(code)
                if eval(test["assertion"]):
                    passed += 1
            except:
                pass
        return passed / len(test_cases)

    def train_with_rlvr(self, model, problems, verifier):
        """使用RLVR训练模型"""
        for problem in problems:
            # 生成多个候选答案
            candidates = model.generate_multiple(problem, n=8)

            # 验证每个答案
            rewards = [verifier(problem, c) for c in candidates]

            # GRPO更新
            self.grpo_update(model, candidates, rewards)

class DAPOTrainer:
    """
    DAPO：解耦裁剪策略优化

    相比GRPO的改进：
    1. 解耦裁剪：clip_higher提升低概率token，鼓励探索
    2. 动态采样：针对中等难度样本，保持多样学习信号
    3. 更稳定的训练过程
    """

    def __init__(self, epsilon_low=0.2, epsilon_high=0.28, delta=1.5):
        self.epsilon_low = epsilon_low    # 下界裁剪
        self.epsilon_high = epsilon_high  # 上界裁剪（更大，鼓励探索）
        self.delta = delta                # 动态采样阈值

    def compute_advantage(self, rewards: list) -> list:
        """计算组内相对优势"""
        mean_r = sum(rewards) / len(rewards)
        std_r = (sum((r - mean_r)**2 for r in rewards) / len(rewards)) ** 0.5
        std_r = max(std_r, 1e-8)  # 避免除零

        advantages = [(r - mean_r) / std_r for r in rewards]
        return advantages

    def dapo_clip(self, ratio: float, advantage: float) -> float:
        """
        解耦裁剪：
        - 正优势时：允许更大范围探索（clip_higher）
        - 负优势时：更严格裁剪
        """
        if advantage > 0:
            # 正优势：使用更宽松的上界
            return min(ratio * advantage,
                      self.clip(ratio, 1 - self.epsilon_low, 1 + self.epsilon_high) * advantage)
        else:
            # 负优势：标准裁剪
            return self.clip(ratio, 1 - self.epsilon_low, 1 + self.epsilon_low) * advantage

    def dynamic_sampling(self, problems: list, model) -> list:
        """
        动态采样：针对中等难度样本

        目标：保持学习信号的多样性
        - 太简单：奖励总是1，无学习信号
        - 太难：奖励总是0，无学习信号
        - 中等：最有价值
        """
        sampled = []
        for problem in problems:
            # 快速评估难度
            candidates = model.generate_multiple(problem, n=4)
            rewards = self.verify(candidates)
            avg_reward = sum(rewards) / len(rewards)

            # 选择中等难度的样本（奖励在0.3-0.7之间）
            if 0.3 <= avg_reward <= 0.7:
                sampled.append(problem)

        return sampled

class ProRLTrainer:
    """
    ProRL：延长强化学习训练

    关键发现：
    - 传统RL训练在几百步后就停止
    - ProRL显示持续训练数千步仍有提升
    - 对OOD（分布外）任务提升尤其明显
    """

    def __init__(self, total_steps=10000):
        self.total_steps = total_steps
        self.curriculum = self.build_curriculum()

    def build_curriculum(self):
        """构建课程学习进度"""
        return {
            "phase1": {"steps": (0, 2000), "difficulty": "easy"},
            "phase2": {"steps": (2000, 5000), "difficulty": "medium"},
            "phase3": {"steps": (5000, 8000), "difficulty": "hard"},
            "phase4": {"steps": (8000, 10000), "difficulty": "expert"}
        }

    def train(self, model, environment):
        """延长训练循环"""
        for step in range(self.total_steps):
            # 根据训练阶段选择难度
            difficulty = self.get_current_difficulty(step)
            problems = environment.sample(difficulty=difficulty)

            # DAPO训练步骤
            self.dapo_step(model, problems)

            # 定期评估
            if step % 500 == 0:
                metrics = self.evaluate(model)
                print(f"Step {step}: Pass@1={metrics['pass@1']:.2%}")

class RULERTrainer:
    """
    RULER：使用LLM作为评判者生成相对奖励

    优势：
    - 不需要手动编写奖励函数
    - 不需要标注数据
    - LLM判断"哪个更好"比"给出0-10分"更可靠
    """

    def __init__(self, judge_model="gpt-4o"):
        self.judge = judge_model

    def llm_judge(self, trajectories: list) -> list:
        """
        使用LLM评判多条轨迹

        Args:
            trajectories: 同一任务的N条不同执行轨迹

        Returns:
            每条轨迹的相对排名分数
        """
        prompt = f"""
        你是一个Agent行为评判专家。请评估以下{len(trajectories)}条执行轨迹，
        判断哪条最能成功完成任务。

        任务：{trajectories[0]['task']}

        请对每条轨迹打分（0-1），考虑：
        1. 是否正确使用了工具
        2. 推理过程是否合理
        3. 是否成功完成任务
        4. 效率如何
        """

        # 调用LLM评判
        scores = self.call_llm_judge(prompt, trajectories)
        return scores

    def train_with_ruler(self, model, tasks):
        """使用RULER训练Agent"""
        for task in tasks:
            # 生成多条轨迹
            trajectories = []
            for _ in range(8):  # G=8
                traj = model.run_agent(task)
                trajectories.append(traj)

            # LLM评判
            scores = self.llm_judge(trajectories)

            # GRPO更新（使用LLM分数作为奖励）
            self.grpo_update(model, trajectories, scores)

class ARTTrainer:
    """
    ART：Agent强化训练器

    特点：
    - 支持任意MCP Server作为训练环境
    - 自动生成训练场景
    - 内置轨迹评估
    """

    def __init__(self, mcp_server_url: str):
        self.mcp_client = MCPClient(mcp_server_url)
        self.tools = self.discover_tools()

    def discover_tools(self) -> list:
        """自动发现MCP Server提供的工具"""
        return self.mcp_client.list_tools()

    def generate_training_scenarios(self, tool_name: str, count: int) -> list:
        """为特定工具生成训练场景"""
        tool_schema = self.get_tool_schema(tool_name)

        scenarios = []
        for _ in range(count):
            scenario = {
                "task": self.generate_task_description(tool_schema),
                "expected_tool": tool_name,
                "expected_params": self.generate_valid_params(tool_schema),
                "success_criteria": self.define_success_criteria(tool_schema)
            }
            scenarios.append(scenario)

        return scenarios

    def train_agent_for_mcp(self, model, mcp_server_url: str, steps: int = 1000):
        """
        训练Agent掌握特定MCP Server

        Args:
            model: 要训练的语言模型
            mcp_server_url: MCP Server地址
            steps: 训练步数
        """
        self.mcp_client = MCPClient(mcp_server_url)
        tools = self.discover_tools()

        for step in range(steps):
            # 为每个工具生成训练场景
            for tool in tools:
                scenarios = self.generate_training_scenarios(tool["name"], 4)

                # Agent尝试执行
                trajectories = []
                for scenario in scenarios:
                    traj = model.execute(scenario["task"])
                    trajectories.append(traj)

                # 评估并更新
                rewards = [self.evaluate_trajectory(t, s)
                          for t, s in zip(trajectories, scenarios)]
                self.grpo_update(model, trajectories, rewards)

# 使用Unsloth快速训练推理模型
from unsloth import FastLanguageModel
from trl import GRPOConfig, GRPOTrainer

# 加载模型
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="Qwen/Qwen2.5-3B",
    max_seq_length=2048,
    load_in_4bit=True,  # 4bit量化节省显存
)

# 配置GRPO
grpo_config = GRPOConfig(
    output_dir="./outputs",
    learning_rate=5e-6,
    num_generations=8,        # 每个prompt生成8个候选
    max_steps=300,            # 训练步数
    epsilon=0.2,              # PPO裁剪参数
    epsilon_high=0.28,        # DAPO解耦裁剪
    loss_type="grpo",         # 可选: grpo, dr_grpo, dapo
    use_vllm=True,            # 使用vLLM加速推理
)

# 定义奖励函数
def reward_function(completions, **kwargs):
    """
    奖励函数：评估模型输出
    """
    rewards = []
    for completion in completions:
        # 检查格式
        if "<reasoning>" in completion and "</reasoning>" in completion:
            reward = 1.0
        else:
            reward = 0.0

        # 检查答案正确性（如果有标准答案）
        if "answer" in kwargs:
            if extract_answer(completion) == kwargs["answer"]:
                reward += 2.0

        rewards.append(reward)
    return rewards

# 训练
trainer = GRPOTrainer(
    model=model,
    args=grpo_config,
    train_dataset=dataset,
    reward_funcs=reward_function,
)

trainer.train()

# 保存模型
model.save_pretrained_gguf("my_reasoning_model", tokenizer)