Package: craftax
CraftaxCraftax is an RL environment written entirely inJAX.  Craftax reimplements and significantly extends the game mechanics ofCrafter, taking inspiration from roguelike games such asNetHack.
Craftax conforms to thegymnaxinterface, allowing easy integration with existing JAX-based frameworks likePureJaxRL.Basic UsageCraftax conforms to the gymnax interface:rng = jax.random.PRNGKey(0)
rng, _rng = jax.random.split(rng)
rngs = jax.random.split(_rng, 3)

# Create environment
env = AutoResetEnvWrapper(CraftaxSymbolicEnv())
env_params = env.default_params

# Get an initial state and observation
obs, state = env.reset(rngs[0], env_params)

# Pick random action
action = env.action_space(env_params).sample(rngs[1])

# Step environment
obs, state, reward, done, info = env.step(rngs[2], state, action, env_params)InstallationThe latest Craftax release can be installed from PyPi:pip install craftaxIf you want the most recent commit instead use:pip install git+https://github.com/MichaelTMatthews/Craftax.git@mainExtending CraftaxIf you want to extend Craftax, run:git clone https://github.com/MichaelTMatthews/Craftax.git
cd Craftax
pip install --editable .GPU-Enabled JAXBy default, both of the above methods will install JAX on the CPU.  If you want to run JAX on a GPU/TPU, you'll need to install the correct wheel for your system fromJAX.
For NVIDIA GPU the command is:pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlPlayTo play Craftax run:python -m src.play_craftaxor to play Craftax-Classic run:python -m src.play_craftax_classicSince Craftax runs entirely in JAX, it will take some time to compile the rendering and step functions - it might take around 30s to render the first frame and then another 20s to take the first action.  After this it should be very quick.  A tutorial for how to beat the game is present intutorial.md.  The controls are printed out at the beginning of play.ExperimentTo run PPO with default hyperparameters run:python -m src.ppoor to run PPO with memory call:python -m src.ppo_rnnTo use ICM or E3B with the default parameters use the--train_icmand--use_e3bflags.
Use theenv_nameparameter to control which environment is used.  It can be set to"Craftax-Symbolic-v1","Craftax-Pixels-v1","Craftax-Classic-Symbolic-v1"or"Craftax-Classic-Pixels-v1"GotchasCraftax provides the option to use optimistic resets to improve performance, which means that (unlike regular gymnax environments) itdoes not auto-resetby default.
This means that the environment should always be wrapped either inEfficientResetVecEnvWrapperorAutoResetEnvWrapper.  Seeppo.pyfor correct usage of both wrappers.ScoreboardIf you would like to add an algorithm please open a PR and provide a reference to the source of the results.Craftax-1BAlgorithmReward (% max)SourcePPO-RNN15.3PureJaxRLPPO11.9PureJaxRLICM11.9ICME3B11.0E3BCraftax-1MAlgorithmReward (% max)SourcePPO-RNN2.3PureJaxRLPPO2.2PureJaxRLICM2.2ICME3B2.2E3BCitationIf you use Craftax in your work please cite it as follows:@article{matthews2024craftax,
  title={Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning},
  author={Michael Matthews and Michael Beukman and Benjamin Ellis and Mikayel Samvelyan and Matthew Jackson and Samuel Coward and Jakob Foerster},
  journal={arXiv preprint},
  year={2024},
}