Paper Notes

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
机构Microsoft Research
发表arXiv · 2025
参考GitHub

28分钟阅读

Agent Lightning: Train ANY AI Agents with Reinforcement Learning

Paper: arXiv:2508.03680 Code: microsoft/agent-lightning Code reference: main @ 0b40cb72 (2026-04-29)

1. 动机与问题

这篇论文的核心问题不是“再提出一个 agent 定义”，而是把现实中的 agent 训练问题重新拆成一个可被 RL 系统消费的接口问题：已有 LLM RL 框架通常假设输入是一个 prompt、输出是一段 response、奖励作用在这段 response 上；真实 agent 却常常由多次 LLM 调用、工具/API 调用、环境状态、动态分支、错误恢复和多 agent 协作组成。如果把整条轨迹强行拼成一条长序列，再用 mask 区分哪些 token 该优化，训练端会被 agent 的执行细节绑死，agent 端也要为训练框架重写大量胶水代码。Agent Lightning 要解决的就是这个耦合：让 agent 按原来的 LangChain、OpenAI Agents SDK、AutoGen 或裸 Python 逻辑运行，同时让训练系统只看到标准化的状态、动作、奖励转移。

论文选择的切入点很实用。它把“agent 是什么”抽象为“一个执行过程中包含一次或多次 LLM 调用的软件系统”。这个定义覆盖了固定 workflow、图式 orchestration、动态规划、多 agent 长程任务和工具增强推理。对 RL 来说，关键不是完整还原软件执行图，而是知道当前可见状态是什么、哪一次 LLM 调用产生了可优化动作、该动作后来得到什么奖励。传统的端到端序列建模会把 tool output、环境反馈、prompt 模板、系统消息和 LLM 生成 token 都塞进同一条上下文；Agent Lightning 则把 agent 执行看成一系列 component invocation，并从中抽取 LLM 相关的 transition。

这带来的现实价值有三层。第一，开发者不必把 agent 迁移到训练框架内部；训练系统只需要通过 OpenAI-like API、trace、reward span 与 agent 通信。第二，训练算法可以选择性优化某些 LLM 调用或某些 agent，而不是被整条上下文绑定；这对 SQL agent 中的 writer/reviewer/rewrite agent、RAG 中的 query generator/answerer、以及 tool-use math agent 都很重要。第三，长上下文压力被拆解：每次 LLM 调用形成独立 transition，训练时只对该调用的 input/output/reward 估计优势，而不是把所有历史 turn 和 tool trace 作为一个超长 response 处理。

论文的主要对比对象是多轮 RL 中常见的“concatenate + mask”路线。那条路线适合简单顺序工作流，但在真实 agent 中会遇到三类问题：轨迹长度随交互增长导致训练成本和上下文长度失控；不同 agent 框架产生的消息格式、tool schema、metadata 难以统一；多 agent 或动态 workflow 中很难在一条序列里清楚指定到底优化哪个决策。Agent Lightning 的立场是：只要把 agent 执行抽象成 MDP，并把每一次 LLM 调用抽成 transition，底层可以继续使用现有 GRPO/PPO 等单轮 LLM RL 方法。

2. 核心想法

最重要的想法是“训练端与 agent 端解耦”。论文把它称为 Training-Agent Disaggregation：训练端负责模型权重、rollout 调度、资源管理和 RL 更新；agent 端负责业务逻辑、工具调用、环境交互和 reward 计算。两者通过统一 trace/store/API 连接，而不是共享同一个进程、同一套框架或同一条拼接序列。这样做的关键不是简单加一个 RPC 层，而是同时定义了数据接口、算法拆解和系统架构。

从数据角度，Agent Lightning 定义了统一数据接口。agent 执行中的 state 是程序计数器、变量值、调用栈、资源上下文等语义变量的快照；call 是 LLM 或 tool 的一次调用；input/output 是某个 call 在某个 state 中可见或写入的语义变量。训练不需要完整 DAG，只需要保留与策略 LLM 决策相关的 input、output 和 reward。也就是说，原始 agent 可能有复杂控制流，但训练样本最终被整理为一组 (input_t, output_t, r_t) transition。

