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

Licencias

• Acceso abierto
• Artículo: CC BY 4.0

Todo lo que sigue se basa en el artículo completo bajo su licencia CC BY 4.0; su figura se muestra con crédito a los autores.

En resumen

Isaac Gym ejecuta la simulación PhysX, las observaciones, las recompensas y el entrenamiento PPO en una sola GPU mediante una API de tensores de PyTorch; la reorientación de un cubo con Shadow Hand pasa de 30 h en un clúster de CPU a menos de 1 h.

Por qué importa

Tareas que antes necesitaban un clúster de CPU se pueden entrenar en una sola GPU en minutos u horas, lo que acorta las iteraciones en la investigación sobre locomoción y manipulación diestra. Los autores también transfieren políticas entrenadas de esta forma a un ANYmal real en terreno irregular y a un TriFinger real remoto, lo que indica que la física de contacto simulada sirve para trabajo sim-to-real.

Problema

El RL para robots suele simular la física en CPU y entrenar las redes en GPU, de modo que los datos se mueven entre dispositivos en cada paso y la velocidad de simulación limita el entrenamiento. El trabajo de OpenAI con el cubo y la Shadow Hand, por ejemplo, usó 384 máquinas con 6144 núcleos de CPU más 8 GPU V100 durante 30 horas, y su continuación con el cubo de Rubik usó casi 30.000 núcleos de CPU.

Solución propuesta

Isaac Gym ejecuta NVIDIA PhysX en la GPU con dos añadidos: los resultados de la simulación ya no se traen a la CPU tras cada paso, y una API directa de GPU lee el estado y aplica los controles en búferes de GPU. Estos búferes se envuelven como tensores de PyTorch sin copia, de modo que miles de entornos en paralelo avanzan, se observan, reciben recompensa y se entrenan (con la implementación de PPO de rl_games) en una sola GPU, con robots cargados desde URDF o MJCF.

Ideas clave

• Mantener en la GPU el estado de la física, el cálculo de observaciones y recompensas y el entrenamiento de la política elimina las copias CPU-GPU que limitan los pipelines de RL convencionales.
• Más entornos en paralelo acortan el entrenamiento hasta que el horizonte por entorno se vuelve demasiado corto: Humanoid no gana nada y aprende marchas peores por encima de 4096 entornos con horizonte 32, pero aprende con 8192 y 16384 entornos cuando se duplican los horizontes.
• El rendimiento de la simulación cae con la complejidad del modelo: hasta unos 700K pasos de entorno por segundo para Ant, 200K para Humanoid y 150K para Shadow Hand en una A100.
• El actor-crítico asimétrico (asymmetric actor-critic), en el que el crítico recibe el estado privilegiado del simulador y la política solo observaciones disponibles en un robot real, se usa en las configuraciones Shadow Hand OpenAI y TriFinger para permitir la transferencia sim-to-real.
• Una muñeca móvil ayuda a la reorientación: la Allegro Hand, con la muñeca fija y dedos algo más largos, logra menos éxitos consecutivos que la Shadow Hand pese a tener menos grados de libertad.

Resultados principales

• En una sola A100: Ant supera una recompensa de 3000 en 20 s, Humanoid corre (recompensa 5000) en menos de 4 min, ANYmal sigue objetivos de velocidad en menos de 2 min y el apilado de cubos con Franka converge en menos de 25 min.
• Shadow Hand con observaciones de OpenAI y aleatorización de dominio: más de 20 éxitos consecutivos en menos de 1 h (feed-forward) y 37 en poco menos de 6 h (LSTM), frente a 30 h y unas 17 h de OpenAI con 6144 núcleos de CPU más 8 GPU V100; la mejor semilla LSTM llega a 37 en 2,5 h.
• Shadow Hand Standard (sin aleatorización): 20 éxitos consecutivos en menos de 35 min; pasar de 256 a 16384 entornos reduce el entrenamiento de 5×10⁴ s (unas 14 h) a 3×10³ s (aproximadamente 1 h).
• Animación de personaje humanoide con AMP: unos 39 millones de muestras en unos 6 min con 4096 entornos, frente a unas 30 h con 16 núcleos de CPU en PyBullet, presentado como una mejora de 300×.
• Sim-to-real: la política de terreno irregular de ANYmal se entrena en una RTX A6000 y se transfiere al robot real en menos de 20 min; el TriFinger real alcanza una tasa de éxito media del 55 %, con la transferencia hecha de forma remota.
• Cifras de aceleración contradictorias: el resumen de la v1 indica 1-2 órdenes de magnitud respecto al entrenamiento con simulador en CPU, mientras que la introducción de la v1 y la Fig. 2 indican 100×-1000×; el resumen actual de arXiv indica 2-3 órdenes de magnitud.

Limitaciones

• Los experimentos con Shadow Hand son solo en simulación; a diferencia del trabajo de OpenAI, no se intenta la transferencia sim-to-real para la mano, y las aleatorizaciones omiten el retardo de las acciones y el juego (backlash) de los motores.
• Shadow Hand Standard usa observaciones que pueden no estar disponibles en un robot real, por lo que los autores afirman que no es adecuada para la transferencia sim-to-real.
• Ingenuity es un modelo simplificado: las fuerzas se aplican directamente a los dos rotores en lugar de simular la aerodinámica.
• Con el motor FleX solo está disponible una funcionalidad limitada de la API de tensores, la configuración de la escena se ejecuta en la CPU y los ejemplos incluidos usan PyTorch, mientras que la integración con TensorFlow requiere personalización adicional.
• Añadir entornos deja de ayudar cuando los horizontes se acortan: Ant no muestra ganancia de 8192 a 16384 entornos, y Humanoid por encima de 4096 entornos con horizonte 32 entrena más despacio y aprende marchas subóptimas.

Impacto en el campo

1469 citas. El artículo publica Isaac Gym como descarga de NVIDIA (developer.nvidia.com/isaac-gym), con una implementación básica de PPO y un sistema de tareas de RL, y aloja resultados y vídeos en el sitio del proyecto.

Resumen

Isaac Gym ofrece una plataforma de aprendizaje de alto rendimiento para entrenar políticas para una amplia variedad de tareas robóticas directamente en la GPU. Tanto la simulación física como el entrenamiento de la política de red neuronal residen en la GPU y se comunican pasando datos directamente de los búferes de física a tensores de PyTorch, sin pasar nunca por cuellos de botella de la CPU. Esto da lugar a tiempos de entrenamiento rapidísimos para tareas robóticas complejas en una sola GPU, con mejoras de 2-3 órdenes de magnitud en comparación con el entrenamiento de RL convencional que usa un simulador basado en CPU y la GPU para las redes neuronales. Alojamos los resultados y los vídeos en https://sites.google.com/view/isaacgym-nvidia e isaac gym se puede descargar en https://developer.nvidia.com/isaac-gym.

Enlaces

• arXiv 2108.10470: https://arxiv.org/abs/2108.10470
• DOI https://doi.org/10.48550/arxiv.2108.10470: https://doi.org/https://doi.org/10.48550/arxiv.2108.10470
• PDF: https://arxiv.org/pdf/2108.10470v1
• Página del proyecto: https://sites.google.com/view/isaacgym-nvidia