Official code repository for Prompt-DT. [website][paper]
Prompt-DT Architecture:
We tested the code in Ubuntu 20.04.
- We recommend using Anaconda to create a virtual environment.
conda create --name prompt-dt python=3.8.5
conda activate prompt-dt
-
Our experiments require MuJoCo as well as mujoco-py. Install them by following the instructions in the mujoco-py repo.
-
Install environments and dependencies with the following commands:
# install dependencies
pip install -r requirements.txt
# install environments
./install_envs.sh
- We log experiments with wandb. Check out the wandb quickstart doc to create an account.
- We share example datasets via this Google Drive link.
- Download the "data" folder.
wget -O data.zip 'https://drive.google.com/uc?export=download&id=1rZufm-XRq1Ig-56DejkQUX1si_WzCGBe&confirm=True'
unzip data.zip
rm data.zip
- Organize folders as follows.
.
├── config
├── data
│ ├── ant_dir
│ ├── cheetah_dir
│ ├── cheetah_vel
│ └── ML1-pick-place-v2
├── envs
├── prompt_dt
└── ...
# Prompt-DT
python pdt_main.py --env cheetah_dir # choices:['cheetah_dir', 'cheetah_vel', 'ant_dir', 'ML1-pick-place-v2']
# Prompt-MT-BC
python pdt_main.py --no-rtg --no-r
# MT-ORL
python pdt_main.py --no-prompt
# MT-BC-Finetune
python pdt_main.py --no-prompt --no-rtg --no-r --finetune
The code for prompt-dt is based on decision-transformer. We build environments based on repos including macaw, rand_param_envs, and metaworld.
If you find our code helpful for your research, please consider citing the paper!
@inproceedings{xu2022prompting,
title={Prompting Decision Transformer for Few-Shot Policy Generalization},
author={Xu, Mengdi and Shen, Yikang and Zhang, Shun and Lu, Yuchen and Zhao, Ding and Tenenbaum, Joshua and Gan, Chuang},
booktitle={International Conference on Machine Learning},
pages={24631--24645},
year={2022},
organization={PMLR}
}
Suggestions for enhancing and improving the code are welcome. Please email mengdixu@andrew.cmu.edu with comments and suggestions.