从算法角度，LightningRL 把一次 agent episode 拆成多个可单独训练的 LLM transition。一次 LLM 调用生成的一串 token 被视为 MDP 中的一个高层 action；底层 token 级优化仍可交给 GRPO/PPO 这类已有 LLM RL 算法。信用分配模块负责把最终 reward 或中间 reward 归到相应 transition 上。这样，复杂 agent 的长程互动先在高层被切成 transition，再在每个 transition 内部用单轮 LLM RL 训练 token policy。

从系统角度，Agent Lightning 引入 Lightning Server / Lightning Client / Store / Runner / Tracer / Adapter 这一组角色。agent 继续调用 OpenAI-compatible endpoint 获得当前 policy LLM 的输出；Tracer 记录 LLM 调用、tool 调用和 reward；Store 存储 rollout、attempt、span、resource；Algorithm 从 Store 查询 span，经 Adapter 转成 triplet，再更新资源或模型。这个设计的实质是把 agent observability 技术用于 RL 数据采集：不是要求开发者手工把每一步写成训练样本，而是尽量从 OpenTelemetry 风格的 spans 中恢复训练需要的信息。

3. 方法细节

3.1 Agent 执行的 MDP 形式化

论文把 agent 的每个任务输入记为 x，同一个任务可以被执行多次，记为第 k 次 execution。执行过程中的 state_t(x,k) 是 agent 在第 t 个时间点的快照。对策略 LLM 来说，真正可见的是某个语义 input，例如用户问题、当前检索结果、数据库 schema、已有错误反馈或工具返回值；LLM 输出的是一串 token，对应一个语义 output，例如 SQL、检索 query、calculator call 或最终答案。

在这个 MDP 中，状态转移由 agent 逻辑和环境共同决定，通常不可显式建模；奖励函数可以只在 episode 末尾给，也可以在中间步骤给。统一数据接口因此包含三类对象：state/call 描述执行；reward/dataset 描述反馈；transition 描述可优化的 LLM 决策。一个 execution 被表示为若干 (call_i, r_i)；进一步抽取与目标 LLM 有关的调用，形成 (input_i, output_i, r_i)。这里的关键是 input_i 不等于整条历史，而是该 LLM 调用在当前 state 下需要看到的上下文，因此可以自然避免拼接整条 agent trace。

这种形式化还支持“单 LLM 多 agent 角色”。例如 RAG 中同一个 LLM 可以先扮演 query generator，再扮演 answer generator；SQL agent 中同一个基础模型可以在 writer、reviewer、rewriter 等不同 prompt 下承担不同角色。transition 里保留 agent name 或 component identity 后，训练时可以只优化特定角色。例如只训练 SQL writer 和 reviewer，不训练 executor；或者只训练 RAG 的 query generation，而保留 retriever 不变。相比 sequence mask，这种选择更贴近软件组件边界。

3.2 LightningRL：高层 transition + 底层单轮 RL

LightningRL 的基本流程是：先让 agent 以原始方式执行任务，收集 trace 和 reward；再由 credit assignment 把 reward 分配到每个 LLM transition；最后把每个 transition 当作单轮 LLM RL 样本。它的层次结构可以理解为两层：高层 MDP 把一次 LLM 调用当作一个 action；低层 RL 在该 action 的 token 序列上计算 logprob、优势和策略梯度。

论文用 GRPO 做主要说明。单轮 GRPO 把同一 task 的多个 sampled outputs 分组，用组内 reward 估计 advantage；多轮拼接方法会把整条 trajectory 作为 response，并 mask 非 LLM token；LightningRL 则把同一 task 的多条 trajectory 拆成 transition，然后对同一 task/同一 transition 组进行 advantage estimation。这样 reward assignment 可以与 agent 的逻辑边界对齐：如果中间有 tool success reward，就可以给对应调用；如果只有 final reward，就可以给最后调用或按策略传播到相关调用。

这个设计有两个直观好处。第一，context construction 变得灵活：每个 transition 的 input 可以由当前 state 中必要变量构造，而不是默认塞入所有历史。第二，多 agent 选择性优化更简单：训练数据层面筛选 transition 即可，而不需要在一条长 token 序列上维护复杂 mask。代价是 credit assignment 的质量会成为瓶颈；如果所有 reward 都只在最终答案上给，早期检索 query 或 SQL rewrite 的学习信号可能仍然稀疏，论文也在 future work 中明确把长程信用分配列为需要改进的方向。

3.3 Training-Agent Disaggregation 系统架构

Training-Agent Disaggregation 的核心是把“计算密集的 LLM 生成/训练”与“轻量但多样的 agent 业务逻辑/工具”分离。训练端可以管理 GPU、模型权重、vLLM/verl 等基础设施，并向 agent 暴露 OpenAI-like API；agent 端可以运行在普通 CPU 服务、开发者本地或已有应用环境里。agent 不需要和训练进程 colocate，也不需要把业务逻辑改写为 RL 框架内部的 rollout 函数。

论文中的 Lightning Server 负责训练流程控制、任务分发、资源更新和模型 API；Lightning Client 负责 agent 运行、trace 采集、环境/奖励服务对接。Store 是中间事实源：algorithm 往里写任务和资源，runner 从里取 rollout，tracer 把 span 写回，algorithm 再查询 span 并训练。这个结构让训练和 agent 具有双向无关性：trainer-agnostic agent 可以在不了解 RL 细节的情况下运行；agent-agnostic trainer 只关心 span、triplet 和 reward。

OpenTelemetry/observability 在这里不是附属功能，而是数据接口的载体。一次 LLM call、tool call、reward emission 都可以被记录成 span；span 中包含 prompt、completion、token ids、agent/component 标识、rollout id、attempt id、reward 等属性。Adapter 的任务是把这些原始 spans 变成 (prompt, response, reward) triplet。论文强调“几乎零代码修改”，但这里更准确的理解是：已有 agent 不需要被训练框架重写；实际落地仍需要让 LLM 调用走可追踪的 proxy/API，并在合适位置提供 reward。

3.4 与当前开源实现的对应关系

Code reference: main @ 0b40cb72 (2026-04-29)

