An open-world environment for training RL agents
Project description
Craftax
Craftax is an RL environment written entirely in JAX. Craftax reimplements and significantly extends the game mechanics of Crafter, taking inspiration from roguelike games such as NetHack. Craftax conforms to the gymnax interface, allowing easy integration with existing JAX-based frameworks like PureJaxRL and JaxUED.
Basic Usage
Craftax 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)
Installation
The latest Craftax release can be installed from PyPi:
pip install craftax
If you want the most recent commit instead use:
pip install git+https://github.com/MichaelTMatthews/Craftax.git@main
Extending Craftax
If you want to extend Craftax, run (make sure you have pip>=23.0
):
git clone https://github.com/MichaelTMatthews/Craftax.git
cd Craftax
pip install --editable .
GPU-Enabled JAX
By 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 from JAX. For NVIDIA GPU the command is:
pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Play
To play Craftax run:
play_craftax
or to play Craftax-Classic run:
play_craftax_classic
Since 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 in tutorial.md
. The controls are printed out at the beginning of play.
Experiment
To run PPO with default hyperparameters run:
python -m src.ppo
or to run PPO with memory call:
python -m src.ppo_rnn
To use ICM or E3B with the default parameters use the --train_icm
and --use_e3b
flags.
Use the env_name
parameter 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"
Gotchas
Optimistic Resets
Craftax provides the option to use optimistic resets to improve performance, which means that (unlike regular gymnax environments) it does not auto-reset by default.
This means that the environment should always be wrapped either in OptimisticResetVecEnvWrapper
(for efficient resets) or AutoResetEnvWrapper
(to recover the default gymnax auto-reset behaviour). See ppo.py
for correct usage of both wrappers.
Texture Caching
We use a texture cache to avoid recreating the texture atlas every time Craftax is imported. If you are just running Craftax as a benchmark this will not affect you. However, if you are editing the game (e.g. adding new blocks, entities etc.) then a stale cache could cause errors. You can export the following environment variable to force textures to be created from scratch.
export CRAFTAX_RELOAD_TEXTURES=true
Scoreboard
If you would like to add an algorithm please open a PR and provide a reference to the source of the results. We report reward as a % of the maximum (226).
Craftax-1B
Algorithm | Reward (% max) | Source |
---|---|---|
PPO-RNN | 15.3 | PureJaxRL |
PPO | 11.9 | PureJaxRL |
ICM | 11.9 | ICM |
E3B | 11.0 | E3B |
Craftax-1M
Algorithm | Reward (% max) | Source |
---|---|---|
PPO-RNN | 2.3 | PureJaxRL |
PPO | 2.2 | PureJaxRL |
ICM | 2.2 | ICM |
E3B | 2.2 | E3B |
See Also
- ⛏️ Crafter The original Crafter benchmark.
- ⚔️ NLE NetHack as an RL environment.
- ⚡ PureJaxRL End-to-end RL implementations in Jax.
- 🌎 JaxUED: CleanRL style UED implementations in Jax.
- 🌍 Minimax: Modular UED implementations in Jax.
- 🏋️ Gymnax: Standard Jax RL interface with classic environments.
- 🧑🤝🧑 JaxMARL: Multi-agent RL in Jax.
Citation
If 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},
}
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.