Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning
This provides a computationally efficient benchmark for researchers to study complex, open-ended reinforcement learning problems, addressing a bottleneck in the field.
The authors tackled the lack of fast and complex benchmarks for open-ended reinforcement learning by introducing Craftax, a JAX-based rewrite of Crafter that runs up to 250x faster and achieves 90% optimal reward with PPO in under an hour on a single GPU, while also extending it with NetHack-inspired mechanics to pose deeper challenges.
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.