使用 Amazon SageMaker AI 上的 NVIDIA Isaac Lab 擴充套件機器人強化學習

物理人工智慧正從研究走向生產。機器人在高保真模擬器中訓練後部署到工廠、倉庫和物流中心，因為現實世界中的訓練速度慢、成本高且通常不安全，而 GPU 加速的模擬可以將數月的學習壓縮到數小時內。這將對計算能力提出更高要求。強化學習（RL）需要大量計算，尤其是像人形機器人在粗糙地形上行走這樣的複雜行為，單節點訓練可能需要數小時甚至數天。機器人團隊需要在研究階段快速迭代，同時執行生產級、長週期的訓練任務，而無需承擔維護計算叢集的運營負擔。

Amazon SageMaker AI 消除了管理機器學習訓練基礎設施的繁重工作。該服務可以自動配置例項、配置驅動程式和網路、監控節點健康，並在作業完成後清理資源，從而使工程師能夠專注於開發機器人策略而非基礎設施。對於機器人策略的強化學習，這種優勢尤為明顯，因為訓練任務往往耗時長、GPU 密集，且通常跨多個節點分散式執行。開發通常包括兩個階段：短期的迭代實驗（調整獎勵函式、觀測空間和模型架構）和長期的生產級訓練（將調整後的配置訓練至收斂）。SageMaker AI 提供了兩種計算選項來適應這些階段：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 核心而不適用。該解決方案已配置彈性結構介面卡（EFA）以實現多節點通訊。

設定步驟包括克隆 GitHub 倉庫、構建 Docker 映象並推送到 Amazon ECR。之後可按需選擇 SageMaker HyperPod 或 SageMaker Training Jobs 後端進行訓練。詳細指南請參考 AWS 官方部落格及 GitHub 倉庫。