⚡ 技術調査 - Agent Lightning

はじめに

LangChain、AutoGen、OpenAI Assistants API などのフレームワークにより、自律型 AI エージェントの構築は容易になりました。一方で、構築したエージェントの性能を体系的に向上させる「最適化」は、手動のプロンプトエンジニアリングに依存しているのが現状です。

Microsoft Research が開発した Agent Lightning は、この課題を解決するオープンソースフレームワークです。エージェントの「実行」と「学習」を構造的に分離し、コード変更をほぼゼロに抑えながら、強化学習（RL）や自動プロンプト最適化（APO）を適用できます。

本記事では、Agent Lightning のアーキテクチャ、データモデル、環境構築、実装、運用までを体系的に解説します。

対象読者: AI エージェント開発の経験があり、エージェントの性能最適化に関心のあるエンジニア
前提知識: Python、LLM の基本概念、強化学習の基礎（報酬・エピソード等）

エージェント最適化の課題

従来の機械学習パイプラインは、モデルと学習ループが密結合した「モノリシック」な構造を前提としています。エージェントシステムでは、この前提が成り立ちません。

既存のエージェントフレームワークとの違いを整理します。

Training-Agent Disaggregation アーキテクチャ

Agent Lightning は Training-Agent Disaggregation というアーキテクチャを採用しています。エージェントの「実行」と「学習」を明確に分離し、非同期に連携させる設計です。

Runner: 実行者

Runner はエージェントを実行するワーカープロセスです。学習アルゴリズムの詳細を知る必要はなく、タスクをエージェントに処理させ、結果（軌跡と報酬）を記録します。エージェントのネイティブ環境（Docker コンテナやローカルマシン）で動作できるため、複雑な環境設定を学習サーバーから切り離せます。

Algorithm: 学習者

Algorithm はシステムの「頭脳」です。Runner から収集されたトレースを分析し、エージェントのリソース（プロンプトテンプレートやモデルの重み）を更新します。エージェントの実装詳細から完全に独立しており、RL、APO、SFT など多様なアルゴリズムを同一エージェントに適用できます。

主なアルゴリズムは以下のとおりです。

LightningStore: 仲介者

Runner と Algorithm の間に直接の通信経路は存在しません。すべてのデータ交換は LightningStore を介して行われます。タスクキュー、結果ストレージ、リソースのバージョン管理を兼ね備えた中央データストアです。

主な REST API エンドポイントは以下のとおりです。

この設計により、Runner がクラッシュしても Algorithm は影響を受けず、Algorithm が停止しても Runner はタスクを処理し続けられます。

LLM Proxy: 推論の仲介

RL トレーニングでは、生成トークンの対数確率（logprobs）や温度パラメータの制御が必要です。Agent Lightning は LLM Proxy を提供し、エージェントと LLM の間に介在します。

最適化ループのデータフロー

Agent Lightning の最適化プロセスは、以下の循環的なデータフローで進行します。

このアーキテクチャの最大の利点は、エージェントを「入力を受け取り、トレースと報酬を出力するブラックボックス」として抽象化できる点です。インターフェースさえ適合させれば、あらゆるエージェントを最適化対象にできます。

データモデル: 統一データインターフェース

Agent Lightning は、エージェントの挙動を最適化アルゴリズムが扱える形式に変換するため、厳密なデータモデル（Unified Data Interface）を定義しています。

コアデータプリミティブ

Task

エージェントに与える入力と正解データを含む最小単位です。従来の RL における「環境のリセット」と「初期状態の提供」に相当します。

class RoomSelectionTask(TypedDict):
    input: str          # "4人用の会議室を10時に予約して"
    expected_choice: str # "会議室B"
    constraints: list    # ["プロジェクターあり", "窓側"]

Rollout

1つの Task に対するエージェントの完全な実行履歴を管理するオブジェクトです。タスク受領から報酬受領までの一連の流れ（エピソード）全体を表します。ステータスは QUEUING → RUNNING → SUCCEEDED / FAILED と遷移します。

Attempt

エージェントは確率的に動作し、外部要因で失敗する可能性があります。1つの Rollout 内に複数の Attempt が含まれる構造により、再試行ロジックをデータ構造レベルでサポートしています。API エラーで1回目が失敗した場合、システムは自動的に2回目の Attempt を作成します。

