NVIDIA Warp와 MuJoCo Warp(MJWarp)로 로봇을 위한 2,048세계 GPU 시뮬레이션 가능

TL;DR

NVIDIA Warp와 그 MuJoCo 호환 레이어인 MJWarp를 사용하면 단일 GPU에서 수천 개의 로봇 시뮬레이션 세계를 배치 처리할 수 있으며, 전통적인 CPU 기반 MuJoCo보다 주어진 시간에 처리 가능한 총량이 수배 이상 증가합니다.

스택 개요

계층 역할
NVIDIA Warp CUDA로 컴파일되는 파이썬 커널 언어로, 자동 미분을 제공하고 PyTorch/JAX와 상호 운용 가능
MJWarp MJCF 형식을 유지하면서 배치 단계 처리를 지원하는 MuJoCo 물리 시뮬레이션의 GPU 구현
당신의 시나리오 (SO‑101) 표준 Menagerie / Robot Studio 자산과 작업 기하학
Next (Newton / Isaac Lab) 다중 솔버 API, USD 자산, 센서, 매니저, 학습 루프

각 구성 요소를 사용할 때

필요 사항 사용할 것
단일 로봇 MPC 또는 원격 조작 전통적인 MuJoCo CPU
최대 MuJoCo 물리 처리량 MJWarp (또는 mjlab)
JAX 기반 학습 레시피 MuJoCo Playground / MJX (구현 = warp)
다중 솔버 + Isaac Lab 통합 Newton (미래 게시물)

NVIDIA Warp: 최소한의 GPU 커널

Warp는 JIT 컴파일로 네이티브 CUDA로 변환되는 정적 타입 파이썬 커널을 작성할 수 있게 해줍니다. 아래 예제는 중력 하에서 점의 위치를 적분하며, 커널 로직을 변경하지 않고 두 점에서 수백만 점으로 확장할 수 있습니다.

import numpy as np, warp as wp

@wp.kernel
def integrate(positions: wp.array[wp.vec3],
                  velocities: wp.array[wp.vec3],
                  dt: float):
    i = wp.tid()
    velocities[i] += wp.vec3(0.0, 0.0, -9.81) * dt
    positions[i] += velocities[i] * dt

wp.init()
device = "cuda:0" if wp.is_cuda_available() else "cpu"
start = np.array([[0.0, 0.0, 0.5], [0.2, 0.0, 0.5]], dtype=np.float32)
positions = wp.array(start, dtype=wp.vec3, device=device)
velocities = wp.zeros_like(positions)

wp.launch(integrate, dim=len(start), inputs=[positions, velocities, 0.01], device=device)
wp.synchronize_device(device)
print(positions.numpy())

로봇 공학에서의 중요성

  • 명시적 병렬 작업 – wp.tid()는 논리적 스레드(점, 접촉, 본체, 또는 세계)를 식별합니다.
  • 명시적 장치 배열 – 배열은 선택된 장치에 존재합니다; .numpy()는 호스트 복사 작업을 트리거하며, DLPack 어댑터를 통해 PyTorch/JAX와 0복사 공유가 가능합니다.
  • 조합 가능한 실행 – 여러 커널을 CUDA 그래프에 캡처할 수 있어 실행 오버헤드를 줄일 수 있습니다.

미분 가능성 및 결정론성 (Warp 기능)

  • 미분 가능한 커널 – wp.Tape는 전방 실행을 기록하고 backward()에서 자동으로 인접값을 계산합니다. 이는 미분 가능한 기하학, CFD 및 사용자 정의 물리 파이프라인을 가능하게 합니다.
  • 결정론적 실행 – Warp 1.15에서 도입; 결정론적 모드는 약간의 성능 저하를 희생하여 재현 가능한 원자적 순서를 보장하며, 회귀 테스트에 유용합니다. 이 기능들은 SO‑101 안내서에서는 선택 사항이며 필수는 아닙니다.

MuJoCo Warp (MJWarp)란?

MJWarp는 Warp 위에 구축된 MuJoCo 물리 파이프라인의 GPU 기반 구현입니다. 동일한 MJCF 모델을 수용하며, 모델과 독립적인 상태의 배치를 GPU에 로드한 후 단일 mjw.step 호출로 전체 배치를 한 번에 진행합니다.

  • 지연 vs 처리량 – MJWarp는 단일 세계 단계를 더 빠르게 보장하지 않습니다. 장점은 집합 처리량입니다: 병렬로 실행되는 많은 세계에서 초당 완료되는 세계 단계의 총 수입니다. 이는 강화 학습 및 대규모 샘플링의 핵심 지표입니다.

