AWS · ML 博客

呵

在 Amazon SageMaker AI 上使用 NVIDIA Isaac Lab 扩展机器人强化学习

Scale Robot Reinforcement Learning with NVIDIA Isaac Lab on Amazon SageMaker AI

二〇二六年六月十日 · 英文原文
摘要

本文展示了如何使用 NVIDIA Isaac Lab 在 Amazon SageMaker AI 上训练 Unitree H1 人形机器人策略，通过两种计算选项实现：SageMaker HyperPod 提供持久 GPU 集群，支持弹性、长时间运行的分布式训练，并集成自动故障恢复与 checkpoint 恢复；SageMaker Training Jobs 提供临时按需运行，适用于实验和超参数扫描。两种方案共享同一容器镜像和配置，训练代码保持不变。完整代码已发布于 GitHub 仓库。

Physical AI 正从研究走向生产。机器人在部署到工厂、仓库和物流中心之前，越来越多地先在高保真仿真中进行训练，因为在现实世界中训练既慢、又贵，而且往往不安全，而 GPU 加速的仿真则能将数月的学习压缩到几小时内完成。这便将挑战转移到了算力上。针对复杂行为（例如人形机器人在崎岖地形上的行走）的强化学习（RL）是计算密集型的，单节点训练运行时间从几小时到几天不等。机器人团队需要在研究阶段快速迭代，同时也要能运行生产级、长周期的训练任务，而无需承担维护计算集群的运维负担。在这篇文章中，我们将展示如何使用 NVIDIA Isaac Lab 在 Amazon SageMaker AI 上，通过两种计算选项——Amazon SageMaker HyperPod 和 Amazon SageMaker Training Jobs——来训练 Unitree H1 人形机器人的策略。本解决方案的完整代码可在随附的 GitHub 仓库中找到。图片来源：NVIDIA

1. 为什么选择 Amazon SageMaker AI 进行 Physical AI 训练

Amazon SageMaker AI 消除了管理机器学习（ML）训练计算基础设施中那些无差异化的繁重工作。该服务负责配置实例、设置驱动程序和网络、监控节点健康状态，并在任务完成后清理资源，从而使工程工作能够集中在开发机器人策略上，而不是底层基础设施上。这对于机器人策略 RL 尤其重要，因为这种任务基础设施负担很重：运行时间长、GPU 密集，并且通常分布在多个节点上。开发通常包含两个阶段：用于调整奖励函数、观测空间和模型架构的短周期迭代实验，以及将调整好的配置训练到收敛的更长周期生产运行。SageMaker AI 提供了适合这两个阶段的两种计算选项。

使用 SageMaker HyperPod 实现集群弹性和控制

SageMaker HyperPod 是一个专为大规模基础模型的分布式训练和推理而构建的托管基础设施。弹性是 SageMaker HyperPod 的核心。硬件故障在规模扩大时会成为问题，而在多节点 RL 运行中，每一次故障都意味着训练进度的损失，外加检测故障、替换节点以及从上一个 checkpoint 重新启动所需的时间。SageMaker HyperPod 在每个节点上运行一个健康监控代理，执行基本和深度健康检查。当检测到故障时，它会自动重启或替换故障实例。借助自动恢复功能，训练任务会在替换节点就绪后从上一个 checkpoint 重新启动，无需人工干预。通过与 Amazon Elastic Kubernetes Service（Amazon EKS）或 Slurm 编排，HyperPod 提供对集群节点的直接访问以及跨运行持久存在的稳定环境。HyperPod 可观测性插件将数百个集群、节点和任务指标发布到 Amazon Managed Service for Prometheus，并在预构建的 Amazon Managed Grafana 仪表板中可视化。团队无需设置指标管道即可获得 GPU 利用率、内存压力、网络吞吐量和任务级性能。HyperPod 任务治理基于 Kueue 构建，允许管理员将集群划分为具有计算配额、优先级和抢占权的命名空间作用域队列。可以按实例、按整个 GPU 或按 NVIDIA Multi-Instance GPU（MIG）的 GPU 分区来定义分配。细粒度的配额涵盖加速器、vCPU 和内存。

