Agent Lightning: Train ANY AI Agents with Reinforcement Learning

2025-08-05

AI Summary

This work addresses key challenges in reinforcement learning (RL) for AI agents: strong inter-agent coupling, poor reusability of existing systems, and limited support for multi-agent coordination and dynamic workflows. To this end, we propose LightningRL—a lightweight, modular RL framework. Methodologically, it introduces a unified trajectory interface and hierarchical MDP modeling to fully decouple agent execution from RL training; pioneers a credit-allocation-driven trajectory decomposition mechanism that automatically converts arbitrary agent outputs into trainable RL samples; and adopts a decentralized training architecture with enhanced observability, enabling zero-code integration with mainstream agent frameworks such as LangChain and OpenAI Agents SDK. Empirically, LightningRL demonstrates consistent performance gains across text-to-SQL generation, retrieval-augmented generation (RAG), and mathematical tool-calling tasks—extending the applicability boundary of RL to complex, real-world agent systems.

Technology Category

• Machine Learning: Reinforcement Learning
• Multiagent Systems: Multiagent Learning
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

Abstract

We present Agent Lightning, a flexible and extensible framework that enables Reinforcement Learning (RL)-based training of Large Language Models (LLMs) for any AI agent. Unlike existing methods that tightly couple RL training with agent or rely on sequence concatenation with masking, Agent Lightning achieves complete decoupling between agent execution and training, allowing seamless integration with existing agents developed via diverse ways (e.g., using frameworks like LangChain, OpenAI Agents SDK, AutoGen, and building from scratch) with almost ZERO code modifications. By formulating agent execution as Markov decision process, we define an unified data interface and propose a hierarchical RL algorithm, LightningRL, which contains a credit assignment module, allowing us to decompose trajectories generated by ANY agents into training transition. This enables RL to handle complex interaction logic, such as multi-agent scenarios and dynamic workflows. For the system design, we introduce a Training-Agent Disaggregation architecture, and brings agent observability frameworks into agent runtime, providing a standardized agent finetuning interface. Experiments across text-to-SQL, retrieval-augmented generation, and math tool-use tasks demonstrate stable, continuous improvements, showcasing the framework's potential for real-world agent training and deployment.

Problem

• Decoupling agent execution and RL training for flexibility
• Unified RL training for diverse AI agent frameworks
• Handling complex agent interactions via hierarchical RL

Innovation

• Decouples agent execution and RL training
• Uses hierarchical RL with credit assignment
• Introduces Training-Agent Disaggregation architecture