This repo contains code for the paper
Learning One Representation to Optimize All Rewards. Ahmed Touati, Yann Ollivier. NeurIPS 2021
pip install 'gym[atari]'
pip install torch
pip install opencv-python
# Baselines for Atari preprocessing
# Tensorflow is a dependency, but you don't need to install the GPU version
conda install tensorflow
pip install git+git://github.com/openai/baselines
# AtariARI (Atari Annotated RAM Interface)
pip install git+git://github.com/mila-iqia/atari-representation-learning.git
If you want to use GPU, just add the flag --cuda
.
- train discrete maze:
python grid_main.py \
--agent FB \
--n-cycles 25 \
--n-test-rollouts 10 \
--num-rollouts-per-cycle 4 \
--update-eps 1 \
--soft-update \
--temp 200 \
--seed 0 \
--gamma 0.99 \
--lr 0.0005 \
--polyak 0.95 \
--embed-dim 100 \
--w-sampling cauchy_ball \
--n-epochs 200 \
- train continuous maze:
python continuous_main.py \
--agent FB \
--n-cycles 25 \
--n-test-rollouts 10 \
--num-rollouts-per-cycle 4 \
--update-eps 1 \
--soft-update \
--temp 200 \
--seed 0 \
--gamma 0.99 \
--lr 0.0005 \
--polyak 0.95 \
--embed-dim 100 \
--w-sampling cauchy_ball \
--n-epochs 200 \
- train atari:
python atari_main.py \
--agent FB \
--n-cycles 25 \
--n-test-rollouts 10 \
--num-rollouts-per-cycle 2 \
--update-eps 0.2 \
--soft-update \
--temp 200 \
--seed 0 \
--gamma 0.9 \
--lr 0.0005 \
--polyak 0.95 \
--embed-dim 100 \
--w-sampling cauchy_ball \
--n-epochs 200 \