Package: robot-awe
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.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: