kevinzakka/mjbatch

A Python library for running thousands of MuJoCo simulations in parallel on CPU

mjbatch – CPU에서 병렬 MuJoCo 시뮬레이션

무엇인가요mjbatch는 Python의 GIL을 해제하는 C++ 스레드 풀을 사용하여 수천 개의 MuJoCo 물리 시뮬레이션을 동시에 실행할 수 있는 작은 Python 패키지입니다. 강화학습, 모델 예측 제어, 시스템 식별, 하드웨어-공학 설계에 필요한 대규모 배치 실행이 필요한 연구자와 엔지니어를 대상으로 합니다.

핵심 아이디어

  • 스레드 풀 실행 – 무거운 작업은 네이티브 C++ 스레드에서 수행되므로, 각 시뮬레이션은 Python의 글로벌 인터프리터 락(GIL)에 의해 차단되지 않고 별도의 코어에서 실행됩니다.
  • 라이브 배열 바인딩qpos, ctrl와 같은 MuJoCo 데이터 필드에 직접 바인딩할 수 있으며, 전체 배치에 걸쳐 인플레이스로 업데이트되는 NumPy 배열처럼 사용할 수 있습니다.
  • 시뮬레이션별 매개변수expand를 통해 geom_friction와 같은 모델 상수를 각 인스턴스별로 변경할 수 있으며, set_const로 유도된 상수를 재계산할 수 있습니다.

일반적인 워크플로우

import mujoco, numpy as np
from mjbatch import Batch

# MuJoCo 모델을 한 번만 로드
model = mujoco.MjModel.from_xml_path('scene.xml')

# 4096개의 병렬 시뮬레이션 배치 생성 (기본적으로 논리 CPU당 1개 스레드)
batch = Batch(model, num_sims=4096)

# 상태 및 제어 배열에 바인딩 – 이들은 전체 배치에 대한 NumPy 뷰입니다
qpos, ctrl = batch.bind('qpos'), batch.bind('ctrl')

# 각 시뮬레이션에 대해 모델 매개변수 무작위화
batch.expand('geom_friction')[:, :, 0] = np.random.uniform(0.4, 1.2, (4096, 1))

for _ in range(1000):
    ctrl[:] = policy(qpos)   # 당신의 컨트롤러가 모든 시뮬레이션에서 동시에 실행
    batch.step()             # 모든 시뮬레이션을 병렬로 단계 진행; qpos는 인플레이스로 업데이트

위의 루프는 단일 시뮬레이션 스크립트와 거의 동일해 보이지만, 내부적으로는 4,000개 환경이 동시에 진행됩니다.

이걸로 할 수 있는 일

  • 강화학습 – 대규모 리플레이 버퍼 생성 또는 정책 빠른 평가 가능 (README에는 오래된 노트북에서 Go1 사족 로봇이 1분 미만에 걷기 학습하는 예시 있음).
  • 모델 예측 제어 / iLQR – 여러 초기 조건에 걸쳐 리카레언트 호라이즌 최적화 실행 가능.
  • 시스템 식별 – 여러 왜곡된 모델을 시뮬레이션하고 합성 데이터와 비교하여 관성 매개변수 피팅.
  • 공설계 – 진화적 또는 기울기 기반 방법을 사용해 로봇 기하학, 기어비, 제어 정책을 동시에 최적화.

리포지토리에 포함된 예제

예제 목적 기술
cartpole_swingup.py 이중 막대 카트 스윙업 iLQR
cartpole_mpc.py 예측 샘플링을 통한 스윙업 MPC
g1_flip.py 인간형 로봇의 백플립 리카레언트 호라이즌 iLQR
go1_joystick.py 사족 로봇 원격 조종 PPO 학습된 정책
arm_throw.py 팔의 비율과 제어 공설계 CEM 최적화
rizon_inertia.py 관성 매개변수 식별 감쇠된 가우스-뉴턴

uv run examples/<file>.py로 예제를 실행할 수 있습니다; 일부는 디스플레이가 필요합니다 (--headless로 비디오 없이도 작동). 리포지토리는 CI 배지와 PyPI 릴리스를 제공하여 패키지가 유지보수되고 pip install mjbatch로 설치 가능함을 나타냅니다.

설치

pip install mjbatch   # 컴파일된 확장과 순수 Python 래퍼를 가져옵니다

이 라이브러리는 공식 MuJoCo Python 바인딩에 의존하므로 MuJoCo 라이선스(학술용 무료)와 mujoco 패키지 설치가 필요합니다.

라이선스 – Apache-2.0, 상용 및 오픈소스 사용 허용.


요약하면, mjbatch는 GPU 클러스터 없이도 MuJoCo 기반 실험을 확장하고 싶은 누구에게나 실용적인 도구입니다. 스레딩 세부 사항을 추상화하면서 전체 배치 시뮬레이션 상태에 직접 NumPy 스타일 접근을 제공합니다.

관련

  • 프로젝트
  • 프로젝트
  • 프로젝트
  • 프로젝트