[논문 리뷰] Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning

Michael R. Matthews, Michael Beukman외 5명|arXiv (Cornell University)|2024. 02. 26.

Reinforcement Learning in Robotics인용 수 4

한 줄 요약

Craftax는 NetHack와 Crafter의 복잡성을 결합한 JAX 기반의 초고속 벤치마크를 소개하며, 원래의 Crafter 대비 250배 빠른 추론 속도를 제공한다. 이는 단일 GPU로도 복잡하고 장기적인 작업에 대해 효율적인 에이전트 학습을 가능하게 하며, PPO는 Craftax-Classic을 한 시간 이내에 해결하지만 더 도전적인 Craftax 환경에선 기존 방법들이 실패한다.

ABSTRACT

Benchmarks play a crucial role in the development and analysis of reinforcement learning (RL) algorithms. We identify that existing benchmarks used for research into open-ended learning fall into one of two categories. Either they are too slow for meaningful research to be performed without enormous computational resources, like Crafter, NetHack and Minecraft, or they are not complex enough to pose a significant challenge, like Minigrid and Procgen. To remedy this, we first present Craftax-Classic: a ground-up rewrite of Crafter in JAX that runs up to 250x faster than the Python-native original. A run of PPO using 1 billion environment interactions finishes in under an hour using only a single GPU and averages 90% of the optimal reward. To provide a more compelling challenge we present the main Craftax benchmark, a significant extension of the Crafter mechanics with elements inspired from NetHack. Solving Craftax requires deep exploration, long term planning and memory, as well as continual adaptation to novel situations as more of the world is discovered. We show that existing methods including global and episodic exploration, as well as unsupervised environment design fail to make material progress on the benchmark. We believe that Craftax can for the first time allow researchers to experiment in a complex, open-ended environment with limited computational resources.

연구 동기 및 목표

• 계산적으로 실행 가능하면서도 도전적인, 개방형 강화학습을 위한 빠른 복잡한 벤치마크의 부족을 해결한다.
• Crafter, NetHack, Minecraft와 같은 기존 개방형 환경들이 느린 추론 속도로 인해 계산적으로 접근이 어려운 문제를 해결한다.
• 딥 탐색, 장기적 계획, 지속적인 적응을 지원하는 고속의 JAX 네이티브 환경을 개발한다.
• 대규모 컴퓨팅 클러스터가 필요 없이도 샘플 효율성과 일반화 능력을 갖춘 강화학습 에이전트에 대한 유의미한 연구를 가능하게 하는 벤치마크를 제공한다.
• 현재 최고 수준의 탐색 및 비지도 환경 설계 방법이 새로운 벤치마크에서 뚜렷한 진전을 이루지 못함을 입증함으로써, 그 도전 수준을 검증한다.

제안 방법

• 원래의 파이썬 기반 버전 대비 최대 250배 빠른 속도를 달성하기 위해 Crafter를 JAX로 재현하여 Craftax-Classic을 구현한다.
• JAX의 JIT 컴파일, GPU 가속, 전체 GPU 기반 학습 파이프라인을 활용해 CPU-GPU 전송 병목 현상을 제거하고 환경 롤아웃의 대량 병렬 처리를 가능하게 한다.
• Crafter-Classic을 확장하여 프로시저럴 생성, 로그루크 메커니즘(예: 영구사망, 턴기반, 격자 기반 이동) 및 NetHack에서 영감을 받은 더 풍부한 월드 동역학을 통합함으로써 Craftax를 구현한다.
• 깊은 탐색, 장기적 기억, 새로운 상황에 대한 지속적 적응이 요구되는 복잡하고 개방형 환경을 설계하여, 세계가 발견됨에 따라 지속적으로 변화하는 환경을 반영한다.
• 1회 실행당 10억 개의 환경 상호작용을 사용하는 새로운 평가 프로토콜을 도입하여 샘플 효율성과 장기적 학습 능력을 의미 있게 평가할 수 있도록 한다.
• RNN 기반 메모리와 커리큘럼 학습을 갖춘 PPO 에이전트를 사용해 Craftax-Classic과 Craftax 양쪽에서 성능을 평가하고, 탐색 및 환경 설계 방법에 대한 아블레이션 스터디를 수행한다.

Figure 1 : Pixel-based view from Craftax .

실험 결과

연구 질문

• RQ1 JAX 기반 환경가 기존 개방형 강화학습 벤치마크보다 수십 배에서 수백 배 빠른 속도를 달성하면서도 복잡성을 유지할 수 있는가?
• RQ2 표준 탐색 및 비지도 환경 설계 방법(예: UED)이 새로운 Craftax 벤치마크를 얼마나 잘 해결할 수 있는가?
• RQ3 Crafter를 넘어서는 Craftax의 복잡성은 현재 최고 수준의 RL 알고리즘에 의미 있는 도전이 되는가?
• RQ4 간단한 PPO 에이전트가 합리적인 컴퓨팅 예산(예: 단일 GPU에서 10억 스텝) 내에서 Craftax-Classic에서 최적 수준의 성능을 달성할 수 있는가?
• RQ5 다양한 커리큘럼 학습 및 환경 설계 전략(예: PLR, ACCEL, DR)은 Craftax에서 학습 효율성과 최종 성능에 어떤 영향을 미치는가?

주요 결과

• Craftax-Classic은 원래의 파이썬 기반 Crafter 대비 250배 더 빠르게 실행되며, 단일 GPU에서 10억 개의 환경 상호작용을 포함한 전체 PPO 학습 런이 한 시간 이내에 완료된다.
• RNN 정책을 갖춘 PPO 에이전트는 51분의 학습 시간 내에 Craftax-Classic에서 최적 수익의 90%를 달성하여, 표준 알고리즘으로도 환경이 해법 가능함을 입증한다.
• 기존의 글로벌 및 에피소드 기반 탐색, 비지도 환경 설계(UED) 방법은 더 복잡한 Craftax 벤치마크에서 거의 또는 전혀 진전이 없다.
• Craftax-Classic과 Craftax 간 성능 격차는 뚜렷하다: PPO는 첫 번째를 빠르게 해결하지만, 두 번째에선 의미 있는 진전이 없어, 후자의 높은 도전 수준을 확인한다.
• JAX의 사용 덕분에 최대 4096개의 환경 워커를 효율적으로 병렬 처리할 수 있어 훈련 시간을 크게 단축하고 소비자 수준의 하드웨어에서도 대규모 RL 실험을 가능하게 한다.
• 기존 방법을 Craftax의 장기적 동역학에 적응시키기 위해 하이퍼파라미터 튜닝 및 아키텍처 수정(예: UED에서 롤아웃 길이를 에피소드 길이에서 분리)이 필요했다.

Figure 2 : Speed comparison with popular benchmarks for open-ended learning. Craftax-Classic and Craftax are 257x and 169x faster than Crafter respectively. Details of the speed test are in Appendix B and best case results are in Table 1 .