MuJoCo 시나리오를 MJWarp로 마이그레이션하기

이 가이드는 세 단계를 거칩니다: CPU 기준 설정, 단일 세계 GPU 실행 검증, 2,048세계로 확장.

1. MuJoCo CPU 기준 설정

  • mujoco.MjModel.from_xml_path로 SO‑101 피크 앤 플레이스 MJCF(표준 Menagerie 자산)를 로드합니다.
  • 50Hz 제어 루프와 10단계 물리 하위 단계(timestep = 0.002 s)를 사용합니다.
  • 큐브가 정지한 후 수평 오차 ≤ 0.015m이고 수직 간격이 0.035m에서 0.055m 사이인지 확인하여 작업 성공 여부를 검증합니다.

2. 단일 세계 MJWarp 일치성 검증

import warp as wp, mujoco_warp as mjw
wp.init()
device = wp.get_device()

# 모델 업로드 및 단일 세계 배치 할당
m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax)

# CPU 실행에서 상태 시드
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)
  • 호스트 루프는 CPU 워크플로우를 반영하며, 각 하위 단계마다 제어를 장치로 복사하고 qpos/qvel을 다시 가져와 후속 계산에 사용합니다.
  • 이 경로는 성능 벤치마크가 아닙니다; GPU 시뮬레이션이 CPU 결과를 재현하는지 검증하기 위한 것입니다.

3. 접촉 및 제약 버퍼 크기 조정

  • MJWarp는 세계당 접촉(nconmax) 및 제약(njmax) 버퍼를 미리 할당합니다. 충분한 용량이 없으면 오버플로우 경고가 발생하고 궤적에 오류가 생길 수 있습니다.
  • mjwarp-testspeed --measure_alloc를 사용해 실제 사용량을 측정하고 한도를 조정합니다 (예: SO‑101 작업에 대해 nconmax=128, njmax=300).

4. 2,048세계로 확장

nworld = 2_048
d = mjw.make_data(m, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax)

# 초기 상태를 배치 전체에 복제
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)

# 단계를 CUDA 그래프에 캡처하여 낮은 오버헤드로 재생
with wp.ScopedCapture() as cap:
    mjw.step(m, d)
step_graph = cap.graph
  • 캡처된 그래프는 동일한 버퍼를 재사용하며, 재생 사이에 d.ctrl만 업데이트하면 됩니다.

5. 집합 처리량 측정

# 워밍업 (컴파일, 할당)
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()
elapsed = time.perf_counter() - t0
world_steps = 200 * nworld
print(f"{world_steps/elapsed:,.0f} world‑steps/second")
  • 배치 크기와 함께 세계 단계/초와 배치 단계당 밀리초를 모두 보고하세요.
  • 제공된 scaling_study.py 스크립트를 사용해 배치 크기(1, 64, 1,024, 2,048, 8,192)를 스윕하고 처리량이 포화되는 지점을 관찰하세요.

시작하기

  • Warp – pip install warp-lang (결정론적 모드는 ≥ 1.15 필요). python -m warp.examples.browse로 예제 탐색.
  • MJWarp – pip install mujoco-warp; mjwarp-viewer path/to/scene.xml로 뷰어 실행.
  • SO‑101 리소스 – NVIDIA의 Sim‑to‑Real 과정 및 Menagerie 자산.
  • MJWarp 위에 학습하기 – mjlab, MuJoCo Playground (MJX impl='warp'), 또는 곧 출시될 Newton 통합을 통한 Isaac Lab.

다음 단계

다음 게시물에서는 MJWarp를 Newton 물리 엔진 내부에 통합하여 다중 형식 자산, 교체 가능한 솔버, 센서, 그리고 Isaac Lab 학습 파이프라인을 추가할 예정입니다. 동일한 SO‑101 작업(옵션으로 reBot 변형 포함)을 사용해 추가 통합 단계를 설명할 것입니다.

참고 자료

  • 물리적 AI를 위한 시뮬레이션의 현황 – 시리즈의 첫 번째 블로그.
  • NVIDIA Warp GitHub 및 문서 (v1.15.0에서 GPU 결정론성 추가).
  • MuJoCo Warp 저장소 및 공식 문서.
  • mjlab (arXiv:2601.22074) 및 MuJoCo Playground.
  • NVIDIA SO‑101 시뮬레이션에서 실제 적용 학습 경로.
  • 곧 출시될 Newton 저장소.

Sources