Agent Lightning：無需重寫程式碼即可為 AI 代理加入強化學習

AI 代理（AI agents）正在重塑軟體開發，從編寫程式碼到執行複雜指令。然而，基於大型語言模型（LLM）的代理容易出錯，且在複雜的多步驟任務中表現往往不佳。強化學習（RL）是一種讓 AI 系統透過接收其行為的獎勵或懲罰來學習做出最佳決策的方法，並在嘗試與錯誤中不斷進步。RL 可以幫助代理提升效能，但通常需要開發者大規模重寫程式碼。儘管這些代理生成的數據可以透過 RL 訓練顯著提升性能，但重寫代碼的負擔阻礙了這項技術的普及。

為了瞭解決這個問題，來自 微軟亞洲研究院 – 上海 的研究團隊推出了 Agent Lightning。這個 開源 (在新分頁中開啟) 框架透過將代理執行任務的方式與模型訓練分離，使 AI 代理能夠進行 RL 訓練，讓開發者幾乎無需修改程式碼即可增加 RL 能力。

捕捉代理行為以進行訓練

Agent Lightning 將代理的執行視為一系列狀態（states）與動作（actions）的序列，藉此將代理的經驗轉換為 RL 可用的格式。其中每個狀態捕捉代理的當前狀況，而每次 LLM 調用則是一個將代理推向新狀態的動作。

這種方法適用於任何工作流，無論多麼複雜。無論是涉及多個代理協作還是動態工具調用，Agent Lightning 都會將其分解為一系列轉換（transitions）。每個轉換都會記錄 LLM 的輸入、輸出和獎勵（圖 1）。這種標準化格式意味著數據無需任何額外步驟即可用於訓練。

圖 1. 使用檢索增強生成（RAG）代理說明 Agent Lightning 的標準化格式。左側：完整的代理工作流，代理狀態在每個組件步驟後更新。綠色區塊顯示已賦值的變數，灰色區塊表示無內容的變數。右側：收集的轉換基於 RL 訓練過程的標準化格式，每個轉換對應一個包含提示（prompt）、結果和即時獎勵的 LLM 步驟。

分層強化學習

對於需要多次 LLM 請求的代理，傳統的 RL 訓練涉及將所有內容拼接成一個長序列，然後在訓練期間識別哪些部分應該學習、哪些應該忽略。這種方法難以實施，且可能產生過長的序列，進而降低模型性能。

相反地，Agent Lightning 的 LightningRL 演算法採用分層方法。在任務完成後，信用分配（credit assignment）模組會確定每個 LLM 請求對結果的貢獻程度，並分配相應的獎勵。這些獨立的步驟現在配有各自的獎勵分數，可以與任何現有的單步 RL 演算法配合使用，例如近端策略優化（PPO）或群體相對策略優化（GRPO）（圖 2）。

圖 2. (a) 單步 GRPO：LLM 在一次調用中完成任務。比較同一任務的多個回應，以確定每個回應應被強化的程度。(b) 先前的多步 GRPO：任務涉及多次 LLM 調用。比較同一任務的多個多步運行，訓練期間忽略非 LLM 生成的標記（灰色框）。(c) LightningRL：將多步運行分解為單個 LLM 調用。比較來自同一任務的調用，以確定每個調用應被強化的程度。每次調用包含其輸入、上下文、輸出以及由信用分配模組分配的獎勵。

這種設計提供了多項優勢。它與廣泛使用的單步 RL 演算法完全相容，允許在不修改的情況下應用現有的訓練方法。將數據組織為一系列獨立的轉換，讓開發者可以根據需要靈活構建 LLM 輸入，支援如使用多個工具或與其他代理協作等複雜行為。此外，透過保持序列簡短，該方法可以清晰地擴展並保持訓練效率。

Agent Lightning 作為中間件

Agent Lightning 作為 RL 演算法與代理環境之間的中間件，提供模組化組件，透過標準化協議和定義良好的接口實現可擴展的 RL。

代理運行器（agent runner） 負責管理代理完成任務。它分配工作並收集、儲存結果與進度數據。它的運作與 LLM 分離，使它們能在不同的資源上運行，並擴展以支援多個代理同時運行。

演算法（algorithm） 負責訓練模型並託管用於推理和訓練的 LLM。它協調整個 RL 週期，管理分配哪些任務、代理如何完成任務，以及如何根據代理學到的內容更新模型。它通常在 GPU 資源上運行，並透過共享協議與代理運行器通訊。

LightningStore (在新分頁中開啟) 作為系統內所有數據交換的中央存儲庫。它提供標準化接口和共享格式，確保不同組件可以協同工作，並使演算法與代理運行器能夠有效溝通。

圖 3. Agent Lightning 框架

所有 RL 週期遵循兩個步驟：(1) Agent Lightning 收集代理執行數據（稱為「spans」）並將其儲存在數據庫中；(2) 隨後檢索所需數據並將其發送至演算法進行訓練。透過這種設計，演算法可以異步地將任務委派給代理運行器，後者完成任務並回報結果（圖 4）。

圖 4. Agent Lightning 的 RL 週期

這種方法的一個關鍵優勢是演算法的靈活性。該系統讓開發者能輕鬆自定義代理的學習方式，無論是定義不同的獎勵、捕捉中間數據，還是嘗試不同的訓練方法。

另一個優勢是資源效率。代理式 RL 系統非常複雜，整合了代理系統、LLM 推理引擎和訓練框架。透過分離這些組件，Agent Lightning 使這種複雜性變得可控，並允許每個部分獨立優化。

解耦設計允許每個組件使用最適合它的硬體。代理運行器可以使用 CPU，而模型訓練則使用 GPU。每個組件也可以獨立擴展，提高效率並使系統更易於維護。在實踐中，開發者可以保留現有的代理框架，並將模型調用切換到 Agent Lightning API，而無需更改其代理程式碼（圖 5）。

圖 5. 左側為開發者實現的代理程式碼。右下角為 Agent Lightning 所需的程式碼。代理程式碼的主體保持不變。

三個真實場景的評估

Agent Lightning 在三個不同的任務上進行了測試，在所有場景中均實現了持續的性能提升（圖 6）：

Text-to-SQL (LangChain)： 在一個包含三個代理（負責 SQL 生成、檢查和重寫）的系統中，Agent Lightning 同時優化了其中兩個代理，顯著提高了從自然語言查詢生成可執行 SQL 的準確性。

檢索增強生成 (OpenAI Agents SDK 實現)： 在多跳問答數據集 MuSiQue（需要查詢大型維基百科數據庫）上，Agent Lightning 幫助代理生成更有效的搜索查詢，並從檢索到的內容中進行更好的推理。

數學問答與工具使用 (AutoGen 實現)： 對於複雜的數學問題，Agent Lightning 訓練 LLM 更準確地確定何時以及如何調用工具，並將結果整合到其推理中，從而提高了準確性。

圖 6. 三個評估場景的獎勵曲線

實現代理的持續改進

透過簡化 RL 的整合，Agent Lightning 可以讓開發者更輕鬆地構建、迭代和部署高性能代理。我們計劃擴展 Agent Lightning 的功能，包括自動提示詞優化（prompt optimization）和更多的 RL 演算法。

該框架旨在作為一個開放平台，讓任何 AI 代理都能透過現實世界的實踐不斷進步。透過將現有的代理系統與強化學習橋接，Agent Lightning 旨在幫助創建能夠從經驗中學習並隨時間改進的 AI 系統。

（在新分頁中開啟）本貼文 Agent Lightning: Adding reinforcement learning to AI agents without code rewrites 首發於 Microsoft Research。