[논문 리뷰] Podracer architectures for scalable Reinforcement Learning

Matteo Hessel, Matthias Kroiß외 6명|arXiv (Cornell University)|2021. 04. 13.

Reinforcement Learning in Robotics참고 문헌 28인용 수 9

한 줄 요약

이 논문은 JAX를 사용하여 TPU 풀에서 강화학습을 효율적으로 스케일링할 수 있도록 설계된 Podracer 아키텍처인 Anakin과 Sebulba를 소개한다. 막대한 병렬 처리, 큰 배치 크기, 최적화된 TPU 활용을 통해 저자들은 전체 TPU 풀에서 최대 4300만 프레임/초를 달성했으며, 이는 한 분 이내로 에이전트를 훈련시켜 이전 시스템에 비해 비용을 절감하고 데이터 효율성을 향상시킨다.

ABSTRACT

Supporting state-of-the-art AI research requires balancing rapid prototyping, ease of use, and quick iteration, with the ability to deploy experiments at a scale traditionally associated with production systems.Deep learning frameworks such as TensorFlow, PyTorch and JAX allow users to transparently make use of accelerators, such as TPUs and GPUs, to offload the more computationally intensive parts of training and inference in modern deep learning systems. Popular training pipelines that use these frameworks for deep learning typically focus on (un-)supervised learning. How to best train reinforcement learning (RL) agents at scale is still an active research area. In this report we argue that TPUs are particularly well suited for training RL agents in a scalable, efficient and reproducible way. Specifically we describe two architectures designed to make the best use of the resources available on a TPU Pod (a special configuration in a Google data center that features multiple TPU devices connected to each other by extremely low latency communication channels).

연구 동기 및 목표

• 확장 가능한 딥 강화학습(Reinforcement Learning, RL) 훈련의 증가하는 계산 요구사항을 해결하기 위해.
• TPU 풀을 활용해 빠른 프로토타이핑과 고처리량 훈련을 가능하게 하여 RL 연구의 효율성을 높이기 위해.
• 사용의 용이성과 프로덕션 수준의 확장성 및 재현 가능성 사이의 균형을 이루는 시스템을 설계하기 위해.
• 최적화된 TPU 활용과 아키텍처 선택을 통해 데이터 효율성을 향상시키고 훈련 비용을 감소시키기 위해.
• MuZero와 같은 모델-프리 및 검색 기반 에이전트를 포함한 다양한 RL 워크로드를 지원하기 위해.

제안 방법

• TPU 풀에서 온라인 및 액터-러닝 RL 에이전트를 훈련하기 위해 Podracer 아키텍처로 Anakin과 Sebulba를 설계하기 위해.
• JAX를 사용해 조합 가능한 프로그램 변환, 자동 미분, 그리고 TPU에서의 하드웨어 가속을 수행하기 위해.
• 더 큰 액터 및 러너 배치 크기, 더 긴 트레이젝터리, 더 깊은 네트워크를 통해 TPU 활용도를 극대화하기 위해 훈련 스케일링을 수행하기 위해.
• 특히 검색 기반 에이전트에 유리한 바탕으로, 액팅과 학습 배치 크기를 분리함으로써 처리량을 선형적으로 확장하기 위해.
• MuZero 스타일 에이전트를 위한 순수 JAX 기반 MCTS를 구현하여 검색을 위한 커스텀 C++ 코드를 회피하기 위해.
• 최대 2048개의 TPU 코어를 활용해 여러 코어에 걸쳐 실험을 복제하고 스케일링하여 고처리량과 낮은 월클록 시간을 달성하기 위해.

실험 결과

연구 질문

• RQ1 연구의 민첩성을 유지하면서 TPU 풀을 효과적으로 활용해 딥 강화학습 훈련을 스케일링할 수 있는 방법은 무엇인가?
• RQ2 대규모 RL 훈련에서 TPU에서 처리량과 데이터 효율성을 극대화하기 위해 어떤 아키텍처 선택이 가장 효과적인가?
• RQ3 복잡한 RL 에이전트인 MuZero와 같은 경우, 하이브리드 C++/파이썬 시스템에 비해 순수 JAX 구현이 성능을 동일하거나 초월할 수 있는가?
• RQ4 배치 크기와 네트워크 용량을 증가시키는 것이 TPU 기반 RL 훈련에서 데이터 효율성과 비용에 어떤 영향을 미치는가?
• RQ5 액터-러닝 RL 프레임워크에서 확장성과 데이터 효율성을 어느 정도 분리할 수 있는가?

주요 결과

• Sebulba는 2048개 코어 TPU 풀에서 최대 4300만 프레임/초를 달성했으며, Pong 에이전트를 한 분 이내로 훈련시켰다.
• 8코어 TPU에서 V-trace 에이전트를 훈련시킬 때, 액터 배치 크기를 32에서 128로 늘림으로써 20만 프레임/초에 도달했다.
• 더 큰 네트워크를 사용함으로써 더 큰 배치 크기를 늘리지 않아도 데이터 효율성이 향상되었으며, TPU 시간이나 비용도 증가하지 않았다.
• Sebulba에서 훈련한 MuZero 에이전트는 16코어 TPU에서 9시간 만에 2억 프레임의 Atari 프레임을 수집했으며, 프리emptible 인스턴스 기준 비용은 약 $40였다.
• MuZero의 경우 처리량이 TPU 코어 수에 따라 선형적으로 증가하여 수평적 확장이 효과적으로 이루어졌음을 입증했다.
• 액팅과 학습 배치 크기를 분리함으로써 데이터 효율성을 유지하면서도 복제를 통해 더 빠른 훈련을 가능하게 했다.