Paper ConceptSource FileKey Class/Function
统一 transition / rollout / attempt 数据结构agentlightning/types/core.pyTriplet, Rollout, AttemptedRollout, Dataset
从 trace/span 恢复训练 tripletagentlightning/adapter/triplet.pyTraceTree, TracerTraceToTriplet.adapt, TraceTree.to_trajectory, span_to_triplet
reward span 与 final reward 读取agentlightning/emitter/reward.pyemit_reward, get_reward_value, find_reward_spans, find_final_reward
Store 的客户端/服务端解耦agentlightning/store/client_server.pyLightningStoreServer, LightningStoreClient, enqueue_rollout, query_spans, update_resources
agent rollout 执行与 tracingagentlightning/runner/agent.pyLitAgentRunner
OpenAI-compatible LLM proxy 与 token/trace 注入agentlightning/llm_proxy.pyLLMProxy, LightningSpanExporter, RolloutAttemptMiddleware, AddReturnTokenIds
顶层训练编排agentlightning/trainer/trainer.pyTrainer
VERL/GRPO 算法入口agentlightning/algorithm/verl/interface.py, agentlightning/verl/trainer.pyVERL, AgentLightningTrainer
论文三类任务示例examples/spider/*, examples/rag/*, examples/calc_x/*LitSQLAgent, RAGAgent, calc_agent, train_sql_agent.py, train_rag.py, train_calc_agent.py

基于当前代码，Agent Lightning 的实际执行可概括为以下伪代码。注意这不是论文中的抽象伪代码，而是对应开源代码结构的路径级实现逻辑：

# Algorithm side: prepare resources and rollout tasks
algorithm = agl.VERL(config)                 # agentlightning/algorithm/verl/interface.py
trainer = agl.Trainer(
    n_runners=10,
    algorithm=algorithm,
    adapter={"agent_match": active_agent},   # filter components/agents of interest
)
trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)
 
# Runner side: execute original agent workflow
runner = LitAgentRunner(tracer=OtelTracer(...))
with runner.run_context(agent=agent, store=LightningStoreClient(...)):
    attempted = store.dequeue_rollout()
    resources = store.get_latest_resources()
    result = agent.training_rollout(attempted.input, resources=resources)
    # tracer streams spans; explicit or final reward becomes reward span
    store.update_attempt(status="completed")
 
# Adapter side: turn observability traces into RL triplets
spans = store.query_spans(rollout_id=...)
trace_tree = TraceTree.from_spans(spans)
transitions = TracerTraceToTriplet(agent_match=active_agent).adapt(spans)
# each transition ~= {prompt/token_ids/raw_content, response/token_ids/raw_content, reward}

当前 repo 中的 examples/spider/train_sql_agent.py 使用 adv_estimator="grpo"，实际训练脚本里有 train_batch_size=32、max_prompt_length=4096、max_response_length=2048 等配置，并通过 Trainer(n_runners=10, algorithm=agl.VERL(config), adapter={"agent_match": active_agent}) 指定并发 rollout 和可优化 agent。

examples/rag/train_rag.py 同样使用 GRPO，默认 multi-GPU 配置中 train_batch_size=16，单卡配置会降到 4，Trainer(n_runners=4, ...) 表明 RAG 示例更保守。examples/calc_x/train_calc_agent.py 的默认配置也使用 GRPO，train_batch_size=32、max_prompt_length=4096、max_response_length=2048，并额外提供 llm proxy、external store、algorithm/runner 分角色运行等选项。

这些配置说明开源实现不是只给 paper demo，而是把论文中的 disaggregation 具体落到了 store、runner、adapter、trainer 的多进程/多角色运行上；本文所有训练脚本引用都使用 repo-root 相对路径，即 examples/spider/train_sql_agent.py、examples/rag/train_rag.py 和 examples/calc_x/train_calc_agent.py。

3.5 为什么这个方法可能有效

Agent Lightning 有效的原因在于它把 agent RL 中最难泛化的部分缩到了“可观测、可筛选、可转换”的 trace 层。不同 agent 框架的业务逻辑可能差异极大，但 LLM 调用大多可以被表达为 prompt/response/token ids，reward 大多可以被表达为 scalar span。只要这两个接口统一，训练算法就不需要理解 LangGraph 节点、AutoGen 消息循环、OpenAI Agents SDK 的工具封装或自定义 Python 状态机。它优化的是 component invocation，而不是整个应用框架。

另一个原因是 credit assignment 的粒度更符合 agent 结构。真实 agent 的错误通常不是“整条 response 都坏了”，而是某一步 query 写错、某次 SQL validation 判断错、某次 tool call 参数错或某个最终答案整合错。transition 级训练让每个错误更容易映射到某个 LLM 决策。即便只有最终 reward，也至少可以把最后一次 LLM 输出作为直接监督对象，并为未来更细粒度 reward 留出接口。

最后，这种系统拆分更接近生产部署。agent 业务逻辑常常依赖数据库、网页、企业 API、模拟器或已有服务，不适合塞进 GPU 训练进程。TA Disaggregation 让 GPU 训练、inference serving、agent runtime、reward/environment service 可以分别扩缩容。论文在 future work 中进一步提出 trainer、rollout inference engine、agent workflow 的更细粒度 disaggregation，本质上是把 agentic RL 变成一个系统工程问题，而不是单一算法问题。

4. 实验与结果

论文在三个任务上验证：Text-to-SQL、Open-domain QA/RAG、Math QA with tool usage。三个任务分别使用不同 agent 框架：LangChain/LangGraph、OpenAI Agents SDK、AutoGen；数据集分别是 Spider、MuSiQue、Calc-X；工具分别是 SQL executor、Wikipedia retriever、calculator。表格设置中，Text-to-SQL 是 3 个 agent、优化 2 个；RAG 和 Math QA 都是单 agent、优化 1 个。这组设计的目的不是刷 SOTA，而是证明同一训练框架能覆盖不同 agent scaffold 和不同工具环境。

4.1 Text-to-SQL via LangChain

Spider 是跨数据库 text-to-SQL benchmark，包含超过一万问题、约 200 个数据库、138 个 domain，测试时要求泛化到未见数据库。论文使用 Llama-3.2-3B-Instruct 作为 base model。agent workflow 是多 agent：SQL writer 先生成 query；SQL executor 执行；如果 query 有错，reviewer 给出反馈；rewriter 根据反馈改 SQL；最后 agent 基于执行结果回答问题。奖励由最终答案是否正确决定，测试指标也是 answer accuracy/reward。

图中训练 reward 从约 0.05 附近快速升到 0.4 以上，后期多在 0.45—0.6 区间波动，并偶有 0.65—0.7 的高点；测试 reward 从约 0.15 提升到 0.5 以上，后续在 0.51—0.57 附近震荡。这个结果说明 transition 级训练能处理 SQL 生成、执行反馈、错误检查和 rewrite 这种多步工具循环。它也暴露了一个现实点：训练 reward 噪声很大，说明最终答案 reward 对中间 SQL 决策仍然是稀疏和高方差的，后续需要更强的过程 reward 或 credit assignment。

4.2 RAG via OpenAI Agents SDK

RAG 实验使用 MuSiQue，多跳问答需要组合多个单跳事实。论文为了模拟真实搜索环境，把检索库设为整个 Wikipedia，约 2100 万文档；retriever 使用 BGE embedding 与 cosine similarity。base model 同样是 Llama-3.2-3B-Instruct。agent 首先生成自然语言检索 query，通过 Wikipedia retriever 找证据，然后基于证据输出答案。奖励由最终答案是否正确决定。

RAG 的绝对 reward 明显低于 SQL 和 calculator：测试曲线从接近 0 提升到约 0.23，训练曲线从约 0.01—0.05 逐步升到 0.25—0.35，并在后期出现约 0.4 的峰值。这个现象符合任务难度：MuSiQue 的多跳搜索空间大，retriever 覆盖整个 Wikipedia，最终答案 reward 对 query 质量、证据召回、答案整合都敏感。Agent Lightning 的贡献在于它仍能在 OpenAI Agents SDK scaffold 下稳定获得提升，说明统一 trace/transition 接口没有被特定框架限制。

4.3 Calculator / Math QA via AutoGen

Calc-X 实验关注 tool-use math reasoning。数据集从 GSM8K、Ape210K 等数学数据修改而来，强调外部 calculator 的调用与结果整合。agent 使用 AutoGen，单个 LLM 需要决定何时调用 calculator、如何组织参数、如何解释工具输出并给出最终答案。奖励同样基于最终答案正确性，测试指标为 answer accuracy/reward。

Calculator 任务的提升最明显：测试 reward 从约 0.05 快速升到 0.56，随后逐步到 0.74—0.77 区间；训练 reward 从约 0.07 升到 0.6—0.7 区间，后期仍有波动但整体平台更高。相比 RAG，这个任务的 reward 更密集可学：calculator call 的格式和数值正确性更容易从最终答案中反推出有效行为，工具环境也更确定。它支持论文的一个重要主张：Agent Lightning 能训练 tool-augmented agent，而不只是在纯文本单轮任务上调模型。

4.4 实验解读

这三组实验共同验证的是框架泛化性，而不是某一个 benchmark 的极限性能。它们覆盖了三个维度：多 agent SQL 工作流、单 agent 大规模检索、多步 tool-use math；也覆盖了三个主流 agent 框架。所有曲线都显示稳定或总体上升，说明把 agent trace 拆成 transition 后，现有 GRPO 类算法可以接上真实 agent 执行。但实验仍有两个不足：论文没有提供更系统的 ablation，例如 transition-based vs concatenate-mask 的同等资源对比，也没有量化不同 credit assignment 策略对早期步骤的影响；曲线多以 reward 展示，缺少更细的 error taxonomy，例如 SQL 语法错、检索召回错、calculator 参数错分别改善多少。

5. 评价、局限与可复现注意点

5.1 主要贡献评价

这篇论文的贡献更像“agentic RL 的接口与系统层标准化”，而不是单一新 RL objective。它的价值在于把三个长期混在一起的问题拆开：agent 业务逻辑如何运行、LLM policy 如何训练、trace/reward 如何变成训练数据。这个拆分让已有 agent 可以以较小改动接入 RL，也让训练算法可以复用 verl/GRPO/PPO 等生态。对实际工程而言，这比提出一个只适用于某个 benchmark 的新算法更重要。

论文中最值得借鉴的是 transition-based modeling。它承认真实 agent 的控制流复杂、非确定、框架多样，因此不试图把所有 agent 统一成一个 DSL，而是在观测层统一为 spans，在学习层统一为 triplets。这个层次选择很合理：业务框架仍保持开放，训练数据接口保持收敛。对于企业内已有 agent、复杂 tool service 或多 agent pipeline，这种“只要求可观测和可奖励”的接入门槛更现实。

5.2 局限

第一，credit assignment 仍是关键瓶颈。论文提出了模块化位置，但实验主要依赖最终答案 reward 展示曲线。对于 RAG 这类长程任务，最终答案错并不能明确说明是 query 错、retriever 召回错、证据阅读错还是答案整合错。transition 化降低了训练样本长度，却没有自动解决稀疏 reward 的语义归因问题。未来需要更强的过程 reward、工具级 verifier、反事实 rollout 或 off-policy credit assignment。

第二，框架的“零代码修改”需要谨慎理解。agent 业务逻辑可以不重写，但要想训练，仍需要 LLM 调用可被代理或 trace，reward 能被 emit，token ids/logprobs 能被训练端拿到，agent/component identity 能被匹配。开源代码里 LLMProxy、LightningSpanExporter、emit_reward、TracerTraceToTriplet 等模块正是为此服务。对复杂生产 agent 来说，接入 observability 和 reward instrumentation 仍是工程工作。

第三，实验没有完整展示系统扩展性边界。TA Disaggregation 理论上支持 trainer、inference engine、agent workflow 分离扩缩容，但论文主要给 reward 曲线，缺少吞吐、延迟、GPU 利用率、runner 数量、store 瓶颈、trace 体积等系统指标。对于 agentic RL，rollout bottleneck 往往比 update bottleneck 更难处理；如果 agent 工具很慢或 reward 服务昂贵，Store 和 runner 调度会成为重要问题。

第四，多 LLM / 多策略联合优化仍停留在讨论层。论文指出可把多个 LLM 视为独立 MDP 或用 MARL/博弈论建模，但当前方法更自然支持“一个 LLM 扮演多个角色”或“选择性优化某些 transition”。如果多 agent 使用不同模型、不同权限、不同 reward，联合训练的非平稳性和协调问题会比当前实验复杂得多。

5.3 Future Work 与可扩展方向

论文的 future work 很明确。第一，支持 RL 之外的优化方法，例如 automatic prompt optimization。通过 Component of Interest，prompt template rendering 也可以被视为可优化 component，从而使用同一 trace/transition 数据结构做 prompt 搜索或其他优化。第二，改进 RL 算法，尤其是长程 credit assignment、exploration、off-policy learning。第三，改进 agentic RL 系统基础设施，进一步分离 trainer、rollout inference engine 和 agent workflow，以缓解 rollout bottleneck。第四，在 serving 侧利用更友好的抽象、长上下文加速和更好的工具/环境资源调度。

对复现者来说，建议优先从开源 repo 的三个示例入手：examples/calc_x 最容易看到快速提升，因为工具确定、奖励相对清晰；examples/spider 能覆盖多 agent 和 SQL executor；examples/rag 更能暴露长程检索与稀疏 reward 的挑战。复现时不要只看 config/base 或 README 默认值，应以对应训练脚本为准；当前代码中的实际训练入口是 agl.Trainer(...).fit(...)，算法多用 agl.VERL(config) 和 adv_estimator="grpo"，adapter 通过 agent_match 选择要优化的 agent/component。

5.4 一句话结论

Agent Lightning 的核心贡献是把“任意 agent 的 RL 训练”变成一个可观测 trace 到 transition/triplet 的标准化问题，并用 Training-Agent Disaggregation 把 agent 业务逻辑与 LLM RL 训练系统解耦；它不是彻底解决长程 agent credit assignment 的终点，但给出了一个现实可接入、可扩展、能覆盖多框架 agent 的工程化起点。