Span

OpenTelemetry 仕様に準拠した実行トレースの最小単位です。親子関係を持つツリー構造（Trace Tree）を形成し、因果関係を復元できます。

Reward Span により、長い推論プロセスの途中経過に報酬を与える「密な報酬」設計が可能です。従来の「疎な報酬」（エピソード終了時のみ）に比べ、学習効率が向上します。例えば、「正しいツールを選択できたか」という中間ステップに報酬を付与できます。

Resource

アルゴリズムによって更新・最適化される資産です。すべてイミュータブルなバージョンとして保存され、Rollout ごとに使用バージョンが紐付けられます。これにより実験の再現性が担保されます。

Trace から Trajectory への変換

階層的な Span Tree を、RL アルゴリズムが理解できる線形な学習データに変換するのが Adapter 層です。

この変換により、LangChain や AutoGen など異種混合のエージェントを統一的な MDP（マルコフ決定過程）として抽象化し、あらゆるエージェントに RL アルゴリズムを適用できます。

環境構築

前提条件

• Python 3.10 以上
• GPU（RL トレーニング時）: NVIDIA GPU + CUDA 対応環境
• Docker（LightningStore の本番運用時）

インストール

Agent Lightning は PyPI から配布されています。

# 基本パッケージ
pip install agent-lightning

# vLLM サポートを含むインストール
pip install agent-lightning[vllm]

# 開発用ツールを含むインストール
pip install agent-lightning[dev]

LightningStore の立ち上げ

Docker による実運用環境

最も推奨される方法は、提供されている Docker Compose ファイルの使用です。

cd /path/to/agent-lightning
docker compose -f docker/compose.store.yml up -d

このコマンドにより、http://localhost:4747 で Store API が利用可能になり、データは MongoDB に永続化されます。

CLI による簡易起動

一時的なテストやデバッグには、CLI コマンドでローカル起動できます。

# メモリ内バックエンド（プロセス終了でデータ消失）
agl store --port 4747 --backend memory

# MongoDB バックエンド
agl store --port 4747 --backend mongo --mongo-uri mongodb://localhost:27017

vLLM 推論サーバーの準備

RL では数千〜数万回のロールアウトが必要です。コストと速度の観点から、ローカル LLM の使用が一般的です。Agent Lightning は高速推論エンジン vLLM と統合されています。

agl vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
    --port 8000 \
    --host 0.0.0.0 \
    --tensor-parallel-size 1

Agent Lightning 経由で vLLM を起動すると、推論リクエストへの自動計装や学習に必要なメタデータの注入が可能になります。

エージェントの実装

エージェントを Agent Lightning で最適化するには、既存のコードを Agent Lightning のインターフェースでラップします。主に2つの方法があります。

方法 A: @rollout デコレータ

シンプルなエージェントには @rollout デコレータが最も手軽です。関数を自動的に FunctionalLitAgent に変換します。

import agentlightning as agl
from openai import OpenAI

# 最適化対象のプロンプトリソース
math_instruction = agl.PromptTemplate(
    template="あなたは数学の専門家です。次の問題を解いてください: {input}",
    name="math_instruction"
)

@agl.rollout
def math_solver_agent(task: agl.Task, resources: agl.NamedResources) -> float:
    prompt_template = resources.get("instruction", math_instruction)
    formatted_prompt = prompt_template.format(input=task.input)

    client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
    response = client.chat.completions.create(
        model="meta-llama/Meta-Llama-3-8B-Instruct",
        messages=[{"role": "user", "content": formatted_prompt}]
    )
    answer = response.choices[0].message.content

    reward = 1.0 if "42" in answer else 0.0
    return reward

方法 B: LitAgent クラス

LangChain や AutoGen など複雑なフレームワークを使用する場合、LitAgent クラスの継承が推奨されます。PyTorch Lightning の LightningModule に相当し、エージェントのライフサイクル全体を制御できます。

import agentlightning as agl
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

