• We've spent months optimizing this code to make diffusion models train fast in CALVIN, and can now offer fully training to 300K gradient steps in 8 hours, and evaluation on 1000 tasks in 6 hours. We get these numbers on the A10G 24GB.

• You can easily adjust the amount of burden to put on the high level vs the low level policy by simply adjusting the temporal stride, or the clip_stride as referred to in the code

• Our high level diffusion policy is independent of the low level policy. This means you can plug and play different low level policies and environments.

• We achieve an average 88.7% success on the horizon length 1 SR of the multitask long-horizon control problem in CALVIN. This is the highest performance yet of any model that incorporates no inductive biases on the problem.

• We offer a daemon script that runs in the background, caching the entire shared memory dataset of HULC in the background to save you ~20 minutes at the start of each training run, which aggregates to hours and possibly days over the course of a full research project

• Supports for Python 3.8 and higher.
• Simple and flexible development interface for fast extensibility to new environments and models
• Poetry as the dependencies manager. See configuration in pyproject.toml and setup.cfg.
• Automatic codestyle with black, isort and pyupgrade.
• Ready-to-use pre-commit hooks with code-formatting.
• Type checks with mypy; docstring checks with darglint; security checks with safety and bandit
• Testing with pytest.
• Ready-to-use .editorconfig, .dockerignore, and .gitignore. You don't have to worry about those things.

An example of a recent LCD training run for your reference can be found here. We will also be releasing the entire experiment logs of all experiments done for this project soon.

We require either Mambaforge or Conda and git lfs. We highly recommend Mambaforge (a faster drop-in replacement of conda). Conda is also supported. To install, simply run

$ make install

This will set up a new conda environment lcd. It will also download all necessary data (~9.5 GB/), including

• HULC seeds
• On-policy offline datasets
• LCD seeds

If you would like to install just the repository without downloading data, run

$ make NO_DATA=1 install

If you just want to download the data, run

$ git clone https://github.com/ezhang7423/hulc-data.git --recurse-submodules

We have thoroughly tested the installation process and environment with NVIDIA GPUs CUDA Versions 11.6, 11.7, and 12.0 on Ubuntu 18.04, 20.04, and AlmaLinux 8.7 (similar to RHEL). Running on windows or macOS will likely present difficulties. If any part of the installation fails, you can activate the environment with mamba/conda activate lcd and try debugging what the issue is by running individual commands in the install section of the Makefile.

If you see this error:

UnsatisfiableError: The following specifications were found
to be incompatible with the existing python installation in your environment:
Specifications:
  - pytorch3d=0.7.2 -> python[version='>=2.7,<2.8.0a0|>=3.11,<3.12.0a0|>=3.7,<3.8.0a0|>=3.6,<3.7.0a0|>=3.5,<3.6.0a0']
Your python: python=3.8

try using mamba instead of conda.

Once the repository is installed, all scripts can be run with lcd. This is also equivalent to directly running python ./src/lcd/__main__.py, which may be preferable when using a debugger.

> lcd                                                          (lcd) 
                                                                     
 Usage: lcd [OPTIONS] COMMAND [ARGS]...                              
                                                                     
╭─ Options ─────────────────────────────────────────────────────────╮
│ --install-completion          Install completion for the current  │
│                               shell.                              │
│ --show-completion             Show completion for the current     │
│                               shell, to copy it or customize the  │
│                               installation.                       │
│ --help                        Show this message and exit.         │
╰───────────────────────────────────────────────────────────────────╯
╭─ Commands ────────────────────────────────────────────────────────╮
│ rollout     Rollout in the environment for evaluation or dataset  │
│             collection                                            │
│ train_hulc  Train the original hulc model                         │
│ train_lcd   Train the original hulc model                         │
╰───────────────────────────────────────────────────────────────────╯

Train HULC

lcd train_hulc

Evaluate HULC

lcd rollout hulc

Generate on-policy data

lcd rollout generate

Train LCD

lcd train_lcd

Evaluate LCD

lcd rollout lcd

It's really that easy! If you would like to pass more training options to hulc or diffuser, both can be done in the exact same way as in the original repositories (see here for hulc and here for diffuser). These scripts will simply pass the arguments to the underlying original code. It should be noted that the dataset key is no longer supported for diffuser as there is only one dataset, and that the task_D_D dataset needs to be downloaded first by following the directions at ./language-control-diffusion/submodules/hulc-baseline/dataset.

If using wandb, please change the key wanb_entity in ./lcd/config/calvin.py and ./submodules/hulc-baseline/conf/logger/wandb.yaml to your team or username.

You can see the list of available releases on the GitHub Releases page.

We follow Semantic Versions specification, and use the Release Drafter. As pull requests are merged, a draft release is kept up-to-date listing the changes, ready to publish when you’re ready. With the categories option, you can categorize pull requests in release notes using labels.

