PaperlayerView exact revision on arXiv

Source-linked AI summary

Agent Lightning: Train ANY AI Agents with Reinforcement Learning

Xufang Luo, Yuge Zhang, Zhiyuan He, Zilong Wang, Siyun Zhao, Dongsheng Li, Luna K. Qiu, Yuqing Yang

arXiv:2508.03680v1cs.AIcs.LG

TL;DR

Agent Lightning addresses the challenge of applying RL to complex, diverse AI agents by decoupling agent execution from training and representing executions through a unified MDP-based interface. Its hierarchical RL and disaggregated system support existing agents with almost ZERO code modifications, while experiments report continuous, stable improvements across multiple tasks.

  • Problem

    Complex agents involve multiple LLM calls, tools, and dynamic workflows, while existing agent frameworks generally lack automatic optimization for demanding real-world settings.

  • Method

    Agent Lightning formulates agent execution as an MDP, converts executions into unified transitions, applies hierarchical RL with credit assignment, and separates training from runtime execution.

  • Results

    Experiments across text-to-SQL, retrieval-augmented generation, and math tool-use agents show continuous and stable performance improvement across diverse scenarios.

  • Takeaways & Limitations

    The framework provides a standardized approach for optimizing agents built with different frameworks and interaction patterns, including multi-agent workflows.

  • Takeaways & Limitations

    Future work includes further separating trainer, rollout, and agent workflows to address rollout bottlenecks and improve scalability, especially for long-horizon tasks.

Abstract

from arXiv · show

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.

1 Introduction

Agent Lightning addresses the difficulty of applying RL to diverse, multi-step agents by decoupling agent execution from training. It combines a unified data interface, hierarchical credit assignment, and disaggregated infrastructure, with stable improvements across several agent tasks.

  • Motivation: LLM agents still struggle with unfamiliar tools, private-domain data, and complex multi-turn workflows, motivating agent-specific training.RL is attractive because it can use outcome-based rewards without requiring costly step-by-step annotations.
  • Core contribution: Agent Lightning fully decouples agent execution from RL training, enabling almost ZERO code modifications across diverse agent implementations.The framework targets agents built with different approaches, including existing development frameworks and custom implementations.
  • Algorithm: The unified interface models agent execution as an MDP and converts each call into transitions containing state, action, and reward.This representation aligns collected execution data with RL training while abstracting away different agent execution logic.
  • System design: Training-Agent Disaggregation separates RL training from agent execution through Lightning Server and Lightning Client components.The client runs agents and collects trajectories, while the server manages training and exposes updated models through a standardized service.
  • Evaluation: Experiments with LangChain, OpenAI Agents SDK, and AutoGen agents show continuous and stable performance enhancement across text-to-SQL, retrieval-augmented generation, and math tool-use tasks.These results demonstrate the framework across multiple agent implementations and task types.
  • Algorithm: LightningRL uses hierarchical RL and credit assignment to allocate trajectory-level returns to responses generated by individual LLM calls.This design supports complex interaction logic rather than requiring simple concatenated sequences.

2 Modern AI Agents

Modern AI agents combine LLMs and tools through customized orchestration that may change dynamically during execution. This flexibility supports diverse workflows but complicates optimization and leaves many agents without automatic self-improvement.

  • Agent definition: AI agents are software systems incorporating one or more LLM calls, spanning fixed workflows and dynamic multi-agent systems.The paper uses this broad formulation rather than introducing a new formal definition.
  • Components: LLMs provide the core reasoning and generation functions, while tools perform tasks such as retrieval, code execution, and API interaction.Tools may be stateless or stateful and can be implemented as APIs, programs, or libraries.
  • Orchestration: Agent components are orchestrated through execution flows, dependencies, and ordering customized to task requirements.Developers may build agents from scratch or use frameworks such as LangChain, OpenAI Agents SDK, and AutoGen.
  • Orchestration: Modern agent orchestration is often dynamic and nondeterministic, with LLMs selecting actions or tools as context evolves.The number of interaction turns can vary, such as when an agent refines a query or answers directly.
  • Optimization challenge: Dynamic interaction creates data-modeling and learning challenges that make implementation with existing RL frameworks impractical and unscalable.Agents also typically lack mechanisms for automatic self-improvement, especially when private data creates demanding performance requirements.

3 Agent Lightning

Agent Lightning is introduced as a framework whose data interface, MDP formulation, hierarchical RL algorithm, and system design jointly support RL-based optimization of diverse agents.

  • Framework overview: The framework first defines a unified execution-data interface, then formulates agents as MDPs, develops hierarchical RL, and describes its tailored system design.These components provide the conceptual and infrastructural path from agent execution data to model updates.

3.1 Unified Data Interface

