[논문 리뷰] EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine

Jiayi Weng, Min Lin외 10명|arXiv (Cornell University)|2022. 06. 21.

Reinforcement Learning in Robotics인용 수 14

한 줄 요약

EnvPool는 C++ 기반 스레드 풀 실행자로 환경 시뮬레이션을 최적화하여 강화학습(Reinforcement Learning, RL) 훈련을 가속화하는 고도로 병렬화된 강화학습 환경 실행 엔진입니다. MuJoCo에서는 최대 300만 개의 물리적 스텝/초, Atari에서는 100만 프레임/초를 기록하여 기준 파이썬 구현 대비 각각 19.2배 빠릅니다. 기존 RL 라이브러리와 원활하게 통합되며, 노트북에서 Atari Pong이나 MuJoCo Ant 훈련을 5분 이내로 완료할 수 있도록 빠른 훈련을 가능하게 합니다.

ABSTRACT

There has been significant progress in developing reinforcement learning (RL) training systems. Past works such as IMPALA, Apex, Seed RL, Sample Factory, and others, aim to improve the system's overall throughput. In this paper, we aim to address a common bottleneck in the RL training system, i.e., parallel environment execution, which is often the slowest part of the whole system but receives little attention. With a curated design for paralleling RL environments, we have improved the RL environment simulation speed across different hardware setups, ranging from a laptop and a modest workstation, to a high-end machine such as NVIDIA DGX-A100. On a high-end machine, EnvPool achieves one million frames per second for the environment execution on Atari environments and three million frames per second on MuJoCo environments. When running EnvPool on a laptop, the speed is 2.8x that of the Python subprocess. Moreover, great compatibility with existing RL training libraries has been demonstrated in the open-sourced community, including CleanRL, rl_games, DeepMind Acme, etc. Finally, EnvPool allows researchers to iterate their ideas at a much faster pace and has great potential to become the de facto RL environment execution engine. Example runs show that it only takes five minutes to train agents to play Atari Pong and MuJoCo Ant on a laptop. EnvPool is open-sourced at https://github.com/sail-sg/envpool.

연구 동기 및 목표

• RL 훈련 시스템에서의 병행 환경 실행에 대한 최적화되지 않은 성능 저하 요인을 해결한다.
• 노트북부터 고성능 DGX-A100 시스템에 이르기까지 다양한 하드웨어에서 시뮬레이션 처리량을 향상시킨다.
• OpenAI Gym 및 dm_env API의 즉시 사용 가능한 대체품을 제공함으로써 연구자들이 최소한의 코드 수정으로 에이전트를 더 빠르게 훈련시킬 수 있도록 한다.
• 동기 및 비동기 실행 모드를 모두 지원하여 하드웨어 활용도를 극대화한다.
• CleanRL, rl_games, DeepMind Acme와 같은 주요 RL 훈련 라이브러리와의 광범위한 호환성을 확보한다.

제안 방법

• 다수의 RL 환경을 병행 실행하기 위해 C++ 기반 스레드 풀 실행자를 구현하여 파이썬 GIL 및 I/O 오버헤드를 최소화한다.
• C++ 측에서 파이썬 래퍼를 최적화하여 환경 상호작용 및 관측 처리의 지연을 감소시킨다.
• 표준 OpenAI Gym 및 dm_env 호환 파이썬 API를 통해 환경 상호작용을 노출하여 원활한 통합을 가능하게 한다.
• 동기 및 비동기 실행 모드를 모두 지원하며, 비동기 모드는 더 높은 처리량과 낮은 데이터 노후화를 제공한다.
• 개발자가 커스텀 C++ 기반 RL 환경을 최소한의 노력으로 EnvPool에 통합할 수 있도록 모듈러 아키텍처를 설계한다.
• 효율적인 메모리 관리와 C++ 및 파이썬 간의 0-복사 데이터 전송을 사용하여 직렬화 오버헤드를 감소시킨다.

실험 결과

연구 질문

• RQ1 고도로 병렬화된 환경 실행 엔진이 RL 훈련에서 시뮬레이션 지연을 크게 감소시킬 수 있는가?
• RQ2 EnvPool의 성능는 노트북부터 고성능 DGX 시스템에 이르기까지 다양한 하드웨어 구성에서 어떻게 스케일링되는가?
• RQ3 표준 RL 벤치마크에서 샘플 효율성을 유지하면서 EnvPool이 훈련 처리량을 얼마나 향상시킬 수 있는가?
• RQ4 EnvPool는 에이전트 코드를 수정하지 않고도 기존 RL 훈련 파이프라인에 효율적으로 통합될 수 있는가?
• RQ5 환경 배치에서 비동기 실행 방식이 동기 실행 방식에 비해 어떤 성능 향상을 이룰 수 있는가?

주요 결과

• NVIDIA DGX-A100(256개 CPU 코어)에서 EnvPool는 Atari에서 100만 프레임/초, MuJoCo에서는 300만 물리적 스텝/초를 기록하여 기준 파이썬 구현 대비 각각 14.9배, 19.2배 빠릅니다.
• 12개 CPU 코어를 가진 노트북에서도 EnvPool는 표준 파이썬 subprocess 기반 환경 실행 대비 2.8배 빠릅니다.
• 노트북에서 Atari Pong 및 MuJoCo Ant에 대한 엔드 투 엔드 훈련이 5분 이내로 완료되어 빠른 연구 반복에 실용적임을 입증합니다.
• MuJoCo HalfCheetah-v3에서 같은 Acme PPO 훈련 설정에서 EnvPool는 DummyVecEnv 대비 벽 시간을 절반 이하로 줄였으며, 유사한 최종 에피소드 수익을 유지합니다.
• 훈련 곡선 분석 결과, num_envs와 batch_size의 곱을 일정하게 유지할 경우 병행 환경 수를 늘릴수록 훈련 시간이 크게 감소하지만 샘플 효율성은 저하되지 않습니다.
• CleanRL, rl_games, Ray, DeepMind Acme와 같은 주요 RL 프레임워크와의 강력한 호환성을 입증하여 즉시 통합이 가능합니다.