LlamaGym：用强化学习优化LLM Agent

通过在线强化学习微调 LLM Agent

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

“Agent”这一概念源自强化学习。在强化学习中，Agent 通过与环境交互并接收奖励信号来学习。然而，如今基于 LLM 的 Agent 并不会通过强化学习进行在线学习（即持续、实时地学习）。

OpenAI 创建了 Gym，用于标准化并简化强化学习环境。但如果你尝试将基于 LLM 的 Agent 直接放入 Gym 环境中训练，就会发现仍然需要编写大量代码，才能处理 LLM 对话上下文、episode 批次、奖励分配、PPO 配置等问题。

LlamaGym 旨在简化使用强化学习微调 LLM Agent 的过程。目前，它只提供了一个 Agent 抽象类，用来处理上述所有问题，让你能够在任意 Gym 环境中快速迭代和实验 Agent 的 prompt 与超参数。

现在，通过强化学习微调基于 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)

最后，像往常一样编写强化学习循环，只需调用 Agent，让它执行动作、接收奖励并终止 episode：

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

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

    while not done:
        action = agent.act(observation) # act based on observation
        observation, reward, terminated, truncated, info = env.step(action)
        agent.assign_reward(reward) # provide reward to agent
        done = terminated or truncated

    train_stats = agent.terminate_episode() # trains if batch is full

上面的代码片段做了少量简化，但 examples/blackjack.py 中提供了一个可以完整运行的示例。

众所周知，让在线强化学习收敛非常困难，因此你需要反复调整超参数，才能看到效果提升。在运行强化学习之前，模型也可能受益于一个针对采样轨迹的监督微调阶段（我们未来可能会加入这一功能）。

在运行强化学习之前，模型也可能受益于一个针对采样轨迹的监督微调阶段（我们未来可能会加入这一功能）。

我们的实现更看重简洁性，因此在计算效率上不如 Lamorel 等方案，但更容易上手实验。

LlamaGym 是一个周末项目，目前仍在开发中，但我们非常欢迎大家贡献代码！

用在线强化学习为大语言模型提供现实基础（Grounding Large Language Models with Online Reinforcement Learning）

Lamorel：用于强化学习的语言模型（Language Models for Reinforcement Learning）

Lamorel：用于强化学习的语言模型（Language Models for Reinforcement Learning）

真知源于实践：通过强化学习让 LLM 与具身环境对齐（True Knowledge Comes from Practice: Aligning LLMs with Embodied Environments via Reinforcement Learning）

bibtex
@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}
}