davrempe/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.

Release status: Diffusion model training and visualization has been merged into the main branch. Maze2d planning is currently in the maze2d branch, but will be merged into main soon. Robotic block-stacking experiments are in the kuka branch. Classifier-guided sampling code will be cleaned up 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.