GraphDiffusionImitate

Caio Freitas

Official implementation of master's thesis "Graph Diffusion for Imitation Leraning in Robotics" at IAS/PEARL - TU Darmstadt, Germany.

Installation

To install the dependencies, run

$ mamba env create -f conda_environment.yaml

On using Robosuite

When using robosuite, the environments are based on v1.4.1 but using the direct joint values as well in the actions. To add that to the robosuite environments, the file robosuite/robosuite/models/robots/robot.py should be modified to include (flag as True) the robot{i}_joint_pos in the actives list in the setup_observables() method (line 204).

Downloading training data

To download the training data used (from the Robomimic project), use the following commands from within the robomimic/robomimic/scrips directory:

# Download from robomimic
python download_datasets.py --tasks lift square transport --dataset_types ph mh --hdf5_types low_dim
# Move to /data directory
cd ../../datasets
mv * ~/workspace/GraphDiffusionImitate/data/

Testing

Tests live in the tests/ directory and use pytest. Make sure the imitation conda environment is active before running any tests.

mamba activate imitation

Run all tests

pytest tests/ -v

Graph Diffusion Policies

GraphDDPM

Denoising diffusion on top of the graph's node features, based on Denoising Diffusion Probabilistic Models (DDPM) from Ho et al..

ARGG (Autorregressive Graph Generative Policy)

Autoregressive graph generation for an graph-structured robot action, based on the GraphARM model from Kong et al..

Acknowledgements

This repo relies on the following existing codebases / datasets:

• Dataset from lift, square and transport tasks from robosuite
• Benchmark model from diffusion_policy
• EGNN implementation from Satorras