LlamaGym

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

关于 LlamaGym

“Agent”一词起源于强化学习领域，Agent 通过与环境交互并接收奖励信号进行学习。然而，目前的基于 LLM 的 Agent 并不支持在线（即实时持续）的强化学习。

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

LlamaGym 旨在简化使用强化学习微调 LLM Agent 的过程。目前，它提供了一个单一的 Agent 抽象类，涵盖了上述所有问题，让你能够快速在任何 Gym 环境中迭代和实验 Agent 的提示词（prompting）及超参数。

使用方法

使用强化学习微调基于 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 进行行动、奖励和终止：

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。
在线强化学习的收敛过程极其困难，你需要反复调整超参数才能看到效果。
- 在运行强化学习之前，模型如果先在采样轨迹上进行监督微调（SFT），效果可能会更好（我们未来可能会添加此功能）。
我们的实现追求简洁，因此在计算效率上不如 Lamorel 等项目，但它更容易上手尝试。
LlamaGym 是一个周末项目，目前仍在进行中，我们非常欢迎贡献！

引用

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

KhoomeiK/LlamaGym

LlamaGym

关于 LlamaGym

使用方法

相关工作

引用