base on Fine-tune LLM agents with online reinforcement learning <p align="center"> <img src="https://raw.githubusercontent.com/khoomeik/LlamaGym/main/llamagym.png" height="250" alt="Llama Gym" /> </p> <p align="center"> <em>Fine-tune LLM agents with online reinforcement learning</em> </p> <p align="center"> <a href="https://pypi.org/project/llamagym/" target="_blank"> <img alt="Python" src="https://img.shields.io/badge/python-3670A0?style=for-the-badge&logo=python&logoColor=ffdd54" /> <img alt="Version" src="https://img.shields.io/pypi/v/llamagym?style=for-the-badge&color=3670A0"> </a> </p> <p align="center"> <a href="https://reworkd.ai/">🔗 Agents for Web Data Extraction</a> <span>&nbsp;&nbsp;•&nbsp;&nbsp;</span> <a href="https://x.com/khoomeik/status/1766805213644800011">🐦 Twitter</a> # LlamaGym "Agents" originated in reinforcement learning, where they learn by interacting with an environment and receiving a reward signal. However, LLM-based agents today do not learn online (i.e. continuously in real time) via reinforcement. OpenAI created [Gym](https://github.com/Farama-Foundation/Gymnasium) to standardize and simplify RL environments, but if you try dropping an LLM-based agent into a Gym environment for training, you'd find it's still quite a bit of code to handle LLM conversation context, episode batches, reward assignment, PPO setup, and more. LlamaGym seeks to simplify fine-tuning LLM agents with RL. Right now, it's a single `Agent` abstract class that handles all the issues mentioned above, letting you quickly iterate and experiment with agent prompting & hyperparameters across any Gym environment. ## Usage Fine-tuning an LLM-based agent to play in a Gym-style environment with RL has never been easier! Once you install LlamaGym... ``` pip install llamagym ``` First, implement 3 abstract methods on the Agent class: ```python 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 ``` Then, define your base LLM (as you would for any fine-tuning job) and instantiate your agent: ```python model = AutoModelForCausalLMWithValueHead.from_pretrained("Llama-2-7b").to(device) tokenizer = AutoTokenizer.from_pretrained("Llama-2-7b") agent = BlackjackAgent(model, tokenizer, device) ``` Finally, write your RL loop as usual and simply call your agent to act, reward, and terminate: ```python 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 ``` Some reminders: - above code snippets are mildly simplified above but a fully working example is available in [`examples/blackjack.py`](https://github.com/KhoomeiK/LlamaGym/blob/main/examples/blackjack.py) - getting online RL to converge is notoriously difficult so you'll have to mess with hyperparameters to see improvement - your model may also benefit from a supervised fine-tuning stage on sampled trajectories before running RL (we may add this feature in the future) - our implementation values simplicity so is not as compute efficient as e.g. [Lamorel](https://github.com/flowersteam/lamorel), but easier to start playing around with - LlamaGym is a weekend project and still a WIP, but we love contributions! ## Relevant Work - [Grounding Large Language Models with Online Reinforcement Learning](https://github.com/flowersteam/Grounding_LLMs_with_online_RL) - [Lamorel: Language Models for Reinforcement Learning](https://github.com/flowersteam/lamorel) - [True Knowledge Comes from Practice: Aligning LLMs with Embodied Environments via Reinforcement Learning](https://github.com/WeihaoTan/TWOSOME) ## Citation ``` 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} } ``` ", Assign "at most 3 tags" to the expected json: {"id":"8493","tags":[]} "only from the tags list I provide: [{"id":77,"name":"3d"},{"id":89,"name":"agent"},{"id":17,"name":"ai"},{"id":54,"name":"algorithm"},{"id":24,"name":"api"},{"id":44,"name":"authentication"},{"id":3,"name":"aws"},{"id":27,"name":"backend"},{"id":60,"name":"benchmark"},{"id":72,"name":"best-practices"},{"id":39,"name":"bitcoin"},{"id":37,"name":"blockchain"},{"id":1,"name":"blog"},{"id":45,"name":"bundler"},{"id":58,"name":"cache"},{"id":21,"name":"chat"},{"id":49,"name":"cicd"},{"id":4,"name":"cli"},{"id":64,"name":"cloud-native"},{"id":48,"name":"cms"},{"id":61,"name":"compiler"},{"id":68,"name":"containerization"},{"id":92,"name":"crm"},{"id":34,"name":"data"},{"id":47,"name":"database"},{"id":8,"name":"declarative-gui "},{"id":9,"name":"deploy-tool"},{"id":53,"name":"desktop-app"},{"id":6,"name":"dev-exp-lib"},{"id":59,"name":"dev-tool"},{"id":13,"name":"ecommerce"},{"id":26,"name":"editor"},{"id":66,"name":"emulator"},{"id":62,"name":"filesystem"},{"id":80,"name":"finance"},{"id":15,"name":"firmware"},{"id":73,"name":"for-fun"},{"id":2,"name":"framework"},{"id":11,"name":"frontend"},{"id":22,"name":"game"},{"id":81,"name":"game-engine "},{"id":23,"name":"graphql"},{"id":84,"name":"gui"},{"id":91,"name":"http"},{"id":5,"name":"http-client"},{"id":51,"name":"iac"},{"id":30,"name":"ide"},{"id":78,"name":"iot"},{"id":40,"name":"json"},{"id":83,"name":"julian"},{"id":38,"name":"k8s"},{"id":31,"name":"language"},{"id":10,"name":"learning-resource"},{"id":33,"name":"lib"},{"id":41,"name":"linter"},{"id":28,"name":"lms"},{"id":16,"name":"logging"},{"id":76,"name":"low-code"},{"id":90,"name":"message-queue"},{"id":42,"name":"mobile-app"},{"id":18,"name":"monitoring"},{"id":36,"name":"networking"},{"id":7,"name":"node-version"},{"id":55,"name":"nosql"},{"id":57,"name":"observability"},{"id":46,"name":"orm"},{"id":52,"name":"os"},{"id":14,"name":"parser"},{"id":74,"name":"react"},{"id":82,"name":"real-time"},{"id":56,"name":"robot"},{"id":65,"name":"runtime"},{"id":32,"name":"sdk"},{"id":71,"name":"search"},{"id":63,"name":"secrets"},{"id":25,"name":"security"},{"id":85,"name":"server"},{"id":86,"name":"serverless"},{"id":70,"name":"storage"},{"id":75,"name":"system-design"},{"id":79,"name":"terminal"},{"id":29,"name":"testing"},{"id":12,"name":"ui"},{"id":50,"name":"ux"},{"id":88,"name":"video"},{"id":20,"name":"web-app"},{"id":35,"name":"web-server"},{"id":43,"name":"webassembly"},{"id":69,"name":"workflow"},{"id":87,"name":"yaml"}]" returns me the "expected json"