Training and visualizing of diffusion models from Planning with Diffusion for Flexible Behavior Synthesis.
The main branch contains code for training diffusion models and planning via value-function guided sampling on the D4RL locomotion environments. The kuka branch contains block-stacking experiments. The maze2d branch contains goal-reaching via inpainting in the Maze2D environments.
Updates
Load a pretrained diffusion model and sample from it in your browser with scripts/diffuser-sample.ipynb.
conda env create -f environment.yml
conda activate diffuser
pip install -e .
Download pretrained diffusion models and value functions with:
./scripts/download_pretrained.sh
This command downloads and extracts a tarfile containing this directory to logs/pretrained
. The models are organized according to the following structure:
└── logs/pretrained
├── ${environment_1}
│ ├── diffusion
│ │ └── ${experiment_name}
│ │ ├── state_${epoch}.pt
│ │ ├── sample-${epoch}-*.png
│ │ └── {dataset, diffusion, model, render, trainer}_config.pkl
│ ├── values
│ │ └── ${experiment_name}
│ │ ├── state_${epoch}.pt
│ │ └── {dataset, diffusion, model, render, trainer}_config.pkl
│ └── plans
│ └── defaults
│ ├── 0
│ ├── 1
│ ├── ...
│ └── 149
│
├── ${environment_2}
│ └── ...
The state_${epoch}.pt
files contain the network weights and the config.pkl
files contain the instantation arguments for the relevant classes.
The png files contain samples from different points during training of the diffusion model.
Within the plans
subfolders, there are the results of 150 evaluation trials for each environment using the default hyperparameters.
logs
folder, run python scripts/read_results.py
. (Expand to view the output of this command on the plans downloaded from Google Drive.)
python plotting/table.py
. This will print a table that can be copied into a Latex document. (Expand to view table source.)To plan with guided sampling, run:
python scripts/plan_guided.py --dataset halfcheetah-medium-expert-v2 --logbase logs/pretrained
The --logbase
flag points the experiment loaders to the folder containing the pretrained models.
You can override planning hyperparameters with flags, such as --batch_size 8
, but the default
hyperparameters are a good starting point.
python scripts/train.py --dataset halfcheetah-medium-expert-v2
The default hyperparameters are listed in locomotion:diffusion.
You can override any of them with flags, eg, --n_diffusion_steps 100
.
Train a value function with:
python scripts/train_values.py --dataset halfcheetah-medium-expert-v2
See locomotion:values for the corresponding default hyperparameters.
Plan using your newly-trained models with the same command as in the pretrained planning section, simply replacing the logbase to point to your new models:
python scripts/plan_guided.py --dataset halfcheetah-medium-expert-v2 --logbase logs
See locomotion:plans for the corresponding default hyperparameters.
Deferred f-strings. Note that some planning script arguments, such as --n_diffusion_steps
or --discount
,
do not actually change any logic during planning, but simply load a different model using a deferred f-string.
For example, the following flags:
---horizon 32 --n_diffusion_steps 20 --discount 0.997
--value_loadpath 'f:values/defaults_H{horizon}_T{n_diffusion_steps}_d{discount}'
will resolve to a value checkpoint path of values/defaults_H32_T20_d0.997
. It is possible to
change the horizon of the diffusion model after training (see here for an example),
but not for the value function.
Build the image:
docker build -f Dockerfile . -t diffuser
Test the image:
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 halfcheetah-medium-expert-v2 --logbase logs"
Build the image:
singularity build --fakeroot diffuser.sif Singularity.def
Test the image:
singularity exec --nv --writable-tmpfs diffuser.sif \
bash -c \
"pip install -e . && \
python scripts/train.py --dataset halfcheetah-medium-expert-v2 --logbase logs"
Tag the Docker image (built in the Docker 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
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
.
Download azcopy
: ./azure/download.sh
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
.
To rsync the results from the Azure storage container, run ./azure/sync.sh
.
To mount the storage container:
./azure/make_fuse_config.sh
./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
@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},
}
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.