GitHub creates the bug, enhancement, and documentation labels for you. Dependabot creates the dependencies label. Create the remaining labels on the Issues tab of your GitHub repository, when you need them.

.
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── LICENSE
├── Makefile
├── README.md
├── SECURITY.md
├── cookiecutter-config-file.yml
├── poetry.lock
├── pyproject.toml
├── requirements.txt
├── setup.cfg
├── submodules
│   ├── hulc-baseline # original hulc code
│   └── hulc-data
│       ├── (indices,ordering,seq).pt # evaluation sequences sorted by task, useful for single task evaluation
│       ├── README.md
│       ├── annotations.json # the training and evaluation language descriptions of all tasks
│       ├── default_1000_sequences.pt # the default 1000 tasks that are generated during evaluation. this is cached to speed up eval by ~3 min per run.
│       ├── hulc-baselines-30 # a collection of 3 seeds of HULC, each run for 30 epochs trained in the same manner with the original code checkpointed at hulc-baseline
│       ├── hulc-trajectories # the offline on-policy datasets used to train lcd. these are saved in the latent space of the LLP, after processing through its encoder
│       ├── lcd-seeds # high level diffusion policies trained on the aforementioned seeds 
│       └── t5-v1_1-xxl_embeddings.pt # t5 XXL embeddings of all annotations
└── src
    └── lcd
        ├── __init__.py
        ├── __main__.py # the primary entrypoint. delegates commands by automatically parsing the apps directory
        ├── __pycache__
        ├── apps # this holds all the ways of interacting with this repo
        │   ├── __pycache__
        │   ├── rollout.py # rollout allows you to evaluate hulc, lcd, and collect on-policy datasets
        │   ├── train_hulc.py # a wrapper around the original hulc training code in hulc-baseline.
        │   └── train_lcd.py # a wrapper around the diffuser training code, which calls ./src/lcd/apps/diffuser.py
        ├── config
        │   ├── __pycache__
        │   └── calvin.py # holds all configuration for training diffuser on the calvin benchmark
        ├── datasets
        │   ├── __init__.py
        │   ├── __pycache__
        │   └── sequence.py # load the prior mentioned on-policy datasets
        ├── models
        │   ├── __init__.py
        │   ├── __pycache__
        │   ├── diffusion.py
        │   ├── helpers.py
        │   └── temporal.py
        ├── scripts # older entrypoints that will be transitioned to apps
        │   ├── diffuser.py # training and evaluation of diffuser
        │   └── generation
        │       ├── dataset.py # generate the on-policy dataset from the goal space
        │       ├── embeddings.py # generate the t5 XXL embeddings
        │       └── task_orderings.py # generate an updated version of the (indices,ordering,seq).pt file in ./submodules/hulc-data
        └── utils
            ├── __init__.py
            ├── __pycache__
            ├── arrays.py
            ├── config.py
            ├── eval.py # the meat of the rollout and evaluation code
            ├── git_utils.py
            ├── serialization.py
            ├── setup.py
            ├── timer.py
            └── training.py # the meat of the diffuser training code

1. Install and initialize poetry and install pre-commit hooks:

make install
make pre-commit-install

2. Run the codestyle:

make codestyle

Makefile contains a lot of functions for faster development.

make poetry-download

make poetry-remove

make install

make pre-commit-install

make codestyle
# or use synonym
make formatting

make check-codestyle

make update-dev-deps

make check-safety

make check-safety

make mypy

make test

make lint

make test && make check-codestyle && make mypy && make check-safety

make docker-build

make docker-build VERSION=latest

make docker-remove

make pycache-remove

make build-remove

make dsstore-remove

make mypycache-remove

make cleanup

Want to know more about Poetry? Check its documentation.

Replanning, further probing the language generalization, trying discrete state-action spaces, scaling the model are all interesting research directions to explore. As far as further improvements to this repository go, we are planning to add

1. Parallel training (multiple runs, one per gpu)
2. Diffuser from pixels
3. BC-Z and RT-1 ablations

This project is licensed under the terms of the MIT license. See LICENSE for more details.

@article{language-control-diffusion,
  author={Zhang, Edwin and Lu, Yujie and Wang, William and Zhang, Amy},
  title={LAD: Language Control Diffusion: efficiently scaling through Space, Time, and Tasks},
  year = {2023},
  journal={arXiv preprint arXiv:2210.15629},
  howpublished = {\url{https://github.com/ezhang7423/language-control-diffusion}}
}

Massive thanks to Oier Mees and Luka Shermann for providing te CALVIN and HULC codebases, and huge thanks to Michael Janner and Yilun Du for providing the Diffuser codebase. This work would not be possible without standing on the shoulders of these giants.

Template: python-package-template