Agent Lightning: Train ANY AI Agents with Reinforcement Learning

1 Introduction

• Agents need training beyond prompt engineering; outcome-based RL avoids costly step-by-step annotations, but existing static, single-call approaches poorly accommodate diverse, multi-call agent workflows.Agents remain error-prone with unfamiliar workflows, private-domain data, and tools, while their execution involves multiple LLM calls and external interactions. Real-world interaction data offers scale and diversity beyond human-curated datasets, motivating fine-tuning for specialized skills and more versatile interactive models.Source: §1 · pp. 1–2
• An MDP-based interface converts agent execution into state–action–reward transitions, while LightningRL assigns trajectory-level credit to individual LLM responses for single-turn RL updates.States describe execution snapshots, and actions are policy-LLM outputs; the interface abstracts away orchestration and framework details. This supports flexible context construction, selective optimization of multiple agents, and avoidance of excessively long accumulated contexts.Source: §1 · pp. 2–3
• Training-Agent Disaggregation separates a Lightning Server controlling RL training from a Lightning Client handling communication, agent execution, and trajectory collection.The server exposes updated models through an OpenAI-like API, providing a standardized training service. The client transparently manages execution and collection without code modifications, supporting integration with diverse agent frameworks.Source: §1 · pp. 2–3
• The runtime reuses observability infrastructure for trajectory collection and provides customizable Automatic Intermediate Rewarding to mitigate sparse rewards using monitoring signals.Observability frameworks such as OpenTelemetry connect monitoring data to RL training. AIR assigns intermediate rewards to transitions using signals such as tool-call return statuses.Source: §1 · p. 2
• Reward curves show continuous, stable performance improvements across text-to-SQL, retrieval-augmented generation, and tool-using math question answering agents.The agents were implemented with LangChain, OpenAI Agents SDK, and AutoGen, respectively, demonstrating applicability across different agent scenarios.Source: §1 · pp. 2–3
• Agent Lightning fully decouples agent execution from RL training, supporting agents regardless of implementation approach with almost no code modifications.The unified interface also collects diverse agent-generated data for improving model capabilities.Source: §1 · pp. 2–3

2 Modern AI Agents

• The paper uses a broad formulation of an AI agent as software incorporating one or more LLM calls, rather than proposing a new definition.This formulation accommodates diverse and rapidly evolving agent designs.Source: §2 · p. 3

2.1 Component

• LLMs provide stateless prompt-to-response reasoning and generation, typically through API endpoints, and agents can use multiple models with distinct endpoints for different tasks.Computational demands typically place models on high-performance servers or cloud platforms.Source: §2.1 · p. 3
• Tools perform tasks such as retrieval, code execution, and API interaction, and may be stateless or maintain context across invocations.Tools can be external APIs, local programs, or libraries; protocols such as MCP (Model Context Protocol) standardize agent–tool interactions.Source: §2.1 · p. 3

2.2 Orchestration

• Dynamic orchestration complicates data modeling and downstream learning, making diverse agents impractical and unscalable to implement with current RL training frameworks.Users customize execution flow, dependencies, and component ordering. Orchestration need not be fixed or deterministic: LLMs can choose subsequent actions and tools based on evolving context.Source: §2.2 · p. 4

2.3 Implementation

• Agents built from scratch or with development frameworks typically lack automatic self-improvement, motivating automated optimization when real-world performance falls short.Frameworks provide modular LLM and tool components and dynamic workflows responsive to task requirements and runtime state. Performance deficiencies are particularly relevant when agents handle private data.Source: §2.3 · p. 4

Section 3.1

• A unified data interface specifies what data to collect during agent execution, independently of agent design.Source: §3 · p. 4

Section 3.3

• A hierarchical reinforcement learning algorithm uses collected agent data to update the LLM.Source: §3 · p. 4

Section 3.2

• A Markov decision process formulation of agents provides the foundation for reinforcement learning.Source: §3 · p. 4

Section 3.4

• Tailored infrastructure introduces design elements that streamline the agent training process.Source: §3 · p. 4

3.1.1 State and Call

• States capture execution snapshots, while optimization focuses on Semantic Variables (Lin et al., 2024) used or modified by components and carrying critical semantics.Snapshots can include program counters, variable values, call stacks, and resource context; repeated executions of the same task can produce different traces.Source: §3.1.1 · p. 5
• Calls record the invoked LLM or tool, its inputs, outputs, and metadata; components need not see the entire state.Metadata includes component identity, version, API endpoint, and runtime parameters such as sampling strategy and temperature.Source: §3.1.1 · p. 5

3.1.2 Reward and Dataset

• Scalar rewards assess invocation quality at intermediate steps or execution completion, with terminal-only feedback supported as a special case.Intermediate rewards assess partial progress, whereas final rewards typically assess overall task success. Training or evaluation datasets comprise tasks and a reward function, allowing repeated, reward-annotated executions.Source: §3.1.2 · p. 5

3.1.3 An Illustrative Example

