Treinamento de Robôs com NVIDIA Isaac Lab no Amazon SageMaker AI

Physical AI: do laboratório para a produção

O treinamento de robôs físicos está deixando de ser um tema de pesquisa acadêmica e entrando de vez nas operações reais de fábricas, armazéns e centros de logística. O motivo é simples: treinar um robô no mundo real é lento, caro e muitas vezes perigoso. A simulação acelerada por GPU resolve esse problema ao comprimir meses de aprendizado em horas de processamento.

O desafio, então, passa a ser de infraestrutura. Aprendizado por Reforço (RL) para comportamentos complexos — como locomoção de robôs humanoides em terrenos irregulares — é extremamente intensivo em GPU, com execuções que podem durar de horas a dias em um único nó. Equipes de robótica precisam iterar rapidamente durante a pesquisa e também executar treinamentos de longa duração em produção, sem precisar manter clusters de computação manualmente.

A AWS publicou um guia técnico detalhado mostrando como fazer exatamente isso: treinar políticas de controle para o robô humanoide Unitree H1 usando o NVIDIA Isaac Lab no Amazon SageMaker AI. O código completo da solução está disponível no repositório GitHub.

Por que usar o Amazon SageMaker AI para treinar robôs?

O SageMaker AI elimina a necessidade de gerenciar a infraestrutura de computação para treinamento de Aprendizado de Máquina (ML). O serviço provisiona instâncias, configura drivers e redes, monitora a saúde dos nós e libera os recursos ao final de cada job — permitindo que a equipe de engenharia foque no desenvolvimento da política do robô, não na infraestrutura.

Isso é especialmente relevante para RL de políticas robóticas, que é pesado em infraestrutura: execuções longas, intensivas em GPU e frequentemente distribuídas em múltiplos nós. O desenvolvimento costuma ter duas fases distintas: experimentos iterativos curtos para ajustar funções de recompensa e arquiteturas de modelo, e execuções de produção mais longas para treinar uma configuração ajustada até a convergência.

Para cobrir essas duas fases, o SageMaker AI oferece duas opções de computação.

Resiliência e controle com o SageMaker HyperPod

O SageMaker HyperPod é uma infraestrutura gerenciada e dedicada para treinamento distribuído de modelos de grande escala. A resiliência é o núcleo do serviço: falhas de hardware em execuções multi-nó significam perda de progresso de treinamento, tempo para detectar a falha, substituir o nó e reiniciar a partir do último checkpoint. O HyperPod executa um agente de monitoramento de saúde em cada nó que realiza verificações básicas e profundas. Quando uma falha é detectada, ele reinicia ou substitui a instância com defeito automaticamente. Com a funcionalidade de retomada automática, o job de treinamento reinicia a partir do último checkpoint sem intervenção manual.

Orquestrado pelo Amazon Elastic Kubernetes Service (Amazon EKS) ou Slurm, o HyperPod oferece acesso direto aos nós do cluster e um ambiente estável que persiste entre execuções. O complemento de observabilidade do HyperPod publica centenas de métricas de cluster, nó e job para o Amazon Managed Service for Prometheus, com dashboards pré-construídos no Amazon Managed Grafana. A governança de tarefas do HyperPod, construída sobre o Kueue, permite que administradores dividam o cluster em filas com cotas de computação, prioridades e preempção por instância, GPU inteira ou partição de GPU com NVIDIA Multi-Instance GPU (MIG).

Computação efêmera com SageMaker Training Jobs

Os SageMaker Training Jobs são uma forma totalmente gerenciada e sob demanda de executar cargas de trabalho de treinamento em contêineres, sem manter computação de longa duração. Cada job provisiona instâncias de GPU, baixa o contêiner do Amazon Elastic Container Registry (Amazon ECR), executa o script de treinamento, faz upload dos artefatos para o Amazon Simple Storage Service (Amazon S3) e encerra as instâncias ao final. Não há custo de computação ociosa entre execuções.

Esse modelo é ideal para a fase de iteração do desenvolvimento de políticas, onde funções de recompensa, espaços de observação e arquiteturas de rede mudam frequentemente entre execuções curtas. Também é adequado para varreduras de ajuste de hiperparâmetros, onde muitas execuções curtas rodam em paralelo e depois liberam a computação.

NVIDIA Isaac Lab: simulação paralela em GPU

O NVIDIA Isaac Lab é um framework open-source de aprendizado robótico construído sobre o NVIDIA Isaac Sim. Ele usa simulação paralela em GPU para executar milhares de instâncias de robôs simultaneamente em uma ou múltiplas GPUs, transformando o que seriam meses de experiência no mundo real em horas de treinamento simulado. O Isaac Lab fornece APIs estruturadas para definir tarefas, espaços de observação e ação, funções de recompensa e loops de treinamento para aprendizado por reforço e aprendizado por imitação.

