├── .gitignore
├── README.md
├── 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
│ ├── noise_sampler.py
│ ├── normalize_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
│ ├── 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_image.yaml
│ │ ├── lift_image_abs.yaml
│ │ ├── lift_lowdim.yaml
│ │ ├── lift_lowdim_abs.yaml
│ │ ├── pusht_image.yaml
│ │ ├── pusht_lowdim.yaml
│ │ ├── real_pusht_image.yaml
│ │ ├── square_image.yaml
│ │ ├── square_image_abs.yaml
│ │ ├── square_lowdim.yaml
│ │ ├── square_lowdim_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
│ ├── 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_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_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
├── 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_action_dataset.py
│ ├── robomimic_replay_image_dataset.py
│ └── robomimic_replay_lowdim_dataset.py
├── env
│ ├── aloha
│ │ ├── aloha_image_wrapper.py
│ │ ├── assets
│ │ │ ├── bimanual_viperx_ee_insertion.xml
│ │ │ ├── bimanual_viperx_ee_transfer_cube.xml
│ │ │ ├── bimanual_viperx_insertion.xml
│ │ │ ├── bimanual_viperx_transfer_cube.xml
│ │ │ ├── scene.xml
│ │ │ ├── singlearm_viperx_ee_transfer_cube.xml
│ │ │ ├── singlearm_viperx_transfer_cube.xml
│ │ │ ├── tabletop.stl
│ │ │ ├── vx300s_10_custom_finger_left.stl
│ │ │ ├── vx300s_10_custom_finger_right.stl
│ │ │ ├── vx300s_10_gripper_finger.stl
│ │ │ ├── vx300s_11_ar_tag.stl
│ │ │ ├── vx300s_1_base.stl
│ │ │ ├── vx300s_2_shoulder.stl
│ │ │ ├── vx300s_3_upper_arm.stl
│ │ │ ├── vx300s_4_upper_forearm.stl
│ │ │ ├── vx300s_5_lower_forearm.stl
│ │ │ ├── vx300s_6_wrist.stl
│ │ │ ├── vx300s_7_gripper.stl
│ │ │ ├── vx300s_8_gripper_prop.stl
│ │ │ ├── vx300s_9_gripper_bar.stl
│ │ │ ├── vx300s_dependencies.xml
│ │ │ ├── vx300s_left.xml
│ │ │ └── vx300s_right.xml
│ │ ├── constants.py
│ │ ├── ee_sim_env.py
│ │ ├── env_utils.py
│ │ ├── scripted_policy.py
│ │ ├── sim_env.py
│ │ └── visualize_episodes.py
│ ├── 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
│ ├── aloha_image_runner.py
│ ├── 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_longhist_image_runner.py
│ ├── robomimic_lowdim_runner.py
│ └── robomimic_square_long_image_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
│ │ │ │ ├── multi_code_model.py
│ │ │ │ ├── trainer.py
│ │ │ │ └── utils.py
│ │ └── utils.py
│ ├── common
│ │ ├── dict_of_tensor_mixin.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_diffusion.py
│ │ └── transformer_for_regression.py
│ └── vision
│ │ ├── crop_randomizer.py
│ │ ├── model_getter.py
│ │ └── multi_image_obs_encoder.py
├── policy
│ ├── base_image_policy.py
│ ├── base_lowdim_policy.py
│ ├── bet_image_policy.py
│ ├── bet_lowdim_policy.py
│ ├── diffusion_transformer_hybrid_image_policy.py
│ ├── diffusion_transformer_lowdim_policy.py
│ ├── diffusion_unet_hybrid_image_policy.py
│ ├── diffusion_unet_image_policy.py
│ ├── diffusion_unet_lowdim_policy.py
│ ├── diffusion_unet_video_policy.py
│ ├── ibc_dfo_hybrid_image_policy.py
│ ├── ibc_dfo_lowdim_policy.py
│ ├── reg_transformer_hybrid_image_policy.py
│ ├── robomimic_image_policy.py
│ └── robomimic_lowdim_policy.py
├── real
│ ├── eval_launcher.py
│ ├── main_eval.py
│ ├── main_nocalib.py
│ ├── policy_wrapper.py
│ ├── policy_wrapper_marcel.py
│ ├── robot_env.py
│ ├── svo_reader.py
│ └── timestep_processing_marcel.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
│ └── 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_image_workspace.py
│ ├── train_bet_lowdim_workspace.py
│ ├── train_diffusion_transformer_hybrid_workspace.py
│ ├── train_diffusion_transformer_lowdim_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.py
├── experiment_configs
├── aloha
│ ├── transformer_aloha.yaml
│ ├── transformer_aloha_emb.yaml
│ └── transformer_aloha_reg_emb.yaml
├── longhist
│ ├── transformer_longhist.yaml
│ ├── transformer_longhist_emb.yaml
│ ├── transformer_longhist_past.yaml
│ ├── transformer_longhist_past_emb.yaml
│ ├── transformer_longhist_reg.yaml
│ └── transformer_longhist_reg_emb.yaml
├── square
│ ├── transformer_square.yaml
│ ├── transformer_square_bc_rnn.yaml
│ ├── transformer_square_emb.yaml
│ ├── transformer_square_past.yaml
│ ├── transformer_square_past_emb.yaml
│ ├── transformer_square_reg.yaml
│ ├── transformer_square_reg_emb.yaml
│ ├── unet_hybrid_square.yaml
│ └── unet_hybrid_square_emb.yaml
├── tool
│ ├── transformer_tool_hang.yaml
│ ├── transformer_tool_hang_emb.yaml
│ ├── transformer_tool_hang_reg.yaml
│ └── transformer_tool_hang_reg_emb.yaml
├── transformer_pusht.yaml
├── transformer_pusht_emb.yaml
└── transport
│ ├── transformer_transport.yaml
│ ├── transformer_transport_emb.yaml
│ ├── transformer_transport_reg.yaml
│ └── transformer_transport_reg_emb.yaml
├── gather_rollouts.py
├── get_action_loss_train.py
├── hsic.py
├── mlp_correlation.py
├── perturbs
├── chunk_1.yaml
├── chunk_2.yaml
├── chunk_4.yaml
├── chunk_8.yaml
└── none.yaml
├── requirements.txt
├── rewrite_with_embeddings.py
├── rollouts
└── merge_actions.py
├── rollouts_via_policy.py
├── run_gather_rollouts.sh
├── setup.py
├── train.py
├── train_rollouts.sh
├── transformer_eval.sh
└── transformer_history.sh
/conda_environment.yaml:
--------------------------------------------------------------------------------
1 | name: robodiff-lh
2 | channels:
3 | - pytorch
4 | - pytorch3d
5 | - nvidia
6 | - conda-forge
7 | dependencies:
8 | - python=3.9
9 | - pip=22.2.2
10 | - cudatoolkit=11.6
11 | - pytorch=1.12.1
12 | - torchvision=0.13.1
13 | - pytorch3d=0.7.0
14 | - numpy=1.23.3
15 | - numba==0.56.4
16 | - scipy==1.9.1
17 | - py-opencv=4.6.0
18 | - cffi=1.15.1
19 | - ipykernel=6.16
20 | - matplotlib=3.6.1
21 | - zarr=2.12.0
22 | - numcodecs=0.10.2
23 | - h5py=3.7.0
24 | - hydra-core=1.2.0
25 | - einops=0.4.1
26 | - tqdm=4.64.1
27 | - dill=0.3.5.1
28 | - scikit-video=1.1.11
29 | - scikit-image=0.19.3
30 | - gym=0.21.0
31 | - pymunk=6.2.1
32 | - wandb=0.13.3
33 | - threadpoolctl=3.1.0
34 | - shapely=1.8.4
35 | - cython=0.29.32
36 | - imageio=2.22.0
37 | - imageio-ffmpeg=0.4.7
38 | - termcolor=2.0.1
39 | - tensorboard=2.10.1
40 | - tensorboardx=2.5.1
41 | - psutil=5.9.2
42 | - click=8.0.4
43 | - boto3=1.24.96
44 | - accelerate=0.13.2
45 | - datasets=2.6.1
46 | - diffusers=0.11.1
47 | - av=10.0.0
48 | - cmake=3.24.3
49 | # trick to avoid cpu affinity issue described in https://github.com/pytorch/pytorch/issues/99625
50 | - llvm-openmp=14
51 | # trick to force reinstall imagecodecs via pip
52 | - imagecodecs==2022.8.8
53 | - pip:
54 | - ray[default,tune]==2.2.0
55 | # requires mujoco py dependencies libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf
56 | - free-mujoco-py==2.1.6
57 | - pygame==2.1.2
58 | - pybullet-svl==3.1.6.4
59 | - robosuite @ https://github.com/cheng-chi/robosuite/archive/277ab9588ad7a4f4b55cf75508b44aa67ec171f0.tar.gz
60 | - robomimic==0.2.0
61 | - pytorchvideo==0.1.5
62 | # pip package required for jpeg-xl
63 | - imagecodecs==2022.9.26
64 | - r3m @ https://github.com/facebookresearch/r3m/archive/b2334e726887fa0206962d7984c69c5fb09cceab.tar.gz
65 | - dm-control==1.0.9
66 |
--------------------------------------------------------------------------------
/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/noise_sampler.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import pdb
3 | class NoiseGenerator:
4 | def __init__(self, noise_strength, correlation_factor=0.9):
5 | self.noise_strength = noise_strength
6 | self.correlation_factor = correlation_factor
7 | self.previous_noise = None
8 | def step(self, pred):
9 | # Generate random noise
10 | noise_step = np.random.randn(*pred.shape) * self.noise_strength * pred
11 | # noise_step = pred * (np.random.rand(pred.shape[0], 1, pred.shape[2]) + 0.5) * np.random.choice([-1, 1], size=(pred.shape[0], 1, pred.shape[2])) * self.noise_strength
12 | # action_step = (pred[:, 1:] - pred[:, :-1])
13 | # noise_step = noise_seed.repeat(action_step.shape[1], axis=1) * action_step * self.noise_strength
14 | # If it's the first time step, there's no previous noise, so use the seed directly
15 | if self.previous_noise is None:
16 | self.previous_noise = noise_step
17 | else:
18 | # Combine the previous noise with new noise to create temporally correlated noise
19 | noise_step = self.correlation_factor * self.previous_noise + (1 - self.correlation_factor) * noise_step
20 | self.previous_noise = noise_step
21 | # noise_cum = np.cumsum(noise_step, axis=1)
22 | return noise_step
23 |
--------------------------------------------------------------------------------
/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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
32 | seed: 42
33 | val_ratio: 0.02
34 |
--------------------------------------------------------------------------------
/diffusion_policy/config/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
38 | seed: 42
39 | val_ratio: 0.02
40 | max_train_episodes: 90
41 |
--------------------------------------------------------------------------------
/diffusion_policy/config/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:'${n_action_steps}-1'}
32 | seed: 42
33 | val_ratio: 0.02
34 | max_train_episodes: 90
35 |
--------------------------------------------------------------------------------
/diffusion_policy/config/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/task/transport_image.yaml:
--------------------------------------------------------------------------------
1 | name: transport_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | shouldercamera0_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 | shouldercamera1_image:
20 | shape: [3, 84, 84]
21 | type: rgb
22 | robot1_eye_in_hand_image:
23 | shape: [3, 84, 84]
24 | type: rgb
25 | robot1_eef_pos:
26 | shape: [3]
27 | # type default: low_dim
28 | robot1_eef_quat:
29 | shape: [4]
30 | robot1_gripper_qpos:
31 | shape: [2]
32 | action:
33 | shape: [14]
34 |
35 | task_name: &task_name transport
36 | dataset_type: &dataset_type ph
37 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image.hdf5
38 | abs_action: &abs_action False
39 |
40 | env_runner:
41 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
42 | dataset_path: *dataset_path
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 | max_steps: 700
51 | n_obs_steps: ${n_obs_steps}
52 | n_action_steps: ${n_action_steps}
53 | render_obs_key: 'shouldercamera0_image'
54 | fps: 10
55 | crf: 22
56 | past_action: ${past_action_visible}
57 | abs_action: *abs_action
58 | tqdm_interval_sec: 1.0
59 | n_envs: 28
60 | # evaluation at this config requires a 16 core 64GB instance.
61 |
62 | dataset:
63 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
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:'${n_action_steps}-1'}
69 | n_obs_steps: ${dataset_obs_steps}
70 | abs_action: *abs_action
71 | rotation_rep: 'rotation_6d'
72 | use_legacy_normalizer: False
73 | use_cache: True
74 | seed: 42
75 | val_ratio: 0.02
76 |
--------------------------------------------------------------------------------
/diffusion_policy/config/task/transport_image_abs.yaml:
--------------------------------------------------------------------------------
1 | name: transport_image
2 |
3 | shape_meta: &shape_meta
4 | # acceptable types: rgb, low_dim
5 | obs:
6 | shouldercamera0_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 | shouldercamera1_image:
20 | shape: [3, 84, 84]
21 | type: rgb
22 | robot1_eye_in_hand_image:
23 | shape: [3, 84, 84]
24 | type: rgb
25 | robot1_eef_pos:
26 | shape: [3]
27 | # type default: low_dim
28 | robot1_eef_quat:
29 | shape: [4]
30 | robot1_gripper_qpos:
31 | shape: [2]
32 | action:
33 | shape: [20]
34 |
35 | task_name: &task_name transport
36 | dataset_type: &dataset_type ph
37 | dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
38 | abs_action: &abs_action True
39 |
40 | env_runner:
41 | _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
42 | dataset_path: *dataset_path
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 | max_steps: 700
51 | n_obs_steps: ${n_obs_steps}
52 | n_action_steps: ${n_action_steps}
53 | render_obs_key: 'shouldercamera0_image'
54 | fps: 10
55 | crf: 22
56 | past_action: ${past_action_visible}
57 | abs_action: *abs_action
58 | tqdm_interval_sec: 1.0
59 | n_envs: 28
60 | # evaluation at this config requires a 16 core 64GB instance.
61 |
62 | dataset:
63 | _target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
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:'${n_action_steps}-1'}
69 | n_obs_steps: ${dataset_obs_steps}
70 | abs_action: *abs_action
71 | rotation_rep: 'rotation_6d'
72 | use_legacy_normalizer: False
73 | use_cache: True
74 | seed: 42
75 | val_ratio: 0.02
76 |
--------------------------------------------------------------------------------
/diffusion_policy/config/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:'${n_action_steps}-1'}
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/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:'${n_action_steps}-1'}
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/train_robomimic_image_workspace.yaml:
--------------------------------------------------------------------------------
1 | defaults:
2 | - _self_
3 | - task: lift_image
4 |
5 | name: train_robomimic_image
6 | _target_: diffusion_policy.workspace.train_robomimic_image_workspace.TrainRobomimicImageWorkspace
7 |
8 | task_name: ${task.name}
9 | shape_meta: ${task.shape_meta}
10 | exp_name: "default"
11 |
12 | horizon: &horizon 10
13 | n_obs_steps: 1
14 | n_action_steps: 1
15 | n_latency_steps: 0
16 | dataset_obs_steps: *horizon
17 | past_action_visible: False
18 | keypoint_visible_rate: 1.0
19 |
20 | policy:
21 | _target_: diffusion_policy.policy.robomimic_image_policy.RobomimicImagePolicy
22 | shape_meta: ${shape_meta}
23 | algo_name: bc_rnn
24 | obs_type: image
25 | # oc.select resolver: key, default
26 | task_name: ${oc.select:task.task_name,lift}
27 | dataset_type: ${oc.select:task.dataset_type,ph}
28 | crop_shape: [76,76]
29 |
30 | dataloader:
31 | batch_size: 64
32 | num_workers: 16
33 | shuffle: True
34 | pin_memory: True
35 | persistent_workers: False
36 |
37 | val_dataloader:
38 | batch_size: 64
39 | num_workers: 16
40 | shuffle: False
41 | pin_memory: True
42 | persistent_workers: False
43 |
44 | training:
45 | device: "cuda:0"
46 | seed: 42
47 | debug: False
48 | resume: True
49 | # optimization
50 | num_epochs: 3050
51 | # training loop control
52 | # in epochs
53 | rollout_every: 50
54 | checkpoint_every: 50
55 | val_every: 1
56 | sample_every: 5
57 | # steps per epoch
58 | max_train_steps: null
59 | max_val_steps: null
60 | # misc
61 | tqdm_interval_sec: 1.0
62 |
63 | logging:
64 | project: diffusion_policy_debug
65 | resume: True
66 | mode: online
67 | name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}
68 | tags: ["${name}", "${task_name}", "${exp_name}"]
69 | id: null
70 | group: null
71 |
72 | checkpoint:
73 | topk:
74 | monitor_key: test_mean_score
75 | mode: max
76 | k: 5
77 | format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
78 | save_last_ckpt: True
79 | save_last_snapshot: False
80 |
81 | multi_run:
82 | run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
83 | wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}
84 |
85 | hydra:
86 | job:
87 | override_dirname: ${name}
88 | run:
89 | dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
90 | sweep:
91 | dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
92 | subdir: ${hydra.job.num}
93 |
--------------------------------------------------------------------------------
/diffusion_policy/config/train_robomimic_lowdim_workspace.yaml:
--------------------------------------------------------------------------------
1 | defaults:
2 | - _self_
3 | - task: pusht_lowdim
4 |
5 | name: train_robomimic_lowdim
6 | _target_: diffusion_policy.workspace.train_robomimic_lowdim_workspace.TrainRobomimicLowdimWorkspace
7 |
8 | obs_dim: ${task.obs_dim}
9 | action_dim: ${task.action_dim}
10 | transition_dim: "${eval: ${task.obs_dim} + ${task.action_dim}}"
11 | task_name: ${task.name}
12 | exp_name: "default"
13 |
14 | horizon: 10
15 | n_obs_steps: 1
16 | n_action_steps: 1
17 | n_latency_steps: 0
18 | past_action_visible: False
19 | keypoint_visible_rate: 1.0
20 |
21 | policy:
22 | _target_: diffusion_policy.policy.robomimic_lowdim_policy.RobomimicLowdimPolicy
23 | action_dim: ${action_dim}
24 | obs_dim: ${obs_dim}
25 | algo_name: bc_rnn
26 | obs_type: low_dim
27 | # oc.select resolver: key, default
28 | task_name: ${oc.select:task.task_name,lift}
29 | dataset_type: ${oc.select:task.dataset_type,ph}
30 |
31 | dataloader:
32 | batch_size: 256
33 | num_workers: 1
34 | shuffle: True
35 | pin_memory: True
36 | persistent_workers: False
37 |
38 | val_dataloader:
39 | batch_size: 256
40 | num_workers: 1
41 | shuffle: False
42 | pin_memory: True
43 | persistent_workers: False
44 |
45 | training:
46 | device: "cuda:0"
47 | seed: 42
48 | debug: False
49 | resume: True
50 | # optimization
51 | num_epochs: 5000
52 | # training loop control
53 | # in epochs
54 | rollout_every: 50
55 | checkpoint_every: 50
56 | val_every: 1
57 | # steps per epoch
58 | max_train_steps: null
59 | max_val_steps: null
60 | # misc
61 | tqdm_interval_sec: 1.0
62 |
63 | logging:
64 | project: diffusion_policy_debug
65 | resume: True
66 | mode: online
67 | name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}
68 | tags: ["${name}", "${task_name}", "${exp_name}"]
69 | id: null
70 | group: null
71 |
72 | checkpoint:
73 | topk:
74 | monitor_key: test_mean_score
75 | mode: max
76 | k: 5
77 | format_str: 'epoch={epoch:04d}-test_mean_score={test_mean_score:.3f}.ckpt'
78 | save_last_ckpt: True
79 | save_last_snapshot: False
80 |
81 | multi_run:
82 | run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
83 | wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}
84 |
85 | hydra:
86 | job:
87 | override_dirname: ${name}
88 | run:
89 | dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
90 | sweep:
91 | dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
92 | subdir: ${hydra.job.num}
93 |
--------------------------------------------------------------------------------
/diffusion_policy/config/train_robomimic_real_image_workspace.yaml:
--------------------------------------------------------------------------------
1 | defaults:
2 | - _self_
3 | - task: real_pusht_image
4 |
5 | name: train_robomimic_image
6 | _target_: diffusion_policy.workspace.train_robomimic_image_workspace.TrainRobomimicImageWorkspace
7 |
8 | task_name: ${task.name}
9 | shape_meta: ${task.shape_meta}
10 | exp_name: "default"
11 |
12 | horizon: &horizon 10
13 | n_obs_steps: 1
14 | n_action_steps: 1
15 | n_latency_steps: 1
16 | dataset_obs_steps: *horizon
17 | past_action_visible: False
18 | keypoint_visible_rate: 1.0
19 |
20 | policy:
21 | _target_: diffusion_policy.policy.robomimic_image_policy.RobomimicImagePolicy
22 | shape_meta: ${shape_meta}
23 | algo_name: bc_rnn
24 | obs_type: image
25 | # oc.select resolver: key, default
26 | task_name: ${oc.select:task.task_name,tool_hang}
27 | dataset_type: ${oc.select:task.dataset_type,ph}
28 | crop_shape: [216, 288] # ch, cw 320x240 90%
29 |
30 | dataloader:
31 | batch_size: 32
32 | num_workers: 8
33 | shuffle: True
34 | pin_memory: True
35 | persistent_workers: True
36 |
37 | val_dataloader:
38 | batch_size: 32
39 | num_workers: 1
40 | shuffle: False
41 | pin_memory: True
42 | persistent_workers: False
43 |
44 | training:
45 | device: "cuda:0"
46 | seed: 42
47 | debug: False
48 | resume: True
49 | # optimization
50 | num_epochs: 1000
51 | # training loop control
52 | # in epochs
53 | rollout_every: 50
54 | checkpoint_every: 50
55 | val_every: 1
56 | sample_every: 5
57 | # steps per epoch
58 | max_train_steps: null
59 | max_val_steps: null
60 | # misc
61 | tqdm_interval_sec: 1.0
62 |
63 | logging:
64 | project: diffusion_policy_debug
65 | resume: True
66 | mode: online
67 | name: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}
68 | tags: ["${name}", "${task_name}", "${exp_name}"]
69 | id: null
70 | group: null
71 |
72 | checkpoint:
73 | topk:
74 | monitor_key: train_loss
75 | mode: min
76 | k: 5
77 | format_str: 'epoch={epoch:04d}-train_loss={train_loss:.3f}.ckpt'
78 | save_last_ckpt: True
79 | save_last_snapshot: False
80 |
81 | multi_run:
82 | run_dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
83 | wandb_name_base: ${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}
84 |
85 | hydra:
86 | job:
87 | override_dirname: ${name}
88 | run:
89 | dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
90 | sweep:
91 | dir: data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}
92 | subdir: ${hydra.job.num}
93 |
--------------------------------------------------------------------------------
/diffusion_policy/dataset/base_dataset.py:
--------------------------------------------------------------------------------
1 | from typing import Dict
2 |
3 | import torch
4 | import torch.nn
5 | from diffusion_policy.model.common.normalizer import LinearNormalizer
6 |
7 | class BaseLowdimDataset(torch.utils.data.Dataset):
8 | def get_validation_dataset(self) -> 'BaseLowdimDataset':
9 | # return an empty dataset by default
10 | return BaseLowdimDataset()
11 |
12 | def get_normalizer(self, **kwargs) -> LinearNormalizer:
13 | raise NotImplementedError()
14 |
15 | def get_all_actions(self) -> torch.Tensor:
16 | raise NotImplementedError()
17 |
18 | def __len__(self) -> int:
19 | return 0
20 |
21 | def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
22 | """
23 | output:
24 | obs: T, Do
25 | action: T, Da
26 | """
27 | raise NotImplementedError()
28 |
29 |
30 | class BaseImageDataset(torch.utils.data.Dataset):
31 | def get_validation_dataset(self) -> 'BaseLowdimDataset':
32 | # return an empty dataset by default
33 | return BaseImageDataset()
34 |
35 | def get_normalizer(self, **kwargs) -> LinearNormalizer:
36 | raise NotImplementedError()
37 |
38 | def get_all_actions(self) -> torch.Tensor:
39 | raise NotImplementedError()
40 |
41 | def __len__(self) -> int:
42 | return 0
43 |
44 | def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
45 | """
46 | output:
47 | obs:
48 | key: T, *
49 | action: T, Da
50 | """
51 | raise NotImplementedError()
52 |
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/bimanual_viperx_ee_transfer_cube.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 |
40 |
41 |
42 |
43 |
44 |
45 |
46 |
47 |
48 |
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/scene.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
17 |
18 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 |
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/singlearm_viperx_ee_transfer_cube.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 |
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/singlearm_viperx_transfer_cube.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/aloha/assets/tabletop.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/tabletop.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_10_custom_finger_left.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_10_custom_finger_left.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_10_custom_finger_right.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_10_custom_finger_right.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_10_gripper_finger.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_10_gripper_finger.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_11_ar_tag.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_11_ar_tag.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_1_base.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_1_base.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_2_shoulder.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_2_shoulder.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_3_upper_arm.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_3_upper_arm.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_4_upper_forearm.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_4_upper_forearm.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_5_lower_forearm.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_5_lower_forearm.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_6_wrist.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_6_wrist.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_7_gripper.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_7_gripper.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_8_gripper_prop.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_8_gripper_prop.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_9_gripper_bar.stl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/diffusion_policy/env/aloha/assets/vx300s_9_gripper_bar.stl
--------------------------------------------------------------------------------
/diffusion_policy/env/aloha/assets/vx300s_dependencies.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
--------------------------------------------------------------------------------
/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/suction/suction-base.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 |
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 |
68 |
69 |
70 |
--------------------------------------------------------------------------------
/diffusion_policy/env/block_pushing/assets/suction/suction-head.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 |
68 |
69 |
70 |
71 |
--------------------------------------------------------------------------------
/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/block_pushing/utils/pose3d.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 | """A simple 6DOF pose container.
17 | """
18 |
19 | import dataclasses
20 | import numpy as np
21 | from scipy.spatial import transform
22 |
23 |
24 | class NoCopyAsDict(object):
25 | """Base class for dataclasses. Avoids a copy in the asdict() call."""
26 |
27 | def asdict(self):
28 | """Replacement for dataclasses.asdict.
29 |
30 | TF Dataset does not handle dataclasses.asdict, which uses copy.deepcopy when
31 | setting values in the output dict. This causes issues with tf.Dataset.
32 | Instead, shallow copy contents.
33 |
34 | Returns:
35 | dict containing contents of dataclass.
36 | """
37 | return {k.name: getattr(self, k.name) for k in dataclasses.fields(self)}
38 |
39 |
40 | @dataclasses.dataclass
41 | class Pose3d(NoCopyAsDict):
42 | """Simple container for translation and rotation."""
43 |
44 | rotation: transform.Rotation
45 | translation: np.ndarray
46 |
47 | @property
48 | def vec7(self):
49 | return np.concatenate([self.translation, self.rotation.as_quat()])
50 |
51 | def serialize(self):
52 | return {
53 | "rotation": self.rotation.as_quat().tolist(),
54 | "translation": self.translation.tolist(),
55 | }
56 |
57 | @staticmethod
58 | def deserialize(data):
59 | return Pose3d(
60 | rotation=transform.Rotation.from_quat(data["rotation"]),
61 | translation=np.array(data["translation"]),
62 | )
63 |
64 | def __eq__(self, other):
65 | return np.array_equal(
66 | self.rotation.as_quat(), other.rotation.as_quat()
67 | ) and np.array_equal(self.translation, other.translation)
68 |
69 | def __ne__(self, other):
70 | return not self.__eq__(other)
71 |
--------------------------------------------------------------------------------
/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/microwave_chain.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 |
--------------------------------------------------------------------------------
/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/assets/slidecabinet_chain.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 |
--------------------------------------------------------------------------------
/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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/pusht/pusht_image_env.py:
--------------------------------------------------------------------------------
1 | from gym import spaces
2 | from diffusion_policy.env.pusht.pusht_env import PushTEnv
3 | import numpy as np
4 | import cv2
5 |
6 | class PushTImageEnv(PushTEnv):
7 | metadata = {"render.modes": ["rgb_array"], "video.frames_per_second": 10}
8 |
9 | def __init__(self,
10 | legacy=False,
11 | block_cog=None,
12 | damping=None,
13 | render_size=96,
14 | perturb=0.0):
15 | super().__init__(
16 | legacy=legacy,
17 | block_cog=block_cog,
18 | damping=damping,
19 | render_size=render_size,
20 | render_action=False,
21 | perturb=perturb)
22 | ws = self.window_size
23 | self.observation_space = spaces.Dict({
24 | 'image': spaces.Box(
25 | low=0,
26 | high=1,
27 | shape=(3,render_size,render_size),
28 | dtype=np.float32
29 | ),
30 | 'agent_pos': spaces.Box(
31 | low=0,
32 | high=ws,
33 | shape=(2,),
34 | dtype=np.float32
35 | )
36 | })
37 | self.render_cache = None
38 |
39 | def _get_obs(self):
40 | img = super()._render_frame(mode='rgb_array')
41 |
42 | agent_pos = np.array(self.agent.position)
43 | img_obs = np.moveaxis(img.astype(np.float32) / 255, -1, 0)
44 | obs = {
45 | 'image': img_obs,
46 | 'agent_pos': agent_pos
47 | }
48 |
49 | # draw action
50 | if self.latest_action is not None:
51 | action = np.array(self.latest_action)
52 | coord = (action / 512 * 96).astype(np.int32)
53 | marker_size = int(8/96*self.render_size)
54 | thickness = int(1/96*self.render_size)
55 | cv2.drawMarker(img, coord,
56 | color=(255,0,0), markerType=cv2.MARKER_CROSS,
57 | markerSize=marker_size, thickness=thickness)
58 | self.render_cache = img
59 |
60 | return obs
61 |
62 | def render(self, mode):
63 | assert mode == 'rgb_array'
64 |
65 | if self.render_cache is None:
66 | self._get_obs()
67 |
68 | return self.render_cache
69 |
--------------------------------------------------------------------------------
/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/long-context-dp/ldp/1cac600a71e09079f49ca09262f2b3d06b9b305a/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()
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/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]) -> Dict[str, torch.Tensor]:
11 | """
12 | obs_dict:
13 | str: B,To,*
14 | return: B,Ta,Da
15 | """
16 | raise NotImplementedError()
17 |
18 | # reset state for stateful policies
19 | def reset(self):
20 | pass
21 |
22 | # ========== training ===========
23 | # no standard training interface except setting normalizer
24 | def set_normalizer(self, normalizer: LinearNormalizer):
25 | raise NotImplementedError()
26 |
--------------------------------------------------------------------------------
/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/main_eval.py:
--------------------------------------------------------------------------------
1 | from r2d2.controllers.oculus_controller import VRPolicy
2 | from r2d2.robot_env import RobotEnv
3 | from r2d2.user_interface.data_collector import DataCollecter
4 | from r2d2.user_interface.gui_nocalib import RobotGUI
5 | from r2d2.evaluation.policy_wrapper_marcel import PolicyWrapper
6 |
7 | # Make the robot env
8 | env = RobotEnv()
9 | controller = VRPolicy()
10 | # Make the data collector
11 | path = ".ckpt"
12 | policy = torch.load(path)
13 |
14 | policy_timestep_filtering_kwargs = {
15 | "action_space": "cartesian_position"
16 | "robot_state_keys": ["cartesian_position", "gripper_position"],
17 |
18 | }
19 |
20 | wrapped_policy = PolicyWrapper(
21 | policy=policy,
22 | timestep_filtering_kwargs=policy_timestep_filtering_kwargs,
23 | image_transform_kwargs=policy_image_transform_kwargs,
24 | eval_mode=True,
25 | )
26 |
27 | # Launch GUI #
28 | data_collector = DataCollecter(
29 | env=env,
30 | controller=controller,
31 | policy=wrapped_policy,
32 | save_traj_dir=log_dir,
33 | save_data=True,
34 | )
35 |
36 | # Make the GUI
37 | user_interface = RobotGUI(robot=data_collector)
38 |
--------------------------------------------------------------------------------
/diffusion_policy/real/main_nocalib.py:
--------------------------------------------------------------------------------
1 | from r2d2.controllers.oculus_controller import VRPolicy
2 | from r2d2.robot_env import RobotEnv
3 | from r2d2.user_interface.data_collector import DataCollecter
4 | from r2d2.user_interface.gui_nocalib import RobotGUI
5 |
6 | # Make the robot env
7 | env = RobotEnv()
8 | controller = VRPolicy()
9 | # Make the data collector
10 | data_collector = DataCollecter(env=env, controller=controller)
11 | #data_collector.collect_trajectory()
12 | # Make the GUI
13 | user_interface = RobotGUI(robot=data_collector)
14 |
15 |
--------------------------------------------------------------------------------
/diffusion_policy/real/policy_wrapper.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import torch
3 |
4 | from r2d2.data_processing.timestep_processing import TimestepProcesser
5 |
6 |
7 | def converter_helper(data, batchify=True):
8 | if torch.is_tensor(data):
9 | pass
10 | elif isinstance(data, np.ndarray):
11 | data = torch.from_numpy(data)
12 | else:
13 | raise ValueError
14 |
15 | if batchify:
16 | data = data.unsqueeze(0)
17 | return data
18 |
19 |
20 | def np_dict_to_torch_dict(np_dict, batchify=True):
21 | torch_dict = {}
22 |
23 | for key in np_dict:
24 | curr_data = np_dict[key]
25 | if isinstance(curr_data, dict):
26 | torch_dict[key] = np_dict_to_torch_dict(curr_data)
27 | elif isinstance(curr_data, np.ndarray) or torch.is_tensor(curr_data):
28 | torch_dict[key] = converter_helper(curr_data, batchify=batchify)
29 | elif isinstance(curr_data, list):
30 | torch_dict[key] = [converter_helper(d, batchify=batchify) for d in curr_data]
31 | else:
32 | raise ValueError
33 |
34 | return torch_dict
35 |
36 |
37 | class PolicyWrapper:
38 | def __init__(self, policy, timestep_filtering_kwargs, image_transform_kwargs, eval_mode=True):
39 | self.policy = policy
40 |
41 | if eval_mode:
42 | self.policy.eval()
43 | else:
44 | self.policy.train()
45 |
46 | self.timestep_processor = TimestepProcesser(
47 | ignore_action=True, **timestep_filtering_kwargs, image_transform_kwargs=image_transform_kwargs
48 | )
49 |
50 | def forward(self, observation):
51 | timestep = {"observation": observation}
52 | processed_timestep = self.timestep_processor.forward(timestep)
53 | torch_timestep = np_dict_to_torch_dict(processed_timestep)
54 | action = self.policy(torch_timestep)[0]
55 | np_action = action.detach().numpy()
56 |
57 | # a_star = np.cumsum(processed_timestep['observation']['state']) / 7
58 | # print('Policy Action: ', np_action)
59 | # print('Expert Action: ', a_star)
60 | # print('Error: ', np.abs(a_star - np_action).mean())
61 |
62 | # import pdb; pdb.set_trace()
63 | return np_action
64 |
--------------------------------------------------------------------------------
/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 | from typing import Dict, Callable, Tuple
2 | import numpy as np
3 | from diffusion_policy.common.cv2_util import get_image_transform
4 |
5 | def get_real_obs_dict(
6 | env_obs: Dict[str, np.ndarray],
7 | shape_meta: dict,
8 | ) -> Dict[str, np.ndarray]:
9 | obs_dict_np = dict()
10 | obs_shape_meta = shape_meta['obs']
11 | for key, attr in obs_shape_meta.items():
12 | type = attr.get('type', 'low_dim')
13 | shape = attr.get('shape')
14 | if type == 'rgb':
15 | this_imgs_in = env_obs[key]
16 | t,hi,wi,ci = this_imgs_in.shape
17 | co,ho,wo = shape
18 | assert ci == co
19 | out_imgs = this_imgs_in
20 | if (ho != hi) or (wo != wi) or (this_imgs_in.dtype == np.uint8):
21 | tf = get_image_transform(
22 | input_res=(wi,hi),
23 | output_res=(wo,ho),
24 | bgr_to_rgb=False)
25 | out_imgs = np.stack([tf(x) for x in this_imgs_in])
26 | if this_imgs_in.dtype == np.uint8:
27 | out_imgs = out_imgs.astype(np.float32) / 255
28 | # THWC to TCHW
29 | obs_dict_np[key] = np.moveaxis(out_imgs,-1,1)
30 | elif type == 'low_dim':
31 | this_data_in = env_obs[key]
32 | if 'pose' in key and shape == (2,):
33 | # take X,Y coordinates
34 | this_data_in = this_data_in[...,[0,1]]
35 | obs_dict_np[key] = this_data_in
36 | return obs_dict_np
37 |
38 |
39 | def get_real_obs_resolution(
40 | shape_meta: dict
41 | ) -> Tuple[int, int]:
42 | out_res = None
43 | obs_shape_meta = shape_meta['obs']
44 | for key, attr in obs_shape_meta.items():
45 | type = attr.get('type', 'low_dim')
46 | shape = attr.get('shape')
47 | if type == 'rgb':
48 | co,ho,wo = shape
49 | if out_res is None:
50 | out_res = (wo, ho)
51 | assert out_res == (wo, ho)
52 | return out_res
53 |
--------------------------------------------------------------------------------
/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/generate_bet_blockpush.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 tqdm import tqdm
15 | from diffusion_policy.common.replay_buffer import ReplayBuffer
16 | from tf_agents.environments.wrappers import TimeLimit
17 | from tf_agents.environments.gym_wrapper import GymWrapper
18 | from tf_agents.trajectories.time_step import StepType
19 | from diffusion_policy.env.block_pushing.block_pushing_multimodal import BlockPushMultimodal
20 | from diffusion_policy.env.block_pushing.block_pushing import BlockPush
21 | from diffusion_policy.env.block_pushing.oracles.multimodal_push_oracle import MultimodalOrientedPushOracle
22 |
23 | @click.command()
24 | @click.option('-o', '--output', required=True)
25 | @click.option('-n', '--n_episodes', default=1000)
26 | @click.option('-c', '--chunk_length', default=-1)
27 | def main(output, n_episodes, chunk_length):
28 |
29 | buffer = ReplayBuffer.create_empty_numpy()
30 | env = TimeLimit(GymWrapper(BlockPushMultimodal()), duration=350)
31 | for i in tqdm(range(n_episodes)):
32 | print(i)
33 | obs_history = list()
34 | action_history = list()
35 |
36 | env.seed(i)
37 | policy = MultimodalOrientedPushOracle(env)
38 | time_step = env.reset()
39 | policy_state = policy.get_initial_state(1)
40 | while True:
41 | action_step = policy.action(time_step, policy_state)
42 | obs = np.concatenate(list(time_step.observation.values()), axis=-1)
43 | action = action_step.action
44 | obs_history.append(obs)
45 | action_history.append(action)
46 |
47 | if time_step.step_type == 2:
48 | break
49 |
50 | # state = env.wrapped_env().gym.get_pybullet_state()
51 | time_step = env.step(action)
52 | obs_history = np.array(obs_history)
53 | action_history = np.array(action_history)
54 |
55 | episode = {
56 | 'obs': obs_history,
57 | 'action': action_history
58 | }
59 | buffer.add_episode(episode)
60 |
61 | buffer.save_to_path(output, chunk_length=chunk_length)
62 |
63 | if __name__ == '__main__':
64 | main()
65 |
--------------------------------------------------------------------------------
/diffusion_policy/scripts/real_dataset_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 | import zarr
13 | import cv2
14 | import threadpoolctl
15 | from diffusion_policy.real_world.real_data_conversion import real_data_to_replay_buffer
16 |
17 | @click.command()
18 | @click.option('--input', '-i', required=True)
19 | @click.option('--output', '-o', default=None)
20 | @click.option('--resolution', '-r', default='640x480')
21 | @click.option('--n_decoding_threads', '-nd', default=-1, type=int)
22 | @click.option('--n_encoding_threads', '-ne', default=-1, type=int)
23 | def main(input, output, resolution, n_decoding_threads, n_encoding_threads):
24 | out_resolution = tuple(int(x) for x in resolution.split('x'))
25 | input = pathlib.Path(os.path.expanduser(input))
26 | in_zarr_path = input.joinpath('replay_buffer.zarr')
27 | in_video_dir = input.joinpath('videos')
28 | assert in_zarr_path.is_dir()
29 | assert in_video_dir.is_dir()
30 | if output is None:
31 | output = input.joinpath(resolution + '.zarr.zip')
32 | else:
33 | output = pathlib.Path(os.path.expanduser(output))
34 |
35 | if output.exists():
36 | click.confirm('Output path already exists! Overrite?', abort=True)
37 |
38 | cv2.setNumThreads(1)
39 | with threadpoolctl.threadpool_limits(1):
40 | replay_buffer = real_data_to_replay_buffer(
41 | dataset_path=str(input),
42 | out_resolutions=out_resolution,
43 | n_decoding_threads=n_decoding_threads,
44 | n_encoding_threads=n_encoding_threads
45 | )
46 |
47 | print('Saving to disk')
48 | if output.suffix == '.zip':
49 | with zarr.ZipStore(output) as zip_store:
50 | replay_buffer.save_to_store(
51 | store=zip_store
52 | )
53 | else:
54 | with zarr.DirectoryStore(output) as store:
55 | replay_buffer.save_to_store(
56 | store=store
57 | )
58 |
59 | if __name__ == '__main__':
60 | main()
61 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/hsic.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | def rbf_kernel(X, sigma=None):
4 | """ Compute RBF kernel matrix for X """
5 | pairwise_sq_dists = torch.cdist(X, X, p=2) ** 2
6 | if sigma is None:
7 | sigma = torch.median(pairwise_sq_dists) # Median heuristic
8 | K = torch.exp(-pairwise_sq_dists / (2 * sigma ** 2))
9 | return K
10 |
11 | def center_kernel(K):
12 | """ Center the kernel matrix K """
13 | n = K.shape[0]
14 | H = torch.eye(n, device=K.device) - (1.0 / n) * torch.ones((n, n), device=K.device)
15 | return H @ K @ H
16 |
17 | def hsic(X, Y):
18 | """ Compute HSIC between X and Y """
19 | Kx = center_kernel(rbf_kernel(X))
20 | Ky = center_kernel(rbf_kernel(Y))
21 | return torch.trace(Kx @ Ky) / (X.shape[0] - 1) ** 2
22 |
23 | def batch_hsic(actions):
24 | """ Compute HSIC for each batch """
25 | B, N, D = actions.shape # B=batch, N=sequence length, D=7
26 | hsic_values = torch.zeros(B, device=actions.device)
27 |
28 | for b in range(B):
29 | X = actions[b] # Select the batch element
30 | Y = X # Can also compare with another variable if needed
31 | hsic_values[b] = hsic(X, Y)
32 |
33 | return hsic_values # Shape: (B,)
34 |
35 |
36 | if __name__ == "__main__":
37 | # Example usage
38 | B, N, D = 8, 50, 7 # Example batch size, sequence length, action dimensions
39 | actions = torch.randn(B, N, D) # Random action tensor
40 | hsic_values = batch_hsic(actions)
41 | print(hsic_values) # Tensor of HSIC values for each batch
42 |
--------------------------------------------------------------------------------
/perturbs/chunk_1.yaml:
--------------------------------------------------------------------------------
1 | chunk: 1
2 |
--------------------------------------------------------------------------------
/perturbs/chunk_2.yaml:
--------------------------------------------------------------------------------
1 | chunk: 2
2 |
--------------------------------------------------------------------------------
/perturbs/chunk_4.yaml:
--------------------------------------------------------------------------------
1 | chunk: 4
2 |
--------------------------------------------------------------------------------
/perturbs/chunk_8.yaml:
--------------------------------------------------------------------------------
1 | chunk: 8
2 |
--------------------------------------------------------------------------------
/perturbs/none.yaml:
--------------------------------------------------------------------------------
1 | dummy: 1
2 |
--------------------------------------------------------------------------------
/rollouts/merge_actions.py:
--------------------------------------------------------------------------------
1 | import pickle
2 | import numpy as np
3 | import os
4 | import argparse
5 |
6 | # Set up argument parser
7 | parser = argparse.ArgumentParser(description='Merge .pkl files by timestep and batch.')
8 | parser.add_argument('folder_path', type=str, help='Path to the folder containing .pkl files')
9 |
10 | # Parse the arguments
11 | args = parser.parse_args()
12 |
13 | # Directory path where the .pkl files are located
14 | folder_path = args.folder_path
15 |
16 | # List to store all the actions from the pkl files
17 | all_actions_by_timestep = []
18 |
19 | # Check if the folder path exists
20 | if not os.path.exists(folder_path):
21 | print(f"Error: The directory {folder_path} does not exist.")
22 | exit(1)
23 |
24 | print_shape = True
25 | # Iterate through the pkl files and load the data
26 | for filename in sorted(os.listdir(folder_path)):
27 | if filename.endswith('.pkl'):
28 | file_path = os.path.join(folder_path, filename)
29 |
30 | # Load the pkl file
31 | with open(file_path, 'rb') as f:
32 | actions = pickle.load(f) # Assuming actions is a list of (28, 1, 10) shaped arrays
33 |
34 | # Convert list to a numpy array of shape (num_timesteps, 1, 10)
35 |
36 | actions = np.array(actions).squeeze() # Shape should be (500, 28, 1, 10)
37 | # B, T, N, D = actions.shape # B=batch, T=time, N=sub-batch, D=dim
38 | # actions = actions.transpose(0, 2, 1, 3).reshape(B * N, T, D)
39 | if print_shape:
40 | print(actions.shape)
41 | print_shape = False
42 |
43 | # Reshape the actions to (28, 500, 10)
44 | actions_by_timestep = actions.transpose(1, 0, 2) # Transpose to shape (28, 500, 10)
45 |
46 | # Append the actions by timestep to the list
47 | all_actions_by_timestep.append(actions_by_timestep)
48 |
49 | # Now concatenate them by batch
50 | merged_actions = np.concatenate(all_actions_by_timestep, axis=0) # Shape: (46, 500, 10)
51 |
52 | # Define the path for the merged actions file in the same folder
53 | output_file_path = os.path.join(folder_path, 'merged_actions.pkl')
54 |
55 | # Save the result to merged_actions.pkl
56 | with open(output_file_path, 'wb') as f:
57 | pickle.dump(merged_actions, f)
58 |
59 | print(f"Merging complete. The result has been saved as '{output_file_path}'.")
--------------------------------------------------------------------------------
/run_gather_rollouts.sh:
--------------------------------------------------------------------------------
1 | checkpoint=$2
2 | name=$1
3 |
4 | if [ -n "$3" ]; then
5 | python gather_rollouts.py -c $2 -n 400 -o "rollouts/$1" -e $3
6 | else
7 | python gather_rollouts.py -c $2 -n 400 -o "rollouts/$1"
8 | fi
--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
1 | from setuptools import setup, find_packages
2 |
3 | setup(
4 | name = 'diffusion_policy',
5 | packages = find_packages(),
6 | )
7 |
--------------------------------------------------------------------------------
/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 | print(OmegaConf.to_yaml(cfg))
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_rollouts.sh:
--------------------------------------------------------------------------------
1 | name=$1
2 | checkpoint=$2
3 |
4 | if [ -n "$3" ]; then
5 | python rollouts_via_policy.py --disable_tqdm --transport $1 $2
6 | else
7 | python rollouts_via_policy.py --disable_tqdm $1 $2
8 | fi
9 |
--------------------------------------------------------------------------------
/transformer_eval.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | # Parse flagged arguments
4 | while [[ "$#" -gt 0 ]]; do
5 | case $1 in
6 | --checkpoint) ckpt="$2"; shift ;;
7 | --name) name="$2"; shift ;;
8 | --perturb) perturb="$2"; shift ;;
9 | --samples) samples="$2"; shift ;;
10 | -h|--help)
11 | echo "Usage: $0 --checkpoint PATH --name NAME --perturb VALUE --samples N"
12 | exit 0
13 | ;;
14 | *)
15 | echo "Unknown parameter passed: $1"
16 | exit 1 ;;
17 | esac
18 | shift
19 | done
20 |
21 | # Validate required arguments
22 | if [[ -z "$ckpt" || -z "$name" || -z "$perturb" || -z "$samples" ]]; then
23 | echo "Missing required arguments."
24 | echo "Run with -h or --help to see usage."
25 | exit 1
26 | fi
27 |
28 | hhmmss=$(date +"%H%M%S")
29 | export MUJOCO_GL="egl"
30 |
31 | python eval.py \
32 | --checkpoint "$ckpt" \
33 | -o "data/${name}_$hhmmss" \
34 | -p "$perturb" \
35 | -n "$samples"
36 |
37 |
--------------------------------------------------------------------------------