• The RAG example represents query generation, retrieval, and answer generation as ordered calls that update semantic variables, followed by an answer-quality reward.The final LLM call uses retrieved passages and the original user input; a reward function such as RAG F1 evaluates its answer.Source: §3.1.3 · pp. 5–6, Figure 2 · p. 6
• Capturing non-LLM state changes supports selective optimization within multi-agent systems and optimization beyond model finetuning, including prompt tuning.The paper focuses on LLM-centric optimization, while the interface permits these broader optimization methods.Source: §3.1.3 · p. 6

3.1 Unified Data Interface

• The unified interface decouples agent execution from RL training by capturing execution context without explicitly parsing the entire execution DAG.The interface applies to data from any AI agent; RL optimization requires identifying current states and the calls driving transitions.Source: §3.1 · p. 4, §3.1.3 · p. 6

3.2.1 Formulation

• A single policy LLM is modeled as a POMDP, with each complete invocation output treated as one action and its input as an observation.Observations expose relevant agent state; actions induce usually unknown state transitions and receive scalar state-action rewards. The policy maximizes return, defined as the sum of rewards.Source: §3.2.1 · pp. 6–7

3.2.2 Data Extraction for RL

• A unified interface extracts policy LLM inputs, outputs, and rewards, separating task-specific agent logic from policy optimization without parsing how inputs were produced.Trajectories retain policy decisions and associated rewards while ignoring task-dependent input construction and framework rendering or parsing.Source: §3.2.2 · p. 7
• One LLM can fill multiple agent roles through different prompts, while selected agents are optimized by including only their corresponding transitions.In the RAG example, the same LLM first generates a search query and then answers the question; answer quality determines the final reward.Source: §3.2.2 · p. 7
• Independent MDP optimization simplifies multi-LLM training but ignores interdependencies; MARL or game theory could capture interactions (Lowe et al., 2017).Separate optimization may yield suboptimal coordination. The alternative treats each LLM as an agent with its own objective and policy, capturing cooperative or competitive behavior.Source: §3.2.2 · p. 7

3.3.1 Preliminary on Single-Turn Reinforcement Learning for LLMs

• Single-turn LLM RL treats tokens as actions and rewards completed responses, with algorithms differing chiefly in advantage estimation.PPO (Ouyang et al., 2022) uses a learned critic, whereas GRPO (Shao et al., 2024) normalizes response rewards using same-prompt group means and standard deviations. REINFORCE++ (Hu, 2025) subtracts the training-batch mean reward from the current response’s reward.Source: §3.3.1 · p. 8

3.3.2 Extend to Agent Scenarios via LightningRL

• LightningRL treats each LLM invocation as an action, assigning episode returns across actions before deriving token-level supervision through existing single-turn RL algorithms.An episode can contain multiple LLM invocations, while optimization operates on tokens within each invocation.Source: §3.3.2 · pp. 8–9
• The current implementation assigns every action the final episode return R, while the credit-assignment module permits future heuristic or learned alternatives.The authors report that identical assignment works across multiple scenarios and datasets. Estimating action-specific expected returns with a high-level value function remains future work.Source: §3.3.2 · p. 9
• LightningRL directly reuses single-turn RL algorithms without modification, including lightweight value-free methods.For GRPO, repeated executions of each task are decomposed into actions and grouped by task to compute advantage statistics. Similar adaptations apply to PPO and REINFORCE++.Source: §3.3.2 · p. 9
• Individual transitions allow flexible policy observations, unlike trajectory concatenation with masking, which cannot support the same modular context construction.Inputs can be LLM-generated summaries of previous steps, template-assembled prompts, or instructions specifying the LLM’s current role.Source: §3.3.2 · p. 9
• Transition-based training eliminates additional masking, simplifies implementation, and alleviates accumulated and excessively long contexts caused by concatenating all turns.It avoids masking-related coupling between training and agent execution, token-continuity disruption, and debugging/kernel complexity. Long trajectories become batches of transitions, enabling batch accumulation for efficient updates.Source: §3.3.2 · p. 9

3.4.1 Training-Agent Disaggregation Architecture

• Training-Agent Disaggregation separates RL-managed LLM generation from independently executed agent logic and tools, making training frameworks agent-agnostic and agents trainer-agnostic.An OpenAI-like API connects the components, and agent logic need not be colocated with GPU resources. The framework unifies agent development with RL and other optimization techniques.Source: §3.4 · p. 9, §3.4.1 · p. 10
• The server coordinates batched training and task dispatch, while clients execute agents and return traces and rewards for model updates.An event-driven system coordinates generation and training stages, assigning available tasks as clients become ready. Each task receives a unique OpenAI-like API endpoint for connection and data tracing.Source: §3.4.1 · p. 10

3.4.2 Agent Runtime

• The client manages execution, captures data without agent-logic modifications, and handles failures without disrupting training.Instrumentation uses OpenTelemetry and AgentOps, with lightweight API-endpoint tracing available for frameworks incompatible with OpenTelemetry. Unhandled crashes and hanging tool calls are detected; failed tasks can be retried or reassigned, with logs retained for debugging.Source: §3.4.2 · pp. 10–11
• Two-level data parallelism runs agent instances across workers within a machine and across clients on multiple machines to increase throughput and reduce rollout latency.Lightweight environments and reward functions run within agent workers, while resource-intensive or expensive-to-initialize components can use pooled shared services.Source: §3.4.2 · pp. 10–11
• Automatic Intermediate Rewarding converts system monitoring data into intermediate rewards to provide feedback beyond delayed, sparse rewards.This mechanism supplies learning signals without relying solely on costly human annotation or complex reward computation logic.Source: §3.4.2 · p. 11

