Automatic Waypoint Extraction (AWE)

[Project website] [Paper]

This repo contains the implementation of Automatic Waypoint Extraction (AWE): a plug-and-play module for selecting waypoints from demonstrations for performant behavioral cloning. This repo also includes instantiations of combining AWE with two state-of-the-art imitation learning methods, Diffusion Policy and Action Chunking with Transformers (ACT), and the respective benchmarking environments, RoboMimic and Bimanual Simulation Suite.

If you encountered any issue, feel free to contact lucyshi (at) stanford (dot) edu

Installation

1. Clone this repository

git clone git@github.com:lucys0/awe.git
cd awe

2. Create a virtual environment

conda create -n awe_venv python=3.9
conda activate awe_venv

3. Install MuJoCo 2.1

• Download the MuJoCo version 2.1 binaries for Linux or OSX.
• Extract the downloaded mujoco210 directory into ~/.mujoco/mujoco210.

4. Install packages

pip install -e .

RoboMimic

Set up the environment

# install robomimic
pip install -e robomimic/

# install robosuite
pip install -e robosuite/

Download data

# download unprocessed data from the robomimic benchmark
python robomimic/robomimic/scripts/download_datasets.py --tasks lift can square  

# download processed image data from diffusion policy (faster)
mkdir data && cd data
wget https://diffusion-policy.cs.columbia.edu/data/training/robomimic_image.zip
unzip robomimic_image.zip && rm -f robomimic_image.zip && cd ..

Usage

Please replace [TASK] with your desired task to train. [TASK]={lift, can, square}

• Convert delta actions to absolute actions

python utils/robomimic_convert_action.py --dataset=robomimic/datasets/[TASK]/ph/low_dim.hdf5

• Save waypoints

python utils/robomimic_save_waypoints.py --dataset=robomimic/datasets/[TASK]/ph/low_dim.hdf5 --err_threshold=0.005

• Replay waypoints (save 3 videos and 3D visualizations by default)

mkdir video
python example/robomimic_waypoint_replay.py --dataset=robomimic/datasets/[TASK]/ph/low_dim.hdf5 \
    --record_video --video_path video/[TASK]_waypoint.mp4 --task=[TASK] \
    --plot_3d --auto_waypoint --err_threshold=0.005

AWE + Diffusion Policy

Install Diffusion Policy

conda env update -f diffusion_policy/conda_environment.yaml

If the installation is too slow, consider using Mambaforge instead of the standard anaconda distribution, as recommended by the Diffusion Policy authors. That is:

mamba env create -f diffusion_policy/conda_environment.yaml

Train policy

python diffusion_policy/train.py --config-dir=config --config-name=waypoint_image_[TASK]_ph_diffusion_policy_transformer.yaml hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'

Bimanual Simulation Suite

Set up the environment

conda env update -f act/conda_env.yaml

Download data

Please download scripted/human demo for simulated environments from here and save them in data/act/.

If you need real robot data, please contact Lucy Shi: lucyshi (at) stanford (dot) edu

Usage

Please replace [TASK] with your desired task to train. [TASK]={sim_transfer_cube_scripted, sim_insertion_scripted, sim_transfer_cube_human, sim_insertion_human}

• Visualize waypoints

python example/act_waypoint.py --dataset=data/act/[TASK] --err_threshold=0.01 --plot_3d --end_idx=0 

• Save waypoints

python example/act_waypoint.py --dataset=data/act/[TASK] --err_threshold=0.01 --save_waypoints 

AWE + ACT

Train policy

python act/imitate_episodes.py \
    --task_name [TASK] \
    --ckpt_dir data/outputs/act_ckpt/[TASK]_waypoint \
    --policy_class ACT --kl_weight 10 --chunk_size 50 --hidden_dim 512 --batch_size 8 --dim_feedforward 3200 \
    --num_epochs 8000  --lr 1e-5 \
    --seed 0 --temporal_agg --use_waypoint

For human datasets, set --kl_weight=80, as suggested by the ACT authors. To evaluate the policy, run the same command with --eval.

Citation

If you find our code useful for your research, please cite: