端到端優化所有能力，字節跳動提出強化學習LLM Agent框架AGILE
大語言模型（Large Language Models, LLMs）的強大能力推動了 LLM Agent 的迅速發展。圍繞增強 LLM Agent 的能力，近期相關研究提出了若干關鍵組件或工作流。然而，如何將核心要素集成到一個統一的框架中，能夠進行端到端優化，仍然是一個亟待解決的問題。
來自字節跳動 ByteDance Research 的研究人員提出了基于強化學習（Reinforcement Learning, RL）的 LLM Agent 框架 ——AGILE。該框架下，Agent 能夠擁有記憶、工具使用、規劃、反思、與外界環境交互、主動求助專家等多種能力，并且通過強化學習實現所有能力的端到端訓練。尤其值得注意的是，AGILE 框架允許 Agent 在不自信時主動向人類專家尋求建議。這帶來了兩大優勢：首先，Agent 在處理復雜問題時能夠持續保持高準確率；其次，通過向人類學習，增強了其快速適應新任務的泛化能力。
- 論文地址：https://arxiv.org/abs/2405.14751
- 代碼地址：https://github.com/bytarnish/AGILE
研究者讓 AGILE 框架在復雜問答任務中進行了驗證。在 ProductQA 與 MedMCQA 任務上，經過 RL 訓練后的 13B 模型或者 7B 模型能夠超越提示工程構建的 GPT-4 Agent。
AGILE Agent 框架
如圖 1 (a) 所示，AGILE 框架包含四個核心模塊：LLM、記憶（Memory）、工具（Tools）和執行器（Executor）。LLM 負責預測動作；記憶模塊記錄 Agent 的軌跡；工具模塊提供 Agent 可以調用的外部 API；執行器會監聽 LLM 的輸出，根據 LLM 的輸出調用相應的模塊執行具體的操作，并將執行操作得到的結果添加到 LLM 的上下文中。
圖 1 (b) 展示了 AGILE Agent 在電商問答場景中的一個示例。電商問答是一個復雜的實際應用場景，面臨的挑戰包括需要海量商品的領域知識、靈活運用商品檢索工具、以及快速適應不斷涌現的新商品。如圖 1 (b) 所示，AGILE Agent 會根據用戶的問題檢索記憶，如果無法確定問題的答案，Agent 會向人類專家尋求幫助。在獲得專家的反饋后，Agent 會反思并將新掌握的領域知識存儲在記憶中。在未來，面對新的問題時，Agent 能夠從記憶中檢索到這條知識，并基于這些知識直接給出準確的答案。除此之外，AGILE Agent 也會根據用戶的問題選擇是否調用外部工具（如搜索、數據庫檢索），輔助生成最終的回答。
強化學習定義：LLM Agents 被定義為一個 token-level MDP（Markov Decision Process）。動作空間（Action space）由 LLM 的詞表構成，LLM 生成的每一個 token 是一個動作，LLM 本身則作為 Agent 的策略模型（Policy model）。Agent 的狀態（State）由 LLM 上下文和記憶組成。在每個時刻，LLM 預測動作，執行器根據預定義的邏輯完成狀態轉移，同時環境給予 Agent 相應的獎勵（Reward）。
在 AGILE 框架下，Agent 有兩種策略學習方法。第一種是模仿學習，通過收集人類軌跡數據或更高級別 Agent 的軌跡數據，對 LLM 進行 SFT 訓練。第二種是強化學習，通過定義獎勵函數，利用強化學習算法來訓練 LLM。
此外，LLM Agent 可能會產生跨越數百萬個 tokens 的軌跡，這為直接訓練帶來了挑戰。為了應對這種長程依賴的問題，研究人員提出了一種片段級別的優化算法。
主動尋求幫助：AGILE 框架允許 Agent 主動向外部的人類專家尋求幫助。這種機制有兩個優勢：首先，當 Agent 遇到不確定的情況時，通過求助人類專家，確保其在實際應用中達到高準確率。其次，Agent 能夠通過對人類的反饋反思并積累知識，從而更快適應新環境，提升其泛化能力。然而決定何時尋求幫助是一個復雜決策，它涉及到 Agent 的自我評估、人類反饋對未來的價值以及人類專家的成本。因此，標注何時應該求助是很難的。但在強化學習框架中，可以通過定義相關獎勵，將這種求助能力作為策略模型的一部分，在端到端訓練中得到提升。
實驗結果
ProductQA
ProductQA 是一個商品問答任務。該數據集包含 26 個對應不同商品類別的 QA 任務，每個任務平均包含 3,393 個問題。ProductQA 包括基于事實的問題、推理問題和商品推薦問題，它能夠全面評估 Agent 處理歷史信息和累積知識、利用工具、向人求助、自我評估和反思的能力。此外，訓練和測試集由不同的任務構成，以評估 Agent 在新商品問答上的泛化能力。
在商品問答（ProductQA）任務上，基于 Vicuna-13b 訓練的 AGILE Agent（agile-vic13b-ppo）表現超過了 GPT-4（gpt4-prompt）與提升工程構建的 GPT-4 Agent（agile-gpt4-prompt）。在使用了相當的求助比例（Advice Rate）的情況下，agile-vic13b-ppo 的 acc 相比于 agile-gpt4-prompt 提升了 7.4%，在 Total Score 上提升了 9.2%。
從上圖可以看出，在執行包含上千個問答的任務整個過程中，agile-vic13b-ppo 的 acc 持續穩定地高于 agile-gpt4-prompt。同時尋求人類幫助的頻率（Advice Rate）隨著問答輪數的增加逐漸下降。此外，通過調整人類的咨詢成本（Seeking Advice Cost）和進行強化學習訓練，AGILE Agent 可以有效實現準確率與專家成本的權衡。
消融實驗結果顯示，記憶、反思、咨詢人類建議、工具使用、RL 訓練在實現高性能的 AGILE Agent 中均具有重要作用。
MedMCQA
MedMCQA 是一個多項選擇的問答數據集，其問題來自醫學院入學考試。在 MedMCQA 任務上，基于 Meerkat-7b 訓練的 AGILE Agent（agile-mek7b-ppo）表現優于 GPT-4 Agent。準確率達到了 85.2%，超過了之前的 SOTA 方法 ——GPT 4-MedPrompt 的 79.1%。消融實驗結果進一步驗證了記憶、咨詢人類建議、反思和 RL 訓練的重要性。
更多研究細節，可參考原論文。
總結
AGILE是一種基于強化學習的LLM Agent框架。AGILE Agent具備擁有長期記憶、使用外部工具、向人類專家進行咨詢、反思等能力，并且所有能力可以進行端到端的統一優化。AGILE的一個重要特點是Agent在遇到無法自行解決的問題時可以向人類專家咨詢，這種機制保證了應用場景中對高準確率的要求，同時增強了Agent的學習與泛化能力。實驗表明，經過強化學習訓練的13B或7B模型的Agent，其能力可以超過GPT-4 Agent。
團隊介紹
ByteDance Research 專注于人工智能領域的前沿技術研究，涵蓋了機器翻譯、視頻生成基礎模型、機器人研究、機器學習公平性、量子化學、AI 制藥、分子動力學等多技術研究領域，同時致力于將研究成果落地，為公司現有的產品和業務提供核心技術支持和服務。