class SQLLitAgent(agl.LitAgent):
    def __init__(self, db_schema: str):
        super().__init__()
        self.db_schema = db_schema

    def rollout(self, task: agl.Task, resources: agl.NamedResources,
                rollout_meta: agl.Rollout) -> float:
        llm_config: agl.LLM = resources["main_llm"]
        llm = ChatOpenAI(
            base_url=llm_config.endpoint,
            api_key=llm_config.api_key,
            model=llm_config.model,
            temperature=llm_config.sampling_params.get("temperature", 0.0)
        )

        prompt = ChatPromptTemplate.from_template(
            "Schema: {schema}\nQuestion: {question}\nSQL Query:"
        )
        chain = prompt | llm

        inputs = {"schema": self.db_schema, "question": task.input}
        try:
            result = chain.invoke(
                inputs,
                config={"callbacks": [self.tracer.get_langchain_handler()]}
            )
            generated_sql = result.content
        except Exception as e:
            agl.emit_exception(e)
            return 0.0

        reward = self.evaluate_sql(generated_sql, task.ground_truth)
        agl.emit_reward(reward)
        return reward

計装（Instrumentation）

Agent Lightning の学習は、エージェントの実行ログ（トレース）の分析に基づいています。正確な記録が極めて重要です。

これらの自動計装は Python のメタプログラミング機能を活用しています。ライブラリのメソッドを動的にパッチ（Monkey Patching）することで、開発者が意識せずに内部通信を傍受します。これが「コード変更ほぼゼロ」の技術的根拠です。

トレーニングループの実行

エージェントの実装が完了したら、Trainer で最適化プロセスを実行します。

import agentlightning as agl
from agentlightning.algorithms import APO
from my_agent import SQLLitAgent

def main():
    # 1. Store への接続
    store = agl.LightningStoreClient("http://localhost:4747")

    # 2. LLM リソースの定義
    llm_resource = agl.LLM(
        model="meta-llama/Meta-Llama-3-8B-Instruct",
        endpoint="http://localhost:8000/v1",
        api_key="EMPTY"
    )

    # 3. アルゴリズムの設定
    algorithm = APO(
        llm=llm_resource,
        num_proposals=4,   # 1ラウンドあたりのプロンプト案の数
        generations=5       # 最適化ラウンド数
    )

    # 4. エージェントの初期化
    agent = SQLLitAgent(db_schema="...")

    # 5. Trainer の構築
    trainer = agl.Trainer(
        agent=agent,
        algorithm=algorithm,
        store=store,
        execution_strategy="client_server",
        default_resources={"main_llm": llm_resource}
    )

    # 6. トレーニング開始
    train_dataset = agl.Dataset.from_json("train_data.json")
    trainer.fit(train_dataset=train_dataset)

if __name__ == "__main__":
    main()

スクリプトを実行すると、以下のループが自動進行します。

1. Trainer がタスクを Store に登録
2. APO が初期プロンプトを生成
3. Runner がタスクを処理し、トレースと報酬を返却
4. Algorithm がフィードバックを基に次世代プロンプトを生成
5. 2〜4 を指定ラウンド数だけ繰り返し

実行戦略

Agent Lightning は開発フェーズに応じて2つの実行トポロジーを提供しています。

Shared-Memory 戦略

単一 Python プロセス内で Algorithm と Runner をスレッドとして実行します。データのシリアライズが不要でデバッグが容易です。プロトタイピングに最適ですが、Python の GIL 制約により大規模並列処理には向きません。

Client-Server 戦略

HTTP/gRPC 通信を介して物理的に異なるマシン間で連携します。Runner を Kubernetes クラスタや複数の GPU ノードに分散配置でき、環境依存の問題を回避できます。128 GPU 規模での RL トレーニング実績があります。

Ray による分散学習

大規模モデル（70B+ パラメータ）の RL では単一 GPU ではメモリ不足になります。Agent Lightning は分散処理フレームワーク Ray と統合されており、Algorithm 側の学習ステップを複数 GPU ノードに分散できます。

高度な最適化

カスタムアルゴリズム

Algorithm 基底クラスを継承し、2つのメソッドをオーバーライドすることで独自ロジックを実装できます。

