Craftax: A Lightning-Fast JAX-Based Benchmark for Open-Ended RL

Craftax is an RL environment entirely written in JAX, reimplementing and significantly extending the game mechanics of Crafter, drawing inspiration from roguelikes like NetHack. Compatible with the gymnax interface, it integrates easily with JAX frameworks like PureJaxRL and JaxUED. Accepted as a spotlight paper at ICML 2024, Craftax offers CPU and GPU/TPU support, along with tutorials and baselines. Installation instructions and usage examples are provided.