This code is a part of supplementary material of NeurIPS submission named 'Diffused Task Agnostic Milestone Planner'
The code contains the authors' implementation of diffused task agnostic milestone planner (DTAMP)

Setup

1. Install D4RL:
The D4RL environments and dataset can be installed by following the instructions presented in the following github repository:

https://github.com/Farama-Foundation/D4RL

The installation of D4RL also requires to install gym, mujoco-py, mjrl and dm_control.


2. Install CALVIN:
The CALVIN environment can be installed by following the instructions presented in the following github repository:

https://github.com/mees/calvin


3. Download CALVIN dataset:
As Rostete-beas et. al. used different data split from the original dataset,
we recomand to use the dataset provided in the following github repository:

https://github.com/ErickRosete/tacorl


4. Install requirements:
All the requirements can be installed by running:

pip install -r requirements.txt


5. Training DTAMP on D4RL dataset:
Run:

export PYTHONPATH=$PYTHONPATH:/path/to/dtamp
python scripts/train.py --algo dtamp --env antmaze-medium-play-v2


6. Training DTAMP on CALVIN dataset:
1) You should first preprocess the data by running:

python data_modules/preprocess_calvin_data.py --source_data_dir /path/to/data/before/processed --target_data_dir /path/to/data/after/processed

2) Train LMP for learning skills:

python scripts/train.py --algo play_lmp --env calvin

3) Add the learned skills to the data:

python data_modules/prepare_skill_dataset.py

4) Train DTAMP:

python scripts/train.py --algo dtamp --env calvin


7. Evaluation
The trained models can be evaluated by running:

1) D4RL experiments:

python scripts/evaluate_d4rl.py --algo dtamp --env antmaze-medium-play-v2 --checkpoint_dir /path/to/checkpoint --multigoal (if you want to run multi-goal antmaze experimments)

2) CALVIN experiments:

python scripts/evaluate_calvin.py --algo dtamp --n_tasks 1 --checkpoint_dir /path/to/checkpoint --lmp_model_dir /path/to/lmp_checkpoint


8. Changing configurations
All the default configurations that used in our paper are presented in 'configs' folder.
You can may change configurations by editing yaml files.


The code was tested on python3.8.

