Package: waypoint-extraction
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 PolicyandAction Chunking with Transformers (ACT), and the respective benchmarking environments,RoboMimicandBimanual Simulation Suite.Given a set of demonstrations, extracting waypoints is as simple as:pipinstallrobot-awe
fromrobot_aweimportextract_waypointswaypoints=extract_waypoints(states,err_threshold)If you encountered any issue, feel free to contact lucyshi (at) stanford (dot) eduInstallationClone this repositorygitclonegit@github.com:lucys0/awe.gitcdaweCreate a virtual environmentcondacreate-nawe_venvpython=3.9
condaactivateawe_venvInstall MuJoCo 2.1Download the MuJoCo version 2.1 binaries forLinuxorOSX.Extract the downloadedmujoco210directory into~/.mujoco/mujoco210.Install packagespipinstall-e.RoboMimicSet up the environment# install robomimicpipinstall-erobomimic/# install robosuitepipinstall-erobosuite/Download data# download unprocessed data from the robomimic benchmarkpythonrobomimic/robomimic/scripts/download_datasets.py--tasksliftcansquare# download processed image data from diffusion policy (faster)mkdirdata&&cddata
wgethttps://diffusion-policy.cs.columbia.edu/data/training/robomimic_image.zip
unziprobomimic_image.zip&&rm-frobomimic_image.zip&&cd..UsagePlease replace[TASK]with your desired task to train.[TASK]={lift, can, square}Convert delta actions to absolute actionspythonutils/robomimic_convert_action.py--dataset=robomimic/datasets/[TASK]/ph/low_dim.hdf5Save waypointspythonutils/robomimic_save_waypoints.py--dataset=robomimic/datasets/[TASK]/ph/low_dim.hdf5--err_threshold=0.005Replay waypoints (save 3 videos and 3D visualizations by default)mkdirvideo
pythonexample/robomimic_waypoint_replay.py--dataset=robomimic/datasets/[TASK]/ph/low_dim.hdf5\--record_video--video_pathvideo/[TASK]_waypoint.mp4--task=[TASK]\--plot_3d--auto_waypoint--err_threshold=0.005AWE + Diffusion PolicyInstall Diffusion Policycondaenvupdate-fdiffusion_policy/conda_environment.yamlIf the installation is too slow, consider usingMambaforgeinstead of the standard anaconda distribution, as recommended by theDiffusion Policyauthors. That is:mambaenvcreate-fdiffusion_policy/conda_environment.yamlTrain policypythondiffusion_policy/train.py--config-dir=config--config-name=waypoint_image_[TASK]_ph_diffusion_policy_transformer.yamlhydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'Bimanual Simulation SuiteSet up the environmentcondaenvupdate-fact/conda_env.yamlDownload dataPlease download scripted/human demo for simulated environments fromhereand save them indata/act/.If you need real robot data, please contact Lucy Shi: lucyshi (at) stanford (dot) eduUsagePlease replace[TASK]with your desired task to train.[TASK]={sim_transfer_cube_scripted, sim_insertion_scripted, sim_transfer_cube_human, sim_insertion_human}Visualize waypointspythonexample/act_waypoint.py--dataset=data/act/[TASK]--err_threshold=0.01--plot_3d--end_idx=0Save waypointspythonexample/act_waypoint.py--dataset=data/act/[TASK]--err_threshold=0.01--save_waypointsAWE + ACTTrain policypythonact/imitate_episodes.py\--task_name[TASK]\--ckpt_dirdata/outputs/act_ckpt/[TASK]_waypoint\--policy_classACT--kl_weight10--chunk_size50--hidden_dim512--batch_size8--dim_feedforward3200\--num_epochs8000--lr1e-5\--seed0--temporal_agg--use_waypointFor human datasets, set--kl_weight=80, as suggested by the ACT authors. To evaluate the policy, run the same command with--eval.CitationIf you find our code useful for your research, please cite: