Agent Lightning: コードを書き換えずに AI エージェントに強化学習を追加

AI エージェントは、コードの作成から複雑な命令の実行に至るまで、ソフトウェア開発を再構築しています。しかし、LLM ベースのエージェントはエラーが発生しやすく、複雑な複数ステップのタスクではパフォーマンスが低下することがよくあります。強化学習 (RL) は、AI システムが自分の行動に対して報酬またはペナルティを受け取ることで最適な意思決定を行うことを学習し、試行錯誤を通じて改善するアプローチです。 RL はエージェントの改善に役立ちますが、通常、開発者はコードを大幅に書き直す必要があります。これらのエージェントが生成するデータが RL トレーニングを通じてパフォーマンスを大幅に向上させる可能性があるにもかかわらず、これにより導入が妨げられます。

これに対処するために、からの研究チームは、 マイクロソフト リサーチ アジア – 上海 導入しました エージェント ライトニング。これ オープンソース (新しいタブで開きます) このフレームワークは、エージェントがタスクを実行する方法をモデルのトレーニングから分離することで、RL を通じて AI エージェントをトレーニングできるようにし、開発者がコードを実質的に変更せずに RL 機能を追加できるようにします。

トレーニングのためにエージェントの行動をキャプチャする

Agent Lightning は、エージェントの実行を一連の状態とアクションとして扱うことにより、エージェントのエクスペリエンスを RL が使用できる形式に変換します。各状態はエージェントのステータスを取得し、各 LLM 呼び出しはエージェントを新しい状態に移行するアクションです。

このアプローチは、どんなに複雑であっても、どのようなワークフローでも機能します。複数のエージェントの連携や動的なツールの使用に関係なく、Agent Lightning はそれを一連の遷移に分割します。各遷移は、LLM の入力、出力、報酬をキャプチャします (図 1)。この標準化された形式は、追加の手順を行わずにデータをトレーニングに使用できることを意味します。

階層型強化学習

複数の LLM リクエストを行うエージェントに対する従来の RL トレーニングでは、すべてのコンテンツを 1 つの長いシーケンスにつなぎ合わせ、トレーニング中にどの部分を学習する必要があり、どの部分を無視するかを特定する必要があります。このアプローチは実装が難しく、モデルのパフォーマンスを低下させる過度に長いシーケンスが作成される可能性があります。

代わりに、Agent Lightning の LightningRL アルゴリズムは階層的なアプローチを採用しています。タスクが完了すると、クレジット割り当てモジュールが各 LLM リクエストが結果にどの程度貢献したかを判断し、それに対応する報酬を割り当てます。これらの独立したステップは、独自の報酬スコアとペアになっており、近接ポリシー最適化 (PPO) やグループ相対ポリシー最適化 (GRPO) などの既存のシングルステップ RL アルゴリズムで使用できます (図 2)。

この設計にはいくつかの利点があります。広く使用されているシングルステップ RL アルゴリズムと完全な互換性を維持しているため、既存のトレーニング方法を変更せずに適用できます。データを独立したトランジションのシーケンスとして整理することで、開発者は必要に応じて LLM 入力を柔軟に構築でき、複数のツールを使用するエージェントや他のエージェントと連携するエージェントなどの複雑な動作をサポートできます。さらに、シーケンスを短く保つことで、アプローチが適切に拡張され、トレーニングの効率が維持されます。

ミドルウェアとしての Agent Lightning

Agent Lightning は、RL アルゴリズムとエージェント環境の間のミドルウェアとして機能し、標準化されたプロトコルと明確に定義されたインターフェイスを通じてスケーラブルな RL を可能にするモジュール式コンポーネントを提供します。

アン エージェントランナー タスクを完了するエージェントを管理します。作業を分散し、結果と進捗データを収集して保存します。 LLM とは別に動作するため、LLM を異なるリソース上で実行し、同時に実行される複数のエージェントをサポートするように拡張できます。

アン アルゴリズム モデルをトレーニングし、推論とトレーニングに使用される LLM をホストします。これは、RL サイクル全体を調整し、どのタスクが割り当てられるか、エージェントがタスクを完了する方法、エージェントが学習した内容に基づいてモデルがどのように更新されるかを管理します。通常、GPU リソース上で実行され、共有プロトコルを通じてエージェント ランナーと通信します。

の ライトニングストア (新しいタブで開きます) システム内のすべてのデータ交換の中央リポジトリとして機能します。標準化されたインターフェイスと共有フォーマットを提供し、さまざまなコンポーネントが連携して動作できるようにし、アルゴリズムとエージェント ランナーが効果的に通信できるようにします。

すべての RL サイクルは 2 つのステップに従います。(1) Agent Lightning はエージェント実行データ (「スパン」と呼ばれる) を収集し、データ ストアに保存します。 (2) 次に、必要なデータを取得し、トレーニングのためにアルゴリズムに送信します。この設計により、アルゴリズムはタスクを非同期にエージェント ランナーに委任することができ、エージェント ランナーはタスクを完了して結果を報告します (図 4)。

このアプローチの重要な利点の 1 つは、アルゴリズムの柔軟性です。このシステムにより、開発者はさまざまな報酬の定義、中間データのキャプチャ、さまざまなトレーニング アプローチの実験など、エージェントの学習方法を簡単にカスタマイズできます。

もう 1 つの利点は、リソース効率です。エージェント RL システムは複雑で、エージェント システム、LLM 推論エンジン、トレーニング フレームワークが統合されています。これらのコンポーネントを分離することで、Agent Lightning はこの複雑さを管理しやすくし、各部分を個別に最適化できるようにします。

分離された設計により、各コンポーネントが最適なハードウェアを使用できるようになります。エージェント ランナーは CPU を使用できますが、モデル トレーニングでは GPU が使用されます。各コンポーネントは個別に拡張することもできるため、効率が向上し、システムの保守が容易になります。実際には、開発者はエージェント コードを変更せずに、既存のエージェント フレームワークを維持し、モデル呼び出しをエージェント Lightning API に切り替えることができます (図 5)。

3 つの現実世界のシナリオにわたる評価

Agent Lightning は 3 つの異なるタスクでテストされ、すべてのシナリオで一貫したパフォーマンスの向上が達成されました (図 6)。

テキストから SQL (LangChain): SQL の生成、チェック、再書き込みを処理する 3 つのエージェントを備えたシステムにおいて、Agent Lightning はそのうちの 2 つを同時に最適化し、自然言語クエリから実行可能な SQL を生成する精度を大幅に向上させました。

取得拡張生成 (OpenAI Agents SDK 実装): 大規模な Wikipedia データベースへのクエリを必要とするマルチホップ質問応答データセット MuSiQue では、エージェント ライトニングは、エージェントがより効果的な検索クエリを生成し、取得したコンテンツからより適切に推論できるように支援しました。

数学的 QA とツールの使用 (AutoGen 実装): 複雑な数学の問題に対して、エージェント ライトニングは LLM をトレーニングして、ツールを呼び出すタイミングと方法をより正確に決定し、その結果を推論に統合して精度を向上させました。

エージェントの継続的な改善を可能にする

Agent Lightning は RL 統合を簡素化することで、開発者が高性能エージェントを構築、反復、デプロイすることを容易にします。 Agent Lightning の機能を拡張して、自動プロンプト最適化と追加の RL アルゴリズムを含める予定です。

このフレームワークは、AI エージェントが実際の実践を通じて改善できるオープン プラットフォームとして機能するように設計されています。 Agent Lightning は、既存のエージェント システムと強化学習を橋渡しすることで、経験から学習し、時間の経過とともに改善する AI システムの作成を支援することを目指しています。

新しいタブで開きます