A tarefa de treinamento usada no guia é Isaac-Velocity-Rough-H1-v0, onde o robô humanoide Unitree H1 aprende a seguir comandos de velocidade enquanto caminha por terrenos irregulares. O robô precisa coordenar seus 19 juntas para manter o equilíbrio em superfícies irregulares geradas proceduralmente. O treinamento usa Otimização de Política Proximal (PPO) através do skrl, um dos vários frameworks de RL suportados pelo Isaac Lab. Escalar para múltiplos nós multiplica o número de ambientes paralelos, produzindo experiências mais diversas por atualização de política e acelerando a convergência.

Visão geral da solução

A solução no repositório GitHub consiste em duas partes principais: (1) uma única imagem Docker que executa o código de treinamento tanto no SageMaker HyperPod quanto no SageMaker Training Jobs, e (2) um script gerador que renderiza os manifestos Kubernetes e o script de lançamento do SageMaker a partir de um arquivo de configuração compartilhado.

As duas opções de serviço diferem apenas na forma como a imagem é lançada: como um PyTorchJob do Kubernetes no SageMaker HyperPod, ou por meio de uma chamada à API CreateTrainingJob para um SageMaker Training Job. A tarefa de locomoção do H1 é a mesma do workshop NVIDIA Isaac Lab na AWS, que executa a carga de trabalho no Amazon Elastic Compute Cloud (Amazon EC2) e no AWS Batch. A migração para o SageMaker AI mantém o código de treinamento inalterado e adiciona clusters gerenciados, recuperação integrada de falhas e execução serverless de jobs de treinamento.

Um ponto importante sobre compatibilidade de hardware: o Isaac Sim é construído sobre o NVIDIA Omniverse e usa o Omniverse RTX Renderer, que exige GPUs com RT Cores de hardware. A família G de instâncias GPU da AWS é adequada para cargas de trabalho do Isaac Lab. A família P não é compatível, pois usa GPUs de data center sem RT Cores. Consulte a página de requisitos do Isaac Sim 5.1 para a lista completa de hardware suportado. Os exemplos usam ml.g6.12xlarge ao longo de todo o guia.

Configuração inicial: repositório e imagem Docker

Dois passos de configuração são compartilhados entre os dois fluxos: clonar o repositório e construir a imagem de treinamento.

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

Em seguida, defina as variáveis de ambiente, verifique ou crie o repositório no Amazon ECR, autentique-se, construa e faça o push da imagem:

export AWS_REGION=us-east-1 # your region
export ACCOUNT=$(aws sts get-caller-identity --query Account --output text)

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"

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

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 push $ACCOUNT.dkr.ecr.$AWS_REGION.amazonaws.com/isaaclab-sagemaker:5.1.0

Treinamento no SageMaker HyperPod com Amazon EKS

O fluxo do HyperPod usa um cluster existente orquestrado pelo Amazon EKS, com um grupo de instâncias de dois nós ml.g6.12xlarge (4× NVIDIA L4 cada, 8 GPUs no total). O objetivo é um job de treinamento distribuído para a tarefa de locomoção do H1, com métricas ao vivo no SageMaker managed MLflow e os checkpoints resultantes gravados no FSx for Lustre.

Pré-requisitos do HyperPod

• Cota de serviço suficiente para o tipo de instância GPU escolhido na região alvo. Solicite um aumento pelo AWS Service Quotas antes de criar ou escalar o cluster.
• Um cluster SageMaker HyperPod orquestrado pelo Amazon EKS. Consulte Criando um cluster SageMaker HyperPod com orquestração Amazon EKS.
• kubectl configurado contra o cluster, com o Kubeflow Training Operator instalado para que os recursos customizados PyTorchJob sejam reconhecidos.
• O FSx for Lustre CSI Driver instalado e um sistema de arquivos Amazon FSx for Lustre na mesma VPC e sub-rede dos nós do HyperPod.

Configuração e geração dos manifestos

Copie o arquivo de exemplo, preencha os valores do seu ambiente e execute o gerador:

cp config.yaml.example config.yaml
python generate.py

O script gera os seguintes arquivos: generated/storage.yaml (PersistentVolume e PersistentVolumeClaim para o FSx), generated/training-job.yaml (PyTorchJob do Kubeflow) e generated/viz-eks-webrtc-pod.yaml (pod de visualização via WebRTC, opcional).

Deploy do armazenamento e lançamento do treinamento

kubectl apply -f generated/storage.yaml
kubectl apply -f generated/training-job.yaml

Quando o job inicia, o Kubeflow Training Operator injeta as variáveis de ambiente PyTorch distribuído (MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE) em cada pod. O script de lançamento do contêiner as repassa ao torchrun, que gerencia o rendezvous e a configuração do grupo de processos. Com dois nós e 4 GPUs cada, o treinamento usa 8 GPUs no total. O manifesto também verifica no FSx a existência de um best_agent.pt de execuções anteriores e, se encontrado, retoma o treinamento a partir daquele ponto.

Rastreamento de experimentos com SageMaker managed MLflow

As métricas de treinamento, parâmetros de execução (tarefa, iterações, seed) e o diretório do checkpoint final são encaminhados para o Amazon SageMaker managed MLflow como store de experimentos persistente e pesquisável. Métricas de sistema como utilização de GPU e uso de CPU/memória são amostradas pela própria thread de background do MLflow. A habilitação do MLflow é opcional: deixar o MLFLOW_TRACKING_URI vazio (padrão) faz o script de treinamento ignorar todas as chamadas MLflow.

