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)

所属 Microsoft Research

投稿日 2025年08月07日

カテゴリ cs.AI, cs.LG

arXiv 2508.03680 ↗

Agent Lightning: Train ANY AI Agents with Reinforcement Learning

基本情報

arXiv ID: 2508.03680 (https://arxiv.org/abs/2508.03680)
著者: Xufang Luo, Yuge Zhang, Zhiyuan He, Zilong Wang, Siyun Zhao, Dongsheng Li, Luna K. Qiu, Yuqing Yang (Microsoft Research)
所属: Microsoft Research
投稿日: 2025年08月07日
カテゴリ: cs.AI, cs.LG

簡単に説明すると

Agent LightningはAIエージェントを強化学習（RL）で訓練するための革新的なフレームワークです。従来のRLフレームワークとは異なり、エージェントの実行と訓練を完全に分離することで、LangChain、OpenAI Agents SDK、AutoGenなど、どのようなフレームワークで作られたエージェントでも、ほぼコード変更なしにRLで最適化できます。GitHubリポジトリ（https://github.com/microsoft/agent-lightning/tree/main/examples/apo）で実装の詳細が公開されています。

1. 研究概要

1.1 背景と動機

近年、大規模言語モデル（LLM）の進歩により、AIエージェントは検索、コード生成、ツール利用などの複雑なタスクで有効性を示しています。しかし、LLMは訓練されていないシナリオ（マルチターンのコーディングワークフロー、プライベートドメインのデータセット、なじみのないツールなど）でエラーを起こしやすく、エンドツーエンドのソフトウェア開発のような現実世界のタスクを確実に解決することが困難です。

エージェント実行中に生成される豊富な相互作用データは、現実世界の問題解決の複雑さを捉えており、人間が作成した従来のデータセットを規模と多様性の両面で上回ります。このデータを活用してファインチューニングすることで、エージェントの専門的なスキルを洗練させるだけでなく、動的でインタラクティブな環境に適したより汎用的なLLMの開発を促進できます。

しかし、既存のRL手法とフレームワークは主に静的な単一呼び出しタスク（選好アライメントや数学的推論など）向けに設計されており、エージェントの複雑性と多様性に対応できません。エージェントの実行は複数のLLM呼び出し、外部ツール、API、環境との相互作用を含み、さまざまなアプリケーションの要件に合わせて異なるエージェントを設計する必要があります。

1.2 主要な貢献

本研究の主要な貢献は以下の3点にまとめられます。

エージェントとRL訓練の完全な分離を実現する初のフレームワークの提案
エージェント実行をマルコフ決定過程（MDP）として定式化し、統一データインターフェースとLightningRLという階層的RLアルゴリズムの開発
Training-Agent Disaggregationアーキテクチャの導入による、スケーラブルで拡張可能なシステム設計

2. 提案手法

2.1 手法の概要

Agent Lightningは、エージェント実行とRL訓練を完全に分離することで、既存のエージェントをほぼコード変更なしに最適化可能にします。エージェント実行をMDPとして定式化し、状態（state）を実行のスナップショット、アクション（action）をポリシーLLMが生成する出力として定義します。

統一データインターフェースでは、エージェントの軌跡を遷移（transition）のシーケンスとして構造化し、各遷移は現在の状態（LLM入力）、アクション（LLM出力）、報酬を含みます。このインターフェースは、基盤となるオーケストレーションロジックやエージェントフレームワークの詳細を抽象化し、あらゆるエージェントに適用可能です。

2.2 技術的詳細

提案されたLightningRLアルゴリズムは、2段階のメカニズムを採用しています。まず、エピソードレベルのリターンをクレジット割り当てモジュールによって各アクションに割り当て、次に各アクション内のトークンにさらに分解してトークンレベルの監視信号を生成します。

システム設計では、Training-Agent Disaggregationアーキテクチャを導入し、計算集約的なLLM生成と軽量だが多様で柔軟なアプリケーションロジック・ツールを分離しています。Lightning ServerとLightning Clientの2つのコンポーネントで構成され、ServerはRL訓練システムのコントローラーとして機能し、Clientはエージェントランタイムとしてエージェントをカプセル化します。

エージェントランタイムは、OpenTelemetryなどの包括的な観測可能性フレームワークを活用して軌跡収集を行い、Automatic Intermediate Rewarding（AIR）メカニズムによって中間報酬の割り当てを可能にします。

2.3 新規性

既存手法との主な違いは以下の点です：

連結とマスキングに依存する従来の手法と異なり、個々の遷移としてデータを整理することで、複雑な相互作用ロジックを柔軟に扱える
エージェント実行と訓練の完全な分離により、多様なエージェントフレームワークとシームレスに統合可能
累積的なコンテキスト長の問題を軽減し、実装の複雑さを削減

3. 実験結果

3.1 実験設定

3つの異なるタスクで評価を行いました：

Text-to-SQL（LangChain実装）: Spider データセット、3エージェントシステム（SQLライター、チェッカー、リライター）
Retrieval-Augmented Generation（OpenAI Agents SDK実装）: MuSiQueデータセット、2100万文書のWikipedia検索
Math QA with Tool Usage（AutoGen実装）: Calc-Xデータセット、計算機ツール使用

いずれもLlama-3.2-3B-Instructをベースモデルとして使用しました。

3.2 主要な結果

すべてのタスクにおいて、Agent Lightningは安定した継続的な性能向上を示しました。Text-to-SQLタスクでは複雑な多段階決定とコード生成・ツール使用の最適化能力を実証し、RAGタスクでは挑戦的なマルチホップ質問応答での有効性を示し、Math QAタスクではツール拡張設定での一貫した改善を達成しました。

3.3 既存手法との比較

従来の連結ベースの手法と比較して、Agent Lightningの遷移ベースのアプローチは以下の利点を示しました：

より広範囲のエージェントアーキテクチャとワークフローのサポート
累積的なコンテキスト長による問題の軽減
カスタムマスキング戦略の不要化による実装の簡素化
階層的RLアルゴリズムなど、より高度なRL手法への拡張可能性

4. 実用性評価

4.1 実装の容易性

Agent Lightningの最大の利点は、既存のエージェントコードをほぼ変更せずに最適化できることです。附録のコード例で示されているように、開発者は簡単な設定でエージェントの最適化を開始できます。また、LangChain、OpenAI Agents SDK、AutoGenなど、主要なエージェント開発フレームワークすべてをサポートしています。

4.2 計算効率

Training-Agent Disaggregationアーキテクチャにより、GPU リソースと従来のプログラミング言語で書かれたアプリケーションロジックを効率的に分離できます。データ並列処理により、大規模バッチサイズでの効率的な訓練が可能で、ノード内およびノード間の並列処理により柔軟なスケーリングが実現されています。

4.3 応用可能性

Agent Lightningは特定のタスクやシナリオに限定されず、あらゆるAIエージェントに適用可能です。また、RLベースの手法以外にも、自動プロンプト最適化などの他の最適化手法もサポートする拡張性があります。将来的には、より効率的なRLアルゴリズム（長期的なクレジット割り当て、探索アルゴリズム、オフポリシーアルゴリズムなど）の統合も可能です。

5. まとめと所感

5.1 論文の意義

Agent Lightningは、AIエージェントの訓練における重要なブレークスルーを達成しています。エージェント実行と訓練の完全な分離という革新的なアプローチにより、実世界のエージェント開発と最適化の間のギャップを埋めることに成功しています。これにより、開発者は既存のエージェントを容易に改善でき、より適応的で学習可能なAIシステムの実現に向けた重要な一歩となっています。

特に注目すべきは、この研究がDeepSeek-R1やKimi k1.5などの最新の推論モデルで成功したRL手法を、より広範なエージェントシナリオに拡張している点です。統一データインターフェースとMDP定式化により、複雑で動的なエージェントの振る舞いを扱える柔軟性を提供しています。

5.2 今後の展望

今後の研究方向として、より洗練されたクレジット割り当て戦略の開発、効率的なサービング手法の探索、そしてRLシステムインフラストラクチャのさらなる改善が挙げられます。また、Component of Interest（CoI）の概念を活用した、より包括的な最適化手法の開発も期待されます。長期的には、このフレームワークが実世界のエージェント展開における標準的なツールとなり、より高度で適応的なAIシステムの開発を加速することが期待されます。