neverix / diffuser

Code for the paper "Planning with Diffusion for Flexible Behavior Synthesis"

Planning with Diffusion   

Training and visualizing of diffusion models from Planning with Diffusion for Flexible Behavior Synthesis. This branch has the Maze2D experiments and will be merged into main shortly.

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 --config config.maze2d --dataset maze2d-large-v1

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

Plan using the diffusion model with:

python scripts/plan_maze2d.py --config config.maze2d --dataset maze2d-large-v1

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.