Skip to main content

research project

Project description

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.

Given a set of demonstrations, extracting waypoints is as simple as:

pip install robot-awe
from robot_awe import extract_waypoints
waypoints = extract_waypoints(states, err_threshold)

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
  1. Create a virtual environment
conda create -n awe_venv python=3.9
conda activate awe_venv
  1. Install MuJoCo 2.1
  • Download the MuJoCo version 2.1 binaries for Linux or OSX.
  • Extract the downloaded mujoco210 directory into ~/.mujoco/mujoco210.
  1. 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:


Project details


Release history Release notifications | RSS feed

This version

0.1

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

waypoint_extraction-0.1.tar.gz (48.3 kB view details)

Uploaded Source

Built Distribution

waypoint_extraction-0.1-py3-none-any.whl (58.7 kB view details)

Uploaded Python 3

File details

Details for the file waypoint_extraction-0.1.tar.gz.

File metadata

  • Download URL: waypoint_extraction-0.1.tar.gz
  • Upload date:
  • Size: 48.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.2 CPython/3.10.10

File hashes

Hashes for waypoint_extraction-0.1.tar.gz
Algorithm Hash digest
SHA256 f81077fae69ede969a3b60f1e8f55fa02327a52e99777e0cdb552b14c355dde7
MD5 d9186b7c8ca729154c822004f35120c3
BLAKE2b-256 daf8feff745f6b3b3443c38d1dd7a9025e2b7bb594eeeaaeafa7448bf76bef25

See more details on using hashes here.

Provenance

File details

Details for the file waypoint_extraction-0.1-py3-none-any.whl.

File metadata

File hashes

Hashes for waypoint_extraction-0.1-py3-none-any.whl
Algorithm Hash digest
SHA256 a827f96d2603c63e5ecaf7ad5f229312b147d2f75c7fc2e995cda52ca1404ce9
MD5 119c41210439f0fe484f9170693f053b
BLAKE2b-256 b954a1d3f5983d06cfafa7f76c01542c6f20c21d4063ded4bc8dc75d26e53ca4

See more details on using hashes here.

Provenance

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page