NVIDIA Warp e MuJoCo Revolucionam Simulação Robótica com Escala em GPU

A simulação robótica clássica com MuJoCo oferece execução rápida baseada em CPU para desenvolvimento, teste e controle de robôs, parallelizando amostras entre núcleos de CPU. No entanto, com o crescimento das cargas de trabalho de aprendizado, a pergunta muda de quão rápido um mundo pode ser executado para quantos mundos podem ser executados simultaneamente. A aceleração via GPU torna possível avançar esses mundos em grandes lotes, mantendo os dados de simulação e aprendizado próximos ao dispositivo.

MuJoCo Warp (MJWarp), construído sobre NVIDIA Warp, leva modelos compatíveis do MuJoCo para esse regime de escala GPU. Neste artigo, um braço seguidor SO-101 foi migrado de um fluxo de trabalho MuJoCo convencional para até 2.048 ambientes paralelos MJWarp, examinando a tecnologia e as etapas de validação que tornam essa transição possível.

A NVIDIA Warp é uma estrutura Python para escrever kernels de alta performance acelerados por GPU. Ela permite que desenvolvedores criem kernels estaticamente tipados em Python e compilem para execução em CPU ou CUDA. A primeira execução compila e armazena em cache um módulo nativo; execuções posteriores reaproveitam esse módulo. A linguagem dos kernels é um subconjunto orientado a desempenho de Python, enquanto o Python convencional permanece responsável pela configuração, alocação e orquestração de lançamento.

O MuJoCo e o MJWarp podem executar o mesmo robô e tarefa compatíveis, mas organizam o trabalho de forma diferente. MuJoCo é ideal para desenvolver e inspecionar um ou alguns mundos em CPU. MJWarp é uma implementação da NVIDIA Warp do pipeline de física do MuJoCo que coloca o modelo e um lote de estados independentes nas GPUs NVIDIA; uma única chamada a mjw.step avança todo o lote.

O valor do MJWarp não está necessariamente em um passo mais rápido para um único mundo, mas na capacidade de avançar centenas ou milhares juntos, fornecendo ao GPU trabalho paralelo suficiente para melhorar a taxa agregada de processamento (total de passos mundiais por segundo). Isso favorece aprendizado por reforço e amostragem em larga escala, onde coletar experiências importa mais que minimizar a latência de um ambiente.

Para migrar workflows MuJoCo para MJWarp, recomenda-se utilizar mjw.make_data() quando se pretende um estado padrão ou novo. Quando o estado MuJoCo exato deve atravessar o limite de migração, usa-se mjw.put_data().

A alocação de recursos em lote exige definição de parâmetros como tamanho do lote, captura de gráficos CUDA e ajuste de buffers de contato e restrições. O comando mjwarp-testspeed auxilia na medição de alocações e detecção de transbordamentos, especialmente em visualizadores MJWarp.

  • CUDA graph capture: mjw.step envolve muitos lançamentos de kernel. Capture uma vez e reproduza frequentemente.
  • Ajuste preciso de nconmax, naconmax e njmax, já que memória e trabalho escalam com eles.
  • Após dimensionar buffers de contato e restrições, teste limites de iteração do solucionador sem alterar o comportamento da tarefa.
  • Geometrias e configurações CCD podem aumentar consumo de memória. Use nccdmax e naccdmax para reduzir alocações de buffers CCD conforme necessário.

O cenário utilizado no experimento inclui um braço SO-101, uma mesa e dois cubos para empilhamento, definidos como MJCF comum. O cubo de 44 mm de aresta foi dimensionado usando valores de meia-extensão (size=\"0.022 …\"), servindo também como referência para limiares de sucesso.

Para instalar e experimentar: use pip install mujoco-warp, execute mjwarp-viewer path/to/scene.xml e consulte o tutorial Colab para exemplos interativos. Recomenda-se a versão ≥ 1.15 do Warp para funcionalidades de determinismo em GPU.

Com informações de: Hugging Face Blog