使用 SageMaker Training Jobs 实现临时计算

SageMaker Training Jobs 是一种完全托管、按需运行容器化训练工作负载的方式，无需维护任何长期存在的计算资源。每个任务都会配置 GPU 实例，从 Amazon Elastic Container Registry（Amazon ECR）拉取容器，运行训练脚本，将产物上传到 Amazon Simple Storage Service（Amazon S3），并在任务完成后终止实例。运行之间没有空闲的计算成本。这种模式适合策略开发的迭代阶段，在此阶段，奖励函数、观测空间和网络架构在短周期运行之间频繁变化。它也适合超参数调优扫描，其中许多短周期运行并行执行，然后释放其计算资源。

2. NVIDIA Isaac Lab 与训练任务

NVIDIA Isaac Lab 是一个基于 NVIDIA Isaac Sim 构建的开源机器人学习框架。它使用 GPU 并行仿真，在一张或多张 GPU 上同时运行数千个机器人实例，将原本需要数月的真实世界经验转化为数小时的模拟训练。Isaac Lab 提供了结构化的 API 来定义任务、观测和动作空间、奖励函数以及用于强化学习和模仿学习的训练循环。图片来源：NVIDIA

本文中的示例训练任务是 Isaac-Velocity-Rough-H1-v0，其中 Unitree H1 人形机器人学习在崎岖地形上行走时跟踪速度指令。机器人必须协调其 19 个关节，以在程序生成的崎岖表面上保持平衡。训练使用通过 skrl（Isaac Lab 支持的多个 RL 框架之一）实现的 Proximal Policy Optimization（PPO）。扩展到多个节点会倍增并行环境的数量，为每次策略更新产生更多样化的经验，从而加速收敛。您可以将此解决方案中提供的脚本和配置扩展到其他机器人学习任务。

3. 解决方案概述

随附 GitHub 仓库中的解决方案包含两个主要部分：（1）一个单一的 Docker 镜像，可在 SageMaker HyperPod 和 SageMaker Training Jobs 上运行训练代码；（2）一个生成器脚本，可根据共享配置文件渲染 Kubernetes 清单和 SageMaker 启动脚本。这两种服务选项的区别仅在于镜像的启动方式：在 SageMaker HyperPod 上作为 Kubernetes PyTorchJob 启动，或通过 SageMaker Training Job 的 CreateTrainingJob API 调用启动。此处使用的 H1 运动任务与 NVIDIA Isaac Lab on AWS workshop 中的任务相同，后者在 Amazon Elastic Compute Cloud（Amazon EC2）和 AWS Batch 上运行工作负载。迁移到 SageMaker AI 后，训练代码保持不变，并增加了托管集群、集成的故障恢复和无服务器训练任务执行。

训练镜像

训练容器镜像基于 nvcr.io/nvidia/isaac-sim:5.1.0 构建。提供的 Dockerfile 克隆了 Isaac Lab v2.3.2，将其安装到 Isaac Sim 捆绑的 Python 环境中，并复制了入口点脚本，该脚本解析 SageMaker Training Jobs 资源配置以启动 torchrun。完整的 Dockerfile 位于 docker/Dockerfile。两种服务选项使用相同的镜像。

实验追踪

当配置了追踪服务器时，训练指标会流式传输到 Amazon SageMaker 托管的 MLflow，以便在两个后端之间进行持久、可搜索的实验追踪。MLflow 是可选的：将追踪 URI 留空即可完全禁用它。第 4.5 节介绍了配置。

配置与生成器脚本

生成器脚本通过 config.yaml 中定义的环境特定变量进行配置。generate.py 脚本读取配置，并将 templates/ 中的模板渲染为 generated/ 下可直接应用的文件。运行生成器只需一条命令：

python generate.py

每个后端使用的具体文件将在第 4 节（SageMaker HyperPod）和第 5 节（SageMaker Training Jobs）的演练中分别介绍。

跨后端的训练拓扑

