From a25030cc478cf9181f0c76e615365f230d8cfbb1 Mon Sep 17 00:00:00 2001 From: Vincent Micheli Date: Mon, 3 Oct 2022 17:21:16 +0200 Subject: [PATCH] Add pretrained models. --- .gitattributes | 1 + README.md | 15 +++++++++++++++ 2 files changed, 16 insertions(+) create mode 100644 .gitattributes diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..e92cf9b --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.pt filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index caace25..06bcf80 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,7 @@ If you find this code or paper useful, please use the following reference: ## Setup +- Clone the repository. If you want to download the [pretrained models](#pretrained-models) (~3.1 GB), install [Git LFS](https://git-lfs.github.com/) before cloning. To prevent Git LFS from downloading models, set `GIT_LFS_SKIP_SMUDGE=1`. - Install [PyTorch](https://pytorch.org/get-started/locally/) (torch and torchvision). Code developed with torch==1.11.0 and torchvision==0.12.0. - Install [other dependencies](requirements.txt): `pip install -r requirements.txt` - Warning: Atari ROMs will be downloaded with the dependencies, which means that you acknowledge that you have the license to use them. @@ -114,6 +115,20 @@ The folder `results/data/` contains raw scores (for each game, and for each trai Use the notebook `results/results_iris.ipynb` to reproduce the figures from the paper. +## Pretrained models + +Pretrained models are available in `pretrained_models/`. + +- To start a training run from one of these checkpoints, in the section `initialization` of `config/trainer.yaml`, set `path_to_checkpoint` to the corresponding path, and `load_tokenizer`, `load_world_model`, and `load_actor_critic` to `True`. + +- To visualize one of these checkpoints, run the following commands: + ```bash + mkdir checkpoints + cp pretrained_models/Breakout.pt checkpoints/last.pt + ./scripts/play.sh + ``` + with optionally the flags described above for `play.sh`. + ## Credits - [https://github.com/pytorch/pytorch](https://github.com/pytorch/pytorch)