PPO code for Procgen (Jax + Flax + Optax)
The repo provides a minimalistic and stable implementation of OpenAI's Procgen PPO algorithm (which is in Tensorflow v.1) in Jax [link to paper] [link to official codebase].
The code's performance was benchmarked on all 16 games using 3 random seeds (see W&B dashboard).
Prerequisites & installation
The only packages needed for the repo are jax, jaxlib, flax and optax.
Jax for CUDA can be installed as follows:
pip install --upgrade "jax[cuda111]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
To install all requirements all at once, run
pip install -r requirements.txt
Use [Training]
After installing Jax, simply run
python train_ppo.py --env_name="bigfish" --seed=31241 --train_steps=200_000_000
to train PPO on bigfish on 200M frames. Right now all returns reported on the W&B dashboard are for training on 200 "easy" levels
and validating both on the same 200 levels as well as the entire "easy" distribution (this can be changed in the train_ppo.py file).
Use [Evaluation]
The code automatically saves the most recent training state checkpoint, which can then be loaded to evaluate the PPO policies. Pre-trained policy weights for 25M frames are included in the repo in the ppo_weights_25M.zip zip file.
To evaluate the pre-trained policies, run
wget https://www.dropbox.com/s/4hskek45cpkbid0/ppo_weights_25M.zip?dl=1 unzip ppo_weights_25M.zip?dl=1 python evaluate_ppo.py --env_name="starpilot" --model_dir="model_weights"
Results
See this W&B report for aggregate performance metrics on all 16 games (easy mode, 200 training levels and all test levels):
Link to results
Cite
To cite this repo in your publications, use
@misc{ppojax, author = {Mazoure, Bogdan}, title = {Jax implementation of PPO}, year = {2021}, publisher = {GitHub}, journal = {GitHub repository}, howpublished = {\url{https://github.com/bmazoure/ppo_jax}}, }