class MyCustomAlgorithm(agl.Algorithm):
    def tell(self, rollouts: list[agl.Rollout]):
        # 報酬の集計、勾配計算、リプレイバッファへの追加など
        for rollout in rollouts:
            self.buffer.append(rollout)

    def ask(self, batch_size: int) -> list[agl.Task]:
        # 蓄積データを分析し、次のタスクと更新リソースを返却
        new_prompt = self.optimize_prompt(self.buffer)
        return self.generate_tasks(new_prompt, batch_size)

VERL との統合

Agent Lightning は、エージェント学習に特化した RL ライブラリ VERL（Volatile Experience Replay Learning）と連携します。Adapter が Store の OTel トレースデータを、RL アルゴリズムが理解できる形式（Prompt, Response, Reward のトリプレット）に変換します。変換されたデータが VERL の Trainer に渡され、PPO や GRPO でモデルの重みが更新されます。

エラー情報のデータ化

Agent Lightning ではエラーも重要な学習データとして扱います。Tracer が例外の種類（ToolTimeoutError、ContextLengthExceeded など）を Span の属性として記録し、報酬関数で特定エラーにペナルティを設定できます。エージェントは「ツール引数を正しく生成する」「コンテキスト長を超えない応答をする」といった行動を自律的に学習します。

運用とモニタリング

ダッシュボード

Agent Lightning には学習状況を可視化する Web ダッシュボードが付属しています。

Trace Viewer は、どのステップで時間がかかっているか、どこでエラーが発生したかを視覚的に特定でき、ハルシネーションやループのデバッグに有用です。

cd dashboard
npm ci
npm run build

デバッグ

学習が停滞した場合、Store のクエリ機能で失敗したロールアウトのみを抽出・分析できます。

failed_rollouts = store.query_rollouts(
    filter={"status": "FAILED"},
    limit=10
)

for rollout in failed_rollouts:
    print(f"Rollout ID: {rollout.rollout_id}")

トラブルシューティング

まとめ

Agent Lightning は、AI エージェント開発を手動のプロンプト調整から、データ駆動型の自動改善サイクルへ移行させるフレームワークです。

従来の「プロンプトを書いて試す」というサイクルに限界を感じている方にとって、Agent Lightning はエージェント最適化の自動化という新たな選択肢を提供します。特に APO によるプロンプト最適化は、GPU を必要とせず手軽に始められるため、最初の一歩として適しています。

次のステップ

1. 公式チュートリアルで最初のエージェントをトレーニング
2. @rollout デコレータで既存エージェントをラップし、APO による自動プロンプト最適化を試行
3. 効果が確認できたら、Client-Server 戦略に移行して本格的な RL トレーニングを実施

この記事が少しでも参考になった、あるいは改善点などがあれば、ぜひリアクションやコメント、SNSでのシェアをいただけると励みになります！

参考リンク

• 公式ドキュメント
  • Agent-lightning
  • Bird's Eye View - Agent-lightning
  • Understanding Store - Agent-lightning
  • Train the First Agent - Agent-lightning
  • Write the First Algorithm - Agent-lightning
  • Examples Catalog - Agent-lightning
  • Work with Traces - Agent-lightning
  • Write Agents - Agent-lightning
  • Installation - Agent-lightning
  • Debugging - Agent-lightning
  • Command Line - Agent-lightning
  • Types - Agent-lightning
  • Agent - Agent-lightning
  • Store - Agent-lightning
  • Algorithm - Agent-lightning
  • Agent Lightning Trainer
• GitHub
  • microsoft/agent-lightning
• 論文・研究
  • Agent Lightning - Microsoft Research
  • Agent Lightning: Train ANY AI Agents with Reinforcement Learning - arXiv
  • Agent Lightning: Train ANY AI Agents with Reinforcement Learning - Hugging Face
• 記事
  • Deep dive into the Training-Agent Disaggregation architecture of Agent Lightning - Medium
  • Tracing + Prompt Optimization for a Tool-Using SQL Agent - Medium
  • Train Your AI Agents with Microsoft Agent Lightning - Analytics Vidhya
  • Microsoft Agent Lightning - Overview and Design Principles - Stackademic
  • Get Started with Distributed Training using PyTorch Lightning - Ray Docs
  • Monitor agents with the Agent Monitoring Dashboard - Microsoft Foundry