Rode 2.048 Simulações MuJoCo em Paralelo numa Única GPU com NVIDIA Warp

Por que simular em escala de GPU? 🚀

O MuJoCo clássico roda rápido na CPU e paraleliza amostragem entre núcleos. Mas cargas modernas de aprendizado de robô precisam de milhares de mundos avançando ao mesmo tempo. O NVIDIA Warp e sua implementação MuJoCo, o MJWarp, permitem manter o mesmo modelo MJCF enquanto avança centenas ou milhares de estados independentes por chamada de GPU.

Neste tutorial, vamos migrar uma cena pick-and-place do SO-101 de MuJoCo CPU para 2.048 ambientes paralelos no MJWarp. Você vai validar um mundo primeiro, depois escalar e medir do jeito certo.

Diferença chave: latência é o tempo de um passo; throughput agregado é o total de world-steps por segundo. O MJWarp otimiza o segundo.

Para uma visão mais ampla de estratégias de utilização de GPU, veja nosso blueprint de ROI e custos em IA.

Passo 1: Valide um mundo na CPU

Comece com o loop familiar do MuJoCo. A tarefa: pegar um cubo de 44 mm e empilhar sobre o cubo azul.

import mujoco

mjm = mujoco.MjModel.from_xml_path("scene_pick_place.xml")
mjd = mujoco.MjData(mjm)
fps = 50
sim_substeps = 10
frame_dt = 1.0 / fps
mjm.opt.timestep = frame_dt / sim_substeps  # 0.002 s

for _ in range(600):
    ctrl = controller.step(mjm, mjd, frame_dt)
    for _ in range(sim_substeps):
        mjd.ctrl[: mjm.nu] = ctrl
        mujoco.mj_step(mjm, mjd)

Critérios de sucesso: xy_err <= 0.015 e 0.035 <= dz <= 0.055 após assentamento.

Passo 2: Migre para o MJWarp com um mundo

import warp as wp
import mujoco_warp as mjw

wp.init()
device = wp.get_device()
m = mjw.put_model(mjm)  # levanta erro se houver recursos não suportados
d = mjw.make_data(mjm, nworld=1, nconmax=128, njmax=300)

wp.copy(d.qpos, wp.array(mjd.qpos[None, :], dtype=wp.float32, device=device))
wp.copy(d.qvel, wp.array(mjd.qvel[None, :], dtype=wp.float32, device=device))
wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device))
mjw.forward(m, d)

Todo array no device tem uma dimensão líder de mundo. Mantenha o host no loop para validação:

def simulate_frame():
    ctrl = controller.step(mjm, mjd, frame_dt)
    for _ in range(sim_substeps):
        mjd.ctrl[: mjm.nu] = ctrl
        wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device))
        mjw.step(m, d)
    mjd.qpos[:] = d.qpos.numpy()[0]
    mjd.qvel[:] = d.qvel.numpy()[0]
    mujoco.mj_forward(mjm, mjd)  # atualiza quantidades derivadas

Fique de olho em avisos de overflow: aumente nconmax/njmax se o MJWarp reportar capacidade insuficiente.

Passo 3: Escale para 2.048 mundos e capture um CUDA Graph

import numpy as np

nworld = 2048
d = mjw.make_data(mjm, nworld=nworld, nconmax=128, njmax=300)
wp.copy(d.qpos, wp.array(np.tile(mjd.qpos, (nworld, 1)), dtype=wp.float32, device=device))
wp.copy(d.qvel, wp.array(np.tile(mjd.qvel, (nworld, 1)), dtype=wp.float32, device=device))
wp.copy(d.ctrl, wp.array(np.tile(mjd.ctrl, (nworld, 1)), dtype=wp.float32, device=device))
mjw.forward(m, d)

with wp.ScopedCapture() as capture:
    mjw.step(m, d)
step_graph = capture.graph

Passo 4: Benchmark do jeito certo

import time

for _ in range(10):
    wp.capture_launch(step_graph)
wp.synchronize()

t0 = time.perf_counter()
for _ in range(200):
    wp.capture_launch(step_graph)
wp.synchronize()  # sem isso você mede a fila, não o trabalho
elapsed = time.perf_counter() - t0
print(f"{200 * nworld / elapsed:,.0f} world-steps/segundo")

Reporte tanto world-steps por segundo quanto milissegundos por passo em lote, junto com o tamanho do batch. Comparar latência de um mundo só não estabelece throughput.

Limitações e cuidados ⚠️

• Latência de um mundo pode piorar. Overhead de launch e cópias host-device podem deixar um mundo na GPU mais lento que o MuJoCo na CPU. O valor do MJWarp está no batching.
• Dimensionar buffers é crítico. Estourar nconmax/njmax invalida o rollout mesmo que a execução continue com aviso. Dimensione pelo momento mais pesado de contato (as duas garras e a mesa tocando o cubo).
• Determinismo é opt-in. Warp 1.15+ suporta execução determinística, mas troca performance por ordenação reprodutível; não é garantido por padrão num rollout inteiro do MJWarp.
• Primeiro launch compila kernels. Espere vários minutos de delay inicial na primeira execução.

Próximos passos

• Porte o mesmo MJCF para o Newton usando newton.solvers.SolverMuJoCo para suporte multi-solver e integração com Isaac Lab.
• Experimente mjlab ou MuJoCo Playground para treinar políticas sobre a física do MJWarp.
• Explore kernels diferenciáveis do Warp para simulação baseada em gradiente e otimização de design.

Para escalar cargas de GPU além da simulação, veja nosso deep dive sobre maximizar utilização de GPU para inferência de LLM com NVIDIA Runai e NIM.

Conclusão

O MJWarp permite manter seu modelo MJCF e escalar de um mundo na CPU para milhares de mundos na GPU. Valide a paridade primeiro, dimensione buffers de contato pelo pior momento da tarefa e meça throughput agregado com sincronização adequada. A migração é pequena em superfície de API, mas exige disciplina em medição e planejamento de capacidade.

Comece com pip install mujoco-warp e o tutorial oficial no Colab, depois escale sua própria cena. 🔥