Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning

Autores: Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, Gavriel State · arXiv 2021

Licenças

Acesso aberto
Artigo: CC BY 4.0

Tudo o que segue se baseia no artigo completo sob a licença CC BY 4.0; sua figura é mostrada com crédito aos autores.

Em resumo

O Isaac Gym executa a simulação PhysX, as observações, as recompensas e o treinamento PPO em uma única GPU por meio de uma API de tensores do PyTorch; a reorientação de cubo com a Shadow Hand cai de 30 h em um cluster de CPU para menos de 1 h.

Por que importa

Tarefas que antes exigiam um cluster de CPU podem ser treinadas em uma única GPU em minutos ou horas, o que encurta as iterações na pesquisa em locomoção e manipulação destra. Os autores também transferem políticas treinadas dessa forma para um ANYmal real em terreno irregular e para um TriFinger real remoto, o que indica que a física de contato simulada sustenta trabalho sim-to-real.

Problema

O RL para robôs geralmente simula a física em CPUs e treina as redes em GPUs, então os dados se movem entre dispositivos a cada passo e a velocidade da simulação limita o treinamento. O trabalho da OpenAI com o cubo e a Shadow Hand, por exemplo, usou 384 máquinas com 6144 núcleos de CPU mais 8 GPUs V100 por 30 horas, e sua continuação com o cubo de Rubik usou quase 30.000 núcleos de CPU.

Solução proposta

O Isaac Gym executa o NVIDIA PhysX na GPU com dois acréscimos: os resultados da simulação não são mais trazidos para a CPU após cada passo, e uma API direta de GPU lê o estado e aplica os controles em buffers da GPU. Esses buffers são encapsulados como tensores do PyTorch sem cópia, de modo que milhares de ambientes em paralelo avançam, são observados, recebem recompensa e são treinados (com a implementação de PPO do rl_games) em uma única GPU, com robôs carregados a partir de URDF ou MJCF.

Principais ideias

Manter na GPU o estado da física, o cálculo de observações e recompensas e o treinamento da política elimina as cópias CPU-GPU que limitam os pipelines de RL convencionais.
Mais ambientes em paralelo encurtam o treinamento até que o horizonte por ambiente fique curto demais: o Humanoid não ganha nada e aprende marchas piores acima de 4096 ambientes com horizonte 32, mas aprende com 8192 e 16384 ambientes quando os horizontes são dobrados.
A vazão da simulação cai com a complexidade do modelo: até cerca de 700K passos de ambiente por segundo para o Ant, 200K para o Humanoid e 150K para a Shadow Hand em uma A100.
O ator-crítico assimétrico (asymmetric actor-critic), em que o crítico recebe o estado privilegiado do simulador e a política apenas observações disponíveis em um robô real, é usado nas configurações Shadow Hand OpenAI e TriFinger para permitir a transferência sim-to-real.
Um punho móvel ajuda na reorientação: a Allegro Hand, com punho fixo e dedos um pouco mais longos, alcança menos sucessos consecutivos que a Shadow Hand apesar de ter menos graus de liberdade.

Principais resultados

Em uma única A100: o Ant passa de uma recompensa de 3000 em 20 s, o Humanoid corre (recompensa 5000) em menos de 4 min, o ANYmal segue alvos de velocidade em menos de 2 min e o empilhamento de cubos com o Franka converge em menos de 25 min.
Shadow Hand com observações da OpenAI e randomização de domínio: mais de 20 sucessos consecutivos em menos de 1 h (feed-forward) e 37 em pouco menos de 6 h (LSTM), contra 30 h e cerca de 17 h da OpenAI com 6144 núcleos de CPU mais 8 GPUs V100; a melhor semente LSTM chega a 37 em 2,5 h.
Shadow Hand Standard (sem randomização): 20 sucessos consecutivos em menos de 35 min; aumentar os ambientes de 256 para 16384 reduz o treinamento de 5×10⁴ s (cerca de 14 h) para 3×10³ s (cerca de 1 h).
Animação de personagem humanoide com AMP: cerca de 39 milhões de amostras em cerca de 6 min com 4096 ambientes, contra cerca de 30 h com 16 núcleos de CPU no PyBullet, relatado como uma melhoria de 300×.
Sim-to-real: a política de terreno irregular do ANYmal é treinada em uma RTX A6000 e transferida para o robô real em menos de 20 min; o TriFinger real atinge uma taxa de sucesso média de 55 %, com a transferência feita remotamente.
Números de aceleração conflitantes: o resumo da v1 indica 1-2 ordens de grandeza em relação ao treinamento com simulador em CPU, enquanto a introdução da v1 e a Fig. 2 indicam 100×-1000×; o resumo atual no arXiv indica 2-3 ordens de grandeza.

Limitações

Os experimentos com a Shadow Hand são apenas em simulação; ao contrário do trabalho da OpenAI, não se tenta a transferência sim-to-real para a mão, e as randomizações omitem o atraso das ações e a folga (backlash) dos motores.
A Shadow Hand Standard usa observações que podem não estar disponíveis em um robô real, por isso os autores afirmam que ela não é adequada para transferência sim-to-real.
O Ingenuity é um modelo simplificado: as forças são aplicadas diretamente aos dois rotores em vez de simular a aerodinâmica.
Com o motor FleX está disponível apenas uma funcionalidade limitada da API de tensores, a configuração da cena roda na CPU e os exemplos incluídos usam PyTorch, sendo que a integração com TensorFlow exige personalização adicional.
Acrescentar ambientes deixa de ajudar quando os horizontes ficam curtos: o Ant não mostra ganho de 8192 para 16384 ambientes, e o Humanoid acima de 4096 ambientes com horizonte 32 treina mais devagar e aprende marchas subótimas.

Impacto na área

1469 citações. O artigo disponibiliza o Isaac Gym como download da NVIDIA (developer.nvidia.com/isaac-gym), incluindo uma implementação básica de PPO e um sistema de tarefas de RL, e hospeda resultados e vídeos no site do projeto.

Resumo

O Isaac Gym oferece uma plataforma de aprendizado de alto desempenho para treinar políticas para uma ampla variedade de tarefas robóticas diretamente na GPU. Tanto a simulação física quanto o treinamento da política de rede neural ficam na GPU e se comunicam passando dados diretamente dos buffers de física para tensores do PyTorch, sem nunca passar por gargalos da CPU. Isso leva a tempos de treinamento extremamente rápidos para tarefas robóticas complexas em uma única GPU, com melhorias de 2-3 ordens de grandeza em comparação com o treinamento de RL convencional que usa um simulador baseado em CPU e a GPU para as redes neurais. Hospedamos os resultados e vídeos em https://sites.google.com/view/isaacgym-nvidia e o isaac gym pode ser baixado em https://developer.nvidia.com/isaac-gym.

Open Source Robotics