robot-awe 0.1

Description:

robotawe 0.1

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

Clone this repository

git clone [email protected]: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: