Learning Coordinated Bimanual Manipulation Policies using State Diffusion and Inverse Dynamics Models

Haonan Chen¹, Jiaming Xu^*1, Lily Sheng^*1, Tianchen Ji¹ Shuijing Liu³ Yunzhu Li², Katherine Driggs-Campbell¹

¹University of Illinois, Urbana-Champaign, ²Columbia University, ³The University of Texas at Austin

Environment Setup

We recommend using Mambaforge over the standard Anaconda distribution for a faster installation process. Create your environment using:

1. Install the necessary dependencies:

   sudo apt install -y libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf libglm-dev
   

2. Clone the repository:

   git clone https://github.com/haonan16/state_diff.git
   cd state_diff/
   

3. Update mamba, and create and activate the environment:

   To update mamba and create the environment, use the following commands:

   mamba install mamba=1.5.1 -n base -c conda-forge
   mamba env create -f conda_environment.yaml
   mamba activate coord_bimanual
   

Data Preparation

mkdir -p data

Place the pushl dataset in the data folder. The directory structure should be:

To obtain the dataset, download the corresponding zip file and unzip it from the following link:

• PushL Dataset

You can automate this process by running the following command:

gdown https://drive.google.com/uc?id=1433HHxOH5nomDZ12aUZ4XZel4nVO_uwL -O data/pushl_dataset.zip
unzip data/pushl_dataset.zip -d data
rm data/pushl_dataset.zip

For datasets from other simulation benchmarks, you can find them here:

• Simulation Benchmark Datasets from Diffusion Policy

Once downloaded, extract the contents into the data folder.

Directory structure:

data/
└── pushl_dataset/

Demo, Training and Eval

Activate conda environment and login to wandb (if you haven't already).

conda activate coord_bimanual
wandb login

Collecting Human Demonstration Data

Run the following script to collect human demonstration:

python demo_push_data_collection.py

The agent agent position can be controlled by the mouse. The following keys can be used to control the environment:

• Space - Pause and step forward (increments plan_idx)
• R - Retry current attempt
• Q - Exit the script

Training

To launch training, run:

python train.py \
  --config-name=train_diffusion_trajectory_unet_lowdim_workspace.yaml 

Acknowledgement

• Policy training implementation is adapted from Diffusion Policy.