기술 요약: MuJoCoUni

문제 정의

로봇 학습 시스템이 물리 시뮬레이터를 훈련 루프에 직접 통합함에 따라, 물리 단계, 센서 읽기, 환경 재설정을 고주파수로 배치 실행해야 하는 요구가 증가하고 있습니다. 상류 (upstream) MuJoCo는 성숙한 에셋과 솔버를 제공하지만, 표준 Python 인터페이스는 고주파수에서 객체 수명 관리, 인터페이스 호출, 출력 구성에 관한 오버헤드를 초래합니다.

기존 솔루션은 다음과 같은 트레이드오프를 제시합니다:

• 상류 mujoco.rollout: 개방 루프 궤적 생성 (계획, 시스템 식별) 에 효율적이지만, 온라인 강화 학습 (RL) 에 필요한 상태 유지 환경 지속성, 희소 재설정 기능, 환경별 무작위화 수명 주기가 결여되어 있습니다.
• GPU-주거 시뮬레이터 (예: Brax, MJX, Isaac Gym): 대규모 병렬성을 제공하지만, 종종 모델과 데이터 레이아웃이 특정 가속기 제약에 맞춰져야 하므로 전체 기능 커버리지, 복잡한 접촉 처리, 또는 기존 CPU 최적화 MuJoCo 에셋의 재사용 능력을 희생할 수 있습니다.
• 표준 CPU 벡터화 런타임: 종종 특정 MuJoCo 솔버 의미론과 깊은 통합이 부족하거나, 모델 지속성과 도메인 무작위화를 효율적으로 유지하기 위해 상당한 엔지니어링이 필요합니다.

모델, 센서, 접촉, 제약 등 상류 MuJoCo 의 전체 의미론적 충실도를 유지하면서, 온라인 RL 및 배치 물리 평가를 위한 고처리량 병렬 실행과 상태 유지 환경 관리를 제공하는 CPU 기반 배치 런타임이 필요합니다.

방법론

MuJoCoUni 는 MuJoCo 의 경량 다운스트림 배포판으로, 엔지니어링 추가 사항을 Python 바인딩 레이어 내에만 집중시켜 핵심 물리 커널, 솔버, 적분기, 접촉 모델을 변경하지 않습니다.

핵심 아키텍처: BatchEnvPool

중앙 구성 요소인 BatchEnvPool 은 다음을 관리하는 C++/pybind11 실행기입니다:

1. 지속적 리소스: 환경당 하나씩 $N$개의 mjModel 복사본과 스레드당 하나씩 $W$개의 mjData 작업자를 소유합니다.
2. 스레드 풀: 내부 스레드 풀이 환경 인덱스를 작업 스레드로 배정하여 병렬 실행을 수행합니다.
3. 상태 유지 실행: 호출마다 상태가 없는 rollout 과 달리, BatchEnvPool 은 호출 간 환경 상태를 유지하여 희소 재설정과 지속적 모델 변형을 가능하게 합니다.

주요 원시 연산

API 는 온라인 로봇 훈련 및 평가를 위해 설계된 여러 원시 연산을 노출합니다:

• step: 모든 환경에서 $n$개의 적분 단계 (mj_step) 를 수행합니다. 전체 궤적을 반환하는 오버헤드를 피하기 위해 최종 상태 (및 선택적으로 최종 센서 데이터) 만 반환합니다.
• forward: 동역학을 진행시키지 않고 현재 상태를 기반으로 센서를 업데이트하기 위해 mj_forward 를 실행합니다. 관측치 구성에 유용합니다.
• reset: 희소 재설정을 구현하여 종료된 환경의 부분집합 (env_ids) 에 대해서만 작동합니다. 재설정 수명 주기 도메인 무작위화를 지원하여 필드 패치 (예: 질량, 관성, 마찰) 를 적용하고 영향을 받는 환경에 대해서만 mj_setConst 를 호출합니다.
• compute_site_jacobians: 특정 사이트에 대한 병렬 이동 및 회전 야코비안의 배치 계산을 수행하며, 최소 운동학 접두사 (mj_kinematics, mj_comPos) 를 실행한 후 mj_jacSite 를 호출합니다.
• sample_hfield_height: MuJoCo 높이장 기하학을 다양한 프레임 정렬 (world, body, yaw) 을 지원하며 배치로 이차 선형 보간합니다.

설계 경계

구현은 MuJoCo 소스 트리를 수정하지 않습니다. 환경 인스턴스 생성에는 mj_copyModel 을, 무작위화에는 mj_setConst 를 의존합니다. 설계는 두 가지 모드를 지원합니다:

1. 단일 공유 모델: 모든 환경이 모델 복사본을 공유합니다 (메모리 효율적, 무작위화 유연성 낮음).
2. 모델 변형: 각 환경이 고유한 모델 복사본을 소유하여, 구성 시 사전 컴파일된 기하학적 수준 무작위화 (예: 링크 길이, 메쉬 스케일) 를 허용합니다.

주요 기여

1. 상태 유지 배치 런타임: 개방 루프 궤적 생성과 온라인 상태 유지 RL 환경 요구 사항 사이의 간극을 해소하기 위해 BatchEnvPool 을 도입합니다.
2. 효율적인 희소 재설정 및 무작위화: 종료된 환경만 재설정하고 필드 패치 및 모델 변형을 포함한 도메인 무작위화를 적용하는 메커니즘을 제공하며, 오버헤드를 최소화하여 재설정 환경 수에 비례하여 선형적으로 확장됩니다.
3. 배치 물리 쿼리: 동역학을 진행시키지 않고 센서 순전달, 사이트 야코비안, 높이장 샘플링의 고성능 배치 구현을 노출합니다.
4. 의미론 보존: 물리 행동이 상류 CPU MuJoCo 와 동일하게 유지되도록 하여 기존 XML/MJB 에셋, 디버깅 도구, 센서 정의와의 호환성을 유지합니다.

결과

저자들은 Intel i9-14900HX CPU(16 스레드) 를 사용하여 4 가지 로봇 모델 (Unitree Go1, Wonik Allegro, Franka Panda, CMU Humanoid) 에서 MuJoCoUni 를 검증했습니다.

• 처리량: 처리량은 약 256~512 개 환경에서 포화됩니다. 포화 시 시스템은 다음과 같은 성능을 달성합니다:
  • Wonik Allegro: 초당 약 180 만 단계.
  • Unitree Go1: 초당 약 120 만 단계.
  • Franka Panda: 초당 약 41 만 단계.
  • CMU Humanoid: 초당 약 29 만 단계.
• 재설정 성능: C++ BatchEnvPool 경로는 Python 루프보다 훨씬 빠릅니다. 4096 개 환경의 전체 재설정의 경우, 3.5 ms가 소요되는 반면 Python for 루프는 53 ms가 소요되어 약 15 배의 속도 향상을 보입니다.
• 야코비안 성능: 4096 개 환경에 대한 사이트 야코비안 계산은 C++ 에서 0.53 ms가 소요되는 반면 Python 에서는 11.9 ms가 소요되어 약 22 배의 속도 향상을 보입니다.
• 높이장 샘플링: 4096 개 환경에 대한 높이장 샘플링은 C++ 에서 0.52 ms가 소요되는 반면 Python 에서는 290 ms가 소요되어 약 555 배의 속도 향상을 보입니다.
• 모델 변형: 포화 상태에서 환경당 고유한 모델 복사본을 사용하는 것 (모델 변형 모드) 과 단일 공유 모델을 사용하는 것 사이의 오버헤드는 무시할 수 있으며, 처리량은 본질적으로 동일해집니다.

중요성 및 범위

본 논문은 MuJoCoUni 를 GPU-주거 시뮬레이터나 상류 rollout 을 대체하는 것이 아닌 보완적 엔지니어링 솔루션으로 위치시킵니다.

• GPU 에 대한 보완: 이는 원시 가속기 거주성보다 기능 커버리지(전체 MuJoCo 의미론, 복잡한 접촉, 디버깅) 와 크로스 플랫폼 이식성이 우선시되는 워크로드를 대상으로 합니다. 시스템을 CPU MuJoCo 솔버의 안정성을 유지하면서 현대 멀티코어 CPU 를 활용한 고처리량 데이터 생성에 활용할 수 있게 합니다.
• rollout 에 대한 보완: 지속적 환경, 희소 재설정, 현재 상태 쿼리가 필요한 온라인 로봇 RL 및 상태 유지 평가 루프의 특정 요구 사항을 해결합니다. 이러한 작업은 전체 궤적 생성을 위해 설계된 rollout 에 최적화되지 않았습니다.
• 시스템 통합: 물리 실행을 상위 수준 훈련 로직, 로깅, 스케줄링에서 분리하는 시뮬레이션 백엔드로 작용합니다. 도메인 무작위화, 지형 인식 보행, 배치 최적화 (예: 진화 전략) 가 필요한 시뮬레이션 - 실제 (sim-to-real) 훈련에 특히 적합합니다.

저자들은 이 기여가 엔지니어링 지향적이며, 상류 MuJoCo 생태계를 보존하는 CPU 기반 배치 실행을 위한 구체적이고 재현 가능한 경로를 제공한다고 강조합니다.