The unified data interface represents arbitrary agent executions as state changes caused by component invocations, with associated rewards. This transition-based representation supports RL optimization without reconstructing the full execution graph.

  • Unified Data Interface: Agent executions can be modeled through relevant states and calls without parsing the complete execution DAG.The MDP formulation identifies the current state and factors influencing transitions.
  • State and Call: A state is a snapshot of execution containing evolving variables such as program counters, values, call stacks, and resource context.The paper emphasizes variables sufficient to describe execution status rather than every intermediate operation.
  • State and Call: Semantic Variables are key variables used or modified by components that represent critical execution semantics.Auxiliary operations such as loop counters and simple string manipulations are not the primary focus for optimization.
  • State and Call: Component invocations capture the called LLM or tool together with its inputs, outputs, and invocation metadata.Metadata can include component identity, version, endpoint, sampling strategy, and temperature.
  • Reward and Dataset: Rewards may be assigned at intermediate steps or at execution completion, with terminal rewards commonly evaluating overall task success.Intermediate rewards can reflect tool success or partial completion, while the final reward assesses the full task outcome.
  • Illustrative Example: In the RAG example, an LLM generates a query, a search tool retrieves passages, and another LLM produces an answer from the passages and user input.Each invocation updates semantic variables and contributes to an ordered execution trajectory.
  • Extensibility: The interface captures state changes from non-LLM components and supports selective multi-agent optimization, prompt tuning, and other downstream methods.Complete execution context enables optimization without explicitly parsing the entire workflow DAG.

3.2 Markov Decision Process in Agents

Agent Lightning models agent execution as a partially observable Markov decision process, representing each policy-LLM invocation as an action within a state transition. This formulation yields a unified transition interface for extracting training data from diverse agent workflows.

  • MDP formulation: Agent execution is modeled as a POMDP with states, observations, actions, transition dynamics, and rewards.The policy LLM observes inputs derived from agent states and produces outputs that update those states.
  • MDP formulation: A complete LLM output sequence from one invocation is treated as a single action, rather than separate token-level decisions.The action is the generated sequence output_t, after which the agent transitions to a new state and may receive a scalar reward.
  • Data extraction: Agent trajectories retain each policy-LLM input, output, and associated reward for reinforcement-learning updates.The extracted data focuses on information relevant to optimizing the policy model rather than reproducing every execution detail.
  • Data extraction: Flexible LLM inputs can encode prompts, histories, queries, reasoning steps, retrieved documents, templates, or structured formats produced by agent logic.This accommodates the varied contexts encountered in dynamic agent executions.
  • Agent scenarios: The unified interface separates task-specific agent design from policy optimization and extends to multi-agent workflows with changing LLM roles.For example, one call can generate a search query and a later call can generate the final answer in a RAG workflow.

3.3 LightningRL: A Hierarchical RL Method for Optimizing LLMs in Agents

LightningRL adapts single-turn LLM reinforcement-learning methods to multi-step agents by decomposing trajectories into transitions and assigning episode returns through hierarchical credit assignment. This avoids concatenated-turn masking while preserving flexible observations and compatibility with established RL optimizers.

  • Motivation: Existing LLM RL methods primarily optimize single-call responses, whereas agents require multiple LLM actions interleaved with tools and environment interactions.LightningRL addresses this mismatch with a hierarchical reinforcement-learning design.
  • Hierarchical credit assignment: LightningRL decomposes each trajectory into transitions containing the current input context, LLM output, and reward.The episode return is first distributed across actions, then across tokens within each action.
  • Hierarchical credit assignment: The current implementation assigns every action the same value, equal to the episode’s final return.Existing single-turn RL algorithms then handle token-level optimization within each action.
  • Compatibility with RL methods: LightningRL reuses single-turn methods such as GRPO by grouping decomposed actions from repeated executions of the same task for advantage estimation.Similar adaptations apply to PPO and REINFORCE++.
  • Design advantages: Transition-level observations support summaries, structured prompts, and role instructions, unlike concatenation-and-masking approaches that restrict modular context construction.Transition decomposition also removes additional masking and helps manage long trajectories through batches and batch accumulation.
  • Limitations and extensions: More sophisticated heuristic, learned, or high-level value-based credit assignment remains a future extension, although identical assignment works across the reported scenarios and datasets.The limitation concerns the granularity of current credit distribution rather than the transition interface itself.

3.4 System Design of Agent Lightning

