Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning
作者: Michael Matthews, Michael Beukman, Benjamin Ellis, Mikayel Samvelyan, Matthew Jackson, Samuel Coward, Jakob Foerster
分类: cs.LG
发布日期: 2024-02-26 (更新: 2024-06-03)
💡 一句话要点
提出Craftax以解决开放式强化学习基准测试效率低下问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 开放式强化学习 基准测试 Craftax Crafter 深度学习 高效计算 环境交互
📋 核心要点
- 现有开放式强化学习基准测试要么计算资源消耗巨大，运行速度慢，要么复杂度不足以构成挑战，限制了研究的深入。
- 提出Craftax-Classic，通过JAX重写Crafter，显著提升运行速度，并引入Craftax基准测试以增加复杂性和挑战性。
- 实验结果显示，使用Craftax进行训练的PPO在1小时内完成1亿次交互，且获得的奖励接近最优水平，展示了其高效性。
📝 摘要（中文）
基准测试在强化学习算法的开发与分析中起着至关重要的作用。现有的开放式学习基准测试要么计算资源消耗巨大，运行速度缓慢，要么复杂度不足以构成挑战。为此，本文首先提出Craftax-Classic，这是Crafter的全新重写版本，运行速度比原Python版本快250倍。使用单个GPU进行1亿次环境交互的PPO训练在一小时内完成，且平均获得90%的最优奖励。此外，Craftax基准测试在Crafter机制的基础上进行了显著扩展，要求深度探索、长期规划和记忆能力，以及对新情况的持续适应。实验表明，现有方法在该基准测试上进展有限。
🔬 方法详解
问题定义：本文旨在解决现有开放式强化学习基准测试在计算效率和复杂性上的不足。现有基准如Crafter、NetHack和Minecraft计算资源消耗巨大，而Minigrid和Procgen则缺乏足够的挑战性。
核心思路：论文的核心思路是通过重写Crafter并引入新的挑战元素，创建一个既高效又复杂的基准测试环境，使研究者能够在有限的计算资源下进行深入实验。
技术框架：整体架构包括Craftax-Classic和主Craftax基准测试。Craftax-Classic是Crafter的高效实现，主Craftax则在此基础上增加了复杂的游戏机制，要求深度探索和适应能力。
关键创新：最重要的创新在于Craftax-Classic的实现，使得基准测试的运行速度提升至原来的250倍，极大降低了计算资源的需求。同时，Craftax基准测试引入了新的挑战元素，推动了开放式学习的研究。
关键设计：在Craftax-Classic中，使用JAX进行高效计算，优化了环境交互的速度。PPO算法的训练过程中，设计了适应性强的奖励机制，以确保在复杂环境中获得良好的学习效果。
🖼️ 关键图片
📊 实验亮点
实验结果显示，使用Craftax-Classic进行PPO训练时，1亿次环境交互的完成时间不足一小时，且平均奖励达到90%的最优水平。这一性能提升展示了Craftax在开放式强化学习研究中的重要性。
🎯 应用场景
该研究的潜在应用领域包括游戏AI、机器人控制和自主系统等。Craftax提供了一个高效且复杂的测试环境，能够帮助研究者在有限的计算资源下探索开放式强化学习的各种策略和算法，推动相关领域的进步。
📄 摘要（原文）
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.