AI brokers are reshaping software program growth, from writing code to finishing up complicated directions. But LLM-based brokers are vulnerable to errors and infrequently carry out poorly on difficult, multi-step duties. Reinforcement studying (RL) is an strategy the place AI methods be taught to make optimum choices by receiving rewards or penalties for his or her actions, enhancing via trial and error. RL can assist brokers enhance, nevertheless it sometimes requires builders to extensively rewrite their code. This discourages adoption, although the info these brokers generate might considerably increase efficiency via RL coaching.
To handle this, a analysis crew from Microsoft Analysis Asia – Shanghai has launched Agent Lightning. This open-source (opens in new tab) framework makes AI brokers trainable via RL by separating how brokers execute duties from mannequin coaching, permitting builders so as to add RL capabilities with nearly no code modification.
Capturing agent habits for coaching
Agent Lightning converts an agent’s expertise right into a format that RL can use by treating the agent’s execution as a sequence of states and actions, the place every state captures the agent’s standing and every LLM name is an motion that strikes the agent to a brand new state.
This strategy works for any workflow, irrespective of how complicated. Whether or not it entails a number of collaborating brokers or dynamic device use, Agent Lightning breaks it down right into a sequence of transitions. Every transition captures the LLM’s enter, output, and reward (Determine 1). This standardized format means the info can be utilized for coaching with none further steps.
Hierarchical reinforcement studying
Conventional RL coaching for brokers that make a number of LLM requests entails stitching collectively all content material into one lengthy sequence after which figuring out which elements ought to be realized and which ignored throughout coaching. This strategy is troublesome to implement and might create excessively lengthy sequences that degrade mannequin efficiency.
As an alternative, Agent Lightning’s LightningRL algorithm takes a hierarchical strategy. After a job completes, a credit score project module determines how a lot every LLM request contributed to the result and assigns it a corresponding reward. These impartial steps, now paired with their very own reward scores, can be utilized with any current single-step RL algorithm, resembling Proximal Coverage Optimization (PPO) or Group Relative Coverage Optimization (GRPO) (Determine 2).
This design provides a number of advantages. It stays absolutely suitable with broadly used single-step RL algorithms, permitting current coaching strategies to be utilized with out modification. Organizing knowledge as a sequence of impartial transitions lets builders flexibly assemble the LLM enter as wanted, supporting complicated behaviors like brokers that use a number of instruments or work with different brokers. Moreover, by holding sequences brief, the strategy scales cleanly and retains coaching environment friendly.
Agent Lightning as middleware
Agent Lightning serves as middleware between RL algorithms and agent environments, offering modular parts that allow scalable RL via standardized protocols and well-defined interfaces.
An agent runner manages the brokers as they full duties. It distributes work and collects and shops the outcomes and progress knowledge. It operates individually from the LLMs, enabling them to run on completely different assets and scale to help a number of brokers operating concurrently.
An algorithm trains the fashions and hosts the LLMs used for inference and coaching. It orchestrates the general RL cycle, managing which duties are assigned, how brokers full them, and the way fashions are up to date based mostly on what the brokers be taught. It sometimes runs on GPU assets and communicates with the agent runner via shared protocols.
The LightningStore (opens in new tab) serves because the central repository for all knowledge exchanges inside the system. It gives standardized interfaces and a shared format, guaranteeing that the completely different parts can work collectively and enabling the algorithm and agent runner to speak successfully.
All RL cycles observe two steps: (1) Agent Lightning collects agent execution knowledge (known as “spans”) and retailer them within the knowledge retailer; (2) it then retrieves the required knowledge and sends it to the algorithm for coaching. Via this design, the algorithm can delegate duties asynchronously to the agent runner, which completes them and studies the outcomes again (Determine 4).
One key benefit of this strategy is its algorithmic flexibility. The system makes it straightforward for builders to customise how brokers be taught, whether or not they’re defining completely different rewards, capturing intermediate knowledge, or experimenting with completely different coaching approaches.
One other benefit is useful resource effectivity. Agentic RL methods are complicated, integrating agentic methods, LLM inference engines, and coaching frameworks. By separating these parts, Agent Lightning makes this complexity manageable and permits every half to be optimized independently
A decoupled design permits every part to make use of the {hardware} that fits it finest. The agent runner can use CPUs whereas mannequin coaching makes use of GPUs. Every part can even scale independently, enhancing effectivity and making the system simpler to keep up. In apply, builders can preserve their current agent frameworks and swap mannequin calls to the Agent Lightning API with out altering their agent code (Determine 5).
Analysis throughout three real-world situations
Agent Lightning was examined on three distinct duties, reaching constant efficiency enhancements throughout all situations (Determine 6):
Textual content-to-SQL (LangChain): In a system with three brokers dealing with SQL technology, checking, and rewriting, Agent Lightning concurrently optimized two of them, considerably enhancing the accuracy of producing executable SQL from pure language queries.
Retrieval-augmented technology (OpenAI Brokers SDK implementation): On the multi-hop question-answering dataset MuSiQue, which requires querying a big Wikipedia database, Agent Lightning helped the agent generate more practical search queries and purpose higher from retrieved content material.
Mathematical QA and power use (AutoGen implementation): For complicated math issues, Agent Lightning skilled LLMs to extra precisely decide when and easy methods to name the device and combine the outcomes into its reasoning, rising accuracy.
Enabling steady agent enchancment
By simplifying RL integration, Agent Lightning could make it simpler for builders to construct, iterate, and deploy high-performance brokers. We plan to broaden Agent Lightning’s capabilities to incorporate computerized immediate optimization and extra RL algorithms.
The framework is designed to function an open platform the place any AI agent can enhance via real-world apply. By bridging current agentic methods with reinforcement studying, Agent Lightning goals to assist create AI methods that be taught from expertise and enhance over time.