在提供的解决方案中，两种路径最终都在同一镜像上以相同的 torchrun 调用 Isaac Lab 的 skrl 训练器结束。主要区别在于每个环境如何向容器提供拓扑信息。在 SageMaker HyperPod 上，Kubeflow Training Operator 将 MASTER_ADDR、MASTER_PORT、RANK 和 WORLD_SIZE 注入到每个 pod 中。这些描述了 pod 级别的拓扑（WORLD_SIZE 是 pod 数量，RANK 是每个 pod 的索引）。入口点将它们转发给 torchrun，后者在每个 pod 内为每个 GPU 生成一个进程。每个 pod 的启动器通过 MASTER_ADDR:MASTER_PORT 进行会合，形成全局进程组。在 SageMaker Training Jobs 上，SageMaker 将主机列表写入 /opt/ml/input/config/resourceconfig.json，容器入口点在启动时解析它。

GPU 实例兼容性

Isaac Sim 基于 NVIDIA Omniverse 构建，并使用 Omniverse RTX 渲染器，这需要具有硬件 RT Core 的 GPU。AWS G 系列的 GPU 实例适合 Isaac Lab 工作负载。P 系列不适合，因为它使用没有 RT Core 的数据中心 GPU。有关支持和不支持硬件的完整列表，请参阅 Isaac Sim 5.1 要求页面。

实例系列 GPU 类型与代际 RT Core / Isaac Sim 兼容性
ml.g5 NVIDIA A10G (Ampere) 是
ml.g6 NVIDIA L4 (Ada Lovelace) 是
ml.g6e NVIDIA L40S (Ada Lovelace) 是
ml.g7e NVIDIA RTX PRO 6000 (Blackwell) 是
ml.p4d, ml.p4de, ml.p5, ml.p5e, ml.p5en, ml.p6-b200, ml.p6-b300, ml.p6e-gb200 NVIDIA A100 (Ampere), H100 / H200 (Hopper), B200 / B300 / GB200 (Blackwell) 否

本文中的示例全程使用 ml.g6.12xlarge。您可以在 config.yaml 中更改实例类型。ml.g6、ml.g6e 和 ml.g7e 系列在 8xlarge 及以上尺寸支持 Elastic Fabric Adapter（EFA），这为 NCCL 提供了用于多节点集合通信的内核旁路、支持 RDMA 的传输。在 HyperPod 上启用 EFA 需要 AWS EFA 设备插件，并在 pod 规范中请求 vpc.amazonaws.com/efa 资源。在 SageMaker Training Jobs 上，您必须在容器镜像和虚拟私有云（VPC）配置中配置 EFA。对于 SageMaker HyperPod 和 SageMaker Training Jobs 后端，解决方案会自动配置 EFA。SageMaker Training Job 的设置请参阅文档。

设置：克隆仓库并构建镜像

两个演练共享两个设置步骤：克隆随附仓库和构建训练镜像。

克隆解决方案的仓库：

git clone https://github.com/awslabs/awsome-distributed-ai.git
cd awsome-distributed-ai/3.test_cases/pytorch/nvidia-isaac-lab

该仓库包含 Dockerfile、配置模板、生成器以及两个后端使用的入口点脚本。

从仓库根目录构建镜像并推送到 Amazon ECR。根据您的设置定义环境变量：

export AWS_REGION=us-east-1 # 您的区域
export ACCOUNT=$(aws sts get-caller-identity --query Account --output text)

检查相应的 ECR 仓库是否存在，如果不存在则创建：

aws ecr describe-repositories --repository-names isaaclab-sagemaker --region "$AWS_REGION" 2>/dev/null || \
aws ecr create-repository --repository-name isaaclab-sagemaker --region "$AWS_REGION"

向 Amazon ECR 进行身份验证：

aws ecr get-login-password --region $AWS_REGION | \
docker login --username AWS --password-stdin \
$ACCOUNT.dkr.ecr.$AWS_REGION.amazonaws.com

构建并标记 Docker 镜像：

docker build -t isaaclab-sagemaker:5.1.0 -f docker/Dockerfile .
docker tag isaaclab-sagemaker:5.1.0 $ACCOUNT.dkr.ecr.$AWS_REGION.amazonaws.com/isaaclab-sagemaker:5.1.0

将 Docker 镜像推送到 Amazon ECR：

