sail-sg/envpool

해결하는 문제

EnvPool은 느린 환경 시뮬레이션으로 인해 발생하는 강화학습(RL)의 성능 저하 문제를 해결합니다. 기존의 파이썬 기반 벡터 환경은 서브프로세스로 인한 높은 오버헤드로 인해 경험 수집의 처리량이 제한됩니다. EnvPool은 고성능 C++ 기반 실행 엔진을 제공하여 표준 파이썬 구현보다 훨씬 높은 프레임당 초(FPS)로 환경을 대규모 병렬화할 수 있습니다.

작동 방식

EnvPool은 C++로 구현되었으며, pybind11를 사용하여 파이썬 인터페이스를 제공합니다. 환경 풀을 관리하고 스레드 풀을 활용해 병렬로 실행합니다. 주로 두 가지 상호작용 모드를 지원합니다:

• 동기 API: Gymnasium 및 dm_env API와 유사하게, step 및 reset 함수는 기본적으로 환경의 배치를 처리합니다.
• 비동기 API: step 함수를 send와 recv 작업으로 분리하여, 사용자가 액션을 입력하고 상태를 수신할 때 전체 배치가 완료되기를 기다리지 않고도 처리할 수 있도록 하여 효율을 더욱 높입니다.

JAX jit 함수를 통해 XLA와 통합되며, Stable-Baselines3, Tianshou, CleanRL과 같은 주요 RL 라이브러리와 호환됩니다.

대상 사용자

에이전트를 훈련하기 위해 대량의 경험 데이터를 빠르게 수집해야 하는 강화학습 연구자 및 개발자에게 적합합니다. 특히 Atari, MuJoCo, 다양한 로봇 시뮬레이터와 같은 계산 비용이 큰 환경을 사용하는 경우에 유용합니다.

주요 특징

• 극도의 높은 처리량: 고사양 하드웨어(DGX-A100)에서 Atari 기준 약 100만 FPS, MuJoCo 기준 약 300만 스텝/초를 달성 가능.
• 광범위한 환경 지원: Atari, MuJoCo, Box2D, DeepMind Control Suite, ViZDoom 등 다수 지원.
• 크로스 플랫폼: Linux, macOS, Windows에서 사용 가능.
• 유연한 API: Gymnasium 및 dm_env와 호환되며, 동기 및 비동기 실행 모두 지원.
• 내장 렌더링: 배치 처리된 rgb_array 출력과 human 표시 모드를 지원.
• C++ 확장성: 사용자 정의 C++ 환경을 통합하기 위한 개발자 API 제공.