Amazon SageMaker AI 上の NVIDIA Isaac Lab でロボット強化学習をスケール

物理 AI は研究から生産へと移行しています。ロボットは、現実世界でのトレーニングが遅く、高コストで、しばしば危険であるため、工場や倉庫、物流センターに配備される前に高忠実度シミュレーションでトレーニングされるようになっています。GPU 高速化シミュレーションは、数ヶ月の学習を数時間に圧縮できます。これにより、計算能力への要求が高まります。強化学習（RL）は、特に不整地での歩行のような複雑な動作では計算集約的であり、単一ノードのトレーニング実行が数時間から数日かかることもあります。ロボットチームは、研究段階で迅速に反復し、同時に計算クラスターの運用負担なく本番環境向けの長期トレーニングジョブを実行する必要があります。

Amazon SageMaker AI は、機械学習トレーニングのためのインフラストラクチャ管理の負担を軽減します。このサービスはインスタンスのプロビジョニング、ドライバーとネットワークの設定、ノードの健全性監視、ジョブ終了後のリソース解放を自動化し、エンジニアはインフラストラクチャではなくロボットポリシーの開発に集中できます。この利点は、ロボットポリシーの RL で特に顕著であり、実行時間が長く、GPU 集約的で、複数ノードに分散されることがよくあります。開発は通常、短期間の反復実験（報酬関数、観測空間、モデルアーキテクチャの調整）と、調整された構成を収束までトレーニングする長期の本番実行の2段階で行われます。SageMaker AI はこれらの段階に適合する2つのコンピュートオプションを提供します。SageMaker HyperPod は、自動ノード健全性チェックと障害回復機能を備えた永続的で弾力性のあるクラスターを提供し、本番トレーニングに適しています。SageMaker Training Jobs は、完全管理型のオンデマンドコンピュートを提供し、アイドルコストが発生しないため、反復実験やハイパーパラメータチューニングに適しています。

NVIDIA Isaac Lab は、NVIDIA Isaac Sim 上に構築されたオープンソースのロボット学習フレームワークです。GPU 並列シミュレーションを利用して、数千のロボットインスタンスを同時に実行し、数ヶ月の現実経験を数時間のシミュレーショントレーニングに変換します。サンプルトレーニングタスクは Isaac-Velocity-Rough-H1-v0 で、Unitree H1 ヒューマノイドが不整地を歩行しながら速度指令を追跡することを学習します。ロボットは19の関節を協調させて不規則な地形でバランスを保つ必要があります。トレーニングは PPO（近接ポリシー最適化）アルゴリズムを skrl を通じて使用します。

ソリューションは、統一された Docker イメージ（nvcr.io/nvidia/isaac-sim:5.1.0 ベース）、生成スクリプト（Kubernetes マニフェストと SageMaker 起動スクリプトを生成）、オプションの MLflow 実験追跡で構成されています。トレーニングトポロジは両方のバックエンドで本質的に同一であり、torchrun を使用して Isaac Lab のトレーナーを起動しますが、各環境がコンテナにトポロジ情報を提供する方法が異なります。GPU インスタンスの互換性については、Isaac Sim はハードウェア RT コアを搭載した GPU を必要とするため、G ファミリーインスタンス（ml.g6、ml.g6e、ml.g7e など）が推奨され、P ファミリーインスタンス（p4d、p5 など）は RT コアがないため使用できません。ソリューションはマルチノード通信のために Elastic Fabric Adapter（EFA）を設定しています。

セットアップ手順には、GitHub リポジトリのクローン、Docker イメージのビルド、Amazon ECR へのプッシュが含まれます。その後、SageMaker HyperPod または SageMaker Training Jobs のいずれかのバックエンドを選択してトレーニングを実行できます。詳細なガイドは AWS 公式ブログと GitHub リポジトリを参照してください。