docker push $ACCOUNT.dkr.ecr.$AWS_REGION.amazonaws.com/isaaclab-sagemaker:5.1.0

如果您想使用 Training Jobs，请跳至第 5 节。

4. 演练：在 SageMaker HyperPod 上使用 Amazon EKS 进行训练

在本演练中，我们使用一个由 Amazon EKS 编排的现有 SageMaker HyperPod 集群，该集群包含一个由两个 ml.g6.12xlarge 节点（每个节点 4× NVIDIA L4，共 8 个 GPU）组成的 GPU 实例组。目标是针对 H1 运动任务进行分布式训练，在 SageMaker 托管的 MLflow 中实时查看指标，并将生成的 checkpoint 写入 FSx for Lustre。

4.1 前提条件

该解决方案需要满足以下前提条件：

4.2 配置并生成清单

复制示例配置：

cp config.yaml.example config.yaml

填写您的环境值以及 AWS 账户 ID、区域和集群详细信息：

aws:
  account_id: " " # 您的 12 位 AWS 账户 ID
  region: " " # 例如 us-east-2
ecr:
  repository: "isaaclab-sagemaker" # 必须与您推送到的仓库匹配
  tag: "5.1.0"
training:
  task: "Isaac-Velocity-Rough-H1-v0"
  max_iterations: 1000 # PPO 迭代次数；生产运行请增加
  framework: "skrl" # skrl | rsl_rl | rl_games | sb3
hyperpod_eks:
  fsx:
    file_system_id: " "
    dns_name: " .fsx. .amazonaws.com"
    mount_name: " " # 8 字符的 FSx 挂载名称
  jobs:
    training_job:
      instance_type: "ml.g6.12xlarge"
      gpus_per_node: 4
      num_nodes: 2 # 单节点训练设为 1
      fsx_log_dir: "/fsx/isaaclab-h1/logs"

重要的配置字段包括：

通过执行脚本生成清单：

python generate.py
# Config: config.yaml
# Image: .dkr.ecr..amazonaws.com/isaaclab-sagemaker:5.1.0
# Task: Isaac-Velocity-Rough-H1-v0
# Iterations:1000
#
# Generated: generated/storage.yaml
# Generated: generated/training-job.yaml
# Generated: generated/launch-sm-training.py
# Generated: generated/viz-eks-webrtc-pod.yaml

以下 Kubernetes 清单已生成，并将在演练的后续部分中使用：

生成的文件 说明 何时应用
storage.yaml PersistentVolume 和 PersistentVolumeClaim，绑定到您的 FSx for Lustre 文件系统，将其挂载到 pod 的 /fsx 路径。 每个集群一次。
training-job.yaml Kubeflow PyTorchJob，包含一个 Master 副本和 num_nodes - 1 个 Worker 副本。当 num_nodes: 1 时在单个节点上运行，否则跨多个节点运行。 每次训练运行。
viz-eks-webrtc-pod.yaml 运行 Isaac Sim 的 Pod，采用无头流式传输模式，并附带一个基于浏览器的 WebRTC 客户端，用于可视化训练好的策略。 可选。在第 6 节中介绍。

剩余文件 launch-sm-training.py 将在 SageMaker Training Jobs 演练（第 5 节）中介绍。

4.3 部署共享存储

FSx for Lustre 是本演练的存储层。它提供并行、高吞吐量的写入，能够处理来自多个 pod 的 checkpoint，而不会成为训练循环的瓶颈，并且允许训练任务和可视化 pod 使用相同的卷。生成的文件 generated/storage.yaml 包含指向您 FSx 文件系统的 PersistentVolume 和 PersistentVolumeClaim。将其应用到您的集群：

kubectl apply -f generated/storage.yaml
kubectl get pvc isaaclab-fsx-pvc
# NAME               STATUS   VOLUME          CAPACITY   ACCESS MODES
# isaaclab-fsx-pvc   Bound    isaaclab-fsx-pv 1200Gi     RWX

4.4 启动训练

生成的文件 generated/training-job.yaml 是一个 Kubeflow PyTorchJob，包含一个 Master 副本和 num_nodes - 1 个 Worker 副本。使用 config.yaml 中的默认值 jobs.training_job.num_nodes: 2，它将跨两个 ml.g6.12xlarge 节点（共 8 个 GPU）运行。将 num_nodes 设置为 1 将生成一个没有 Worker 副本的单节点任务。

当任务启动时，Kubeflow Training Operator 将标准的 PyTorch 分布式环境变量（MASTER_ADDR、MASTER_PORT、RANK、WORLD_SIZE）注入到每个 pod 中。容器启动脚本将它们传递给 torchrun，后者负责处理会合和进程组设置：

# 摘自 generated/training-job.yaml
/isaac-sim/python.sh -m torch.distributed.run \
    --nproc_per_node=4 \
    --nnodes=2 \
    --node_rank=$RANK \
    --rdzv_id=isaaclab-job \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    scripts/reinforcement_learning/skrl/train.py \
    --distributed --task=Isaac-Velocity-Rough-H1-v0 \
    --max_iterations=1000 --headless

应用清单：

kubectl apply -f generated/training-job.yaml

观察任务状态：

kubectl get pytorchjobs
# NAME          STATE     AGE
# isaaclab-h1   Running   3m
kubectl logs -f isaaclab-h1-master-0

早期日志会显示每个 pod 打印其 rank、master 地址以及 nvidia-smi 的输出。在 workers 连接到 master 后，Isaac Lab 加载场景（每个节点首次加载时，在资产缓存预热期间需要一两分钟），在所有可用的 GPU 上生成并行环境，并开始每隔几秒记录奖励和价值损失指标。入口点在将控制权交给 torchrun 之前，会打印 Kubeflow 注入的 pod 级拓扑。训练清单还会在启动时检查 FSx 上是否存在现有的 best_agent.pt。如果从之前的运行中找到，它会将 --checkpoint 传递给 train.py，以便训练从该点恢复，而不是从头开始。由 HyperPod 健康监控触发的 pod 重启或节点替换会自动从上一个 checkpoint 继续。

=== Master Node Info ===
Hostname: isaaclab-h1-master-0
MASTER_ADDR: isaaclab-h1-master-0
WORLD_SIZE: 2
RANK: 0
GPU 0: NVIDIA L4 (UUID: GPU-dd5102c3-be08-...)
GPU 1: NVIDIA L4 (UUID: GPU-3c0d70fe-519f-...)
GPU 2: NVIDIA L4 (UUID: GPU-42485aa5-6a2a-...)
GPU 3: NVIDIA L4 (UUID: GPU-18f62bef-4155-...)
=== Starting Master (2 nodes, 8 GPUs total, 1000 iterations) ===
...
[INFO][AppLauncher]: Using device: cuda:0
[INFO]: Scene manager: Number of environments: 4096
  Environment spacing : 2.5
[INFO]: Time taken for scene creation: 16.30 seconds
 10%|▉         | 118/24000 [00:09

这里 WORLD_SIZE: 2 是 pod 数量，而不是全局进程数。torchrun 在启动后总共生成 8 个进程（2 个 pod 各 4 个），在这些进程内部，WORLD_SIZE 变为 8。

4.5 使用 SageMaker 托管的 MLflow 追踪实验

训练指标、运行参数（任务、迭代次数、种子）以及最终的 checkpoint 目录会被转发到 Amazon SageMaker 托管的 MLflow，以提供持久、可搜索的实验存储。系统指标（如 GPU 利用率和 CPU/内存使用率）由 MLflow 自身的后台线程采样。启用 MLflow 是可选的。当 MLFLOW_TRACKING_URI 为空（默认值）时，训练脚本会跳过所有 MLflow 调用。在 config.yaml 中设置追踪 URI 和实验名称：

mlflow:
  tracking_uri: "arn:aws:sagemaker::: :mlflow-app/ "
  experiment_name: "isaaclab-h1"
  # 仅当追踪 URI 是 Studio 作用域的 MLflow App 且
  # 训练角色在该 Studio 域之外时才需要。见下文。
  assume_role_arn: ""

重新生成清单并重新启动训练任务。训练 pod 日志会在启动后几秒钟内打印运行 URL：

INFO mlflow.tracking.fluent: Experiment with name 'isaaclab-h1' does not exist. Creating a new experiment.
INFO mlflow.system_metrics.system_metrics_monitor: Started monitoring system metrics.
...
View run 2026-04-27_18-59-29_ppo_torch at:
https://mlflow.sagemaker.us-east-2.app.aws/#/experiments/1/runs/
View experiment at:
https://mlflow.sagemaker.us-east-2.app.aws/#/experiments/1

当运行完成时，MLflow 会关闭系统指标线程，并且来自 skrl 的 Training time: 行会出现：

Training time: 91.9 seconds
INFO mlflow.system_metrics.system_metrics_monitor: Stopping system metrics monitoring...
INFO mlflow.system_metrics.system_metrics_monitor: Successfully terminated system metrics monitoring!

打开 URL，或从 SageMaker Studio 导航到 MLflow UI，即可实时查看奖励曲线和价值损失更新以及 GPU 利用率。SageMaker 托管的 MLflow 存在两种授权模型：

5. 演练：在 SageMaker Training Jobs 上进行训练

SageMaker Training Jobs 通过不同的生命周期运行相同的镜像。每个任务都会配置请求的 GPU 实例，从 Amazon ECR 拉取镜像，运行入口点，将训练脚本复制到 /opt/ml/model/ 的所有文件上传到 S3 输出路径，然后终止实例。

5.1 前提条件

5.2 配置

config.yaml 的 Training Jobs 部分捕获了执行角色、实例类型和数量以及输出位置。当 scripts S3 URI 和 output S3 路径留空时，它们会从顶层的 s3.bucket 值自动派生。

s3:
  bucket: " "
sagemaker_training:
  role_arn: "arn:aws:iam:::role/ "
  instance_type: "ml.g6.12xlarge"
  instance_count: 2 # SageMaker 处理多节点连接
  volume_size_gb: 200
  max_runtime_seconds: 7200 # 任务的硬性上限

重要的配置字段包括：

5.3 上传入口点

SageMaker 在任务启动时将入口点脚本从 Amazon S3 拉取到每个训练实例中。上传一次。后续每个任务都会从同一位置读取，直到上传新版本。将存储桶名称替换为您选择的 Amazon S3 存储桶：

aws s3 cp scripts/sm-train-entrypoint.sh \
    s3:///scripts/sm-train-entrypoint.sh

5.4 生成并启动

运行 python generate.py 会在 generated/ 目录下生成一个 launch-sm-training.py 脚本，其中包含从 config.yaml 预填充的镜像 URI、IAM 角色、S3 路径和实例配置。该脚本提供了一个小型 CLI，用于在运行之间覆盖值：

python generate.py # 刷新 generated/
python generated/launch-sm-training.py # 默认：1000 次迭代
python generated/launch-sm-training.py --iterations 1000
python generated/launch-sm-training.py --dry-run # 打印任务配置，不启动

启动器会调用 CreateTrainingJob，并附带一个带有时间戳后缀的任务名称、Amazon ECR 镜像和 S3 入口点位置。它还会传递容器所需的 Isaac Sim 环境变量（ACCEPT_EULA、NVIDIA_VISIBLE_DEVICES=all、MAX_ITERATIONS 等）。成功后，它会打印任务名称和一个用于监控进度的 describe-training-job 命令。

5.5 监控

aws sagemaker describe-training-job \
    --training-job-name \
    --query '{Status: TrainingJobStatus, Secondary: SecondaryStatus}'

SecondaryStatus 字段会依次经历 Pending → Downloading → Training → Uploading → Completed：

{
    "Status": "InProgress",
    "SecondaryStatus": "Training"
}

训练日志会流式传输到 /aws/sagemaker/TrainingJobs 日志组下的 Amazon CloudWatch Logs，每个实例对应一个日志流。如果您更喜欢使用 UI，SageMaker 控制台会直接从任务页面链接到该流。一个成功的 rank 0 流会以入口点的自检开始：

=== SageMaker Training Job ===
Hostname: ip-10-0-195-224.us-east-2.compute.internal
GPU 0: NVIDIA L4 (UUID: GPU-66e3a452-...)
GPU 1: NVIDIA L4 (UUID: GPU-a075bb9c-...)
GPU 2: NVIDIA L4 (UUID: GPU-2ba15062-...)
GPU 3: NVIDIA L4 (UUID: GPU-e25c05a4-...)
=== Resource Config ===
{"current_host":"algo-1","hosts":["algo-1","algo-2"],"network_interface_name":"eth0"}
=== Training Configuration ===
CURRENT_HOST=algo-1
MASTER_HOST=algo-1
NNODES=2
NODE_RANK=0
NPROC=4
MAX_ITERATIONS=1000
=== Starting Isaac Lab H1 Training ===

当任务完成时，SageMaker 会将入口点复制到 /opt/ml/model/ 的所有内容打包成一个 model.tar.gz，并上传到输出 S3 路径。对于 H1，该归档文件包含 skrl 的 logs/ 目录，其中包含训练 checkpoint 和 best_agent.pt。当您在 config.yaml 中设置 sagemaker_training.checkpoint_s3_path 时，启动器会包含一个 CheckpointConfig，告诉 SageMaker 在训练期间持续将 /opt/ml/checkpoints 同步到 Amazon S3。入口点会将 skrl 的日志目录符号链接到该路径，因此 skrl 写入的每个 checkpoint 都会近乎实时地备份到 Amazon S3。如果任务失败或中断，使用相同的 checkpoint_prefix 重新启动它将恢复最新的 checkpoint 并自动恢复训练。

第 4.5 节中描述的相同 MLflow 集成也适用于 Training Jobs。当您设置 mlflow.tracking_uri 时，生成的 launch-sm-training.py 会将 MLflow 环境变量作为 CreateTrainingJob 请求的一部分转发到训练容器，训练代码会将指标写入同一个实验。对于 Studio MLflow Apps，训练任务的执行角色必须列在 Studio 执行角色的信任策略中，并且在其自身权限中携带 sts:AssumeRole。

6. 可视化训练好的策略

该仓库为 SageMaker HyperPod 提供了一个可视化 pod，它通过 WebRTC 将 Isaac Sim GUI 直接流式传输到浏览器中，使用与训练任务相同的 FSx 卷，因此 HyperPod 上生成的任何 checkpoint 都可以重放。

HyperPod 集群上的 WebRTC 流式传输

Isaac Sim 包含内置的无头 WebRTC 流式传输。可视化 pod 捆绑了两个共享 pod 网络的容器：

可视化 pod 在同一个集群的 GPU 节点上运行，并挂载训练任务使用的 FSx 卷，因此 HyperPod 运行产生的 checkpoint 可直接用于重放。清单由 generate.py 与训练清单一起渲染到 generated/viz-eks-webrtc-pod.yaml 中：

kubectl apply -f generated/viz-eks-webrtc-pod.yaml
kubectl logs -f isaacsim-webrtc -c isaacsim # 等待 "app ready"

连接需要额外一步：kubectl port-forward 仅支持 TCP，但 WebRTC 媒体需要 UDP。krelay 是一个添加了 UDP 转发的 kubectl 插件。使用以下命令安装它：

kubectl krew install relay

通过运行以下命令启动端口转发：

kubectl relay pod/isaacsim-webrtc \
    8210:8210 49100:49100 47998:47998@udp
# 在 Chromium 中打开

浏览器连接到 web viewer sidecar，后者与 Isaac Sim 容器协商一个 WebRTC 会话，并显示实时视口。要重放不同的 checkpoint，请编辑可视化 pod 清单中的 isaacsim 容器参数（或者在更新 config.yaml 中的 training.task 后删除 pod 并重新生成）。对于团队可访问的部署，请将本地 relay 替换为一个同时暴露 TCP 和 UDP 端口的 AWS Network Load Balancer，并将 Isaac Sim 的 publicIp 标志设置为 NLB 的公共地址。提供的可视化 pod 是 skrl 专用的。如果您在 config.yaml 中更改了框架，请相应地更新 checkpoint 路径和 play 脚本。

替代方案：使用 NICE DCV 的 Amazon EC2

如果完整的 Linux 桌面比基于浏览器的查看器更可取（例如，为了在 Isaac Sim 旁边运行终端和文件浏览器），NVIDIA Isaac Lab on AWS workshop 介绍了如何设置一个独立的 Amazon EC2 GPU 实例，该实例带有 NICE DCV，并通过低延迟的远程桌面流式传输交互式运行相同的 Isaac Lab 镜像。本文中 SageMaker 任务产生的 checkpoint 可以通过挂载 FSx 文件系统或从 S3 存储桶下载，在该实例上重放。

7. 成本考量与清理

两种计算选项具有不同的成本形态。SageMaker HyperPod 是一个持久集群：实例在作为集群一部分期间计费。FSx for Lustre 按每预置容量每小时计费，第 6 节中的可视化 pod 在运行期间会占用一个 GPU 节点。SageMaker Training Jobs 仅按每个任务的运行时间计费。有关当前费率，请参阅 SageMaker AI 和 FSx for Lustre 定价页面。

7.1 清理

SageMaker HyperPod

# 删除训练任务和可视化 pod
kubectl delete pytorchjob isaaclab-h1
kubectl delete -f generated/viz-eks-webrtc-pod.yaml

将会话之间的 GPU 实例组缩容到零，以在保持集群配置的同时暂停实例成本，或者完全删除集群。请参阅管理 SageMaker HyperPod 集群。删除 FSx 文件系统将永久删除其上存储的所有训练 checkpoint 和日志。在继续之前，请下载您想要保留的任何 checkpoint。

# 删除 FSx 文件系统（替换为您的文件系统 ID）
aws fsx delete-file-system --file-system-id

SageMaker Training Jobs

Training Jobs 在任务完成或失败时自动终止。无需进行计算清理。以下命令会永久删除训练产物和 checkpoint。在运行之前，请下载您想要保留的任何文件。

# 从 S3 删除训练产物和 checkpoint
aws s3 rm s3:///sm-training-output/ --recursive
aws s3 rm s3:///sm-training-checkpoints/ --recursive

Amazon ECR

# 从 ECR 删除训练镜像（替换为您的区域和账户）
aws ecr batch-delete-image \
    --repository-name isaaclab-sagemaker \
    --image-ids imageTag=5.1.0 \
    --region $AWS_REGION

8. 结论

随着 Physical AI 工作负载进入生产阶段，团队需要在不承担管理计算基础设施运维负担的情况下扩展策略训练。在本文中，我们展示了 SageMaker HyperPod 和 SageMaker Training Jobs 如何让机器人团队在托管的 GPU 基础设施上运行分布式 Isaac Lab 训练，使用单个容器镜像和跨两种计算模型的共享配置。SageMaker HyperPod 提供持久 GPU 集群，支持弹性、长时间运行的训练。SageMaker Training Jobs 提供临时的按需运行，适用于实验和超参数扫描。两者都运行相同的容器镜像和相同的 skrl 训练器的 torchrun 调用，因此在它们之间切换只是一个配置更改。要开始使用，请探索随附的仓库以启动您的第一次 H1 训练运行，并将该模式扩展到其他 Isaac Lab 任务（人形机器人操作、四足机器人、灵巧手）。要了解更多信息，请参阅 Amazon SageMaker HyperPod 文档和 Amazon SageMaker Training Jobs 文档。

关于作者

Roy Allela Roy 是 AWS 的高级 AI/ML 专家解决方案架构师。他帮助 AWS 客户（从初创公司到大型企业）在 AWS 上高效地训练和部署基础模型。他拥有微处理器工程背景，热衷于计算优化问题和提升 AI 工作负载的性能。

Nicolas Jourdan Nicolas 是 AWS 的专家解决方案架构师，他帮助客户在云中释放 AI 和 ML 的全部潜力。Nicolas 在包括自动驾驶、无人机和制造业在内的多个行业拥有丰富的实践经验，曾担任从研究科学家到工程经理等职务。他为获奖研究做出了贡献，拥有目标检测和异常检测方面的专利，并热衷于应用尖端 AI 解决复杂的现实世界问题。

译自 AWS · ML 博客 · 录于 二〇二六年六月十日