Add pretrained models.

.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

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)