Treinamento com SageMaker Training Jobs

Os SageMaker Training Jobs executam a mesma imagem por um ciclo de vida diferente: cada job provisiona as instâncias de GPU solicitadas, baixa a imagem do Amazon ECR, executa o entrypoint, faz upload de tudo que o script copiou para /opt/ml/model/ como model.tar.gz para o caminho de saída no S3 e encerra as instâncias.

Pré-requisitos dos Training Jobs

• Cota de serviço suficiente para o tipo de instância GPU. Solicite aumento pelo AWS Service Quotas.
• Uma role IAM que o SageMaker possa assumir para o job de treinamento, com permissões para baixar do Amazon ECR, ler o entrypoint do S3 e gravar artefatos de volta no S3. Consulte a documentação de execution role do SageMaker.
• Um bucket S3 para o script de entrypoint (entrada) e os artefatos de treinamento (saída).
• boto3 instalado localmente (pip install boto3).

Geração e lançamento

Após configurar a seção sagemaker_training no config.yaml, execute o gerador e lance o job:

python generate.py
python generated/launch-sm-training.py --iterations 1000

O launcher chama CreateTrainingJob com um nome de job com sufixo de timestamp, a imagem do Amazon ECR e a localização do entrypoint no S3. Em caso de sucesso, imprime o nome do job e um comando describe-training-job para monitorar o progresso.

Para monitorar o status do job:

aws sagemaker describe-training-job \
  --training-job-name <TRAINING-JOB-NAME> \
  --query '{Status: TrainingJobStatus, Secondary: SecondaryStatus}'

Os logs de treinamento são transmitidos para o Amazon CloudWatch Logs no grupo /aws/sagemaker/TrainingJobs, com um stream de log por instância. A mesma integração com MLflow descrita para o HyperPod se aplica aos Training Jobs.

Visualização de políticas treinadas

O repositório inclui um pod de visualização para o SageMaker HyperPod que transmite a interface gráfica do Isaac Sim diretamente para um navegador via WebRTC, usando o mesmo volume FSx dos jobs de treinamento para que qualquer checkpoint produzido no HyperPod possa ser reproduzido.

kubectl apply -f generated/viz-eks-webrtc-pod.yaml

Como o kubectl port-forward suporta apenas TCP e a mídia WebRTC exige UDP, o plugin krelay adiciona o encaminhamento UDP necessário. Para equipes que preferem um desktop Linux completo, o workshop NVIDIA Isaac Lab na AWS descreve como configurar uma instância Amazon EC2 GPU com NICE DCV para executar a mesma imagem interativamente via streaming de desktop remoto de baixa latência.

Considerações de custo e limpeza

As duas opções de computação têm perfis de custo diferentes. O SageMaker HyperPod é um cluster persistente: as instâncias são cobradas enquanto fazem parte do cluster. O FSx for Lustre cobra por hora por capacidade provisionada. Os SageMaker Training Jobs cobram apenas pelo tempo de execução de cada job — sem custo de computação ociosa entre execuções. Consulte as páginas de preços do SageMaker AI e do FSx for Lustre para as tarifas atuais.

Para limpeza do HyperPod:

kubectl delete pytorchjob isaaclab-h1
kubectl delete -f generated/viz-eks-webrtc-pod.yaml

Escale o grupo de instâncias GPU para zero entre sessões para pausar os custos de instância mantendo o cluster configurado, ou exclua o cluster inteiramente. Consulte Gerenciar um cluster SageMaker HyperPod. Os Training Jobs encerram automaticamente ao concluir ou falhar — nenhuma limpeza de computação é necessária.

Conclusão

À medida que as cargas de trabalho de Physical AI avançam para a produção, equipes de robótica precisam escalar o treinamento de políticas sem a sobrecarga operacional de gerenciar infraestrutura de computação. A solução publicada pela AWS demonstra como o SageMaker HyperPod e os SageMaker Training Jobs permitem executar treinamento distribuído do Isaac Lab em infraestrutura de GPU gerenciada, usando uma única imagem de contêiner e uma configuração compartilhada entre os dois modelos de computação.

O HyperPod oferece clusters de GPU persistentes com treinamento resiliente de longa duração. Os Training Jobs oferecem execuções efêmeras e sob demanda adequadas para experimentos e varreduras de hiperparâmetros. Ambos executam a mesma imagem e a mesma invocação do torchrun do trainer skrl — alternar entre eles é apenas uma mudança de configuração.

Para começar, explore o repositório para lançar seu primeiro treinamento do H1 e estender o padrão para outras tarefas do Isaac Lab (manipulação de humanoides, quadrúpedes, mãos dexterosas). Para saber mais, consulte a documentação do Amazon SageMaker HyperPod e a documentação do Amazon SageMaker Training Jobs.

Fonte

Scale Robot Reinforcement Learning with NVIDIA Isaac Lab on Amazon SageMaker AI (https://aws.amazon.com/blogs/machine-learning/scale-robot-reinforcement-learning-with-nvidia-isaac-lab-on-amazon-sagemaker-ai/)