Repository files navigation

URLB provides a set of leading algorithms for unsupervised reinforcement learning where agents first pre-train without access to extrinsic rewards and then are finetuned to downstream tasks.

This codebase was adapted fromDrQv2. The DDPG agent and training scripts were developed by Denis Yarats. All authors contributed to developing individual baselines for URLB.

We assume you have access to a GPU that can run CUDA 10.2 and CUDNN 8. 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 urlb

We support the following domains.

Each domain supports two observation modes: states and pixels.

To run pre-training use thepretrain.py script

python pretrain.py agent=icm domain=walker

or, if you want to train a skill-based agent, like DIAYN, run:

python pretrain.py agent=diayn domain=walker

This script will produce several agent snapshots after training for100k,500k,1M, and2M frames. The snapshots will be stored under the following directory:

./pretrained_models/<obs_type>/<domain>/<agent>/

For example:

./pretrained_models/states/walker/icm/

Once you have pre-trained your method, you can use the saved snapshots to initialize theDDPG agent and fine-tune it on a downstream task. For example, let's say you have pre-trainedICM, you can fine-tune it onwalker_run by running the following command:

python finetune.py pretrained_agent=icm task=walker_run snapshot_ts=1000000 obs_type=states

This will load a snapshot stored in./pretrained_models/states/walker/icm/snapshot_1000000.pt, initializeDDPG with it (both the actor and critic), and start training onwalker_run using the extrinsic reward of the task.

For methods that use skills, include the agent, and thereward_free tag to false.

python finetune.py pretrained_agent=smm task=walker_run snapshot_ts=1000000 obs_type=states agent=smm reward_free=false

Logs are stored in theexp_local folder. To launch tensorboard run:

tensorboard --logdir exp_local

The console output is also available in a form:

| train | F: 6000 | S: 3000 | E: 6 | L: 1000 | R: 5.5177 | FPS: 96.7586 | T: 0:00:42

a training entry decodes as

F  : total number of environment framesS  : total number of agent stepsE  : total number of episodesR  : episode returnFPS: training throughput (frames per second)T  : total training time