flairox / jafar

JAX reimplementation of the DeepMind paper "Genie: Generative Interactive Environments"

Jafar: A JAX-based Genie Implementation 🧞

Jafar is a JAX-based implementation of the DeepMind paper "Genie: Generative Interactive Environments" (Bruce et al., 2024).

Jafar supports training of all Genie components and can complete the CoinRun reproducibility experiment (Appendix F) on a single L40S GPU in under a week.

Setup 🧗

Jafar was built with python 3.10 and jax 0.4.30. To install requirements, run:

pip install -r requirements.txt

Before training the models, generate the CoinRun dataset by running:

python generate_dataset.py --num_episodes 10000

Note: this is a large dataset (around 100GB) and may take a while to generate.

Quick Start 🚀

Genie has three components: a video tokenizer, a latent action model, and a dynamics model. Each of these components are trained separately, however, the dynamics model requires a pre-trained video tokenizer and latent action model.

To train the video tokenizer (similar for the LAM), run:

python train_tokenizer.py --ckpt_dir <path>

Once the tokenizer and LAM are trained, the dynamics model can be trained with:

python train_dynamics.py --tokenizer_checkpoint <path> --lam_checkpoint <path>

Logging with wandb is supported. To enable logging, set the WANDB_API_KEY environment variable or run:

wandb login

Training can then be logged by setting the --log flag:

python train_tokenizer.py --log --entity <wandb-entity> --project <wandb-project>

Citing Jafar 📜

Jafar was built by Matthew Jackson and Timon Willi.

If you use Jafar in your work, please cite us and the original Genie paper as follows:

@inproceedings{
    willi2024jafar,
    title={Jafar: An Open-Source Genie Reimplemention in Jax},
    author={Timon Willi and Matthew Thomas Jackson and Jakob Nicolaus Foerster},
    booktitle={First Workshop on Controllable Video Generation @ ICML 2024},
    year={2024},
    url={https://openreview.net/forum?id=ZZGaQHs9Jb}
}

@inproceedings{
    bruce2024genie,
    title={Genie: Generative Interactive Environments},
    author={Jake Bruce and Michael D Dennis and Ashley Edwards and Jack Parker-Holder and Yuge Shi and Edward Hughes and Matthew Lai and Aditi Mavalankar and Richie Steigerwald and Chris Apps and Yusuf Aytar and Sarah Maria Elisabeth Bechtle and Feryal Behbahani and Stephanie C.Y. Chan and Nicolas Heess and Lucy Gonzalez and Simon Osindero and Sherjil Ozair and Scott Reed and Jingwei Zhang and Konrad Zolna and Jeff Clune and Nando de Freitas and Satinder Singh and Tim Rockt{\"a}schel},
    booktitle={Forty-first International Conference on Machine Learning},
    year={2024},
    url={https://openreview.net/forum?id=bJbSbJskOS}
}