/diffusion-planning

Planning in large domains with continuous diffusion and discrete planners - Andreas Burger, Jack Sun, Yifan Ruan

Primary LanguagePythonMIT LicenseMIT

Planning in large domains with continuous diffusion and discrete planners

Andreas Burger, Jack Sun, Yifan Ruan

(Work in progress) Go to scripts/run_diffusion_planner.py! \

Structure:

  • running diffusion planner: scripts/run_diffusion_planner.py
  • original diffuser: scripts/train_diffuser.py (to train) and scripts/run_diffuser_maze2d.py (to run)
  • original diffusion model: diffuser/models/diffusion.py and diffuser/utils/training.py
  • our large maze: diffuser/planning/largemaze2d.py
  • our planner diffuser/planning/planner.py

Installation

First you need to install MuJoCo 2.0.
Follow section 1 here!
See the official instructions for more infos.

Set up the environment:

conda env create -f environment.yml # this will partly fail, but it will create the environment
conda activate diffuser
pip install -e .

pip install setuptools==65.5.0 pip==21  # gym 0.21 installation is broken with more recent versions
pip install wheel==0.38.0
pip install -r requirements.txt 

# check your cuda version with nvidia-smi
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# on Ubuntu 22.04
# https://github.com/openai/mujoco-py/issues/492#issuecomment-607688770
sudo apt-get install patchelf
sudo apt-get install libglu1-mesa-dev mesa-common-dev
# install old gcc version. Ideally gcc-7
gcc --version # check your gcc version
sudo apt-get install build-essential
sudo apt install gcc-9 g++-9
sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-9 9
sudo update-alternatives --config gcc # select gcc-9
pip install mujoco_py==2.0.2.8
pip install "cython<3"
sudo update-alternatives --config gcc # select whatever gcc version you had before

# ./scripts/download_pretrained.sh

Run

Our diffusion planner:

python scripts/run_diffusion_planner.py
# python scripts/run_diffusion_planner.py --config config.maze2d --dataset maze2d-large-v1-3x3

Original diffuser:

# training # ~20h on RTX 3060
python scripts/train_diffuser.py --config config.maze2d --dataset maze2d-large-v1 
# inference
python scripts/run_diffuser_maze2d.py --config config.maze2d --dataset maze2d-large-v1

Maze2d

In this repo

  • Using the env: scripts/run_diffuser_maze2d.py
  • Loading the env from d4rl: diffuser/datasets/d4rl.py
  • Rendering, maze size: diffuser/utils/rendering.py

Outside this repo

Resources

This codebase is based on