[도구]엔비디아 Warp 기반 MJWarp, MuJoCo 로봇 시뮬레이션 GPU로 확장

무슨 일이야

기존의 MuJoCo는 CPU 코어 기반 병렬 샘플링을 지원하며 로봇 개발 및 제어 테스트에 널리 쓰여 왔습니다. 하지만 강화학습 등 대규모 로봇 학습 워크로드에서는 단일 환경의 실행 속도보다 수많은 환경을 동시에 실행하는 능력이 중요해졌습니다.

엔비디아가 공개한 MJWarp(MuJoCo Warp)는 파이썬 기반 GPU 커널 프레임워크인 NVIDIA Warp 위에서 MuJoCo의 물리 파이프라인을 재구현한 도구입니다. 기존의 MJCF 모델 호환성을 유지하면서, 모델과 독립적인 다중 환경 상태를 NVIDIA GPU에 올려 단 한 번의 호출로 전체 배치를 전진시키는 구조를 취합니다. 예시로 다룬 SO-101 로봇 팔 워크플로의 경우 최대 2,048개의 병렬 환경까지 확장해 실행할 수 있습니다.

기반 기술인 NVIDIA Warp는 파이썬 문법으로 정적 타입 커널을 작성하면 이를 CPU 또는 CUDA 실행 코드로 컴파일해 주는 도구입니다. DLPack 호환 기능 및 프레임워크 어댑터를 통해 PyTorch, JAX 같은 딥러닝 라이브러리와 GPU 메모리 상에서 직접 데이터를 주고받을 수 있으며, Warp 1.15 버전부터는 성능을 일부 희생하는 대신 시뮬레이션 재현성을 보장하는 결정론적(deterministic) 실행 모드도 지원합니다.

왜 중요해

로봇 강화학습에서 시뮬레이션 데이터를 CPU에서 GPU로 복사하는 과정은 큰 오버헤드로 작용해 왔습니다. MJWarp를 사용하면 시뮬레이션 연산 자체가 GPU에서 병렬로 수행되므로, PyTorch나 JAX 학습 파이프라인과 제로 카피(zero-copy) 방식으로 데이터를 연계해 처리량을 극대화할 수 있습니다.

짚고 넘어갈 점

NVIDIA Warp 프레임워크 자체는 자동 미분과 결정론적 실행을 지원하지만, 원문에 따르면 이것이 MJWarp 시뮬레이션 롤아웃 전체의 미분 가능성이나 완전한 결정론을 보장하는 것은 아닙니다. 또한 이번 내용은 환경 준비와 스케일링에 초점을 맞추고 있어 실제 정책(policy) 학습 루프나 복잡한 센서 통합은 후속 연동 레이어(Newton, Isaac Lab 등)를 거쳐야 할 것으로 보입니다.

참고: Hugging Face