mjbatch

CPU에서 수천 개의 MuJoCo 시뮬레이션을 병렬로 돌린다

라이선스

Apache-2.0

개요

mjbatch는 수천 개의 MuJoCo 시뮬레이션을 CPU에서 병렬로 돌리는 파이썬 라이브러리다. 스텝 진행은 C++ 스레드 풀이 GIL을 놓은 채 수행하고, 배치는 살아 있는 NumPy 뷰로 파이썬에 노출된다. bind는 모든 시뮬레이션에 걸친 MjData 필드 배열을 돌려주므로 ctrl에 쓰고 qpos를 읽는 데 복사도, 시뮬레이션마다 도는 루프도 필요 없다. expand는 MjModel 필드에 같은 일을 해서 시뮬레이션마다 제 물리를 갖게 하고 — 마찰이 다르고 링크 길이가 다르다 — 그 뒤 set_const가 파생 상수를 다시 계산한다. 저장소에는 강화학습, 모델 예측 제어, 시스템 식별, 하드웨어 협업 설계를 아우르는 자족적인 예제 여섯 개가 들어 있다.

핵심은 일부러 작게 두었다. MjModel 하나와 시뮬레이션 개수로 Batch를 만들면 스레드 풀이 기본으로 모든 논리 CPU를 쓰고, batch.step()이 전부를 진행시킨다. 이것을 단순한 for 루프 이상으로 만드는 것은 배열이 모아 온 것이 아니라 살아 있다는 점이다. qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl")이 건네는 것은 뷰이고, 제어기는 배치 전체에 한 번의 대입으로 써 넣을 수 있으며, 배치가 진행되는 동안 그 자리에서 갱신된다. 시뮬레이션별 모델 파라미터가 나머지 절반이다 — batch.expand("geom_friction")은 시뮬레이션마다 무작위화할 수 있는 배열을 돌려주는데, 그것이 곧 도메인 무작위화이자 시스템 식별이며 별도의 장치가 필요 없다.

예제가 곧 이 라이브러리의 논거다. 하나하나가 시연이 아니라 완결된 방법이기 때문이다. iLQR로 세우는 카트폴, 예측 샘플링을 쓴 카트폴 MPC, 후퇴 지평 iLQR로 기준 백플립을 추종하는 G1 휴머노이드, PPO로 학습한 Go1 사족 조이스틱 제어기, 던지는 팔의 비율과 기어와 제어를 CEM이 함께 최적화하는 예제, 그리고 감쇠 가우스-뉴턴으로 Rizon 팔의 관성 파라미터를 맞추는 예제다. 저자는 다섯 해 된 M1 노트북에서 Go1 제어기가 1분이 채 안 되어 걷는 법을 배웠다고 보고한다. PyPI에서 설치되고 mujoco 3.13.0을 고정하며 scikit-build-core와 nanobind로 빌드되고, CI와 pytest 스위트, 확장성 벤치마크를 갖췄다.

중요한 이유

대규모 병렬 시뮬레이션은 GPU 이야기가 되었고 — MJX, Isaac, Warp — 그것은 워크스테이션 카드가 없는 사람의 진입 비용을 조용히 올렸다. mjbatch는 한 번에 한 코어가 아니라 CPU를 제대로 씀으로써 같은 모양의 작업을 노트북 위로 되돌린다. 게다가 재구현이 아니라 순정 MuJoCo를 통과하므로 모델도 물리도 이미 신뢰하던 그대로다. 독자 대부분이 클러스터가 아니라 노트북과 값싼 팔을 가진 이 카탈로그에는 이쪽이 더 쓸모 있는 절충이다.

해결하는 문제

수천 번의 MuJoCo 롤아웃을 얻으려면 대개 GPU 스택으로 옮기고 그 하드웨어 비용과 물리의 재구현을 받아들여야 했다. mjbatch는 순정 MuJoCo를 그대로 두고 이미 손에 있는 CPU에서 병렬성을 끌어낸다. 노트북 규모의 제어·식별·협업 설계 작업에는 그것으로 충분하다.

핵심 요점

• C++ 스레드 풀이 GIL을 놓은 채 배치 전체를 진행시키므로 병렬성은 파이썬의 교차 실행이 아니라 진짜다
• bind는 MjData 필드를 모든 시뮬레이션에 걸친 살아 있는 배열로 노출한다 — 제어기는 ctrl에 한 번 쓰고 qpos를 그 자리에서 읽으며, 모으는 단계가 없다
• expand는 시뮬레이션마다 제 MjModel 파라미터를 주고 set_const가 파생 상수를 다시 계산한다 — 별도 장치 없이 도메인 무작위화와 시스템 식별이 된다
• 여섯 개의 완성된 예제는 저마다 하나의 방법이다: iLQR, 예측 샘플링 MPC, PPO, CEM 하드웨어 협업 설계, 감쇠 가우스-뉴턴 관성 추정
• 저자는 다섯 해 된 M1 노트북에서 Go1 조이스틱 제어기가 1분이 채 안 되어 걷는 법을 배웠다고 보고한다
• PyPI에서 설치되고 mujoco 3.13.0을 고정하며 CI와 pytest 스위트, 확장성 벤치마크를 갖췄다

대상

• 스텝마다 많은 롤아웃이 필요한 MPC·궤적 최적화·CEM 연구자
• GPU 클러스터 없이 작은 강화학습 제어기를 훈련하는 사람
• 모델 파라미터 위에서 시스템 식별이나 도메인 무작위화를 하는 실무자
• MJX나 Warp로 옮기지 않고 배치 시뮬레이션을 원하는 MuJoCo 사용자

한계

• 설계상 CPU 전용 — GPU 규모의 배치에는 여전히 MJX나 MuJoCo Warp가 길이며 이것이 그것들을 대체하지 않는다
• 작성 시점에 생후 엿새, v0.1.1이며 Beta로 분류되어 있다. 별 521개는 라이브러리만큼이나 저자의 이력을 반영한다
• mujoco==3.13.0을 정확히 고정하므로 저자가 올려야만 따라 움직인다
• 소스에서 빌드하려면 C++ 툴체인이 필요하고, 창을 여는 예제는 --headless를 주지 않는 한 디스플레이가 있어야 한다
• API 표면이 일부러 좁다 — Batch, bind, expand, set_const — 그 모양을 벗어나는 일은 MuJoCo로 곧장 내려간다는 뜻이다
• 공개된 분류자에는 macOS와 리눅스만 있고 윈도우 휠은 없다