TLDR: AgentFly is a novel, scalable, and extensible framework designed to train language model (LM) agents using reinforcement learning (RL). It addresses challenges in multi-turn interactions by applying token-level masking, unifies external interfaces as ‘tools’ (both stateful and non-stateful), and ensures high-throughput training through asynchronous execution and centralized resource management. The framework also offers a developer-friendly interface for defining tools and rewards. Experiments demonstrate its effectiveness in training LM agents across various complex tasks and environments, showing improved performance and reduced ‘hallucinations’ as agents learn.
Language model (LM) agents are becoming increasingly popular for their ability to perform tasks autonomously by interacting with various environments, tools, and APIs. Traditionally, these agents are built using prompt engineering or supervised fine-tuning. However, integrating reinforcement learning (RL) with LM agents, a field known as Agent-RL, has remained largely unexplored and lacks systematic study. This is where AgentFly comes in.
AgentFly is a novel, scalable, and extensible framework designed to empower LM agents with a wide range of reinforcement learning algorithms. It tackles several key challenges in Agent-RL, such as managing multi-turn interactions, handling expensive rollout phases involving concurrent tool interactions, and striking a balance between modularity and training efficiency.
How AgentFly Works
One of AgentFly’s core innovations is its approach to multi-turn interactions. In typical RL, models learn from their own outputs. AgentFly adapts this for multi-turn scenarios by using token-level masking, ensuring the model only learns from the tokens it generates, rather than being penalized for information introduced by the environment or other components.
The framework unifies all external interfaces an LM agent interacts with under the concept of “tools.” Whether it’s a simple function, an API, or a complex environment, AgentFly treats every interaction as a tool invocation. This simplifies integration and makes the system highly generalizable. Tools can be either “non-stateful,” like a calculator or a web searcher that don’t maintain any environment state, or “stateful,” such as an operating system or a web browser that require dedicated, isolated environment instances for their use.
To support high-throughput training, AgentFly implements asynchronous execution for tool calls and reward computations. This means multiple operations can happen simultaneously without blocking each other. It also features a centralized resource management system that efficiently coordinates environment instances, allowing for scalable parallel rollouts and reward calculations. For developers, AgentFly offers a user-friendly, decorator-based interface, enabling them to define tools and reward functions easily using standard Python functions.
Also Read:
- AlphaAlign: A New Approach to Safer Language Models Through Self-Awareness
- LLMs Learn to Think Smarter with Hierarchical Budget Policy Optimization
Real-World Applications and Performance
The effectiveness of AgentFly has been demonstrated through extensive experiments. The framework successfully integrates four widely used RL algorithms and comes with a suite of pre-built tools and environments. These include a Code Interpreter for solving math problems, Search and Retrieve tools for online information, and complex simulated environments like WebShop (an e-commerce site), ALFWorld (a household manipulation simulator), and ScienceWorld (a text-based science sandbox).
Experiments showed that both 3B and 7B language models could be effectively trained with AgentFly, with the larger 7B model generally achieving better performance. For more complex tasks requiring multiple turns and diverse tool usage, such as ALFWorld and ScienceWorld, training progressed steadily, indicating the framework’s capability to handle intricate scenarios. The research also observed that as agents learned, their “hallucination” rate (generating invalid actions or formats) decreased, and they became more efficient in their use of tools, adapting their strategies over time. For more technical details, you can refer to the full research paper.
In conclusion, AgentFly provides a robust, flexible, and scalable solution for training language model agents using reinforcement learning. By addressing key challenges in multi-turn interactions, tool integration, and efficient resource management, it paves the way for more capable and autonomous LM agents in various applications.
