[論文レビュー] Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning

Michael R. Matthews, Michael Beukman+5 more|arXiv (Cornell University)|Feb 26, 2024

Reinforcement Learning in Robotics被引用数 4

ひとこと要約

Craftaxは、NetHackとCrafterの複雑さを組み合わせた、JAXベースの高速なベンチマークを提供する。元のCrafterに比べて250倍高速な推論を実現し、単一のGPUでのみ複雑で長時間にわたるタスクの強化学習エージェントの効率的訓練を可能にする。PPOはCraftax-Classicを1時間未塔で解消可能であり、既存の手法はより挑戦的な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ネイティブで、深く探索的で、長期計画と継続的適応を可能にする高機能な環境を構築し、最小限のハードウェアリソースで運用可能にすること。
• 大規模な計算クラスタを必要とせず、サンプル効率的かつ汎化可能な強化学習エージェントに関する有意義な研究を可能にするベンチマークを提供すること。
• 現在の最先端の探索法および無教師環境設計手法が、新しいベンチマークで顕著な進展を示さないことを示し、その挑戦レベルの妥当性を検証すること。

提案手法

• CrafterをJAXで再実装し、元のPythonベースのバージョンに比べて最大250倍の高速化を達成するCraftax-Classicを構築する。
• JAXのJITコンパイル、GPUアクセラレーション、フルGPUトレーニングパイプラインを活用し、CPU-GPU転送のボトル neck を排除し、環境ロールアウトの並列処理を大規模に可能にする。
• Craftax-Classicを拡張してCraftaxにし、手続き的生成、Roguelikeメカニクス（例：パーマデス、ターンベース、グリッドベース移動）を統合し、NetHackをインspiredしたより洗練された世界ダイナミクスを導入する。
• 深く探索的で、長期記憶と継続的適応を要する複雑でオープンエンドな環境を設計し、世界が発見されるにつれて新たな状況に適応する必要がある。
• 1回の実行で10億の環境インタラクションを用いた新しい評価プロトコルを導入し、サンプル効率性および長時間にわたる学習の評価を可能にする。
• RNNベースの記憶を持つPPOエージェントとカリキュラム学習を用いて、Craftax-ClassicおよびCraftaxにおける性能を評価し、探索および環境設計手法に関するアブレーションスタディを実施する。

Figure 1 : Pixel-based view from Craftax .

実験結果

リサーチクエスチョン

• RQ1 JAXベースの環境は、複雑さを損なわずに、既存のオープンエンド強化学習ベンチマークに比べて桁違いの高速化を達成できるか？
• RQ2 標準的な探索法および無教師環境設計手法（例：UED）は、新しいCraftaxベンチマークをどの程度解消できるか？
• RQ3 Crafterをはるかに超える複雑さを持つCraftaxは、現在のSOTA強化学習アルゴリズムにとって意味のある挑戦を提供するか？
• RQ4 シンプルなPPOエージェントは、合理的な計算リソース（例：単一GPUで10億ステップ）の制約内でCraftax-Classicでほぼ最適なパフォーマンスを達成できるか？
• RQ5 異なるカリキュラム学習および環境設計戦略（例：PLR, ACCEL, DR）は、Craftaxにおける学習効率および最終パフォーマンスにどのように影響するか？

主な発見

• Craftax-Classicは元のPythonベースのCrafterに比べて250倍高速に実行され、単一GPUで10億の環境インタラクションを含む完全なPPOトレーニングが1時間未塔で完了する。
• RNNポリシーを備えたPPOエージェントは、51分のトレーニングでCraftax-Classicで最適報酬の90％を達成し、標準的なアルゴリズムでもトレーニング可能であることを示した。
• グローバルおよびエピソードベース探索、および無教師環境設計（UED）の最先端手法は、より複雑なCraftaxベンチマークではほとんど進展を示さない。
• Craftax-ClassicとCraftaxのパフォーマンスギャップは顕著である。PPOは前者を素早く解消できるが、後者では意味のある進展を示さないため、後者の高い挑戦レベルを示している。
• JAXの使用により、最大4096の環境ワーカーを並列処理可能となり、トレーニング時間を大幅に短縮し、消費者用ハードウェアでも大規模なRL実験を可能にする。
• 既存の手法をCraftaxの長時間にわたるダイナミクスに適応させるために、ハイパーパramータチューニングおよびアーキテクチャの変更（例：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 .