Amazon SageMaker AIにおけるマルチターン強化学習のベストプラクティス
この記事では、Amazon SageMaker AIで信頼性の高いマルチターン強化学習トレーニングを実施するためのベストプラクティスを紹介します。信頼できるトレーニング環境の構築、外部評価の設定、最終タスクに合わせた報酬設計、エージェントが複数ターン実行した際の変化の管理、反復のタイミングを示すメトリクスの監視について説明します。
Amazon SageMaker AIでマルチターンエージェントをトレーニングしてサポートチケットの解決やコンテンツモデレーションを行う場合、単一の応答ではなく、一連の依存ステップを処理する必要があります。エージェントは指示を読み、ツール呼び出しを行い、結果を読み取り、次のアクションを決定し、最終回答を確定する前にエラーから回復します。この柔軟性こそが、エージェント強化学習（RL）を困難にしています。行動の選択肢が増えるということは、タスクを実行せずに報酬を満たす方法も増えることを意味し、エージェントがトレーニングする環境は静かにトレーニング信号を汚染する可能性があります。
この記事では、信頼性の高いマルチターンRLトレーニングのベストプラクティスを共有します。信頼できるトレーニング環境の構築方法、外部評価の設定方法、最終タスクに合わせた報酬の設計方法、エージェントが複数ターン実行した際の変化の管理方法、反復のタイミングを示すメトリクスの監視方法について説明します。例として、12のビジネスドメインにわたる複雑な標準運用手順（SOP）に基づいてエージェントのタスク解決能力を評価するAmazon ScienceのベンチマークであるSOP-Benchデータセットを使用します。
SageMaker AIマルチターン強化学習（SageMaker AI MTRL）は、エージェントタスクのトレーニングループを提供します。エージェントはAmazon Bedrock AgentCore、Amazon EKS、Amazon EC2、AWS Fargate、またはユーザーが選択したインフラストラクチャ上で実行できます。小さなアダプターを介してツールサーフェスをロールアウトサーバーに公開し、SageMaker AI MTRLが残りを処理します。モジュール式のエージェント-環境インターフェースにより、低コードで統合を維持しながら完全なアルゴリズム制御を実現します。カスタム報酬、カスタムツールループ、マルチターン会話形状を自由に定義できます。サーバーレス実行によりインフラストラクチャの懸念を簡素化し、GPUクラスターのプロビジョニングや管理なしで、トークン単位の価格設定でプロダクション規模のエージェントRLを実現します。非同期ロールアウトと軌跡収集により、生成と勾配更新が並行して実行され、トレーニングを高速化します。ネイティブアルゴリズムライブラリはPPO、CISPO、重要性サンプリング損失を網羅し、複数のグループベースのアドバンテージ推定器と組み合わせます。シーケンス拡張トレーニングにより、長いマルチターン軌跡の壁時計時間を短縮します。SageMaker AI管理のMLflowでの軌跡と報酬の可観測性により、エージェントのターンごとの動作をトレーニングステップ全体で確認できます。評価ジョブは、SageMaker AIエンドポイントまたはAmazon Bedrockにデプロイする前に、報酬、pass@k、軌跡メトリクスなどを報告します。
サービスはトレーニングループ、ハードウェア、オーケストレーションを提供します。信頼性の高いエージェントを得るかどうかを決定する選択はユーザーに委ねられています。エージェントがトレーニングする環境を構築し、報酬以外で成功を測定し、報酬自体を設計し、曲線が停滞したときに反復する方法を決定します。
安価で再現可能かつ代表的なトレーニング環境の構築 シングルターンRLにはプロンプトと報酬関数が必要です。マルチターンRLでは、エージェントがターン全体で行動するための環境（ツール呼び出しとその背後にあるシステム）が追加されます。この環境はトレーニング設定の一部であり、その構築方法はモデルが学習できる内容とメトリクスの信頼性の両方を形作ります。エージェントをトレーニングする際は、本番環境に類似しているがライブトラフィックから隔離されたサンドボックスまたはシミュレーション環境を構築します。ツール呼び出しと応答は同じスキーマとビジネスロジックを維持し、ライブコールの代わりに記録された応答または分離された状態によって駆動されます。
シミュレーション環境が推奨される出発点です。典型的な実行では数千のロールアウトが生成され、各ロールアウトで複数のツール呼び出しが行われます。たとえば、バッチサイズ128、グループサイズ8の場合、1ステップあたり1024のロールアウトになります。このトラフィックをライブシステムに向けると、顧客への影響が生じる可能性があります。シミュレーション環境がない場合、探索によって実際の副作用が発生する可能性があります。たとえば、試行錯誤によって学習するエージェントは、意図しない返金、レコードの削除、ワークフローのトリガーを行う可能性があります。さらに、ライブデータは動的に変化するため、同じ軌跡でも実行ごとにスコアが異なります。報酬を計算するには正しい結果を知る必要があるため、ツール呼び出しの宛先に関係なく、固定されたラベル付きタスクセットが必要です。
シミュレーション環境の構築方法はツールの機能によって異なりますが、ほとんどのユースケースをカバーする3つのパターンがあります：読み取り専用ツール、ステートフルツール、検証可能な結果です。どのパターンを選択しても、再現性（同じ引数でのツール呼び出しが同じ結果を返す）と代表性（実際のスキーマとデータ分布から環境を構築し、学習した動作が本番に移行できるようにする）の2つのプロパティを固定します。
トレーニング前の外部評価の設定 環境が整って検証されたら、報酬関数を書く前に成功を測定する方法を構築します。その測定は最終目標を直接捉えるべきです。RLは報酬信号を文字通り最適化するため、報酬だけが監視する数値である場合、タスクの進捗と報酬基準を満たす進捗を区別できません。報酬、環境シード、ハイパーパラメータを反復する際にガイドとなる、信頼できる外部評価が必要です。
パターン：デプロイ時に重視する結果を、報酬とは独立して計算するホールドアウト評価を設定します。実際には、モデルを取得し、固定されたテスト分割でロールアウトサーバーを介して実行し、単一のタスク成功率を返す小さなコードです。SOP-Benchの場合、評価は最終JSONオブジェクトの完全一致です。報酬関数は部分点と重み付けコンポーネントを計算できますが、評価はしません。
トレーニング前に、ベースラインを確立します。ベースモデルとリファレンスモデル（Amazon Bedrockでホストされるフロンティアモデルが適しています）を同じ評価で実行します。これにより、ベースモデルが到達すべき距離と、このタスクで良好な状態がどのようなものかがわかります。
アンチパターン：トレーニング報酬またはその派生メトリクスを成功の尺度として扱うこと。マルチターンエージェントは特別な考慮が必要です。ツール呼び出しに対して支払う報酬は、エージェントにできるだけ多くのツールを呼び出すように教えます。ターン数を罰する報酬は、必要な情報を得る前に回答をコミットするように教えます。どちらの場合も、トレーニング報酬は上昇しますが、実際のタスク成功は低下します。
優れたマルチターンRL報酬関数の設計 報酬設計はRLで最も難しい未解決問題の1つです。同じ柔軟性により、エージェントは実際のタスクを解決できる一方で、タスクを実行せずに報酬を満たす方法を見つけることもできます。デフォルトでは、トレーニングと評価に同じスコアリングルールを使用し、具体的な理由がある場合にのみ逸脱します。
SOP-Benchでは、ベンチマークは最終JSONオブジェクトを期待します。トレーニングと評価は通常このスコアリングルールを共有し、その周囲で観察するものだけが異なります。デフォルトのベンチマークスコアリングルールから逸脱する正当な理由は2つあり、どちらもより密な報酬を必要とします。1つ目はアルゴリズム的な理由です。RLはグループベースのアドバンテージ法を使用してグループ内の分散から学習信号を計算します。バイナリスコアはその分散を崩壊させる可能性があります。グループ内のすべてのロールアウトが同じスコアの場合、相対信号はゼロになり、グループは勾配に貢献しません。2つ目は収束速度です。グループ分散が健全であっても、密な報酬は完全に成功したロールアウトだけでなく、すべてのロールアウトに対して部分的な進捗への勾配をモデルに与えます。
結論として、これらのベストプラクティスに従うことで、開発者はマルチターンRLエージェントをより確実にトレーニングし、本番環境で複雑なタスクを効果的に実行できるようになります。