[论文解读] 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 引入了一项基于 JAX 的闪电般快速的开放式强化学习基准测试，结合了 NetHack 和 Crafter 的复杂性，推理速度比原始的 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 的环境，支持深度探索、长期规划和持续适应，且对硬件要求极低。
• 提供一个基准测试，使研究人员能够在不依赖大规模计算集群的前提下，开展关于样本效率和泛化能力强化学习智能体的有意义研究。
• 证明当前最先进的探索方法和无监督环境设计方法在新基准测试上无法取得显著进展，从而验证其挑战水平。

提出的方法

• 在 JAX 中重实现 Crafter 为 Craftax-Classic，实现相比原始基于 Python 的版本最高达 250 倍的速度提升。
• 利用 JAX 的 JIT 编译、GPU 加速以及全 GPU 训练流水线，消除 CPU-GPU 数据传输瓶颈，实现环境轨迹的海量并行化。
• 在 Craftax-Classic 的基础上扩展出 Craftax，引入程序化生成、Roguelike 机制（如永久死亡、回合制、基于网格的移动）以及受 NetHack 启发的更丰富的世界动态。
• 设计一个复杂且开放的环境，要求智能体具备深度探索、长期记忆以及对新发现世界情境的持续适应能力。
• 引入一种新的评估协议，每轮训练使用 10 亿次环境交互，以实现对样本效率和长时程学习的有意义评估。
• 使用带有 RNN 记忆的 PPO 智能体评估在 Craftax-Classic 和 Craftax 上的性能，并对探索和环境设计方法进行消融研究。

Figure 1 : Pixel-based view from Craftax .

实验结果

研究问题

• RQ1 基于 JAX 的环境是否能在不牺牲复杂性的情况下，实现相比现有开放式强化学习基准测试数量级的速度提升？
• RQ2 标准探索方法和无监督环境设计方法（如 UED）在多大程度上能够解决新的 Craftax 基准测试？
• RQ3 相比 Crafter，Craftax 的更高复杂性是否对当前最先进强化学习算法构成有意义的挑战？
• RQ4 一个简单的 PPO 智能体是否能在合理计算预算（例如单张 GPU 上训练 10 亿步）内于 Craftax-Classic 上实现接近最优的性能？
• RQ5 不同的课程学习和环境设计策略（如 PLR、ACCEL、DR）在 Craftax 上如何影响学习效率和最终性能？

主要发现

• Craftax-Classic 的运行速度比原始基于 Python 的 Crafter 快 250 倍，使得在单张 GPU 上完成 10 亿次环境交互的完整 PPO 训练过程可在一小时内完成。
• 采用 RNN 策略的 PPO 智能体在训练 51 分钟内即可在 Craftax-Classic 上达到最优回报的 90%，证明该环境对标准算法是可解的。
• 现有的全局和回合内探索最先进方法，以及无监督环境设计（UED）方法，在更复杂的 Craftax 基准测试上几乎无法取得进展。
• Craftax-Classic 与 Craftax 之间的性能差距显著：虽然 PPO 能快速解决前者，但在后者上几乎无法取得有意义的进展，表明后者具有极高的挑战性。
• JAX 的使用使得最多可并行 4096 个环境工作进程，极大缩短训练时间，使在消费级硬件上开展大规模强化学习实验成为可能。
• 为适应 Craftax 的长时程动态，必须对现有方法进行超参数调优和网络架构修改（例如，在 UED 中将 rollout 长度与 episode 长度解耦）。

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 .