Planning with Diffusion   

Training and visualizing of diffusion models from Planning with Diffusion for Flexible Behavior Synthesis. Guided sampling code to come soon!

Quickstart

Load a pretrained diffusion model and sample from it in your browser with scripts/diffuser-sample.ipynb.

Installation

conda env create -f environment.yml
conda activate diffusion
pip install -e .

Usage

Train a diffusion model with:

python scripts/train.py --dataset hopper-medium-replay-v2 \
    --horizon 512 --n_diffusion_steps 200

The default hyperparameters are listed in config/locomotion.py. You can override any of them with runtime flags, eg --batch_size 64.

Docker

1. Build the container:

docker build -f azure/Dockerfile . -t diffuser

2. Test the container:

docker run -it --rm --gpus all \
    --mount type=bind,source=$PWD,target=/home/code \
    --mount type=bind,source=$HOME/.d4rl,target=/root/.d4rl \
    diffuser \
    bash -c \
    "export PYTHONPATH=$PYTHONPATH:/home/code && \
    python /home/code/scripts/train.py --dataset hopper-medium-expert-v2 --logbase logs/docker"

Running on Azure

Setup

1. Launching jobs on Azure requires one more python dependency:

pip install git+https://github.com/JannerM/doodad.git@janner

2. Tag the image built in the previous section and push it to Docker Hub:

export DOCKER_USERNAME=$(docker info | sed '/Username:/!d;s/.* //')
docker tag diffuser ${DOCKER_USERNAME}/diffuser:latest
docker image push ${DOCKER_USERNAME}/diffuser

3. Update azure/config.py, either by modifying the file directly or setting the relevant environment variables. To set the AZURE_STORAGE_CONNECTION variable, navigate to the Access keys section of your storage account. Click Show keys and copy the Connection string.

4. Download azcopy: ./azure/download.sh

Usage

Launch training jobs with python azure/launch.py. The launch script takes no command-line arguments; instead, it launches a job for every combination of hyperparameters in params_to_sweep.

Viewing results

To rsync the results from the Azure storage container, run ./azure/sync.sh.

To mount the storage container:

1. Create a blobfuse config with ./azure/make_fuse_config.sh
2. Run ./azure/mount.sh to mount the storage container to ~/azure_mount

To unmount the container, run sudo umount -f ~/azure_mount; rm -r ~/azure_mount

Reference

@inproceedings{janner2022diffuser,
  title = {Planning with Diffusion for Flexible Behavior Synthesis},
  author = {Michael Janner and Yilun Du and Joshua B. Tenenbaum and Sergey Levine},
  booktitle = {International Conference on Machine Learning},
  year = {2022},
}

Acknowledgements

The diffusion model implementation is based on Phil Wang's denoising-diffusion-pytorch repo. The organization of this repo and remote launcher is based on the trajectory-transformer repo.