Craftax: a JAX roguelike benchmark that runs entirely on accelerator hardware

The gap Craftax fills between Crafter and NetHack

Crafter is a 2D open-world survival benchmark with a fixed achievement tree. NetHack is a much older, much deeper roguelike whose state space and item interactions are notoriously large. Craftax sits between them. The README describes it as reimplementing and significantly extending Crafter's game mechanics while taking inspiration from roguelike games such as NetHack. The target audience is reinforcement learning researchers who want a procedurally generated, long-horizon environment that runs inside JAX rather than through a Python or C++ simulator. The repository carries the ICML 2024 Spotlight label and links to arXiv paper 2402.16801. Because the environment is written entirely in JAX and conforms to the gymnax interface, it can be dropped into existing JAX training stacks, and the README names PureJaxRL and JaxUED as integration targets. The problem it solves is not 'make a game'. It is 'make a hard game that an accelerator can step through thousands of times in parallel without leaving the device'.

What actually runs when you call env.step

The mechanism is functional and stateless from the caller's perspective. You split a PRNG key, build an environment with make_craftax_env_from_name, read env.default_params, and call env.reset(rng, env_params) to get an observation and a state. Actions come from env.action_space(env_params).sample(rng), and env.step(rng, state, action, env_params) returns observation, state, reward, done and info. The state is carried explicitly rather than hidden inside the environment object, which is what makes the whole thing jit-compatible and vectorizable. The README shows exactly this five-call sequence. Observations come in at least two flavours: the example uses the name Craftax-Symbolic-v1, and the play instructions separately mention Craftax-Classic, so the observation representation is a selectable axis rather than a single fixed format. Rendering is also JAX, which is why the README warns that compiling the render and step functions takes roughly 30 seconds for the first frame and another 20 seconds for the first action. That compile cost is the visible signature of the design: the environment is a traced computation, not an interpreter loop.

Installation, GPU wheels, and the two ways to get the code

For benchmark use, pip install craftax pulls the current release from PyPI. If you need the tip of main, the README gives pip install git+https://github.com/MichaelTMatthews/Craftax.git@main. Both routes install JAX for CPU by default, and the README is explicit that running on GPU or TPU requires installing the correct JAX wheel yourself; for NVIDIA it gives pip install -U "jax[cuda12]". If you intend to modify the environment rather than consume it, the documented path is git clone, cd Craftax, pip install -e ".[dev]", then pre-commit install, with a stated requirement of pip>=23.0. The dev extra and the pre-commit hook indicate that contributions are expected to pass linting; the README displays a ruff badge. Two console entry points ship with the package: play_craftax and play_craftax_classic. Neither is a benchmark harness. For training runs the README points outward to the separate Craftax_Baselines repository rather than bundling example training code here.

Optimistic resets are the sharpest edge in the API

The README's gotchas section is the most consequential part of the documentation. Environments created with auto_reset=False do not reset themselves, and the README states plainly that if you do not handle this, episodes will continue into invalid states. The intended pattern is to wrap such environments in OptimisticResetVecEnvWrapper for efficient resets, or in AutoResetEnvWrapper to recover default gymnax behaviour. The reference implementation for correct usage is ppo.py in Craftax_Baselines, not anything in this repository. This is a real failure mode, not a theoretical one: a training loop that ignores it will silently produce garbage trajectories rather than raising an error, because from the environment's point of view it is simply being stepped with a state that should have been discarded. If you want the safe default, auto_reset=True returns an ordinary auto-reset environment that behaves like any other gymnax environment. The performance-oriented path is the one that requires discipline.

The texture cache will bite you the moment you edit the game

Craftax caches a texture atlas to avoid rebuilding it on every import. For benchmark users this is invisible. For anyone adding blocks or entities, the README warns that a stale cache can cause errors, and the documented escape hatch is the environment variable CRAFTAX_RELOAD_TEXTURES, set to true to force textures to be created from scratch on every run. The v1.6.1 release note, 'Lazily load textures', indicates this area was still being adjusted as recently as the latest release, so anyone maintaining a fork should treat the texture pipeline as moving. This is a small detail with outsized cost: a contributor who does not know about the cache can spend time debugging an error that has nothing to do with their change. It is also the kind of issue that a benchmark-only user will never encounter, which is a reasonable argument for keeping the two audiences separated.

Where Craftax is the wrong tool

Craftax is a poor fit if your workflow is not already JAX-shaped. The compile latency documented in the README, roughly 30 seconds to render and another 20 to take a first action, is a fixed tax per process, and it is paid again whenever functions are recompiled. Interactive experimentation and quick debugging loops suffer accordingly. If your environment stack is Gymnasium-based and your agents are PyTorch, porting to a gymnax interface is a rewrite, not a wrapper. The scoreboard is a second limitation, and the README says so itself: scores from outside the original Craftax paper are reported and have not been verified. That means the table mixing PPO-GTrXL at 18.3 percent, PQN-RNN at 16.0, PPO-RNN at 15.3 and RND at 12.0 is a convenience listing, not a controlled comparison. Reward is reported as a percentage of a maximum of 226, which is a specific normalisation choice worth checking against your own reporting. If you need a benchmark with a stable, audited leaderboard, Craftax does not currently provide one.

Crafter is the honest comparison

Crafter is the direct alternative and the project Craftax is built on top of. The difference is architectural. Crafter is a Python environment with a rendering stack designed for human viewing and single-instance stepping; Craftax is a traced JAX computation designed for batched, on-device rollouts. That difference changes what you can do: with Craftax you can split keys, vmap across environments, and keep the whole rollout inside the accelerator, which is the pattern PureJaxRL and JaxUED are built around. The cost is that Craftax inherits the constraints of JAX, including compilation time, static shapes, and the explicit state-passing discipline. If you only need a survival benchmark for a small number of parallel environments on CPU, Crafter's simpler execution model is easier to reason about. If your bottleneck is environment throughput and you are already committed to JAX, Craftax is the version of that idea that scales.

Maintenance, licensing, and what to check before you commit

Craftax is MIT licensed, which permits commercial and closed-source use with the usual attribution requirement; that is a summary of the identifier, not legal advice, and you should read the LICENSE file in the repository before relying on it. Maintenance activity is visible in the release history: v1.5.0 fixed a plant growing bug in July 2025, v1.6.0 was a bug fixes and refactoring release, and v1.6.1 addressed lazy texture loading in June 2026. The cadence is irregular and the changes are sometimes behavioural, which means pinning a version is safer than tracking main for reproducible experiments. Two practical checks before adopting: confirm your training loop wraps auto_reset=False environments correctly, using ppo.py from Craftax_Baselines as the reference, and confirm your JAX installation matches your accelerator, since the default pip install will silently give you a CPU build. If your experiments depend on the exact game dynamics, pin the version in your requirements and note it alongside your results.

Editorial conclusion

Adopt Craftax if you already train JAX agents and want a harder, procedurally generated benchmark that stays on-device from reset to step. Do not adopt it if you need a CPU-friendly environment, a stable API without wrapper discipline, or a benchmark with independently verified scoreboards. Before committing, verify that your training loop handles auto_reset=False correctly, check whether CRAFTAX_RELOAD_TEXTURES matters for your fork, and confirm your JAX wheel matches your accelerator.

Official sources

1. Issues
2. License: MIT
3. MichaelTMatthews/Craftax on GitHub
4. README
5. Releases