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
Hashes for waypoint_extraction-0.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | a827f96d2603c63e5ecaf7ad5f229312b147d2f75c7fc2e995cda52ca1404ce9 |
|
MD5 | 119c41210439f0fe484f9170693f053b |
|
BLAKE2b-256 | b954a1d3f5983d06cfafa7f76c01542c6f20c21d4063ded4bc8dc75d26e53ca4 |