4.1 Text-to-SQL via LangChain

• The LangChain experiment simultaneously optimizes SQL-writing and rewriting agents while leaving the checking agent unoptimized, demonstrating selective training within a multi-agent workflow.Three roles use the same LLM with different prompts to generate, check, and revise SQL before answering questions on Spider. Training rewards final-answer correctness, and test performance uses answer accuracy.Source: §4.1 · pp. 11–12

4 Results

• Agent Lightning achieves stable training improvements across text-to-SQL, retrieval-augmented generation, and calculator-assisted math QA implemented with three different agent frameworks.The supplied passages report improving training curves without numerical benchmark comparisons.Source: §4 · p. 11, Figure 5 · p. 12, §4.1 · p. 12

4.2 Retrieval-Augmented Generation via OpenAI Agents SDK

• The OpenAI Agents SDK experiment uses one agent to generate free-text retrieval queries and decide whether to refine them or answer from retrieved documents.MuSiQue requires multi-hop reasoning over 21 million Wikipedia documents, retrieved using BGE embeddings (Zhang et al., 2023) and cosine similarity. Training uses R = 0.9 × Rcorrectness + 0.1 × Rformat, combining answer word-level F1 with adherence to required output tags.Source: §4.2 · pp. 12–13

4.3 Math QA with Tool Usage via AutoGen

• The AutoGen experiment trains a single LLM to choose calculator calls, interpret tool outputs, and integrate intermediate computations into final mathematical answers.Calc-X (Kadlˇc´ık et al., 2023) tests arithmetic and symbolic reasoning with external computation. Final-answer correctness determines training reward, while test performance is measured by answer accuracy.Source: §4.3 · p. 13

5.1 Related Work

• Transition-based modeling supports complex workflows and multi-agent orchestration while alleviating accumulated context length, eliminating custom masks, and enabling hierarchical RL.The MDP formulation, unified interface, and credit assignment module organize execution into individual transitions. Each training sample includes only the current LLM input, avoiding concatenation-specific input, loss, and attention masks.Source: §5.1 · p. 14
• Decoupling RL training from agent execution eliminates training-engine data-collection logic and requires almost no agent code changes, supporting agents built with diverse frameworks.Existing systems often require labor-intensive agent migration and specialized expertise, while embedded tools and APIs add maintenance and scalability overhead. Integration with diverse RL training frameworks is theoretically supported.Source: §5.1 · pp. 14–15
• Unlike task-specific approaches with predefined training and inference workflows, Agent Lightning provides a unified method for general-purpose agent training across tasks and scenarios.The distinction is general-purpose agent training rather than optimization restricted to particular applications.Source: §5.1 · p. 15

5.2 Future Work

• The Component of Interest concept identifies execution-trace components for optimization beyond RL, with broader optimization-method support planned.Prompt-template rendering can be represented as a tool call and selected as a Component of Interest for automatic prompt optimization. The extensible data structure also supports downstream agent-behavior analysis.Source: §5.2 · p. 15
• Future RL algorithm development targets long-horizon credit assignment, exploration, and off-policy learning, with transition-organized data facilitating integration.More efficient algorithms are identified as necessary for optimization in increasingly complex agent scenarios.Source: §5.2 · p. 15
• Further separating trainers, rollout engines, and agent workflows could address rollout bottlenecks, while serving improvements target resource utilization, response times, and long-context processing.Serving directions include LLM-friendly abstractions such as Parrot (Lin et al., 2024), sparse long-context computation such as Minference (Jiang et al., 2024), and improved resource and tool scheduling. Long-horizon training also calls for coordinated algorithm and system advances.Source: §5.2 · p. 15

A Example to Optimize an Existing Agent with Agent Lightning

• A separate train.py script optimizes an existing game-playing agent with Agent Lightning without modifying the agent’s code.The example covers agents interacting with games such as 20 questions or number guessing.Source: A Example to Optimize an Existing… · p. 19
• The agent_run wrapper creates a seeded game, calls the existing agent using the resource’s LLM endpoint, and returns a reward scored against task ground truth.Tasks provide game seeds and ground truth from data/train.jsonl; agent.py contains the existing agent logic.Source: A Example to Optimize an Existing… · p. 19
• The client connects through the AgentLightningServerUrl environment variable, uploads data/train.jsonl, and starts training with client.train(agent_run, nworkers=-1).The Agent Lightning server must be started beforehand; the script also permits specifying a pre-uploaded dataset ID.Source: A Example to Optimize an Existing… · p. 19
• The open-source API may change, so the example directs readers to the GitHub repository for up-to-date information.Source: A Example to Optimize an Existing… · p. 19

• Figure 8 provides the process diagram.Source: Figure 8 · p. 20