Craftax Baselines

This repository contains the code for running the baselines from the Craftax paper. For packaging reasons, this is separate to the main repository.

Installation

git clone https://github.com/MichaelTMatthews/Craftax_Baselines.git
cd Craftax_Baselines
pip install -r requirements.txt -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pre-commit install

Run Experiments

PPO

python ppo.py

PPO-RNN

python ppo_rnn.py

ICM

python ppo.py --train_icm

E3B

python ppo.py --train_icm --use_e3b --icm_reward_coeff 0

RND

python ppo_rnd.py

Visualisation

You can save trained policies with the --save_policy flag. These can then be viewed with the view_ppo_agent script (pass in the path up to the files directory).