🐪如何用 PPO 微调 LLM 作为 Agent

智能体
约 17 分钟阅读

阅读量:15次
Agent

使用 PPO（Proximal Policy Optimization）微调大语言模型（LLM）作为智能体（Agent），适用于任务型对话、工具调用、自主规划等场景。我们将结合 Hugging Face Transformers + TRL（Transformer Reinforcement Learning）库 + 自定义环境，以一个简化但完整的例子说明整个流程。

🎯 目标场景示例

任务：让 LLM Agent 在一个模拟环境中回答用户问题，通过调用工具（如计算器、搜索API）获取信息，并最终输出正确答案。
奖励信号：根据最终答案是否正确、是否高效使用工具、是否遵守格式等给出标量奖励。
方法：使用 PPO 对 LLM 策略进行微调，使其学会“何时调用工具”“如何推理”“如何输出”。

🧰 技术栈

  • Python ≥ 3.10
  • PyTorch
  • Hugging Face Transformers
  • TRL（Transformer Reinforcement Learning）
  • accelerate（用于分布式训练）
  • 自定义环境（或使用 WebArena、BabyAGI 等简化版）

📦 第一步：安装依赖

pip install torch transformers datasets accelerate peft trl
注：TRL 是 Hugging Face 官方维护的 RL 微调 LLM 的库，原生支持 PPO。

🧪 第二步：定义任务与环境（简化版）

我们构建一个 “数学问答 + 工具调用” 环境：

  • 用户提问："What is 123 * 456?"

  • Agent 应输出：

    Thought: I need to calculate 123 * 456.
Action: calculator(123 * 456)
Observation: 56088
Answer: The result is 56088.

✅ 环境模拟器（伪代码）

def simulate_env(prompt: str, response: str) -> float:
    # 从 response 中解析是否调用了 calculator
    if "calculator(" in response:
        # 提取表达式并计算
        expr = extract_expression(response)  # e.g., "123 * 456"
        try:
            correct = eval(expr)
            # 检查最终答案是否包含正确结果
            if str(correct) in response:
                return 1.0  # 完全正确
            else:
                return 0.2  # 调用了工具但答错
        except:
            return -0.5  # 工具调用格式错误
    else:
        return -1.0  # 未调用工具，直接瞎猜
实战中可用更复杂的环境（如 WebShop、WebArena），但原理相同。

🤖 第三步：准备 LLM 策略模型

我们使用 LLaMA-2-7b-chat 或 Mistral-7B-Instruct（需有 Hugging Face token 权限），或开源模型如 Qwen-1.8B-Chat。

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "Qwen/Qwen2.5-7B-Chat"  # 或其他支持对话的模型
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)
注意：为节省显存，可使用 LoRA（PEFT） 进行参数高效微调。

🔧 第四步：使用 TRL 的 PPOTrainer

TRL 提供了 PPOTrainer，可直接对接 LLM 和奖励函数。

4.1 准备数据集（Prompt）

from datasets import Dataset

prompts = [
    "Calculate 123 * 456.",
    "What is the square root of 1444?",
    "Compute 999 + 1001."
]
dataset = Dataset.from_dict({"query": prompts})

4.2 配置 PPO

from trl import PPOConfig

config = PPOConfig(
    model_name=model_name,
    learning_rate=1e-5,
    batch_size=1,          # 小 batch 适合 LLM
    mini_batch_size=1,
    gradient_accumulation_steps=4,
    log_with="wandb",      # 可选：集成 Weights & Biases
)

4.3 初始化 PPO Trainer

from trl import PPOTrainer
from peft import LoraConfig, get_peft_model

# 使用 LoRA 减少显存
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

ppo_trainer = PPOTrainer(
    config=config,
    model=model,
    ref_model=None,  # 可选：使用固定参考模型
    tokenizer=tokenizer,
    dataset=dataset,
)

🎁 第五步：训练循环（核心）

generation_kwargs = {
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
    "max_new_tokens": 128,
}

for epoch in range(3):
    for batch in ppo_trainer.dataloader:
        query_tensors = batch["input_ids"]

        # 1. 生成响应
        response_tensors = ppo_trainer.generate(
            query_tensors, 
            return_prompt=False, 
            **generation_kwargs
        )
        batch["response"] = [tokenizer.decode(r) for r in response_tensors]

        # 2. 计算奖励
        rewards = []
        for prompt, response in zip(batch["query"], batch["response"]):
            reward = simulate_env(prompt, response)
            rewards.append(torch.tensor(reward, dtype=torch.float32))

        # 3. PPO 更新
        stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
        ppo_trainer.log_stats(stats, batch, rewards)

✅ 第六步：评估与部署

  • 保存 LoRA 适配器：

    model.save_pretrained("ppo-agent-lora")

  • 推理时加载：

    from peft import PeftModel
base_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-1.8B-Chat")
model = PeftModel.from_pretrained(base_model, "ppo-agent-lora")

⚠️ 关键挑战与技巧

挑战解决方案
奖励稀疏使用 LLM 自动生成子目标奖励（如“是否调用了工具？”）
训练不稳定降低学习率、使用 KL 散度约束（TRL 默认启用）
显存不足使用 LoRA + gradient checkpointing + bf16
动作空间大限制输出格式（如强制以 Thought: 开头）
评估困难构建自动化测试集 + 人工审核

📚 扩展阅读 & 参考项目

  1. TRL 官方 PPO 示例：
    https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py
  2. WebArena（真实环境 + LLM Agent + RL）：
    https://webarena.dev/
  3. Voyager（Minecraft 中的 LLM Agent + PPO）：
    https://voyager.minedojo.org/

  4. 论文：

    • “Large Language Models as Optimizers” (LLM + RL)
    • “Reflexion: Language Agents with Verbal Reinforcement Learning”

💡 总结

PPO 微调 LLM 作为 Agent 的核心思想：
将 LLM 视为策略网络 π(a|s)，通过与环境交互获得奖励，用 PPO 更新其参数，使其学会在复杂任务中做出正确决策序列。

虽然目前仍面临样本效率低、奖励设计难等问题，但这是构建 自主、可学习、可进化 Agent 的关键路径。

添加新评论