Engenheiros da NVIDIA detalham escalonamento em GPU para simulação de robótica no MuJoCo usando o Warp

Escalando a simulação robótica de mundos únicos na CPU para lotes massivos na GPU

Em 23 de setembro de 2026, os pesquisadores da NVIDIA Johnny Nuñez Cano, Asier Arranz, Rishabh Chadha e Ben Oliveri publicaram um guia de implementação mostrando como desenvolvedores de robótica podem migrar simulações tradicionais do MuJoCo baseadas em CPU para cargas de trabalho de alto rendimento em GPU. À medida que os pipelines de IA física e aprendizado por reforço se expandem, a velocidade de treinamento depende cada vez mais da execução de centenas ou milhares de ambientes simultâneos, em vez de focar apenas na otimização da latência de um único mundo.

Para preencher essa lacuna, a NVIDIA introduziu um fluxo de trabalho utilizando o MuJoCo Warp (MJWarp), uma implementação acelerada por GPU do motor de física MuJoCo desenvolvida sobre o NVIDIA Warp. Como o segundo artigo da série "State of Simulation for Physical AI" da NVIDIA, o guia demonstra a transição de um braço robótico seguidor SO-101 executando uma tarefa de empilhamento de blocos de um ambiente padrão em Python na CPU diretamente para 2.048 estados de simulação paralelizados em uma GPU NVIDIA.

Portando o braço SO-101: dimensionamento de buffers, CUDA Graphs e validação de paridade

Na base dessa arquitetura está o NVIDIA Warp, um framework baseado em Python que compila kernels em Python de tipagem estática diretamente para código CUDA nativo, com suporte a compilação just-in-time (JIT), fusão de kernels, diferenciação automática via wp.Tape e modos de execução determinísticos introduzidos na versão 1.15. O MJWarp aplica esse framework ao motor de física do MuJoCo, processando descrições de cena padrão em XML MJCF enquanto transfere os arrays para a GPU com uma dimensão de lote (batch dimension) inicial adicionada.

A migração tem início com uma tarefa de pick-and-place do SO-101 configurada com frequência de controle de 50 Hz e 10 subetapas de física por quadro, visando um intervalo de tempo (timestep) de física de 0,002 segundos. O objetivo do braço é pegar um cubo vermelho de 44 mm e posicioná-lo sobre um cubo azul de tamanho idêntico. A transição desse fluxo de trabalho para o MJWarp exige uma mudança de API: os desenvolvedores carregam o modelo usando mjw.put_model(), alocam o estado residente na GPU com mjw.make_data() ou mjw.put_data() e avançam todos os mundos paralelos simultaneamente via mjw.step().

Um requisito operacional crítico é o gerenciamento dos buffers de contato e restrições (constraints) na GPU. Os desenvolvedores precisam definir limites de capacidade como nconmax (máximo de contatos por ambiente), naconmax (teto global de contatos) e njmax (máximo de restrições por ambiente). Como o MJWarp emite apenas avisos em vez de interromper a execução imediatamente quando as colisões em fase estreita (narrowphase) excedem a capacidade, os experimentos correm o risco de invalidação silenciosa caso não sejam verificados. A NVIDIA recomenda dimensionar esses limites no momento de maior contato da tarefa — por exemplo, quando ambas as garras e a mesa tocam um cubo simultaneamente — e diagnosticar o uso de memória com utilitários como mjwarp-testspeed --measure_alloc.

Com a paridade física de um único mundo validada em relação à base de referência na CPU, os ambientes podem ser escalados para 2.048 mundos ou mais. Para eliminar o overhead de despacho (dispatch latency) ao longo de milhares de etapas em lote, os desenvolvedores envolvem mjw.step() em um bloco wp.ScopedCapture(), criando CUDA Graphs reutilizáveis. O perfilamento de desempenho deve levar em consideração a execução assíncrona da GPU, inserindo chamadas a wp.synchronize() antes e depois de loops cronometrados e mantendo os dados da simulação na memória da placa de vídeo, evitando transferências para a memória host via .numpy() durante as etapas de rollout.

Expandindo a pilha de simulação de IA física para Newton e Isaac Lab

A transição para o MJWarp evidencia uma divisão cada vez mais clara no setor entre o controle preditivo baseado em modelo (MPC) voltado a robôs individuais e o aprendizado por reforço em lotes massivos. Enquanto o MuJoCo clássico em CPU segue como o padrão para teleoperação, depuração em ambiente único e MPC de baixa latência, os ambientes massivamente paralelizados em GPU atendem às demandas de coleta de dados em larga escala.

O MJWarp foi arquitetado para se integrar a diversos ecossistemas modernos de IA física. Ele atua como backend de execução para o mjlab (que integra o MJWarp diretamente ao PyTorch), para o MuJoCo Playground por meio da implementação Warp do MJX, e para pipelines de aprendizado mais amplos. A NVIDIA adiantou que a próxima edição da série de simulações demonstrará como importar esse mesmo ambiente do SO-101 para seu framework multi-solver, o Newton, no qual o MJWarp operará como o solucionador de corpos rígidos subjacente (newton.solvers.SolverMuJoCo), conectando as cenas diretamente aos fluxos de trabalho de treinamento do Isaac Lab.