LlamaGym

使用在线强化学习（Online Reinforcement Learning）微调 LLM Agent。

🔗 用于 Web 数据提取的 Agent   •   🐦 Twitter

简介

“Agent”一词起源于强化学习（Reinforcement Learning），其核心在于通过与环境交互并接收奖励信号来学习。然而，现今基于 LLM 的 Agent 并不支持通过强化学习进行在线（即实时持续）学习。

OpenAI 创建了 Gym 以标准化和简化 RL 环境，但如果你尝试将基于 LLM 的 Agent 放入 Gym 环境进行训练，你会发现处理 LLM 对话上下文、episode 批次、奖励分配、PPO 设置等问题依然非常繁琐。

LlamaGym 旨在简化使用 RL 微调 LLM Agent 的过程。目前，它通过一个单一的 Agent 抽象类来处理上述所有复杂问题，让你能够快速在任何 Gym 环境中迭代和试验 Agent 的提示词（Prompting）及超参数。

使用方法

使用 RL 微调 LLM Agent 以在 Gym 风格的环境中运行从未如此简单！安装 LlamaGym 后：

pip install llamagym

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

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:
        return f"Your current total is {observation[0]}"

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

然后，定义你的基础 LLM（就像进行任何微调任务一样）并实例化你的 Agent：

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

最后，像往常一样编写 RL 循环，只需调用你的 Agent 执行动作、分配奖励并终止：

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) # 向 Agent 提供奖励
        done = terminated or truncated

    train_stats = agent.terminate_episode() # 如果批次已满则进行训练

一些提示：

• 上述代码片段经过了简化，完整的运行示例可参考 examples/blackjack.py。
• 在线 RL 收敛难度众所周知，因此你需要调整超参数以观察性能提升。
  • 在运行 RL 之前，模型通过对采样轨迹进行监督微调（SFT）可能会有所助益（我们未来可能会添加此功能）。
• 我们的实现追求简洁，因此在计算效率上不如 Lamorel 等项目，但它更容易上手尝试。
• LlamaGym 是一个周末项目，目前仍在开发中（WIP），欢迎提交贡献！

相关工作

• Grounding Large Language Models with Online Reinforcement Learning
  • Lamorel: Language Models for Reinforcement Learning
• True Knowledge Comes from Practice: Aligning LLMs with Embodied Environments via Reinforcement Learning

引用

@misc{pandey2024llamagym,
  title        = {LlamaGym: Fine-tune LLM agents with Online Reinforcement Learning},
  author       = {Rohan Pandey},
  year         = {2024},
  howpublished = {GitHub},
  url          = {https://github.com/KhoomeiK/LlamaGym}
}