├── .gitignore
├── LICENSE
├── README.md
├── assets
├── scaling-law-curve.png
├── wild-data.gif
└── wild-eval.gif
├── conda_environment.yaml
├── diffusion_policy
├── codecs
│ └── imagecodecs_numcodecs.py
├── common
│ ├── checkpoint_util.py
│ ├── cv2_util.py
│ ├── env_util.py
│ ├── json_logger.py
│ ├── nested_dict_util.py
│ ├── normalize_util.py
│ ├── pose_repr_util.py
│ ├── pose_trajectory_interpolator.py
│ ├── precise_sleep.py
│ ├── pymunk_override.py
│ ├── pymunk_util.py
│ ├── pytorch_util.py
│ ├── replay_buffer.py
│ ├── robomimic_config_util.py
│ ├── robomimic_util.py
│ ├── sampler.py
│ └── timestamp_accumulator.py
├── config
│ ├── legacy
│ │ ├── task
│ │ │ ├── blockpush_lowdim_seed.yaml
│ │ │ ├── blockpush_lowdim_seed_abs.yaml
│ │ │ ├── can_image.yaml
│ │ │ ├── can_image_abs.yaml
│ │ │ ├── can_lowdim.yaml
│ │ │ ├── can_lowdim_abs.yaml
│ │ │ ├── kitchen_lowdim.yaml
│ │ │ ├── kitchen_lowdim_abs.yaml
│ │ │ ├── lift_abs.yaml
│ │ │ ├── lift_image.yaml
│ │ │ ├── lift_image_abs.yaml
│ │ │ ├── lift_lowdim.yaml
│ │ │ ├── lift_lowdim_abs.yaml
│ │ │ ├── pusht_image.yaml
│ │ │ ├── pusht_lowdim.yaml
│ │ │ ├── real_pusht_image.yaml
│ │ │ ├── square_abs.yaml
│ │ │ ├── square_image.yaml
│ │ │ ├── square_image_abs.yaml
│ │ │ ├── square_lowdim.yaml
│ │ │ ├── square_lowdim_abs.yaml
│ │ │ ├── tool_hang_abs.yaml
│ │ │ ├── tool_hang_image.yaml
│ │ │ ├── tool_hang_image_abs.yaml
│ │ │ ├── tool_hang_lowdim.yaml
│ │ │ ├── tool_hang_lowdim_abs.yaml
│ │ │ ├── transport_image.yaml
│ │ │ ├── transport_image_abs.yaml
│ │ │ ├── transport_lowdim.yaml
│ │ │ ├── transport_lowdim_abs.yaml
│ │ │ ├── umi.yaml
│ │ │ ├── umi_image.yaml
│ │ │ └── umi_teleop_image.yaml
│ │ ├── train_bet_lowdim_workspace.yaml
│ │ ├── train_diffusion_transformer_hybrid_workspace.yaml
│ │ ├── train_diffusion_transformer_lowdim_kitchen_workspace.yaml
│ │ ├── train_diffusion_transformer_lowdim_pusht_workspace.yaml
│ │ ├── train_diffusion_transformer_lowdim_workspace.yaml
│ │ ├── train_diffusion_transformer_real_hybrid_workspace.yaml
│ │ ├── train_diffusion_transformer_umi_teleop_hybrid_workspace.yaml
│ │ ├── train_diffusion_unet_ddim_hybrid_workspace.yaml
│ │ ├── train_diffusion_unet_ddim_lowdim_workspace.yaml
│ │ ├── train_diffusion_unet_hybrid_workspace.yaml
│ │ ├── train_diffusion_unet_image_pretrained_workspace.yaml
│ │ ├── train_diffusion_unet_image_workspace.yaml
│ │ ├── train_diffusion_unet_lowdim_workspace.yaml
│ │ ├── train_diffusion_unet_real_hybrid_workspace.yaml
│ │ ├── train_diffusion_unet_real_image_workspace.yaml
│ │ ├── train_diffusion_unet_real_pretrained_workspace.yaml
│ │ ├── train_diffusion_unet_timm_workspace.yaml
│ │ ├── train_diffusion_unet_umi_hybrid_workspace.yaml
│ │ ├── train_diffusion_unet_umi_teleop_hybrid_workspace.yaml
│ │ ├── train_diffusion_unet_umi_timm_workspace.yaml
│ │ ├── train_diffusion_unet_umi_workspace.yaml
│ │ ├── train_diffusion_unet_video_workspace.yaml
│ │ ├── train_ibc_dfo_hybrid_workspace.yaml
│ │ ├── train_ibc_dfo_lowdim_workspace.yaml
│ │ ├── train_ibc_dfo_real_hybrid_workspace.yaml
│ │ ├── train_robomimic_image_workspace.yaml
│ │ ├── train_robomimic_lowdim_workspace.yaml
│ │ └── train_robomimic_real_image_workspace.yaml
│ ├── task
│ │ ├── lift.yaml
│ │ ├── square.yaml
│ │ ├── tool_hang.yaml
│ │ ├── tool_hang_abs.yaml
│ │ ├── umi.yaml
│ │ ├── umi_bimanual.yaml
│ │ ├── umi_image.yaml
│ │ └── umi_teleop_image.yaml
│ ├── train_diffusion_transformer_umi_bimanual_workspace.yaml
│ ├── train_diffusion_transformer_umi_workspace.yaml
│ ├── train_diffusion_unet_image_workspace.yaml
│ ├── train_diffusion_unet_timm_robomimic_workspace.yaml
│ ├── train_diffusion_unet_timm_umi_workspace.yaml
│ └── train_diffusion_unet_umi_bimanual_workspace.yaml
├── dataset
│ ├── base_dataset.py
│ ├── blockpush_lowdim_dataset.py
│ ├── kitchen_lowdim_dataset.py
│ ├── kitchen_mjl_lowdim_dataset.py
│ ├── pusht_dataset.py
│ ├── pusht_image_dataset.py
│ ├── real_pusht_image_dataset.py
│ ├── robomimic_replay_dataset.py
│ ├── robomimic_replay_image_dataset.py
│ ├── robomimic_replay_lowdim_dataset.py
│ ├── umi_dataset.py
│ ├── umi_image_dataset.py
│ └── umi_teleop_image_dataset.py
├── env
│ ├── block_pushing
│ │ ├── assets
│ │ │ ├── block.urdf
│ │ │ ├── block2.urdf
│ │ │ ├── blocks
│ │ │ │ ├── blue_cube.urdf
│ │ │ │ ├── cube.obj
│ │ │ │ ├── green_star.urdf
│ │ │ │ ├── moon.obj
│ │ │ │ ├── pentagon.obj
│ │ │ │ ├── red_moon.urdf
│ │ │ │ ├── star.obj
│ │ │ │ └── yellow_pentagon.urdf
│ │ │ ├── insert.urdf
│ │ │ ├── plane.obj
│ │ │ ├── suction
│ │ │ │ ├── base.obj
│ │ │ │ ├── cylinder.urdf
│ │ │ │ ├── cylinder_real.urdf
│ │ │ │ ├── head.obj
│ │ │ │ ├── mid.obj
│ │ │ │ ├── suction-base.urdf
│ │ │ │ ├── suction-head-long.urdf
│ │ │ │ ├── suction-head.urdf
│ │ │ │ └── tip.obj
│ │ │ ├── workspace.urdf
│ │ │ ├── workspace_real.urdf
│ │ │ ├── zone.obj
│ │ │ ├── zone.urdf
│ │ │ └── zone2.urdf
│ │ ├── block_pushing.py
│ │ ├── block_pushing_discontinuous.py
│ │ ├── block_pushing_multimodal.py
│ │ ├── oracles
│ │ │ ├── discontinuous_push_oracle.py
│ │ │ ├── multimodal_push_oracle.py
│ │ │ ├── oriented_push_oracle.py
│ │ │ ├── pushing_info.py
│ │ │ └── reach_oracle.py
│ │ └── utils
│ │ │ ├── pose3d.py
│ │ │ ├── utils_pybullet.py
│ │ │ └── xarm_sim_robot.py
│ ├── kitchen
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── kitchen_lowdim_wrapper.py
│ │ ├── kitchen_util.py
│ │ ├── relay_policy_learning
│ │ │ ├── adept_envs
│ │ │ │ ├── .pylintrc
│ │ │ │ ├── .style.yapf
│ │ │ │ └── adept_envs
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ ├── base_robot.py
│ │ │ │ │ ├── franka
│ │ │ │ │ ├── __init__.py
│ │ │ │ │ ├── assets
│ │ │ │ │ │ └── franka_kitchen_jntpos_act_ab.xml
│ │ │ │ │ ├── kitchen_multitask_v0.py
│ │ │ │ │ └── robot
│ │ │ │ │ │ ├── franka_config.xml
│ │ │ │ │ │ └── franka_robot.py
│ │ │ │ │ ├── mujoco_env.py
│ │ │ │ │ ├── robot_env.py
│ │ │ │ │ ├── simulation
│ │ │ │ │ ├── module.py
│ │ │ │ │ ├── renderer.py
│ │ │ │ │ └── sim_robot.py
│ │ │ │ │ └── utils
│ │ │ │ │ ├── config.py
│ │ │ │ │ ├── configurable.py
│ │ │ │ │ ├── constants.py
│ │ │ │ │ ├── parse_demos.py
│ │ │ │ │ └── quatmath.py
│ │ │ ├── adept_models
│ │ │ │ ├── .gitignore
│ │ │ │ ├── CONTRIBUTING.public.md
│ │ │ │ ├── LICENSE
│ │ │ │ ├── README.public.md
│ │ │ │ ├── __init__.py
│ │ │ │ ├── kitchen
│ │ │ │ │ ├── assets
│ │ │ │ │ │ ├── backwall_asset.xml
│ │ │ │ │ │ ├── backwall_chain.xml
│ │ │ │ │ │ ├── counters_asset.xml
│ │ │ │ │ │ ├── counters_chain.xml
│ │ │ │ │ │ ├── hingecabinet_asset.xml
│ │ │ │ │ │ ├── hingecabinet_chain.xml
│ │ │ │ │ │ ├── kettle_asset.xml
│ │ │ │ │ │ ├── kettle_chain.xml
│ │ │ │ │ │ ├── microwave_asset.xml
│ │ │ │ │ │ ├── microwave_chain.xml
│ │ │ │ │ │ ├── oven_asset.xml
│ │ │ │ │ │ ├── oven_chain.xml
│ │ │ │ │ │ ├── slidecabinet_asset.xml
│ │ │ │ │ │ └── slidecabinet_chain.xml
│ │ │ │ │ ├── counters.xml
│ │ │ │ │ ├── hingecabinet.xml
│ │ │ │ │ ├── kettle.xml
│ │ │ │ │ ├── kitchen.xml
│ │ │ │ │ ├── meshes
│ │ │ │ │ │ ├── burnerplate.stl
│ │ │ │ │ │ ├── burnerplate_mesh.stl
│ │ │ │ │ │ ├── cabinetbase.stl
│ │ │ │ │ │ ├── cabinetdrawer.stl
│ │ │ │ │ │ ├── cabinethandle.stl
│ │ │ │ │ │ ├── countertop.stl
│ │ │ │ │ │ ├── faucet.stl
│ │ │ │ │ │ ├── handle2.stl
│ │ │ │ │ │ ├── hingecabinet.stl
│ │ │ │ │ │ ├── hingedoor.stl
│ │ │ │ │ │ ├── hingehandle.stl
│ │ │ │ │ │ ├── hood.stl
│ │ │ │ │ │ ├── kettle.stl
│ │ │ │ │ │ ├── kettlehandle.stl
│ │ │ │ │ │ ├── knob.stl
│ │ │ │ │ │ ├── lightswitch.stl
│ │ │ │ │ │ ├── lightswitchbase.stl
│ │ │ │ │ │ ├── micro.stl
│ │ │ │ │ │ ├── microbutton.stl
│ │ │ │ │ │ ├── microdoor.stl
│ │ │ │ │ │ ├── microefeet.stl
│ │ │ │ │ │ ├── microfeet.stl
│ │ │ │ │ │ ├── microhandle.stl
│ │ │ │ │ │ ├── microwindow.stl
│ │ │ │ │ │ ├── oven.stl
│ │ │ │ │ │ ├── ovenhandle.stl
│ │ │ │ │ │ ├── oventop.stl
│ │ │ │ │ │ ├── ovenwindow.stl
│ │ │ │ │ │ ├── slidecabinet.stl
│ │ │ │ │ │ ├── slidedoor.stl
│ │ │ │ │ │ ├── stoverim.stl
│ │ │ │ │ │ ├── tile.stl
│ │ │ │ │ │ └── wall.stl
│ │ │ │ │ ├── microwave.xml
│ │ │ │ │ ├── oven.xml
│ │ │ │ │ ├── slidecabinet.xml
│ │ │ │ │ └── textures
│ │ │ │ │ │ ├── marble1.png
│ │ │ │ │ │ ├── metal1.png
│ │ │ │ │ │ ├── tile1.png
│ │ │ │ │ │ └── wood1.png
│ │ │ │ └── scenes
│ │ │ │ │ ├── basic_scene.xml
│ │ │ │ │ └── textures
│ │ │ │ │ ├── white_marble_tile.png
│ │ │ │ │ └── white_marble_tile2.png
│ │ │ └── third_party
│ │ │ │ └── franka
│ │ │ │ ├── LICENSE
│ │ │ │ ├── README.md
│ │ │ │ ├── assets
│ │ │ │ ├── actuator0.xml
│ │ │ │ ├── actuator1.xml
│ │ │ │ ├── assets.xml
│ │ │ │ ├── basic_scene.xml
│ │ │ │ ├── chain0.xml
│ │ │ │ ├── chain0_overlay.xml
│ │ │ │ ├── chain1.xml
│ │ │ │ └── teleop_actuator.xml
│ │ │ │ ├── bi-franka_panda.xml
│ │ │ │ ├── franka_panda.png
│ │ │ │ ├── franka_panda.xml
│ │ │ │ ├── franka_panda_teleop.xml
│ │ │ │ └── meshes
│ │ │ │ ├── collision
│ │ │ │ ├── finger.stl
│ │ │ │ ├── hand.stl
│ │ │ │ ├── link0.stl
│ │ │ │ ├── link1.stl
│ │ │ │ ├── link2.stl
│ │ │ │ ├── link3.stl
│ │ │ │ ├── link4.stl
│ │ │ │ ├── link5.stl
│ │ │ │ ├── link6.stl
│ │ │ │ └── link7.stl
│ │ │ │ └── visual
│ │ │ │ ├── finger.stl
│ │ │ │ ├── hand.stl
│ │ │ │ ├── link0.stl
│ │ │ │ ├── link1.stl
│ │ │ │ ├── link2.stl
│ │ │ │ ├── link3.stl
│ │ │ │ ├── link4.stl
│ │ │ │ ├── link5.stl
│ │ │ │ ├── link6.stl
│ │ │ │ └── link7.stl
│ │ └── v0.py
│ ├── pusht
│ │ ├── __init__.py
│ │ ├── pusht_env.py
│ │ ├── pusht_image_env.py
│ │ ├── pusht_keypoints_env.py
│ │ ├── pymunk_keypoint_manager.py
│ │ └── pymunk_override.py
│ └── robomimic
│ │ ├── robomimic_image_wrapper.py
│ │ └── robomimic_lowdim_wrapper.py
├── env_runner
│ ├── base_image_runner.py
│ ├── base_lowdim_runner.py
│ ├── blockpush_lowdim_runner.py
│ ├── kitchen_lowdim_runner.py
│ ├── pusht_image_runner.py
│ ├── pusht_keypoints_runner.py
│ ├── real_pusht_image_runner.py
│ ├── robomimic_image_runner.py
│ ├── robomimic_lowdim_runner.py
│ └── robomimic_runner.py
├── gym_util
│ ├── async_vector_env.py
│ ├── multistep_wrapper.py
│ ├── sync_vector_env.py
│ ├── video_recording_wrapper.py
│ └── video_wrapper.py
├── model
│ ├── bet
│ │ ├── action_ae
│ │ │ ├── __init__.py
│ │ │ └── discretizers
│ │ │ │ └── k_means.py
│ │ ├── latent_generators
│ │ │ ├── latent_generator.py
│ │ │ ├── mingpt.py
│ │ │ └── transformer.py
│ │ ├── libraries
│ │ │ ├── loss_fn.py
│ │ │ └── mingpt
│ │ │ │ ├── LICENSE
│ │ │ │ ├── __init__.py
│ │ │ │ ├── model.py
│ │ │ │ ├── trainer.py
│ │ │ │ └── utils.py
│ │ └── utils.py
│ ├── common
│ │ ├── dict_of_tensor_mixin.py
│ │ ├── lr_decay.py
│ │ ├── lr_scheduler.py
│ │ ├── module_attr_mixin.py
│ │ ├── normalizer.py
│ │ ├── rotation_transformer.py
│ │ ├── shape_util.py
│ │ └── tensor_util.py
│ ├── diffusion
│ │ ├── conditional_unet1d.py
│ │ ├── conv1d_components.py
│ │ ├── ema_model.py
│ │ ├── mask_generator.py
│ │ ├── positional_embedding.py
│ │ ├── transformer_for_action_diffusion.py
│ │ └── transformer_for_diffusion.py
│ └── vision
│ │ ├── choice_randomizer.py
│ │ ├── crop_randomizer.py
│ │ ├── model_getter.py
│ │ ├── multi_image_obs_encoder.py
│ │ ├── timm_obs_encoder.py
│ │ └── transformer_obs_encoder.py
├── policy
│ ├── base_image_policy.py
│ ├── base_lowdim_policy.py
│ ├── bet_lowdim_policy.py
│ ├── diffusion_transformer_hybrid_image_policy.py
│ ├── diffusion_transformer_lowdim_policy.py
│ ├── diffusion_transformer_timm_policy.py
│ ├── diffusion_unet_hybrid_image_policy.py
│ ├── diffusion_unet_image_policy.py
│ ├── diffusion_unet_lowdim_policy.py
│ ├── diffusion_unet_timm_policy.py
│ ├── diffusion_unet_video_policy.py
│ ├── ibc_dfo_hybrid_image_policy.py
│ ├── ibc_dfo_lowdim_policy.py
│ ├── robomimic_image_policy.py
│ └── robomimic_lowdim_policy.py
├── real_world
│ ├── keystroke_counter.py
│ ├── multi_camera_visualizer.py
│ ├── multi_realsense.py
│ ├── real_data_conversion.py
│ ├── real_env.py
│ ├── real_inference_util.py
│ ├── realsense_config
│ │ ├── 415_high_accuracy_mode.json
│ │ └── 435_high_accuracy_mode.json
│ ├── rtde_interpolation_controller.py
│ ├── single_realsense.py
│ ├── spacemouse.py
│ ├── spacemouse_shared_memory.py
│ ├── umi_env.py
│ └── video_recorder.py
├── scripts
│ ├── bet_blockpush_conversion.py
│ ├── blockpush_abs_conversion.py
│ ├── episode_lengths.py
│ ├── generate_bet_blockpush.py
│ ├── real_dataset_conversion.py
│ ├── real_pusht_metrics.py
│ ├── real_pusht_successrate.py
│ ├── robomimic_dataset_action_comparison.py
│ └── robomimic_dataset_conversion.py
├── shared_memory
│ ├── shared_memory_queue.py
│ ├── shared_memory_ring_buffer.py
│ ├── shared_memory_util.py
│ └── shared_ndarray.py
└── workspace
│ ├── base_workspace.py
│ ├── train_bet_lowdim_workspace.py
│ ├── train_diffusion_transformer_hybrid_workspace.py
│ ├── train_diffusion_transformer_lowdim_workspace.py
│ ├── train_diffusion_transformer_timm_workspace.py
│ ├── train_diffusion_unet_hybrid_workspace.py
│ ├── train_diffusion_unet_image_workspace.py
│ ├── train_diffusion_unet_lowdim_workspace.py
│ ├── train_diffusion_unet_video_workspace.py
│ ├── train_ibc_dfo_hybrid_workspace.py
│ ├── train_ibc_dfo_lowdim_workspace.py
│ ├── train_robomimic_image_workspace.py
│ └── train_robomimic_lowdim_workspace.py
├── eval.sh
├── eval_real.py
├── example
├── calibration
│ ├── aruco_config.yaml
│ └── gopro_intrinsics_2_7k.json
└── eval_robots_config.yaml
├── franka_instruction.md
├── run_generate_dataset.sh
├── run_slam.sh
├── run_slam_pipeline.py
├── scripts
├── calibrate_gripper_latency.py
├── calibrate_gripper_range.py
├── calibrate_robot_latency.py
├── calibrate_robot_world_hand_eye.py
├── calibrate_slam_tag.py
├── calibrate_timecode_offset.py
├── calibrate_uvc_camera_latency.py
├── check_gopro_orientation.py
├── control_robot_spacemouse_traj.py
├── count_dataset_videos.py
├── detect_aruco.py
├── detect_bad_rotation.py
├── draw_aurco_detection.py
├── draw_charuco_detection.py
├── eval_real_robot.py
├── filter_videos.py
├── gen_aruco_tag_pdf.py
├── gen_charuco_board.py
├── gen_image_mask.py
├── generate_video_thumbnail.py
├── generate_video_thumbnail_demos.py
├── record_robot_world_hand_eye.py
├── replay_real_robot.py
├── reset_usb.py
├── slam_pose_replay.py
├── uvc_camera.py
└── uvc_camera_test.py
├── scripts_real
├── control_franka.py
├── control_robot_spacemouse.py
├── control_robots.py
├── control_wsg_spacemouse.py
├── demo_real_bimanual_robots.py
├── demo_real_robot.py
├── demo_real_umi.py
├── eval_real_bimanual_umi.py
├── eval_real_robot.py
├── eval_real_umi.py
├── eval_replay_real_robot.py
├── latency_test_wsg.py
├── launch_franka_interface_server.py
└── replay_real_bimanual_umi.py
├── scripts_slam_pipeline
├── 00_process_videos.py
├── 01_extract_gopro_imu.py
├── 02_create_map.py
├── 03_batch_slam.py
├── 04_detect_aruco.py
├── 05_run_calibrations.py
├── 06_generate_dataset_plan.py
└── 07_generate_replay_buffer.py
├── tests
├── test_multi_uvc_camera.py
├── test_pose_util.py
└── test_uvc_camera.py
├── train.py
├── train_scripts
├── arrange_mouse.sh
├── fold_towel.sh
├── pour_water.sh
└── unplug_charger.sh
├── umi
├── asset
│ └── mask.json
├── common
│ ├── cv_util.py
│ ├── exiftool_util.py
│ ├── interpolation_util.py
│ ├── k3d_util.py
│ ├── latency_util.py
│ ├── mocap_util.py
│ ├── nested_dict_util.py
│ ├── orb_slam_util.py
│ ├── pose_trajectory_interpolator.py
│ ├── pose_util.py
│ ├── precise_sleep.py
│ ├── timecode_util.py
│ ├── timestamp_accumulator.py
│ └── usb_util.py
├── pipeline
│ └── aruco_detection.py
├── real_world
│ ├── bimanual_umi_env.py
│ ├── cmd_measure.lua
│ ├── franka_interpolation_controller.py
│ ├── keystroke_counter.py
│ ├── multi_camera_visualizer.py
│ ├── multi_uvc_camera.py
│ ├── real_env.py
│ ├── real_inference_util.py
│ ├── rtde_interpolation_controller.py
│ ├── spacemouse_shared_memory.py
│ ├── umi_env.py
│ ├── uvc_camera.py
│ ├── video_recorder.py
│ ├── wsg_binary_driver.py
│ └── wsg_controller.py
├── shared_memory
│ ├── shared_memory_queue.py
│ ├── shared_memory_ring_buffer.py
│ ├── shared_memory_util.py
│ └── shared_ndarray.py
└── traj_eval
│ ├── align_trajectory.py
│ ├── align_utils.py
│ ├── compute_trajectory_errors.py
│ ├── results_writer.py
│ ├── trajectory_utils.py
│ └── transformations.py
└── visualize_dataset.py
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2023 Columbia Artificial Intelligence and Robotics Lab
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
--------------------------------------------------------------------------------
/assets/scaling-law-curve.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/assets/scaling-law-curve.png
--------------------------------------------------------------------------------
/assets/wild-data.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/assets/wild-data.gif
--------------------------------------------------------------------------------
/assets/wild-eval.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/assets/wild-eval.gif
--------------------------------------------------------------------------------
/conda_environment.yaml:
--------------------------------------------------------------------------------
1 | name: umi
2 | channels:
3 | - pytorch
4 | - nvidia
5 | - conda-forge
6 | dependencies:
7 | - python=3.9
8 | - pip=22.1
9 | - pytorch=2.1.0
10 | - pytorch-cuda=12.1
11 | - torchvision=0.16.0
12 | - numpy=1.24
13 | - scipy=1.11
14 | - numba=0.57
15 | - pandas=2.1
16 | - py-opencv=4.7
17 | - ipykernel=6.23
18 | - matplotlib=3.7
19 | - av=10.0
20 | - k3d=2.15
21 | - pyyaml=6.0
22 | - tqdm=4.65
23 | - zarr=2.16
24 | - numcodecs=0.11
25 | - hydra-core=1.2.0
26 | - wandb=0.15.8
27 | - dill=0.3.7
28 | - einops=0.6.1
29 | - diffusers=0.18.2
30 | - scikit-video=1.1
31 | - gym=0.21.0
32 | # trick to avoid cpu affinity issue described in https://github.com/pytorch/pytorch/issues/99625
33 | - llvm-openmp=14
34 | - python-lmdb=1.4
35 | - timm=0.9.10
36 | - scikit-fda=0.9
37 | - cython=0.29
38 | - accelerate=0.24
39 | - boto3=1.33
40 | - exiftool=12
41 | - aiohttp=3.9
42 | - aiofiles=23.2
43 | - pip:
44 | - peft==0.11.1
45 | - ray[default]==2.8.0
46 | - threadpoolctl==3.2
47 | - v4l2py==2.2
48 | - ur-rtde==1.5.6
49 | # requires apt install libspnav-dev spacenavd; systemctl start spacenavd
50 | - spnav @ https://github.com/cheng-chi/spnav/archive/c1c938ebe3cc542db4685e0d13850ff1abfdb943.tar.gz
51 | - atomics==1.0.2
52 | - pynput==1.7.6
53 | - fastcrc==0.2
54 | - imagecodecs==2023.9.18
55 | - robomimic==0.2.0
56 | # requires mujoco py dependencies libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf
57 | - free-mujoco-py==2.1.6
58 | - robosuite @ https://github.com/cheng-chi/robosuite/archive/3f2c3116f1c0a1916e94bfa28da4d055927d1ab3.tar.gz
59 | # implicitly depend on conda package exiftool
60 | - PyExifTool==0.5
61 |
--------------------------------------------------------------------------------
/diffusion_policy/common/checkpoint_util.py:
--------------------------------------------------------------------------------
1 | from typing import Optional, Dict
2 | import os
3 |
4 | class TopKCheckpointManager:
5 | def __init__(self,
6 | save_dir,
7 | monitor_key: str,
8 | mode='min',
9 | k=1,
10 | format_str='epoch={epoch:03d}-train_loss={train_loss:.3f}.ckpt'
11 | ):
12 | assert mode in ['max', 'min']
13 | assert k >= 0
14 |
15 | self.save_dir = save_dir
16 | self.monitor_key = monitor_key
17 | self.mode = mode
18 | self.k = k
19 | self.format_str = format_str
20 | self.path_value_map = dict()
21 |
22 | def get_ckpt_path(self, data: Dict[str, float]) -> Optional[str]:
23 | if self.k == 0:
24 | return None
25 |
26 | value = data[self.monitor_key]
27 | ckpt_path = os.path.join(
28 | self.save_dir, self.format_str.format(**data))
29 |
30 | if len(self.path_value_map) < self.k:
31 | # under-capacity
32 | self.path_value_map[ckpt_path] = value
33 | return ckpt_path
34 |
35 | # at capacity
36 | sorted_map = sorted(self.path_value_map.items(), key=lambda x: x[1])
37 | min_path, min_value = sorted_map[0]
38 | max_path, max_value = sorted_map[-1]
39 |
40 | delete_path = None
41 | if self.mode == 'max':
42 | if value > min_value:
43 | delete_path = min_path
44 | else:
45 | if value < max_value:
46 | delete_path = max_path
47 |
48 | if delete_path is None:
49 | return None
50 | else:
51 | del self.path_value_map[delete_path]
52 | self.path_value_map[ckpt_path] = value
53 |
54 | if not os.path.exists(self.save_dir):
55 | os.mkdir(self.save_dir)
56 |
57 | if os.path.exists(delete_path):
58 | os.remove(delete_path)
59 | return ckpt_path
60 |
--------------------------------------------------------------------------------
/diffusion_policy/common/env_util.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import numpy as np
3 |
4 |
5 | def render_env_video(env, states, actions=None):
6 | observations = states
7 | imgs = list()
8 | for i in range(len(observations)):
9 | state = observations[i]
10 | env.set_state(state)
11 | if i == 0:
12 | env.set_state(state)
13 | img = env.render()
14 | # draw action
15 | if actions is not None:
16 | action = actions[i]
17 | coord = (action / 512 * 96).astype(np.int32)
18 | cv2.drawMarker(img, coord,
19 | color=(255,0,0), markerType=cv2.MARKER_CROSS,
20 | markerSize=8, thickness=1)
21 | imgs.append(img)
22 | imgs = np.array(imgs)
23 | return imgs
24 |
--------------------------------------------------------------------------------
/diffusion_policy/common/nested_dict_util.py:
--------------------------------------------------------------------------------
1 | import functools
2 |
3 | def nested_dict_map(f, x):
4 | """
5 | Map f over all leaf of nested dict x
6 | """
7 |
8 | if not isinstance(x, dict):
9 | return f(x)
10 | y = dict()
11 | for key, value in x.items():
12 | y[key] = nested_dict_map(f, value)
13 | return y
14 |
15 | def nested_dict_reduce(f, x):
16 | """
17 | Map f over all values of nested dict x, and reduce to a single value
18 | """
19 | if not isinstance(x, dict):
20 | return x
21 |
22 | reduced_values = list()
23 | for value in x.values():
24 | reduced_values.append(nested_dict_reduce(f, value))
25 | y = functools.reduce(f, reduced_values)
26 | return y
27 |
28 |
29 | def nested_dict_check(f, x):
30 | bool_dict = nested_dict_map(f, x)
31 | result = nested_dict_reduce(lambda x, y: x and y, bool_dict)
32 | return result
33 |
--------------------------------------------------------------------------------
/diffusion_policy/common/precise_sleep.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | def precise_sleep(dt: float, slack_time: float=0.001, time_func=time.monotonic):
4 | """
5 | Use hybrid of time.sleep and spinning to minimize jitter.
6 | Sleep dt - slack_time seconds first, then spin for the rest.
7 | """
8 | t_start = time_func()
9 | if dt > slack_time:
10 | time.sleep(dt - slack_time)
11 | t_end = t_start + dt
12 | while time_func() < t_end:
13 | pass
14 | return
15 |
16 | def precise_wait(t_end: float, slack_time: float=0.001, time_func=time.monotonic):
17 | t_start = time_func()
18 | t_wait = t_end - t_start
19 | if t_wait > 0:
20 | t_sleep = t_wait - slack_time
21 | if t_sleep > 0:
22 | time.sleep(t_sleep)
23 | while time_func() < t_end:
24 | pass
25 | return
26 |
--------------------------------------------------------------------------------
/diffusion_policy/common/pymunk_util.py:
--------------------------------------------------------------------------------
1 | import pygame
2 | import pymunk
3 | import pymunk.pygame_util
4 | import numpy as np
5 |
6 | COLLTYPE_DEFAULT = 0
7 | COLLTYPE_MOUSE = 1
8 | COLLTYPE_BALL = 2
9 |
10 | def get_body_type(static=False):
11 | body_type = pymunk.Body.DYNAMIC
12 | if static:
13 | body_type = pymunk.Body.STATIC
14 | return body_type
15 |
16 |
17 | def create_rectangle(space,
18 | pos_x,pos_y,width,height,
19 | density=3,static=False):
20 | body = pymunk.Body(body_type=get_body_type(static))
21 | body.position = (pos_x,pos_y)
22 | shape = pymunk.Poly.create_box(body,(width,height))
23 | shape.density = density
24 | space.add(body,shape)
25 | return body, shape
26 |
27 |
28 | def create_rectangle_bb(space,
29 | left, bottom, right, top,
30 | **kwargs):
31 | pos_x = (left + right) / 2
32 | pos_y = (top + bottom) / 2
33 | height = top - bottom
34 | width = right - left
35 | return create_rectangle(space, pos_x, pos_y, width, height, **kwargs)
36 |
37 | def create_circle(space, pos_x, pos_y, radius, density=3, static=False):
38 | body = pymunk.Body(body_type=get_body_type(static))
39 | body.position = (pos_x, pos_y)
40 | shape = pymunk.Circle(body, radius=radius)
41 | shape.density = density
42 | shape.collision_type = COLLTYPE_BALL
43 | space.add(body, shape)
44 | return body, shape
45 |
46 | def get_body_state(body):
47 | state = np.zeros(6, dtype=np.float32)
48 | state[:2] = body.position
49 | state[2] = body.angle
50 | state[3:5] = body.velocity
51 | state[5] = body.angular_velocity
52 | return state
53 |
--------------------------------------------------------------------------------
/diffusion_policy/common/robomimic_config_util.py:
--------------------------------------------------------------------------------
1 | from omegaconf import OmegaConf
2 | from robomimic.config import config_factory
3 | import robomimic.scripts.generate_paper_configs as gpc
4 | from robomimic.scripts.generate_paper_configs import (
5 | modify_config_for_default_image_exp,
6 | modify_config_for_default_low_dim_exp,
7 | modify_config_for_dataset,
8 | )
9 |
10 | def get_robomimic_config(
11 | algo_name='bc_rnn',
12 | hdf5_type='low_dim',
13 | task_name='square',
14 | dataset_type='ph'
15 | ):
16 | base_dataset_dir = '/tmp/null'
17 | filter_key = None
18 |
19 | # decide whether to use low-dim or image training defaults
20 | modifier_for_obs = modify_config_for_default_image_exp
21 | if hdf5_type in ["low_dim", "low_dim_sparse", "low_dim_dense"]:
22 | modifier_for_obs = modify_config_for_default_low_dim_exp
23 |
24 | algo_config_name = "bc" if algo_name == "bc_rnn" else algo_name
25 | config = config_factory(algo_name=algo_config_name)
26 | # turn into default config for observation modalities (e.g.: low-dim or rgb)
27 | config = modifier_for_obs(config)
28 | # add in config based on the dataset
29 | config = modify_config_for_dataset(
30 | config=config,
31 | task_name=task_name,
32 | dataset_type=dataset_type,
33 | hdf5_type=hdf5_type,
34 | base_dataset_dir=base_dataset_dir,
35 | filter_key=filter_key,
36 | )
37 | # add in algo hypers based on dataset
38 | algo_config_modifier = getattr(gpc, f'modify_{algo_name}_config_for_dataset')
39 | config = algo_config_modifier(
40 | config=config,
41 | task_name=task_name,
42 | dataset_type=dataset_type,
43 | hdf5_type=hdf5_type,
44 | )
45 | return config
46 |
47 |
48 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/blockpush_lowdim_seed.yaml:
--------------------------------------------------------------------------------
1 | name: blockpush_lowdim_seed
2 |
3 | obs_dim: 16
4 | action_dim: 2
5 | keypoint_dim: 2
6 | obs_eef_target: True
7 |
8 | env_runner:
9 | _target_: diffusion_policy.env_runner.blockpush_lowdim_runner.BlockPushLowdimRunner
10 | n_train: 6
11 | n_train_vis: 2
12 | train_start_seed: 0
13 | n_test: 50
14 | n_test_vis: 4
15 | test_start_seed: 100000
16 | max_steps: 350
17 | n_obs_steps: ${n_obs_steps}
18 | n_action_steps: ${n_action_steps}
19 | fps: 5
20 | past_action: ${past_action_visible}
21 | abs_action: False
22 | obs_eef_target: ${task.obs_eef_target}
23 | n_envs: null
24 |
25 | dataset:
26 | _target_: diffusion_policy.dataset.blockpush_lowdim_dataset.BlockPushLowdimDataset
27 | zarr_path: data/block_pushing/multimodal_push_seed.zarr
28 | horizon: ${horizon}
29 | pad_before: ${eval:'${n_obs_steps}-1'}
30 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
31 | obs_eef_target: ${task.obs_eef_target}
32 | use_manual_normalizer: False
33 | seed: 42
34 | val_ratio: 0.02
35 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/blockpush_lowdim_seed_abs.yaml:
--------------------------------------------------------------------------------
1 | name: blockpush_lowdim_seed_abs
2 |
3 | obs_dim: 16
4 | action_dim: 2
5 | keypoint_dim: 2
6 | obs_eef_target: True
7 |
8 | env_runner:
9 | _target_: diffusion_policy.env_runner.blockpush_lowdim_runner.BlockPushLowdimRunner
10 | n_train: 6
11 | n_train_vis: 2
12 | train_start_seed: 0
13 | n_test: 50
14 | n_test_vis: 4
15 | test_start_seed: 100000
16 | max_steps: 350
17 | n_obs_steps: ${n_obs_steps}
18 | n_action_steps: ${n_action_steps}
19 | fps: 5
20 | past_action: ${past_action_visible}
21 | abs_action: True
22 | obs_eef_target: ${task.obs_eef_target}
23 | n_envs: null
24 |
25 | dataset:
26 | _target_: diffusion_policy.dataset.blockpush_lowdim_dataset.BlockPushLowdimDataset
27 | zarr_path: data/block_pushing/multimodal_push_seed_abs.zarr
28 | horizon: ${horizon}
29 | pad_before: ${eval:'${n_obs_steps}-1'}
30 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
31 | obs_eef_target: ${task.obs_eef_target}
32 | use_manual_normalizer: False
33 | seed: 42
34 | val_ratio: 0.02
35 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/can_image.yaml:
--------------------------------------------------------------------------------
1 | name: can_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | agentview_image:
7 | shape: [3, 84, 84]
8 | type: rgb
9 | robot0_eye_in_hand_image:
10 | shape: [3, 84, 84]
11 | type: rgb
12 | robot0_eef_pos:
13 | shape: [3]
14 | # type default: low_dim
15 | robot0_eef_quat:
16 | shape: [4]
17 | robot0_gripper_qpos:
18 | shape: [2]
19 | action:
20 | shape: [7]
21 |
22 | task_name: &task_name can
23 | dataset_type: &dataset_type ph
24 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5
25 | abs_action: &abs_action False
26 |
27 | env_runner:
28 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
29 | dataset_path: *dataset_path
30 | shape_meta: *shape_meta
31 | # costs 1GB per env
32 | n_train: 6
33 | n_train_vis: 2
34 | train_start_idx: 0
35 | n_test: 50
36 | n_test_vis: 4
37 | test_start_seed: 100000
38 | # use python's eval function as resolver, single-quoted string as argument
39 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
40 | n_obs_steps: ${n_obs_steps}
41 | n_action_steps: ${n_action_steps}
42 | render_obs_key: 'agentview_image'
43 | fps: 10
44 | crf: 22
45 | past_action: ${past_action_visible}
46 | abs_action: *abs_action
47 | tqdm_interval_sec: 1.0
48 | n_envs: 28
49 | # evaluation at this config requires a 16 core 64GB instance.
50 |
51 | dataset:
52 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
53 | shape_meta: *shape_meta
54 | dataset_path: *dataset_path
55 | horizon: ${horizon}
56 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
57 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
58 | n_obs_steps: ${dataset_obs_steps}
59 | abs_action: *abs_action
60 | rotation_rep: 'rotation_6d'
61 | use_legacy_normalizer: False
62 | use_cache: True
63 | seed: 42
64 | val_ratio: 0.02
65 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/can_image_abs.yaml:
--------------------------------------------------------------------------------
1 | name: can_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | agentview_image:
7 | shape: [3, 84, 84]
8 | type: rgb
9 | robot0_eye_in_hand_image:
10 | shape: [3, 84, 84]
11 | type: rgb
12 | robot0_eef_pos:
13 | shape: [3]
14 | # type default: low_dim
15 | robot0_eef_quat:
16 | shape: [4]
17 | robot0_gripper_qpos:
18 | shape: [2]
19 | action:
20 | shape: [10]
21 |
22 | task_name: &task_name can
23 | dataset_type: &dataset_type ph
24 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
25 | abs_action: &abs_action True
26 |
27 | env_runner:
28 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
29 | dataset_path: *dataset_path
30 | shape_meta: *shape_meta
31 | # costs 1GB per env
32 | n_train: 6
33 | n_train_vis: 2
34 | train_start_idx: 0
35 | n_test: 50
36 | n_test_vis: 4
37 | test_start_seed: 100000
38 | # use python's eval function as resolver, single-quoted string as argument
39 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
40 | n_obs_steps: ${n_obs_steps}
41 | n_action_steps: ${n_action_steps}
42 | render_obs_key: 'agentview_image'
43 | fps: 10
44 | crf: 22
45 | past_action: ${past_action_visible}
46 | abs_action: *abs_action
47 | tqdm_interval_sec: 1.0
48 | n_envs: 28
49 | # evaluation at this config requires a 16 core 64GB instance.
50 |
51 | dataset:
52 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
53 | shape_meta: *shape_meta
54 | dataset_path: *dataset_path
55 | horizon: ${horizon}
56 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
57 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
58 | n_obs_steps: ${dataset_obs_steps}
59 | abs_action: *abs_action
60 | rotation_rep: 'rotation_6d'
61 | use_legacy_normalizer: False
62 | use_cache: True
63 | seed: 42
64 | val_ratio: 0.02
65 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/can_lowdim.yaml:
--------------------------------------------------------------------------------
1 | name: can_lowdim
2 |
3 | obs_dim: 23
4 | action_dim: 7
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos']
8 | task_name: &task_name can
9 | dataset_type: &dataset_type ph
10 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5
11 | abs_action: &abs_action False
12 |
13 | env_runner:
14 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
15 | dataset_path: *dataset_path
16 | obs_keys: *obs_keys
17 | n_train: 6
18 | n_train_vis: 2
19 | train_start_idx: 0
20 | n_test: 50
21 | n_test_vis: 4
22 | test_start_seed: 100000
23 | # use python's eval function as resolver, single-quoted string as argument
24 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
25 | n_obs_steps: ${n_obs_steps}
26 | n_action_steps: ${n_action_steps}
27 | n_latency_steps: ${n_latency_steps}
28 | render_hw: [128,128]
29 | fps: 10
30 | crf: 22
31 | past_action: ${past_action_visible}
32 | abs_action: *abs_action
33 | n_envs: 28
34 |
35 | dataset:
36 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
37 | dataset_path: *dataset_path
38 | horizon: ${horizon}
39 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
40 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
41 | obs_keys: *obs_keys
42 | abs_action: *abs_action
43 | use_legacy_normalizer: False
44 | seed: 42
45 | val_ratio: 0.02
46 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/can_lowdim_abs.yaml:
--------------------------------------------------------------------------------
1 | name: can_lowdim
2 |
3 | obs_dim: 23
4 | action_dim: 10
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos']
8 | task_name: &task_name can
9 | dataset_type: &dataset_type ph
10 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5
11 | abs_action: &abs_action True
12 |
13 | env_runner:
14 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
15 | dataset_path: *dataset_path
16 | obs_keys: *obs_keys
17 | n_train: 6
18 | n_train_vis: 2
19 | train_start_idx: 0
20 | n_test: 50
21 | n_test_vis: 4
22 | test_start_seed: 100000
23 | # use python's eval function as resolver, single-quoted string as argument
24 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
25 | n_obs_steps: ${n_obs_steps}
26 | n_action_steps: ${n_action_steps}
27 | n_latency_steps: ${n_latency_steps}
28 | render_hw: [128,128]
29 | fps: 10
30 | crf: 22
31 | past_action: ${past_action_visible}
32 | abs_action: *abs_action
33 | n_envs: 28
34 |
35 | dataset:
36 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
37 | dataset_path: *dataset_path
38 | horizon: ${horizon}
39 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
40 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
41 | obs_keys: *obs_keys
42 | abs_action: *abs_action
43 | use_legacy_normalizer: False
44 | rotation_rep: rotation_6d
45 | seed: 42
46 | val_ratio: 0.02
47 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/kitchen_lowdim.yaml:
--------------------------------------------------------------------------------
1 | name: kitchen_lowdim
2 |
3 | obs_dim: 60
4 | action_dim: 9
5 | keypoint_dim: 3
6 |
7 | dataset_dir: &dataset_dir data/kitchen
8 |
9 | env_runner:
10 | _target_: diffusion_policy.env_runner.kitchen_lowdim_runner.KitchenLowdimRunner
11 | dataset_dir: *dataset_dir
12 | n_train: 6
13 | n_train_vis: 2
14 | train_start_seed: 0
15 | n_test: 50
16 | n_test_vis: 4
17 | test_start_seed: 100000
18 | max_steps: 280
19 | n_obs_steps: ${n_obs_steps}
20 | n_action_steps: ${n_action_steps}
21 | render_hw: [240, 360]
22 | fps: 12.5
23 | past_action: ${past_action_visible}
24 | n_envs: null
25 |
26 | dataset:
27 | _target_: diffusion_policy.dataset.kitchen_lowdim_dataset.KitchenLowdimDataset
28 | dataset_dir: *dataset_dir
29 | horizon: ${horizon}
30 | pad_before: ${eval:'${n_obs_steps}-1'}
31 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
32 | seed: 42
33 | val_ratio: 0.02
34 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/kitchen_lowdim_abs.yaml:
--------------------------------------------------------------------------------
1 | name: kitchen_lowdim
2 |
3 | obs_dim: 60
4 | action_dim: 9
5 | keypoint_dim: 3
6 |
7 | abs_action: True
8 | robot_noise_ratio: 0.1
9 |
10 | env_runner:
11 | _target_: diffusion_policy.env_runner.kitchen_lowdim_runner.KitchenLowdimRunner
12 | dataset_dir: data/kitchen
13 | n_train: 6
14 | n_train_vis: 2
15 | train_start_seed: 0
16 | n_test: 50
17 | n_test_vis: 4
18 | test_start_seed: 100000
19 | max_steps: 280
20 | n_obs_steps: ${n_obs_steps}
21 | n_action_steps: ${n_action_steps}
22 | render_hw: [240, 360]
23 | fps: 12.5
24 | past_action: ${past_action_visible}
25 | abs_action: ${task.abs_action}
26 | robot_noise_ratio: ${task.robot_noise_ratio}
27 | n_envs: null
28 |
29 | dataset:
30 | _target_: diffusion_policy.dataset.kitchen_mjl_lowdim_dataset.KitchenMjlLowdimDataset
31 | dataset_dir: data/kitchen/kitchen_demos_multitask
32 | horizon: ${horizon}
33 | pad_before: ${eval:'${n_obs_steps}-1'}
34 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
35 | abs_action: ${task.abs_action}
36 | robot_noise_ratio: ${task.robot_noise_ratio}
37 | seed: 42
38 | val_ratio: 0.02
39 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/lift_abs.yaml:
--------------------------------------------------------------------------------
1 | name: lift_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | agentview_image:
7 | shape: [3, 84, 84]
8 | horizon: 2
9 | type: rgb
10 | robot0_eye_in_hand_image:
11 | shape: [3, 84, 84]
12 | horizon: 2
13 | type: rgb
14 | robot0_eef_pos:
15 | shape: [3]
16 | horizon: 2
17 | type: low_dim
18 | robot0_eef_quat:
19 | raw_shape: [4]
20 | shape: [6]
21 | horizon: 2
22 | type: low_dim
23 | rotation_rep: rotation_6d
24 | robot0_gripper_qpos:
25 | shape: [2]
26 | horizon: 2
27 | type: low_dim
28 | action:
29 | shape: [10]
30 | rotation_rep: rotation_6d
31 |
32 | task_name: &task_name lift
33 | dataset_type: &dataset_type ph
34 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
35 | pose_repr: &pose_repr
36 | obs_pose_repr: rel # abs or rel
37 | action_pose_repr: rel
38 |
39 | env_runner:
40 | _target_: diffusion_policy.env_runner.robomimic_runner.RobomimicRunner
41 | dataset_path: *dataset_path
42 | pose_repr: *pose_repr
43 | shape_meta: *shape_meta
44 | n_train: 6
45 | n_train_vis: 2
46 | train_start_idx: 0
47 | n_test: 50
48 | n_test_vis: 4
49 | test_start_seed: 100000
50 | # use python's eval function as resolver, single-quoted string as argument
51 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
52 | n_obs_steps: ${n_obs_steps}
53 | n_action_steps: ${n_action_steps}
54 | render_obs_key: 'agentview_image'
55 | fps: 10
56 | crf: 22
57 | tqdm_interval_sec: 1.0
58 | n_envs: 28
59 | # evaluation at this config requires a 16 core 64GB instance.
60 |
61 | dataset:
62 | _target_: diffusion_policy.dataset.robomimic_replay_dataset.RobomimicReplayDataset
63 | pose_repr: *pose_repr
64 | shape_meta: *shape_meta
65 | dataset_path: *dataset_path
66 | horizon: ${horizon}
67 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
68 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
69 | n_obs_steps: ${dataset_obs_steps}
70 | use_cache: True
71 | seed: 42
72 | val_ratio: 0.02
73 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/lift_image.yaml:
--------------------------------------------------------------------------------
1 | name: lift_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | agentview_image:
7 | shape: [3, 84, 84]
8 | type: rgb
9 | robot0_eye_in_hand_image:
10 | shape: [3, 84, 84]
11 | type: rgb
12 | robot0_eef_pos:
13 | shape: [3]
14 | # type default: low_dim
15 | robot0_eef_quat:
16 | shape: [4]
17 | robot0_gripper_qpos:
18 | shape: [2]
19 | action:
20 | shape: [7]
21 |
22 | task_name: &task_name lift
23 | dataset_type: &dataset_type ph
24 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5
25 | abs_action: &abs_action False
26 |
27 | env_runner:
28 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
29 | dataset_path: *dataset_path
30 | shape_meta: *shape_meta
31 | # costs 1GB per env
32 | n_train: 6
33 | n_train_vis: 1
34 | train_start_idx: 0
35 | n_test: 50
36 | n_test_vis: 3
37 | test_start_seed: 100000
38 | # use python's eval function as resolver, single-quoted string as argument
39 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
40 | n_obs_steps: ${n_obs_steps}
41 | n_action_steps: ${n_action_steps}
42 | render_obs_key: 'agentview_image'
43 | fps: 10
44 | crf: 22
45 | past_action: ${past_action_visible}
46 | abs_action: *abs_action
47 | tqdm_interval_sec: 1.0
48 | n_envs: 28
49 | # evaluation at this config requires a 16 core 64GB instance.
50 |
51 | dataset:
52 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
53 | shape_meta: *shape_meta
54 | dataset_path: *dataset_path
55 | horizon: ${horizon}
56 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
57 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
58 | n_obs_steps: ${dataset_obs_steps}
59 | abs_action: *abs_action
60 | rotation_rep: 'rotation_6d'
61 | use_legacy_normalizer: False
62 | use_cache: True
63 | seed: 42
64 | val_ratio: 0.02
65 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/lift_image_abs.yaml:
--------------------------------------------------------------------------------
1 | name: lift_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | agentview_image:
7 | shape: [3, 84, 84]
8 | type: rgb
9 | robot0_eye_in_hand_image:
10 | shape: [3, 84, 84]
11 | type: rgb
12 | robot0_eef_pos:
13 | shape: [3]
14 | # type default: low_dim
15 | robot0_eef_quat:
16 | shape: [4]
17 | robot0_gripper_qpos:
18 | shape: [2]
19 | action:
20 | shape: [10]
21 |
22 | task_name: &task_name lift
23 | dataset_type: &dataset_type ph
24 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
25 | abs_action: &abs_action True
26 |
27 | env_runner:
28 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
29 | dataset_path: *dataset_path
30 | shape_meta: *shape_meta
31 | n_train: 6
32 | n_train_vis: 2
33 | train_start_idx: 0
34 | n_test: 50
35 | n_test_vis: 4
36 | test_start_seed: 100000
37 | # use python's eval function as resolver, single-quoted string as argument
38 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
39 | n_obs_steps: ${n_obs_steps}
40 | n_action_steps: ${n_action_steps}
41 | render_obs_key: 'agentview_image'
42 | fps: 10
43 | crf: 22
44 | past_action: ${past_action_visible}
45 | abs_action: *abs_action
46 | tqdm_interval_sec: 1.0
47 | n_envs: 28
48 | # evaluation at this config requires a 16 core 64GB instance.
49 |
50 | dataset:
51 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
52 | shape_meta: *shape_meta
53 | dataset_path: *dataset_path
54 | horizon: ${horizon}
55 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
56 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
57 | n_obs_steps: ${dataset_obs_steps}
58 | abs_action: *abs_action
59 | rotation_rep: 'rotation_6d'
60 | use_legacy_normalizer: False
61 | use_cache: True
62 | seed: 42
63 | val_ratio: 0.02
64 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/lift_lowdim.yaml:
--------------------------------------------------------------------------------
1 | name: lift_lowdim
2 |
3 | obs_dim: 19
4 | action_dim: 7
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos']
8 | task_name: &task_name lift
9 | dataset_type: &dataset_type ph
10 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5
11 | abs_action: &abs_action False
12 |
13 | env_runner:
14 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
15 | dataset_path: *dataset_path
16 | obs_keys: *obs_keys
17 | n_train: 6
18 | n_train_vis: 2
19 | train_start_idx: 0
20 | n_test: 50
21 | n_test_vis: 4
22 | test_start_seed: 100000
23 | # use python's eval function as resolver, single-quoted string as argument
24 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
25 | n_obs_steps: ${n_obs_steps}
26 | n_action_steps: ${n_action_steps}
27 | n_latency_steps: ${n_latency_steps}
28 | render_hw: [128,128]
29 | fps: 10
30 | crf: 22
31 | past_action: ${past_action_visible}
32 | abs_action: *abs_action
33 | tqdm_interval_sec: 1.0
34 | n_envs: 28
35 |
36 | dataset:
37 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
38 | dataset_path: *dataset_path
39 | horizon: ${horizon}
40 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
41 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
42 | obs_keys: *obs_keys
43 | abs_action: *abs_action
44 | use_legacy_normalizer: False
45 | seed: 42
46 | val_ratio: 0.02
47 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/lift_lowdim_abs.yaml:
--------------------------------------------------------------------------------
1 | name: lift_lowdim
2 |
3 | obs_dim: 19
4 | action_dim: 10
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos']
8 | task_name: &task_name lift
9 | dataset_type: &dataset_type ph
10 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5
11 | abs_action: &abs_action True
12 |
13 | env_runner:
14 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
15 | dataset_path: *dataset_path
16 | obs_keys: *obs_keys
17 | n_train: 6
18 | n_train_vis: 2
19 | train_start_idx: 0
20 | n_test: 50
21 | n_test_vis: 3
22 | test_start_seed: 100000
23 | # use python's eval function as resolver, single-quoted string as argument
24 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
25 | n_obs_steps: ${n_obs_steps}
26 | n_action_steps: ${n_action_steps}
27 | n_latency_steps: ${n_latency_steps}
28 | render_hw: [128,128]
29 | fps: 10
30 | crf: 22
31 | past_action: ${past_action_visible}
32 | abs_action: *abs_action
33 | tqdm_interval_sec: 1.0
34 | n_envs: 28
35 |
36 | dataset:
37 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
38 | dataset_path: *dataset_path
39 | horizon: ${horizon}
40 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
41 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
42 | obs_keys: *obs_keys
43 | abs_action: *abs_action
44 | use_legacy_normalizer: False
45 | rotation_rep: rotation_6d
46 | seed: 42
47 | val_ratio: 0.02
48 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/pusht_image.yaml:
--------------------------------------------------------------------------------
1 | name: pusht_image
2 |
3 | image_shape: &image_shape [3, 96, 96]
4 | shape_meta: &shape_meta
5 | # acceptable types: rgb, low_dim
6 | obs:
7 | image:
8 | shape: *image_shape
9 | type: rgb
10 | agent_pos:
11 | shape: [2]
12 | type: low_dim
13 | action:
14 | shape: [2]
15 |
16 | env_runner:
17 | _target_: diffusion_policy.env_runner.pusht_image_runner.PushTImageRunner
18 | n_train: 6
19 | n_train_vis: 2
20 | train_start_seed: 0
21 | n_test: 50
22 | n_test_vis: 4
23 | legacy_test: True
24 | test_start_seed: 100000
25 | max_steps: 300
26 | n_obs_steps: ${n_obs_steps}
27 | n_action_steps: ${n_action_steps}
28 | fps: 10
29 | past_action: ${past_action_visible}
30 | n_envs: null
31 |
32 | dataset:
33 | _target_: diffusion_policy.dataset.pusht_image_dataset.PushTImageDataset
34 | zarr_path: data/pusht/pusht_cchi_v7_replay.zarr
35 | horizon: ${horizon}
36 | pad_before: ${eval:'${n_obs_steps}-1'}
37 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
38 | seed: 42
39 | val_ratio: 0.02
40 | max_train_episodes: 90
41 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/pusht_lowdim.yaml:
--------------------------------------------------------------------------------
1 | name: pusht_lowdim
2 |
3 | obs_dim: 20 # 9*2 keypoints + 2 state
4 | action_dim: 2
5 | keypoint_dim: 2
6 |
7 | env_runner:
8 | _target_: diffusion_policy.env_runner.pusht_keypoints_runner.PushTKeypointsRunner
9 | keypoint_visible_rate: ${keypoint_visible_rate}
10 | n_train: 6
11 | n_train_vis: 2
12 | train_start_seed: 0
13 | n_test: 50
14 | n_test_vis: 4
15 | legacy_test: True
16 | test_start_seed: 100000
17 | max_steps: 300
18 | n_obs_steps: ${n_obs_steps}
19 | n_action_steps: ${n_action_steps}
20 | n_latency_steps: ${n_latency_steps}
21 | fps: 10
22 | agent_keypoints: False
23 | past_action: ${past_action_visible}
24 | n_envs: null
25 |
26 | dataset:
27 | _target_: diffusion_policy.dataset.pusht_dataset.PushTLowdimDataset
28 | zarr_path: data/pusht/pusht_cchi_v7_replay.zarr
29 | horizon: ${horizon}
30 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
31 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
32 | seed: 42
33 | val_ratio: 0.02
34 | max_train_episodes: 90
35 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/real_pusht_image.yaml:
--------------------------------------------------------------------------------
1 | name: real_image
2 |
3 | image_shape: [3, 240, 320]
4 | dataset_path: data/pusht_real/real_pusht_20230105
5 |
6 | shape_meta: &shape_meta
7 | # acceptable types: rgb, low_dim
8 | obs:
9 | # camera_0:
10 | # shape: ${task.image_shape}
11 | # type: rgb
12 | camera_1:
13 | shape: ${task.image_shape}
14 | type: rgb
15 | # camera_2:
16 | # shape: ${task.image_shape}
17 | # type: rgb
18 | camera_3:
19 | shape: ${task.image_shape}
20 | type: rgb
21 | # camera_4:
22 | # shape: ${task.image_shape}
23 | # type: rgb
24 | robot_eef_pose:
25 | shape: [2]
26 | type: low_dim
27 | action:
28 | shape: [2]
29 |
30 | env_runner:
31 | _target_: diffusion_policy.env_runner.real_pusht_image_runner.RealPushTImageRunner
32 |
33 | dataset:
34 | _target_: diffusion_policy.dataset.real_pusht_image_dataset.RealPushTImageDataset
35 | shape_meta: *shape_meta
36 | dataset_path: ${task.dataset_path}
37 | horizon: ${horizon}
38 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
39 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
40 | n_obs_steps: ${dataset_obs_steps}
41 | n_latency_steps: ${n_latency_steps}
42 | use_cache: True
43 | seed: 42
44 | val_ratio: 0.00
45 | max_train_episodes: null
46 | delta_action: False
47 |
48 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/square_abs.yaml:
--------------------------------------------------------------------------------
1 | name: square_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | agentview_image:
7 | shape: [3, 84, 84]
8 | horizon: 2
9 | type: rgb
10 | robot0_eye_in_hand_image:
11 | shape: [3, 84, 84]
12 | horizon: 2
13 | type: rgb
14 | robot0_eef_pos:
15 | shape: [3]
16 | horizon: 2
17 | type: low_dim
18 | robot0_eef_quat:
19 | raw_shape: [4]
20 | shape: [6]
21 | horizon: 2
22 | type: low_dim
23 | rotation_rep: rotation_6d
24 | robot0_gripper_qpos:
25 | shape: [2]
26 | horizon: 2
27 | type: low_dim
28 | action:
29 | shape: [10]
30 | rotation_rep: rotation_6d
31 |
32 | task_name: &task_name square
33 | dataset_type: &dataset_type ph
34 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
35 | pose_repr: &pose_repr
36 | obs_pose_repr: rel # abs or rel
37 | action_pose_repr: rel
38 |
39 | env_runner:
40 | _target_: diffusion_policy.env_runner.robomimic_runner.RobomimicRunner
41 | dataset_path: *dataset_path
42 | pose_repr: *pose_repr
43 | shape_meta: *shape_meta
44 | n_train: 6
45 | n_train_vis: 2
46 | train_start_idx: 0
47 | n_test: 50
48 | n_test_vis: 4
49 | test_start_seed: 100000
50 | # use python's eval function as resolver, single-quoted string as argument
51 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
52 | n_obs_steps: ${n_obs_steps}
53 | n_action_steps: ${n_action_steps}
54 | render_obs_key: 'agentview_image'
55 | fps: 10
56 | crf: 22
57 | tqdm_interval_sec: 1.0
58 | n_envs: 28
59 | # evaluation at this config requires a 16 core 64GB instance.
60 |
61 | dataset:
62 | _target_: diffusion_policy.dataset.robomimic_replay_dataset.RobomimicReplayDataset
63 | pose_repr: *pose_repr
64 | shape_meta: *shape_meta
65 | dataset_path: *dataset_path
66 | horizon: ${horizon}
67 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
68 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
69 | n_obs_steps: ${dataset_obs_steps}
70 | use_cache: True
71 | seed: 42
72 | val_ratio: 0.02
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/square_image.yaml:
--------------------------------------------------------------------------------
1 | name: square_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | agentview_image:
7 | shape: [3, 84, 84]
8 | type: rgb
9 | robot0_eye_in_hand_image:
10 | shape: [3, 84, 84]
11 | type: rgb
12 | robot0_eef_pos:
13 | shape: [3]
14 | # type default: low_dim
15 | robot0_eef_quat:
16 | shape: [4]
17 | robot0_gripper_qpos:
18 | shape: [2]
19 | action:
20 | shape: [7]
21 |
22 | task_name: &task_name square
23 | dataset_type: &dataset_type ph
24 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5
25 | abs_action: &abs_action False
26 |
27 | env_runner:
28 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
29 | dataset_path: *dataset_path
30 | shape_meta: *shape_meta
31 | # costs 1GB per env
32 | n_train: 6
33 | n_train_vis: 2
34 | train_start_idx: 0
35 | n_test: 50
36 | n_test_vis: 4
37 | test_start_seed: 100000
38 | # use python's eval function as resolver, single-quoted string as argument
39 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
40 | n_obs_steps: ${n_obs_steps}
41 | n_action_steps: ${n_action_steps}
42 | render_obs_key: 'agentview_image'
43 | fps: 10
44 | crf: 22
45 | past_action: ${past_action_visible}
46 | abs_action: *abs_action
47 | tqdm_interval_sec: 1.0
48 | n_envs: 28
49 | # evaluation at this config requires a 16 core 64GB instance.
50 |
51 | dataset:
52 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
53 | shape_meta: *shape_meta
54 | dataset_path: *dataset_path
55 | horizon: ${horizon}
56 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
57 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
58 | n_obs_steps: ${dataset_obs_steps}
59 | abs_action: *abs_action
60 | rotation_rep: 'rotation_6d'
61 | use_legacy_normalizer: False
62 | use_cache: True
63 | seed: 42
64 | val_ratio: 0.02
65 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/square_image_abs.yaml:
--------------------------------------------------------------------------------
1 | name: square_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | agentview_image:
7 | shape: [3, 84, 84]
8 | type: rgb
9 | robot0_eye_in_hand_image:
10 | shape: [3, 84, 84]
11 | type: rgb
12 | robot0_eef_pos:
13 | shape: [3]
14 | # type default: low_dim
15 | robot0_eef_quat:
16 | shape: [4]
17 | robot0_gripper_qpos:
18 | shape: [2]
19 | action:
20 | shape: [10]
21 |
22 | task_name: &task_name square
23 | dataset_type: &dataset_type ph
24 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
25 | abs_action: &abs_action True
26 |
27 | env_runner:
28 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
29 | dataset_path: *dataset_path
30 | shape_meta: *shape_meta
31 | # costs 1GB per env
32 | n_train: 6
33 | n_train_vis: 2
34 | train_start_idx: 0
35 | n_test: 50
36 | n_test_vis: 4
37 | test_start_seed: 100000
38 | # use python's eval function as resolver, single-quoted string as argument
39 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
40 | n_obs_steps: ${n_obs_steps}
41 | n_action_steps: ${n_action_steps}
42 | render_obs_key: 'agentview_image'
43 | fps: 10
44 | crf: 22
45 | past_action: ${past_action_visible}
46 | abs_action: *abs_action
47 | tqdm_interval_sec: 1.0
48 | n_envs: 28
49 | # evaluation at this config requires a 16 core 64GB instance.
50 |
51 | dataset:
52 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
53 | shape_meta: *shape_meta
54 | dataset_path: *dataset_path
55 | horizon: ${horizon}
56 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
57 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
58 | n_obs_steps: ${dataset_obs_steps}
59 | abs_action: *abs_action
60 | rotation_rep: 'rotation_6d'
61 | use_legacy_normalizer: False
62 | use_cache: True
63 | seed: 42
64 | val_ratio: 0.02
65 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/square_lowdim.yaml:
--------------------------------------------------------------------------------
1 | name: square_lowdim
2 |
3 | obs_dim: 23
4 | action_dim: 7
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos']
8 | task_name: &task_name square
9 | dataset_type: &dataset_type ph
10 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5
11 | abs_action: &abs_action False
12 |
13 | env_runner:
14 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
15 | dataset_path: *dataset_path
16 | obs_keys: *obs_keys
17 | n_train: 6
18 | n_train_vis: 2
19 | train_start_idx: 0
20 | n_test: 50
21 | n_test_vis: 4
22 | test_start_seed: 100000
23 | # use python's eval function as resolver, single-quoted string as argument
24 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
25 | n_obs_steps: ${n_obs_steps}
26 | n_action_steps: ${n_action_steps}
27 | n_latency_steps: ${n_latency_steps}
28 | render_hw: [128,128]
29 | fps: 10
30 | crf: 22
31 | past_action: ${past_action_visible}
32 | abs_action: *abs_action
33 | n_envs: 28
34 |
35 | dataset:
36 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
37 | dataset_path: *dataset_path
38 | horizon: ${horizon}
39 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
40 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
41 | obs_keys: *obs_keys
42 | abs_action: *abs_action
43 | use_legacy_normalizer: False
44 | seed: 42
45 | val_ratio: 0.02
46 | max_train_episodes: null
47 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/square_lowdim_abs.yaml:
--------------------------------------------------------------------------------
1 | name: square_lowdim
2 |
3 | obs_dim: 23
4 | action_dim: 10
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos']
8 | task_name: &task_name square
9 | dataset_type: &dataset_type ph
10 | abs_action: &abs_action True
11 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5
12 |
13 |
14 | env_runner:
15 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
16 | dataset_path: *dataset_path
17 | obs_keys: *obs_keys
18 | n_train: 6
19 | n_train_vis: 2
20 | train_start_idx: 0
21 | n_test: 50
22 | n_test_vis: 4
23 | test_start_seed: 100000
24 | # use python's eval function as resolver, single-quoted string as argument
25 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
26 | n_obs_steps: ${n_obs_steps}
27 | n_action_steps: ${n_action_steps}
28 | n_latency_steps: ${n_latency_steps}
29 | render_hw: [128,128]
30 | fps: 10
31 | crf: 22
32 | past_action: ${past_action_visible}
33 | abs_action: *abs_action
34 | n_envs: 28
35 |
36 | dataset:
37 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
38 | dataset_path: *dataset_path
39 | horizon: ${horizon}
40 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
41 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
42 | obs_keys: *obs_keys
43 | abs_action: *abs_action
44 | use_legacy_normalizer: False
45 | seed: 42
46 | val_ratio: 0.02
47 | max_train_episodes: null
48 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/tool_hang_abs.yaml:
--------------------------------------------------------------------------------
1 | name: tool_hang_image_abs
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | sideview_image:
7 | shape: [3, 240, 240]
8 | horizon: 2
9 | type: rgb
10 | robot0_eye_in_hand_image:
11 | shape: [3, 240, 240]
12 | horizon: 2
13 | type: rgb
14 | robot0_eef_pos:
15 | shape: [3]
16 | type: low_dim
17 | robot0_eef_quat:
18 | raw_shape: [4]
19 | shape: [6]
20 | horizon: 2
21 | type: low_dim
22 | rotation_rep: rotation_6d
23 | robot0_gripper_qpos:
24 | shape: [2]
25 | horizon: 2
26 | type: low_dim
27 | action:
28 | shape: [10]
29 | rotation_rep: rotation_6d
30 |
31 | task_name: &task_name tool_hang
32 | dataset_type: &dataset_type ph
33 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
34 | pose_repr: &pose_repr
35 | obs_pose_repr: rel # abs or rel
36 | action_pose_repr: rel
37 |
38 | env_runner:
39 | _target_: diffusion_policy.env_runner.robomimic_runner.RobomimicRunner
40 | dataset_path: *dataset_path
41 | pose_repr: *pose_repr
42 | shape_meta: *shape_meta
43 | n_train: 6
44 | n_train_vis: 2
45 | train_start_idx: 0
46 | n_test: 50
47 | n_test_vis: 4
48 | test_start_seed: 100000
49 | # use python's eval function as resolver, single-quoted string as argument
50 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
51 | n_obs_steps: ${n_obs_steps}
52 | n_action_steps: ${n_action_steps}
53 | render_obs_key: 'sideview_image'
54 | fps: 10
55 | crf: 22
56 | tqdm_interval_sec: 1.0
57 | n_envs: 28
58 | # evaluation at this config requires a 16 core 64GB instance.
59 |
60 | dataset:
61 | _target_: diffusion_policy.dataset.robomimic_replay_dataset.RobomimicReplayDataset
62 | pose_repr: *pose_repr
63 | shape_meta: *shape_meta
64 | dataset_path: *dataset_path
65 | horizon: ${horizon}
66 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
67 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
68 | n_obs_steps: ${dataset_obs_steps}
69 | use_cache: True
70 | seed: 42
71 | val_ratio: 0.02
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/tool_hang_image.yaml:
--------------------------------------------------------------------------------
1 | name: tool_hang_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | sideview_image:
7 | shape: [3, 240, 240]
8 | type: rgb
9 | robot0_eye_in_hand_image:
10 | shape: [3, 240, 240]
11 | type: rgb
12 | robot0_eef_pos:
13 | shape: [3]
14 | # type default: low_dim
15 | robot0_eef_quat:
16 | shape: [4]
17 | robot0_gripper_qpos:
18 | shape: [2]
19 | action:
20 | shape: [7]
21 |
22 | task_name: &task_name tool_hang
23 | dataset_type: &dataset_type ph
24 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5
25 | abs_action: &abs_action False
26 |
27 | env_runner:
28 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
29 | dataset_path: *dataset_path
30 | shape_meta: *shape_meta
31 | # costs 1GB per env
32 | n_train: 6
33 | n_train_vis: 2
34 | train_start_idx: 0
35 | n_test: 50
36 | n_test_vis: 4
37 | test_start_seed: 100000
38 | max_steps: 700
39 | n_obs_steps: ${n_obs_steps}
40 | n_action_steps: ${n_action_steps}
41 | render_obs_key: 'sideview_image'
42 | fps: 10
43 | crf: 22
44 | past_action: ${past_action_visible}
45 | abs_action: *abs_action
46 | tqdm_interval_sec: 1.0
47 | n_envs: 28
48 | # evaluation at this config requires a 16 core 64GB instance.
49 |
50 | dataset:
51 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
52 | shape_meta: *shape_meta
53 | dataset_path: *dataset_path
54 | horizon: ${horizon}
55 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
56 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
57 | n_obs_steps: ${dataset_obs_steps}
58 | abs_action: *abs_action
59 | rotation_rep: 'rotation_6d'
60 | use_legacy_normalizer: False
61 | use_cache: True
62 | seed: 42
63 | val_ratio: 0.02
64 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/tool_hang_image_abs.yaml:
--------------------------------------------------------------------------------
1 | name: tool_hang_image_abs
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | sideview_image:
7 | shape: [3, 240, 240]
8 | type: rgb
9 | robot0_eye_in_hand_image:
10 | shape: [3, 240, 240]
11 | type: rgb
12 | robot0_eef_pos:
13 | shape: [3]
14 | # type default: low_dim
15 | robot0_eef_quat:
16 | shape: [4]
17 | robot0_gripper_qpos:
18 | shape: [2]
19 | action:
20 | shape: [10]
21 |
22 | task_name: &task_name tool_hang
23 | dataset_type: &dataset_type ph
24 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
25 | abs_action: &abs_action True
26 |
27 | env_runner:
28 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
29 | dataset_path: *dataset_path
30 | shape_meta: *shape_meta
31 | # costs 1GB per env
32 | n_train: 6
33 | n_train_vis: 2
34 | train_start_idx: 0
35 | n_test: 50
36 | n_test_vis: 4
37 | test_start_seed: 100000
38 | max_steps: 700
39 | n_obs_steps: ${n_obs_steps}
40 | n_action_steps: ${n_action_steps}
41 | render_obs_key: 'sideview_image'
42 | fps: 10
43 | crf: 22
44 | past_action: ${past_action_visible}
45 | abs_action: *abs_action
46 | tqdm_interval_sec: 1.0
47 | n_envs: 28
48 | # evaluation at this config requires a 16 core 64GB instance.
49 |
50 | dataset:
51 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
52 | shape_meta: *shape_meta
53 | dataset_path: *dataset_path
54 | horizon: ${horizon}
55 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
56 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
57 | n_obs_steps: ${dataset_obs_steps}
58 | abs_action: *abs_action
59 | rotation_rep: 'rotation_6d'
60 | use_legacy_normalizer: False
61 | use_cache: True
62 | seed: 42
63 | val_ratio: 0.02
64 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/tool_hang_lowdim.yaml:
--------------------------------------------------------------------------------
1 | name: tool_hang_lowdim
2 |
3 | obs_dim: 53
4 | action_dim: 7
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos']
8 | task_name: &task_name tool_hang
9 | dataset_type: &dataset_type ph
10 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5
11 | abs_action: &abs_action False
12 |
13 | env_runner:
14 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
15 | dataset_path: *dataset_path
16 | obs_keys: *obs_keys
17 | n_train: 6
18 | n_train_vis: 2
19 | train_start_idx: 0
20 | n_test: 50
21 | n_test_vis: 4
22 | test_start_seed: 100000
23 | max_steps: 700
24 | n_obs_steps: ${n_obs_steps}
25 | n_action_steps: ${n_action_steps}
26 | n_latency_steps: ${n_latency_steps}
27 | render_hw: [128,128]
28 | fps: 10
29 | crf: 22
30 | past_action: ${past_action_visible}
31 | abs_action: *abs_action
32 | n_envs: 28
33 | # seed 42 will crash MuJoCo for some reason.
34 |
35 | dataset:
36 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
37 | dataset_path: *dataset_path
38 | horizon: ${horizon}
39 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
40 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
41 | obs_keys: *obs_keys
42 | abs_action: *abs_action
43 | use_legacy_normalizer: False
44 | seed: 42
45 | val_ratio: 0.02
46 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/tool_hang_lowdim_abs.yaml:
--------------------------------------------------------------------------------
1 | name: tool_hang_lowdim
2 |
3 | obs_dim: 53
4 | action_dim: 10
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys ['object', 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos']
8 | task_name: &task_name tool_hang
9 | dataset_type: &dataset_type ph
10 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5
11 | abs_action: &abs_action True
12 |
13 | env_runner:
14 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
15 | dataset_path: *dataset_path
16 | obs_keys: *obs_keys
17 | n_train: 6
18 | n_train_vis: 2
19 | train_start_idx: 0
20 | n_test: 50
21 | n_test_vis: 4
22 | test_start_seed: 100000
23 | max_steps: 700
24 | n_obs_steps: ${n_obs_steps}
25 | n_action_steps: ${n_action_steps}
26 | n_latency_steps: ${n_latency_steps}
27 | render_hw: [128,128]
28 | fps: 10
29 | crf: 22
30 | past_action: ${past_action_visible}
31 | abs_action: *abs_action
32 | n_envs: 28
33 | # seed 42 will crash MuJoCo for some reason.
34 |
35 | dataset:
36 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
37 | dataset_path: *dataset_path
38 | horizon: ${horizon}
39 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
40 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
41 | obs_keys: *obs_keys
42 | abs_action: *abs_action
43 | use_legacy_normalizer: False
44 | rotation_rep: rotation_6d
45 | seed: 42
46 | val_ratio: 0.02
47 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/transport_lowdim.yaml:
--------------------------------------------------------------------------------
1 | name: transport_lowdim
2 |
3 | obs_dim: 59 # 41+(3+4+2)*2
4 | action_dim: 14 # 7*2
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys [
8 | 'object',
9 | 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos',
10 | 'robot1_eef_pos', 'robot1_eef_quat', 'robot1_gripper_qpos'
11 | ]
12 | task_name: &task_name transport
13 | dataset_type: &dataset_type ph
14 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim.hdf5
15 | abs_action: &abs_action False
16 |
17 | env_runner:
18 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
19 | dataset_path: *dataset_path
20 | obs_keys: *obs_keys
21 | n_train: 6
22 | n_train_vis: 2
23 | train_start_idx: 0
24 | n_test: 50
25 | n_test_vis: 5
26 | test_start_seed: 100000
27 | max_steps: 700
28 | n_obs_steps: ${n_obs_steps}
29 | n_action_steps: ${n_action_steps}
30 | n_latency_steps: ${n_latency_steps}
31 | render_hw: [128,128]
32 | fps: 10
33 | crf: 22
34 | past_action: ${past_action_visible}
35 | abs_action: *abs_action
36 | n_envs: 28
37 | # evaluation at this config requires a 16 core 64GB instance.
38 |
39 | dataset:
40 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
41 | dataset_path: *dataset_path
42 | horizon: ${horizon}
43 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
44 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
45 | obs_keys: *obs_keys
46 | abs_action: *abs_action
47 | use_legacy_normalizer: False
48 | seed: 42
49 | val_ratio: 0.02
50 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/transport_lowdim_abs.yaml:
--------------------------------------------------------------------------------
1 | name: transport_lowdim
2 |
3 | obs_dim: 59 # 41+(3+4+2)*2
4 | action_dim: 20 # 10*2
5 | keypoint_dim: 3
6 |
7 | obs_keys: &obs_keys [
8 | 'object',
9 | 'robot0_eef_pos', 'robot0_eef_quat', 'robot0_gripper_qpos',
10 | 'robot1_eef_pos', 'robot1_eef_quat', 'robot1_gripper_qpos'
11 | ]
12 | task_name: &task_name transport
13 | dataset_type: &dataset_type ph
14 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/low_dim_abs.hdf5
15 | abs_action: &abs_action True
16 |
17 | env_runner:
18 | _target_: diffusion_policy.env_runner.robomimic_lowdim_runner.RobomimicLowdimRunner
19 | dataset_path: *dataset_path
20 | obs_keys: *obs_keys
21 | n_train: 6
22 | n_train_vis: 2
23 | train_start_idx: 0
24 | n_test: 50
25 | n_test_vis: 4
26 | test_start_seed: 100000
27 | max_steps: 700
28 | n_obs_steps: ${n_obs_steps}
29 | n_action_steps: ${n_action_steps}
30 | n_latency_steps: ${n_latency_steps}
31 | render_hw: [128,128]
32 | fps: 10
33 | crf: 22
34 | past_action: ${past_action_visible}
35 | abs_action: *abs_action
36 | n_envs: 28
37 | # evaluation at this config requires a 16 core 64GB instance.
38 |
39 | dataset:
40 | _target_: diffusion_policy.dataset.robomimic_replay_lowdim_dataset.RobomimicReplayLowdimDataset
41 | dataset_path: *dataset_path
42 | horizon: ${horizon}
43 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
44 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
45 | obs_keys: *obs_keys
46 | abs_action: *abs_action
47 | use_legacy_normalizer: False
48 | seed: 42
49 | val_ratio: 0.02
50 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/umi.yaml:
--------------------------------------------------------------------------------
1 | name: umi
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | img:
7 | shape: [3, 224, 224]
8 | horizon: 2
9 | type: rgb
10 | robot0_eef_pos:
11 | shape: [3]
12 | horizon: 2
13 | type: low_dim
14 | robot0_eef_rot_axis_angle:
15 | raw_shape: [3]
16 | shape: [6]
17 | horizon: 2
18 | type: low_dim
19 | rotation_rep: rotation_6d
20 | robot0_gripper_width:
21 | shape: [1]
22 | horizon: 2
23 | type: low_dim
24 | action:
25 | shape: [10]
26 | horizon: 32
27 | rotation_rep: rotation_6d
28 |
29 | task_name: &task_name umi
30 | dataset_path: &dataset_path /local/real/datasets/umi/move_cube_20231003/dataset.zarr.zip
31 | pose_repr: &pose_repr
32 | obs_pose_repr: abs # abs or rel
33 | action_pose_repr: abs
34 |
35 | env_runner:
36 | _target_: diffusion_policy.env_runner.real_pusht_image_runner.RealPushTImageRunner
37 |
38 | dataset:
39 | _target_: diffusion_policy.dataset.umi_dataset.UmiDataset
40 | shape_meta: *shape_meta
41 | dataset_path: *dataset_path
42 | pose_repr: *pose_repr
43 | temporally_independent_normalization: False
44 | seed: 42
45 | val_ratio: 0.02
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/umi_image.yaml:
--------------------------------------------------------------------------------
1 | name: real_image
2 |
3 | image_shape: [3, 224, 224]
4 | dataset_path: data/gopro_mug_10hz.zarr.zip
5 |
6 | shape_meta: &shape_meta
7 | # acceptable types: rgb, low_dim
8 | obs:
9 | img:
10 | shape: ${task.image_shape}
11 | type: rgb
12 | robot_eef_pose:
13 | shape: [6]
14 | type: low_dim
15 | gripper_width:
16 | shape: [1]
17 | type: low_dim
18 | action:
19 | shape: [10]
20 |
21 | env_runner:
22 | _target_: diffusion_policy.env_runner.real_pusht_image_runner.RealPushTImageRunner
23 |
24 | dataset:
25 | _target_: diffusion_policy.dataset.umi_image_dataset.UmiImageDataset
26 | shape_meta: *shape_meta
27 | dataset_path: ${task.dataset_path}
28 | horizon: ${horizon}
29 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
30 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
31 | n_obs_steps: ${dataset_obs_steps}
32 | n_latency_steps: ${n_latency_steps}
33 | use_cache: True
34 | seed: 42
35 | val_ratio: 0.00
36 | max_train_episodes: null
37 |
--------------------------------------------------------------------------------
/diffusion_policy/config/legacy/task/umi_teleop_image.yaml:
--------------------------------------------------------------------------------
1 | name: real_image
2 |
3 | image_shape: [3, 224, 224]
4 | dataset_path: data/lift_cube_0921
5 |
6 | shape_meta: &shape_meta
7 | # acceptable types: rgb, low_dim
8 | obs:
9 | camera_0:
10 | shape: ${task.image_shape}
11 | type: rgb
12 | camera_1:
13 | shape: ${task.image_shape}
14 | type: rgb
15 | robot_eef_pose:
16 | shape: [6]
17 | type: low_dim
18 | gripper_position:
19 | shape: [1]
20 | type: low_dim
21 | action:
22 | shape: [7]
23 |
24 | env_runner:
25 | _target_: diffusion_policy.env_runner.real_pusht_image_runner.RealPushTImageRunner
26 |
27 | dataset:
28 | _target_: diffusion_policy.dataset.umi_teleop_image_dataset.UmiTeleopImageDataset
29 | shape_meta: *shape_meta
30 | dataset_path: ${task.dataset_path}
31 | horizon: ${horizon}
32 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
33 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
34 | n_obs_steps: ${dataset_obs_steps}
35 | n_latency_steps: ${n_latency_steps}
36 | use_cache: True
37 | seed: 42
38 | val_ratio: 0.00
39 | max_train_episodes: null
40 | delta_action: False
41 |
42 |
--------------------------------------------------------------------------------
/diffusion_policy/config/task/tool_hang_abs.yaml:
--------------------------------------------------------------------------------
1 | name: tool_hang_image_abs
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | sideview_image:
7 | shape: [3, 240, 240]
8 | horizon: 2
9 | type: rgb
10 | robot0_eye_in_hand_image:
11 | shape: [3, 240, 240]
12 | horizon: 2
13 | type: rgb
14 | robot0_eef_pos:
15 | shape: [3]
16 | type: low_dim
17 | robot0_eef_quat:
18 | raw_shape: [4]
19 | shape: [6]
20 | horizon: 2
21 | type: low_dim
22 | rotation_rep: rotation_6d
23 | robot0_gripper_qpos:
24 | shape: [2]
25 | horizon: 2
26 | type: low_dim
27 | action:
28 | shape: [10]
29 | rotation_rep: rotation_6d
30 |
31 | task_name: &task_name tool_hang
32 | dataset_type: &dataset_type ph
33 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
34 | pose_repr: &pose_repr
35 | obs_pose_repr: rel # abs or rel
36 | action_pose_repr: rel
37 |
38 | env_runner:
39 | _target_: diffusion_policy.env_runner.robomimic_runner.RobomimicRunner
40 | dataset_path: *dataset_path
41 | pose_repr: *pose_repr
42 | shape_meta: *shape_meta
43 | n_train: 6
44 | n_train_vis: 2
45 | train_start_idx: 0
46 | n_test: 50
47 | n_test_vis: 4
48 | test_start_seed: 100000
49 | # use python's eval function as resolver, single-quoted string as argument
50 | max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
51 | n_obs_steps: ${n_obs_steps}
52 | n_action_steps: ${n_action_steps}
53 | render_obs_key: 'sideview_image'
54 | fps: 10
55 | crf: 22
56 | tqdm_interval_sec: 1.0
57 | n_envs: 28
58 | # evaluation at this config requires a 16 core 64GB instance.
59 |
60 | dataset:
61 | _target_: diffusion_policy.dataset.robomimic_replay_dataset.RobomimicReplayDataset
62 | pose_repr: *pose_repr
63 | shape_meta: *shape_meta
64 | dataset_path: *dataset_path
65 | horizon: ${horizon}
66 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
67 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
68 | n_obs_steps: ${dataset_obs_steps}
69 | use_cache: True
70 | seed: 42
71 | val_ratio: 0.02
--------------------------------------------------------------------------------
/diffusion_policy/config/task/umi_image.yaml:
--------------------------------------------------------------------------------
1 | name: real_image
2 |
3 | image_shape: [3, 224, 224]
4 | dataset_path: data/gopro_mug_10hz.zarr.zip
5 |
6 | shape_meta: &shape_meta
7 | # acceptable types: rgb, low_dim
8 | obs:
9 | img:
10 | shape: ${task.image_shape}
11 | type: rgb
12 | robot_eef_pose:
13 | shape: [6]
14 | type: low_dim
15 | gripper_width:
16 | shape: [1]
17 | type: low_dim
18 | action:
19 | shape: [10]
20 |
21 | env_runner:
22 | _target_: diffusion_policy.env_runner.real_pusht_image_runner.RealPushTImageRunner
23 |
24 | dataset:
25 | _target_: diffusion_policy.dataset.umi_image_dataset.UmiImageDataset
26 | shape_meta: *shape_meta
27 | dataset_path: ${task.dataset_path}
28 | horizon: ${horizon}
29 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
30 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
31 | n_obs_steps: ${dataset_obs_steps}
32 | n_latency_steps: ${n_latency_steps}
33 | use_cache: True
34 | seed: 42
35 | val_ratio: 0.00
36 | max_train_episodes: null
37 |
--------------------------------------------------------------------------------
/diffusion_policy/config/task/umi_teleop_image.yaml:
--------------------------------------------------------------------------------
1 | name: real_image
2 |
3 | image_shape: [3, 224, 224]
4 | dataset_path: data/lift_cube_0921
5 |
6 | shape_meta: &shape_meta
7 | # acceptable types: rgb, low_dim
8 | obs:
9 | camera_0:
10 | shape: ${task.image_shape}
11 | type: rgb
12 | camera_1:
13 | shape: ${task.image_shape}
14 | type: rgb
15 | robot_eef_pose:
16 | shape: [6]
17 | type: low_dim
18 | gripper_position:
19 | shape: [1]
20 | type: low_dim
21 | action:
22 | shape: [7]
23 |
24 | env_runner:
25 | _target_: diffusion_policy.env_runner.real_pusht_image_runner.RealPushTImageRunner
26 |
27 | dataset:
28 | _target_: diffusion_policy.dataset.umi_teleop_image_dataset.UmiTeleopImageDataset
29 | shape_meta: *shape_meta
30 | dataset_path: ${task.dataset_path}
31 | horizon: ${horizon}
32 | pad_before: ${eval:'${n_obs_steps}-1+${n_latency_steps}'}
33 | pad_after: ${eval:'${horizon}-${n_obs_steps}'}
34 | n_obs_steps: ${dataset_obs_steps}
35 | n_latency_steps: ${n_latency_steps}
36 | use_cache: True
37 | seed: 42
38 | val_ratio: 0.00
39 | max_train_episodes: null
40 | delta_action: False
41 |
42 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/block.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/block2.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/blocks/blue_cube.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/blocks/green_star.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/blocks/red_moon.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/blocks/yellow_pentagon.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/insert.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 |
39 |
40 |
41 |
42 |
43 |
44 |
45 |
46 |
47 |
48 |
49 |
50 |
51 |
52 |
53 |
54 |
55 |
56 |
57 |
58 |
59 |
60 |
61 |
62 |
63 |
64 |
65 |
66 |
67 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/plane.obj:
--------------------------------------------------------------------------------
1 | # Blender v2.66 (sub 1) OBJ File: ''
2 | # www.blender.org
3 | mtllib plane.mtl
4 | o Plane
5 | v 15.000000 -15.000000 0.000000
6 | v 15.000000 15.000000 0.000000
7 | v -15.000000 15.000000 0.000000
8 | v -15.000000 -15.000000 0.000000
9 |
10 | vt 15.000000 0.000000
11 | vt 15.000000 15.000000
12 | vt 0.000000 15.000000
13 | vt 0.000000 0.000000
14 |
15 | usemtl Material
16 | s off
17 | f 1/1 2/2 3/3
18 | f 1/1 3/3 4/4
19 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/workspace.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
28 |
29 |
30 |
31 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/workspace_real.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
28 |
29 |
30 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/zone.obj:
--------------------------------------------------------------------------------
1 | # Object Export From Tinkercad Server 2015
2 |
3 | mtllib obj.mtl
4 |
5 | o obj_0
6 | v 10 -10 20
7 | v 10 -10 0
8 | v 10 10 0
9 | v 10 10 20
10 | v 9.002 9.003 20
11 | v 9.002 -9.002 20
12 | v -10 10 0
13 | v -10 10 20
14 | v -9.003 9.003 20
15 | v -9.003 9.003 0
16 | v 9.002 9.003 0
17 | v 9.002 -9.002 0
18 | v -9.003 -9.002 0
19 | v -9.003 -9.002 20
20 | v -10 -10 0
21 | v -10 -10 20
22 | # 16 vertices
23 |
24 | g group_0_15277357
25 |
26 | usemtl color_15277357
27 | s 0
28 |
29 | f 1 2 3
30 | f 1 3 4
31 | f 4 5 6
32 | f 4 6 1
33 | f 9 10 11
34 | f 9 11 5
35 | f 6 12 13
36 | f 6 13 14
37 | f 10 9 14
38 | f 10 14 13
39 | f 7 10 13
40 | f 7 13 15
41 | f 4 8 5
42 | f 9 5 8
43 | f 8 7 15
44 | f 8 15 16
45 | f 10 7 11
46 | f 3 11 7
47 | f 11 3 12
48 | f 2 12 3
49 | f 14 16 6
50 | f 1 6 16
51 | f 16 15 2
52 | f 16 2 1
53 | f 9 8 14
54 | f 16 14 8
55 | f 7 8 3
56 | f 4 3 8
57 | f 2 15 12
58 | f 13 12 15
59 | f 12 6 5
60 | f 12 5 11
61 | # 32 faces
62 |
63 | #end of obj_0
64 |
65 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/zone.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/zone2.urdf:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/oracles/pushing_info.py:
--------------------------------------------------------------------------------
1 | # coding=utf-8
2 | # Copyright 2022 The Reach ML Authors.
3 | #
4 | # Licensed under the Apache License, Version 2.0 (the "License");
5 | # you may not use this file except in compliance with the License.
6 | # You may obtain a copy of the License at
7 | #
8 | # http://www.apache.org/licenses/LICENSE-2.0
9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 |
16 | """Dataclass holding info needed for pushing oracles."""
17 | import dataclasses
18 | from typing import Any
19 |
20 |
21 | @dataclasses.dataclass
22 | class PushingInfo:
23 | """Holds onto info necessary for pushing state machine."""
24 |
25 | xy_block: Any = None
26 | xy_ee: Any = None
27 | xy_pre_block: Any = None
28 | xy_delta_to_nexttoblock: Any = None
29 | xy_delta_to_touchingblock: Any = None
30 | xy_dir_block_to_ee: Any = None
31 | theta_threshold_to_orient: Any = None
32 | theta_threshold_flat_enough: Any = None
33 | theta_error: Any = None
34 | obstacle_poses: Any = None
35 | distance_to_target: Any = None
36 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/__init__.py:
--------------------------------------------------------------------------------
1 | """Environments using kitchen and Franka robot."""
2 | from gym.envs.registration import register
3 |
4 | register(
5 | id="kitchen-microwave-kettle-light-slider-v0",
6 | entry_point="diffusion_policy.env.kitchen.v0:KitchenMicrowaveKettleLightSliderV0",
7 | max_episode_steps=280,
8 | reward_threshold=1.0,
9 | )
10 |
11 | register(
12 | id="kitchen-microwave-kettle-burner-light-v0",
13 | entry_point="diffusion_policy.env.kitchen.v0:KitchenMicrowaveKettleBottomBurnerLightV0",
14 | max_episode_steps=280,
15 | reward_threshold=1.0,
16 | )
17 |
18 | register(
19 | id="kitchen-kettle-microwave-light-slider-v0",
20 | entry_point="diffusion_policy.env.kitchen.v0:KitchenKettleMicrowaveLightSliderV0",
21 | max_episode_steps=280,
22 | reward_threshold=1.0,
23 | )
24 |
25 | register(
26 | id="kitchen-all-v0",
27 | entry_point="diffusion_policy.env.kitchen.v0:KitchenAllV0",
28 | max_episode_steps=280,
29 | reward_threshold=1.0,
30 | )
31 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/kitchen_lowdim_wrapper.py:
--------------------------------------------------------------------------------
1 | from typing import List, Dict, Optional, Optional
2 | import numpy as np
3 | import gym
4 | from gym.spaces import Box
5 | from diffusion_policy.env.kitchen.base import KitchenBase
6 |
7 | class KitchenLowdimWrapper(gym.Env):
8 | def __init__(self,
9 | env: KitchenBase,
10 | init_qpos: Optional[np.ndarray]=None,
11 | init_qvel: Optional[np.ndarray]=None,
12 | render_hw = (240,360)
13 | ):
14 | self.env = env
15 | self.init_qpos = init_qpos
16 | self.init_qvel = init_qvel
17 | self.render_hw = render_hw
18 |
19 | @property
20 | def action_space(self):
21 | return self.env.action_space
22 |
23 | @property
24 | def observation_space(self):
25 | return self.env.observation_space
26 |
27 | def seed(self, seed=None):
28 | return self.env.seed(seed)
29 |
30 | def reset(self):
31 | if self.init_qpos is not None:
32 | # reset anyway to be safe, not very expensive
33 | _ = self.env.reset()
34 | # start from known state
35 | self.env.set_state(self.init_qpos, self.init_qvel)
36 | obs = self.env._get_obs()
37 | return obs
38 | # obs, _, _, _ = self.env.step(np.zeros_like(
39 | # self.action_space.sample()))
40 | # return obs
41 | else:
42 | return self.env.reset()
43 |
44 | def render(self, mode='rgb_array'):
45 | h, w = self.render_hw
46 | return self.env.render(mode=mode, width=w, height=h)
47 |
48 | def step(self, a):
49 | return self.env.step(a)
50 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/kitchen_util.py:
--------------------------------------------------------------------------------
1 | import struct
2 | import numpy as np
3 |
4 | def parse_mjl_logs(read_filename, skipamount):
5 | with open(read_filename, mode='rb') as file:
6 | fileContent = file.read()
7 | headers = struct.unpack('iiiiiii', fileContent[:28])
8 | nq = headers[0]
9 | nv = headers[1]
10 | nu = headers[2]
11 | nmocap = headers[3]
12 | nsensordata = headers[4]
13 | nuserdata = headers[5]
14 | name_len = headers[6]
15 | name = struct.unpack(str(name_len) + 's', fileContent[28:28+name_len])[0]
16 | rem_size = len(fileContent[28 + name_len:])
17 | num_floats = int(rem_size/4)
18 | dat = np.asarray(struct.unpack(str(num_floats) + 'f', fileContent[28+name_len:]))
19 | recsz = 1 + nq + nv + nu + 7*nmocap + nsensordata + nuserdata
20 | if rem_size % recsz != 0:
21 | print("ERROR")
22 | else:
23 | dat = np.reshape(dat, (int(len(dat)/recsz), recsz))
24 | dat = dat.T
25 |
26 | time = dat[0,:][::skipamount] - 0*dat[0, 0]
27 | qpos = dat[1:nq + 1, :].T[::skipamount, :]
28 | qvel = dat[nq+1:nq+nv+1,:].T[::skipamount, :]
29 | ctrl = dat[nq+nv+1:nq+nv+nu+1,:].T[::skipamount,:]
30 | mocap_pos = dat[nq+nv+nu+1:nq+nv+nu+3*nmocap+1,:].T[::skipamount, :]
31 | mocap_quat = dat[nq+nv+nu+3*nmocap+1:nq+nv+nu+7*nmocap+1,:].T[::skipamount, :]
32 | sensordata = dat[nq+nv+nu+7*nmocap+1:nq+nv+nu+7*nmocap+nsensordata+1,:].T[::skipamount,:]
33 | userdata = dat[nq+nv+nu+7*nmocap+nsensordata+1:,:].T[::skipamount,:]
34 |
35 | data = dict(nq=nq,
36 | nv=nv,
37 | nu=nu,
38 | nmocap=nmocap,
39 | nsensordata=nsensordata,
40 | name=name,
41 | time=time,
42 | qpos=qpos,
43 | qvel=qvel,
44 | ctrl=ctrl,
45 | mocap_pos=mocap_pos,
46 | mocap_quat=mocap_quat,
47 | sensordata=sensordata,
48 | userdata=userdata,
49 | logName = read_filename
50 | )
51 | return data
52 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_envs/adept_envs/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/python
2 | #
3 | # Copyright 2020 Google LLC
4 | #
5 | # Licensed under the Apache License, Version 2.0 (the "License");
6 | # you may not use this file except in compliance with the License.
7 | # You may obtain a copy of the License at
8 | #
9 | # http://www.apache.org/licenses/LICENSE-2.0
10 | #
11 | # Unless required by applicable law or agreed to in writing, software
12 | # distributed under the License is distributed on an "AS IS" BASIS,
13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 | # See the License for the specific language governing permissions and
15 | # limitations under the License.
16 |
17 | import adept_envs.franka
18 |
19 | from adept_envs.utils.configurable import global_config
20 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_envs/adept_envs/franka/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/python
2 | #
3 | # Copyright 2020 Google LLC
4 | #
5 | # Licensed under the Apache License, Version 2.0 (the "License");
6 | # you may not use this file except in compliance with the License.
7 | # You may obtain a copy of the License at
8 | #
9 | # http://www.apache.org/licenses/LICENSE-2.0
10 | #
11 | # Unless required by applicable law or agreed to in writing, software
12 | # distributed under the License is distributed on an "AS IS" BASIS,
13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 | # See the License for the specific language governing permissions and
15 | # limitations under the License.
16 |
17 | from gym.envs.registration import register
18 |
19 | # Relax the robot
20 | register(
21 | id='kitchen_relax-v1',
22 | entry_point='adept_envs.franka.kitchen_multitask_v0:KitchenTaskRelaxV1',
23 | max_episode_steps=280,
24 | )
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_envs/adept_envs/utils/constants.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/python
2 | #
3 | # Copyright 2020 Google LLC
4 | #
5 | # Licensed under the Apache License, Version 2.0 (the "License");
6 | # you may not use this file except in compliance with the License.
7 | # You may obtain a copy of the License at
8 | #
9 | # http://www.apache.org/licenses/LICENSE-2.0
10 | #
11 | # Unless required by applicable law or agreed to in writing, software
12 | # distributed under the License is distributed on an "AS IS" BASIS,
13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 | # See the License for the specific language governing permissions and
15 | # limitations under the License.
16 |
17 | import os
18 |
19 | ENVS_ROOT_PATH = os.path.abspath(os.path.join(
20 | os.path.dirname(os.path.abspath(__file__)),
21 | "../../"))
22 |
23 | MODELS_PATH = os.path.abspath(os.path.join(ENVS_ROOT_PATH, "../adept_models/"))
24 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/.gitignore:
--------------------------------------------------------------------------------
1 | # General
2 | .DS_Store
3 | *.swp
4 | *.profraw
5 |
6 | # Editors
7 | .vscode
8 | .idea
9 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/CONTRIBUTING.public.md:
--------------------------------------------------------------------------------
1 | # How to Contribute
2 |
3 | We'd love to accept your patches and contributions to this project. There are
4 | just a few small guidelines you need to follow.
5 |
6 | ## Contributor License Agreement
7 |
8 | Contributions to this project must be accompanied by a Contributor License
9 | Agreement. You (or your employer) retain the copyright to your contribution;
10 | this simply gives us permission to use and redistribute your contributions as
11 | part of the project. Head over to to see
12 | your current agreements on file or to sign a new one.
13 |
14 | You generally only need to submit a CLA once, so if you've already submitted one
15 | (even if it was for a different project), you probably don't need to do it
16 | again.
17 |
18 | ## Code reviews
19 |
20 | All submissions, including submissions by project members, require review. We
21 | use GitHub pull requests for this purpose. Consult
22 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
23 | information on using pull requests.
24 |
25 | ## Community Guidelines
26 |
27 | This project follows
28 | [Google's Open Source Community Guidelines](https://opensource.google.com/conduct/).
29 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/README.public.md:
--------------------------------------------------------------------------------
1 | # D'Suite Scenes
2 |
3 | This repository is based on a collection of [MuJoCo](http://www.mujoco.org/) simulation
4 | scenes and common assets for D'Suite environments. Based on code in the ROBEL suite
5 | https://github.com/google-research/robel
6 |
7 | ## Disclaimer
8 |
9 | This is not an official Google product.
10 |
11 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/__init__.py
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/assets/backwall_asset.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/assets/backwall_chain.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/assets/counters_asset.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
33 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/assets/hingecabinet_asset.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/assets/kettle_asset.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/assets/kettle_chain.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/assets/microwave_asset.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/assets/slidecabinet_asset.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/counters.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/hingecabinet.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/kettle.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/kitchen.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 |
39 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/burnerplate.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/burnerplate.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/burnerplate_mesh.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/burnerplate_mesh.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/cabinetbase.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/cabinetbase.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/cabinetdrawer.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/cabinetdrawer.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/cabinethandle.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/cabinethandle.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/countertop.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/countertop.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/faucet.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/faucet.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/handle2.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/handle2.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/hingecabinet.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/hingecabinet.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/hingedoor.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/hingedoor.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/hingehandle.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/hingehandle.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/hood.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/hood.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/kettle.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/kettle.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/kettlehandle.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/kettlehandle.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/knob.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/knob.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/lightswitch.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/lightswitch.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/lightswitchbase.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/lightswitchbase.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/micro.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/micro.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microbutton.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microbutton.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microdoor.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microdoor.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microefeet.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microefeet.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microfeet.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microfeet.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microhandle.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microhandle.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microwindow.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/microwindow.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/oven.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/oven.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/ovenhandle.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/ovenhandle.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/oventop.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/oventop.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/ovenwindow.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/ovenwindow.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/slidecabinet.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/slidecabinet.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/slidedoor.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/slidedoor.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/stoverim.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/stoverim.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/tile.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/tile.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/wall.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/meshes/wall.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/microwave.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/oven.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/slidecabinet.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/textures/marble1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/textures/marble1.png
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/textures/metal1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/textures/metal1.png
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/textures/tile1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/textures/tile1.png
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/textures/wood1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/kitchen/textures/wood1.png
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/scenes/basic_scene.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/scenes/textures/white_marble_tile.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/scenes/textures/white_marble_tile.png
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/scenes/textures/white_marble_tile2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/adept_models/scenes/textures/white_marble_tile2.png
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/README.md:
--------------------------------------------------------------------------------
1 | # franka
2 | Franka panda mujoco models
3 |
4 |
5 | # Environment
6 |
7 | franka_panda.xml | coming soon
8 | :-------------------------:|:-------------------------:
9 |  | coming soon
10 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/assets/actuator0.xml:
--------------------------------------------------------------------------------
1 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/assets/actuator1.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/assets/basic_scene.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/assets/teleop_actuator.xml:
--------------------------------------------------------------------------------
1 |
2 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/franka_panda.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/franka_panda.png
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/franka_panda.xml:
--------------------------------------------------------------------------------
1 |
4 |
5 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 |
39 |
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/finger.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/finger.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/hand.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/hand.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link0.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link0.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link1.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link1.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link2.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link2.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link3.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link3.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link4.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link4.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link5.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link5.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link6.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link6.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link7.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/collision/link7.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/finger.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/finger.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/hand.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/hand.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link0.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link0.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link1.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link1.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link2.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link2.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link3.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link3.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link4.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link4.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link5.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link5.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link6.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link6.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link7.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/env/kitchen/relay_policy_learning/third_party/franka/meshes/visual/link7.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/kitchen/v0.py:
--------------------------------------------------------------------------------
1 | from diffusion_policy.env.kitchen.base import KitchenBase
2 |
3 |
4 | class KitchenMicrowaveKettleBottomBurnerLightV0(KitchenBase):
5 | TASK_ELEMENTS = ["microwave", "kettle", "bottom burner", "light switch"]
6 | COMPLETE_IN_ANY_ORDER = False
7 |
8 |
9 | class KitchenMicrowaveKettleLightSliderV0(KitchenBase):
10 | TASK_ELEMENTS = ["microwave", "kettle", "light switch", "slide cabinet"]
11 | COMPLETE_IN_ANY_ORDER = False
12 |
13 |
14 | class KitchenKettleMicrowaveLightSliderV0(KitchenBase):
15 | TASK_ELEMENTS = ["kettle", "microwave", "light switch", "slide cabinet"]
16 | COMPLETE_IN_ANY_ORDER = False
17 |
18 |
19 | class KitchenAllV0(KitchenBase):
20 | TASK_ELEMENTS = KitchenBase.ALL_TASKS
21 |
--------------------------------------------------------------------------------
/diffusion_policy/env/pusht/__init__.py:
--------------------------------------------------------------------------------
1 | from gym.envs.registration import register
2 | import diffusion_policy.env.pusht
3 |
4 | register(
5 | id='pusht-keypoints-v0',
6 | entry_point='envs.pusht.pusht_keypoints_env:PushTKeypointsEnv',
7 | max_episode_steps=200,
8 | reward_threshold=1.0
9 | )
--------------------------------------------------------------------------------
/diffusion_policy/env_runner/base_image_runner.py:
--------------------------------------------------------------------------------
1 | from typing import Dict
2 | from diffusion_policy.policy.base_image_policy import BaseImagePolicy
3 |
4 | class BaseImageRunner:
5 | def __init__(self, output_dir):
6 | self.output_dir = output_dir
7 |
8 | def run(self, policy: BaseImagePolicy) -> Dict:
9 | raise NotImplementedError()
10 |
--------------------------------------------------------------------------------
/diffusion_policy/env_runner/base_lowdim_runner.py:
--------------------------------------------------------------------------------
1 | from typing import Dict
2 | from diffusion_policy.policy.base_lowdim_policy import BaseLowdimPolicy
3 |
4 | class BaseLowdimRunner:
5 | def __init__(self, output_dir):
6 | self.output_dir = output_dir
7 |
8 | def run(self, policy: BaseLowdimPolicy) -> Dict:
9 | raise NotImplementedError()
10 |
--------------------------------------------------------------------------------
/diffusion_policy/env_runner/real_pusht_image_runner.py:
--------------------------------------------------------------------------------
1 | from diffusion_policy.policy.base_image_policy import BaseImagePolicy
2 | from diffusion_policy.env_runner.base_image_runner import BaseImageRunner
3 |
4 | class RealPushTImageRunner(BaseImageRunner):
5 | def __init__(self,
6 | output_dir):
7 | super().__init__(output_dir)
8 |
9 | def run(self, policy: BaseImagePolicy):
10 | return dict()
11 |
--------------------------------------------------------------------------------
/diffusion_policy/gym_util/video_recording_wrapper.py:
--------------------------------------------------------------------------------
1 | import gym
2 | import numpy as np
3 | from diffusion_policy.real_world.video_recorder import VideoRecorder
4 |
5 | class VideoRecordingWrapper(gym.Wrapper):
6 | def __init__(self,
7 | env,
8 | video_recoder: VideoRecorder,
9 | mode='rgb_array',
10 | file_path=None,
11 | steps_per_render=1,
12 | **kwargs
13 | ):
14 | """
15 | When file_path is None, don't record.
16 | """
17 | super().__init__(env)
18 |
19 | self.mode = mode
20 | self.render_kwargs = kwargs
21 | self.steps_per_render = steps_per_render
22 | self.file_path = file_path
23 | self.video_recoder = video_recoder
24 |
25 | self.step_count = 0
26 |
27 | def reset(self, **kwargs):
28 | obs = super().reset(**kwargs)
29 | self.frames = list()
30 | self.step_count = 1
31 | self.video_recoder.stop()
32 | return obs
33 |
34 | def step(self, action):
35 | result = super().step(action)
36 | self.step_count += 1
37 | if self.file_path is not None \
38 | and ((self.step_count % self.steps_per_render) == 0):
39 | if not self.video_recoder.is_ready():
40 | self.video_recoder.start(self.file_path)
41 |
42 | frame = self.env.render(
43 | mode=self.mode, **self.render_kwargs)
44 | assert frame.dtype == np.uint8
45 | self.video_recoder.write_frame(frame)
46 | return result
47 |
48 | def render(self, mode='rgb_array', **kwargs):
49 | if self.video_recoder.is_ready():
50 | self.video_recoder.stop()
51 | return self.file_path
52 |
--------------------------------------------------------------------------------
/diffusion_policy/gym_util/video_wrapper.py:
--------------------------------------------------------------------------------
1 | import gym
2 | import numpy as np
3 |
4 | class VideoWrapper(gym.Wrapper):
5 | def __init__(self,
6 | env,
7 | mode='rgb_array',
8 | enabled=True,
9 | steps_per_render=1,
10 | **kwargs
11 | ):
12 | super().__init__(env)
13 |
14 | self.mode = mode
15 | self.enabled = enabled
16 | self.render_kwargs = kwargs
17 | self.steps_per_render = steps_per_render
18 |
19 | self.frames = list()
20 | self.step_count = 0
21 |
22 | def reset(self, **kwargs):
23 | obs = super().reset(**kwargs)
24 | self.frames = list()
25 | self.step_count = 1
26 | if self.enabled:
27 | frame = self.env.render(
28 | mode=self.mode, **self.render_kwargs)
29 | assert frame.dtype == np.uint8
30 | self.frames.append(frame)
31 | return obs
32 |
33 | def step(self, action):
34 | result = super().step(action)
35 | self.step_count += 1
36 | if self.enabled and ((self.step_count % self.steps_per_render) == 0):
37 | frame = self.env.render(
38 | mode=self.mode, **self.render_kwargs)
39 | assert frame.dtype == np.uint8
40 | self.frames.append(frame)
41 | return result
42 |
43 | def render(self, mode='rgb_array', **kwargs):
44 | return self.frames
45 |
--------------------------------------------------------------------------------
/diffusion_policy/model/bet/action_ae/__init__.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | from torch.utils.data import DataLoader
4 | import abc
5 |
6 | from typing import Optional, Union
7 |
8 | import diffusion_policy.model.bet.utils as utils
9 |
10 |
11 | class AbstractActionAE(utils.SaveModule, abc.ABC):
12 | @abc.abstractmethod
13 | def fit_model(
14 | self,
15 | input_dataloader: DataLoader,
16 | eval_dataloader: DataLoader,
17 | obs_encoding_net: Optional[nn.Module] = None,
18 | ) -> None:
19 | pass
20 |
21 | @abc.abstractmethod
22 | def encode_into_latent(
23 | self,
24 | input_action: torch.Tensor,
25 | input_rep: Optional[torch.Tensor],
26 | ) -> torch.Tensor:
27 | """
28 | Given the input action, discretize it.
29 |
30 | Inputs:
31 | input_action (shape: ... x action_dim): The input action to discretize. This can be in a batch,
32 | and is generally assumed that the last dimnesion is the action dimension.
33 |
34 | Outputs:
35 | discretized_action (shape: ... x num_tokens): The discretized action.
36 | """
37 | raise NotImplementedError
38 |
39 | @abc.abstractmethod
40 | def decode_actions(
41 | self,
42 | latent_action_batch: Optional[torch.Tensor],
43 | input_rep_batch: Optional[torch.Tensor] = None,
44 | ) -> torch.Tensor:
45 | """
46 | Given a discretized action, convert it to a continuous action.
47 |
48 | Inputs:
49 | latent_action_batch (shape: ... x num_tokens): The discretized action
50 | generated by the discretizer.
51 |
52 | Outputs:
53 | continuous_action (shape: ... x action_dim): The continuous action.
54 | """
55 | raise NotImplementedError
56 |
57 | @property
58 | @abc.abstractmethod
59 | def num_latents(self) -> Union[int, float]:
60 | """
61 | Number of possible latents for this generator, useful for state priors that use softmax.
62 | """
63 | return float("inf")
64 |
--------------------------------------------------------------------------------
/diffusion_policy/model/bet/libraries/mingpt/LICENSE:
--------------------------------------------------------------------------------
1 | The MIT License (MIT) Copyright (c) 2020 Andrej Karpathy
2 |
3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
4 |
5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
6 |
7 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
8 |
9 |
--------------------------------------------------------------------------------
/diffusion_policy/model/bet/libraries/mingpt/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/diffusion_policy/model/bet/libraries/mingpt/__init__.py
--------------------------------------------------------------------------------
/diffusion_policy/model/bet/libraries/mingpt/utils.py:
--------------------------------------------------------------------------------
1 | import random
2 | import numpy as np
3 | import torch
4 | import torch.nn as nn
5 | from torch.nn import functional as F
6 |
7 |
8 | def set_seed(seed):
9 | random.seed(seed)
10 | np.random.seed(seed)
11 | torch.manual_seed(seed)
12 | torch.cuda.manual_seed_all(seed)
13 |
14 |
15 | def top_k_logits(logits, k):
16 | v, ix = torch.topk(logits, k)
17 | out = logits.clone()
18 | out[out < v[:, [-1]]] = -float("Inf")
19 | return out
20 |
21 |
22 | @torch.no_grad()
23 | def sample(model, x, steps, temperature=1.0, sample=False, top_k=None):
24 | """
25 | take a conditioning sequence of indices in x (of shape (b,t)) and predict the next token in
26 | the sequence, feeding the predictions back into the model each time. Clearly the sampling
27 | has quadratic complexity unlike an RNN that is only linear, and has a finite context window
28 | of block_size, unlike an RNN that has an infinite context window.
29 | """
30 | block_size = model.get_block_size()
31 | model.eval()
32 | for k in range(steps):
33 | x_cond = (
34 | x if x.size(1) <= block_size else x[:, -block_size:]
35 | ) # crop context if needed
36 | logits, _ = model(x_cond)
37 | # pluck the logits at the final step and scale by temperature
38 | logits = logits[:, -1, :] / temperature
39 | # optionally crop probabilities to only the top k options
40 | if top_k is not None:
41 | logits = top_k_logits(logits, top_k)
42 | # apply softmax to convert to probabilities
43 | probs = F.softmax(logits, dim=-1)
44 | # sample from the distribution or take the most likely
45 | if sample:
46 | ix = torch.multinomial(probs, num_samples=1)
47 | else:
48 | _, ix = torch.topk(probs, k=1, dim=-1)
49 | # append to the sequence and continue
50 | x = torch.cat((x, ix), dim=1)
51 |
52 | return x
53 |
--------------------------------------------------------------------------------
/diffusion_policy/model/common/dict_of_tensor_mixin.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 |
4 | class DictOfTensorMixin(nn.Module):
5 | def __init__(self, params_dict=None):
6 | super().__init__()
7 | if params_dict is None:
8 | params_dict = nn.ParameterDict()
9 | self.params_dict = params_dict
10 |
11 | @property
12 | def device(self):
13 | return next(iter(self.parameters())).device
14 |
15 | def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
16 | def dfs_add(dest, keys, value: torch.Tensor):
17 | if len(keys) == 1:
18 | dest[keys[0]] = value
19 | return
20 |
21 | if keys[0] not in dest:
22 | dest[keys[0]] = nn.ParameterDict()
23 | dfs_add(dest[keys[0]], keys[1:], value)
24 |
25 | def load_dict(state_dict, prefix):
26 | out_dict = nn.ParameterDict()
27 | for key, value in state_dict.items():
28 | value: torch.Tensor
29 | if key.startswith(prefix):
30 | param_keys = key[len(prefix):].split('.')[1:]
31 | # if len(param_keys) == 0:
32 | # import pdb; pdb.set_trace()
33 | dfs_add(out_dict, param_keys, value.clone())
34 | return out_dict
35 |
36 | self.params_dict = load_dict(state_dict, prefix + 'params_dict')
37 | self.params_dict.requires_grad_(False)
38 | return
39 |
--------------------------------------------------------------------------------
/diffusion_policy/model/common/lr_scheduler.py:
--------------------------------------------------------------------------------
1 | from diffusers.optimization import (
2 | Union, SchedulerType, Optional,
3 | Optimizer, TYPE_TO_SCHEDULER_FUNCTION
4 | )
5 |
6 | def get_scheduler(
7 | name: Union[str, SchedulerType],
8 | optimizer: Optimizer,
9 | num_warmup_steps: Optional[int] = None,
10 | num_training_steps: Optional[int] = None,
11 | **kwargs
12 | ):
13 | """
14 | Added kwargs vs diffuser's original implementation
15 |
16 | Unified API to get any scheduler from its name.
17 |
18 | Args:
19 | name (`str` or `SchedulerType`):
20 | The name of the scheduler to use.
21 | optimizer (`torch.optim.Optimizer`):
22 | The optimizer that will be used during training.
23 | num_warmup_steps (`int`, *optional*):
24 | The number of warmup steps to do. This is not required by all schedulers (hence the argument being
25 | optional), the function will raise an error if it's unset and the scheduler type requires it.
26 | num_training_steps (`int``, *optional*):
27 | The number of training steps to do. This is not required by all schedulers (hence the argument being
28 | optional), the function will raise an error if it's unset and the scheduler type requires it.
29 | """
30 | name = SchedulerType(name)
31 | schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
32 | if name == SchedulerType.CONSTANT:
33 | return schedule_func(optimizer, **kwargs)
34 |
35 | # All other schedulers require `num_warmup_steps`
36 | if num_warmup_steps is None:
37 | raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
38 |
39 | if name == SchedulerType.CONSTANT_WITH_WARMUP:
40 | return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, **kwargs)
41 |
42 | # All other schedulers require `num_training_steps`
43 | if num_training_steps is None:
44 | raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
45 |
46 | return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, **kwargs)
47 |
--------------------------------------------------------------------------------
/diffusion_policy/model/common/module_attr_mixin.py:
--------------------------------------------------------------------------------
1 | import torch.nn as nn
2 |
3 | class ModuleAttrMixin(nn.Module):
4 | def __init__(self):
5 | super().__init__()
6 | self._dummy_variable = nn.Parameter(requires_grad=False)
7 |
8 | @property
9 | def device(self):
10 | return next(iter(self.parameters())).device
11 |
12 | @property
13 | def dtype(self):
14 | return next(iter(self.parameters())).dtype
15 |
--------------------------------------------------------------------------------
/diffusion_policy/model/common/shape_util.py:
--------------------------------------------------------------------------------
1 | from typing import Dict, List, Tuple, Callable
2 | import torch
3 | import torch.nn as nn
4 |
5 | def get_module_device(m: nn.Module):
6 | device = torch.device('cpu')
7 | try:
8 | param = next(iter(m.parameters()))
9 | device = param.device
10 | except StopIteration:
11 | pass
12 | return device
13 |
14 | @torch.no_grad()
15 | def get_output_shape(
16 | input_shape: Tuple[int],
17 | net: Callable[[torch.Tensor], torch.Tensor]
18 | ):
19 | device = get_module_device(net)
20 | test_input = torch.zeros((1,)+tuple(input_shape), device=device)
21 | test_output = net(test_input)
22 | output_shape = tuple(test_output.shape[1:])
23 | return output_shape
24 |
--------------------------------------------------------------------------------
/diffusion_policy/model/diffusion/conv1d_components.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 | # from einops.layers.torch import Rearrange
5 |
6 |
7 | class Downsample1d(nn.Module):
8 | def __init__(self, dim):
9 | super().__init__()
10 | self.conv = nn.Conv1d(dim, dim, 3, 2, 1)
11 |
12 | def forward(self, x):
13 | return self.conv(x)
14 |
15 | class Upsample1d(nn.Module):
16 | def __init__(self, dim):
17 | super().__init__()
18 | self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1)
19 |
20 | def forward(self, x):
21 | return self.conv(x)
22 |
23 | class Conv1dBlock(nn.Module):
24 | '''
25 | Conv1d --> GroupNorm --> Mish
26 | '''
27 |
28 | def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8):
29 | super().__init__()
30 |
31 | self.block = nn.Sequential(
32 | nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2),
33 | # Rearrange('batch channels horizon -> batch channels 1 horizon'),
34 | nn.GroupNorm(n_groups, out_channels),
35 | # Rearrange('batch channels 1 horizon -> batch channels horizon'),
36 | nn.Mish(),
37 | )
38 |
39 | def forward(self, x):
40 | return self.block(x)
41 |
42 |
43 | def test():
44 | cb = Conv1dBlock(256, 128, kernel_size=3)
45 | x = torch.zeros((1,256,16))
46 | o = cb(x)
47 |
--------------------------------------------------------------------------------
/diffusion_policy/model/diffusion/positional_embedding.py:
--------------------------------------------------------------------------------
1 | import math
2 | import torch
3 | import torch.nn as nn
4 |
5 | class SinusoidalPosEmb(nn.Module):
6 | def __init__(self, dim):
7 | super().__init__()
8 | self.dim = dim
9 |
10 | def forward(self, x):
11 | device = x.device
12 | half_dim = self.dim // 2
13 | emb = math.log(10000) / (half_dim - 1)
14 | emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
15 | emb = x[:, None] * emb[None, :]
16 | emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
17 | return emb
18 |
--------------------------------------------------------------------------------
/diffusion_policy/model/vision/choice_randomizer.py:
--------------------------------------------------------------------------------
1 | import random
2 | from collections.abc import Sequence
3 | import torch.nn as nn
4 |
5 |
6 | class RandomChoice(nn.Module):
7 | """Apply a single transformation randomly picked from a list. This transform does not support torchscript."""
8 |
9 | def __init__(self, transforms, p=None):
10 | super().__init__()
11 | self.transforms = transforms
12 | if p is not None and not isinstance(p, Sequence):
13 | raise TypeError("Argument p should be a sequence")
14 | self.p = p
15 |
16 | def forward(self, *args):
17 | t = random.choices(self.transforms, weights=self.p)[0]
18 | return t(*args)
19 |
20 | def __repr__(self) -> str:
21 | return f"{self.__class__.__name__}(transforms={self.transforms}, p={self.p})"
--------------------------------------------------------------------------------
/diffusion_policy/model/vision/model_getter.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torchvision
3 |
4 | def get_resnet(name, weights=None, **kwargs):
5 | """
6 | name: resnet18, resnet34, resnet50
7 | weights: "IMAGENET1K_V1", "r3m"
8 | """
9 | # load r3m weights
10 | if (weights == "r3m") or (weights == "R3M"):
11 | return get_r3m(name=name, **kwargs)
12 |
13 | func = getattr(torchvision.models, name)
14 | resnet = func(weights=weights, **kwargs)
15 | resnet.fc = torch.nn.Identity()
16 | return resnet
17 |
18 | def get_r3m(name, **kwargs):
19 | """
20 | name: resnet18, resnet34, resnet50
21 | """
22 | import r3m
23 | r3m.device = 'cpu'
24 | model = r3m.load_r3m(name)
25 | r3m_model = model.module
26 | resnet_model = r3m_model.convnet
27 | resnet_model = resnet_model.to('cpu')
28 | return resnet_model
29 |
--------------------------------------------------------------------------------
/diffusion_policy/policy/base_image_policy.py:
--------------------------------------------------------------------------------
1 | from typing import Dict
2 | import torch
3 | import torch.nn as nn
4 | from diffusion_policy.model.common.module_attr_mixin import ModuleAttrMixin
5 | from diffusion_policy.model.common.normalizer import LinearNormalizer
6 |
7 | class BaseImagePolicy(ModuleAttrMixin):
8 | # init accepts keyword argument shape_meta, see config/task/*_image.yaml
9 |
10 | def predict_action(self, obs_dict: Dict[str, torch.Tensor], fixed_action_prefix: torch.Tensor=None) -> Dict[str, torch.Tensor]:
11 | """
12 | obs_dict:
13 | str: B,To,*
14 | fixed_action_prefix:
15 | B, Tp, Da
16 | return: B,Ta,Da
17 | """
18 | raise NotImplementedError()
19 |
20 | # reset state for stateful policies
21 | def reset(self):
22 | pass
23 |
24 | # ========== training ===========
25 | # no standard training interface except setting normalizer
26 | def set_normalizer(self, normalizer: LinearNormalizer):
27 | raise NotImplementedError()
28 |
--------------------------------------------------------------------------------
/diffusion_policy/policy/base_lowdim_policy.py:
--------------------------------------------------------------------------------
1 | from typing import Dict
2 | import torch
3 | import torch.nn as nn
4 | from diffusion_policy.model.common.module_attr_mixin import ModuleAttrMixin
5 | from diffusion_policy.model.common.normalizer import LinearNormalizer
6 |
7 | class BaseLowdimPolicy(ModuleAttrMixin):
8 | # ========= inference ============
9 | # also as self.device and self.dtype for inference device transfer
10 | def predict_action(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
11 | """
12 | obs_dict:
13 | obs: B,To,Do
14 | return:
15 | action: B,Ta,Da
16 | To = 3
17 | Ta = 4
18 | T = 6
19 | |o|o|o|
20 | | | |a|a|a|a|
21 | |o|o|
22 | | |a|a|a|a|a|
23 | | | | | |a|a|
24 | """
25 | raise NotImplementedError()
26 |
27 | # reset state for stateful policies
28 | def reset(self):
29 | pass
30 |
31 | # ========== training ===========
32 | # no standard training interface except setting normalizer
33 | def set_normalizer(self, normalizer: LinearNormalizer):
34 | raise NotImplementedError()
35 |
36 |
--------------------------------------------------------------------------------
/diffusion_policy/real_world/keystroke_counter.py:
--------------------------------------------------------------------------------
1 | from pynput.keyboard import Key, KeyCode, Listener
2 | from collections import defaultdict
3 | from threading import Lock
4 |
5 | class KeystrokeCounter(Listener):
6 | def __init__(self):
7 | self.key_count_map = defaultdict(lambda:0)
8 | self.key_press_list = list()
9 | self.lock = Lock()
10 | super().__init__(on_press=self.on_press, on_release=self.on_release)
11 |
12 | def on_press(self, key):
13 | with self.lock:
14 | self.key_count_map[key] += 1
15 | self.key_press_list.append(key)
16 |
17 | def on_release(self, key):
18 | pass
19 |
20 | def clear(self):
21 | with self.lock:
22 | self.key_count_map = defaultdict(lambda:0)
23 | self.key_press_list = list()
24 |
25 | def __getitem__(self, key):
26 | with self.lock:
27 | return self.key_count_map[key]
28 |
29 | def get_press_events(self):
30 | with self.lock:
31 | events = list(self.key_press_list)
32 | self.key_press_list = list()
33 | return events
34 |
35 | if __name__ == '__main__':
36 | import time
37 | with KeystrokeCounter() as counter:
38 | try:
39 | while True:
40 | print('Space:', counter[Key.space])
41 | print('q:', counter[KeyCode(char='q')])
42 | time.sleep(1/60)
43 | except KeyboardInterrupt:
44 | events = counter.get_press_events()
45 | print(events)
46 |
--------------------------------------------------------------------------------
/diffusion_policy/real_world/real_inference_util.py:
--------------------------------------------------------------------------------
1 | c
2 | def get_real_obs_dict(
3 | env_obs: Dict[str, np.ndarray],
4 | shape_meta: dict,
5 | ) -> Dict[str, np.ndarray]:
6 | obs_dict_np = dict()
7 | obs_shape_meta = shape_meta['obs']
8 | for key, attr in obs_shape_meta.items():
9 | type = attr.get('type', 'low_dim')
10 | shape = attr.get('shape')
11 | if type == 'rgb':
12 | this_imgs_in = env_obs[key]
13 | t,hi,wi,ci = this_imgs_in.shape
14 | co,ho,wo = shape
15 | assert ci == co
16 | out_imgs = this_imgs_in
17 | if (ho != hi) or (wo != wi) or (this_imgs_in.dtype == np.uint8):
18 | tf = get_image_transform(
19 | input_res=(wi,hi),
20 | output_res=(wo,ho),
21 | bgr_to_rgb=False)
22 | out_imgs = np.stack([tf(x) for x in this_imgs_in])
23 | if this_imgs_in.dtype == np.uint8:
24 | out_imgs = out_imgs.astype(np.float32) / 255
25 | # THWC to TCHW
26 | obs_dict_np[key] = np.moveaxis(out_imgs,-1,1)
27 | elif type == 'low_dim':
28 | this_data_in = env_obs[key]
29 | if 'pose' in key and shape == (2,):
30 | # take X,Y coordinates
31 | this_data_in = this_data_in[...,[0,1]]
32 | obs_dict_np[key] = this_data_in
33 | return obs_dict_np
34 |
35 |
36 | def get_real_obs_resolution(
37 | shape_meta: dict
38 | ) -> Tuple[int, int]:
39 | out_res = None
40 | obs_shape_meta = shape_meta['obs']
41 | for key, attr in obs_shape_meta.items():
42 | type = attr.get('type', 'low_dim')
43 | shape = attr.get('shape')
44 | if type == 'rgb':
45 | co,ho,wo = shape
46 | if out_res is None:
47 | out_res = (wo, ho)
48 | assert out_res == (wo, ho)
49 | return out_res
50 |
--------------------------------------------------------------------------------
/diffusion_policy/scripts/bet_blockpush_conversion.py:
--------------------------------------------------------------------------------
1 | if __name__ == "__main__":
2 | import sys
3 | import os
4 | import pathlib
5 |
6 | ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
7 | sys.path.append(ROOT_DIR)
8 |
9 |
10 | import os
11 | import click
12 | import pathlib
13 | import numpy as np
14 | from diffusion_policy.common.replay_buffer import ReplayBuffer
15 |
16 | @click.command()
17 | @click.option('-i', '--input', required=True, help='input dir contains npy files')
18 | @click.option('-o', '--output', required=True, help='output zarr path')
19 | @click.option('--abs_action', is_flag=True, default=False)
20 | def main(input, output, abs_action):
21 | data_directory = pathlib.Path(input)
22 | observations = np.load(
23 | data_directory / "multimodal_push_observations.npy"
24 | )
25 | actions = np.load(data_directory / "multimodal_push_actions.npy")
26 | masks = np.load(data_directory / "multimodal_push_masks.npy")
27 |
28 | buffer = ReplayBuffer.create_empty_numpy()
29 | for i in range(len(masks)):
30 | eps_len = int(masks[i].sum())
31 | obs = observations[i,:eps_len].astype(np.float32)
32 | action = actions[i,:eps_len].astype(np.float32)
33 | if abs_action:
34 | prev_eef_target = obs[:,8:10]
35 | next_eef_target = prev_eef_target + action
36 | action = next_eef_target
37 | data = {
38 | 'obs': obs,
39 | 'action': action
40 | }
41 | buffer.add_episode(data)
42 |
43 | buffer.save_to_path(zarr_path=output, chunk_length=-1)
44 |
45 | if __name__ == '__main__':
46 | main()
47 |
--------------------------------------------------------------------------------
/diffusion_policy/scripts/blockpush_abs_conversion.py:
--------------------------------------------------------------------------------
1 | if __name__ == "__main__":
2 | import sys
3 | import os
4 | import pathlib
5 |
6 | ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
7 | sys.path.append(ROOT_DIR)
8 |
9 | import os
10 | import click
11 | import pathlib
12 | from diffusion_policy.common.replay_buffer import ReplayBuffer
13 |
14 |
15 | @click.command()
16 | @click.option('-i', '--input', required=True)
17 | @click.option('-o', '--output', required=True)
18 | @click.option('-t', '--target_eef_idx', default=8, type=int)
19 | def main(input, output, target_eef_idx):
20 | buffer = ReplayBuffer.copy_from_path(input)
21 | obs = buffer['obs']
22 | action = buffer['action']
23 | prev_eef_target = obs[:,target_eef_idx:target_eef_idx+action.shape[1]]
24 | next_eef_target = prev_eef_target + action
25 | action[:] = next_eef_target
26 | buffer.save_to_path(zarr_path=output, chunk_length=-1)
27 |
28 | if __name__ == '__main__':
29 | main()
30 |
--------------------------------------------------------------------------------
/diffusion_policy/scripts/episode_lengths.py:
--------------------------------------------------------------------------------
1 | if __name__ == "__main__":
2 | import sys
3 | import os
4 | import pathlib
5 |
6 | ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
7 | sys.path.append(ROOT_DIR)
8 |
9 | import click
10 | import numpy as np
11 | import json
12 | from diffusion_policy.common.replay_buffer import ReplayBuffer
13 |
14 | @click.command()
15 | @click.option('--input', '-i', required=True)
16 | @click.option('--dt', default=0.1, type=float)
17 | def main(input, dt):
18 | buffer = ReplayBuffer.create_from_path(input)
19 | lengths = buffer.episode_lengths
20 | durations = lengths * dt
21 | result = {
22 | 'duration/mean': np.mean(durations)
23 | }
24 |
25 | text = json.dumps(result, indent=2)
26 | print(text)
27 |
28 | if __name__ == '__main__':
29 | main()
30 |
--------------------------------------------------------------------------------
/diffusion_policy/scripts/robomimic_dataset_action_comparison.py:
--------------------------------------------------------------------------------
1 | if __name__ == "__main__":
2 | import sys
3 | import os
4 | import pathlib
5 |
6 | ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
7 | sys.path.append(ROOT_DIR)
8 |
9 | import os
10 | import click
11 | import pathlib
12 | import h5py
13 | import numpy as np
14 | from tqdm import tqdm
15 | from scipy.spatial.transform import Rotation
16 |
17 | def read_all_actions(hdf5_file, metric_skip_steps=1):
18 | n_demos = len(hdf5_file['data'])
19 | all_actions = list()
20 | for i in tqdm(range(n_demos)):
21 | actions = hdf5_file[f'data/demo_{i}/actions'][:]
22 | all_actions.append(actions[metric_skip_steps:])
23 | all_actions = np.concatenate(all_actions, axis=0)
24 | return all_actions
25 |
26 |
27 | @click.command()
28 | @click.option('-i', '--input', required=True, help='input hdf5 path')
29 | @click.option('-o', '--output', required=True, help='output hdf5 path. Parent directory must exist')
30 | def main(input, output):
31 | # process inputs
32 | input = pathlib.Path(input).expanduser()
33 | assert input.is_file()
34 | output = pathlib.Path(output).expanduser()
35 | assert output.is_file()
36 |
37 | input_file = h5py.File(str(input), 'r')
38 | output_file = h5py.File(str(output), 'r')
39 |
40 | input_all_actions = read_all_actions(input_file)
41 | output_all_actions = read_all_actions(output_file)
42 | pos_dist = np.linalg.norm(input_all_actions[:,:3] - output_all_actions[:,:3], axis=-1)
43 | rot_dist = (Rotation.from_rotvec(input_all_actions[:,3:6]
44 | ) * Rotation.from_rotvec(output_all_actions[:,3:6]).inv()
45 | ).magnitude()
46 |
47 | print(f'max pos dist: {pos_dist.max()}')
48 | print(f'max rot dist: {rot_dist.max()}')
49 |
50 | if __name__ == "__main__":
51 | main()
52 |
--------------------------------------------------------------------------------
/diffusion_policy/shared_memory/shared_memory_util.py:
--------------------------------------------------------------------------------
1 | from typing import Tuple
2 | from dataclasses import dataclass
3 | import numpy as np
4 | from multiprocessing.managers import SharedMemoryManager
5 | from atomics import atomicview, MemoryOrder, UINT
6 |
7 | @dataclass
8 | class ArraySpec:
9 | name: str
10 | shape: Tuple[int]
11 | dtype: np.dtype
12 |
13 |
14 | class SharedAtomicCounter:
15 | def __init__(self,
16 | shm_manager: SharedMemoryManager,
17 | size :int=8 # 64bit int
18 | ):
19 | shm = shm_manager.SharedMemory(size=size)
20 | self.shm = shm
21 | self.size = size
22 | self.store(0) # initialize
23 |
24 | @property
25 | def buf(self):
26 | return self.shm.buf[:self.size]
27 |
28 | def load(self) -> int:
29 | with atomicview(buffer=self.buf, atype=UINT) as a:
30 | value = a.load(order=MemoryOrder.ACQUIRE)
31 | return value
32 |
33 | def store(self, value: int):
34 | with atomicview(buffer=self.buf, atype=UINT) as a:
35 | a.store(value, order=MemoryOrder.RELEASE)
36 |
37 | def add(self, value: int):
38 | with atomicview(buffer=self.buf, atype=UINT) as a:
39 | a.add(value, order=MemoryOrder.ACQ_REL)
40 |
--------------------------------------------------------------------------------
/eval.sh:
--------------------------------------------------------------------------------
1 | task_name="pour_water" # or "arrange_mouse" "fold_towel" "unplug_charger"
2 |
3 | python eval_real.py \
4 | --robot_config=example/eval_robots_config.yaml \
5 | -i data/checkpoints/${task_name}/latest.ckpt \
6 | -o data/eval_${task_name} \
7 | --frequency 5 \
8 | --temporal_agg -si 1
--------------------------------------------------------------------------------
/example/calibration/aruco_config.yaml:
--------------------------------------------------------------------------------
1 | aruco_dict:
2 | predefined: DICT_4X4_50
3 | marker_size_map:
4 | default: 0.16
5 | 0: 0.016
6 | 1: 0.016
7 | 2: 0.06
8 | 3: 0.06
9 | 4: 0.06
10 | 5: 0.06
11 | 6: 0.016
12 | 7: 0.016
13 | 8: 0.06
14 | 9: 0.06
15 | 10: 0.06
16 | 11: 0.06
17 | 12: 0.16
18 | 13: 0.16
19 | 14: 0.16
20 | 15: 0.16
21 | 16: 0.16
22 | 17: 0.16
23 | 18: 0.16
24 | 19: 0.16
25 |
--------------------------------------------------------------------------------
/example/calibration/gopro_intrinsics_2_7k.json:
--------------------------------------------------------------------------------
1 | {
2 | "final_reproj_error": 0.2916398582648,
3 | "fps": 59.94005994005994,
4 | "image_height": 2028,
5 | "image_width": 2704,
6 | "intrinsic_type": "FISHEYE",
7 | "intrinsics": {
8 | "aspect_ratio": 1.0029788958491257,
9 | "focal_length": 796.8544625226342,
10 | "principal_pt_x": 1354.4265245977356,
11 | "principal_pt_y": 1011.4847310011687,
12 | "radial_distortion_1": -0.02196117964405394,
13 | "radial_distortion_2": -0.018959717016668237,
14 | "radial_distortion_3": 0.001693880829392453,
15 | "radial_distortion_4": -0.00016807228608000285,
16 | "skew": 0.0
17 | },
18 | "nr_calib_images": 59,
19 | "stabelized": false
20 | }
--------------------------------------------------------------------------------
/example/eval_robots_config.yaml:
--------------------------------------------------------------------------------
1 | {
2 | # robot configurations, 0->right 1->left
3 | # comment out robot #1 if only using 1 arm
4 | "robots": [
5 | {
6 | "robot_type": "franka",
7 | "robot_ip": "172.16.0.1",
8 | "robot_obs_latency": 0.0001, "robot_action_latency": 0.1, "tcp_offset": 0.235,
9 | # table height for finger-table collision avoidance
10 | # -inf to disable
11 | "height_threshold": -1,
12 | # sphere collision primitive size and location for inter-gripper collision avoidance
13 | # 0 to disable
14 | "sphere_radius": 0.1, "sphere_center": [0, -0.06, -0.185],
15 | "joints_init": [0.00010400846076663584, -0.24481917917728424, -0.0936817154288292, -2.7792861461639404, -0.0670328289270401, 1.998847246170044, 0.7243571877479553]
16 | },
17 | # {
18 | # "robot_type": "ur5",
19 | # "robot_ip": "192.168.0.9",
20 | # "robot_obs_latency": 0.0001, "robot_action_latency": 0.1, "tcp_offset": 0.235,
21 | # "height_threshold": 0.022,
22 | # "sphere_radius": 0.1, "sphere_center": [0, -0.06, -0.185]
23 | # }
24 | ],
25 | # gripper configurations, 0->right 1->left
26 | "grippers": [
27 | {
28 | "gripper_ip": "192.168.1.20", # wsg: 192.168.0.18
29 | "gripper_port": 1000, "gripper_obs_latency": 0.01, "gripper_action_latency": 0.1
30 | },
31 | # {
32 | # "gripper_ip": "192.168.0.27",
33 | # "gripper_port": 1000, "gripper_obs_latency": 0.01, "gripper_action_latency": 0.1
34 | # }
35 | ],
36 | # transformation right arm's base frame in left arm's base coordinate frame
37 | "tx_left_right": [
38 | [ 0.99996206, 0.00661996, 0.00566226, -0.01676012],
39 | [-0.00663261, 0.99997554, 0.0022186 , -0.60552492],
40 | [-0.00564743, -0.00225607, 0.99998151, -0.007277 ],
41 | [ 0. , 0. , 0. , 1. ]
42 | ]
43 | }
44 |
--------------------------------------------------------------------------------
/franka_instruction.md:
--------------------------------------------------------------------------------
1 | # Deploy UMI Policy on Franka Robot
2 |
3 | ## System Overview
4 | 
5 |
6 | * FrankaInterface (Server): [scripts_real/launch_franka_interface_server.py](scripts_real/launch_franka_interface_server.py) (L7)
7 |
8 | * FrankaInterface (Client): [umi/real_world/franka_interpolation_controller.py](umi/real_world/franka_interpolation_controller.py) (L36)
9 |
10 | * FrankaInterpolationController: [umi/real_world/franka_interpolation_controller.py](umi/real_world/franka_interpolation_controller.py) (L71)
11 |
12 |
13 | ## Instructions
14 | * Follow the [documentation](https://facebookresearch.github.io/fairo/polymetis/prereq.html#franka-panda-hardware-setup) to install *Polymetis* on a computer with realtime kernel (i.e., NUC in our example).
15 | * Launch FrankaInterface Server on NUC.
16 |
17 | `python scripts_real/launch_franka_interface_server.py`
18 | * (optional) Now you should be able to control the Franka arm using a space mouse on another desktop, the one you are going to run the robot policy on.
19 |
20 | `python scripts_real/control_franka.py`
21 | * Change [eval_robots_config.yaml](example/eval_robots_config.yaml)/robots/robot_type to `'franka'`. This will change the robot controller in [umi_env](umi/real_world/bimanual_umi_env.py) (L233).
22 |
23 | * Obtain IP address and update [eval_robots_config.yaml](example/eval_robots_config.yaml)/robots/robot_ip.
24 | * On the Franka interface website
25 | * Set mass to 1.8 kg
26 | * Set Flange to Center of Mass of Load Vector to (0.064, -0.06, 0.03)m.
27 |
28 | * Then you should be able to launch the evaluation on the franka arm.
29 |
30 | `python eval_real.py --robot_config=example/eval_robots_config.yaml -i cup_wild_vit_l.ckpt -o data/eval_cup_wild_example`
--------------------------------------------------------------------------------
/run_generate_dataset.sh:
--------------------------------------------------------------------------------
1 | task_name="pour_water" # or "arrange_mouse" "fold_towel" "unplug_charger"
2 |
3 | python scripts_slam_pipeline/07_generate_replay_buffer.py \
4 | -o data/dataset/${task_name}/dataset.zarr.zip \
5 | data/dataset/${task_name}/
--------------------------------------------------------------------------------
/run_slam.sh:
--------------------------------------------------------------------------------
1 | task_name="pour_water" # or "arrange_mouse" "fold_towel" "unplug_charger"
2 |
3 | python run_slam_pipeline.py data/dataset/${task_name}/
4 | sleep 3s
5 | docker stop $(docker ps -aq)
6 | sleep 3s
7 |
--------------------------------------------------------------------------------
/scripts/calibrate_timecode_offset.py:
--------------------------------------------------------------------------------
1 | # %%
2 | import sys
3 | import os
4 |
5 | ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
6 | sys.path.append(ROOT_DIR)
7 | os.chdir(ROOT_DIR)
8 |
9 | # %%
10 | import click
11 | import os
12 | import av
13 | import cv2
14 | import datetime
15 | from tqdm import tqdm
16 | import json
17 | from umi.common.timecode_util import stream_get_start_datetime
18 |
19 | # %%
20 | def parse_qr_datetime(qr_txt):
21 | return datetime.datetime.strptime(qr_txt, r"#%y%m%d%H%M%S.%f")
22 |
23 | # %%
24 | @click.command()
25 | @click.option('-i', '--input', required=True, help='GoPro MP4 file')
26 | def main(input):
27 | input = os.path.expanduser(input)
28 |
29 | # find the first QR code timestamp
30 | detector = cv2.QRCodeDetector()
31 | qr_datetime = None
32 | tc_datetime = None
33 | with av.open(input) as container:
34 | stream = container.streams.video[0]
35 | tc_datetime = stream_get_start_datetime(stream=stream)
36 | for i, frame in tqdm(enumerate(container.decode(stream))):
37 | img = frame.to_ndarray(format='rgb24')
38 | frame_cts_sec = frame.pts * stream.time_base
39 |
40 | qr_txt, _, _ =detector.detectAndDecodeCurved(img)
41 | if qr_txt.startswith('#'):
42 | # found
43 | qr_datetime = parse_qr_datetime(qr_txt) \
44 | - datetime.timedelta(seconds=float(frame_cts_sec))
45 | break
46 |
47 | if not qr_datetime:
48 | raise RuntimeError("No valid QR code found.")
49 |
50 | dt = (qr_datetime - tc_datetime).total_seconds()
51 |
52 | print("time = date + timecode + dt")
53 | print(f"dt = {dt}")
54 |
55 | if __name__ == '__main__':
56 | main()
57 |
58 |
--------------------------------------------------------------------------------
/scripts/check_gopro_orientation.py:
--------------------------------------------------------------------------------
1 | # %%
2 | import sys
3 | import os
4 |
5 | ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
6 | sys.path.append(ROOT_DIR)
7 | os.chdir(ROOT_DIR)
8 |
9 | # %%
10 | import click
11 | import json
12 | import pathlib
13 | import numpy as np
14 | from tqdm import tqdm
15 | from exiftool import ExifToolHelper
16 |
17 | # %%
18 | @click.command(help='Session directories. Assumming mp4 videos are in /raw_videos')
19 | @click.argument('session_dir', nargs=-1)
20 | def main(session_dir):
21 | with ExifToolHelper() as et:
22 | for session in tqdm(session_dir):
23 | session = pathlib.Path(os.path.expanduser(session)).absolute()
24 | # hardcode subdirs
25 | demos_dir = session.joinpath('demos')
26 |
27 | mp4_paths = list(demos_dir.glob("demo_*/raw_video.mp4"))
28 | if len(mp4_paths) < 1:
29 | continue
30 | all_meta = et.get_tags(mp4_paths, ['QuickTime:AutoRotation'])
31 | for mp4_path, meta in zip(mp4_paths, all_meta):
32 | rot = meta['QuickTime:AutoRotation']
33 | if rot != 'U':
34 | demo_dir = mp4_path.parent
35 | print(f"Found rotated video: {session.name} {demo_dir.name}")
36 |
37 | if __name__ == '__main__':
38 | if len(sys.argv) == 1:
39 | main.main(['--help'])
40 | else:
41 | main()
42 |
--------------------------------------------------------------------------------
/scripts/gen_charuco_board.py:
--------------------------------------------------------------------------------
1 | # %%
2 | import sys
3 | import os
4 |
5 | ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
6 | sys.path.append(ROOT_DIR)
7 | os.chdir(ROOT_DIR)
8 |
9 | # %%
10 | import click
11 | from PIL import Image
12 | from umi.common.cv_util import get_charuco_board, draw_charuco_board
13 |
14 | # %%
15 | @click.command()
16 | @click.option('-o', '--output', required=True, help='Output pdf path')
17 | @click.option('-to', '--tag_id_offset', type=int, default=50)
18 | def main(output, tag_id_offset):
19 | dpi = 300
20 | board = get_charuco_board(tag_id_offset=tag_id_offset)
21 | board_img = draw_charuco_board(board=board, dpi=dpi)
22 | im = Image.fromarray(board_img)
23 | im.save(output, resolution=dpi)
24 |
25 | # %%
26 | if __name__ == '__main__':
27 | main()
28 |
--------------------------------------------------------------------------------
/scripts/gen_image_mask.py:
--------------------------------------------------------------------------------
1 | # %%
2 | import sys
3 | import os
4 |
5 | ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
6 | sys.path.append(ROOT_DIR)
7 | os.chdir(ROOT_DIR)
8 |
9 | # %%
10 | import click
11 | import imageio
12 | import pickle
13 | import json
14 | import numpy as np
15 | from umi.common.cv_util import get_gripper_with_finger_mask, get_mirror_mask
16 |
17 | # %%
18 | @click.command()
19 | @click.option('-c', '--config', required=False, default='umi/asset/mask.json', help='Mask config path')
20 | @click.option('-o', '--output', required=True, help='Output png path')
21 | @click.option('-h', '--height', required=False, default=2028, help='Image height')
22 | @click.option('-w', '--width', required=False, default=2704, help='Image width')
23 | def main(config, output, height, width):
24 | json_data = json.load(open(config, 'r'))
25 | mask_image = np.zeros([height, width, 3], dtype=np.uint8)
26 | mask_image = get_mirror_mask(json_data, mask_image, (255, 255, 255))
27 | mask_image = get_gripper_with_finger_mask(mask_image, color=(255, 255, 255))
28 |
29 | imageio.imwrite(output, mask_image)
30 |
31 | # # load gripper image
32 | # pkl_path = '/local/real/datasets/umi/toss_orange_20231018/calibration/robot_world_hand_eye.pkl'
33 | # pkl_data = pickle.load(open(pkl_path, 'rb'))
34 | # img_gripper = pkl_data[0]['img']
35 | # img_gripper = (img_gripper * 0.6 + img_gripper * (mask_image/255.0) * 0.4).astype(np.uint8)
36 | # imageio.imwrite('missor_mask_new.png', img_gripper)
37 | # %%
38 | if __name__ == '__main__':
39 | main()
40 |
--------------------------------------------------------------------------------
/scripts/uvc_camera_test.py:
--------------------------------------------------------------------------------
1 | # %%
2 | import sys
3 | import os
4 |
5 | ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
6 | sys.path.append(ROOT_DIR)
7 | os.chdir(ROOT_DIR)
8 |
9 | # %%
10 | import cv2
11 | import time
12 | from multiprocessing.managers import SharedMemoryManager
13 | from umi.real_world.uvc_camera import UvcCamera
14 | from umi.common.usb_util import create_usb_list
15 | from umi.common.precise_sleep import precise_wait
16 |
17 | # %%
18 | def main():
19 | cv2.setNumThreads(1)
20 |
21 | dev_video_path = '/dev/video0'
22 |
23 | # enumerate UBS device to find Elgato Capture Card
24 | device_list = create_usb_list()
25 | dev_usb_path = None
26 | for dev in device_list:
27 | if 'Elgato' in dev['description']:
28 | dev_usb_path = dev['path']
29 | print('Found :', dev['description'])
30 | break
31 |
32 | fps = 60
33 | dt = 1 / fps
34 |
35 | with SharedMemoryManager() as shm_manager:
36 | with UvcCamera(
37 | shm_manager=shm_manager,
38 | dev_video_path=dev_video_path,
39 | dev_usb_path=dev_usb_path
40 | ) as camera:
41 | print('Ready!')
42 | t_start = time.monotonic()
43 | iter_idx = 0
44 | while True:
45 | t_cycle_end = t_start + (iter_idx + 1) * dt
46 |
47 | data = camera.get()
48 | img = data['color']
49 |
50 | # Display the resulting frame
51 | cv2.imshow('frame', img)
52 | if cv2.pollKey() & 0xFF == ord('q'):
53 | break
54 |
55 | precise_wait(t_cycle_end)
56 | iter_idx += 1
57 |
58 | # %%
59 | if __name__ == '__main__':
60 | main()
61 |
--------------------------------------------------------------------------------
/scripts_real/launch_franka_interface_server.py:
--------------------------------------------------------------------------------
1 | import zerorpc
2 | from polymetis import RobotInterface
3 | import scipy.spatial.transform as st
4 | import numpy as np
5 | import torch
6 |
7 | class FrankaArmInterface:
8 | def __init__(self):
9 | self.robot = RobotInterface(ip_address="172.16.0.3")
10 | print('FrankaArmInterface initialized')
11 |
12 | def get_ee_pose(self):
13 | data = self.robot.get_ee_pose()
14 | pos = data[0].numpy()
15 | quat_xyzw = data[1].numpy()
16 | rot_vec = st.Rotation.from_quat(quat_xyzw).as_rotvec()
17 | return np.concatenate([pos, rot_vec]).tolist()
18 |
19 | def get_joint_positions(self):
20 | return self.robot.get_joint_positions().numpy().tolist()
21 |
22 | def get_joint_velocities(self):
23 | return self.robot.get_joint_velocities().numpy().tolist()
24 |
25 | def move_to_joint_positions(self, positions, time_to_go):
26 | self.robot.move_to_joint_positions(
27 | positions=torch.Tensor(positions),
28 | time_to_go=time_to_go
29 | )
30 |
31 | def move_to_ee_pose(self, pose, time_to_go):
32 | self.robot.move_to_ee_pose(
33 | position=torch.Tensor(pose[:3]),
34 | orientation=torch.Tensor(st.Rotation.from_rotvec(pose[3:]).as_quat()),
35 | time_to_go=time_to_go
36 | )
37 |
38 | def start_cartesian_impedance(self, Kx, Kxd):
39 | self.robot.start_cartesian_impedance(
40 | Kx=torch.Tensor(Kx),
41 | Kxd=torch.Tensor(Kxd)
42 | )
43 |
44 | def update_desired_ee_pose(self, pose):
45 | pose = np.asarray(pose)
46 | self.robot.update_desired_ee_pose(
47 | position=torch.Tensor(pose[:3]),
48 | orientation=torch.Tensor(st.Rotation.from_rotvec(pose[3:]).as_quat())
49 | )
50 |
51 | def terminate_current_policy(self):
52 | self.robot.terminate_current_policy()
53 |
54 | arm_server = zerorpc.Server(FrankaArmInterface())
55 | arm_server.bind("tcp://0.0.0.0:4242")
56 | arm_server.run()
--------------------------------------------------------------------------------
/tests/test_pose_util.py:
--------------------------------------------------------------------------------
1 | # %%
2 | import sys
3 | import os
4 |
5 | ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
6 | sys.path.append(ROOT_DIR)
7 | os.chdir(ROOT_DIR)
8 |
9 | # %%
10 | import numpy as np
11 | from diffusion_policy.model.common.rotation_transformer import RotationTransformer
12 | from umi.common.pose_util import rot6d_to_mat, mat_to_rot6d
13 |
14 | # %%
15 | def test():
16 | N = 100
17 | d6 = np.random.normal(size=(N,6))
18 | rt = RotationTransformer(from_rep='rotation_6d', to_rep='matrix')
19 | gt_mat = rt.forward(d6)
20 | mat = rot6d_to_mat(d6)
21 | assert np.allclose(gt_mat, mat)
22 |
23 | to_d6 = mat_to_rot6d(mat)
24 | to_d6_gt = rt.inverse(mat)
25 | assert np.allclose(to_d6, to_d6_gt)
26 |
27 | gt_mat = rt.forward(d6[0])
28 | mat = rot6d_to_mat(d6[0])
29 | assert np.allclose(gt_mat, mat)
30 |
31 | to_d6 = mat_to_rot6d(mat)
32 | to_d6_gt = rt.inverse(mat)
33 | assert np.allclose(to_d6, to_d6_gt)
34 |
35 |
36 | if __name__ == "__main__":
37 | test()
38 |
--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
1 | """
2 | Usage:
3 | Training:
4 | python train.py --config-name=train_diffusion_lowdim_workspace
5 | """
6 |
7 | import sys
8 | # use line-buffering for both stdout and stderr
9 | sys.stdout = open(sys.stdout.fileno(), mode='w', buffering=1)
10 | sys.stderr = open(sys.stderr.fileno(), mode='w', buffering=1)
11 |
12 | import hydra
13 | from omegaconf import OmegaConf
14 | import pathlib
15 | from diffusion_policy.workspace.base_workspace import BaseWorkspace
16 |
17 | # allows arbitrary python code execution in configs using the ${eval:''} resolver
18 | OmegaConf.register_new_resolver("eval", eval, replace=True)
19 |
20 | @hydra.main(
21 | version_base=None,
22 | config_path=str(pathlib.Path(__file__).parent.joinpath(
23 | 'diffusion_policy','config'))
24 | )
25 | def main(cfg: OmegaConf):
26 | # resolve immediately so all the ${now:} resolvers
27 | # will use the same time.
28 | OmegaConf.resolve(cfg)
29 |
30 | cls = hydra.utils.get_class(cfg._target_)
31 | workspace: BaseWorkspace = cls(cfg)
32 | workspace.run()
33 |
34 | if __name__ == "__main__":
35 | main()
36 |
--------------------------------------------------------------------------------
/train_scripts/arrange_mouse.sh:
--------------------------------------------------------------------------------
1 | task_name="arrange_mouse"
2 | logging_time=$(date "+%d-%H.%M.%S")
3 | now_seconds="${logging_time: -8}"
4 | now_date=$(date "+%Y.%m.%d")
5 | run_dir="data/outputs/${now_date}/${now_seconds}"
6 | echo ${run_dir}
7 |
8 | # accelerate launch --mixed_precision 'bf16' ../train.py \
9 | python ../train.py \
10 | --config-name=train_diffusion_unet_timm_umi_workspace \
11 | multi_run.run_dir=${run_dir} multi_run.wandb_name_base=${logging_time} hydra.run.dir=${run_dir} hydra.sweep.dir=${run_dir} \
12 | task.dataset_path=../data/dataset/${task_name}/dataset.zarr.zip \
13 | training.num_epochs=150 \
14 | dataloader.batch_size=32 \
15 | logging.name="${logging_time}_${task_name}" \
16 | policy.obs_encoder.model_name='vit_large_patch14_dinov2.lvd142m' \
17 | task.dataset.dataset_idx=\'1-32\' \
18 | task.dataset.use_ratio=0.5
--------------------------------------------------------------------------------
/train_scripts/fold_towel.sh:
--------------------------------------------------------------------------------
1 | task_name="fold_towel"
2 | logging_time=$(date "+%d-%H.%M.%S")
3 | now_seconds="${logging_time: -8}"
4 | now_date=$(date "+%Y.%m.%d")
5 | run_dir="data/outputs/${now_date}/${now_seconds}"
6 | echo ${run_dir}
7 |
8 | # accelerate launch --mixed_precision 'bf16' ../train.py \
9 | python ../train.py \
10 | --config-name=train_diffusion_unet_timm_umi_workspace \
11 | multi_run.run_dir=${run_dir} multi_run.wandb_name_base=${logging_time} hydra.run.dir=${run_dir} hydra.sweep.dir=${run_dir} \
12 | task.dataset_path=../data/dataset/${task_name}/dataset.zarr.zip \
13 | training.num_epochs=150 \
14 | dataloader.batch_size=32 \
15 | logging.name="${logging_time}_${task_name}" \
16 | policy.obs_encoder.model_name='vit_large_patch14_dinov2.lvd142m' \
17 | task.dataset.dataset_idx=\'1-32\' \
18 | task.dataset.use_ratio=1
--------------------------------------------------------------------------------
/train_scripts/pour_water.sh:
--------------------------------------------------------------------------------
1 | task_name="pour_water"
2 | logging_time=$(date "+%d-%H.%M.%S")
3 | now_seconds="${logging_time: -8}"
4 | now_date=$(date "+%Y.%m.%d")
5 | run_dir="data/outputs/${now_date}/${now_seconds}"
6 | echo ${run_dir}
7 |
8 | # accelerate launch --mixed_precision 'bf16' ../train.py \
9 | python ../train.py \
10 | --config-name=train_diffusion_unet_timm_umi_workspace \
11 | multi_run.run_dir=${run_dir} multi_run.wandb_name_base=${logging_time} hydra.run.dir=${run_dir} hydra.sweep.dir=${run_dir} \
12 | task.dataset_path=../data/dataset/${task_name}/dataset.zarr.zip \
13 | training.num_epochs=150 \
14 | dataloader.batch_size=32 \
15 | logging.name="${logging_time}_${task_name}" \
16 | policy.obs_encoder.model_name='vit_large_patch14_dinov2.lvd142m' \
17 | task.dataset.dataset_idx=\'1-32\' \
18 | task.dataset.use_ratio=0.5 \
19 | task.obs_down_sample_steps=[3,15] \
20 | task.action_down_sample_steps=3 \
21 | task.low_dim_obs_horizon=3 \
22 | task.img_obs_horizon=3
--------------------------------------------------------------------------------
/train_scripts/unplug_charger.sh:
--------------------------------------------------------------------------------
1 | task_name="unplug_charger"
2 | logging_time=$(date "+%d-%H.%M.%S")
3 | now_seconds="${logging_time: -8}"
4 | now_date=$(date "+%Y.%m.%d")
5 | run_dir="data/outputs/${now_date}/${now_seconds}"
6 | echo ${run_dir}
7 |
8 | # accelerate launch --mixed_precision 'bf16' ../train.py \
9 | python ../train.py \
10 | --config-name=train_diffusion_unet_timm_umi_workspace \
11 | multi_run.run_dir=${run_dir} multi_run.wandb_name_base=${logging_time} hydra.run.dir=${run_dir} hydra.sweep.dir=${run_dir} \
12 | task.dataset_path=../data/dataset/${task_name}/dataset.zarr.zip \
13 | training.num_epochs=150 \
14 | dataloader.batch_size=32 \
15 | logging.name="${logging_time}_${task_name}" \
16 | policy.obs_encoder.model_name='vit_large_patch14_dinov2.lvd142m' \
17 | task.dataset.dataset_idx=\'1-32\' \
18 | task.dataset.use_ratio=1 \
19 | task.obs_down_sample_steps=[3,30] \
20 | task.action_down_sample_steps=3 \
21 | task.low_dim_obs_horizon=3 \
22 | task.img_obs_horizon=3
--------------------------------------------------------------------------------
/umi/asset/mask.json:
--------------------------------------------------------------------------------
1 | {
2 | "mirror_mask_pts":[
3 | [540, 1700],
4 | [680, 1450],
5 | [590, 1070],
6 | [290, 1130],
7 | [290, 1770],
8 | [550, 1770]
9 | ],
10 | "gripper_mask_pts": [
11 | [1100, 1700],
12 | [650, 1500],
13 | [0, 1350],
14 | [520, 2028]
15 | ],
16 | "resolution": [2028, 2704]
17 | }
--------------------------------------------------------------------------------
/umi/common/exiftool_util.py:
--------------------------------------------------------------------------------
1 | from exiftool import ExifToolHelper
2 |
3 | def get_videos_metadata(video_paths, keys=['QuickTime:CameraSerialNumber', 'QuickTime:Model']):
4 | results = dict()
5 | with ExifToolHelper() as et:
6 | for meta in et.get_metadata(video_paths):
7 | result = dict()
8 | for key in keys:
9 | result[key] = meta[key]
10 | results[meta['SourceFile']] = result
11 | return results
12 |
--------------------------------------------------------------------------------
/umi/common/interpolation_util.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import scipy.interpolate as si
3 | import scipy.spatial.transform as st
4 |
5 |
6 | def get_interp1d(t, x):
7 | gripper_interp = si.interp1d(
8 | t, x,
9 | axis=0, bounds_error=False,
10 | fill_value=(x[0], x[-1]))
11 | return gripper_interp
12 |
13 |
14 | class PoseInterpolator:
15 | def __init__(self, t, x):
16 | pos = x[:,:3]
17 | rot = st.Rotation.from_rotvec(x[:,3:])
18 | self.pos_interp = get_interp1d(t, pos)
19 | self.rot_interp = st.Slerp(t, rot)
20 |
21 | @property
22 | def x(self):
23 | return self.pos_interp.x
24 |
25 | def __call__(self, t):
26 | min_t = self.pos_interp.x[0]
27 | max_t = self.pos_interp.x[-1]
28 | t = np.clip(t, min_t, max_t)
29 |
30 | pos = self.pos_interp(t)
31 | rot = self.rot_interp(t)
32 | rvec = rot.as_rotvec()
33 | pose = np.concatenate([pos, rvec], axis=-1)
34 | return pose
35 |
36 | def get_gripper_calibration_interpolator(
37 | aruco_measured_width,
38 | aruco_actual_width):
39 | """
40 | Assumes the minimum width in aruco_actual_width
41 | is measured when the gripper is fully closed
42 | and maximum width is when the gripper is fully opened
43 | """
44 | aruco_measured_width = np.array(aruco_measured_width)
45 | aruco_actual_width = np.array(aruco_actual_width)
46 | assert len(aruco_measured_width) == len(aruco_actual_width)
47 | assert len(aruco_actual_width) >= 2
48 | aruco_min_width = np.min(aruco_actual_width)
49 | gripper_actual_width = aruco_actual_width - aruco_min_width
50 | interp = get_interp1d(aruco_measured_width, gripper_actual_width)
51 | return interp
52 |
--------------------------------------------------------------------------------
/umi/common/k3d_util.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import numba
3 |
4 | @numba.jit()
5 | def k3d_get_pose_axis(poses, axis_size_m = 0.1):
6 | # points in camera frame
7 | points = np.zeros((4,3), dtype=poses.dtype)
8 | points[1,0] = axis_size_m
9 | points[2,1] = axis_size_m
10 | points[3,2] = axis_size_m
11 |
12 | n_poses = poses.shape[0]
13 | out_verts = np.zeros((n_poses * 4, 3), dtype=poses.dtype)
14 | out_idxs = np.zeros((n_poses * 3, 2), dtype=np.int64)
15 | out_colors = np.zeros((n_poses * 4,), dtype=np.int64)
16 | for i in range(n_poses):
17 | this_pose = poses[i]
18 | # convert points to world frame
19 | this_verts = points @ this_pose[:3,:3].T + this_pose[:3,3]
20 | # fill in vert array
21 | vert_idx_start = i * 4
22 | out_verts[vert_idx_start:vert_idx_start+4] = this_verts
23 | # draw 3 lines for x,y,z axis
24 | this_idxs = out_idxs[i*3:(i+1)*3]
25 | this_idxs[0] = [0,1]
26 | this_idxs[1] = [0,2]
27 | this_idxs[2] = [0,3]
28 | this_idxs += vert_idx_start
29 | # fill out vertex colors, rgb for xyz
30 | out_colors[i*4:(i+1)*4] = [0xffffff, 0xff0000, 0x00ff00, 0x0000ff]
31 |
32 | return out_verts, out_idxs, out_colors
--------------------------------------------------------------------------------
/umi/common/latency_util.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from scipy.interpolate import interp1d
3 | import scipy.signal as ss
4 |
5 | def regular_sample(x, t, t_samples):
6 | spline = interp1d(x=t, y=x, bounds_error=False, fill_value=(x[0], x[-1]))
7 | result = spline(t_samples)
8 | return result
9 |
10 | def get_latency(
11 | x_target, t_target,
12 | x_actual, t_actual,
13 | t_start=None, t_end=None,
14 | resample_dt=1/1000,
15 | force_positive=False
16 | ):
17 | assert len(x_target) == len(t_target)
18 | assert len(x_actual) == len(t_actual)
19 | if t_start is None:
20 | t_start = max(t_target[0], t_actual[0])
21 | if t_end is None:
22 | t_end = min(t_target[-1], t_actual[-1])
23 | n_samples = int((t_end - t_start) / resample_dt)
24 | t_samples = np.arange(n_samples) * resample_dt + t_start
25 | target_samples = regular_sample(x_target, t_target, t_samples)
26 | actual_samples = regular_sample(x_actual, t_actual, t_samples)
27 |
28 | # normalize samples to zero mean unit std
29 | mean = np.mean(np.concatenate([target_samples, actual_samples]))
30 | std = np.std(np.concatenate([target_samples, actual_samples]))
31 | target_samples = (target_samples - mean) / std
32 | actual_samples = (actual_samples - mean) / std
33 |
34 | # cross correlation
35 | correlation = ss.correlate(actual_samples, target_samples)
36 | lags = ss.correlation_lags(len(actual_samples), len(target_samples))
37 | t_lags = lags * resample_dt
38 |
39 | latency = None
40 | if force_positive:
41 | latency = t_lags[np.argmax(correlation[t_lags >= 0])]
42 | else:
43 | latency = t_lags[np.argmax(correlation)]
44 |
45 | info = {
46 | 't_samples': t_samples,
47 | 'x_target': target_samples,
48 | 'x_actual': actual_samples,
49 | 'correlation': correlation,
50 | 'lags': t_lags
51 | }
52 |
53 | return latency, info
54 |
--------------------------------------------------------------------------------
/umi/common/mocap_util.py:
--------------------------------------------------------------------------------
1 | from datetime import datetime
2 | import pandas as pd
3 | import numpy as np
4 | from scipy.spatial.transform import Rotation
5 |
6 | def get_mocap_start_datetime(csv_path):
7 | with open(csv_path, 'r') as f:
8 | first_row = f.readline().split(',')
9 | meta_dict = dict()
10 | for i in range(len(first_row) // 2):
11 | start = i * 2
12 | end = start + 1
13 | meta_dict[first_row[start]] = first_row[end]
14 | start_timestamp_str = meta_dict['Capture Start Time']
15 | start_date = datetime.strptime(start_timestamp_str, r"%Y-%m-%d %I.%M.%S.%f %p")
16 | return start_date
17 |
18 | def get_mocap_data(csv_path, rigid_body_name):
19 | mocap_df = pd.read_csv(csv_path, skiprows=2, index_col=0, header=[1,3,4])
20 | assert mocap_df.index[0] == 0
21 | assert mocap_df.index[-1] == (len(mocap_df) - 1)
22 | assert mocap_df.columns[0][-1] == 'Time (Seconds)'
23 |
24 | time_since_start = mocap_df.iloc[:,0].to_numpy()
25 | pos = np.zeros((len(mocap_df), 3))
26 | pos[:,0] = mocap_df[(rigid_body_name, 'Position', 'X')]
27 | pos[:,1] = mocap_df[(rigid_body_name, 'Position', 'Y')]
28 | pos[:,2] = mocap_df[(rigid_body_name, 'Position', 'Z')]
29 |
30 | rot_quat = np.zeros((len(mocap_df), 4))
31 | rot_quat[:,0] = mocap_df[(rigid_body_name, 'Rotation', 'X')]
32 | rot_quat[:,1] = mocap_df[(rigid_body_name, 'Rotation', 'Y')]
33 | rot_quat[:,2] = mocap_df[(rigid_body_name, 'Rotation', 'Z')]
34 | rot_quat[:,3] = mocap_df[(rigid_body_name, 'Rotation', 'W')]
35 | rot = Rotation.from_quat(rot_quat)
36 |
37 | pose = np.zeros((pos.shape[0], 4, 4), dtype=pos.dtype)
38 | pose[:,3,3] = 1
39 | pose[:,:3,:3] = rot.as_matrix()
40 | pose[:,:3,3] = pos
41 |
42 | result = {
43 | 'time_since_start': time_since_start,
44 | 'pose': pose,
45 | }
46 | return result
47 |
--------------------------------------------------------------------------------
/umi/common/nested_dict_util.py:
--------------------------------------------------------------------------------
1 | import functools
2 |
3 | def nested_dict_map(f, x):
4 | """
5 | Map f over all leaf of nested dict x
6 | """
7 |
8 | if not isinstance(x, dict):
9 | return f(x)
10 | y = dict()
11 | for key, value in x.items():
12 | y[key] = nested_dict_map(f, value)
13 | return y
14 |
15 | def nested_dict_reduce(f, x):
16 | """
17 | Map f over all values of nested dict x, and reduce to a single value
18 | """
19 | if not isinstance(x, dict):
20 | return x
21 |
22 | reduced_values = list()
23 | for value in x.values():
24 | reduced_values.append(nested_dict_reduce(f, value))
25 | y = functools.reduce(f, reduced_values)
26 | return y
27 |
28 |
29 | def nested_dict_check(f, x):
30 | bool_dict = nested_dict_map(f, x)
31 | result = nested_dict_reduce(lambda x, y: x and y, bool_dict)
32 | return result
33 |
--------------------------------------------------------------------------------
/umi/common/orb_slam_util.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import pandas as pd
3 | from scipy.spatial.transform import Rotation
4 |
5 | def load_tum_trajectory(tum_txt_path):
6 | tum_traj_raw = np.loadtxt(tum_txt_path, delimiter=' ', dtype=np.float32)
7 | if len(tum_traj_raw) == 0:
8 | return {
9 | 'timestamp': np.array([]),
10 | 'pose': np.array([]),
11 | }
12 |
13 | timestamp_sec = tum_traj_raw[:,0]
14 | cam_pos = tum_traj_raw[:,1:4]
15 | cam_rot_quat_xyzw = tum_traj_raw[:,4:8]
16 | cam_rot = Rotation.from_quat(cam_rot_quat_xyzw)
17 |
18 | cam_pose = np.zeros((cam_pos.shape[0], 4, 4), dtype=np.float32)
19 | cam_pose[:,3,3] = 1
20 | cam_pose[:,:3,3] = cam_pos
21 | cam_pose[:,:3,:3] = cam_rot.as_matrix()
22 |
23 | result = {
24 | 'timestamp': timestamp_sec,
25 | 'pose': cam_pose
26 | }
27 | return result
28 |
29 | def load_csv_trajectory(csv_path):
30 | df = pd.read_csv(csv_path)
31 | if (~df.is_lost).sum() == 0:
32 | return {
33 | 'raw_data': df
34 | }
35 |
36 | valid_df = df.loc[~df.is_lost]
37 |
38 | timestamp_sec = valid_df['timestamp'].to_numpy()
39 | cam_pos = valid_df[['x', 'y', 'z']].to_numpy()
40 | cam_rot_quat_xyzw = valid_df[['q_x', 'q_y', 'q_z', 'q_w']].to_numpy()
41 | cam_rot = Rotation.from_quat(cam_rot_quat_xyzw)
42 |
43 | cam_pose = np.zeros((cam_pos.shape[0], 4, 4), dtype=np.float32)
44 | cam_pose[:,3,3] = 1
45 | cam_pose[:,:3,3] = cam_pos
46 | cam_pose[:,:3,:3] = cam_rot.as_matrix()
47 |
48 | result = {
49 | 'timestamp': timestamp_sec,
50 | 'pose': cam_pose,
51 | 'raw_data': df
52 | }
53 | return result
54 |
--------------------------------------------------------------------------------
/umi/common/precise_sleep.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | def precise_sleep(dt: float, slack_time: float=0.001, time_func=time.monotonic):
4 | """
5 | Use hybrid of time.sleep and spinning to minimize jitter.
6 | Sleep dt - slack_time seconds first, then spin for the rest.
7 | """
8 | t_start = time_func()
9 | if dt > slack_time:
10 | time.sleep(dt - slack_time)
11 | t_end = t_start + dt
12 | while time_func() < t_end:
13 | pass
14 | return
15 |
16 | def precise_wait(t_end: float, slack_time: float=0.001, time_func=time.monotonic):
17 | t_start = time_func()
18 | t_wait = t_end - t_start
19 | if t_wait > 0:
20 | t_sleep = t_wait - slack_time
21 | if t_sleep > 0:
22 | time.sleep(t_sleep)
23 | while time_func() < t_end:
24 | pass
25 | return
26 |
--------------------------------------------------------------------------------
/umi/common/timecode_util.py:
--------------------------------------------------------------------------------
1 | from typing import Union
2 | from fractions import Fraction
3 | import datetime
4 | import av
5 |
6 |
7 | def timecode_to_seconds(
8 | timecode: str, frame_rate: Union[int, float, Fraction]
9 | ) -> Union[float, Fraction]:
10 | """
11 | Convert non-skip frame timecode into seconds since midnight
12 | """
13 | # calculate whole frame rate
14 | # 29.97 -> 30, 59.94 -> 60
15 | int_frame_rate = round(frame_rate)
16 |
17 | # parse timecode string
18 | h, m, s, f = [int(x) for x in timecode.split(':')]
19 |
20 | # calculate frames assuming whole frame rate (i.e. non-drop frame)
21 | frames = (3600 * h + 60 * m + s) * int_frame_rate + f
22 |
23 | # convert to seconds
24 | seconds = frames / frame_rate
25 | return seconds
26 |
27 |
28 | def stream_get_start_datetime(stream: av.stream.Stream) -> datetime.datetime:
29 | """
30 | Combines creation time and timecode to get high-precision
31 | time for the first frame of a video.
32 | """
33 | # read metadata
34 | frame_rate = stream.average_rate
35 | tc = stream.metadata['timecode']
36 | creation_time = stream.metadata['creation_time']
37 |
38 | # get time within the day
39 | seconds_since_midnight = float(timecode_to_seconds(timecode=tc, frame_rate=frame_rate))
40 | delta = datetime.timedelta(seconds=seconds_since_midnight)
41 |
42 | # get dates
43 | create_datetime = datetime.datetime.strptime(creation_time, r"%Y-%m-%dT%H:%M:%S.%fZ")
44 | create_datetime = create_datetime.replace(hour=0, minute=0, second=0, microsecond=0)
45 | start_datetime = create_datetime + delta
46 | return start_datetime
47 |
48 |
49 | def mp4_get_start_datetime(mp4_path: str) -> datetime.datetime:
50 | with av.open(mp4_path) as container:
51 | stream = container.streams.video[0]
52 | return stream_get_start_datetime(stream=stream)
53 |
--------------------------------------------------------------------------------
/umi/pipeline/aruco_detection.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Fanqi-Lin/Data-Scaling-Laws/bd6941edea44a72611f585a23cf219f402bbaafc/umi/pipeline/aruco_detection.py
--------------------------------------------------------------------------------
/umi/real_world/keystroke_counter.py:
--------------------------------------------------------------------------------
1 | from pynput.keyboard import Key, KeyCode, Listener
2 | from collections import defaultdict
3 | from threading import Lock
4 |
5 | class KeystrokeCounter(Listener):
6 | def __init__(self):
7 | self.key_count_map = defaultdict(lambda:0)
8 | self.key_press_list = list()
9 | self.lock = Lock()
10 | super().__init__(on_press=self.on_press, on_release=self.on_release)
11 |
12 | def on_press(self, key):
13 | with self.lock:
14 | self.key_count_map[key] += 1
15 | self.key_press_list.append(key)
16 |
17 | def on_release(self, key):
18 | pass
19 |
20 | def clear(self):
21 | with self.lock:
22 | self.key_count_map = defaultdict(lambda:0)
23 | self.key_press_list = list()
24 |
25 | def __getitem__(self, key):
26 | with self.lock:
27 | return self.key_count_map[key]
28 |
29 | def get_press_events(self):
30 | with self.lock:
31 | events = list(self.key_press_list)
32 | self.key_press_list = list()
33 | return events
34 |
35 | if __name__ == '__main__':
36 | import time
37 | with KeystrokeCounter() as counter:
38 | try:
39 | while True:
40 | print('Space:', counter[Key.space])
41 | print('q:', counter[KeyCode(char='q')])
42 | time.sleep(1/60)
43 | except KeyboardInterrupt:
44 | events = counter.get_press_events()
45 | print(events)
46 |
--------------------------------------------------------------------------------
/umi/shared_memory/shared_memory_util.py:
--------------------------------------------------------------------------------
1 | from typing import Tuple
2 | from dataclasses import dataclass
3 | import numpy as np
4 | from multiprocessing.managers import SharedMemoryManager
5 | from atomics import atomicview, MemoryOrder, UINT
6 |
7 | @dataclass
8 | class ArraySpec:
9 | name: str
10 | shape: Tuple[int]
11 | dtype: np.dtype
12 |
13 |
14 | class SharedAtomicCounter:
15 | def __init__(self,
16 | shm_manager: SharedMemoryManager,
17 | size :int=8 # 64bit int
18 | ):
19 | shm = shm_manager.SharedMemory(size=size)
20 | self.shm = shm
21 | self.size = size
22 | self.store(0) # initialize
23 |
24 | @property
25 | def buf(self):
26 | return self.shm.buf[:self.size]
27 |
28 | def load(self) -> int:
29 | with atomicview(buffer=self.buf, atype=UINT) as a:
30 | value = a.load(order=MemoryOrder.ACQUIRE)
31 | return value
32 |
33 | def store(self, value: int):
34 | with atomicview(buffer=self.buf, atype=UINT) as a:
35 | a.store(value, order=MemoryOrder.RELEASE)
36 |
37 | def add(self, value: int):
38 | with atomicview(buffer=self.buf, atype=UINT) as a:
39 | a.add(value, order=MemoryOrder.ACQ_REL)
40 |
--------------------------------------------------------------------------------
/umi/traj_eval/trajectory_utils.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python2
2 | """
3 | @author: Christian Forster
4 | """
5 |
6 | import os
7 | import numpy as np
8 | import umi.traj_eval.transformations as tf
9 | import numba
10 |
11 | def get_rigid_body_trafo(quat, trans):
12 | T = tf.quaternion_matrix(quat)
13 | T[0:3, 3] = trans
14 | return T
15 |
16 | def get_distance_from_start(gt_translation):
17 | distances = np.diff(gt_translation[:, 0:3], axis=0)
18 | distances = np.sqrt(np.sum(np.multiply(distances, distances), 1))
19 | distances = np.cumsum(distances)
20 | distances = np.concatenate(([0], distances))
21 | return distances
22 |
23 | @numba.jit(nopython=True)
24 | def compute_comparison_indices_length(distances, dist, max_dist_diff):
25 | max_idx = len(distances)
26 | comparisons = []
27 | for idx, d in enumerate(distances):
28 | best_idx = -1
29 | error = max_dist_diff
30 | for i in range(idx, max_idx):
31 | if np.abs(distances[i]-(d+dist)) < error:
32 | best_idx = i
33 | error = np.abs(distances[i] - (d+dist))
34 | if best_idx != -1:
35 | comparisons.append(best_idx)
36 | return comparisons
37 |
38 |
39 | def compute_angle(transform):
40 | """
41 | Compute the rotation angle from a 4x4 homogeneous matrix.
42 | """
43 | # an invitation to 3-d vision, p 27
44 | return np.arccos(
45 | min(1, max(-1, (np.trace(transform[0:3, 0:3]) - 1)/2)))*180.0/np.pi
46 |
--------------------------------------------------------------------------------
/visualize_dataset.py:
--------------------------------------------------------------------------------
1 | import zarr
2 | from diffusion_policy.common.replay_buffer import ReplayBuffer
3 | from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs
4 | from PIL import Image
5 | import os
6 | import numpy as np
7 |
8 | register_codecs()
9 |
10 | def read_dataset(dataset_path):
11 | with zarr.ZipStore(dataset_path, mode='r') as zip_store:
12 | replay_buffer = ReplayBuffer.copy_from_store(
13 | src_store=zip_store,
14 | store=zarr.MemoryStore()
15 | )
16 | return replay_buffer
17 |
18 | task_name = 'pour_water' # or 'arrange_mouse' 'fold_towel' 'unplug_charger'
19 | dataset_dir = f'data/dataset/{task_name}'
20 | save_dir = f'data/data_vis/{task_name}'
21 | os.makedirs(save_dir, exist_ok=True)
22 |
23 | replay_buffer = read_dataset(os.path.join(dataset_dir, 'dataset.zarr.zip'))
24 | with open(os.path.join(dataset_dir, 'count.txt'), 'r') as f:
25 | counts = f.readlines()
26 | counts = [int(count[:-1]) if count[-1] == '\n' else int(count) for count in counts]
27 | counts = [0] + list(np.cumsum(counts))[:-1]
28 |
29 | episode_ends = replay_buffer.episode_ends[:]
30 | episode_start = [0] + list(episode_ends)
31 |
32 | save_img_idx = [episode_start[count] for count in counts]
33 | for i, single_episode in enumerate(save_img_idx):
34 | episode_img = replay_buffer.data.camera0_rgb[single_episode]
35 | episode_img = Image.fromarray(episode_img)
36 | episode_img.save(os.path.join(save_dir, f'{i}.png'))
37 |
--------------------------------------------------------------------------------