Skip to content


Repository files navigation

Proto-RL: Reinforcement Learning with Prototypical Representations

This is a PyTorch implementation of Proto-RL from

Reinforcement Learning with Prototypical Representations by

Denis Yarats, Rob Fergus, Alessandro Lazaric, Lerrel Pinto.



If you use this repo in your research, please consider citing the paper as follows

    title={Reinforcement Learning with Prototypical Representations},
    author={Denis Yarats and Rob Fergus and Alessandro Lazaric and Lerrel Pinto},


We assume you have access to a gpu that can run CUDA 11. Then, the simplest way to install all required dependencies is to create an anaconda environment by running

conda env create -f conda_env.yml

After the instalation ends you can activate your environment with

conda activate proto


In order to pretrain the agent you need to specify the number of task-agnostic environment steps by setting num_expl_steps, after that many steps, the agent will start receving the downstream task reward until it takes num_train_steps in total. For example, to pre-train the Proto-RL agent on Cheetah Run task unsupervisely for 500k environment steps and then train it further with the downstream reward for another 500k steps, you can run:

python env=cheetah_run num_expl_steps=250000 num_train_steps=500000

Note that we divide the number of steps by action repeat, which is set to 2 for all the environments.

This will produce the exp_local folder, where all the outputs are going to be stored including train/eval logs, tensorboard blobs, and evaluation episode videos. To launch tensorboard run

tensorboard --logdir exp_local