├── .gitignore ├── .gitmodules ├── LICENSE ├── README.md ├── benchmark └── cfgs │ ├── cfg1.yaml │ ├── cfg2.yaml │ ├── cfg3.yaml │ └── cfg4.yaml ├── datasets └── load_test_dataset.py ├── diffusion ├── __init__.py ├── diffusion.py ├── gaussian.py └── models │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-38.pyc │ ├── blocks.cpython-38.pyc │ └── temporalunet.cpython-38.pyc │ ├── blocks.py │ └── temporalunet.py ├── guides ├── cfgs │ ├── guide1.yaml │ ├── guide10.yaml │ ├── guide11.yaml │ ├── guide12.yaml │ ├── guide13.yaml │ ├── guide14.yaml │ ├── guide15.yaml │ ├── guide16.yaml │ ├── guide17.yaml │ ├── guide18.yaml │ ├── guide2.yaml │ ├── guide21.yaml │ ├── guide3.yaml │ ├── guide4.yaml │ ├── guide5.yaml │ └── guide9.yaml └── guide_cfg.py ├── infer_serial.py ├── lib ├── __init__.py ├── environment.py ├── guide.py └── metrics.py ├── mpinets ├── __init__.py ├── geometry.py ├── loss.py ├── metrics.py ├── model.py ├── third_party │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── sparc.cpython-37.pyc │ │ └── sparc.cpython-38.pyc │ └── sparc.py ├── types.py └── utils.py ├── results ├── teaser_compressed.gif └── teaser_compressed.mp4 └── urdfs ├── panda.urdf └── workspace.urdf /.gitignore: -------------------------------------------------------------------------------- 1 | models/* 2 | datasets/global_solvable_problems.pkl 3 | datasets/hybrid_solvable_problems.pkl 4 | datasets/both_solvable_problems.pkl 5 | datasets/__pycache__/* 6 | wandb/* 7 | mpinets/__pycache__/* 8 | mpinets/__pycache__/* 9 | diffusion/__pycache__/* 10 | lib/__pycache__/* 11 | mpinets/__pycache__/* 12 | *.cpython-310.pyc 13 | *.zip -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "robofin"] 2 | path = robofin 3 | url = https://github.com/fishbotics/robofin.git 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Vishal Reddy Mandadi 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # EDMP: Ensemble-of-Costs Guided Diffusion For Motion Planning 2 | 3 | [Kallol Saha](https://www.linkedin.com/in/kallolsaha/) \*1, 4 | [Vishal Mandadi](https://vishal-2000.github.io/) \*1, 5 | [Jayaram Reddy](http://www.linkedin.com/in/jayaram6111997) \*1, 6 | [Ajit Srikanth](https://twitter.com/ajitsrikanth) 1, 7 | [Aditya Agarwal](https://skymanaditya1.github.io/) 2, 8 | [Bipasha Sen](https://bipashasen.github.io/) 2, 9 | [Arun Kumar Singh](https://tuit.ut.ee/en/content/arun-kumar-singh) 3, 10 | [Madhava Krishna](https://www.iiit.ac.in/people/faculty/mkrishna/) 1, 11 | 12 | 1International Institute of Information Technology, Hyderabad, 2MIT, 3University of Tartu 13 | 14 | \*denotes equal contribution 15 | 16 | This is the official implementation of the paper "EDMP: Ensemble-of-Costs Guided Diffusion For Motion Planning", which is currently under review 17 | 18 | ![teaser_compressed (1)](https://github.com/vishal-2000/EDMP/assets/56213799/1ea9d638-1231-4e7e-b283-f8a5eaa7ba73) 19 | 20 | 21 | 22 | For more details and results, visit our [project page](https://ensemble-of-costs-diffusion.github.io/) and read our [paper](https://arxiv.org/abs/2309.11414). 23 | 24 | ## Code Release plan 25 | 1. Serial Version 26 | 1. All the classifiers are run in serial 27 | 2. Release Status: Beta version out 28 | 2. Parallel version 29 | 1. Classifiers will be run in parallel 30 | 2. Release status: Will be out by 25th December, 2023 31 | 32 | ## Setting up 33 | 1. Clone this repository with all the submodules 34 | ```bash 35 | git clone --recurse-submodules https://github.com/vishal-2000/EDMP.git 36 | ``` 37 | 2. Move to robofin folder and change the branch to v0.0.1 38 | ```bash 39 | cd robofin 40 | git checkout v0.0.1 41 | ``` 42 | 3. Install robofin 43 | ```bash 44 | pip install geometrout==0.0.3.4 45 | pip install urchin 46 | cd robofin 47 | pip install -e . 48 | ``` 49 | 4. Install other necessary packages 50 | ```bash 51 | pip install torch torchvision h5py einops autolab_core wandb scipy 52 | ``` 53 | 5. Download the datasets from [link](https://drive.google.com/drive/folders/1PhNjMhYHWwq9IjHTeyR2ydqEhaHxBdUW?usp=drive_link) to './datasets' folder. This would look like: 54 | ```bash 55 | ./datasets/*_solvable_problems.pkl 56 | ``` 57 | Where * is a placeholder for all the three datasets - global, hybrid, and both 58 | 6. Download the models folder from [link](https://drive.google.com/drive/folders/10FAqqfazU35eLAs3wb_iGKcRs8e_t4gG?usp=sharing) and unzip it in the main folder. 59 | 7. Final directory structure must look like: 60 | ```bash 61 | . 62 | ├── benchmark 63 | │   └── cfgs 64 | ├── datasets 65 | │   └── __pycache__ 66 | ├── diffusion 67 | │   ├── models 68 | │   │   └── __pycache__ 69 | │   └── __pycache__ 70 | ├── guides 71 | │   └── cfgs 72 | ├── lib 73 | │   └── __pycache__ 74 | ├── models 75 | │   └── TemporalUNetModel255_N50 76 | ├── mpinets 77 | │   ├── __pycache__ 78 | │   └── third_party 79 | │   └── __pycache__ 80 | ├── robofin 81 | │   └── robofin 82 | │   ├── kinematics 83 | │   ├── pointcloud 84 | │   │   └── cache 85 | │   │   └── franka 86 | │   ├── standalone_meshes 87 | │   └── urdf 88 | │   ├── franka_panda 89 | │   │   ├── hd_meshes 90 | │   │   │   ├── collision 91 | │   │   │   └── visual 92 | │   │   └── meshes 93 | │   │   ├── collision 94 | │   │   └── visual 95 | │   └── panda_hand 96 | │   └── meshes 97 | │   ├── collision 98 | │   └── visual 99 | └── urdfs 100 | 101 | ``` 102 | ## Running Inference 103 | Inference configurations can be set using the config files placed in the benchmark folder (```./benchmark/cfgs/cfg1.yaml```). A custom config file can be created following syntax similar to the files in the ```benchmark/cfgs``` directory 104 | 105 | To run inference: 106 | ```bash 107 | python infer_serial.py -c 108 | ``` 109 | 110 | For example: 111 | ```bash 112 | python infer_serial.py -c ./benchmark/cfgs/cfg1.yaml 113 | ``` 114 | 115 | ## Adding custom guides 116 | New custom guides can be added into the ```./guides/cfgs/``` folder, following syntax similar to the other files in that folder. In order to include this guide during inference, please add the guide number (where the number must match the number on the file name, i.e., ./guides/cfgs/guide10.yaml has guide number == 10 (also change the index number in the file according to the guide number)). 117 | 118 | You can control the number of guides you want to run in parallel, and what guides you want to run in parallel using the config file in benchmark folder. For example, in ```./benchmark/cfgs/cfg1.yaml```, setting the guides parameter to the following will run the inference script with the guides 1, 2, and 3. 119 | ```yaml 120 | guides: [1, 2, 3] 121 | ``` 122 | 123 | ## Results 124 | - For replicating the results shown in the paper, please use cfg ```./benchmark/cfgs/cfg1.yaml``` 125 | 126 | ## Citation 127 | If you find our work useful in your research, please cite: 128 | ``` 129 | @misc{saha2023edmp, 130 | title={EDMP: Ensemble-of-costs-guided Diffusion for Motion Planning}, 131 | author={Kallol Saha and Vishal Mandadi and Jayaram Reddy and Ajit Srikanth and Aditya Agarwal and Bipasha Sen and Arun Singh and Madhava Krishna}, 132 | year={2023}, 133 | eprint={2309.11414}, 134 | archivePrefix={arXiv}, 135 | primaryClass={cs.RO} 136 | } 137 | ``` 138 | 139 | ## Contact 140 | 141 | Kallol Saha: ksaha@andrew.cmu.edu
142 | Vishal Mandadi: vishal.mandadi@alumni.iiit.ac.in
143 | Jayaram Reddy: jayaram.reddy@research.iiit.ac.in
144 | 145 | -------------------------------------------------------------------------------- /benchmark/cfgs/cfg1.yaml: -------------------------------------------------------------------------------- 1 | guide: 2 | guides: [1, 2, 3, 4, 5, 10, 11, 13, 14, 16, 18, 21] # [1, 2, 3] 3 | batch_size_per_guide: 10 4 | guide_path: './guides/' 5 | dataset: 6 | path: './datasets/' 7 | dataset_type: 'hybrid' # 'global' # or 'hybrid' 8 | scene_types: ['tabletop', 'cubby', 'merged_cubby', 'dresser'] # 'tabletop' 9 | num_scenes_per_type: -1 # -1 implies all scenes per type 10 | random_scenes: False # If true, samples 'num_scenes_per_type' number of scenes randomly 11 | save_scene_indices: True # Saves the scene indices in their corresponding order 12 | model: 13 | model_dir: './models/' # Directory where the model isi 14 | device: 'cuda:0' # or 'cuda:x' or 'cpu' 15 | T: 255 16 | traj_len: 50 17 | num_channels: 7 18 | general: 19 | gui: True # Set as false when running on servers 20 | save_dir: './results/' # Directory to save the scenes 21 | wandb: 22 | enable_wandb: True 23 | run_num: 1 24 | project_name: 'Video Submission' 25 | -------------------------------------------------------------------------------- /benchmark/cfgs/cfg2.yaml: -------------------------------------------------------------------------------- 1 | guide: 2 | guides: [4, 5, 10] 3 | batch_size_per_guide: 10 4 | guide_path: './guides/' 5 | dataset: 6 | path: './datasets/' 7 | dataset_type: 'hybrid' # 'global' # or 'hybrid' 8 | scene_types: ['tabletop', 'cubby', 'merged_cubby', 'dresser'] # 'tabletop' 9 | num_scenes_per_type: -1 # -1 implies all scenes per type 10 | random_scenes: False # If true, samples 'num_scenes_per_type' number of scenes randomly 11 | save_scene_indices: True # Saves the scene indices in their corresponding order 12 | model: 13 | model_dir: './models/' # Directory where the model is 14 | device: 'cuda:0' 15 | T: 255 16 | traj_len: 50 17 | num_channels: 7 18 | general: 19 | gui: False # Set as false when running on servers 20 | save_dir: './results2/' # Directory to save the scenes 21 | wandb: 22 | enable_wandb: True 23 | run_num: 1 24 | project_name: 'Test Set Diffusion Benchmark Hybrid 2' 25 | 26 | -------------------------------------------------------------------------------- /benchmark/cfgs/cfg3.yaml: -------------------------------------------------------------------------------- 1 | guide: 2 | guides: [11, 13, 14] 3 | batch_size_per_guide: 10 4 | guide_path: './guides/' 5 | dataset: 6 | path: './datasets/' 7 | dataset_type: 'hybrid' # 'global' # or 'hybrid' 8 | scene_types: ['tabletop', 'cubby', 'merged_cubby', 'dresser'] # 'tabletop' 9 | num_scenes_per_type: -1 # -1 implies all scenes per type 10 | random_scenes: False # If true, samples 'num_scenes_per_type' number of scenes randomly 11 | save_scene_indices: True # Saves the scene indices in their corresponding order 12 | model: 13 | model_dir: './models/' # Directory where the model isi 14 | device: 'cuda:1' 15 | T: 255 16 | traj_len: 50 17 | num_channels: 7 18 | general: 19 | gui: False # Set as false when running on servers 20 | save_dir: './results3/' # Directory to save the scenes 21 | wandb: 22 | enable_wandb: True 23 | run_num: 1 24 | project_name: 'Test Set Diffusion Benchmark Hybrid 2' 25 | 26 | -------------------------------------------------------------------------------- /benchmark/cfgs/cfg4.yaml: -------------------------------------------------------------------------------- 1 | guide: 2 | guides: [16, 18, 21] 3 | batch_size_per_guide: 10 4 | guide_path: './guides/' 5 | dataset: 6 | path: './datasets/' 7 | dataset_type: 'hybrid' 8 | scene_types: ['tabletop', 'cubby', 'merged_cubby', 'dresser'] # 'tabletop' 9 | num_scenes_per_type: -1 # -1 implies all scenes per type 10 | random_scenes: False # If true, samples 'num_scenes_per_type' number of scenes randomly 11 | save_scene_indices: True # Saves the scene indices in their corresponding order 12 | model: 13 | model_dir: './models/' # Directory where the model isi 14 | device: 'cuda:1' 15 | T: 255 16 | traj_len: 50 17 | num_channels: 7 18 | general: 19 | gui: False # Set as false when running on servers 20 | save_dir: './results3/' # Directory to save the scenes 21 | wandb: 22 | enable_wandb: True 23 | run_num: 1 24 | project_name: 'Test Set Diffusion Benchmark Hybrid 2' 25 | 26 | -------------------------------------------------------------------------------- /datasets/load_test_dataset.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import torch 3 | import sys 4 | import numpy as np 5 | import os 6 | 7 | current_dir = os.getcwd() 8 | 9 | sys.path.append(f'{current_dir}/motion-policy-networks/') # path to mpinets 10 | 11 | from mpinets.types import PlanningProblem, ProblemSet 12 | from geometrout.primitive import Cuboid, Cylinder 13 | from robofin.robots import FrankaRobot 14 | 15 | class TestDataset: 16 | def __init__(self, type='global', d_path='./datasets/') -> None: 17 | ''' 18 | 19 | Parameters: 20 | type: 'global' or 'hybrid' or 'both' 21 | d_path: path to the datasets directory 22 | ''' 23 | self.type = type 24 | 25 | if d_path[-1]!='/': 26 | d_path += '/' 27 | 28 | if type=='global': 29 | # self.data = np.load(d_path + 'global_solvable_problems.pkl', allow_pickle=True) 30 | with open(d_path + f'global_solvable_problems.pkl', 'rb') as f: 31 | self.data = pickle.load(f) 32 | elif type=='hybrid': 33 | with open(d_path + f'hybrid_solvable_problems.pkl', 'rb') as f: 34 | self.data = pickle.load(f) 35 | elif type=='both': 36 | with open(d_path + f'both_solvable_problems.pkl', 'rb') as f: 37 | self.data = pickle.load(f) 38 | else: 39 | raise ModuleNotFoundError("No such dataset exists") 40 | 41 | print('-'*50) 42 | print(f"Loaded the dataset with {self.data.keys()} scene types, each consisting of {self.data['tabletop'].keys()} problem types") 43 | print('-'*50) 44 | print("Dataset statistics:"+'*'*20) 45 | total_data = 0 46 | for key1 in self.data.keys(): 47 | for key2 in self.data[key1]: 48 | print(f"{key1}: {key2}: {len(self.data[key1][key2])} problems") 49 | total_data += len(self.data[key1][key2]) 50 | print('-'*50) 51 | print(f"Total data: {total_data}") 52 | print('-'*50) 53 | 54 | self.tabletop_data = np.hstack((self.data['tabletop']['task_oriented'], self.data['tabletop']['neutral_start'], self.data['tabletop']['neutral_goal'])) 55 | self.cubby_data = np.hstack((self.data['cubby']['task_oriented'], self.data['cubby']['neutral_start'], self.data['cubby']['neutral_goal'])) 56 | self.merged_cubby_data = np.hstack((self.data['merged_cubby']['task_oriented'], self.data['merged_cubby']['neutral_start'], self.data['merged_cubby']['neutral_goal'])) 57 | self.dresser_data = np.hstack((self.data['dresser']['task_oriented'], self.data['dresser']['neutral_start'], self.data['dresser']['neutral_goal'])) 58 | 59 | self.data_nums = {} 60 | self.data_nums['tabletop'] = len(self.tabletop_data) 61 | self.data_nums['cubby'] = len(self.cubby_data) 62 | self.data_nums['merged_cubby'] = len(self.cubby_data) 63 | self.data_nums['dresser'] = len(self.dresser_data) 64 | 65 | 66 | def fetch_batch(self, problem_type='random', task_type='random'): 67 | ''' 68 | Parameters: 69 | scene_type: 'all' or 'tabletop' or 'cubby' or 'merged_cubby' or 'dresser' 70 | problem_type: 'random' or 'task_oriented' or 'neutral_start' or 'neutral_goal' 71 | ''' 72 | raise NotImplementedError("Not implemented yet!") 73 | 74 | 75 | 76 | def fetch_data(self, scene_num, scene_type='tabletop'): 77 | ''' 78 | Parameters: 79 | scene_type: tabletop or cubby or merged_cubby or dresser 80 | ''' 81 | if scene_type=='tabletop': 82 | data = self.tabletop_data[scene_num] 83 | elif scene_type=='cubby': 84 | data = self.cubby_data[scene_num] 85 | elif scene_type=='merged_cubby': 86 | data = self.merged_cubby_data[scene_num] 87 | elif scene_type=='dresser': 88 | data = self.dresser_data[scene_num] 89 | else: 90 | raise ModuleNotFoundError("What are you looking for? This dataset only has 4 options. Try to choose one of them and retry!") 91 | 92 | # Initialize lists to store centers and quaternions 93 | cuboid_centers = [] 94 | cuboid_dims = [] 95 | cuboid_quaternions = [] 96 | 97 | cylinder_centers = [] 98 | cylinder_heights = [] 99 | cylinder_radii = [] 100 | cylinder_quaternions = [] 101 | 102 | num_cuboids = 0 103 | num_cylinders = 0 104 | 105 | # Extract centers and quaternions based on object type 106 | for obstacle in data.obstacles: 107 | if isinstance(obstacle, Cuboid): 108 | cuboid_centers.append(np.array(obstacle.center)) 109 | cuboid_quaternions.append(np.array(list(obstacle._pose._so3._quat))) 110 | cuboid_dims.append(np.array(obstacle.dims)) 111 | num_cuboids += 1 112 | elif isinstance(obstacle, Cylinder): 113 | cylinder_centers.append(np.array(obstacle.center)) 114 | cylinder_heights.append(np.array(obstacle.height)) 115 | cylinder_radii.append(np.array(obstacle.radius)) 116 | cylinder_quaternions.append(np.array(list(obstacle._pose._so3._quat))) 117 | num_cylinders += 1 118 | 119 | # Convert the lists to NumPy arrays 120 | cuboid_config = [] 121 | cylinder_config = [] 122 | if(num_cuboids >= 1): 123 | # Convert the lists to NumPy arrays 124 | cuboid_centers = np.array(cuboid_centers) 125 | cuboid_dims = np.array(cuboid_dims) 126 | cuboid_quaternions = np.roll(np.array(cuboid_quaternions), -1, axis = 1) 127 | cuboid_config = np.concatenate([cuboid_centers, cuboid_quaternions, cuboid_dims], axis = 1) 128 | 129 | if(num_cylinders >= 1): 130 | cylinder_centers = np.array(cylinder_centers) 131 | cylinder_heights = np.array(cylinder_heights).reshape(-1, 1) 132 | cylinder_radii = np.array(cylinder_radii).reshape(-1, 1) 133 | cylinder_quaternions = np.roll(np.array(cylinder_quaternions), -1, axis = 1) 134 | cylinder_config = np.concatenate([cylinder_centers, cylinder_quaternions, cylinder_radii, cylinder_heights], axis = 1) 135 | 136 | cylinder_cuboid_dims = np.zeros((num_cylinders, 3)) 137 | cylinder_cuboid_dims[:, 0] = cylinder_radii[:, 0] 138 | cylinder_cuboid_dims[:, 1] = cylinder_radii[:, 0] 139 | cylinder_cuboid_dims[:, 2] = cylinder_heights[:, 0] 140 | 141 | if(num_cylinders >= 1): 142 | obstacle_centers = np.concatenate([cuboid_centers, cylinder_centers], axis = 0) 143 | obstacle_dims = np.concatenate([cuboid_dims, cylinder_cuboid_dims], axis = 0) 144 | obstacle_quaternions = np.concatenate([cuboid_quaternions, cylinder_quaternions], axis = 0) 145 | else: 146 | obstacle_centers = cuboid_centers 147 | obstacle_dims = cuboid_dims 148 | obstacle_quaternions = cuboid_quaternions 149 | 150 | obstacle_config = np.concatenate([obstacle_centers, obstacle_quaternions, obstacle_dims], axis = 1) 151 | # cuboid_centers = np.array(cuboid_centers) 152 | # cuboid_dims = np.array(cuboid_dims) 153 | # cuboid_quaternions = np.array(cuboid_quaternions) 154 | # cuboid_config = np.concatenate([cuboid_centers, cuboid_quaternions, cuboid_dims], axis = 1) 155 | 156 | # cylinder_centers = np.array(cylinder_centers) 157 | # cylinder_heights = np.array(cylinder_heights).reshape(-1, 1) 158 | # cylinder_radii = np.array(cylinder_radii).reshape(-1, 1) 159 | # cylinder_quaternions = np.array(cylinder_quaternions) 160 | # cylinder_config = np.concatenate([cylinder_centers, cylinder_quaternions, cylinder_radii, cylinder_heights], axis = 1) 161 | 162 | # cylinder_cuboid_dims = np.zeros((num_cylinders, 3)) 163 | # cylinder_cuboid_dims[:, 0] = cylinder_radii[:, 0] 164 | # cylinder_cuboid_dims[:, 1] = cylinder_radii[:, 0] 165 | # cylinder_cuboid_dims[:, 2] = cylinder_heights[:, 0] 166 | 167 | # obstacle_centers = np.concatenate([cuboid_centers, cylinder_centers], axis = 0) 168 | # obstacle_dims = np.concatenate([cuboid_dims, cylinder_cuboid_dims], axis = 0) 169 | # obstacle_quaternions = np.concatenate([cuboid_quaternions, cylinder_quaternions], axis = 0) 170 | 171 | # obstacle_config = np.concatenate([obstacle_centers, obstacle_quaternions, obstacle_dims], axis = 1) 172 | 173 | start = data.q0 174 | goal_ee = data.target 175 | 176 | joint_7_samples = np.concatenate((np.random.uniform(-2.8973, 2.8973, 50), np.linspace(-2.8973, 2.8973, 50))) 177 | all_ik_goals = [] 178 | for ik_sol_num, joint_ang in enumerate(joint_7_samples): 179 | ik_solutions = np.array(FrankaRobot.ik(goal_ee, joint_ang, 'right_gripper')).reshape((-1, 7)) 180 | if len(ik_solutions)==0: 181 | continue 182 | 183 | if len(all_ik_goals)==0: 184 | all_ik_goals = ik_solutions.reshape((-1, 7)) 185 | else: 186 | all_ik_goals = np.vstack((all_ik_goals, ik_solutions[0].reshape((-1, 7)))) # n*7 187 | 188 | # print(f"Total IK Solutions: {len(all_ik_goals)}") 189 | return obstacle_config, cuboid_config, cylinder_config, num_cuboids, num_cylinders, start, all_ik_goals 190 | 191 | def random(self): 192 | # Initialize lists to store centers and quaternions 193 | cuboid_centers = [] 194 | cuboid_dims = [] 195 | cuboid_quaternions = [] 196 | 197 | cylinder_centers = [] 198 | cylinder_heights = [] 199 | cylinder_radii = [] 200 | cylinder_quaternions = [] 201 | 202 | num_cuboids = 0 203 | num_cylinders = 0 204 | 205 | cuboid_config, cylinder_config = None, None 206 | 207 | # Extract centers and quaternions based on object type 208 | for obstacle in obstacles: 209 | if isinstance(obstacle, Cuboid): 210 | cuboid_centers.append(np.array(obstacle.center)) 211 | cuboid_quaternions.append(np.array(list(obstacle._pose._so3._quat))) 212 | cuboid_dims.append(np.array(obstacle.dims)) 213 | num_cuboids += 1 214 | elif isinstance(obstacle, Cylinder): 215 | cylinder_centers.append(np.array(obstacle.center)) 216 | cylinder_heights.append(np.array(obstacle.height)) 217 | cylinder_radii.append(np.array(obstacle.radius)) 218 | cylinder_quaternions.append(np.array(list(obstacle._pose._so3._quat))) 219 | num_cylinders += 1 220 | 221 | if(num_cuboids >= 1): 222 | # Convert the lists to NumPy arrays 223 | cuboid_centers = np.array(cuboid_centers) 224 | cuboid_dims = np.array(cuboid_dims) 225 | cuboid_quaternions = np.roll(np.array(cuboid_quaternions), -1, axis = 1) 226 | cuboid_config = np.concatenate([cuboid_centers, cuboid_quaternions, cuboid_dims], axis = 1) 227 | 228 | if(num_cylinders >= 1): 229 | cylinder_centers = np.array(cylinder_centers) 230 | cylinder_heights = np.array(cylinder_heights).reshape(-1, 1) 231 | cylinder_radii = np.array(cylinder_radii).reshape(-1, 1) 232 | cylinder_quaternions = np.roll(np.array(cylinder_quaternions), -1, axis = 1) 233 | # print('cyl centres shape: {}'.format(cylinder_centers.shape)) 234 | cylinder_config = np.concatenate([cylinder_centers, cylinder_quaternions, cylinder_radii, cylinder_heights], axis = 1) 235 | 236 | cylinder_cuboid_dims = np.zeros((num_cylinders, 3)) 237 | cylinder_cuboid_dims[:, 0] = cylinder_radii[:, 0] 238 | cylinder_cuboid_dims[:, 1] = cylinder_radii[:, 0] 239 | cylinder_cuboid_dims[:, 2] = cylinder_heights[:, 0] 240 | 241 | if(num_cylinders >= 1): 242 | obstacle_centers = np.concatenate([cuboid_centers, cylinder_centers], axis = 0) 243 | obstacle_dims = np.concatenate([cuboid_dims, cylinder_cuboid_dims], axis = 0) 244 | obstacle_quaternions = np.concatenate([cuboid_quaternions, cylinder_quaternions], axis = 0) 245 | else: 246 | obstacle_centers = cuboid_centers 247 | obstacle_dims = cuboid_dims 248 | obstacle_quaternions = cuboid_quaternions 249 | 250 | obstacle_config = np.concatenate([obstacle_centers, obstacle_quaternions, obstacle_dims], axis = 1) 251 | 252 | return obstacle_config, cuboid_config, cylinder_config, num_cuboids, num_cylinders 253 | 254 | 255 | if __name__=='__main__': 256 | tData = TestDataset() 257 | 258 | tData.fetch_data(scene_num=5) 259 | -------------------------------------------------------------------------------- /diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | from .models import * 2 | from .diffusion import * 3 | from .gaussian import * -------------------------------------------------------------------------------- /diffusion/diffusion.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import copy 4 | import time 5 | 6 | from .gaussian import Gaussian 7 | 8 | class Diffusion(Gaussian): 9 | 10 | def __init__(self, T, device, variance_thresh = 0.02): 11 | 12 | self.T = T 13 | self.beta = self.schedule_variance(variance_thresh) 14 | 15 | self.alpha = 1 - self.beta 16 | self.alpha_bar = np.array(list(map(lambda t:np.prod(self.alpha[:t]), np.arange(self.T+1)[1:]))) 17 | 18 | self.device = device 19 | 20 | Gaussian.__init__(self) 21 | 22 | def cosine_func(self, t): 23 | """ 24 | Implements the cosine function at a timestep t 25 | Inputs: 26 | t -> Integer ; current timestep 27 | T -> Integer ; total number of timesteps 28 | Outputs: 29 | out -> Float ; cosine function output 30 | """ 31 | 32 | s = 1e-10 33 | out = np.cos((t/self.T + s) * (np.pi/2) / (1 + s)) ** 0.15 34 | 35 | return out 36 | 37 | def schedule_variance(self, thresh = 0.02): 38 | """ 39 | Schedules the variance for the diffuser 40 | Inputs: 41 | T -> Integer ; total number of timesteps 42 | thresh -> Float ; variance threshold at the last step 43 | Outputs: 44 | schedule -> Numpy array of shape (2,) ; the variance schedule 45 | """ 46 | 47 | schedule = np.linspace(0, thresh, self.T + 1)[1:] 48 | 49 | return schedule 50 | 51 | def q_sample(self, x, t, eps = None): 52 | """ 53 | Generates q(xt+1/xt) 54 | Inputs: 55 | x0 -> Numpy array of shape (num_samples, num_channels, trajectory_length) 56 | t -> Numpy array of length (num_samples, ); timesteps (these are the output timesteps) 57 | Outputs: 58 | xt -> Numpy array of shape (num_samples, num_channels, trajectory_length) 59 | mean -> Numpy array of shape (num_samples, num_channels, trajectory_length); mean from which sample is taken 60 | var-> Numpy array of length (num_samples); variance from which sample is taken 61 | """ 62 | 63 | # Gather trajectory params: 64 | b = x.shape[0] 65 | c = x.shape[1] 66 | n = x.shape[2] 67 | 68 | if type(eps) == type(None): 69 | eps = np.random.multivariate_normal(mean = np.zeros((n*c)), 70 | cov = np.eye(n*c), 71 | size = (b,)).reshape(b, c, n) 72 | 73 | xt = np.sqrt(self.alpha[t-1, np.newaxis, np.newaxis]) * x + np.sqrt(1 - self.alpha[t-1, np.newaxis, np.newaxis]) * eps 74 | mean = (np.sqrt(self.alpha[t-1, np.newaxis, np.newaxis]) * x) 75 | var = np.sqrt(1 - self.alpha[t-1]) 76 | 77 | return xt, mean, var 78 | 79 | def q_sample_from_x0(self, x0, t, eps = None): 80 | """ 81 | Generates q(xt/x0) 82 | Inputs: 83 | x0 -> Numpy array of shape (num_samples, num_channels, trajectory_length) 84 | t -> Numpy array of length (num_samples, ); timesteps (these are the output timesteps) 85 | Outputs: 86 | xt -> Numpy array of shape (num_samples, num_channels, trajectory_length) 87 | mean -> Numpy array of shape (num_samples, num_channels, trajectory_length); mean from which sample is taken 88 | var-> Numpy array of length (num_samples); variance from which sample is taken 89 | """ 90 | 91 | # Gather trajectory params: 92 | b = x0.shape[0] 93 | c = x0.shape[1] 94 | n = x0.shape[2] 95 | 96 | if type(eps) == type(None): 97 | eps = np.random.multivariate_normal(mean = np.zeros((n*c)), 98 | cov = np.eye(n*c), 99 | size = (b,)).reshape(b, c, n) 100 | 101 | xt = np.sqrt(self.alpha_bar[t-1, np.newaxis, np.newaxis]) * x0 + np.sqrt(1 - self.alpha_bar[t-1, np.newaxis, np.newaxis]) * eps 102 | mean = (np.sqrt(self.alpha_bar[t-1, np.newaxis, np.newaxis]) * x0) 103 | var = np.sqrt(1 - self.alpha_bar[t-1, np.newaxis, np.newaxis]) 104 | 105 | return xt, mean, var 106 | 107 | def p_sample(self, xt, t, eps): 108 | """ 109 | Generates reverse probability p(xt-1/xt, eps) with no dependence on x0. 110 | """ 111 | 112 | xt_prev = (xt - np.sqrt(1 - self.alpha[t-1]) * eps) / np.sqrt(self.alpha[t-1]) 113 | 114 | return xt_prev 115 | 116 | def p_sample_using_posterior(self, xt, t, eps): 117 | """ 118 | Generates reverse probability p(xt-1/xt, x0) using posterior mean and posterior variance. 119 | """ 120 | 121 | # Gather trajectory params: 122 | b = xt.shape[0] 123 | c = xt.shape[1] 124 | n = xt.shape[2] 125 | 126 | z = np.random.multivariate_normal(mean = np.zeros(n), cov = np.eye(n), size = (b, c)) 127 | z[np.where(t == 1), :, :] = 0 128 | 129 | alpha = self.alpha[t-1, np.newaxis, np.newaxis] 130 | alpha_bar = self.alpha_bar[t-1, np.newaxis, np.newaxis] 131 | beta = self.beta[t-1, np.newaxis, np.newaxis] 132 | 133 | xt_prev = (xt - ((1 - alpha) / np.sqrt(1 - alpha_bar)) * eps) / np.sqrt(alpha) + beta*z 134 | 135 | return xt_prev.copy() 136 | 137 | def forward_diffuse(self, x0, condition = True): 138 | """ 139 | Forward diffuses the trajectory till the last timestep T 140 | Inputs: 141 | x -> Numpy array of shape (num_samples, num_channels, trajectory_length) 142 | Outputs: 143 | diffusion_trajs -> Numpy array of shape (Timesteps, num_samples, num_channels, trajectory_length); diffused trajectories 144 | eps -> Numpy array of shape (Timesteps, num_samples, num_channels, trajectory_length); Noise added at each timestep 145 | kl_divs -> Numpy array of shape (Timesteps, num_samples); KL Divergence at each timestep 146 | """ 147 | 148 | # Gather trajectory params: 149 | b = x0.shape[0] 150 | c = x0.shape[1] 151 | n = x0.shape[2] 152 | 153 | # Initialize the diffusion trajectories across time steps: 154 | diffusion_trajs = np.zeros((self.T + 1, b, c, n)) 155 | diffusion_trajs[0] = x0.copy() 156 | 157 | # Calculate epsilons to forward diffuse current trajectory: 158 | mean = np.zeros(n*c) 159 | cov = np.eye(n*c) 160 | eps = np.random.multivariate_normal(mean, cov, size = (self.T, b)) 161 | kl_divs = np.zeros((self.T, b)) 162 | 163 | # ??? This loop can probably be vectorized 164 | 165 | for t in range(1, self.T + 1): 166 | 167 | diffusion_trajs[t] = self.q_sample(diffusion_trajs[t-1], t, eps[t-1]) 168 | if condition: 169 | diffusion_trajs[t, :, :, 0] = x0[:, :, 0].copy() 170 | diffusion_trajs[t, :, :, -1] = x0[:, :, -1].copy() 171 | 172 | # ??? 173 | kl_divs[t-1, :] = self.KL_divergence_against_gaussian(diffusion_trajs[t].flatten()) 174 | 175 | return diffusion_trajs.copy(), eps.copy(), kl_divs.copy() 176 | 177 | def reverse_diffuse(self, xT, eps): 178 | """ 179 | Reverse diffuses the trajectory from last timestep T to first timestep 0 180 | Inputs: 181 | xT -> Numpy array of shape (num_samples, num_channels, trajectory_length) 182 | eps -> Numpy array of shape (Timesteps, num_samples, num_channels, trajectory_length); Noise added at each timestep 183 | Outputs: 184 | diffusion_trajs -> Numpy array of shape (Timesteps, num_samples, num_channels, trajectory_length); reverse diffused trajectories 185 | """ 186 | 187 | # Gather trajectory params: 188 | b = xT.shape[0] 189 | c = xT.shape[1] 190 | n = xT.shape[2] 191 | 192 | diffusion_trajs = np.zeros((self.T + 1, b, c, n)) 193 | diffusion_trajs[self.T] = xT.copy() 194 | 195 | for t in range(self.T, 0, -1): 196 | 197 | diffusion_trajs[t-1] = self.p_sample(diffusion_trajs[t], t, eps[t-1]) 198 | 199 | return diffusion_trajs.copy() 200 | 201 | def generate_q_sample(self, x0, time_steps = None, condition = True, return_type = "tensor"): 202 | """ 203 | Generates q samples for a random set of timesteps, useful for training\n\n 204 | Inputs:\n 205 | x0 -> Numpy array of shape (num_samples, num_channels, trajectory_length)\n 206 | condition -> Bool ; whether to apply conditioning or not\n 207 | return_type -> String (either "tensor" or "numpy") ; whether to return a tensor or numpy array\n\n 208 | Outputs:\n 209 | X -> Numpy array of shape (num_samples, num_channels, trajectory_length); xt from each x0 diffused to a random timestep t\n 210 | Y -> Numpy array of shape (num_samples, num_channels, trajectory_length); the noise added to each x0\n 211 | time_step -> Numpy array of shape (num_samples, ); the timestep to which each x0 is diffused\n 212 | means -> Numpy array of shape (num_samples, num_channels, trajectory_length); mean of each x0 diffused to a random timestep t\n 213 | vars -> Numpy array of shape (num_samples, ); the variance at the timestep to which x0 is diffused\n 214 | """ 215 | 216 | # Refer to the Training Algorithm in Ho et al 2020 for the psuedo code of this function 217 | 218 | # Gather trajectory params: 219 | b = x0.shape[0] 220 | c = x0.shape[1] 221 | n = x0.shape[2] 222 | 223 | if type(time_steps) == type(None): 224 | time_steps = np.random.randint(1, self.T + 1, size = (b, )) 225 | 226 | # Remember, for each channel, we get a multivariate normal distribution. 227 | eps = np.random.multivariate_normal(mean = np.zeros((n,)), cov = np.eye(n), size = (b, c)) 228 | 229 | # Size chart: 230 | # x0 => (num_samples, 2, traj_len) 231 | # xt => (num_samples, 2, traj_len) 232 | # alpha_bar => (num_samples, 1, 1) 233 | # eps => (num_samples, 2, traj_len) 234 | 235 | xt, means, vars = self.q_sample_from_x0(x0, time_steps, eps) 236 | 237 | # CONDITIONING: 238 | if condition: 239 | xt[:, :, 0] = x0[:, :, 0].copy() 240 | xt[:, :, -1] = x0[:, :, -1].copy() 241 | 242 | # Type Conversion: 243 | if return_type == "tensor": 244 | X = torch.tensor(xt, dtype = torch.float32) 245 | Y = torch.tensor(eps, dtype = torch.float32) 246 | time_steps = torch.tensor(time_steps, dtype = torch.float32) 247 | elif return_type == "numpy": 248 | X = xt.copy() 249 | Y = eps.copy() 250 | 251 | return X, Y, time_steps, means, vars 252 | 253 | def denoise(self, model, traj_len, num_channels, start = None, goal = None, condition = True): 254 | 255 | X_t = np.random.multivariate_normal(mean = np.zeros(traj_len), cov = np.eye(traj_len), size = (1, num_channels)) 256 | 257 | if condition: 258 | X_t[:, :, 0] = start[:] 259 | X_t[:, :, -1] = goal[:] 260 | 261 | model.train(False) 262 | 263 | for t in range(self.T, 0, -1): 264 | 265 | X_input = torch.tensor(X_t, dtype = torch.float32).to(self.device) 266 | time_in = torch.tensor([t], dtype = torch.float32).to(self.device) 267 | 268 | epsilon = model(X_input, time_in).numpy(force=True) 269 | 270 | X_t = self.p_sample_using_posterior(X_t, t, epsilon) 271 | 272 | if condition: 273 | X_t[:, :, 0] = start[:] 274 | X_t[:, :, -1] = goal[:] 275 | 276 | print("\rDenoised " + str(t-1) + " timesteps", end="", flush=True) 277 | 278 | return X_t[0] 279 | 280 | def clip_joints(self, joints): 281 | 282 | joint_lower_limits = np.array([-166*(np.pi/180), 283 | -101*(np.pi/180), 284 | -166*(np.pi/180), 285 | -176*(np.pi/180), 286 | -166*(np.pi/180), 287 | -1*(np.pi/180), 288 | -166*(np.pi/180)]) 289 | 290 | joint_upper_limits = np.array([166*(np.pi/180), 291 | 101*(np.pi/180), 292 | 166*(np.pi/180), 293 | -4*(np.pi/180), 294 | 166*(np.pi/180), 295 | 215*(np.pi/180), 296 | 166*(np.pi/180)]) 297 | 298 | return np.clip(joints, joint_lower_limits[np.newaxis, :, np.newaxis], joint_upper_limits[np.newaxis, :, np.newaxis]) 299 | 300 | def denoise_guided(self, model, guide, traj_len, num_channels, guidance_schedule, batch_size = 1, start = None, goal = None, 301 | condition = True, benchmarking = False): 302 | 303 | X_t = np.random.multivariate_normal(mean = np.zeros(traj_len), cov = np.eye(traj_len), size = (batch_size, num_channels)) 304 | 305 | if condition: 306 | X_t[:, :, 0] = start[:] 307 | X_t[:, :, -1] = goal[:] 308 | 309 | model.train(False) 310 | 311 | period = 2 312 | 313 | # full_info = np.zeros((self.T+1, batch_size, 7, 50)) 314 | for t in range(self.T, 0, -1): 315 | print(f"\rDenoising: " + str(t) + " ", end="") 316 | 317 | # full_info[t, :, : ,:] = X_t.copy() # copy.deepcopy(X_t) 318 | 319 | X_input = torch.tensor(X_t, dtype = torch.float32).to(self.device) 320 | time_in = torch.tensor([t], dtype = torch.float32).to(self.device) 321 | 322 | epsilon = model(X_input, time_in).numpy(force=True) 323 | 324 | X_t = self.p_sample_using_posterior(X_t, t, epsilon) 325 | 326 | if (t%period) < (period/2): 327 | if t >= 5: 328 | clipped_joints = self.clip_joints(X_t[:, :, 1:-1]) 329 | # st = time.time() 330 | gradient = guide.get_gradient(clipped_joints, start[:], goal[:], t) 331 | # print(f"\nGradient time: {time.time() - st}") 332 | 333 | # if guidance_schedule == 'constant': 334 | # X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - guidance_scale * gradient 335 | # elif guidance_schedule == 'varying': 336 | # # Remember to index t-1 while parallelizing 337 | # X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - (1.4 + (t/self.T)) * gradient 338 | # else: 339 | # raise NotImplementedError("This type of guidance schedule is not implemented!!") 340 | 341 | X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - guidance_schedule[:, t-1, np.newaxis, np.newaxis] * gradient 342 | 343 | # weight = np.clip(np.log(1+((t-2)/self.T)*(np.exp(1) - 1)), 0.005, 1) * 2.2 344 | # X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - weight * gradient 345 | # X_t[:, :, 1:-1] = X_t[:, :, 1:-1] - 2. * gradient 346 | 347 | if condition: 348 | X_t[:, :, 0] = start[:] 349 | X_t[:, :, -1] = goal[:] 350 | 351 | if not benchmarking: 352 | print("\rDenoised " + str(t-1) + " timesteps", end="", flush=True) 353 | 354 | # full_info[0, :, :, :] = X_t.copy() 355 | 356 | return X_t.copy() #, full_info.copy() 357 | 358 | 359 | 360 | 361 | 362 | -------------------------------------------------------------------------------- /diffusion/gaussian.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | 4 | class Gaussian: 5 | 6 | def pdf(self, x, mu, sigma): 7 | """ 8 | Returns the probability density function given some mu, sigma and a list of points x. 9 | """ 10 | out = (1 / (sigma * np.sqrt(2*np.pi))) * np.exp(-0.5 * (((x - mu) / sigma) ** 2)) 11 | 12 | return out 13 | 14 | def KL_divergence(self, distr1, distr2): 15 | """ 16 | Computes the Kullback-Leibler Divergence between 2 distributions 17 | """ 18 | 19 | return np.mean(distr1 * np.log(distr1 / distr2)) 20 | 21 | def KL_divergence_against_gaussian(self, sample): 22 | """ 23 | Computes the Kullback-Leibler Divergence of a given set of samples against a standard normal gaussian. Useful in analysing how 24 | the distribution shifts along the timesteps of diffusion. 25 | """ 26 | 27 | mu = np.mean(sample) 28 | sigma = np.var(sample) 29 | 30 | x = np.linspace(-10, 10, int(1e+6)) 31 | 32 | distr1 = self.pdf(x, 0, 1) 33 | distr2 = self.pdf(x, mu, sigma) 34 | 35 | return self.KL_divergence(distr1, distr2) 36 | 37 | def get_limits(self, mu, sigma, edge_factor = 0.01): 38 | """ 39 | Gets the upper and lower limit of the bell curve of the distribution, used for fixing limits for plotting the distribution. 40 | """ 41 | 42 | p_min = edge_factor * self.pdf(mu, mu, sigma) 43 | x_1 = -1 * sigma * np.sqrt(2 * np.log(1 / (p_min * sigma * np.sqrt(2*np.pi)))) + mu 44 | x_2 = sigma * np.sqrt(2 * np.log(1 / (p_min * sigma * np.sqrt(2*np.pi)))) + mu 45 | 46 | return x_1, x_2 47 | 48 | def multivariate_pdf(self, mean, var, size = 1024, limits = (-1, 1)): 49 | """ 50 | Inputs: 51 | mean => 1D array of size k, in a 2D trajectory it is k=2 and for 3D k=3 52 | covariance => 2D array of (k x k) 53 | size => Integer; Per axis length of output 54 | limits => Upper and lower limits for each axis 55 | Output: 56 | Array of 'k' dimensions of shape (size x size x .... x size) 57 | """ 58 | 59 | mean = np.array(mean) 60 | k = mean.size 61 | 62 | axis_div = np.repeat([np.linspace(limits[0], limits[1], size)], repeats = k, axis = 0) 63 | x = np.array(np.meshgrid(*axis_div)) 64 | 65 | mu = np.reshape(mean, (k, *np.ones(k, dtype = np.int32))) 66 | 67 | pdf = (1 / (((2 * np.pi) ** (k/2)) * np.sqrt(var**k))) * np.exp(-0.5 * np.sum((x - mu)**2, axis = 0) / var) 68 | 69 | return pdf 70 | 71 | -------------------------------------------------------------------------------- /diffusion/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .blocks import * 2 | from .temporalunet import * -------------------------------------------------------------------------------- /diffusion/models/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/diffusion/models/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /diffusion/models/__pycache__/blocks.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/diffusion/models/__pycache__/blocks.cpython-38.pyc -------------------------------------------------------------------------------- /diffusion/models/__pycache__/temporalunet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/diffusion/models/__pycache__/temporalunet.cpython-38.pyc -------------------------------------------------------------------------------- /diffusion/models/blocks.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from einops.layers.torch import Rearrange 4 | import numpy as np 5 | import einops 6 | 7 | ################################################################################################################# 8 | #---------------------------------------------- LOWER LEVEL BLOCKS ---------------------------------------------# 9 | ################################################################################################################# 10 | 11 | #--------------------------------------------- 1D CONVOLUTION BLOCK --------------------------------------------# 12 | 13 | class Conv1dBlock(nn.Module): 14 | ''' 15 | Conv1d --> GroupNorm --> Mish 16 | ''' 17 | 18 | def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): 19 | 20 | super().__init__() 21 | 22 | self.block = nn.Sequential( 23 | nn.Conv1d(inp_channels, out_channels, kernel_size, padding=kernel_size // 2), 24 | Rearrange('batch channels horizon -> batch channels 1 horizon'), 25 | nn.GroupNorm(n_groups, out_channels), 26 | Rearrange('batch channels 1 horizon -> batch channels horizon'), 27 | nn.Mish(), 28 | ) 29 | # Why is padding half of kernel size? This makes sure that in the first convolution, half the kernel has zero elements and the other 30 | # half has the first few elements of the input. The same applies to the last convolution. 31 | 32 | def forward(self, x): 33 | 34 | return self.block(x) 35 | 36 | #----------------------------------------- SINUSOIDAL POSITION EMBEDDING ---------------------------------------# 37 | 38 | class SinusoidalPosEmb(nn.Module): 39 | 40 | def __init__(self, dim, device): 41 | 42 | super().__init__() 43 | self.dim = dim 44 | self.device = device 45 | 46 | def forward(self, x): 47 | 48 | half_dim = self.dim // 2 49 | emb = np.log(10000) / (half_dim - 1) 50 | emb = torch.exp(torch.arange(half_dim, device=self.device) * -emb) 51 | emb = x[:, None] * emb[None, :] 52 | emb = torch.cat((emb.sin(), emb.cos()), dim=-1) 53 | 54 | return emb 55 | 56 | #------------------------------------------------- TEMPORAL MLP ------------------------------------------------# 57 | 58 | class TimeMLP(nn.Module): 59 | 60 | def __init__(self, time_embed_dim, out_channels): 61 | 62 | super().__init__() 63 | 64 | self.time_mlp = nn.Sequential( 65 | nn.Mish(), 66 | nn.Linear(time_embed_dim, out_channels), 67 | Rearrange('batch t -> batch t 1'), # This is probably extending by a dimension 68 | ) 69 | 70 | def forward(self, t): 71 | 72 | return self.time_mlp(t) 73 | 74 | #-------------------------------------------- INITIAL TIME EMBEDDING -------------------------------------------# 75 | 76 | class TimeEmbedding(nn.Module): 77 | 78 | def __init__(self, dim, device): 79 | 80 | super().__init__() 81 | self.device = device 82 | 83 | self.time_mlp = nn.Sequential( 84 | SinusoidalPosEmb(dim, self.device), 85 | nn.Linear(dim, dim * 4), 86 | nn.Mish(), 87 | nn.Linear(dim * 4, dim), 88 | ) 89 | 90 | def forward(self, t): 91 | 92 | return self.time_mlp(t) 93 | 94 | ################################################################################################################# 95 | #------------------------------------------ INTERMEDIATE LEVEL BLOCKS ------------------------------------------# 96 | ################################################################################################################# 97 | 98 | #---------------------------------------------- LINEAR ATTENTION -----------------------------------------------# 99 | 100 | class LinearAttention(nn.Module): 101 | 102 | def __init__(self, dim, heads=4, dim_head=32): 103 | 104 | super().__init__() 105 | 106 | self.scale = dim_head ** -0.5 107 | self.heads = heads 108 | hidden_dim = dim_head * heads 109 | 110 | self.to_qkv = nn.Conv1d(dim, hidden_dim * 3, 1, bias=False) 111 | self.to_out = nn.Conv1d(hidden_dim, dim, 1) 112 | 113 | def forward(self, x): 114 | 115 | qkv = self.to_qkv(x).chunk(3, dim = 1) 116 | q, k, v = map(lambda t: einops.rearrange(t, 'b (h c) d -> b h c d', h=self.heads), qkv) 117 | 118 | q = q * self.scale 119 | k = k.softmax(dim = -1) 120 | 121 | # Produces the context of each element in k w.r.t all other elements in v. (weighted sum of every pair of rows possible) 122 | context = torch.einsum('b h d n, b h e n -> b h d e', k, v) 123 | 124 | # Weighted sum of every pair of columns possible: 125 | out = torch.einsum('b h d e, b h d n -> b h e n', context, q) 126 | 127 | # Recombine the the 3 chunks that where seperated to get back the same dimension as input 128 | out = einops.rearrange(out, 'b h c d -> b (h c) d') 129 | 130 | # Convolve back from hidden channels to the original number of channels (dim) 131 | out = self.to_out(out) 132 | 133 | return out 134 | 135 | #------------------------------------------ RESIDUAL CONVOLUTION BLOCK ------------------------------------------# 136 | 137 | class ResidualConvolutionBlock(nn.Module): 138 | 139 | def __init__(self, inp_channels, out_channels, time_embed_dim, kernel_size=5): 140 | super().__init__() 141 | 142 | self.blocks = nn.ModuleList([ 143 | Conv1dBlock(inp_channels, out_channels, kernel_size), 144 | Conv1dBlock(out_channels, out_channels, kernel_size), 145 | ]) 146 | 147 | self.time_mlp = TimeMLP(time_embed_dim, out_channels) 148 | 149 | # The convolution that forms the residual connection between input and output 150 | # If the input and the output have the same number of channels, this is an identity matrix. 151 | self.residual_conv = nn.Conv1d(inp_channels, out_channels, 1) \ 152 | if inp_channels != out_channels else nn.Identity() 153 | 154 | def forward(self, x, t): 155 | ''' 156 | x : [ batch_size x inp_channels x horizon ] 157 | t : [ batch_size x time_embed_dim ] 158 | returns: 159 | out : [ batch_size x out_channels x horizon ] 160 | ''' 161 | 162 | out = self.blocks[0](x) + self.time_mlp(t) 163 | out = self.blocks[1](out) 164 | out = out + self.residual_conv(x) 165 | 166 | return out 167 | 168 | #------------------------------------------- RESIDUAL ATTENTION BLOCK -------------------------------------------# 169 | 170 | class ResidualAttentionBlock(nn.Module): 171 | 172 | def __init__(self, dim, eps = 1e-5): 173 | 174 | super().__init__() 175 | 176 | self.attention = LinearAttention(dim) 177 | 178 | # Layer Norm Parameters: 179 | self.eps = eps 180 | self.g = nn.Parameter(torch.ones(1, dim, 1)) 181 | self.b = nn.Parameter(torch.zeros(1, dim, 1)) 182 | 183 | def forward(self, x): 184 | 185 | # Layer Norm 186 | var = torch.var(x, dim=1, unbiased=False, keepdim=True) 187 | mean = torch.mean(x, dim=1, keepdim=True) 188 | out = (x - mean) / (var + self.eps).sqrt() * self.g + self.b 189 | 190 | # Attention Layer 191 | out = self.attention(out) 192 | 193 | # Residual Connection 194 | out = out + x 195 | 196 | return out 197 | 198 | ################################################################################################################# 199 | #---------------------------------------------- HIGH LEVEL BLOCKS ----------------------------------------------# 200 | ################################################################################################################# 201 | 202 | class DownSampler(nn.Module): 203 | 204 | def __init__(self, dim_in, dim_out, time_dim, is_last = False): 205 | 206 | super().__init__() 207 | 208 | self.down = nn.ModuleList([ResidualConvolutionBlock(dim_in, dim_out, time_embed_dim = time_dim), 209 | ResidualConvolutionBlock(dim_out, dim_out, time_embed_dim = time_dim), 210 | nn.Identity(), # Replace this with ResidualAttentionBlock 211 | nn.Conv1d(dim_out, dim_out, kernel_size = 3, stride = 2, padding = 1) if not is_last else nn.Identity()]) 212 | 213 | def forward(self, x, t): 214 | 215 | x = self.down[0](x, t) 216 | x = self.down[1](x, t) 217 | h = self.down[2](x) 218 | out = self.down[3](h) 219 | 220 | return out, h 221 | 222 | class MiddleBlock(nn.Module): 223 | 224 | def __init__(self, mid_dim, time_dim): 225 | 226 | super().__init__() 227 | 228 | self.middle = nn.ModuleList([ResidualConvolutionBlock(mid_dim, mid_dim, time_embed_dim = time_dim), 229 | nn.Identity(), # Replace this with ResidualAttentionBlock 230 | ResidualConvolutionBlock(mid_dim, mid_dim, time_embed_dim = time_dim)]) 231 | 232 | def forward(self, x, t): 233 | 234 | x = self.middle[0](x, t) 235 | x = self.middle[1](x) 236 | out = self.middle[2](x, t) 237 | 238 | return out 239 | 240 | class UpSampler(nn.Module): 241 | 242 | def __init__(self, dim_in, dim_out, time_dim, is_last = False): 243 | 244 | super().__init__() 245 | 246 | self.up = nn.ModuleList([ResidualConvolutionBlock(dim_out * 2, dim_in, time_embed_dim = time_dim), 247 | ResidualConvolutionBlock(dim_in, dim_in, time_embed_dim = time_dim), 248 | nn.Identity(), # Replace this with ResidualAttentionBlock 249 | nn.ConvTranspose1d(dim_in, dim_in, kernel_size = 4, stride = 2, padding = 1) if not is_last else nn.Identity()]) 250 | 251 | def forward(self, x, h, t): 252 | 253 | x = torch.cat([x, h], dim = 1) 254 | 255 | x = self.up[0](x, t) 256 | x = self.up[1](x, t) 257 | x = self.up[2](x) 258 | out = self.up[3](x) 259 | 260 | return out -------------------------------------------------------------------------------- /diffusion/models/temporalunet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | import os 5 | import torchvision.transforms.functional as tvtf 6 | 7 | from diffusion.models.blocks import * 8 | 9 | class TemporalUNet(nn.Module): 10 | 11 | def __init__(self, model_name, input_dim, time_dim, device, dims = (32, 64, 128, 256)): 12 | 13 | super(TemporalUNet, self).__init__() 14 | 15 | dims = [input_dim, *dims] # length of dims is 5 16 | 17 | # Initial Time Embedding: 18 | self.time_embedding = TimeEmbedding(time_dim, device) 19 | 20 | # Down Sampling: 21 | self.down_samplers = nn.ModuleList([]) 22 | for i in range(len(dims) - 2): # Loops 0, 1, 2 23 | self.down_samplers.append(DownSampler(dims[i], dims[i+1], time_dim)) 24 | self.down_samplers.append(DownSampler(dims[-2], dims[-1], time_dim, is_last = True)) # 3 -> 4 25 | 26 | # Middle Block: 27 | self.middle_block = MiddleBlock(dims[-1], time_dim) 28 | 29 | # Up Sampling: 30 | self.up_samplers = nn.ModuleList([]) 31 | for i in range(len(dims) - 1, 1, -1): # Loops 4, 3, 2 since the last one is a seperate convolution 32 | self.up_samplers.append(UpSampler(dims[i-1], dims[i], time_dim)) 33 | 34 | # Final Convolution: 35 | self.final_conv = nn.Sequential(Conv1dBlock(dims[1], dims[1], kernel_size = 5), 36 | nn.Conv1d(dims[1], input_dim, kernel_size = 1)) 37 | 38 | self.model_name = model_name 39 | if not os.path.exists(model_name): 40 | os.mkdir(model_name) 41 | self.losses = np.array([]) 42 | else: 43 | self.load() 44 | 45 | _ = self.to(device) 46 | 47 | def forward(self, x, t): 48 | """ 49 | x => Tensor of size (batch_size, traj_len*2) 50 | t => Integer representing the diffusion timestep of x 51 | """ 52 | 53 | # Get the time embedding from t: 54 | time_emb = self.time_embedding(t) 55 | 56 | # Down Sampling Layers: 57 | h_list = [] 58 | for i in range(len(self.down_samplers)): 59 | x, h = self.down_samplers[i](x, time_emb) 60 | h_list.append(h) 61 | 62 | # Middle Layer: 63 | x = self.middle_block(x, time_emb) 64 | 65 | # Up Sampling Layers: 66 | for i in range(len(self.up_samplers)): 67 | h_temp = h_list.pop() 68 | # print(f"Shape of x: {x.shape}\t Shape of h_list: {h_temp.shape}") 69 | x = self.up_samplers[i](x, h_temp, time_emb) # How does pop work and not h_list[i] 70 | if x.shape[2] == 8 or x.shape[2] == 14 or x.shape[2] == 26 or x.shape[2] == 8: # Upsampling doubles the dimensions of the input. So, we are manually cropping the extra size to match the size of it's corresponding h (context/residual from the downsampling layer) 71 | x = tvtf.crop(x, 0, 0, x.shape[1], x.shape[2] - 1) 72 | 73 | # Final Convolution 74 | out = self.final_conv(x) 75 | 76 | return out 77 | 78 | def save(self): 79 | 80 | torch.save(self.state_dict(), self.model_name + "/weights_latest.pt") 81 | np.save(self.model_name + "/losses.npy", self.losses) 82 | 83 | def save_checkpoint(self, checkpoint): 84 | 85 | torch.save(self.state_dict(), self.model_name + "/weights_" + str(checkpoint) + ".pt") 86 | np.save(self.model_name + "/latest_checkpoint.npy", checkpoint) 87 | 88 | def load(self): 89 | 90 | self.losses = np.load(self.model_name + "/losses.npy") 91 | self.load_state_dict(torch.load(self.model_name + "/weights_latest.pt")) 92 | print("Loaded Model at " + str(self.losses.size) + " epochs") 93 | 94 | def load_checkpoint(self, checkpoint): 95 | 96 | _ = input("Press Enter if you are running the model for inference, or Ctrl+C\n(Never load a checkpoint for training! This will overwrite progress)") 97 | 98 | latest_checkpoint = np.load(self.model_name + "/latest_checkpoint.npy") 99 | self.load_state_dict(torch.load(self.model_name + "/weights_" + str(checkpoint) + ".pt")) 100 | self.losses = np.load(self.model_name + "/losses.npy")[:checkpoint] 101 | 102 | 103 | 104 | -------------------------------------------------------------------------------- /guides/cfgs/guide1.yaml: -------------------------------------------------------------------------------- 1 | index: 1 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.1, 0.1] 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.0, 0.0] 10 | isr2: [20, 150] 11 | val2: [0.0, 0.0] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: False # True 16 | 17 | guidance_schedule: 18 | type: 'varying' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide10.yaml: -------------------------------------------------------------------------------- 1 | index: 10 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.06, 0.06] # linearly increasing 7 | obstacle_expansion: 8 | isr1: [80, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 80] 11 | val2: [0.0, 0.0] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'sv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: False 16 | 17 | guidance_schedule: 18 | type: 'varying' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide11.yaml: -------------------------------------------------------------------------------- 1 | index: 11 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.0, 0.0] # linearly increasing 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 150] 11 | val2: [0.0, 0.4] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'sv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide12.yaml: -------------------------------------------------------------------------------- 1 | index: 1 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0., 0.] 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 150] 11 | val2: [0.4, 0.] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide13.yaml: -------------------------------------------------------------------------------- 1 | index: 13 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.0, 0.0] # linearly increasing 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 150] 11 | val2: [0.0, 0.4] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'sv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.01 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide14.yaml: -------------------------------------------------------------------------------- 1 | index: 14 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.02, 0.02] # linearly increasing 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 150] 11 | val2: [0.0, 0.4] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'sv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.1 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide15.yaml: -------------------------------------------------------------------------------- 1 | index: 1 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0., 0.] 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 150] 11 | val2: [0.4, 0.] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide16.yaml: -------------------------------------------------------------------------------- 1 | index: 16 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.1, 0.1] # linearly increasing 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 150] 11 | val2: [0.0, 0.4] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'sv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.1 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide17.yaml: -------------------------------------------------------------------------------- 1 | index: 1 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0., 0.] 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 150] 11 | val2: [0.4, 0.] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide18.yaml: -------------------------------------------------------------------------------- 1 | index: 18 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.05, 0.05] # linearly increasing 7 | obstacle_expansion: 8 | isr1: [40, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [10, 40] 11 | val2: [0.0, 0.4] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'sv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide2.yaml: -------------------------------------------------------------------------------- 1 | index: 2 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.05, 0.05] 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.0, 0.0] 10 | isr2: [20, 150] 11 | val2: [0.0, 0.0] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: False 16 | 17 | guidance_schedule: 18 | type: 'varying' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide21.yaml: -------------------------------------------------------------------------------- 1 | index: 21 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.05, 0.05] # linearly increasing 7 | obstacle_expansion: 8 | isr1: [40, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [10, 40] 11 | val2: [0.0, 0.4] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'sv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.1 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide3.yaml: -------------------------------------------------------------------------------- 1 | index: 3 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.01, 0.01] 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.0, 0.0] 10 | isr2: [20, 150] 11 | val2: [0.0, 0.0] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: False 16 | 17 | guidance_schedule: 18 | type: 'varying' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide4.yaml: -------------------------------------------------------------------------------- 1 | index: 4 2 | 3 | hyperparameters: 4 | obstacle_clearance: 5 | range: [0.15, 0.15] 6 | obstacle_expansion: 7 | isr1: [150, 255] 8 | val1: [0.0, 0.0] 9 | isr2: [20, 150] 10 | val2: [0.0, 0.0] 11 | isr3: [0, 20] 12 | val3: [0., 0.] 13 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 14 | grad_norm: False 15 | 16 | guidance_schedule: 17 | type: 'varying' # 'constant' or 'varying' (1.4 + (t/self.T)) 18 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 19 | volume_trust_region: 0.0008 # 0.001 20 | 21 | 22 | -------------------------------------------------------------------------------- /guides/cfgs/guide5.yaml: -------------------------------------------------------------------------------- 1 | index: 5 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0.01, 0.15] # linearly increasing 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.0, 0.0] 10 | isr2: [20, 150] 11 | val2: [0.0, 0.0] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: False 16 | 17 | guidance_schedule: 18 | type: 'varying' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/cfgs/guide9.yaml: -------------------------------------------------------------------------------- 1 | index: 1 2 | 3 | hyperparameters: 4 | batch_size: 10 5 | obstacle_clearance: 6 | range: [0., 0.] 7 | obstacle_expansion: 8 | isr1: [150, 255] 9 | val1: [0.4, 0.4] 10 | isr2: [20, 150] 11 | val2: [0.4, 0.] 12 | isr3: [0, 20] 13 | val3: [0., 0.] 14 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 15 | grad_norm: True 16 | 17 | guidance_schedule: 18 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 19 | scale_val: 0.05 # Value used to scale the gradients (if type=='constant') 20 | volume_trust_region: 0.0008 # 0.001 21 | 22 | 23 | -------------------------------------------------------------------------------- /guides/guide_cfg.py: -------------------------------------------------------------------------------- 1 | '''Guide Class 2 | 3 | self.cfg is a dictionary. Example of self.cfg: 4 | 5 | index: 1 6 | 7 | hyperparameters: 8 | batch_size: 10 9 | obstacle_clearance: 10 | range: [0., 0.] 11 | obstacle_expansion: 12 | ind_set_range1: [255, 150] 13 | val_range1: [0.4, 0.4] 14 | ind_set_range2: [150, 20] 15 | val_range2: [0.4, 0.] 16 | ind_set_range3: [20, 0] 17 | val_range3: [0., 0.] 18 | guidance_method: 'iv' # 'iv': intersection_volume; 'sv': swept_volume 19 | grad_norm: True 20 | 21 | guidance_schedule: 22 | type: 'constant' # 'constant' or 'varying' (1.4 + (t/self.T)) 23 | scale_val: 0.05 24 | ''' 25 | from autolab_core import YamlConfig 26 | 27 | class Guide: 28 | def __init__(self, path_to_yaml='./guides/cfgs/guide1.yaml') -> None: 29 | self.cfg = YamlConfig(path_to_yaml) # Dictionary 30 | 31 | if __name__=='__main__': 32 | classifier = Guide() 33 | -------------------------------------------------------------------------------- /infer_serial.py: -------------------------------------------------------------------------------- 1 | from lib import * 2 | from diffusion import * 3 | from datasets.load_test_dataset import TestDataset 4 | 5 | import argparse 6 | 7 | from autolab_core import YamlConfig 8 | 9 | import os 10 | import wandb 11 | 12 | import time 13 | 14 | if __name__=='__main__': 15 | parser = argparse.ArgumentParser( 16 | prog='Benchmarking Diffusion', 17 | description='Benchmarking with IK on Test sets', 18 | ) 19 | parser.add_argument('-c', '--cfg_path', type=str, default='./benchmark/cfgs/cfg1.yaml') # access as args.cfg_path 20 | 21 | 22 | args = parser.parse_args() 23 | 24 | # Load config file 25 | benchmark_cfg = YamlConfig(args.cfg_path) 26 | 27 | # Load guide params 28 | guides = benchmark_cfg['guide']['guides'] # list 29 | guide_dpath = benchmark_cfg['guide']['guide_path'] 30 | 31 | # Load model params 32 | device = benchmark_cfg['model']['device'] if torch.cuda.is_available() else 'cpu' 33 | traj_len = benchmark_cfg['model']['traj_len'] 34 | T = benchmark_cfg['model']['T'] 35 | num_channels = benchmark_cfg['model']['num_channels'] 36 | 37 | # Load dataset 38 | dataset = TestDataset(benchmark_cfg['dataset']['dataset_type'], d_path=benchmark_cfg['dataset']['path']) # Dataset: global, hybrid, or both; d_path: path to dataset 39 | print(f"Benchmarking dataset: {benchmark_cfg['dataset']['dataset_type']} \tScene types: {benchmark_cfg['dataset']['scene_types']}\tguides: {benchmark_cfg['guide']['guides']}") 40 | 41 | # Load models 42 | # Load Models: 43 | env = RobotEnvironment(gui = benchmark_cfg['general']['gui']) 44 | diffuser = Diffusion(T=benchmark_cfg['model']['T'], device = device) 45 | model_name = benchmark_cfg['model']['model_dir'] + "TemporalUNetModel" + str(T) + "_N" + str(traj_len) 46 | if not os.path.exists(model_name): 47 | print("Model does not exist for these parameters. Train a model first.") 48 | _ = input("Press anything to exit") 49 | exit() 50 | denoiser = TemporalUNet(model_name = model_name, input_dim = num_channels, time_dim = 32, dims=(32, 64, 128, 256, 512, 512), 51 | device = device) 52 | 53 | 54 | enable_wandb = benchmark_cfg['general']['wandb']['enable_wandb'] 55 | 56 | num_guides = len(guides) 57 | batch_size_per_guide = benchmark_cfg['guide']['batch_size_per_guide'] 58 | total_batch_size = int(num_guides * batch_size_per_guide) 59 | guide_cfgs = {'batch_size_per_guide': batch_size_per_guide, 60 | 'total_batch_size': total_batch_size, 61 | 'clearance': np.zeros((total_batch_size, T)), 62 | 'expansion': np.zeros((total_batch_size, T)), 63 | 'guidance_method': np.zeros((total_batch_size,)), 64 | 'grad_norm': np.zeros((total_batch_size,)), 65 | 'guidance_schedule': np.zeros((total_batch_size, T)), 66 | 'volume_trust_region': np.zeros((total_batch_size,)) 67 | } 68 | 69 | 70 | for i in range(len(guides)): 71 | 72 | print(f"Loading Guide {guides[i]}" + "+-+-"*10) 73 | g_cfg = YamlConfig(benchmark_cfg['guide']['guide_path'] + f'cfgs/guide{guides[i]}.yaml') 74 | 75 | guide_cfgs['clearance'][i * batch_size_per_guide : (i+1) * batch_size_per_guide, :] = np.linspace(g_cfg['hyperparameters']['obstacle_clearance']['range'][0], g_cfg['hyperparameters']['obstacle_clearance']['range'][1], T) 76 | 77 | o_e_cfg = g_cfg['hyperparameters']['obstacle_expansion'] 78 | guide_cfgs['expansion'][i * batch_size_per_guide : (i+1) * batch_size_per_guide, o_e_cfg['isr1'][0] : o_e_cfg['isr1'][1]] = np.linspace(o_e_cfg['val1'][0], o_e_cfg['val1'][1], num=abs(o_e_cfg['isr1'][1] - o_e_cfg['isr1'][0])) 79 | guide_cfgs['expansion'][i * batch_size_per_guide : (i+1) * batch_size_per_guide, o_e_cfg['isr2'][0] : o_e_cfg['isr2'][1]] = np.linspace(o_e_cfg['val2'][0], o_e_cfg['val2'][1], num=abs(o_e_cfg['isr2'][1] - o_e_cfg['isr2'][0])) 80 | guide_cfgs['expansion'][i * batch_size_per_guide : (i+1) * batch_size_per_guide, o_e_cfg['isr3'][0] : o_e_cfg['isr3'][1]] = np.linspace(o_e_cfg['val3'][0], o_e_cfg['val3'][1], num=abs(o_e_cfg['isr3'][1] - o_e_cfg['isr3'][0])) 81 | 82 | # expansion_schedule = np.ones(shape = (255,)) 83 | # expansion_schedule[o_e_cfg['isr1'][0]: o_e_cfg['isr1'][1]] = np.linspace(o_e_cfg['val1'][0], o_e_cfg['val1'][1], num=abs(o_e_cfg['isr1'][1] - o_e_cfg['isr1'][0])) 84 | # expansion_schedule[o_e_cfg['isr2'][0]: o_e_cfg['isr2'][1]] = np.linspace(o_e_cfg['val2'][0], o_e_cfg['val2'][1], num=abs(o_e_cfg['isr2'][1] - o_e_cfg['isr2'][0])) 85 | # expansion_schedule[o_e_cfg['isr3'][0]: o_e_cfg['isr3'][1]] = np.linspace(o_e_cfg['val3'][0], o_e_cfg['val3'][1], num=abs(o_e_cfg['isr3'][1] - o_e_cfg['isr3'][0])) 86 | 87 | guide_cfgs['guidance_method'][i * batch_size_per_guide : (i+1) * batch_size_per_guide] = 1 if g_cfg['hyperparameters']['guidance_method'] == "sv" else 0 88 | guide_cfgs['grad_norm'][i * batch_size_per_guide : (i+1) * batch_size_per_guide] = 1 if g_cfg['hyperparameters']['grad_norm'] else 0 89 | 90 | guide_cfgs['guidance_schedule'][i * batch_size_per_guide : (i+1) * batch_size_per_guide, :] = (1.4 + np.arange(T) / T) if g_cfg['hyperparameters']['guidance_schedule']['type'] == 'varying' else g_cfg['hyperparameters']['guidance_schedule']['scale_val'] 91 | guide_cfgs['volume_trust_region'][i * batch_size_per_guide : (i+1) * batch_size_per_guide] = g_cfg['hyperparameters']['volume_trust_region'] 92 | 93 | 94 | t_success = 0 95 | for scene_type in benchmark_cfg['dataset']['scene_types']: 96 | print(f"Loading Scene Type: {scene_type}" + "+-+-"*10) 97 | i = 0 98 | for scene_num in range(dataset.data_nums[scene_type]): 99 | print(f"Scene num: {i+1}\tSuccess_rate: {t_success}/{i}") 100 | env.clear_obstacles() 101 | env.go_home() 102 | 103 | obstacle_config, cuboid_config, cylinder_config, num_cuboids, num_cylinders, start_joints, all_ik_goals = dataset.fetch_data(scene_num=scene_num, scene_type=scene_type) 104 | 105 | guide_set = [] 106 | final_trajectories = [] 107 | 108 | start_time = time.time() 109 | st_time2 = time.time() 110 | 111 | # print(f"Guide: {guide_cfg['index']}, Previous planning time: {time.time() - st_time2}") #, end = "") 112 | guide = IntersectionVolumeGuide(obstacle_config = obstacle_config, device = device, guide_cfgs = guide_cfgs, batch_size = guide_cfgs['total_batch_size']) 113 | 114 | metrics_calculator = MetricsCalculator(guide) 115 | guide_set.append([guide, metrics_calculator]) 116 | 117 | # Filter final IKs 118 | st3 = time.time() 119 | volumes = guide.cost(torch.tensor(all_ik_goals.reshape((-1, 7, 1)), device=device), 0, batch_size = all_ik_goals.shape[0]).sum(axis=(1, 2)).cpu().numpy() 120 | # Sort joints and volumes in the increasing order of volumes 121 | min_volume = np.min(volumes) 122 | indices = np.argsort(volumes) #[:10] 123 | rearranged_volumes = volumes[indices] 124 | goal_joints = all_ik_goals[indices] 125 | volume_trust_region = 0.0008 # Overriding volume trust region as it's not changing with guides 126 | goal_joints = goal_joints[rearranged_volumes < min_volume + volume_trust_region] 127 | ideal_ind = np.argmin(np.linalg.norm(start_joints - goal_joints, axis=1)) 128 | 129 | goal_joints = goal_joints[ideal_ind] 130 | print(f"IK time: {time.time() - st3}") 131 | 132 | st4 = time.time() 133 | 134 | trajectories = diffuser.denoise_guided(model = denoiser, 135 | guide = guide, 136 | batch_size = total_batch_size, 137 | traj_len = traj_len, 138 | num_channels = num_channels, 139 | condition = True, 140 | benchmarking = True, 141 | start = start_joints, 142 | goal = goal_joints, 143 | guidance_schedule = guide_cfgs['guidance_schedule']) 144 | 145 | print(f"Denoiser time: {time.time() - st4}") 146 | 147 | trajectory = guide.choose_best_trajectory(start_joints, goal_joints, trajectories) 148 | final_trajectories.append(trajectory.copy()) 149 | # trajectory is (7, 50) numpy array 150 | end_time = time.time() 151 | 152 | # main_guide = IntersectionVolumeGuide(obstacle_config = obstacle_config, device = device, clearance = guide_cfgs[0]['clearance'], expansion = guide_cfgs[0]['expansion'], 153 | # guidance_method=guide_cfgs[0]['guidance_method'], grad_norm=guide_cfgs[0]['grad_norm']) 154 | 155 | # trajectory = main_guide.choose_best_trajectory_final(np.array(final_trajectories)) 156 | 157 | print(f"\nPlanning Time: {time.time() - start_time} seconds\n") 158 | 159 | if num_cuboids > 0: 160 | env.spawn_collision_cuboids(cuboid_config) 161 | if num_cylinders > 0: 162 | env.spawn_collision_cylinders(cylinder_config) 163 | 164 | 165 | success = env.benchmark_trajectory(trajectory) 166 | 167 | t_success += success 168 | print(f"Success: {success}") 169 | 170 | i+=1 171 | 172 | 173 | -------------------------------------------------------------------------------- /lib/__init__.py: -------------------------------------------------------------------------------- 1 | from .environment import * 2 | from .guide import * 3 | from .metrics import * -------------------------------------------------------------------------------- /lib/environment.py: -------------------------------------------------------------------------------- 1 | import pybullet as p 2 | import pybullet_utils.bullet_client as bc 3 | import pybullet_data 4 | 5 | # from mpinets.types import PlanningProblem, ProblemSet 6 | # from geometrout.primitive import Cuboid, Cylinder 7 | from robofin.robots import FrankaRobot 8 | 9 | import numpy as np 10 | import time 11 | import re 12 | import os 13 | import h5py 14 | 15 | class RobotEnvironment: 16 | 17 | def __init__(self, gui = True, timestep = 1/480, manipulator = True, benchmarking = False): 18 | 19 | self.client_id = bc.BulletClient(p.GUI if gui else p.DIRECT) # Initialize the bullet client 20 | self.client_id.setAdditionalSearchPath(pybullet_data.getDataPath()) # Add pybullet's data package to path 21 | self.client_id.setTimeStep(timestep) # Set simulation timestep 22 | self.client_id.configureDebugVisualizer(p.COV_ENABLE_SHADOWS, 0) # Disable Shadows 23 | 24 | self.timestep = timestep 25 | self.benchmarking = benchmarking 26 | 27 | p.setAdditionalSearchPath(pybullet_data.getDataPath()) # Add pybullet's data package to path 28 | 29 | self.colors = {'blue': np.array([78, 121, 167]) / 255.0, # blue 30 | 'green': np.array([89, 161, 79]) / 255.0, # green 31 | 'brown': np.array([156, 117, 95]) / 255.0, # brown 32 | 'orange': np.array([242, 142, 43]) / 255.0, # orange 33 | 'yellow': np.array([237, 201, 72]) / 255.0, # yellow 34 | 'gray': np.array([186, 176, 172]) / 255.0, # gray 35 | 'red': np.array([255, 87, 89]) / 255.0, # red 36 | 'purple': np.array([176, 122, 161]) / 255.0, # purple 37 | 'cyan': np.array([118, 183, 178]) / 255.0, # cyan 38 | 'pink': np.array([255, 157, 167]) / 255.0} # pink 39 | 40 | target = self.client_id.getDebugVisualizerCamera()[11] # Get cartesian coordinates of the camera's focus 41 | self.client_id.resetDebugVisualizerCamera( # Reset initial camera position 42 | cameraDistance=1.5, 43 | cameraYaw=90, 44 | cameraPitch=-25, 45 | cameraTargetPosition=target, 46 | ) 47 | 48 | p.resetSimulation() 49 | self.client_id.setGravity(0, 0, -9.8) # Set Gravity 50 | 51 | # self.plane = self.client_id.loadURDF("plane.urdf", basePosition=(0, 0, 0), useFixedBase=True) # Load a floor 52 | 53 | # self.client_id.changeDynamics( # Set physical properties of the floor 54 | # self.plane, 55 | # -1, 56 | # lateralFriction=1.1, 57 | # restitution=0.5, 58 | # linearDamping=0.5, 59 | # angularDamping=0.5, 60 | # ) 61 | 62 | if manipulator: 63 | self.initialize_manipulator() 64 | 65 | self.obs_ids = [] 66 | 67 | def initialize_manipulator(self, urdf_file = "franka_panda/panda.urdf", base_position = (0, 0, 0)): 68 | 69 | self.manipulator = self.client_id.loadURDF(urdf_file, basePosition = base_position, useFixedBase = True) 70 | self.joints = [] 71 | 72 | for i in range(self.client_id.getNumJoints(self.manipulator)): 73 | 74 | info = self.client_id.getJointInfo(self.manipulator, i) 75 | 76 | joint_id = info[0] 77 | joint_name = info[1].decode("utf-8") 78 | joint_type = info[2] 79 | 80 | if joint_name == "panda_grasptarget_hand": 81 | self.end_effector = joint_id 82 | 83 | if joint_type == self.client_id.JOINT_REVOLUTE: 84 | self.joints.append(joint_id) 85 | 86 | self.joint_lower_limits = np.array([-166*(np.pi/180), 87 | -101*(np.pi/180), 88 | -166*(np.pi/180), 89 | -176*(np.pi/180), 90 | -166*(np.pi/180), 91 | -1*(np.pi/180), 92 | -166*(np.pi/180)]) 93 | 94 | self.joint_upper_limits = np.array([166*(np.pi/180), 95 | 101*(np.pi/180), 96 | 166*(np.pi/180), 97 | -4*(np.pi/180), 98 | 166*(np.pi/180), 99 | 215*(np.pi/180), 100 | 166*(np.pi/180)]) 101 | 102 | links_folder_path = f'{pybullet_data.getDataPath()}/franka_panda/meshes/collision/' 103 | try: 104 | link_file_names = os.listdir(links_folder_path) 105 | except OSError as e: 106 | print(f"Error reading files in folder: {e}") 107 | 108 | self.link_meshes = {} 109 | self.link_dimensions = {} 110 | self.link_centers = {} 111 | 112 | self.link_index_to_name = ['link0', 'link1', 'link2', 'link3', 'link4', 'link5', 'link6', 'link7', 113 | 'hand', 'finger', 'finger', 'finger', 'finger'] 114 | 115 | for file_name in link_file_names: 116 | 117 | if file_name[-4:] == ".obj": 118 | vertices = [] 119 | link_name = file_name[:-4] 120 | with open(links_folder_path + file_name, 'r') as f: 121 | for line in f: 122 | line = line.strip() 123 | if line.startswith('v '): 124 | vertex = re.split(r'\s+', line)[1:4] 125 | vertex = np.array([float(coord) for coord in vertex]) 126 | vertices.append(vertex) 127 | self.link_meshes[link_name] = np.array(vertices) 128 | max_point = np.max(self.link_meshes[link_name], axis = 0) 129 | min_point = np.min(self.link_meshes[link_name], axis = 0) 130 | self.link_dimensions[link_name] = max_point - min_point 131 | self.link_centers[link_name] = self.link_dimensions[link_name]/2 + min_point 132 | 133 | def get_mpinet_scene(self, index, data = 'train'): 134 | 135 | d_path = '/home/vishal/Volume_E/Active/Undergrad_research/CoRL2023/datasets/nvidia_MPI/mpinets_hybrid_training_data/val/' 136 | # with h5py.File("mpinet_dataset/" + data + ".hdf5", "r") as f: 137 | with h5py.File(d_path + data + ".hdf5", "r") as f: 138 | 139 | cuboid_centers = f['cuboid_centers'][index] 140 | cuboid_dims = f['cuboid_dims'][index] 141 | cuboid_quaternions = np.roll(f['cuboid_quaternions'][index], -1, axis = 1) 142 | 143 | cylinder_centers = f['cylinder_centers'][index] 144 | cylinder_heights = f['cylinder_heights'][index] 145 | cylinder_quaternions = np.roll(f['cylinder_quaternions'][index], -1, axis = 1) 146 | cylinder_radii = f['cylinder_radii'][index] 147 | 148 | num_cuboids = np.argmax(np.any(cuboid_dims == 0, axis = 1)) 149 | num_cylinders = np.argmax(np.any(cylinder_heights == 0, axis = 1)) 150 | 151 | cuboid_centers = cuboid_centers[:num_cuboids] 152 | cuboid_dims = cuboid_dims[:num_cuboids] 153 | cuboid_quaternions = cuboid_quaternions[:num_cuboids] 154 | 155 | cuboid_config = np.concatenate([cuboid_centers, cuboid_quaternions, cuboid_dims], axis = 1) 156 | 157 | cylinder_centers = cylinder_centers[:num_cylinders] 158 | cylinder_heights = cylinder_heights[:num_cylinders] 159 | cylinder_quaternions = cylinder_quaternions[:num_cylinders] 160 | cylinder_radii = cylinder_radii[:num_cylinders] 161 | 162 | cylinder_config = np.concatenate([cylinder_centers, cylinder_quaternions, cylinder_radii, cylinder_heights], axis = 1) 163 | 164 | cylinder_cuboid_dims = np.zeros((num_cylinders, 3)) 165 | cylinder_cuboid_dims[:, 0] = cylinder_radii[:, 0] 166 | cylinder_cuboid_dims[:, 1] = cylinder_radii[:, 0] 167 | cylinder_cuboid_dims[:, 2] = cylinder_heights[:, 0] 168 | 169 | obstacle_centers = np.concatenate([cuboid_centers, cylinder_centers], axis = 0) 170 | obstacle_dims = np.concatenate([cuboid_dims, cylinder_cuboid_dims], axis = 0) 171 | obstacle_quaternions = np.concatenate([cuboid_quaternions, cylinder_quaternions], axis = 0) 172 | 173 | obstacle_config = np.concatenate([obstacle_centers, obstacle_quaternions, obstacle_dims], axis = 1) 174 | 175 | start = f['hybrid_solutions'][index][0, :] 176 | goal = f['hybrid_solutions'][index][-1, :] 177 | 178 | print(f"Actual Start: {start}") 179 | print(f"Actual Goal: {goal}") 180 | 181 | goal_ee = FrankaRobot.fk(goal, eff_frame="right_gripper") 182 | 183 | # goal_ee = data.target # SE3 Pose 184 | 185 | joint_7_samples = np.concatenate((np.random.uniform(-2.8973, 2.8973, 50), np.linspace(-2.8973, 2.8973, 50))) 186 | all_ik_goals = [] 187 | for ik_sol_num, joint_ang in enumerate(joint_7_samples): 188 | ik_solutions = np.array(FrankaRobot.ik(goal_ee, joint_ang, 'right_gripper')).reshape((-1, 7)) 189 | if len(ik_solutions)==0: 190 | continue 191 | 192 | if len(all_ik_goals)==0: 193 | all_ik_goals = ik_solutions.reshape((-1, 7)) 194 | else: 195 | all_ik_goals = np.vstack((all_ik_goals, ik_solutions[0].reshape((-1, 7)))) # n*7 196 | 197 | print(f"Total IK Solutions: {len(all_ik_goals)}") 198 | 199 | return obstacle_config, cuboid_config, cylinder_config, start, all_ik_goals # goal 200 | 201 | def spawn_cuboids(self, cuboid_config): 202 | 203 | for i in range(cuboid_config.shape[0]): 204 | 205 | vuid = self.client_id.createVisualShape(p.GEOM_BOX, 206 | halfExtents = cuboid_config[i, 7:]/2, 207 | rgbaColor = np.hstack([self.colors['yellow'], np.array([1.0])])) 208 | 209 | obs_id = self.client_id.createMultiBody(baseVisualShapeIndex = vuid, 210 | basePosition = cuboid_config[i, :3], 211 | baseOrientation = cuboid_config[i, 3:7]) 212 | 213 | self.obs_ids.append(obs_id) 214 | 215 | def spawn_cylinders(self, cylinder_config): 216 | 217 | for i in range(cylinder_config.shape[0]): 218 | 219 | vuid = self.client_id.createVisualShape(p.GEOM_CYLINDER, 220 | radius = cylinder_config[i, 7], 221 | length = cylinder_config[i, 8], 222 | rgbaColor = np.hstack([self.colors['yellow'], np.array([1.0])])) 223 | 224 | obs_id = self.client_id.createMultiBody(baseVisualShapeIndex = vuid, 225 | basePosition = cylinder_config[i, :3], 226 | baseOrientation = cylinder_config[i, 3:7]) 227 | 228 | self.obs_ids.append(obs_id) 229 | 230 | def spawn_collision_cuboids(self, cuboid_config): 231 | 232 | for i in range(cuboid_config.shape[0]): 233 | 234 | cuid = self.client_id.createCollisionShape(p.GEOM_BOX, 235 | halfExtents = cuboid_config[i, 7:]/2) 236 | 237 | vuid = self.client_id.createVisualShape(p.GEOM_BOX, 238 | halfExtents = cuboid_config[i, 7:]/2, 239 | rgbaColor = np.hstack([self.colors['yellow'], np.array([1.0])])) 240 | 241 | obs_id = self.client_id.createMultiBody(baseMass = 0., 242 | baseCollisionShapeIndex = cuid, 243 | baseVisualShapeIndex = vuid, 244 | basePosition = cuboid_config[i, :3], 245 | baseOrientation = cuboid_config[i, 3:7]) 246 | 247 | self.obs_ids.append(obs_id) 248 | 249 | def spawn_collision_cylinders(self, cylinder_config): 250 | 251 | for i in range(cylinder_config.shape[0]): 252 | 253 | cuid = self.client_id.createCollisionShape(p.GEOM_CYLINDER, 254 | radius = cylinder_config[i, 7], 255 | height = cylinder_config[i, 8]) 256 | 257 | vuid = self.client_id.createVisualShape(p.GEOM_CYLINDER, 258 | radius = cylinder_config[i, 7], 259 | length = cylinder_config[i, 8], 260 | rgbaColor = np.hstack([self.colors['yellow'], np.array([1.0])])) 261 | 262 | obs_id = self.client_id.createMultiBody(baseMass = 0., 263 | baseCollisionShapeIndex = cuid, 264 | baseVisualShapeIndex = vuid, 265 | basePosition = cylinder_config[i, :3], 266 | baseOrientation = cylinder_config[i, 3:7]) 267 | 268 | self.obs_ids.append(obs_id) 269 | 270 | def clear_obstacles(self): 271 | 272 | for id in self.obs_ids: 273 | self.client_id.removeBody(id) 274 | self.obs_ids = [] 275 | 276 | def clip_joints(self, joints): 277 | 278 | return np.clip(joints, self.joint_lower_limits[np.newaxis, :, np.newaxis], self.joint_upper_limits[np.newaxis, :, np.newaxis]) 279 | 280 | def draw_link_bounding_boxes(self): 281 | 282 | self.link_poses = [] 283 | self.link_bounding_vertices = [] 284 | 285 | self.link_bounding_objs = [] 286 | 287 | for link_index in range(0, 11): # The 12th link (i.e. 11th index) is the grasp target and is not needed 288 | 289 | if link_index not in [8, 9]: 290 | 291 | link_name = self.link_index_to_name[link_index+1] 292 | 293 | l, b, h = self.link_dimensions[link_name][0], self.link_dimensions[link_name][1], self.link_dimensions[link_name][2] 294 | 295 | if link_index == 7: 296 | frame_pos, _ = self.client_id.getLinkState(self.manipulator, 7)[4:6] 297 | _, frame_ori = self.client_id.getLinkState(self.manipulator, 10)[4:6] 298 | elif link_index != -1: 299 | frame_pos, frame_ori = self.client_id.getLinkState(self.manipulator, link_index)[4:6] 300 | else: 301 | frame_pos, frame_ori = self.client_id.getBasePositionAndOrientation(self.manipulator) 302 | world_transform = self.pose_to_transformation(np.array([*frame_pos, *frame_ori])) 303 | 304 | link_dimensions = self.link_dimensions[link_name].copy() 305 | if link_index == 10: 306 | link_dimensions[1] *= 4 307 | 308 | world_link_center = (world_transform @ np.vstack((np.expand_dims(self.link_centers[link_name], 1), 1)))[:-1, 0] 309 | 310 | vertices = np.array([[-l/2, -b/2, -h/2], 311 | [ l/2, -b/2, -h/2], 312 | [ l/2, b/2, -h/2], 313 | [-l/2, b/2, -h/2], 314 | [-l/2, -b/2, h/2], 315 | [ l/2, -b/2, h/2], 316 | [ l/2, b/2, h/2], 317 | [-l/2, b/2, h/2]]) 318 | vertices = vertices + np.array([self.link_centers[link_name]]) 319 | vertices = world_transform @ np.vstack((vertices.T, np.ones(8))) 320 | self.link_bounding_vertices.append(vertices.T[:, :-1]) 321 | 322 | self.link_poses.append(np.array([*world_link_center, *frame_ori])) 323 | 324 | vuid = self.client_id.createVisualShape(p.GEOM_BOX, 325 | halfExtents = link_dimensions/2, 326 | rgbaColor = np.hstack([self.colors['red'], np.array([1.0])])) 327 | 328 | obj_id = self.client_id.createMultiBody(baseVisualShapeIndex = vuid, 329 | basePosition = world_link_center, 330 | baseOrientation = frame_ori) 331 | 332 | self.link_bounding_objs.append(obj_id) 333 | 334 | def clear_bounding_boxes(self): 335 | 336 | for obj_id in self.link_bounding_objs: 337 | self.client_id.removeBody(obj_id) 338 | 339 | def get_joint_positions(self): 340 | 341 | return np.array([self.client_id.getJointState(self.manipulator, i)[0] for i in self.joints]) 342 | 343 | def get_joint_velocities(self): 344 | 345 | return np.array([self.client_id.getJointState(self.manipulator, i)[1] for i in self.joints]) 346 | 347 | def get_tf_mat(self, i, joint_angles): 348 | 349 | dh_params = np.array([[0, 0.333, 0, joint_angles[0]], 350 | [0, 0, -np.pi / 2, joint_angles[1]], 351 | [0, 0.316, np.pi / 2, joint_angles[2]], 352 | [0.0825, 0, np.pi / 2, joint_angles[3]], 353 | [-0.0825, 0.384, -np.pi / 2, joint_angles[4]], 354 | [0, 0, np.pi / 2, joint_angles[5]], 355 | [0.088, 0, np.pi / 2, joint_angles[6]], 356 | [0, 0.107, 0, 0], 357 | [0, 0, 0, -np.pi / 4], 358 | [0.0, 0.1034, 0, 0]], dtype=np.float64) 359 | 360 | a = dh_params[i][0] 361 | d = dh_params[i][1] 362 | alpha = dh_params[i][2] 363 | theta = dh_params[i][3] 364 | q = theta 365 | 366 | return np.array([[np.cos(q), -np.sin(q), 0, a], 367 | [np.sin(q) * np.cos(alpha), np.cos(q) * np.cos(alpha), -np.sin(alpha), -np.sin(alpha) * d], 368 | [np.sin(q) * np.sin(alpha), np.cos(q) * np.sin(alpha), np.cos(alpha), np.cos(alpha) * d], 369 | [0, 0, 0, 1]]) 370 | 371 | def draw_frame(self, transform, scale_factor = 0.2): 372 | 373 | unit_axes_world = np.array([[scale_factor, 0, 0], 374 | [0, scale_factor, 0], 375 | [0, 0, scale_factor], 376 | [1, 1, 1]]) 377 | axis_points = ((transform @ unit_axes_world)[:3, :]).T 378 | axis_center = transform[:3, 3] 379 | 380 | l1 = self.client_id.addUserDebugLine(axis_center, axis_points[0], self.colors['red'], lineWidth = 4) 381 | l2 = self.client_id.addUserDebugLine(axis_center, axis_points[1], self.colors['green'], lineWidth = 4) 382 | l3 = self.client_id.addUserDebugLine(axis_center, axis_points[2], self.colors['blue'], lineWidth = 4) 383 | 384 | frame_id = [l1, l2, l3] 385 | 386 | return frame_id[:] 387 | 388 | def remove_frame(self, frame_id): 389 | 390 | for id in frame_id: 391 | self.client_id.removeUserDebugItem(id) 392 | 393 | def inverse_of_transform(self, matrix): 394 | 395 | # Extract the rotation part and translation part of the matrix 396 | rotation_part = matrix[:3, :3] 397 | translation_part = matrix[:3, 3] 398 | 399 | # Calculate the inverse of the rotation part 400 | inverse_rotation = np.linalg.inv(rotation_part) 401 | 402 | # Calculate the new translation by applying the inverse rotation 403 | inverse_translation = -inverse_rotation.dot(translation_part) 404 | 405 | # Create the inverse transformation matrix 406 | inverse_matrix = np.zeros_like(matrix) 407 | inverse_matrix[:3, :3] = inverse_rotation 408 | inverse_matrix[:3, 3] = inverse_translation 409 | inverse_matrix[3, 3] = 1.0 410 | 411 | return inverse_matrix.copy() 412 | 413 | def forward_kinematics(self, joint_angles): 414 | 415 | T_EE = np.identity(4) 416 | for i in range(7 + 3): 417 | T_EE = T_EE @ self.get_tf_mat(i, joint_angles) 418 | 419 | return T_EE 420 | 421 | def get_jacobian(self, joint_angles): 422 | 423 | T_EE = self.get_ee_tf_mat(joint_angles) 424 | 425 | J = np.zeros((6, 10)) 426 | T = np.identity(4) 427 | for i in range(7 + 3): 428 | T = T @ self.get_tf_mat(i, joint_angles) 429 | 430 | p = T_EE[:3, 3] - T[:3, 3] 431 | z = T[:3, 2] 432 | 433 | J[:3, i] = np.cross(z, p) 434 | J[3:, i] = z 435 | 436 | return J[:, :7] 437 | 438 | def inverse_kinematics(self, point, euler_orientation = None): 439 | 440 | if type(euler_orientation) == type(None): 441 | point = self.client_id.calculateInverseKinematics(self.manipulator, self.end_effector, point) 442 | else: 443 | quat = self.client_id.getQuaternionFromEuler(euler_orientation) 444 | point = self.client_id.calculateInverseKinematics(self.manipulator, self.end_effector, point, quat) 445 | 446 | return point 447 | 448 | def get_end_effector_pose(self): 449 | 450 | pos, ori = self.client_id.getLinkState(self.manipulator, 11, computeForwardKinematics = 1)[:2] 451 | pose = np.array([*pos, *ori]) 452 | 453 | return pose 454 | 455 | def get_end_effector_transformation(self): 456 | 457 | pose = self.get_end_effector_pose() 458 | transform = self.pose_to_transformation(pose) 459 | 460 | return transform 461 | 462 | def pose_to_transformation(self, pose): 463 | 464 | pos = pose[:3] 465 | quat = pose[3:] 466 | 467 | rotation_matrix = self.quaternion_to_rotation_matrix(quat) 468 | 469 | transform = np.zeros((4, 4)) 470 | transform[:3, :3] = rotation_matrix.copy() 471 | transform[:3, 3] = pos.copy() 472 | transform[3, 3] = 1 473 | 474 | return transform 475 | 476 | def euler_to_rotation_matrix(yaw, pitch, roll): 477 | 478 | Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0], 479 | [np.sin(yaw), np.cos(yaw), 0], 480 | [0, 0, 1]]) 481 | 482 | Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)], 483 | [0, 1, 0], 484 | [-np.sin(pitch), 0, np.cos(pitch)]]) 485 | 486 | Rx = np.array([[1, 0, 0], 487 | [0, np.cos(roll), -np.sin(roll)], 488 | [0, np.sin(roll), np.cos(roll)]]) 489 | 490 | R = Rz @ (Ry @ Rx) 491 | 492 | return R 493 | 494 | def quaternion_to_rotation_matrix(self, quat): 495 | """ 496 | Convert a quaternion to a rotation matrix. 497 | 498 | :param q: Quaternion [w, x, y, z] 499 | :return: 3x3 rotation matrix 500 | """ 501 | # w, x, y, z = quat 502 | # rotation_matrix = np.array([[1 - 2*y**2 - 2*z**2, 2*x*y - 2*z*w, 2*x*z + 2*y*w], 503 | # [2*x*y + 2*z*w, 1 - 2*x**2 - 2*z**2, 2*y*z - 2*x*w], 504 | # [2*x*z - 2*y*w, 2*y*z + 2*x*w, 1 - 2*x**2 - 2*y**2]]) 505 | 506 | mat = np.array(self.client_id.getMatrixFromQuaternion(quat)) 507 | rotation_matrix = np.reshape(mat, (3, 3)) 508 | 509 | return rotation_matrix 510 | 511 | def rotation_matrix_to_quaternion(self, R): 512 | 513 | trace = np.trace(R) 514 | 515 | if trace > 0: 516 | S = np.sqrt(trace + 1.0) * 2 517 | qw = 0.25 * S 518 | qx = (R[2, 1] - R[1, 2]) / S 519 | qy = (R[0, 2] - R[2, 0]) / S 520 | qz = (R[1, 0] - R[0, 1]) / S 521 | elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]): 522 | S = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2 523 | qw = (R[2, 1] - R[1, 2]) / S 524 | qx = 0.25 * S 525 | qy = (R[0, 1] + R[1, 0]) / S 526 | qz = (R[0, 2] + R[2, 0]) / S 527 | elif R[1, 1] > R[2, 2]: 528 | S = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2 529 | qw = (R[0, 2] - R[2, 0]) / S 530 | qx = (R[0, 1] + R[1, 0]) / S 531 | qy = 0.25 * S 532 | qz = (R[1, 2] + R[2, 1]) / S 533 | else: 534 | S = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2 535 | qw = (R[1, 0] - R[0, 1]) / S 536 | qx = (R[0, 2] + R[2, 0]) / S 537 | qy = (R[1, 2] + R[2, 1]) / S 538 | qz = 0.25 * S 539 | 540 | return np.array([qx, qy, qz, qw]) 541 | 542 | def move_joints(self, target_joint_pos, speed = 0.01, timeout = 3): 543 | 544 | t0 = time.time() 545 | all_joints = [] 546 | 547 | while (time.time() - t0) < timeout: 548 | 549 | current_joint_pos = self.get_joint_positions() 550 | error = target_joint_pos - current_joint_pos 551 | 552 | if all(np.abs(error) < 1e-2): 553 | for _ in range(10): 554 | self.client_id.stepSimulation() # Give time to stop 555 | if self.benchmarking: 556 | self.check_collisions() 557 | if self.num_collisions > 0: 558 | break 559 | if self.benchmarking and (self.num_collisions > 0): 560 | break 561 | return True, all_joints 562 | 563 | norm = np.linalg.norm(error) 564 | vel = error / norm if norm > 0 else 0 565 | next_joint_pos = current_joint_pos + vel * speed 566 | all_joints.append(next_joint_pos) 567 | 568 | self.client_id.setJointMotorControlArray( # Move with constant velocity 569 | bodyIndex = self.manipulator, 570 | jointIndices = self.joints, 571 | controlMode = p.POSITION_CONTROL, 572 | targetPositions = next_joint_pos, 573 | positionGains = np.ones(len(self.joints)), 574 | ) 575 | 576 | self.client_id.stepSimulation() 577 | if self.benchmarking: 578 | self.check_collisions() 579 | if self.num_collisions > 0: 580 | break 581 | 582 | print(f"Warning: move_joints exceeded {timeout} second timeout. Skipping.") 583 | 584 | return False, all_joints 585 | 586 | def go_home(self): 587 | for i, joint_ind in enumerate(self.joints): 588 | self.client_id.resetJointState(self.manipulator, joint_ind, 0) 589 | # self.move_joints(np.zeros((7,))) 590 | 591 | def check_collisions(self): 592 | 593 | # There are 12 links. 594 | 595 | # Collision Info format: 596 | # ((0, 0, 7, 6, -1, (0.6231421349501027, -0.20546030368880996, 0.4579269918307523), 597 | # (0.6231421349501027, -0.20546030368881, 0.45900286827236414), (-4.511815760890017e-15, 2.2559078804450087e-14, -1.0), 598 | # 0.0010758764416118338, 5678.289055003103, 185.98474464217657, (0.0, 1.0, 2.2559078804450087e-14), 1407.3361669376961, 599 | # (1.0, 1.0178240730107782e-28, -4.511815760890017e-15)),) 600 | 601 | self.num_timesteps += 1 602 | 603 | for obs_id in range(len(self.obs_ids)): 604 | 605 | info = self.client_id.getContactPoints(self.manipulator, self.obs_ids[obs_id]) #, link_index) 606 | if len(info) > 0: 607 | self.num_collisions += 1 608 | break 609 | 610 | def execute_trajectory(self, trajectory): 611 | 612 | for i, joint_ind in enumerate(self.joints): 613 | self.client_id.resetJointState(self.manipulator, joint_ind, trajectory[i, 0]) 614 | 615 | _ = input("Press Enter to execute trajectory") 616 | 617 | for i in range(1, trajectory.shape[-1]): 618 | time.sleep(0.4) 619 | # current_joints = np.array([self.client_id.getJointState(self.manipulator, i)[0] for i in self.joints]) 620 | target_joints = trajectory[:, i] 621 | # print(f"Current Joints: {current_joints}") 622 | # print(f"Target Joints: {target_joints}") 623 | # print(f"Itr number: {i}") 624 | 625 | if any(target_joints <= self.joint_lower_limits) or any(target_joints >= self.joint_upper_limits): 626 | 627 | print("Joint Limits Exceeded") 628 | 629 | self.move_joints(target_joints) 630 | self.client_id.stepSimulation() 631 | 632 | def benchmark_trajectory(self, trajectory): # , guide, metrics_calculator): 633 | 634 | self.benchmarking = True 635 | 636 | for i, joint_ind in enumerate(self.joints): 637 | self.client_id.resetJointState(self.manipulator, joint_ind, trajectory[i, 0]) 638 | 639 | self.client_id.setJointMotorControlArray( # Stabilize Robot 640 | bodyIndex = self.manipulator, 641 | jointIndices = self.joints, 642 | controlMode = p.POSITION_CONTROL, 643 | targetPositions = trajectory[:, 0], 644 | positionGains = np.zeros((len(self.joints))), 645 | ) 646 | 647 | for i in range(10): 648 | self.client_id.stepSimulation() 649 | 650 | self.num_timesteps = 0 651 | self.num_collisions = 0 652 | 653 | dt = [0] 654 | 655 | for i in range(1, trajectory.shape[-1]): 656 | time.sleep(0.4) 657 | target_joints = trajectory[:, i] 658 | 659 | if any(target_joints <= self.joint_lower_limits) or any(target_joints >= self.joint_upper_limits): 660 | 661 | print("Joint Limits Exceeded") 662 | 663 | self.move_joints(target_joints) 664 | self.client_id.stepSimulation() 665 | self.check_collisions() 666 | if self.num_collisions > 0: 667 | break 668 | dt.append(self.num_timesteps * self.timestep) 669 | 670 | dt = np.diff(dt, n=1)[:, np.newaxis] 671 | 672 | collision_percentage = np.round((self.num_collisions / self.num_timesteps) * 100, 2) 673 | success = 1 if self.num_collisions == 0 else 0 674 | 675 | # joint_smoothness, end_eff_smoothness = metrics_calculator.smoothness_metric(trajectory, self.timestep) 676 | # joint_path_length, end_eff_path_length = metrics_calculator.path_length_metric(trajectory) 677 | 678 | self.benchmarking = False 679 | 680 | return success #, joint_path_length, end_eff_path_length, joint_smoothness[0], end_eff_smoothness[0] 681 | 682 | 683 | 684 | 685 | 686 | 687 | 688 | 689 | 690 | -------------------------------------------------------------------------------- /lib/metrics.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | class MetricsCalculator: 5 | 6 | def __init__(self, guide): 7 | 8 | self.guide = guide 9 | self.device = self.guide.device 10 | 11 | def smoothness_metric(self, joints, dt): 12 | """ 13 | dt is time taken between each step 14 | joints is (7, 50) numpy array 15 | """ 16 | 17 | joint_tensor = self.guide.rearrange_joints(torch.tensor(joints, dtype = torch.float32, device = self.device).unsqueeze(0)) 18 | end_eff_transforms = self.guide.get_end_effector_transform(joint_tensor) 19 | 20 | end_eff_positions = (end_eff_transforms[0, :, :3, 3]).numpy(force=True) 21 | # end_eff_positions is (50, 3) numpy array 22 | 23 | reshaped_joints = joints.T 24 | # reshaped_joints is (50, 7) numpy array 25 | joint_smoothness = np.linalg.norm(np.diff(reshaped_joints, n=1, axis=0) / dt, axis=1) 26 | joint_smoothness = self.sparc(joint_smoothness, 1. / dt) 27 | 28 | end_eff_smoothness = np.linalg.norm(np.diff(end_eff_positions, n=1, axis=0) / dt, axis=1) 29 | end_eff_smoothness = self.sparc(end_eff_smoothness, 1. / dt) 30 | 31 | return joint_smoothness, end_eff_smoothness 32 | 33 | def path_length_metric(self, joints): 34 | 35 | joint_tensor = self.guide.rearrange_joints(torch.tensor(joints, dtype = torch.float32, device = self.device).unsqueeze(0)) 36 | end_eff_transforms = self.guide.get_end_effector_transform(joint_tensor) 37 | end_eff_positions = (end_eff_transforms[0, :, :3, 3]).numpy(force=True) 38 | # end_eff_positions is (50, 3) numpy array 39 | 40 | reshaped_joints = joints.T 41 | 42 | end_eff_path_length = np.sum(np.linalg.norm(np.diff(end_eff_positions, 1, axis=0), axis=1)) 43 | joint_path_length = np.sum(np.linalg.norm(np.diff(reshaped_joints, 1, axis=0), axis=1)) 44 | 45 | return joint_path_length, end_eff_path_length 46 | 47 | def sparc(self, movement, fs, padlevel=4, fc=10.0, amp_th=0.05): 48 | """ 49 | Calculates the smoothness of the given speed profile using the modified 50 | spectral arc length metric. 51 | Parameters 52 | ---------- 53 | movement : np.array 54 | The array containing the movement speed profile. 55 | fs : float 56 | The sampling frequency of the data. 57 | padlevel : integer, optional 58 | Indicates the amount of zero padding to be done to the movement 59 | data for estimating the spectral arc length. [default = 4] 60 | fc : float, optional 61 | The max. cut off frequency for calculating the spectral arc 62 | length metric. [default = 10.] 63 | amp_th : float, optional 64 | The amplitude threshold to used for determing the cut off 65 | frequency upto which the spectral arc length is to be estimated. 66 | [default = 0.05] 67 | Returns 68 | ------- 69 | sal : float 70 | The spectral arc length estimate of the given movement's 71 | smoothness. 72 | (f, Mf) : tuple of two np.arrays 73 | This is the frequency(f) and the magntiude spectrum(Mf) of the 74 | given movement data. This spectral is from 0. to fs/2. 75 | (f_sel, Mf_sel) : tuple of two np.arrays 76 | This is the portion of the spectrum that is selected for 77 | calculating the spectral arc length. 78 | Notes 79 | ----- 80 | This is the modfieid spectral arc length metric, which has been tested only 81 | for discrete movements. 82 | 83 | Examples 84 | -------- 85 | >>> t = np.arange(-1, 1, 0.01) 86 | >>> move = np.exp(-5*pow(t, 2)) 87 | >>> sal, _, _ = sparc(move, fs=100.) 88 | >>> '%.5f' % sal 89 | '-1.41403' 90 | """ 91 | if np.allclose(movement, 0): 92 | print("All movement was 0, returning 0") 93 | return 0, None, None 94 | # Number of zeros to be padded. 95 | nfft = int(pow(2, np.ceil(np.log2(len(movement))) + padlevel)) 96 | 97 | # Frequency 98 | f = np.arange(0, fs, fs / nfft) 99 | # Normalized magnitude spectrum 100 | Mf = abs(np.fft.fft(movement, nfft)) 101 | Mf = Mf / max(Mf) 102 | 103 | # Indices to choose only the spectrum within the given cut off frequency 104 | # Fc. 105 | # NOTE: This is a low pass filtering operation to get rid of high frequency 106 | # noise from affecting the next step (amplitude threshold based cut off for 107 | # arc length calculation). 108 | fc_inx = ((f <= fc) * 1).nonzero() 109 | f_sel = f[fc_inx] 110 | Mf_sel = Mf[fc_inx] 111 | 112 | # Choose the amplitude threshold based cut off frequency. 113 | # Index of the last point on the magnitude spectrum that is greater than 114 | # or equal to the amplitude threshold. 115 | inx = ((Mf_sel >= amp_th) * 1).nonzero()[0] 116 | fc_inx = range(inx[0], inx[-1] + 1) 117 | f_sel = f_sel[fc_inx] 118 | Mf_sel = Mf_sel[fc_inx] 119 | 120 | # Calculate arc length 121 | new_sal = -sum( 122 | np.sqrt( 123 | pow(np.diff(f_sel) / (f_sel[-1] - f_sel[0]), 2) + pow(np.diff(Mf_sel), 2) 124 | ) 125 | ) 126 | return new_sal, (f, Mf), (f_sel, Mf_sel) -------------------------------------------------------------------------------- /mpinets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/mpinets/__init__.py -------------------------------------------------------------------------------- /mpinets/geometry.py: -------------------------------------------------------------------------------- 1 | # MIT License 2 | # 3 | # Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved. 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a 6 | # copy of this software and associated documentation files (the "Software"), 7 | # to deal in the Software without restriction, including without limitation 8 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 9 | # and/or sell copies of the Software, and to permit persons to whom the 10 | # Software is furnished to do so, subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in 13 | # all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 18 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 20 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 21 | # DEALINGS IN THE SOFTWARE. 22 | 23 | from typing import Union, Sequence 24 | import numpy as np 25 | import random 26 | import torch 27 | from geometrout.primitive import Sphere, Cuboid, Cylinder 28 | 29 | 30 | class TorchSpheres: 31 | """ 32 | A Pytorch representation of a batch of M spheres (i.e. B elements in the batch, 33 | M spheres per element). Any of these spheres can have zero volume (these 34 | will be masked out during calculation of the various functions in this 35 | class, such as sdf). 36 | """ 37 | 38 | def __init__(self, centers: torch.Tensor, radii: torch.Tensor): 39 | """ 40 | :param centers torch.Tensor: a set of centers, has dim [B, M, 3] 41 | :param radii torch.Tensor: a set of radii, has dim [B, M, 1] 42 | """ 43 | assert centers.ndim == 3 44 | assert radii.ndim == 3 45 | 46 | # TODO It would be more memory efficient to rely more heavily on broadcasting 47 | # in some cases where multiple spheres have the same size 48 | assert centers.ndim == radii.ndim 49 | 50 | # This logic is to determine the batch size. Either batch sizes need to 51 | # match or if only one variable has a batch size, that one is assumed to 52 | # be the batch size 53 | 54 | self.centers = centers 55 | self.radii = radii 56 | self.mask = ~torch.isclose(self.radii, torch.zeros(1).type_as(centers)).squeeze( 57 | -1 58 | ) 59 | 60 | def surface_area(self) -> torch.Tensor: 61 | """ 62 | Calculates the surface area of the spheres 63 | 64 | :rtype torch.Tensor: A tensor of the surface areas of the spheres 65 | """ 66 | area = 4 * np.pi * torch.pow(self.radii, 3) 67 | return area 68 | 69 | def sample_surface(self, num_points: int) -> torch.Tensor: 70 | """ 71 | Samples points from all spheres, including ones with zero volume 72 | 73 | :param num_points int: The number of points to sample per sphere 74 | :rtype torch.Tensor: The points, has dim [B, M, N] 75 | """ 76 | B, M, _ = self.centers.shape 77 | unnormalized_points = torch.rand((B, M, num_points, 3)) 78 | normalized = ( 79 | unnormalized_points 80 | / torch.linalg.norm(unnormalized_points, dim=-1)[:, :, :, None] 81 | ) 82 | random_points = ( 83 | normalized * self.radii[:, :, None, :] + self.centers[:, :, None, :] 84 | ) 85 | return random_points 86 | 87 | def sdf(self, points: torch.Tensor) -> torch.Tensor: 88 | """ 89 | :param points torch.Tensor: The points with which to calculate the 90 | SDF, has dim [B, N, 3] (N is the number of points) 91 | :rtype torch.Tensor: The scene SDF value for each point (i.e. the minimum SDF 92 | value for each of the M spheres), has dim [B, N] 93 | """ 94 | assert points.ndim == 3 95 | B, M, _ = self.radii.shape 96 | _, N, _ = points.shape 97 | all_sdfs = float("inf") * torch.ones(B, M, N).type_as(points) 98 | distances = points[:, None, :, :] - self.centers[:, :, None, :] 99 | all_sdfs[self.mask] = ( 100 | torch.linalg.norm(distances[self.mask], dim=-1) - self.radii[self.mask] 101 | ) 102 | return torch.min(all_sdfs, dim=1)[0] 103 | 104 | def sdf_sequence(self, points: torch.Tensor) -> torch.Tensor: 105 | """ 106 | Calculates SDF values for a time sequence of point clouds 107 | :param points torch.Tensor: The batched sequence of point clouds with 108 | dimension [B, T, N, 3] (T in sequence length, 109 | N is number of points) 110 | :rtype torch.Tensor: The scene SDF for each point at each timestep (i.e. the minimum 111 | SDF value across the M spheres at each timestep), 112 | has dim [B, T, N] 113 | """ 114 | assert points.ndim == 4 115 | B, M, _ = self.radii.shape 116 | _, T, N, _ = points.shape 117 | all_sdfs = float("inf") * torch.ones(B, M, T, N).type_as(points) 118 | distances = points[:, None, :, :] - self.centers[:, :, None, None, :] 119 | all_sdfs[self.mask] = ( 120 | torch.linalg.norm(distances[self.mask], dim=-1) 121 | - self.radii[:, :, None, :][self.mask] 122 | ) 123 | return torch.min(all_sdfs, dim=1)[0] 124 | 125 | 126 | class TorchCuboids: 127 | """ 128 | A Pytorch representation of a batch of M cuboids (i.e. B elements in the batch, 129 | M cuboids per element). Any of these cuboids can have zero volume (these 130 | will be masked out during calculation of the various functions in this 131 | class, such as sdf). 132 | """ 133 | 134 | def __init__( 135 | self, centers: torch.Tensor, dims: torch.Tensor, quaternions: torch.Tensor 136 | ): 137 | """ 138 | :param centers torch.Tensor: Has dim [B, M, 3] 139 | :param dims torch.Tensor: Has dim [B, M, 3] 140 | :param quaternions torch.Tensor: Has dim [B, M, 4] with quaternions formatted as 141 | w, x, y, z 142 | """ 143 | assert centers.ndim == 3 144 | assert dims.ndim == 3 145 | assert quaternions.ndim == 3 146 | 147 | self.dims = dims 148 | self.centers = centers 149 | 150 | # It's helpful to ensure the quaternions are normalized 151 | self.quats = quaternions / torch.linalg.norm(quaternions, dim=2)[:, :, None] 152 | 153 | self._init_frames() 154 | # Mask for nonzero volumes 155 | self.mask = ~torch.any( 156 | torch.isclose(self.dims, torch.zeros(1).type_as(centers)), dim=-1 157 | ) 158 | 159 | def geometrout(self): 160 | """ 161 | Helper method to convert this into geometrout primitives 162 | """ 163 | B, M, _ = self.centers.shape 164 | return [ 165 | [ 166 | Cuboid( 167 | center=self.centers[bidx, midx, :].detach().cpu().numpy(), 168 | dims=self.dims[bidx, midx, :].detach().cpu().numpy(), 169 | quaternion=self.quats[bidx, midx, :].detach().cpu().numpy(), 170 | ) 171 | for midx in range(M) 172 | if self.mask[bidx, midx] 173 | ] 174 | for bidx in range(B) 175 | ] 176 | 177 | def _init_frames(self): 178 | """ 179 | In order to calculate the SDF, we need to calculate the inverse 180 | transformation of the cuboid. This is because we are transforming points 181 | in the world frame into the cuboid frame. 182 | """ 183 | 184 | # Initialize the inverse rotation 185 | w = self.quats[:, :, 0] 186 | x = -self.quats[:, :, 1] 187 | y = -self.quats[:, :, 2] 188 | z = -self.quats[:, :, 3] 189 | 190 | # Naming is a little disingenuous here because I'm multiplying everything by two, 191 | # but can't put numbers at the beginning of variable names. 192 | xx = 2 * torch.pow(x, 2) 193 | yy = 2 * torch.pow(y, 2) 194 | zz = 2 * torch.pow(z, 2) 195 | 196 | wx = 2 * w * x 197 | wy = 2 * w * y 198 | wz = 2 * w * z 199 | xy = 2 * x * y 200 | xz = 2 * x * z 201 | yz = 2 * y * z 202 | 203 | B, M, _ = self.centers.shape 204 | B = self.centers.size(0) 205 | M = self.centers.size(1) 206 | self.inv_frames = torch.zeros((B, M, 4, 4)).type_as(self.centers) 207 | self.inv_frames[:, :, 3, 3] = 1 208 | 209 | R = torch.stack( 210 | [ 211 | torch.stack([1 - yy - zz, xy - wz, xz + wy], dim=2), 212 | torch.stack([xy + wz, 1 - xx - zz, yz - wx], dim=2), 213 | torch.stack([xz - wy, yz - wx, 1 - xx - yy], dim=2), 214 | ], 215 | dim=2, 216 | ) 217 | Rt = torch.matmul(R, -1 * self.centers.unsqueeze(3)).squeeze(3) 218 | 219 | # Fill in the rotation matrices 220 | self.inv_frames[:, :, :3, :3] = R 221 | 222 | # Invert the transform by multiplying the inverse translation by the inverse rotation 223 | self.inv_frames[:, :, :3, 3] = Rt 224 | 225 | def surface_area(self) -> torch.Tensor: 226 | """ 227 | Calculates the surface area of the cuboids 228 | 229 | :rtype torch.Tensor: A tensor of the surface areas of the cuboids 230 | """ 231 | area = 2 * ( 232 | self.dims[:, :, 0] * self.dims[:, :, 1] 233 | + self.dims[:, :, 0] * self.dims[:, :, 2] 234 | + self.dims[:, :, 1] * self.dims[:, :, 2] 235 | ) 236 | return area 237 | 238 | def sdf(self, points: torch.Tensor) -> torch.Tensor: 239 | """ 240 | :param points torch.Tensor: The points with which to calculate the SDF, has 241 | dim [B, N, 3] (N is the number of points) 242 | :rtype torch.Tensor: The scene SDF value for each point (i.e. the minimum SDF 243 | value for each of the M cuboids), has dim [B, N] 244 | """ 245 | assert points.ndim == 3 246 | # We are going to need to map points in the global frame to the cuboid frame 247 | # First take the points and make them homogeneous by adding a one to the end 248 | 249 | # points_from_volumes = points[self.nonzero_volumes, :, :] 250 | assert points.size(0) == self.centers.size(0) 251 | if torch.all(~self.mask): 252 | return float("inf") * torch.ones(points.size(0), points.size(1)).type_as( 253 | points 254 | ) 255 | 256 | homog_points = torch.cat( 257 | ( 258 | points, 259 | torch.ones((points.size(0), points.size(1), 1)).type_as(points), 260 | ), 261 | dim=2, 262 | ) 263 | # Next, project these points into their respective cuboid frames 264 | # Will return [B x M x N x 3] 265 | 266 | points_proj = torch.matmul( 267 | self.inv_frames[:, :, None, :, :], homog_points[:, None, :, :, None] 268 | ).squeeze(-1)[:, :, :, :3] 269 | B, M, N, _ = points_proj.shape 270 | masked_points = points_proj[self.mask] 271 | 272 | # The follow computations are adapted from here 273 | # https://github.com/fogleman/sdf/blob/main/sdf/d3.py 274 | # Move points to top corner 275 | 276 | distances = torch.abs(masked_points) - (self.dims[self.mask] / 2)[:, None, :] 277 | # This is distance only for points outside the box, all points inside return zero 278 | # This probably needs to be fixed or else there will be a nan gradient 279 | 280 | outside = torch.linalg.norm( 281 | torch.maximum(distances, torch.zeros_like(distances)), dim=-1 282 | ) 283 | # This is distance for points inside the box, all others return zero 284 | inner_max_distance = torch.max(distances, dim=-1).values 285 | inside = torch.minimum(inner_max_distance, torch.zeros_like(inner_max_distance)) 286 | all_sdfs = float("inf") * torch.ones(B, M, N).type_as(points) 287 | all_sdfs[self.mask] = outside + inside 288 | return torch.min(all_sdfs, dim=1)[0] 289 | 290 | def sdf_sequence(self, points: torch.Tensor) -> torch.Tensor: 291 | """ 292 | Calculates SDF values for a time sequence of point clouds 293 | :param points torch.Tensor: The batched sequence of point clouds with 294 | dimension [B, T, N, 3] (T in sequence length, 295 | N is number of points) 296 | :rtype torch.Tensor: The scene SDF for each point at each timestep 297 | (i.e. the minimum SDF value across the M cuboids 298 | at each timestep), has dim [B, T, N] 299 | """ 300 | assert points.ndim == 4 301 | 302 | # We are going to need to map points in the global frame to the cuboid frame 303 | # First take the points and make them homogeneous by adding a one to the end 304 | 305 | # points_from_volumes = points[self.nonzero_volumes, :, :] 306 | assert points.size(0) == self.centers.size(0) 307 | if torch.all(~self.mask): 308 | return float("inf") * torch.ones(points.shape[:-1]).type_as(points) 309 | 310 | homog_points = torch.cat( 311 | ( 312 | points, 313 | torch.ones((*points.shape[:-1], 1)).type_as(points), 314 | ), 315 | dim=3, 316 | ) 317 | # Next, project these points into their respective cuboid frames 318 | # Will return [B x M x N x 3] 319 | 320 | points_proj = torch.matmul( 321 | self.inv_frames[:, :, None, None, :, :], 322 | homog_points[:, None, :, :, :, None], 323 | ).squeeze(-1)[:, :, :, :, :3] 324 | B, M, T, N, _ = points_proj.shape 325 | assert T == points.size(1) 326 | assert N == points.size(2) 327 | masked_points = points_proj[self.mask] 328 | 329 | # The follow computations are adapted from here 330 | # https://github.com/fogleman/sdf/blob/main/sdf/d3.py 331 | # Move points to top corner 332 | 333 | distances = ( 334 | torch.abs(masked_points) - (self.dims[self.mask] / 2)[:, None, None, :] 335 | ) 336 | # This is distance only for points outside the box, all points inside return zero 337 | # This probably needs to be fixed or else there will be a nan gradient 338 | 339 | outside = torch.linalg.norm( 340 | torch.maximum(distances, torch.zeros_like(distances)), dim=-1 341 | ) 342 | # This is distance for points inside the box, all others return zero 343 | inner_max_distance = torch.max(distances, dim=-1).values 344 | inside = torch.minimum(inner_max_distance, torch.zeros_like(inner_max_distance)) 345 | all_sdfs = float("inf") * torch.ones(B, M, T, N).type_as(points) 346 | all_sdfs[self.mask] = outside + inside 347 | return torch.min(all_sdfs, dim=1)[0] 348 | 349 | 350 | class TorchCylinders: 351 | """ 352 | A Pytorch representation of a batch of M cylinders (i.e. B elements in the batch, 353 | M cylinders per element). Any of these cylinders can have zero volume (these 354 | will be masked out during calculation of the various functions in this 355 | class, such as sdf). 356 | """ 357 | 358 | def __init__( 359 | self, 360 | centers: torch.Tensor, 361 | radii: torch.Tensor, 362 | heights: torch.Tensor, 363 | quaternions: torch.Tensor, 364 | ): 365 | """ 366 | :param centers torch.Tensor: Has dim [B, M, 3] 367 | :param radii torch.Tensor: Has dim [B, M, 1] 368 | :param heights torch.Tensor: Has dim [B, M, 1] 369 | :param quaternions torch.Tensor: Has dim [B, M, 4] with quaternions formatted as 370 | (w, x, y, z) 371 | """ 372 | assert centers.ndim == 3 373 | assert radii.ndim == 3 374 | assert heights.ndim == 3 375 | assert quaternions.ndim == 3 376 | 377 | self.radii = radii 378 | self.heights = heights 379 | self.centers = centers 380 | 381 | # It's helpful to ensure the quaternions are normalized 382 | self.quats = quaternions / torch.linalg.norm(quaternions, dim=2)[:, :, None] 383 | self._init_frames() 384 | # Mask for nonzero volumes 385 | self.mask = ~torch.logical_or( 386 | torch.isclose(self.radii, torch.zeros(1).type_as(centers)).squeeze(-1), 387 | torch.isclose(self.heights, torch.zeros(1).type_as(centers)).squeeze(-1), 388 | ) 389 | 390 | def geometrout(self): 391 | """ 392 | Helper method to convert this into geometrout primitives 393 | """ 394 | B, M, _ = self.centers.shape 395 | return [ 396 | [ 397 | Cylinder( 398 | center=self.centers[bidx, midx, :].detach().cpu().numpy(), 399 | radius=self.radii[bidx, midx, 0].detach().cpu().numpy(), 400 | height=self.heights[bidx, midx, 0].detach().cpu().numpy(), 401 | quaternion=self.quats[bidx, midx, :].detach().cpu().numpy(), 402 | ) 403 | for midx in range(M) 404 | if self.mask[bidx, midx] 405 | ] 406 | for bidx in range(B) 407 | ] 408 | 409 | def _init_frames(self): 410 | """ 411 | In order to calculate the SDF, we need to calculate the inverse 412 | transformation of the cylinder. This is because we are transforming 413 | points in the world frame into the cylinder frame. 414 | """ 415 | # Initialize the inverse rotation 416 | w = self.quats[:, :, 0] 417 | x = -self.quats[:, :, 1] 418 | y = -self.quats[:, :, 2] 419 | z = -self.quats[:, :, 3] 420 | 421 | # Naming is a little disingenuous here because I'm multiplying everything by two, 422 | # but can't put numbers at the beginning of variable names. 423 | xx = 2 * torch.pow(x, 2) 424 | yy = 2 * torch.pow(y, 2) 425 | zz = 2 * torch.pow(z, 2) 426 | 427 | wx = 2 * w * x 428 | wy = 2 * w * y 429 | wz = 2 * w * z 430 | xy = 2 * x * y 431 | xz = 2 * x * z 432 | yz = 2 * y * z 433 | 434 | B, M, _ = self.centers.shape 435 | B = self.centers.size(0) 436 | M = self.centers.size(1) 437 | self.inv_frames = torch.zeros((B, M, 4, 4)).type_as(self.centers) 438 | self.inv_frames[:, :, 3, 3] = 1 439 | 440 | R = torch.stack( 441 | [ 442 | torch.stack([1 - yy - zz, xy - wz, xz + wy], dim=2), 443 | torch.stack([xy + wz, 1 - xx - zz, yz - wx], dim=2), 444 | torch.stack([xz - wy, yz - wx, 1 - xx - yy], dim=2), 445 | ], 446 | dim=2, 447 | ) 448 | Rt = torch.matmul(R, -1 * self.centers.unsqueeze(3)).squeeze(3) 449 | 450 | # Fill in the rotation matrices 451 | self.inv_frames[:, :, :3, :3] = R 452 | 453 | # Invert the transform by multiplying the inverse translation by the inverse rotation 454 | self.inv_frames[:, :, :3, 3] = Rt 455 | 456 | def sdf(self, points: torch.Tensor) -> torch.Tensor: 457 | """ 458 | :param points torch.Tensor: The points with which to calculate the SDF, has 459 | dim [B, N, 3] (N is the number of points) 460 | :rtype torch.Tensor: The scene SDF value for each point (i.e. the minimum SDF 461 | value for each of the M cylinders), has dim [B, N] 462 | """ 463 | assert points.ndim == 3 464 | assert points.size(0) == self.centers.size(0) 465 | if torch.all(~self.mask): 466 | return float("inf") * torch.ones(points.size(0), points.size(1)).type_as( 467 | points 468 | ) 469 | 470 | homog_points = torch.cat( 471 | ( 472 | points, 473 | torch.ones((points.size(0), points.size(1), 1)).type_as(points), 474 | ), 475 | dim=2, 476 | ) 477 | # Next, project these points into their respective cylinder frames 478 | # Will return [B x M x N x 3] 479 | 480 | points_proj = torch.matmul( 481 | self.inv_frames[:, :, None, :, :], homog_points[:, None, :, :, None] 482 | ).squeeze(-1)[:, :, :, :3] 483 | B, M, N, _ = points_proj.shape 484 | masked_points = points_proj[self.mask] 485 | 486 | surface_distance_xy = torch.linalg.norm(masked_points[:, :, :2], dim=2) 487 | z_distance = masked_points[:, :, 2] 488 | 489 | half_extents_2d = torch.stack( 490 | (self.radii[self.mask], self.heights[self.mask] / 2), dim=2 491 | ) 492 | points_2d = torch.stack((surface_distance_xy, z_distance), dim=2) 493 | distances_2d = torch.abs(points_2d) - half_extents_2d 494 | 495 | # This is distance only for points outside the box, all points inside return zero 496 | outside = torch.linalg.norm( 497 | torch.maximum(distances_2d, torch.zeros_like(distances_2d)), dim=2 498 | ) 499 | # This is distance for points inside the box, all others return zero 500 | inner_max_distance_2d = torch.max(distances_2d, dim=2).values 501 | inside = torch.minimum( 502 | inner_max_distance_2d, torch.zeros_like(inner_max_distance_2d) 503 | ) 504 | 505 | all_sdfs = float("inf") * torch.ones(B, M, N).type_as(points) 506 | all_sdfs[self.mask] = outside + inside 507 | return torch.min(all_sdfs, dim=1)[0] 508 | 509 | def sdf_sequence(self, points: torch.Tensor) -> torch.Tensor: 510 | """ 511 | Calculates SDF values for a time sequence of point clouds 512 | :param points torch.Tensor: The batched sequence of point clouds with 513 | dimension [B, T, N, 3] (T in sequence length, 514 | N is number of points) 515 | :rtype torch.Tensor: The scene SDF for each point at each timestep 516 | (i.e. the minimum SDF value across the M cylinders at 517 | each timestep), has dim [B, T, N] 518 | """ 519 | assert points.ndim == 4 520 | 521 | # We are going to need to map points in the global frame to the cylinder frame 522 | # First take the points and make them homogeneous by adding a one to the end 523 | 524 | assert points.size(0) == self.centers.size(0) 525 | if torch.all(~self.mask): 526 | return float("inf") * torch.ones(points.shape[:-1]).type_as(points) 527 | 528 | homog_points = torch.cat( 529 | ( 530 | points, 531 | torch.ones((*points.shape[:-1], 1)).type_as(points), 532 | ), 533 | dim=3, 534 | ) 535 | # Next, project these points into their respective cylinder frames 536 | # Will return [B x M x N x 3] 537 | 538 | points_proj = torch.matmul( 539 | self.inv_frames[:, :, None, None, :, :], 540 | homog_points[:, None, :, :, :, None], 541 | ).squeeze(-1)[:, :, :, :, :3] 542 | B, M, T, N, _ = points_proj.shape 543 | assert T == points.size(1) 544 | assert N == points.size(2) 545 | masked_points = points_proj[self.mask] 546 | 547 | surface_distance_xy = torch.linalg.norm(masked_points[:, :, :, :2], dim=-1) 548 | z_distance = masked_points[:, :, :, 2] 549 | 550 | half_extents_2d = torch.stack( 551 | (self.radii[self.mask], self.heights[self.mask] / 2), dim=2 552 | )[:, :, None, :] 553 | points_2d = torch.stack((surface_distance_xy, z_distance), dim=3) 554 | distances_2d = torch.abs(points_2d) - half_extents_2d 555 | 556 | # This is distance only for points outside the box, all points inside return zero 557 | outside = torch.linalg.norm( 558 | torch.maximum(distances_2d, torch.zeros_like(distances_2d)), dim=3 559 | ) 560 | # This is distance for points inside the box, all others return zero 561 | inner_max_distance_2d = torch.max(distances_2d, dim=3).values 562 | inside = torch.minimum( 563 | inner_max_distance_2d, torch.zeros_like(inner_max_distance_2d) 564 | ) 565 | 566 | all_sdfs = float("inf") * torch.ones(B, M, T, N).type_as(points) 567 | all_sdfs[self.mask] = outside + inside 568 | return torch.min(all_sdfs, dim=1)[0] 569 | 570 | 571 | def construct_mixed_point_cloud( 572 | obstacles: Sequence[Union[Sphere, Cuboid, Cylinder]], num_points: int 573 | ) -> np.ndarray: 574 | """ 575 | Creates a random point cloud from a collection of obstacles. The points in 576 | the point cloud should be fairly(-ish) distributed amongst the obstacles based 577 | on their surface area. 578 | 579 | :param obstacles Sequence[Union[Sphere, Cuboid, Cylinder]]: The obstacles in the scene 580 | :param num_points int: The total number of points in the samples scene (not 581 | the number of points per obstacle) 582 | :rtype np.ndarray: Has dim [N, 3] where N is num_points 583 | """ 584 | point_set = [] 585 | total_obstacles = len(obstacles) 586 | if total_obstacles == 0: 587 | return np.array([[]]) 588 | 589 | # Allocate points based on obstacle surface area for even sampling 590 | surface_areas = np.array([o.surface_area for o in obstacles]) 591 | total_area = np.sum(surface_areas) 592 | proportions = (surface_areas / total_area).tolist() 593 | 594 | indices = list(range(1, total_obstacles + 1)) 595 | random.shuffle(indices) 596 | idx = 0 597 | 598 | for o, prop in zip(obstacles, proportions): 599 | sample_number = int(prop * num_points) + 500 600 | samples = o.sample_surface(sample_number) 601 | _points = indices[idx] * np.ones((sample_number, 4)) 602 | _points[:, :3] = samples 603 | point_set.append(_points) 604 | idx += 1 605 | points = np.concatenate(point_set, axis=0) 606 | 607 | # Downsample to the desired number of points 608 | return points[np.random.choice(points.shape[0], num_points, replace=False), :] 609 | -------------------------------------------------------------------------------- /mpinets/loss.py: -------------------------------------------------------------------------------- 1 | # MIT License 2 | # 3 | # Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved. 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a 6 | # copy of this software and associated documentation files (the "Software"), 7 | # to deal in the Software without restriction, including without limitation 8 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 9 | # and/or sell copies of the Software, and to permit persons to whom the 10 | # Software is furnished to do so, subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in 13 | # all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 18 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 20 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 21 | # DEALINGS IN THE SOFTWARE. 22 | 23 | from typing import Tuple 24 | from mpinets import utils 25 | from mpinets.geometry import TorchCuboids, TorchCylinders 26 | import torch.nn.functional as F 27 | import torch 28 | from robofin.pointcloud.torch import FrankaSampler 29 | 30 | 31 | def point_match_loss(input_pc: torch.Tensor, target_pc: torch.Tensor) -> torch.Tensor: 32 | """ 33 | A combination L1 and L2 loss to penalize large and small deviations between 34 | two point clouds 35 | 36 | :param input_pc torch.Tensor: Point cloud sampled from the network's output. 37 | Has dim [B, N, 3] 38 | :param target_pc torch.Tensor: Point cloud sampled from the supervision 39 | Has dim [B, N, 3] 40 | :rtype torch.Tensor: The single loss value 41 | """ 42 | return F.mse_loss(input_pc, target_pc, reduction="mean") + F.l1_loss( 43 | input_pc, target_pc, reduction="mean" 44 | ) 45 | 46 | 47 | def collision_loss( 48 | input_pc: torch.Tensor, 49 | cuboid_centers: torch.Tensor, 50 | cuboid_dims: torch.Tensor, 51 | cuboid_quaternions: torch.Tensor, 52 | cylinder_centers: torch.Tensor, 53 | cylinder_radii: torch.Tensor, 54 | cylinder_heights: torch.Tensor, 55 | cylinder_quaternions: torch.Tensor, 56 | ) -> torch.Tensor: 57 | """ 58 | Calculates the hinge loss, calculating whether the robot (represented as a 59 | point cloud) is in collision with any obstacles in the scene. Collision 60 | here actually means within 3cm of the obstacle--this is to provide stronger 61 | gradient signal to encourage the robot to move out of the way. Also, some of the 62 | primitives can have zero volume (i.e. a dim is zero for cuboids or radius or height is zero for cylinders). 63 | If these are zero volume, they will have infinite sdf values (and therefore be ignored by the loss). 64 | 65 | :param input_pc torch.Tensor: Points sampled from the robot's surface after it 66 | is placed at the network's output prediction. Has dim [B, N, 3] 67 | :param cuboid_centers torch.Tensor: Has dim [B, M1, 3] 68 | :param cuboid_dims torch.Tensor: Has dim [B, M1, 3] 69 | :param cuboid_quaternions torch.Tensor: Has dim [B, M1, 4]. Quaternion is formatted as w, x, y, z. 70 | :param cylinder_centers torch.Tensor: Has dim [B, M2, 3] 71 | :param cylinder_radii torch.Tensor: Has dim [B, M2, 1] 72 | :param cylinder_heights torch.Tensor: Has dim [B, M2, 1] 73 | :param cylinder_quaternions torch.Tensor: Has dim [B, M2, 4]. Quaternion is formatted as w, x, y, z. 74 | :rtype torch.Tensor: Returns the loss value aggregated over the batch 75 | """ 76 | 77 | cuboids = TorchCuboids( 78 | cuboid_centers, 79 | cuboid_dims, 80 | cuboid_quaternions, 81 | ) 82 | cylinders = TorchCylinders( 83 | cylinder_centers, 84 | cylinder_radii, 85 | cylinder_heights, 86 | cylinder_quaternions, 87 | ) 88 | sdf_values = torch.minimum(cuboids.sdf(input_pc), cylinders.sdf(input_pc)) 89 | return F.hinge_embedding_loss( 90 | sdf_values, 91 | -torch.ones_like(sdf_values), 92 | margin=0.03, 93 | reduction="mean", 94 | ) 95 | 96 | 97 | class CollisionAndBCLossContainer: 98 | """ 99 | A container class to hold the various losses. This is structured as a 100 | container because that allows it to cache the robot pointcloud sampler 101 | object. By caching this, we reduce parsing time when processing the URDF 102 | and allow for a consistent random pointcloud (consistent per-GPU, that is) 103 | """ 104 | 105 | def __init__( 106 | self, 107 | ): 108 | self.fk_sampler = None 109 | self.num_points = 1024 110 | 111 | def __call__( 112 | self, 113 | input_normalized: torch.Tensor, 114 | cuboid_centers: torch.Tensor, 115 | cuboid_dims: torch.Tensor, 116 | cuboid_quaternions: torch.Tensor, 117 | cylinder_centers: torch.Tensor, 118 | cylinder_radii: torch.Tensor, 119 | cylinder_heights: torch.Tensor, 120 | cylinder_quaternions: torch.Tensor, 121 | target_normalized: torch.Tensor, 122 | ) -> Tuple[torch.Tensor, torch.Tensor]: 123 | """ 124 | This method calculates both constituent loss function after loading, 125 | and then caching, a fixed robot point cloud sampler (i.e. the task 126 | spaces sampled are always the same, as opposed to a random point cloud). 127 | The fixed point cloud is important for loss calculation so that 128 | it's possible to take mse between the two pointclouds. 129 | 130 | :param input_normalized torch.Tensor: Has dim [B, 7] and is always between -1 and 1 131 | :param cuboid_centers torch.Tensor: Has dim [B, M1, 3] 132 | :param cuboid_dims torch.Tensor: Has dim [B, M1, 3] 133 | :param cuboid_quaternions torch.Tensor: Has dim [B, M1, 4]. Quaternion is formatted as w, x, y, z. 134 | :param cylinder_centers torch.Tensor: Has dim [B, M2, 3] 135 | :param cylinder_radii torch.Tensor: Has dim [B, M2, 1] 136 | :param cylinder_heights torch.Tensor: Has dim [B, M2, 1] 137 | :param cylinder_quaternions torch.Tensor: Has dim [B, M2, 4]. Quaternion is formatted as w, x, y, z. 138 | :param target_normalized torch.Tensor: Has dim [B, 7] and is always between -1 and 1 139 | :rtype Tuple[torch.Tensor, torch.Tensor]: The two losses aggregated over the batch 140 | """ 141 | if self.fk_sampler is None: 142 | self.fk_sampler = FrankaSampler( 143 | input_normalized.device, 144 | num_fixed_points=self.num_points, 145 | use_cache=True, 146 | with_base_link=False, # Remove base link because this isn't controllable anyway 147 | ) 148 | input_pc = self.fk_sampler.sample( 149 | utils.unnormalize_franka_joints(input_normalized), 150 | ) 151 | target_pc = self.fk_sampler.sample( 152 | utils.unnormalize_franka_joints(target_normalized), 153 | ) 154 | return ( 155 | collision_loss( 156 | input_pc, 157 | cuboid_centers, 158 | cuboid_dims, 159 | cuboid_quaternions, 160 | cylinder_centers, 161 | cylinder_radii, 162 | cylinder_heights, 163 | cylinder_quaternions, 164 | ), 165 | point_match_loss(input_pc, target_pc), 166 | ) -------------------------------------------------------------------------------- /mpinets/model.py: -------------------------------------------------------------------------------- 1 | # MIT License 2 | # 3 | # Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved. 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a 6 | # copy of this software and associated documentation files (the "Software"), 7 | # to deal in the Software without restriction, including without limitation 8 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 9 | # and/or sell copies of the Software, and to permit persons to whom the 10 | # Software is furnished to do so, subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in 13 | # all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 18 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 20 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 21 | # DEALINGS IN THE SOFTWARE. 22 | 23 | import torch 24 | from torch import nn 25 | from robofin.pointcloud.torch import FrankaSampler, FrankaCollisionSampler 26 | import pytorch_lightning as pl 27 | from pointnet2_ops.pointnet2_modules import PointnetSAModule 28 | 29 | from mpinets import loss 30 | from mpinets.utils import unnormalize_franka_joints 31 | from mpinets.geometry import TorchCuboids, TorchCylinders 32 | from typing import List, Tuple, Sequence, Dict, Callable 33 | 34 | 35 | class MotionPolicyNetwork(pl.LightningModule): 36 | """ 37 | The architecture laid out here is the default architecture laid out in the 38 | Motion Policy Networks paper (Fishman, et. al, 2022). 39 | """ 40 | 41 | def __init__(self): 42 | """ 43 | Constructs the model 44 | """ 45 | super().__init__() 46 | self.point_cloud_encoder = MPiNetsPointNet() 47 | self.feature_encoder = nn.Sequential( 48 | nn.Linear(7, 32), 49 | nn.LeakyReLU(), 50 | nn.Linear(32, 64), 51 | nn.LeakyReLU(), 52 | nn.Linear(64, 128), 53 | nn.LeakyReLU(), 54 | nn.Linear(128, 128), 55 | nn.LeakyReLU(), 56 | nn.Linear(128, 64), 57 | ) 58 | self.decoder = nn.Sequential( 59 | nn.Linear(2048 + 64, 512), 60 | nn.LeakyReLU(), 61 | nn.Linear(512, 256), 62 | nn.LeakyReLU(), 63 | nn.Linear(256, 128), 64 | nn.LeakyReLU(), 65 | nn.Linear(128, 7), 66 | ) 67 | 68 | def configure_optimizers(self): 69 | """ 70 | A standard method in PyTorch lightning to set the optimizer 71 | """ 72 | optimizer = torch.optim.Adam(self.parameters(), lr=1e-4) 73 | return optimizer 74 | 75 | def forward(self, xyz: torch.Tensor, q: torch.Tensor) -> torch.Tensor: # type: ignore[override] 76 | """ 77 | Passes data through the network to produce an output 78 | 79 | :param xyz torch.Tensor: Tensor representing the point cloud. Should 80 | have dimensions of [B x N x 4] where B is the batch 81 | size, N is the number of points and 4 is because there 82 | are three geometric dimensions and a segmentation mask 83 | :param q torch.Tensor: The current robot configuration normalized to be between 84 | -1 and 1, according to each joint's range of motion 85 | :rtype torch.Tensor: The displacement to be applied to the current configuration to get 86 | the position at the next step (still in normalized space) 87 | """ 88 | pc_encoding = self.point_cloud_encoder(xyz) 89 | feature_encoding = self.feature_encoder(q) 90 | x = torch.cat((pc_encoding, feature_encoding), dim=1) 91 | return self.decoder(x) 92 | 93 | 94 | class TrainingMotionPolicyNetwork(MotionPolicyNetwork): 95 | """ 96 | An version of the MotionPolicyNetwork model that has additional attributes 97 | necessary during training (or using the validation step outside of the 98 | training process). This class is a valid model, but it's overkill when 99 | doing real robot inference and, for example, point cloud sampling is 100 | done by an outside process (such as downsampling point clouds from a point cloud). 101 | """ 102 | 103 | def __init__( 104 | self, 105 | num_robot_points: int, 106 | point_match_loss_weight: float, 107 | collision_loss_weight: float, 108 | ): 109 | """ 110 | Creates the network and assigns additional parameters for training 111 | 112 | 113 | :param num_robot_points int: The number of robot points used when resampling 114 | the robot points during rollouts (used in validation) 115 | :param point_match_loss_weight float: The weight assigned to the behavior 116 | cloning loss. 117 | :param collision_loss_weight float: The weight assigned to the collision loss 118 | :rtype Self: An instance of the network 119 | """ 120 | super().__init__() 121 | self.num_robot_points = num_robot_points 122 | self.point_match_loss_weight = point_match_loss_weight 123 | self.collision_loss_weight = collision_loss_weight 124 | self.fk_sampler = None 125 | self.collision_sampler = None 126 | self.loss_fun = loss.CollisionAndBCLossContainer() 127 | 128 | def rollout( 129 | self, 130 | batch: Dict[str, torch.Tensor], 131 | rollout_length: int, 132 | sampler: Callable[[torch.Tensor], torch.Tensor], 133 | unnormalize: bool = False, 134 | ) -> List[torch.Tensor]: 135 | """ 136 | Rolls out the policy an arbitrary length by calling it iteratively 137 | 138 | :param batch Dict[str, torch.Tensor]: A data batch coming from the 139 | data loader--should already be 140 | on the correct device 141 | :param rollout_length int: The number of steps to roll out (not including the start) 142 | :param sampler Callable[[torch.Tensor], torch.Tensor]: A function that takes a batch of robot 143 | configurations [B x 7] and returns a batch of 144 | point clouds samples on the surface of that robot 145 | :param unnormalize bool: Whether to return the whole trajectory unnormalized 146 | (i.e. converted back into joint space) 147 | :rtype List[torch.Tensor]: The entire trajectory batch, i.e. a list of 148 | configuration batches including the starting 149 | configurations where each element in the list 150 | corresponds to a timestep. For example, the 151 | first element of each batch in the list would 152 | be a single trajectory. 153 | """ 154 | xyz, q = ( 155 | batch["xyz"], 156 | batch["configuration"], 157 | ) 158 | # This block is to adapt for the case where we only want to roll out a 159 | # single trajectory 160 | if q.ndim == 1: 161 | xyz = xyz.unsqueeze(0) 162 | q = q.unsqueeze(0) 163 | if unnormalize: 164 | q_unnorm = unnormalize_franka_joints(q) 165 | assert isinstance(q_unnorm, torch.Tensor) 166 | trajectory = [q_unnorm] 167 | else: 168 | trajectory = [q] 169 | 170 | for i in range(rollout_length): 171 | q = torch.clamp(q + self(xyz, q), min=-1, max=1) 172 | q_unnorm = unnormalize_franka_joints(q) 173 | assert isinstance(q_unnorm, torch.Tensor) 174 | q_unnorm = q_unnorm.type_as(q) 175 | if unnormalize: 176 | trajectory.append(q_unnorm) 177 | else: 178 | trajectory.append(q) 179 | 180 | samples = sampler(q_unnorm).type_as(xyz) 181 | xyz[:, : samples.shape[1], :3] = samples 182 | 183 | return trajectory 184 | 185 | def training_step( # type: ignore[override] 186 | self, batch: Dict[str, torch.Tensor], batch_idx: int 187 | ) -> torch.Tensor: 188 | """ 189 | A function called automatically by Pytorch Lightning during training. 190 | This function handles the forward pass, the loss calculation, and what to log 191 | 192 | :param batch Dict[str, torch.Tensor]: A data batch coming from the 193 | data loader--should already be 194 | on the correct device 195 | :param batch_idx int: The index of the batch (not used by this function) 196 | :rtype torch.Tensor: The overall weighted loss (used for backprop) 197 | """ 198 | xyz, q = ( 199 | batch["xyz"], 200 | batch["configuration"], 201 | ) 202 | y_hat = torch.clamp(q + self(xyz, q), min=-1, max=1) 203 | ( 204 | cuboid_centers, 205 | cuboid_dims, 206 | cuboid_quats, 207 | cylinder_centers, 208 | cylinder_radii, 209 | cylinder_heights, 210 | cylinder_quats, 211 | supervision, 212 | ) = ( 213 | batch["cuboid_centers"], 214 | batch["cuboid_dims"], 215 | batch["cuboid_quats"], 216 | batch["cylinder_centers"], 217 | batch["cylinder_radii"], 218 | batch["cylinder_heights"], 219 | batch["cylinder_quats"], 220 | batch["supervision"], 221 | ) 222 | collision_loss, point_match_loss = self.loss_fun( 223 | y_hat, 224 | cuboid_centers, 225 | cuboid_dims, 226 | cuboid_quats, 227 | cylinder_centers, 228 | cylinder_radii, 229 | cylinder_heights, 230 | cylinder_quats, 231 | supervision, 232 | ) 233 | self.log("point_match_loss", point_match_loss) 234 | self.log("collision_loss", collision_loss) 235 | val_loss = ( 236 | self.point_match_loss_weight * point_match_loss 237 | + self.collision_loss_weight * collision_loss 238 | ) 239 | self.log("val_loss", val_loss) 240 | return val_loss 241 | 242 | def sample(self, q: torch.Tensor) -> torch.Tensor: 243 | """ 244 | Samples a point cloud from the surface of all the robot's links 245 | 246 | :param q torch.Tensor: Batched configuration in joint space 247 | :rtype torch.Tensor: Batched point cloud of size [B, self.num_robot_points, 3] 248 | """ 249 | assert self.fk_sampler is not None 250 | return self.fk_sampler.sample(q, self.num_robot_points) 251 | 252 | def validation_step( # type: ignore[override] 253 | self, batch: Dict[str, torch.Tensor], batch_idx: int 254 | ) -> torch.Tensor: 255 | """ 256 | This is a Pytorch Lightning function run automatically across devices 257 | during the validation loop 258 | 259 | :param batch Dict[str, torch.Tensor]: The batch coming from the dataloader 260 | :param batch_idx int: The index of the batch (not used by this function) 261 | :rtype torch.Tensor: The loss values which are to be collected into summary stats 262 | """ 263 | 264 | # These are defined here because they need to be set on the correct devices. 265 | # The easiest way to do this is to do it at call-time 266 | if self.fk_sampler is None: 267 | self.fk_sampler = FrankaSampler(self.device, use_cache=True) 268 | if self.collision_sampler is None: 269 | self.collision_sampler = FrankaCollisionSampler( 270 | self.device, with_base_link=False 271 | ) 272 | rollout = self.rollout(batch, 69, self.sample, unnormalize=True) 273 | 274 | assert self.fk_sampler is not None # Necessary for mypy to type properly 275 | eff = self.fk_sampler.end_effector_pose(rollout[-1]) 276 | position_error = torch.linalg.vector_norm( 277 | eff[:, :3, -1] - batch["target_position"], dim=1 278 | ) 279 | avg_target_error = torch.mean(position_error) 280 | 281 | cuboids = TorchCuboids( 282 | batch["cuboid_centers"], 283 | batch["cuboid_dims"], 284 | batch["cuboid_quats"], 285 | ) 286 | cylinders = TorchCylinders( 287 | batch["cylinder_centers"], 288 | batch["cylinder_radii"], 289 | batch["cylinder_heights"], 290 | batch["cylinder_quats"], 291 | ) 292 | 293 | B = batch["cuboid_centers"].size(0) 294 | rollout = torch.stack(rollout, dim=1) 295 | # Here is some Pytorch broadcasting voodoo to calculate whether each 296 | # rollout has a collision or not (looking to calculate the collision rate) 297 | assert rollout.shape == (B, 70, 7) 298 | rollout = rollout.reshape(-1, 7) 299 | has_collision = torch.zeros(B, dtype=torch.bool, device=self.device) 300 | collision_spheres = self.collision_sampler.compute_spheres(rollout) 301 | for radius, spheres in collision_spheres: 302 | num_spheres = spheres.shape[-2] 303 | sphere_sequence = spheres.reshape((B, -1, num_spheres, 3)) 304 | sdf_values = torch.minimum( 305 | cuboids.sdf_sequence(sphere_sequence), 306 | cylinders.sdf_sequence(sphere_sequence), 307 | ) 308 | assert sdf_values.shape == (B, 70, num_spheres) 309 | radius_collisions = torch.any( 310 | sdf_values.reshape((sdf_values.size(0), -1)) <= radius, dim=-1 311 | ) 312 | has_collision = torch.logical_or(radius_collisions, has_collision) 313 | 314 | avg_collision_rate = torch.count_nonzero(has_collision) / B 315 | return { 316 | "avg_target_error": avg_target_error, 317 | "avg_collision_rate": avg_collision_rate, 318 | } 319 | 320 | def validation_step_end( # type: ignore[override] 321 | self, batch_parts: Dict[str, torch.Tensor] 322 | ) -> Dict[str, torch.Tensor]: 323 | """ 324 | Called by Pytorch Lightning at the end of each validation step to 325 | aggregate across devices 326 | 327 | :param batch_parts Dict[str, torch.Tensor]: The parts accumulated from all devices 328 | :rtype Dict[str, torch.Tensor]: The average values across the devices 329 | """ 330 | return { 331 | "avg_target_error": torch.mean(batch_parts["avg_target_error"]), 332 | "avg_collision_rate": torch.mean(batch_parts["avg_collision_rate"]), 333 | } 334 | 335 | def validation_epoch_end( # type: ignore[override] 336 | self, validation_step_outputs: Sequence[Dict[str, torch.Tensor]] 337 | ): 338 | """ 339 | Pytorch lightning method that aggregates stats from the validation loop and logs 340 | 341 | :param validation_step_outputs Sequence[Dict[str, torch.Tensor]]: The outputs from each 342 | validation step 343 | """ 344 | avg_target_error = torch.mean( 345 | torch.stack([x["avg_target_error"] for x in validation_step_outputs]) 346 | ) 347 | self.log("avg_target_error", avg_target_error) 348 | 349 | avg_collision_rate = torch.mean( 350 | torch.stack([x["avg_collision_rate"] for x in validation_step_outputs]) 351 | ) 352 | self.log("avg_collision_rate", avg_collision_rate) 353 | 354 | 355 | class MPiNetsPointNet(pl.LightningModule): 356 | def __init__(self): 357 | super().__init__() 358 | self._build_model() 359 | 360 | def _build_model(self): 361 | """ 362 | Assembles the model design into a ModuleList 363 | """ 364 | self.SA_modules = nn.ModuleList() 365 | self.SA_modules.append( 366 | PointnetSAModule( 367 | npoint=512, 368 | radius=0.05, 369 | nsample=128, 370 | mlp=[1, 64, 64, 64], 371 | bn=False, 372 | ) 373 | ) 374 | self.SA_modules.append( 375 | PointnetSAModule( 376 | npoint=128, 377 | radius=0.3, 378 | nsample=128, 379 | mlp=[64, 128, 128, 256], 380 | bn=False, 381 | ) 382 | ) 383 | self.SA_modules.append(PointnetSAModule(mlp=[256, 512, 512, 1024], bn=False)) 384 | 385 | self.fc_layer = nn.Sequential( 386 | nn.Linear(1024, 4096), 387 | nn.GroupNorm(16, 4096), 388 | nn.LeakyReLU(inplace=True), 389 | nn.Linear(4096, 2048), 390 | nn.GroupNorm(16, 2048), 391 | nn.LeakyReLU(inplace=True), 392 | nn.Linear(2048, 2048), 393 | ) 394 | 395 | @staticmethod 396 | def _break_up_pc(pc: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: 397 | """ 398 | Breaks up the point cloud into the xyz coordinates and segmentation mask 399 | 400 | :param pc torch.Tensor: Tensor with shape [B, N, M] where M is larger than 3. 401 | The first three dimensions along the last axis will be x, y, z 402 | :rtype Tuple[torch.Tensor, torch.Tensor]: Two tensors, one with just xyz 403 | and one with the corresponding features 404 | """ 405 | xyz = pc[..., 0:3].contiguous() 406 | features = pc[..., 3:].transpose(1, 2).contiguous() 407 | return xyz, features 408 | 409 | def forward(self, point_cloud: torch.Tensor) -> torch.Tensor: # type: ignore[override] 410 | """ 411 | Forward pass of the network 412 | 413 | :param point_cloud torch.Tensor: Has dimensions (B, N, 4) 414 | B is the batch size 415 | N is the number of points 416 | 4 is x, y, z, segmentation_mask 417 | This tensor must be on the GPU (CPU tensors not supported) 418 | :rtype torch.Tensor: The output from the network 419 | """ 420 | assert point_cloud.size(2) == 4 421 | 422 | xyz, features = self._break_up_pc(point_cloud) 423 | # print(xyz.shape) 424 | # print(features.shape) 425 | for module in self.SA_modules: 426 | xyz, features = module(xyz, features) 427 | 428 | return self.fc_layer(features.squeeze(-1)) -------------------------------------------------------------------------------- /mpinets/third_party/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/mpinets/third_party/__init__.py -------------------------------------------------------------------------------- /mpinets/third_party/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/mpinets/third_party/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /mpinets/third_party/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/mpinets/third_party/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /mpinets/third_party/__pycache__/sparc.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/mpinets/third_party/__pycache__/sparc.cpython-37.pyc -------------------------------------------------------------------------------- /mpinets/third_party/__pycache__/sparc.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/mpinets/third_party/__pycache__/sparc.cpython-38.pyc -------------------------------------------------------------------------------- /mpinets/third_party/sparc.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2015, Sivakumar Balasubramanian 2 | # 3 | # Permission to use, copy, modify, and/or distribute this software for any 4 | # purpose with or without fee is hereby granted, provided that the above 5 | # copyright notice and this permission notice appear in all copies. 6 | # 7 | # THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES 8 | # WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF 9 | # MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR 10 | # ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES 11 | # WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN 12 | # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF 13 | # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 14 | 15 | # Not a contribution 16 | # Changes made by NVIDIA CORPORATION & AFFILIATES or University of Washington enabling SPARC computation 17 | # or otherwise documented as NVIDIA-proprietary are not a contribution and subject 18 | # to the following terms and conditions: 19 | # 20 | # MIT License 21 | # 22 | # Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved. 23 | # 24 | # Permission is hereby granted, free of charge, to any person obtaining a 25 | # copy of this software and associated documentation files (the "Software"), 26 | # to deal in the Software without restriction, including without limitation 27 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 28 | # and/or sell copies of the Software, and to permit persons to whom the 29 | # Software is furnished to do so, subject to the following conditions: 30 | # 31 | # The above copyright notice and this permission notice shall be included in 32 | # all copies or substantial portions of the Software. 33 | # 34 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 35 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 36 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 37 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 38 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 39 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 40 | # DEALINGS IN THE SOFTWARE. 41 | 42 | 43 | # Note: original code can be found on Github at 44 | # https://github.com/siva82kb/SPARC/blob/master/scripts/smoothness.py 45 | 46 | import numpy as np 47 | 48 | 49 | def sparc(movement, fs, padlevel=4, fc=10.0, amp_th=0.05): 50 | """ 51 | Calculates the smoothness of the given speed profile using the modified 52 | spectral arc length metric. 53 | Parameters 54 | ---------- 55 | movement : np.array 56 | The array containing the movement speed profile. 57 | fs : float 58 | The sampling frequency of the data. 59 | padlevel : integer, optional 60 | Indicates the amount of zero padding to be done to the movement 61 | data for estimating the spectral arc length. [default = 4] 62 | fc : float, optional 63 | The max. cut off frequency for calculating the spectral arc 64 | length metric. [default = 10.] 65 | amp_th : float, optional 66 | The amplitude threshold to used for determing the cut off 67 | frequency upto which the spectral arc length is to be estimated. 68 | [default = 0.05] 69 | Returns 70 | ------- 71 | sal : float 72 | The spectral arc length estimate of the given movement's 73 | smoothness. 74 | (f, Mf) : tuple of two np.arrays 75 | This is the frequency(f) and the magntiude spectrum(Mf) of the 76 | given movement data. This spectral is from 0. to fs/2. 77 | (f_sel, Mf_sel) : tuple of two np.arrays 78 | This is the portion of the spectrum that is selected for 79 | calculating the spectral arc length. 80 | Notes 81 | ----- 82 | This is the modfieid spectral arc length metric, which has been tested only 83 | for discrete movements. 84 | 85 | Examples 86 | -------- 87 | >>> t = np.arange(-1, 1, 0.01) 88 | >>> move = np.exp(-5*pow(t, 2)) 89 | >>> sal, _, _ = sparc(move, fs=100.) 90 | >>> '%.5f' % sal 91 | '-1.41403' 92 | """ 93 | if np.allclose(movement, 0): 94 | print("All movement was 0, returning 0") 95 | return 0, None, None 96 | # Number of zeros to be padded. 97 | nfft = int(pow(2, np.ceil(np.log2(len(movement))) + padlevel)) 98 | 99 | # Frequency 100 | f = np.arange(0, fs, fs / nfft) 101 | # Normalized magnitude spectrum 102 | Mf = abs(np.fft.fft(movement, nfft)) 103 | Mf = Mf / max(Mf) 104 | 105 | # Indices to choose only the spectrum within the given cut off frequency 106 | # Fc. 107 | # NOTE: This is a low pass filtering operation to get rid of high frequency 108 | # noise from affecting the next step (amplitude threshold based cut off for 109 | # arc length calculation). 110 | fc_inx = ((f <= fc) * 1).nonzero() 111 | f_sel = f[fc_inx] 112 | Mf_sel = Mf[fc_inx] 113 | 114 | # Choose the amplitude threshold based cut off frequency. 115 | # Index of the last point on the magnitude spectrum that is greater than 116 | # or equal to the amplitude threshold. 117 | inx = ((Mf_sel >= amp_th) * 1).nonzero()[0] 118 | fc_inx = range(inx[0], inx[-1] + 1) 119 | f_sel = f_sel[fc_inx] 120 | Mf_sel = Mf_sel[fc_inx] 121 | 122 | # Calculate arc length 123 | new_sal = -sum( 124 | np.sqrt( 125 | pow(np.diff(f_sel) / (f_sel[-1] - f_sel[0]), 2) + pow(np.diff(Mf_sel), 2) 126 | ) 127 | ) 128 | return new_sal, (f, Mf), (f_sel, Mf_sel) 129 | -------------------------------------------------------------------------------- /mpinets/types.py: -------------------------------------------------------------------------------- 1 | # MIT License 2 | # 3 | # Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved. 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a 6 | # copy of this software and associated documentation files (the "Software"), 7 | # to deal in the Software without restriction, including without limitation 8 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 9 | # and/or sell copies of the Software, and to permit persons to whom the 10 | # Software is furnished to do so, subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in 13 | # all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 18 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 20 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 21 | # DEALINGS IN THE SOFTWARE. 22 | 23 | from dataclasses import dataclass, field 24 | from typing import List, Union, Optional, Dict, Sequence 25 | import numpy as np 26 | from geometrout.transform import SE3 27 | from geometrout.primitive import Cuboid, Cylinder, Sphere 28 | 29 | 30 | Obstacles = List[Union[Cuboid, Cylinder, Sphere]] 31 | Trajectory = Sequence[Union[Sequence, np.ndarray]] 32 | 33 | 34 | @dataclass 35 | class PlanningProblem: 36 | """ 37 | Defines a common interface to describe planning problems 38 | """ 39 | 40 | target: SE3 # The target in the `right_gripper` frame 41 | target_volume: Union[Cuboid, Cylinder] 42 | q0: np.ndarray # The starting configuration 43 | obstacles: Optional[Obstacles] = None # The obstacles in the scene 44 | obstacle_point_cloud: Optional[np.ndarray] = None 45 | target_negative_volumes: Obstacles = field(default_factory=lambda: []) 46 | 47 | 48 | ProblemSet = Dict[str, Dict[str, List[PlanningProblem]]] 49 | -------------------------------------------------------------------------------- /mpinets/utils.py: -------------------------------------------------------------------------------- 1 | # MIT License 2 | # 3 | # Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES, University of Washington. All rights reserved. 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a 6 | # copy of this software and associated documentation files (the "Software"), 7 | # to deal in the Software without restriction, including without limitation 8 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 9 | # and/or sell copies of the Software, and to permit persons to whom the 10 | # Software is furnished to do so, subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in 13 | # all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 18 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 20 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 21 | # DEALINGS IN THE SOFTWARE. 22 | 23 | from typing import Union, Tuple 24 | 25 | import numpy as np 26 | import torch 27 | from robofin.robots import FrankaRobot, FrankaRealRobot 28 | 29 | 30 | def _normalize_franka_joints_numpy( 31 | batch_trajectory: np.ndarray, 32 | limits: Tuple[float, float] = (-1, 1), 33 | use_real_constraints: bool = True, 34 | ) -> np.ndarray: 35 | """ 36 | Normalizes joint angles to be within a specified range according to the Franka's 37 | joint limits. This is the numpy version 38 | 39 | :param batch_trajectory np.ndarray: A batch of trajectories. Can have dims 40 | [7] if a single configuration 41 | [B, 7] if a batch of configurations 42 | [B, T, 7] if a batched time-series of configurations 43 | :param limits Tuple[float, float]: The new limits to map to 44 | :param use_real_constraints bool: If true, use the empirically determined joint limits 45 | (this is unpublished--just found by monkeying around 46 | with the robot). 47 | If false, use the published joint limits from Franka 48 | :rtype np.ndarray: An array with the same dimensions as the input 49 | """ 50 | robot = FrankaRealRobot if use_real_constraints else FrankaRobot 51 | franka_limits = robot.JOINT_LIMITS 52 | assert ( 53 | (batch_trajectory.ndim == 1 and batch_trajectory.shape[0] == robot.DOF) 54 | or (batch_trajectory.ndim == 2 and batch_trajectory.shape[1] == robot.DOF) 55 | or (batch_trajectory.ndim == 3 and batch_trajectory.shape[2] == robot.DOF) 56 | ) 57 | normalized = (batch_trajectory - franka_limits[:, 0]) / ( 58 | franka_limits[:, 1] - franka_limits[:, 0] 59 | ) * (limits[1] - limits[0]) + limits[0] 60 | return normalized 61 | 62 | 63 | def _normalize_franka_joints_torch( 64 | batch_trajectory: torch.Tensor, 65 | limits: Tuple[float, float] = (-1, 1), 66 | use_real_constraints: bool = True, 67 | ) -> torch.Tensor: 68 | """ 69 | Normalizes joint angles to be within a specified range according to the Franka's 70 | joint limits. This is the torch version 71 | 72 | :param batch_trajectory torch.Tensor: A batch of trajectories. Can have dims 73 | [7] if a single configuration 74 | [B, 7] if a batch of configurations 75 | [B, T, 7] if a batched time-series of configurations 76 | :param limits Tuple[float, float]: The new limits to map to 77 | :param use_real_constraints bool: If true, use the empirically determined joint limits 78 | (this is unpublished--just found by monkeying around 79 | with the robot). 80 | If false, use the published joint limits from Franka 81 | :rtype torch.Tensor: A tensor with the same dimensions as the input 82 | """ 83 | assert isinstance(batch_trajectory, torch.Tensor) 84 | robot = FrankaRealRobot if use_real_constraints else FrankaRobot 85 | franka_limits = torch.as_tensor(robot.JOINT_LIMITS).type_as(batch_trajectory) 86 | assert ( 87 | (batch_trajectory.ndim == 1 and batch_trajectory.size(0) == robot.DOF) 88 | or (batch_trajectory.ndim == 2 and batch_trajectory.size(1) == robot.DOF) 89 | or (batch_trajectory.ndim == 3 and batch_trajectory.size(2) == robot.DOF) 90 | ) 91 | return (batch_trajectory - franka_limits[:, 0]) / ( 92 | franka_limits[:, 1] - franka_limits[:, 0] 93 | ) * (limits[1] - limits[0]) + limits[0] 94 | 95 | 96 | def normalize_franka_joints( 97 | batch_trajectory: Union[np.ndarray, torch.Tensor], 98 | limits: Tuple[float, float] = (-1, 1), 99 | use_real_constraints: bool = True, 100 | ) -> Union[np.ndarray, torch.Tensor]: 101 | """ 102 | Normalizes joint angles to be within a specified range according to the Franka's 103 | joint limits. This is semantic sugar that dispatches to the correct implementation. 104 | 105 | :param batch_trajectory Union[np.ndarray, torch.Tensor]: A batch of trajectories. Can have dims 106 | [7] if a single configuration 107 | [B, 7] if a batch of configurations 108 | [B, T, 7] if a batched time-series of configurations 109 | :param limits Tuple[float, float]: The new limits to map to 110 | :param use_real_constraints bool: If true, use the empirically determined joint limits 111 | (this is unpublished--just found by monkeying around 112 | with the robot). 113 | If false, use the published joint limits from Franka 114 | :rtype Union[np.ndarray, torch.Tensor]: A tensor or numpy array with the same dimensions 115 | and type as the input 116 | :raises NotImplementedError: Raises an error if another data type (e.g. a list) is passed in 117 | """ 118 | if isinstance(batch_trajectory, torch.Tensor): 119 | return _normalize_franka_joints_torch( 120 | batch_trajectory, limits=limits, use_real_constraints=True 121 | ) 122 | elif isinstance(batch_trajectory, np.ndarray): 123 | return _normalize_franka_joints_numpy( 124 | batch_trajectory, limits=limits, use_real_constraints=True 125 | ) 126 | else: 127 | raise NotImplementedError("Only torch.Tensor and np.ndarray implemented") 128 | 129 | 130 | def _unnormalize_franka_joints_numpy( 131 | batch_trajectory: np.ndarray, 132 | limits: Tuple[float, float] = (-1, 1), 133 | use_real_constraints: bool = True, 134 | ) -> np.ndarray: 135 | """ 136 | Unnormalizes joint angles from a specified range back into the Franka's joint limits. 137 | This is the numpy version and the inverse of `_normalize_franka_joints_numpy`. 138 | 139 | :param batch_trajectory np.ndarray: A batch of trajectories. Can have dims 140 | [7] if a single configuration 141 | [B, 7] if a batch of configurations 142 | [B, T, 7] if a batched time-series of configurations 143 | :param limits Tuple[float, float]: The current limits to map to the joint limits 144 | :param use_real_constraints bool: If true, use the empirically determined joint limits 145 | (this is unpublished--just found by monkeying around 146 | with the robot). 147 | If false, use the published joint limits from Franka 148 | :rtype np.ndarray: An array with the same dimensions as the input 149 | """ 150 | robot = FrankaRealRobot if use_real_constraints else FrankaRobot 151 | franka_limits = robot.JOINT_LIMITS 152 | assert ( 153 | (batch_trajectory.ndim == 1 and batch_trajectory.shape[0] == robot.DOF) 154 | or (batch_trajectory.ndim == 2 and batch_trajectory.shape[1] == robot.DOF) 155 | or (batch_trajectory.ndim == 3 and batch_trajectory.shape[2] == robot.DOF) 156 | ) 157 | assert np.all(batch_trajectory >= limits[0]) 158 | assert np.all(batch_trajectory <= limits[1]) 159 | franka_limit_range = franka_limits[:, 1] - franka_limits[:, 0] 160 | franka_lower_limit = franka_limits[:, 0] 161 | for _ in range(batch_trajectory.ndim - 1): 162 | franka_limit_range = franka_limit_range[np.newaxis, ...] 163 | franka_lower_limit = franka_lower_limit[np.newaxis, ...] 164 | unnormalized = (batch_trajectory - limits[0]) * franka_limit_range / ( 165 | limits[1] - limits[0] 166 | ) + franka_lower_limit 167 | 168 | return unnormalized 169 | 170 | 171 | def _unnormalize_franka_joints_torch( 172 | batch_trajectory: torch.Tensor, 173 | limits: Tuple[float, float] = (-1, 1), 174 | use_real_constraints: bool = True, 175 | ) -> torch.Tensor: 176 | """ 177 | Unnormalizes joint angles from a specified range back into the Franka's joint limits. 178 | This is the torch version and the inverse of `_normalize_franka_joints_torch`. 179 | 180 | :param batch_trajectory torch.Tensor: A batch of trajectories. Can have dims 181 | [7] if a single configuration 182 | [B, 7] if a batch of configurations 183 | [B, T, 7] if a batched time-series of configurations 184 | :param limits Tuple[float, float]: The current limits to map to the joint limits 185 | :param use_real_constraints bool: If true, use the empirically determined joint limits 186 | (this is unpublished--just found by monkeying around 187 | with the robot). 188 | If false, use the published joint limits from Franka 189 | :rtype torch.Tensor: A tensor with the same dimensions as the input 190 | """ 191 | assert isinstance(batch_trajectory, torch.Tensor) 192 | robot = FrankaRealRobot if use_real_constraints else FrankaRobot 193 | franka_limits = torch.as_tensor(robot.JOINT_LIMITS).type_as(batch_trajectory) 194 | dof = franka_limits.size(0) 195 | assert ( 196 | (batch_trajectory.ndim == 1 and batch_trajectory.size(0) == dof) 197 | or (batch_trajectory.ndim == 2 and batch_trajectory.size(1) == dof) 198 | or (batch_trajectory.ndim == 3 and batch_trajectory.size(2) == dof) 199 | ) 200 | assert torch.all(batch_trajectory >= limits[0]) 201 | assert torch.all(batch_trajectory <= limits[1]) 202 | franka_limit_range = franka_limits[:, 1] - franka_limits[:, 0] 203 | franka_lower_limit = franka_limits[:, 0] 204 | for _ in range(batch_trajectory.ndim - 1): 205 | franka_limit_range = franka_limit_range.unsqueeze(0) 206 | franka_lower_limit = franka_lower_limit.unsqueeze(0) 207 | return (batch_trajectory - limits[0]) * franka_limit_range / ( 208 | limits[1] - limits[0] 209 | ) + franka_lower_limit 210 | 211 | 212 | def unnormalize_franka_joints( 213 | batch_trajectory: Union[np.ndarray, torch.Tensor], 214 | limits: Tuple[float, float] = (-1, 1), 215 | use_real_constraints: bool = True, 216 | ) -> Union[np.ndarray, torch.Tensor]: 217 | """ 218 | Unnormalizes joint angles from a specified range back into the Franka's joint limits. 219 | This is semantic sugar that dispatches to the correct implementation, the inverse of 220 | `normalize_franka_joints`. 221 | 222 | :param batch_trajectory Union[np.ndarray, torch.Tensor]: A batch of trajectories. Can have dims 223 | [7] if a single configuration 224 | [B, 7] if a batch of configurations 225 | [B, T, 7] if a batched time-series of configurations 226 | :param limits Tuple[float, float]: The current limits to map to the joint limits 227 | :param use_real_constraints bool: If true, use the empirically determined joint limits 228 | (this is unpublished--just found by monkeying around 229 | with the robot). 230 | If false, use the published joint limits from Franka 231 | :rtype Union[np.ndarray, torch.Tensor]: A tensor or numpy array with the same dimensions 232 | and type as the input 233 | :raises NotImplementedError: Raises an error if another data type (e.g. a list) is passed in 234 | """ 235 | if isinstance(batch_trajectory, torch.Tensor): 236 | return _unnormalize_franka_joints_torch( 237 | batch_trajectory, limits=limits, use_real_constraints=use_real_constraints 238 | ) 239 | elif isinstance(batch_trajectory, np.ndarray): 240 | return _unnormalize_franka_joints_numpy( 241 | batch_trajectory, limits=limits, use_real_constraints=use_real_constraints 242 | ) 243 | else: 244 | raise NotImplementedError("Only torch.Tensor and np.ndarray implemented") 245 | -------------------------------------------------------------------------------- /results/teaser_compressed.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/results/teaser_compressed.gif -------------------------------------------------------------------------------- /results/teaser_compressed.mp4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vishal-2000/EDMP/558482f97d075bc02eef5452760b3a3a17860573/results/teaser_compressed.mp4 -------------------------------------------------------------------------------- /urdfs/panda.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 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | 139 | 140 | 141 | 142 | 143 | 144 | 145 | 146 | 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | 160 | 161 | 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | 211 | 212 | 213 | 214 | 215 | 216 | 217 | 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 | 234 | 235 | 236 | 237 | 238 | 239 | 240 | 241 | 242 | 243 | 244 | 245 | 246 | 247 | 248 | 249 | 250 | 251 | 252 | 253 | 254 | 255 | 256 | 257 | 258 | 259 | 260 | 261 | 262 | 263 | 264 | 265 | 266 | 267 | 268 | 269 | 270 | 271 | 272 | 273 | 274 | 275 | 276 | 277 | 278 | 279 | 280 | 281 | 282 | 283 | 284 | 285 | 286 | 287 | 288 | 289 | 290 | 291 | 292 | 293 | 294 | 295 | 296 | 297 | 298 | 299 | 300 | 301 | 302 | 303 | 304 | 305 | 306 | 307 | 308 | 309 | 310 | 311 | 312 | 313 | 314 | 315 | 316 | 317 | 318 | 319 | 320 | 321 | 322 | 323 | 324 | 325 | 326 | 327 | 328 | 329 | 330 | 331 | 332 | 333 | 334 | 335 | 336 | 337 | 338 | 339 | -------------------------------------------------------------------------------- /urdfs/workspace.urdf: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 25 | 26 | 27 | 28 | --------------------------------------------------------------------------------