利用在线强化学习对大语言模型代理进行微调

🔗 用于网页数据提取的代理   •   🐦 推特

LlamaGym

“代理”一词最早源于强化学习，其通过与环境互动并接收奖励信号来学习。然而，如今基于大语言模型的代理并不能通过强化学习实现在线学习（即实时持续学习）。

OpenAI 创建了 Gym 以标准化并简化强化学习环境，但如果你尝试将基于大语言模型的代理放入 Gym 环境中进行训练，就会发现，要处理大语言模型的对话上下文、情节批次、奖励分配、PPO 设置等，仍然需要编写相当多的代码。

LlamaGym 旨在通过强化学习简化大语言模型代理的微调过程。目前，它仅提供一个名为 Agent 的抽象类，负责处理上述所有问题，让你能够快速迭代并试验各种代理提示和超参数，适用于任意 Gym 环境。

使用方法

将大语言模型代理微调至 Gym 风格的环境中，并借助强化学习，从未如此简单！一旦你安装了 LlamaGym……

pip install llamagym

首先，在 Agent 类中实现 3 个抽象方法：

from llamagym import Agent

class BlackjackAgent(Agent):
    def get_system_prompt(self) -> str:
        return "你是一位经验丰富的二十一点玩家。"

    def format_observation(self, observation) -> str:
        return f"你的当前总金额为 {observation[0]}"

    def extract_action(self, response: str):
        return 0 if "stay" in response else 1

然后，定义你的基础大语言模型（就像在任何微调任务中一样），并实例化你的代理：

model = AutoModelForCausalLMWithValueHead.from_pretrained("Llama-2-7b").to(device)
tokenizer = AutoTokenizer.from_pretrained("Llama-2-7b")
agent = BlackjackAgent(model, tokenizer, device)

最后，按照常规方式编写你的强化学习循环，并只需调用代理执行动作、获取奖励并完成训练：

env = gym.make("Blackjack-v1")

for episode in trange(5000):
    observation, info = env.reset()
    done = False

    while not done:
        action = agent.act(observation) # 根据观察结果采取行动
        observation, reward, terminated, truncated, info = env.step(action)
        agent.assign_reward(reward) # 向代理提供奖励
        done = terminated or truncated

    train_stats = agent.terminate_episode() # 如果整个批次已完成，则进行训练

一些注意事项：

• 上述代码片段虽经过了略微简化，但完整的运行示例可在此处找到：examples/blackjack.py
• 要让在线强化学习收敛，往往难度极高，因此你可能需要反复调整超参数，才能看到显著的改进。
• 你的模型在运行强化学习之前，或许还可以先通过监督微调阶段，对采样轨迹进行优化（我们未来可能会加入这一功能）。
• 我们的实现注重简洁性，因此计算效率不如例如 Lamorel 这样的工具，但更易于入门和实验。
• LlamaGym 是一个周末项目，目前仍处于开发中，但我们非常欢迎各位的贡献！

相关工作

• 利用在线强化学习为大型语言模型赋能
  • Lamorel：面向强化学习的语言模型
• 真知源自实践：通过强化学习使大语言模型与现实环境相契合

引用格式

bibtex
@misc{pandey2024llamagym,
  title        = {LlamaGym：利用在线强化学习对大语言模型代理进行微调},
  author       = {Rohan Pandey},
  year         = {2024},
  howpublished = {GitHub},
  url          = {https://github.com/KhoomeiK/LlamaGym}
}

LlamaGym 快速上手指南

LlamaGym 是一个旨在简化使用在线强化学习（Online RL）微调大语言模型（LLM）智能体的开源工具。它通过一个抽象的 Agent 类，自动处理对话上下文、回合批次管理、奖励分配及 PPO 设置等复杂逻辑，让开发者能专注于智能体的提示词设计与超参数调整。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux 或 macOS（Windows 用户建议使用 WSL2）。
• Python 版本：3.8 及以上。
• 硬件要求：推荐配备 NVIDIA GPU 以加速训练过程（CPU 亦可运行但速度较慢）。
• 前置依赖：
  • gymnasium (或 gym)：用于提供强化学习环境。
  • transformers & trl：用于加载模型和执行强化学习算法（LlamaGym 安装时会自动处理大部分依赖）。

国内加速建议： 如果遇到 PyPI 源下载缓慢的问题，建议使用清华或阿里镜像源进行安装：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple llamagym

模型下载若受阻，可配置 HF_ENDPOINT 环境变量使用 Hugging Face 中国镜像：

export HF_ENDPOINT=https://hf-mirror.com

安装步骤

通过 pip 直接安装 LlamaGym：

pip install llamagym

如需安装包含额外开发依赖的版本：

pip install llamagym[dev]

基本使用

使用 LlamaGym 的核心流程分为三步：定义智能体类、加载模型、运行强化学习循环。以下以经典的 Blackjack（21 点）游戏为例：

1. 定义智能体类

继承 Agent 类并实现三个核心抽象方法，用于定义系统提示、观测值格式化以及动作提取：

from llamagym import Agent

class BlackjackAgent(Agent):
    def get_system_prompt(self) -> str:
        return "You are an expert blackjack player."

    def format_observation(self, observation) -> str:
        # 将环境观测值转换为 LLM 可理解的文本
        return f"Your current total is {observation[0]}"

    def extract_action(self, response: str):
        # 从 LLM 的回复中解析出具体的动作 ID
        return 0 if "stay" in response else 1

2. 加载模型并实例化智能体

加载预训练的基础模型（如 Llama-2），并将其注入到智能体中：

from transformers import AutoModelForCausalLMWithValueHead, AutoTokenizer
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"

# 加载带有价值头的模型 (PPO 所需)
model = AutoModelForCausalLMWithValueHead.from_pretrained("Llama-2-7b").to(device)
tokenizer = AutoTokenizer.from_pretrained("Llama-2-7b")

# 实例化智能体
agent = BlackjackAgent(model, tokenizer, device)

3. 运行强化学习循环

创建 Gym 环境，并在标准的 RL 循环中调用智能体的方法来进行交互、奖励分配和训练更新：

import gymnasium as gym
from tqdm import trange

env = gym.make("Blackjack-v1")

for episode in trange(5000):
    observation, info = env.reset()
    done = False

    while not done:
        # 智能体根据观测值采取行动
        action = agent.act(observation)
        
        # 环境执行动作并返回结果
        observation, reward, terminated, truncated, info = env.step(action)
        
        # 将奖励反馈给智能体
        agent.assign_reward(reward)
        
        done = terminated or truncated

    # 回合结束，若批次已满则触发训练
    train_stats = agent.terminate_episode()

注意：

• 完整可运行的示例代码请参考 examples/blackjack.py。
• 在线强化学习收敛难度较大，初次尝试时请务必调整超参数以获得最佳效果。
• 建议在正式运行 RL 前，先对模型进行基于采样轨迹的监督微调（SFT），这通常能提升稳定性。