Data-Efficient Reinforcement Learning with Self-Predictive Representations

October 1, 2020 ยท View on GitHub

Max Schwarzer*, Ankesh Anand*, Rishab Goel, R Devon Hjelm, Aaron Courville, Philip Bachman

This repo provides code for implementing the SPR paper

Install

To install the requirements, follow these steps:

# PyTorch
conda install pytorch torchvision -c pytorch
export LC_ALL=C.UTF-8
export LANG=C.UTF-8

# Install requirements
pip install -r requirements.txt

# Finally, clone the project
git clone https://github.com/mila-iqia/spr

Usage:

The default branch for the latest and stable changes is release.

  • To run SPR with augmentation
python -m scripts.run --public --game pong --momentum-tau 1.
  • To run SPR without augmentation
python -m scripts.run --public --game pong --augmentation none --target-augmentation 0 --momentum-tau 0.01 --dropout 0.5

When reporting scores, we average across 10 seeds.

What does each file do?

.
โ”œโ”€โ”€ scripts
โ”‚   โ””โ”€โ”€ run.py                # The main runner script to launch jobs.
โ”œโ”€โ”€ src                     
โ”‚   โ”œโ”€โ”€ agent.py              # Implements the Agent API for action selection 
โ”‚   โ”œโ”€โ”€ algos.py              # Distributional RL loss
โ”‚   โ”œโ”€โ”€ models.py             # Network architecture and forward passes.
โ”‚   โ”œโ”€โ”€ rlpyt_atari_env.py    # Slightly modified Atari env from rlpyt
โ”‚   โ”œโ”€โ”€ rlpyt_utils.py        # Utility methods that we use to extend rlpyt's functionality
โ”‚   โ””โ”€โ”€ utils.py              # Command line arguments and helper functions 
โ”‚
โ””โ”€โ”€ requirements.txt          # Dependencies