Continual Policy Distillation from Distributed Reinforcement Learning Teachers

Contributor: Yuxuan Li, Qijun He

Brief Introduction

Official Repo for our paper Continual Policy Distillation from Distributed Reinforcement Learning Teachers . Our work focuses on a teacher-student training architecture, in which we train single-task teacher models in a distributed way, and then sequentially distill them to a central model.

Setup

1. We use Metaworld as the benchmark, so you can refer to metaworld algorithm as the settings. (About setup of metaworld_algorithms, see here)
2. We also offer a list of required packages to run our codes. (See ./environment.yml)

How to run

• We have 2 main experiments in our work:
  1. 2-phase continual distillation for MT25 (from 10 tasks to 25 tasks);
  2. 5-phase continual distillation for MT25 (5 tasks per phase).
• We have extra codes for the following purposes:
  • distribute_run: training single-task teacher models
  • train_task_embedder: use trained teachers to pre-train task embedder for the central model

Check the file tree to find our codes:

.
|__examples
|  |__multi_task
|     |__dtkw_ablation_cfg.py     # run for baselines/5-phase distill
|     |__dtkw_distill_mt25.py     # run for 2-phase distill
|     |__dtkw_distribute_mt25.py  # run for training teachers
|     |__dtkw_task_embedder.py    # run for training task embedder
|
|__metaworld_algorithms
|  |__config
|  |  |__networks.py              # define config of metaworld_algorithms/rl/networks
|  |  |__nn.py                    # define config of metaworld_algorithms/nn
|  |
|  |__nn
|  |  |__switch_transformer.py    # define inner layers of network
|  |
|  |__rl
|  |  |__algorithms
|  |  |  |__dtkw.py               # algorithms of the central model
|  |  |  |__taskembedder.py       # algorithms of the task embedder
|  |  |  |__worker_ppo.py         # algorithms of training the teachers
|  |  |__networks.py              # define networks for the central model and task embedder
|  |
|  |__distill_run.py              # run for 2-phase distill
|  |__distribute_run.py           # run for training teachers
|  |__mt25_ours.py
|  |__...                         # run for baselines/5-phase distill
|  |__train_task_embedder.py      # run for training task embedder
|
|__run.sh

What is needed to replicate our work?

• dtkw_distribute_mt25.py: train from scratch
• dtkw_task_embedder.py: all teacher models for training task embedder are needed. You need to save the trained teachers into one directory (refer to it by argument data_dir in run.sh)
• dtkw_distill_mt25.py: all teacher models and the task embedder are needed. You need to save the trained teachers into one directory (refer to it by argument data_dir in run.sh). You need to save task embedder and configurate its path through function _init_task_embedder() in /metaworld_algorithms/distill_run.py.
• dtkw_ablation_cfg.py: same as dtkw_distill_mt25.py

More requirements

Please refer to the recommended hyper-params we defined in /examples.

License