Agent Lightning uses Training-Agent Disaggregation to separate RL training from agent execution across a server-client architecture. The design supports independent agent development, trajectory capture, parallel execution, robustness mechanisms, and flexible reward or environment services.

  • Motivation: Training-agent integration is difficult because RL and agent-development frameworks are complex, rapidly evolving, and fragmented.Agent Lightning addresses this with a unified framework for developing and training agents.
  • Disaggregation architecture: Training-Agent Disaggregation separates compute-intensive LLM generation from diverse application logic and tools written in traditional programming languages.Agent logic can run independently from GPU resources through an OpenAI-like model interface.
  • Server-client design: The Lightning Server manages RL training, model optimization, task orchestration, data, and client communication, while the Lightning Client encapsulates and executes agents.This mutual independence makes the trainer agent-agnostic and the agent runtime trainer-agnostic.
  • Runtime capabilities: The client runtime orchestrates agent execution, captures trajectories, handles errors, and communicates with the server.It is designed to support robust, scalable, and efficient training workflows.
  • Runtime capabilities: Multiple agent instances run concurrently through data parallelism to increase rollout throughput and reduce latency.Large batches are important for resource utilization because rollout latency can bottleneck RL training.
  • Data capture: Trajectory capture can use OpenTelemetry and AgentOps instrumentation or a lightweight tracing mechanism embedded in the OpenAI-like API endpoint.These techniques are intended to integrate with existing agent codebases without modifying agent logic.
  • Rewards and environments: Automatic Intermediate Rewarding targets sparse and delayed rewards by mining intermediate feedback from agent execution without relying solely on costly annotation or complex reward logic.Resource-intensive environments and reward functions can be hosted as shared services rather than run locally.

4 Results

Agent Lightning is evaluated on three distinct agent tasks implemented with different frameworks, showing continuous and stable performance improvement across diverse scenarios.

  • Experimental coverage: The evaluation covers text-to-SQL, retrieval-augmented generation, and math question answering with tool usage across distinct agent frameworks.The tasks use LangChain, the OpenAI Agents SDK, and AutoGen, respectively.
  • Text-to-SQL via LangChain: Stable reward improvement demonstrates optimization of complex multi-step text-to-SQL decisions involving code generation and tool use.The task requires generating, executing, checking, and potentially rewriting SQL before answering.
  • Retrieval-Augmented Generation via OpenAI Agents SDK: Stable performance improvement on the challenging RAG task demonstrates effectiveness in complex and open-ended retrieval scenarios.The agent formulates free-text queries, searches a 21-million-document Wikipedia database, and reasons over retrieved information.
  • Math QA with Tool Usage via AutoGen: Performance improves consistently in tool-augmented math tasks requiring precise calculator calls and reasoning.The AutoGen agent decides when to invoke a calculator, interprets tool outputs, and produces the final answer.

5 Discussion

Agent Lightning positions transition-based, decoupled training as a way to support diverse agent workflows while reducing context-growth problems. The discussion also identifies future work in broader optimization, more efficient RL algorithms, infrastructure, and serving.

  • Related work: Transition-based modeling supports diverse agent architectures, including multi-agent orchestration, rather than only simple sequential workflows.It organizes trajectories as individual transitions instead of concatenated turns.
  • Related work: Existing RL systems often require developers to rebuild agents inside the training system, coupling agent execution logic to data organization.This complicates integration across customizable agents built with different frameworks and external tools.
  • Related work: Agent Lightning completely decouples RL training from agents, requiring almost no agent-side code modification and avoiding data-collection logic inside the training engine.The framework is intended to integrate agents developed with different frameworks or built from scratch.
  • Future work: The framework supports optimization beyond RL, including automatic prompt optimization, through a Component of Interest abstraction for selected execution components.Extending support to more optimization methods remains part of the future development roadmap.
  • Future work: More efficient RL algorithms remain important for complex agent scenarios, especially long-horizon credit assignment, exploration, and off-policy learning.Transition-based organization is designed to make additional algorithms easier to integrate.
  • Future work: Future infrastructure work includes further trainer, rollout, and agent-workflow disaggregation, alongside coordinated algorithm and system advances for long-horizon tasks.Efficient serving and long-context acceleration are also identified as directions for improving resource use, response times, and scalability.

A Example to Optimize an Existing Agent with Agent Lightning

The example shows how Agent Lightning can optimize an existing game-interacting agent without changing its agent code. Training uses a separate script that connects the agent, environment, task data, and framework resources.

  • Example setup: The example uses a simple game environment, such as 20 Questions or Guess Number, to demonstrate optimization of an existing agent.The stated goal is to improve performance without modifying the agent code itself.
  • Example setup: The agent project is organized around training data, a game server environment, and an agent.py implementation file.The listed folder structure contains data/train.jsonl, environments/game_server.py, and agent.py.
  • Example setup: The agent function receives an LLM endpoint and game environment, while the game server uses game seeds and ground truth stored in the training data.This separates agent logic from environment setup and task-specific records.
  • Training script: A separate train.py script imports the existing agent function, game server, and Agent Lightning client, resource, and task interfaces.The script creates game instances and maps task data to the training workflow.
  • Training script: The example documentation warns that the open-source API may undergo modifications, directing users to the GitHub repository for current information.This caveat applies to the illustrated API usage.

B Process Diagram of Agent Lightning

Figure 8 is identified as a process diagram for Agent Lightning.

  • Process Diagram: Figure 8 presents a process diagram.
  • Process Diagram: The diagram belongs to the Agent Lightning framework description.
  • Process Diagram: The supplied passage provides the figure label but no further process details.
Loading 2508.03680v1…