NVIDIA 엔지니어링 팀, Warp 활용한 MuJoCo 로보틱스 시뮬레이션 GPU 확장 기술 공개

단일 CPU 환경에서 대규모 GPU 배치로의 로봇 시뮬레이션 확장

2026년 9월 23일, NVIDIA 연구진 조니 누녜스 카노(Johnny Nuñez Cano), 아시에르 아란스(Asier Arranz), 리샤브 차다(Rishabh Chadha), 벤 올리베리(Ben Oliveri)는 로보틱스 개발자가 기존 CPU 기반 MuJoCo 시뮬레이션을 고처리량 GPU 워크로드로 전환하는 방법을 다룬 구현 가이드를 발표했다. 물리적 AI(Physical AI) 및 강화학습 파이프라인이 확장됨에 따라, 훈련 속도는 단순히 단일 환경의 지연 시간을 최적화하는 것보다 수백 또는 수천 개의 환경을 동시에 구동하는 것에 점점 더 크게 의존하고 있다.

이러한 간극을 메우기 위해 NVIDIA는 NVIDIA Warp 기반으로 구축된 MuJoCo 물리 엔진의 GPU 가속 구현체인 MuJoCo Warp(MJWarp)를 활용한 워크플로우를 선보였다. NVIDIA의 '물리적 AI를 위한 시뮬레이션 현황(State of Simulation for Physical AI)' 시리즈의 두 번째 편인 이번 가이드는 블록 쌓기 작업을 수행하는 SO-101 팔로워 로봇 팔을 표준 Python 구동 CPU 환경에서 NVIDIA GPU 상의 2,048개 병렬 시뮬레이션 상태로 직접 이전하는 과정을 보여준다.

SO-101 로봇 팔 포팅: 버퍼 크기 설정, CUDA Graphs 및 동등성 검증

이 구성의 기반에는 정적 타입 Python 커널을 네이티브 CUDA 코드로 컴파일하는 Python 기반 프레임워크인 NVIDIA Warp가 있다. 이 프레임워크는 JIT(Just-In-Time) 컴파일, 커널 융합, wp.Tape를 통한 자동 미분, 버전 1.15에서 도입된 결정론적 실행 모드를 지원한다. MJWarp는 이 프레임워크를 MuJoCo 물리 엔진에 적용하여 표준 MJCF XML 씬 기술을 처리하는 동시에, 선두 배치 차원(leading batch dimension)을 추가해 배열을 GPU로 이동시킨다.

마이그레이션은 0.002초의 물리 타임스텝을 목표로 프레임당 10회의 물리 서브스텝과 50Hz 제어 주파수로 구성된 SO-101의 픽 앤 플레이스(pick-and-place) 작업에서 시작된다. 로봇 팔의 임무는 44mm 크기의 빨간색 큐브를 집어 동일한 크기의 파란색 큐브 위에 올려놓는 것이다. 이 워크플로우를 MJWarp로 이전하는 데는 API 전환이 수반된다. 개발자는 mjw.put_model()을 사용해 모델을 업로드하고, mjw.make_data() 또는 mjw.put_data()로 GPU 상주 상태를 할당하며, mjw.step()을 통해 모든 병렬 환경을 동시에 진행시킨다.

핵심적인 운영 요건 중 하나는 디바이스 측 접촉(contact) 및 제약 조건(constraint) 버퍼를 관리하는 것이다. 개발자는 nconmax(환경당 최대 접촉 수), naconmax(글로벌 접촉 상한선), njmax(환경당 최대 제약 조건 수) 등의 용량을 정의해야 한다. 협소 단계(narrowphase) 충돌이 용량을 초과할 때 MJWarp는 실행을 즉시 중단하지 않고 경고만 출력하므로, 사전에 확인하지 않으면 시뮬레이션 결과가 조용히 무효화될 위험이 있다. NVIDIA는 그리퍼 조(jaw) 양쪽과 테이블이 큐브에 동시에 닿는 순간처럼 작업 중 접촉이 가장 많은 상황을 기준으로 한도를 설정하고, mjwarp-testspeed --measure_alloc과 같은 도구를 사용해 할당 상태를 진단할 것을 권장한다.

단일 환경에서의 물리 동등성이 CPU 기준점과 일치하는 것으로 검증되면, 환경을 2,048개 이상으로 확장할 수 있다. 수천 개의 배치 스텝 전반에 걸친 디스패치 지연을 제거하기 위해 개발자는 mjw.step()을 wp.ScopedCapture() 내에 래핑하여 재생 가능한 CUDA Graphs를 생성한다. 프로파일링 시에는 시간 측정 루프 앞뒤에 wp.synchronize()를 삽입해 GPU의 비동기 실행을 고려해야 하며, 롤아웃 단계 중 .numpy()를 통해 배열을 호스트 메모리로 다시 가져오지 않고 시뮬레이션 데이터를 디바이스 상에 유지해야 한다.

Newton 및 Isaac Lab으로의 물리적 AI 시뮬레이션 스택 확장

MJWarp로의 전환은 단일 로봇 모델 예측 제어(MPC)와 배치 강화학습 간에 점차 뚜렷해지는 업계의 분화를 보여준다. 기존 CPU 기반 MuJoCo는 원격 조작, 단일 환경 디버깅, 저지연 MPC의 표준으로 남아 있지만, 배치 GPU 환경은 대규모 데이터 수집을 목표로 한다.

MJWarp는 여러 최신 물리적 AI 프레임워크와 연동되도록 설계됐다. MJWarp는 mjlab(MJWarp를 PyTorch와 직접 결합)과 MJX의 Warp 구현체를 통한 MuJoCo Playground, 그리고 더 광범위한 학습 파이프라인의 실행 백엔드 역할을 한다. NVIDIA는 다음 시뮬레이션 시리즈에서 동일한 SO-101 환경을 자사의 다중 솔버 프레임워크인 Newton으로 가져오는 방법을 다룰 예정이라고 밝혔다. 이 프레임워크에서 MJWarp는 기반 강체 솔버(newton.solvers.SolverMuJoCo)로 작동하여 씬을 Isaac Lab 훈련 워크플로우와 직접 연결하게 된다.