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
- Clone this repository
git clone git@github.com:lucys0/awe.git
cd awe
- Create a virtual environment
conda create -n awe_venv python=3.9
conda activate awe_venv
- Install MuJoCo 2.1
- Download the MuJoCo version 2.1 binaries for Linux or OSX.
- Extract the downloaded
mujoco210
directory into~/.mujoco/mujoco210
.
- 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
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
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
Algorithm | Hash digest | |
---|---|---|
SHA256 | f81077fae69ede969a3b60f1e8f55fa02327a52e99777e0cdb552b14c355dde7 |
|
MD5 | d9186b7c8ca729154c822004f35120c3 |
|
BLAKE2b-256 | daf8feff745f6b3b3443c38d1dd7a9025e2b7bb594eeeaaeafa7448bf76bef25 |
Provenance
File details
Details for the file waypoint_extraction-0.1-py3-none-any.whl
.
File metadata
- Download URL: waypoint_extraction-0.1-py3-none-any.whl
- Upload date:
- Size: 58.7 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.10.10
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | a827f96d2603c63e5ecaf7ad5f229312b147d2f75c7fc2e995cda52ca1404ce9 |
|
MD5 | 119c41210439f0fe484f9170693f053b |
|
BLAKE2b-256 | b954a1d3f5983d06cfafa7f76c01542c6f20c21d4063ded4bc8dc75d26e53ca4 |