From 0a8bc6163f4048bd7e6fce5ead010b84a300795a Mon Sep 17 00:00:00 2001 From: Silverster98 <1539168414@qq.com> Date: Wed, 19 Apr 2023 17:32:49 +0800 Subject: [PATCH 1/7] Update code for grasp generation and arm motion planning --- .gitignore | 9 +- assets/.gitkeep | 0 configs/model/unet_fk2.yaml | 23 + configs/model/unet_grasp.yaml | 3 +- configs/optimizer/grasp_with_object.yaml | 13 +- configs/planner/greedy_fk2_planning.yaml | 9 + configs/task/franka_planning.yaml | 55 +++ configs/task/grasp_gen_ur.yaml | 26 +- datasets/__init__.py | 1 + datasets/fk2plan.py | 206 ++++++++ datasets/multidex_shadowhand_ur.py | 105 ++-- datasets/transforms.py | 2 +- envs/assets/.gitkeep | 0 envs/tasks/base_task.py | 210 ++++++++ envs/tasks/franka_panda.py | 460 ++++++++++++++++++ envs/tasks/franka_panda.yaml | 49 ++ models/__init__.py | 1 + models/environment.py | 242 +++++++++ models/model/pointnet2/pointnet2_modules.py | 513 ++++++++++++++++++++ models/model/pointnet2/pointnet2_semseg.py | 321 ++++++++++++ models/model/pointnet2/pointnet2_utils.py | 427 ++++++++++++++++ models/model/pointnet2/pytorch_utils.py | 297 ++++++++++++ models/model/scene_model.py | 5 + models/model/unet.py | 8 +- models/optimizer/grasp_with_object.py | 125 ++++- models/planner/fk2planning.py | 69 +++ models/visualizer.py | 87 ++++ plan.py | 2 + scripts/franka_planning/plan.sh | 12 + scripts/franka_planning/train.sh | 11 + scripts/franka_planning/train_ddm.sh | 12 + scripts/grasp_gen/sample.sh | 8 - scripts/grasp_gen/train.sh | 7 - scripts/grasp_gen/train.slurm | 29 -- scripts/grasp_gen_ur/sample.sh | 35 ++ scripts/grasp_gen_ur/test.sh | 0 scripts/grasp_gen_ur/train.sh | 11 + utils/misc.py | 4 + utils/rot6d.py | 15 +- 39 files changed, 3300 insertions(+), 112 deletions(-) create mode 100644 assets/.gitkeep create mode 100644 configs/model/unet_fk2.yaml create mode 100644 configs/planner/greedy_fk2_planning.yaml create mode 100644 configs/task/franka_planning.yaml create mode 100644 datasets/fk2plan.py create mode 100644 envs/assets/.gitkeep create mode 100644 envs/tasks/base_task.py create mode 100644 envs/tasks/franka_panda.py create mode 100644 envs/tasks/franka_panda.yaml create mode 100644 models/model/pointnet2/pointnet2_modules.py create mode 100644 models/model/pointnet2/pointnet2_semseg.py create mode 100644 models/model/pointnet2/pointnet2_utils.py create mode 100644 models/model/pointnet2/pytorch_utils.py create mode 100644 models/planner/fk2planning.py create mode 100644 scripts/franka_planning/plan.sh create mode 100644 scripts/franka_planning/train.sh create mode 100644 scripts/franka_planning/train_ddm.sh delete mode 100755 scripts/grasp_gen/sample.sh delete mode 100755 scripts/grasp_gen/train.sh delete mode 100755 scripts/grasp_gen/train.slurm create mode 100755 scripts/grasp_gen_ur/sample.sh create mode 100755 scripts/grasp_gen_ur/test.sh create mode 100755 scripts/grasp_gen_ur/train.sh diff --git a/.gitignore b/.gitignore index eb86812..e358e11 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,9 @@ __pycache__ -outputs/ \ No newline at end of file +outputs/ + +assets/franka_description/ +assets/object/ +assets/urdf/ + +envs/assets/franka_description/ +envs/assets/scene_description/ \ No newline at end of file diff --git a/assets/.gitkeep b/assets/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/configs/model/unet_fk2.yaml b/configs/model/unet_fk2.yaml new file mode 100644 index 0000000..bff2783 --- /dev/null +++ b/configs/model/unet_fk2.yaml @@ -0,0 +1,23 @@ +name: UNetModel + +d_x: 0 # placeholder +d_model: 512 +time_embed_mult: 2 +nblocks: 4 +resblock_dropout: 0.0 +transformer_num_heads: 8 +transformer_dim_head: 64 +transformer_dropout: 0.1 +transformer_depth: 1 +transformer_mult_ff: 2 +context_dim: 512 +use_position_embedding: false # for input x + +scene_model: + name: PointTransformer + use_color: ${task.dataset.use_color} + use_normal: ${task.dataset.use_normal} + num_points: ${task.dataset.num_points} + pretrained_weights: null + pretrained_weights_slurm: null +freeze_scene_model: false \ No newline at end of file diff --git a/configs/model/unet_grasp.yaml b/configs/model/unet_grasp.yaml index f5cd5b8..ff9b17c 100644 --- a/configs/model/unet_grasp.yaml +++ b/configs/model/unet_grasp.yaml @@ -14,12 +14,11 @@ context_dim: 512 use_position_embedding: false # for input x scene_model: - name: PointNet + name: PointNet2 use_color: ${task.dataset.use_color} use_normal: ${task.dataset.use_normal} num_points: ${task.dataset.num_points} num_tokens: 8 -# pretrained_weights: /home/wangzan/Outputs/point_transformer.scannet/outputs/2022-04-13_18-29-56_POINTTRANS_C_32768/model.pth pretrained_weights: null pretrained_weights_slurm: null freeze_scene_model: false \ No newline at end of file diff --git a/configs/optimizer/grasp_with_object.yaml b/configs/optimizer/grasp_with_object.yaml index 5afd914..2f37785 100644 --- a/configs/optimizer/grasp_with_object.yaml +++ b/configs/optimizer/grasp_with_object.yaml @@ -2,16 +2,17 @@ name: GraspWithObject scale: 1.0 -palm_alignment: true -palm_alignment_weight: 1.0 +collision: true +collision_weight: 1.0 clip_grad_by_value: min: -0.1 max: 0.1 +asset_dir: ${task.dataset.asset_dir} +asset_dir_slurm: ${task.dataset.asset_dir_slurm} + batch_size: ${task.test.batch_size} modeling_keys: ${task.dataset.modeling_keys} - - - - +normalize_x: ${task.dataset.normalize_x} +normalize_x_trans: ${task.dataset.normalize_x_trans} \ No newline at end of file diff --git a/configs/planner/greedy_fk2_planning.yaml b/configs/planner/greedy_fk2_planning.yaml new file mode 100644 index 0000000..748de0b --- /dev/null +++ b/configs/planner/greedy_fk2_planning.yaml @@ -0,0 +1,9 @@ +name: GreedyFK2Planner + +scale: 1.0 + +greedy_type: 'all_frame_exp' # ['last_frame', 'all_frame', 'all_frame_exp'] + +clip_grad_by_value: + min: -0.1 + max: 0.1 \ No newline at end of file diff --git a/configs/task/franka_planning.yaml b/configs/task/franka_planning.yaml new file mode 100644 index 0000000..a806474 --- /dev/null +++ b/configs/task/franka_planning.yaml @@ -0,0 +1,55 @@ +# task: franka motion generation +name: franka_planning +lr: 1e-4 +eval_interval: 1 +eval_visualize: 1 + +train: + batch_size: 32 + num_workers: 4 + num_epochs: 5000 + log_step: 100 + +test: + epoch: null + batch_size: 16 + num_workers: 0 + +has_observation: true # for franka planing, use observation in default +dataset: + planner_batch_size: ${task.test.batch_size} + name: FK2Plan + normalize_x: true + desc: '[FK2Plan] -- dataset used for forward-kinematics franka planning generation condition on 3D scene' + horizon: 64 + modeling_keys: ['allDoFs'] + num_points: 4096 + frame_interval: 16 + sample_trajs_per_scene: 2 + sample_frame_interval: 32 + device: cuda + use_color: false + use_normal: false + is_downsample: true + data_dir: /home/wangzan/Data/SceneDiffuser/grasp_and_armmotion/datasets/FK2PlanDataset + data_dir_slurm: /home/wangzan/Data/SceneDiffuser/grasp_and_armmotion/datasets/FK2PlanDataset + + train_transforms: ['NumpyToTensorPath'] + test_transforms: ['NumpyToTensorPath'] + transform_cfg: {} + +visualizer: + name: FK2PlanVisualizer + visualize: false + ksample: 2 + horizon: ${task.dataset.horizon} + cnt: 40 + +## for planning task +env: + name: FK2PlanningEnvWrapper + max_sample_each_step: 1 # currently only support max_sample_each_step=1 + max_trajectory_length: 300 + arrive_threshold: 0.2 + sims_per_step: 3 + eval_case_num_per_scene: 64 \ No newline at end of file diff --git a/configs/task/grasp_gen_ur.yaml b/configs/task/grasp_gen_ur.yaml index 839d17c..abad9b3 100644 --- a/configs/task/grasp_gen_ur.yaml +++ b/configs/task/grasp_gen_ur.yaml @@ -1,22 +1,25 @@ # task: pose generation -name: grasp_gen +name: grasp_gen_ur lr: 1e-4 eval_interval: 1 eval_visualize: 1 train: - batch_size: 64 + batch_size: 128 # num_workers: 4 num_workers: 0 - num_epochs: 1000 + num_epochs: 3000 log_step: 100 test: - batch_size: 2 + epoch: null + batch_size: 32 num_workers: 0 dataset: name: MultiDexShadowHandUR + normalize_x: true + normalize_x_trans: true desc: '[MultiDex ShadowHand] -- dataset used for grasp pose generation condition on 3D object' modeling_keys: ['allDoFs'] num_points: 2048 @@ -26,17 +29,18 @@ dataset: use_color: false use_normal: false is_downsample: true -# data_dir: /home/wangzan/Data/LEMO/PROX_temporal/PROX_temporal/PROXD_temp - data_dir: /home/puhao/data/MultiDex_UR/shadowhand/ - data_dir_slurm: /home/lipuhao/data/MultiDex_UR/shadowhand/ - scene_path: /home/puhao/data/MultiDex_UR/object_pcds.pkl - scene_path_slurm: /home/lipuhao/data/MultiDex_UR/object_pcds.pkl + asset_dir: /home/wangzan/Data/SceneDiffuser/grasp_and_armmotion/datasets/MultiDex_UR/ + asset_dir_slurm: /home/wangzan/Data/SceneDiffuser/grasp_and_armmotion/datasets/MultiDex_UR/ train_transforms: ['NumpyToTensor'] test_transforms: ['NumpyToTensor'] transform_cfg: {} visualizer: - name: GraspGenVisualizer + name: GraspGenURVisualizer + ksample: 64 # sample k case in each case + vis_denoising: false # visualize denoising process + save_mesh: false + ## visualization config used in training visualize: false - ksample: 4 + interval: 1 diff --git a/datasets/__init__.py b/datasets/__init__.py index 11f4e51..79f2c91 100644 --- a/datasets/__init__.py +++ b/datasets/__init__.py @@ -2,3 +2,4 @@ from .multidex_shadowhand_ur import MultiDexShadowHandUR from .lemo_motion import LEMOMotion from .scannet_path import ScanNetPath +from .fk2plan import FK2Plan diff --git a/datasets/fk2plan.py b/datasets/fk2plan.py new file mode 100644 index 0000000..7a60e9b --- /dev/null +++ b/datasets/fk2plan.py @@ -0,0 +1,206 @@ +import time +from typing import Any, Tuple, Dict +import os +import json +import pickle +import torch +import numpy as np +from torch.utils.data import Dataset, DataLoader +from omegaconf import DictConfig, OmegaConf + +from datasets.misc import collate_fn_general +from datasets.transforms import make_default_transform +from datasets.base import DATASET + +@DATASET.register() +class FK2Plan(Dataset): + """ Dataset for fk franka planning, training with FK2PLAN Dataset + """ + + _scene_pre_code = 'dthvl15jruz9i2fok6bsy3qamp8c4nex' + _train_split = [f"dthvl15jruz9i2fok6bsy3qamp8c4nex{str(i).zfill(3)}" for i in range(160)] + # _test_split = [f"dthvl15jruz9i2fok6bsy3qamp8c4nex{str(i).zfill(3)}" for i in range(160, 200)] + _test_split = [f"dthvl15jruz9i2fok6bsy3qamp8c4nex{str(i).zfill(3)}" for i in range(160, 200)] ##todo: here to test + _all_split = [f"dthvl15jruz9i2fok6bsy3qamp8c4nex{str(i).zfill(3)}" for i in range(200)] + # _test_split = [f"dthvl15jruz9i2fok6bsy3qamp8c4nex{str(i).zfill(3)}" for i in range(200)] + + _joint_angle_lower = np.array([-2.9671, -1.8326, -2.9671, -3.1416, -2.9671, -0.0873, -2.9671]) + _joint_angle_upper = np.array([2.9671, 1.8326, 2.9671, 0.0873, 2.9671, 3.8223, 2.9671]) + + _NORMALIZE_LOWER = -1. + _NORMALIZE_UPPER = 1. + + def __init__(self, cfg: DictConfig, phase: str, slurm: bool, case_only: bool=False, **kwargs: Dict) -> None: + super(FK2Plan, self).__init__() + self.phase = phase + self.slurm = slurm + if self.phase == 'train': + self.split = self._train_split + elif self.phase == 'test': + self.split = self._test_split + elif self.phase == 'all': + self.split = self._all_split + else: + raise Exception('Unsupported phase.') + self.device = cfg.device + self.is_downsample = cfg.is_downsample + self.modeling_keys = cfg.modeling_keys + self.num_points = cfg.num_points + self.use_color = cfg.use_color + self.use_normal = cfg.use_normal + self.normalize_x = cfg.normalize_x + self.horizon = cfg.horizon + self.frame_interval = cfg.frame_interval # interval sampling + self.sample_trajs_per_scene = cfg.sample_trajs_per_scene + self.sample_frame_interval = cfg.sample_frame_interval + self.planner_batch_size = cfg.planner_batch_size + self.obj_dim = int(3 + 3 * self.use_color + 3 * self.use_normal) + self.transform = make_default_transform(cfg, phase) + + ## resource folders + self.data_dir = cfg.data_dir_slurm if self.slurm else cfg.data_dir + with open(os.path.join(self.data_dir, 'desc.json'), 'r') as f: + self.dataset_desc = json.load(f) + + ## load data + self._pre_load_data(case_only) + + def _pre_load_data(self, case_only: bool) -> None: + """ Load dataset + Args: + case_only: only load single case for testing, if ture, the dataset will be smaller. + This is useful in after-training visual evaluation. + """ + self.trajectories = [] + self.indices = [] + self.scene_pcds_nors = {} + + fk2plan_dataset = pickle.load(open(os.path.join(self.data_dir, 'fk2plan_dataset.pkl'), 'rb')) + self.scene_pcds_nors = pickle.load(open(os.path.join(self.data_dir, 'scene_pcds_nors.pkl'), 'rb')) + # todo: length the scene points cloud + num_points = np.max([a.shape[0] for a in list(self.scene_pcds_nors.values())]) + for scene_id in self.scene_pcds_nors.keys(): + if self.scene_pcds_nors[scene_id].shape[0] < num_points: + self.scene_pcds_nors[scene_id] = np.concatenate([self.scene_pcds_nors[scene_id], + self.scene_pcds_nors[scene_id][0:num_points-self.scene_pcds_nors[scene_id].shape[0]]], axis=0) + self.dataset_info = fk2plan_dataset['info'] + self.dataset_desc = json.load(open(os.path.join(self.data_dir, 'desc.json'), 'r')) + + ## load paths + for mdata in fk2plan_dataset['metadata']: + mdata_scene_id = mdata[0] + mdata_start_goal_pose = mdata[1] + mdata_tra_qpos = mdata[2] + if self.normalize_x: + mdata_tra_qpos = self.angle_normalize(mdata_tra_qpos) + if mdata_scene_id in self.split: + self.trajectories.append((mdata_scene_id, mdata_start_goal_pose, mdata_tra_qpos)) + + ## segment path to fixed horizon for training + if case_only: + loaded_counter = {s: 0 for s in self.split} + for i, traj in enumerate(self.trajectories): + mdata_scene_id, mdata_start_goal_pose, mdata_tra_qpos = traj + max_start = mdata_tra_qpos.shape[0] - self.horizon + if max_start <= 0: + continue + + if case_only: + loaded_counter[mdata_scene_id] += 1 + if loaded_counter[mdata_scene_id] > self.planner_batch_size: + continue + # self.indices.append((i, 0, self.horizon + 1)) + # continue + + if case_only: + self.indices.append((i, 0, self.horizon + 1)) + else: + for start in range(0, max_start, self.frame_interval): + end = start + self.horizon + 1 + self.indices.append((i, start, end)) + print('Finishing Pre-load in FK2Plan') + + def angle_normalize(self, joint_angle: np.ndarray): + joint_angle_norm = np.divide((joint_angle - self._joint_angle_lower), + (self._joint_angle_upper - self._joint_angle_lower)) + joint_angle_norm = joint_angle_norm * (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) - (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) / 2 + if type(joint_angle_norm) == torch.Tensor: + return joint_angle_norm.to(torch.float32) + else: + return joint_angle_norm.astype(np.float32) + + def angle_denormalize(self, joint_angle: np.ndarray): + joint_angle_denorm = joint_angle + (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) / 2 + joint_angle_denorm /= (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) + joint_angle_denorm = joint_angle_denorm * (self._joint_angle_upper - self._joint_angle_lower) + self._joint_angle_lower + if type(joint_angle_denorm) == torch.Tensor: + return joint_angle_denorm.to(torch.float32) + else: + return joint_angle_denorm.astype(np.float32) + + def __len__(self): + return len(self.indices) + + def __getitem__(self, index: Any) -> Tuple: + traj_idx, start, end = self.indices[index] + scene_id, start_goal_pose, tra_qpos = self.trajectories[traj_idx] + + ## load trajectory segment + traj_seg = tra_qpos[start:end] + + ## load data + scene_pc = self.scene_pcds_nors[scene_id] + + ## randomly resample points + if self.phase != 'train': + np.random.seed(0) # resample point cloud with a fixed random seed + resample_indices = np.random.permutation(len(scene_pc)) + scene_pc = scene_pc[resample_indices[:self.num_points]] + + ## format point cloud xyz and feature + xyz = scene_pc[:, 0:3] + feat = scene_pc[:, 3:3] # empty array with shape + + if self.use_color: + raise NotImplementedError + + if self.use_normal: + normal = scene_pc[:, 3:6] + feat = np.concatenate([feat, normal], axis=-1) + + data = { + 'x': traj_seg, + 'start': traj_seg[0:1, :], + 'target': tra_qpos[-1], + 'start_goal_pose': start_goal_pose, + 'pos': xyz, + 'feat': feat, + 'scene_id': scene_id, + } + if self.transform is not None: + data = self.transform(data, modeling_keys=self.modeling_keys) + + return data + + def get_dataloader(self, **kwargs): + return DataLoader(self, **kwargs) + + +if __name__ == '__main__': + config_path = "../configs/task/franka_planning.yaml" + cfg = OmegaConf.load(config_path) + dataloader = FK2Plan(cfg.dataset, 'train', False).get_dataloader(batch_size=128, + collate_fn=collate_fn_general, + num_workers=4, + pin_memory=True, + shuffle=True,) + + device = 'cuda' + st = time.time() + print(len(dataloader.dataset)) + for it, d in enumerate(dataloader): + for key in d: + if torch.is_tensor(d[key]): + d[key] = d[key].to(device) + print(f'{time.time() - st}') + st = time.time() diff --git a/datasets/multidex_shadowhand_ur.py b/datasets/multidex_shadowhand_ur.py index 9b068e7..c3b5cfb 100644 --- a/datasets/multidex_shadowhand_ur.py +++ b/datasets/multidex_shadowhand_ur.py @@ -1,18 +1,12 @@ from typing import Any, Tuple, Dict import os -import json -import glob -from tqdm import tqdm import pickle -import trimesh import torch import numpy as np from torch.utils.data import Dataset, DataLoader from omegaconf import DictConfig, OmegaConf from datasets.misc import collate_fn_squeeze_pcd_batch_grasp -from utils.smplx_utils import convert_smplx_verts_transfomation_matrix_to_body -from utils.smplx_utils import SMPLXWrapper from datasets.transforms import make_default_transform from datasets.base import DATASET @@ -22,18 +16,18 @@ class MultiDexShadowHandUR(Dataset): """ _train_split = ["contactdb+alarm_clock", "contactdb+banana", "contactdb+binoculars", - "contactdb+cell_phone", "contactdb+cube_large", "contactdb+cube_medium", - "contactdb+cube_small", "contactdb+cylinder_large", "contactdb+cylinder_small", - "contactdb+elephant", "contactdb+flashlight", "contactdb+hammer", - "contactdb+light_bulb", "contactdb+mouse", "contactdb+piggy_bank", "contactdb+ps_controller", - "contactdb+pyramid_large", "contactdb+pyramid_medium", "contactdb+pyramid_small", - "contactdb+stanford_bunny", "contactdb+stapler", "contactdb+toothpaste", "contactdb+torus_large", - "contactdb+torus_medium", "contactdb+torus_small", "contactdb+train", - "ycb+bleach_cleanser", "ycb+cracker_box", "ycb+foam_brick", "ycb+gelatin_box", "ycb+hammer", - "ycb+lemon", "ycb+master_chef_can", "ycb+mini_soccer_ball", "ycb+mustard_bottle", "ycb+orange", - "ycb+peach", "ycb+pitcher_base", "ycb+plum", "ycb+power_drill", "ycb+pudding_box", - "ycb+rubiks_cube", "ycb+sponge", "ycb+strawberry", "ycb+sugar_box", "ycb+toy_airplane", - "ycb+tuna_fish_can", "ycb+wood_block"] + "contactdb+cell_phone", "contactdb+cube_large", "contactdb+cube_medium", + "contactdb+cube_small", "contactdb+cylinder_large", "contactdb+cylinder_small", + "contactdb+elephant", "contactdb+flashlight", "contactdb+hammer", + "contactdb+light_bulb", "contactdb+mouse", "contactdb+piggy_bank", "contactdb+ps_controller", + "contactdb+pyramid_large", "contactdb+pyramid_medium", "contactdb+pyramid_small", + "contactdb+stanford_bunny", "contactdb+stapler", "contactdb+toothpaste", "contactdb+torus_large", + "contactdb+torus_medium", "contactdb+torus_small", "contactdb+train", + "ycb+bleach_cleanser", "ycb+cracker_box", "ycb+foam_brick", "ycb+gelatin_box", "ycb+hammer", + "ycb+lemon", "ycb+master_chef_can", "ycb+mini_soccer_ball", "ycb+mustard_bottle", "ycb+orange", + "ycb+peach", "ycb+pitcher_base", "ycb+plum", "ycb+power_drill", "ycb+pudding_box", + "ycb+rubiks_cube", "ycb+sponge", "ycb+strawberry", "ycb+sugar_box", "ycb+toy_airplane", + "ycb+tuna_fish_can", "ycb+wood_block"] _test_split = ["contactdb+apple", "contactdb+camera", "contactdb+cylinder_medium", "contactdb+rubber_duck", "contactdb+door_knob", "contactdb+water_bottle", "ycb+baseball", "ycb+pear", "ycb+potted_meat_can", "ycb+tomato_soup_can"] @@ -53,6 +47,20 @@ class MultiDexShadowHandUR(Dataset): "contactdb+door_knob", "contactdb+water_bottle", "ycb+baseball", "ycb+pear", "ycb+potted_meat_can", "ycb+tomato_soup_can"] + _joint_angle_lower = torch.tensor([-0.5235988, -0.7853982, -0.43633232, 0., 0., 0., -0.43633232, 0., 0., 0., + -0.43633232, 0., 0., 0., 0., -0.43633232, 0., 0., 0., -1.047, 0., -0.2618, + -0.5237, 0.]) + _joint_angle_upper = torch.tensor([0.17453292, 0.61086524, 0.43633232, 1.5707964, 1.5707964, 1.5707964, 0.43633232, + 1.5707964, 1.5707964, 1.5707964, 0.43633232, 1.5707964, 1.5707964, 1.5707964, + 0.6981317, 0.43633232, 1.5707964, 1.5707964, 1.5707964, 1.047, 1.309, 0.2618, + 0.5237, 1.]) + + _global_trans_lower = torch.tensor([-0.13128923, -0.10665303, -0.45753425]) + _global_trans_upper = torch.tensor([0.12772022, 0.22954416, -0.21764427]) + + _NORMALIZE_LOWER = -1. + _NORMALIZE_UPPER = 1. + def __init__(self, cfg: DictConfig, phase: str, slurm: bool, case_only: bool=False, **kwargs: Dict) -> None: super(MultiDexShadowHandUR, self).__init__() self.phase = phase @@ -71,12 +79,19 @@ def __init__(self, cfg: DictConfig, phase: str, slurm: bool, case_only: bool=Fal self.num_points = cfg.num_points self.use_color = cfg.use_color self.use_normal = cfg.use_normal + self.normalize_x = cfg.normalize_x + self.normalize_x_trans = cfg.normalize_x_trans self.obj_dim = int(3 + 3 * self.use_color + 3 * self.use_normal) self.transform = make_default_transform(cfg, phase) ## resource folders - self.data_dir = cfg.data_dir_slurm if self.slurm else cfg.data_dir - self.scene_path = cfg.scene_path_slurm if self.slurm else cfg.scene_path + self.asset_dir = cfg.asset_dir_slrum if self.slurm else cfg.asset_dir + self.data_dir = os.path.join(self.asset_dir, 'shadowhand') + self.scene_path = os.path.join(self.asset_dir, 'object_pcds.pkl') + self._joint_angle_lower = self._joint_angle_lower.cpu() + self._joint_angle_upper = self._joint_angle_upper.cpu() + self._global_trans_lower = self._global_trans_lower.cpu() + self._global_trans_upper = self._global_trans_upper.cpu() ## load data self._pre_load_data(case_only) @@ -101,13 +116,43 @@ def _pre_load_data(self, case_only: bool) -> None: # for obj in self.scene_pcds.keys(): # self.scene_pcds[obj] = torch.tensor(self.scene_pcds[obj], device=self.device) for mdata in grasp_dataset['metadata']: + mdata_qpos = mdata[0].cpu() + joint_angle = mdata_qpos.clone().detach()[9:] + global_trans = mdata_qpos.clone().detach()[:3] + if self.normalize_x: + joint_angle = self.angle_normalize(joint_angle) + if self.normalize_x_trans: + global_trans = self.trans_normalize(global_trans) + mdata_qpos = torch.cat([global_trans, joint_angle], dim=0).requires_grad_(True) if mdata[2] in self.split: self.frames.append({'robot_name': 'shadowhand', 'object_name': mdata[2], 'object_rot_mat': mdata[1].clone().detach().numpy(), - 'qpos': mdata[0].clone().detach().requires_grad_(True)}) + 'qpos': mdata_qpos}) # print('Finishing Pre-load in MultiDexShadowHand') - + + def trans_normalize(self, global_trans: torch.Tensor): + global_trans_norm = torch.div((global_trans - self._global_trans_lower), (self._global_trans_upper - self._global_trans_lower)) + global_trans_norm = global_trans_norm * (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) - (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) / 2 + return global_trans_norm + + def trans_denormalize(self, global_trans: torch.Tensor): + global_trans_denorm = global_trans + (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) / 2 + global_trans_denorm /= (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) + global_trans_denorm = global_trans_denorm * (self._global_trans_upper - self._global_trans_lower) + self._global_trans_lower + return global_trans_denorm + + def angle_normalize(self, joint_angle: torch.Tensor): + joint_angle_norm = torch.div((joint_angle - self._joint_angle_lower), (self._joint_angle_upper - self._joint_angle_lower)) + joint_angle_norm = joint_angle_norm * (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) - (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) / 2 + return joint_angle_norm + + def angle_denormalize(self, joint_angle: torch.Tensor): + joint_angle_denorm = joint_angle + (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) / 2 + joint_angle_denorm /= (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) + joint_angle_denorm = joint_angle_denorm * (self._joint_angle_upper - self._joint_angle_lower) + self._joint_angle_lower + return joint_angle_denorm + def __len__(self): return len(self.frames) @@ -125,8 +170,10 @@ def __getitem__(self, index: Any) -> Tuple: ## randomly resample points if self.phase != 'train': np.random.seed(0) # resample point cloud with a fixed random seed - np.random.shuffle(scene_pc) - scene_pc = scene_pc[:self.num_points] + # np.random.shuffle(scene_pc) + # scene_pc = scene_pc[:self.num_points] + resample_indices = np.random.permutation(len(scene_pc)) + scene_pc = scene_pc[resample_indices[:self.num_points]] ## format point cloud xyz and feature xyz = scene_pc[:, 0:3] @@ -165,14 +212,14 @@ def get_dataloader(self, **kwargs): config_path = "../configs/task/grasp_gen.yaml" cfg = OmegaConf.load(config_path) dataloader = MultiDexShadowHandUR(cfg.dataset, 'train', False).get_dataloader(batch_size=4, - collate_fn=collate_fn_squeeze_pcd_batch_grasp, - num_workers=0, - pin_memory=True, - shuffle=True,) + collate_fn=collate_fn_squeeze_pcd_batch_grasp, + num_workers=0, + pin_memory=True, + shuffle=True,) device = 'cuda' for it, data in enumerate(dataloader): for key in data: if torch.is_tensor(data[key]): data[key] = data[key].to(device) - print() + print() \ No newline at end of file diff --git a/datasets/transforms.py b/datasets/transforms.py index a1a1f87..28e908f 100644 --- a/datasets/transforms.py +++ b/datasets/transforms.py @@ -177,7 +177,7 @@ def __init__(self, **kwargs) -> None: def __call__(self, data: Dict, *args: List, **kwargs: Dict) -> Dict: for key in data.keys(): if key in ['x', 'start', 'target', 'pos', 'feat', 's_grid_map', 's_grid_dim', 's_grid_min', 's_grid_max'] and not torch.is_tensor(data[key]): - data[key] = torch.tensor(np.array(data[key])) + data[key] = torch.tensor(np.array(data[key]), dtype=torch.float32) return data diff --git a/envs/assets/.gitkeep b/envs/assets/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/envs/tasks/base_task.py b/envs/tasks/base_task.py new file mode 100644 index 0000000..531fe8c --- /dev/null +++ b/envs/tasks/base_task.py @@ -0,0 +1,210 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +import sys +import os +import operator +from copy import deepcopy +import random + +from isaacgym import gymapi +from isaacgym.gymutil import get_property_setter_map, get_property_getter_map, get_default_setter_args, apply_random_samples, check_buckets, generate_random_samples +from isaacgym import gymtorch +import numpy as np +import torch + + +# Base class for RL tasks +class BaseTask(): + def __init__(self, cfg, enable_camera_sensors=False, + cam_pos=(10, 10, 3.0), cam_target=(0, 0, 0.0)): + self.gym = gymapi.acquire_gym() + + self.device_type = cfg.get("self.device_type", "cuda") + self.device_id = cfg.get("self.device_id", 0) + + self.device = "cpu" + if self.device_type == "cuda" or self.device_type == "GPU": + self.device = "cuda" + ":" + str(self.device_id) + + self.headless = cfg["headless"] + + # double check! + self.graphics_device_id = self.device_id + if enable_camera_sensors == False and self.headless == True: + self.graphics_device_id = -1 + + self.num_envs = cfg["env"]["numEnvs"] + + self.control_freq_inv = cfg["env"].get("controlFrequencyInv", 1) + + # optimization flags for pytorch JIT + torch._C._jit_set_profiling_mode(False) + torch._C._jit_set_profiling_executor(False) + + self.original_props = {} + self.dr_randomizations = {} + self.first_randomization = True + self.actor_params_generator = None + self.extern_actor_params = {} + for env_id in range(self.num_envs): + self.extern_actor_params[env_id] = None + + self.last_step = -1 + self.last_rand_step = -1 + + # create envs, sim and viewer + + self.create_sim() + self.gym.prepare_sim(self.sim) + + # todo: read from config + self.enable_viewer_sync = True + self.viewer = None + self.need_update = False + + # if running with a viewer, set up keyboard shortcuts and camera + if self.headless == False: + # subscribe to keyboard shortcuts + self.viewer = self.gym.create_viewer( + self.sim, gymapi.CameraProperties()) + self.gym.subscribe_viewer_keyboard_event( + self.viewer, gymapi.KEY_ESCAPE, "QUIT") + self.gym.subscribe_viewer_keyboard_event( + self.viewer, gymapi.KEY_V, "toggle_viewer_sync") + + # set the camera position based on up axis + sim_params = self.gym.get_sim_params(self.sim) + if sim_params.up_axis == gymapi.UP_AXIS_Z: + cam_pos = gymapi.Vec3(cam_pos[0], cam_pos[1], cam_pos[2]) + cam_target = gymapi.Vec3(cam_target[0], cam_target[1], cam_target[2]) + else: + cam_pos = gymapi.Vec3(cam_pos[0], cam_pos[1], cam_pos[2]) + cam_target = gymapi.Vec3(cam_target[0], cam_target[1], cam_target[2]) + + self.gym.viewer_camera_look_at( + self.viewer, None, cam_pos, cam_target) + + + + # set gravity based on up axis and return axis index + def set_sim_params_up_axis(self, sim_params, axis): + if axis == 'z': + sim_params.up_axis = gymapi.UP_AXIS_Z + sim_params.gravity.x = 0 + sim_params.gravity.y = 0 + sim_params.gravity.z = -9.81 + return 2 + return 1 + + def create_sim(self, compute_device, graphics_device, physics_engine, sim_params): + sim = self.gym.create_sim(compute_device, graphics_device, physics_engine, sim_params) + if sim is None: + print("*** Failed to create sim") + quit() + + return sim + + def step(self, actions): + if self.dr_randomizations.get('actions', None): + actions = self.dr_randomizations['actions']['noise_lambda'](actions) + + # apply actions + self.pre_physics_step(actions) + + # step physics and render each frame + for _ in range(self.control_freq_inv): + self.render() + self.gym.simulate(self.sim) + + # to fix! + if self.device == 'cpu': + self.gym.fetch_results(self.sim, True) + + # compute observations, rewards, resets, ... + self.post_physics_step() + + if self.dr_randomizations.get('observations', None): + self.obs_buf = self.dr_randomizations['observations']['noise_lambda'](self.obs_buf) + + + + # def control_ik(self, dpose): + # j_eef, self.num_envs + # # solve damped least squares + # j_eef_T = torch.transpose(j_eef, 1, 2) + # lmbda = torch.eye(6, self.device=self.device) * (self.damping ** 2) + # u = (j_eef_T @ torch.inverse(j_eef @ j_eef_T + lmbda) @ dpose).view(self.num_envs, 7) + # return u + + def render(self, sync_frame_time=False): + if self.viewer: + # check for window closed + if self.gym.query_viewer_has_closed(self.viewer): + sys.exit() + + # check for keyboard events + for evt in self.gym.query_viewer_action_events(self.viewer): + if evt.action == "QUIT" and evt.value > 0: + sys.exit() + elif evt.action == "toggle_viewer_sync" and evt.value > 0: + self.enable_viewer_sync = not self.enable_viewer_sync + + # fetch results + if self.device != 'cpu': + self.gym.fetch_results(self.sim, True) + + # step graphics + if self.enable_viewer_sync: + self.gym.step_graphics(self.sim) + self.gym.draw_viewer(self.viewer, self.sim, True) + else: + self.gym.poll_viewer_events(self.viewer) + + def get_actor_params_info(self, dr_params, env): + """Returns a flat array of actor params, their names and ranges.""" + if "actor_params" not in dr_params: + return None + params = [] + names = [] + lows = [] + highs = [] + param_getters_map = get_property_getter_map(self.gym) + for actor, actor_properties in dr_params["actor_params"].items(): + handle = self.gym.find_actor_handle(env, actor) + for prop_name, prop_attrs in actor_properties.items(): + if prop_name == 'color': + continue # this is set randomly + props = param_getters_map[prop_name](env, handle) + if not isinstance(props, list): + props = [props] + for prop_idx, prop in enumerate(props): + for attr, attr_randomization_params in prop_attrs.items(): + name = prop_name+'_'+str(prop_idx)+'_'+attr + lo_hi = attr_randomization_params['range'] + distr = attr_randomization_params['distribution'] + if 'uniform' not in distr: + lo_hi = (-1.0*float('Inf'), float('Inf')) + if isinstance(prop, np.ndarray): + for attr_idx in range(prop[attr].shape[0]): + params.append(prop[attr][attr_idx]) + names.append(name+'_'+str(attr_idx)) + lows.append(lo_hi[0]) + highs.append(lo_hi[1]) + else: + params.append(getattr(prop, attr)) + names.append(name) + lows.append(lo_hi[0]) + highs.append(lo_hi[1]) + return params, names, lows, highs + + def pre_physics_step(self, actions): + raise NotImplementedError + + def post_physics_step(self): + raise NotImplementedError + diff --git a/envs/tasks/franka_panda.py b/envs/tasks/franka_panda.py new file mode 100644 index 0000000..58e08df --- /dev/null +++ b/envs/tasks/franka_panda.py @@ -0,0 +1,460 @@ +import os +from requests import get +from isaacgym import gymutil, gymtorch, gymapi +from isaacgym.torch_utils import * +import yaml +from envs.tasks.base_task import BaseTask +from tqdm import tqdm +import numpy as np +import torch +import trimesh as tm + + +class FrankaMotion_Player(BaseTask): + + def __init__(self, config, sim_params, physics_engine, + device_type, device_id, headless, scene_id, + init_qpos: torch.tensor): + self.gym = None + self.viewer = None + self.num_sim = 0 + self.cfg = config + self.scene_id = scene_id + self.sim_params = sim_params + self.device = device_type + self.init_qpos = init_qpos.clone() + self.physics_engine = physics_engine + self.up_axis = 'z' + self.cfg["device_type"] = device_type + self.cfg["device_id"] = device_id + self.cfg["headless"] = headless + self.device_type = device_type + self.device_id = device_id + self.headless = headless + + self.device = "cuda" + if self.device_type == "cuda" or self.device_type == "GPU": + self.device = "cuda" + ":" + str(self.device_id) + + config["env"]["numTrain"] = init_qpos.shape[0] + config["env"]["numEnvs"] = init_qpos.shape[0] + self.env_num_train = config["env"]["numTrain"] + self.env_num = self.env_num_train + self.asset_root = config["env"]["asset"]["assetRoot"] + self.num_train = config["env"]["numTrain"] + self.tot_num = 1 + self.exp_name = config['env']["env_name"] + + print("Simulator: number of objects", self.tot_num) + print("Simulator: number of environments", self.env_num) + if self.num_train: + assert (self.env_num_train % self.num_train == 0) + + # the number of used length must less than real length + # each object should have equal number envs + assert (self.env_num % self.tot_num == 0) + self.env_per_object = self.env_num // self.tot_num + + self.env_ptr_list = [] + self.obj_loaded = False + self.dexterous_loaded = False + + super().__init__(cfg=self.cfg, enable_camera_sensors=config["env"]["enableCameraSensors"], + cam_pos=(-1.2, -1.2, 1.2), cam_target=(0., 0., 0.5)) + + # acquire tensors + self.root_tensor = gymtorch.wrap_tensor( + self.gym.acquire_actor_root_state_tensor(self.sim)).to(self.device) + self.dof_state_tensor = gymtorch.wrap_tensor( + self.gym.acquire_dof_state_tensor(self.sim)).to(self.device) + # self.rigid_body_tensor = gymtorch.wrap_tensor( + # self.gym.acquire_rigid_body_state_tensor(self.sim)).to(self.device).reshape(self.num_envs, -1, 13) + + self.gym.refresh_actor_root_state_tensor(self.sim) + self.gym.refresh_dof_state_tensor(self.sim) + self.gym.refresh_rigid_body_state_tensor(self.sim) + + self.root_tensor = self.root_tensor.view(self.num_envs, -1, 13) + self.dof_state_tensor = self.dof_state_tensor.view( + self.num_envs, -1, 2) + self.initial_dof_states = self.dof_state_tensor.clone() + self.initial_root_states = self.root_tensor.clone() + # self.initial_obj_state = self.initial_root_states[:, 1, :3] + + self.dexterous_dof_tensor = self.dof_state_tensor[:, :self.dexterous_num_dofs, :] + + self.dexterous_root_tensor = self.root_tensor[:, 0, :] + # self.object_root_tensor = self.root_tensor[:, 1, :] + + self.dof_dim = self.dexterous_num_dofs + self.pos_act = self.initial_dof_states[:, :self.dexterous_num_dofs, 0].clone() + self.eff_act = torch.zeros( + (self.num_envs, self.dof_dim), device=self.device) + self.damping = 0.05 + self.prepare_ik() + + def __del__(self): + if self.gym is not None: + self.gym.destroy_sim(self.sim) + if self.viewer is not None: + self.gym.destroy_viewer(self.viewer) + del self.gym + + def step_sim_q(self, sim_q): + sim_q = sim_q.to(self.device) + self.gym.set_dof_position_target_tensor(self.sim, gymtorch.unwrap_tensor(sim_q)) + self.gym.set_dof_actuation_force_tensor(self.sim, gymtorch.unwrap_tensor(self.effort_action)) + + self.gym.simulate(self.sim) + self.gym.fetch_results(self.sim, True) + + # refresh tensors + self.gym.refresh_rigid_body_state_tensor(self.sim) + self.gym.refresh_dof_state_tensor(self.sim) + self.gym.refresh_jacobian_tensors(self.sim) + self.gym.refresh_mass_matrix_tensors(self.sim) + + # update viewer + self.gym.step_graphics(self.sim) + self.gym.draw_viewer(self.viewer, self.sim, True) + self.gym.sync_frame_time(self.sim) + + def create_sim(self): + self.dt = self.sim_params.dt + self.up_axis_idx = self.set_sim_params_up_axis( + self.sim_params, self.up_axis) + + self.sim = super().create_sim(self.device_id, self.graphics_device_id, + self.physics_engine, self.sim_params) + self._create_ground_plane() + self._place_agents( + self.cfg["env"]["numTrain"], self.cfg["env"]["envSpacing"]) + + def _load_agent(self, env_ptr, env_id): + + if self.dexterous_loaded == False: + self.dexterous_actor_list = [] + asset_root = self.asset_root + dexterous_asset_file = "franka_description/robots/panda.urdf" + # dexterous_asset_file = "franka_description/robots/movable_dexterous_hand_fine_collision.urdf" + asset_options = gymapi.AssetOptions() + asset_options.density = self.cfg['agent']['density'] + asset_options.fix_base_link = True + asset_options.disable_gravity = True + # Switch Meshes from Z-up left-handed system to Y-up Right-handed coordinate system. + asset_options.flip_visual_attachments = True + asset_options.armature = 0.01 + asset_options.use_mesh_materials = True + asset_options.mesh_normal_mode = gymapi.COMPUTE_PER_VERTEX + asset_options.override_com = True # recompute center of mesh + asset_options.override_inertia = True # recompute inertia + asset_options.vhacd_enabled = True + asset_options.vhacd_params = gymapi.VhacdParams() + # asset_options.vhacd_params.resolution = 3000000 + asset_options.vhacd_params.resolution = 3000000 + + self.dexterous_asset = self.gym.load_asset( + self.sim, asset_root, dexterous_asset_file, asset_options) + self.dexterous_loaded = True + + dexterous_dof_max_torque, self.dexterous_dof_lower_limits, self.dexterous_dof_upper_limits = self._get_dof_property( + self.dexterous_asset) + + dof_props = self.gym.get_asset_dof_properties(self.dexterous_asset) + if self.cfg["env"]["driveMode"] in ["pos", "ik"]: + dof_props["driveMode"][:].fill(gymapi.DOF_MODE_POS) + dof_props["stiffness"][:].fill(self.cfg['agent']['dof_props']['stiffness']) + dof_props["velocity"][:].fill(self.cfg['agent']['dof_props']['velocity']) + dof_props["damping"][:].fill(self.cfg['agent']['dof_props']['damping']) + else: # osc + dof_props["driveMode"][:].fill(gymapi.DOF_MODE_EFFORT) + dof_props["stiffness"][:].fill(0.0) + dof_props["damping"][:].fill(0.0) + # root pose + initial_dexterous_pose = gymapi.Transform() + initial_dexterous_pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0) + initial_dexterous_pose.p = gymapi.Vec3(0.0, 0.0, 0.0) + + # set start dof + self.dexterous_num_dofs = self.gym.get_asset_dof_count(self.dexterous_asset) + + # default_dof_pos = np.zeros(self.dexterous_num_dofs, dtype=np.float32) + default_dof_pos = self.init_qpos[env_id, :].cpu() + + # initialize for the pose and rotation + # default_dof_pos = self.compute_pos_action(self.translation, self.rpy, self.new_joint_angle, env_id) + + dexterous_dof_state = np.zeros_like( + dexterous_dof_max_torque, gymapi.DofState.dtype) + dexterous_dof_state["pos"] = default_dof_pos + + dexterous_actor = self.gym.create_actor( + env_ptr, + self.dexterous_asset, + initial_dexterous_pose, + "dexterous", + env_id, + 0, + 0) + for i_body in range(len(self.gym.get_actor_rigid_body_names(env_ptr, dexterous_actor))): + color_vec = np.random.uniform(0., 1., 3) + self.gym.set_rigid_body_color(env_ptr, 0, i_body, gymapi.MESH_VISUAL, + gymapi.Vec3(color_vec[0], color_vec[1], color_vec[2])) + dexterous_shape_props = self.gym.get_actor_rigid_shape_properties(env_ptr, dexterous_actor) + for shape in dexterous_shape_props: + shape.friction = self.cfg['agent']['shape']['friction'] + + self.gym.set_actor_rigid_shape_properties(env_ptr, dexterous_actor, dexterous_shape_props) + self.gym.set_actor_dof_properties(env_ptr, dexterous_actor, dof_props) + self.gym.set_actor_dof_states( + env_ptr, dexterous_actor, dexterous_dof_state, gymapi.STATE_ALL) + self.dexterous_actor_list.append(dexterous_actor) + self.dexterous_link_dict = self.gym.get_asset_rigid_body_dict(self.dexterous_asset) + + def prepare_ik(self): + + # get dof state tensor + self.dof_pos = self.dexterous_dof_tensor[:, :, 0].view(self.num_envs, self.dexterous_num_dofs, 1) + self.dof_vel = self.dexterous_dof_tensor[:, :, 1].view(self.num_envs, self.dexterous_num_dofs, 1) + + # Set action tensors + self.pos_action = torch.zeros_like(self.dof_pos).squeeze(-1) + self.effort_action = torch.zeros_like(self.pos_action) + + def step_qpos(self, qpos): + qpos = qpos.clone().to(self.device) + self.gym.simulate(self.sim) + self.gym.fetch_results(self.sim, True) + + self.pos_action = qpos.clone() + self.gym.set_dof_position_target_tensor(self.sim, gymtorch.unwrap_tensor(self.pos_action)) + self.gym.set_dof_actuation_force_tensor(self.sim, gymtorch.unwrap_tensor(self.effort_action)) + + # refresh tensors + self.gym.refresh_rigid_body_state_tensor(self.sim) + self.gym.refresh_dof_state_tensor(self.sim) + self.gym.refresh_jacobian_tensors(self.sim) + self.gym.refresh_mass_matrix_tensors(self.sim) + + self.num_sim += 1 + # update viewer + self.gym.step_graphics(self.sim) + self.gym.draw_viewer(self.viewer, self.sim, True) + self.gym.sync_frame_time(self.sim) + + def get_cur_qpos(self): + sim_q_state = self.dof_state_tensor[:, :, 0].clone() + cur_qpos = torch.tensor(sim_q_state, device=self.device) + return cur_qpos + + def _get_dof_property(self, asset): + dof_props = self.gym.get_asset_dof_properties(asset) + dof_num = self.gym.get_asset_dof_count(asset) + dof_lower_limits = [] + dof_upper_limits = [] + dof_max_torque = [] + for i in range(dof_num): + dof_max_torque.append(dof_props['effort'][i]) + dof_lower_limits.append(dof_props['lower'][i]) + dof_upper_limits.append(dof_props['upper'][i]) + dof_max_torque = np.array(dof_max_torque) + dof_lower_limits = np.array(dof_lower_limits) + dof_upper_limits = np.array(dof_upper_limits) + return dof_max_torque, dof_lower_limits, dof_upper_limits + + def _load_obj_asset(self): + + self.obj_name_list = [] + self.obj_asset_list = [] + self.table_asset_list = [] + self.obj_pose_list = [] + self.table_pose_list = [] + self.obj_actor_list = [] + self.table_actor_list = [] + + used_len = 1 + with tqdm(total=used_len) as pbar: + pbar.set_description('Loading assets:') + cur = 0 + + obj_asset_list = [] + # prepare the assets to be used + if self.scene_id is None: + self.num_scenes = self.cfg['env']['numScenes'] + scene_hash = 'dthvl15jruz9i2fok6bsy3qamp8c4nex' + for scene_idx in range(self.num_scenes): + scene_urdf_path = f"{scene_hash}{str(scene_idx).zfill(3)}/fine_urdf/fine_scene.urdf" + print(f'load object asset into IsaacGym | Scene ID: {scene_urdf_path}') + object_asset_options = gymapi.AssetOptions() + object_asset_options.density = self.cfg['object']['density'] + # update + object_asset_options.linear_damping = self.cfg['object']['damping']['linear'] + object_asset_options.angular_damping = self.cfg['object']['damping']['angular'] + + object_asset_options.fix_base_link = True + object_asset_options.disable_gravity = True + object_asset_options.armature = 0.01 + object_asset_options.use_mesh_materials = True + object_asset_options.mesh_normal_mode = gymapi.COMPUTE_PER_VERTEX + object_asset_options.override_com = True # recompute center of mesh + object_asset_options.override_inertia = True # recompute inertia + object_asset_options.vhacd_enabled = True + object_asset_options.vhacd_params = gymapi.VhacdParams() + object_asset_options.vhacd_params.resolution = 3000000 + + obj_asset = self.gym.load_asset( + self.sim, './envs/assets/scene_description', scene_urdf_path, object_asset_options) + self.obj_asset_list.append(obj_asset) + obj_start_pose = gymapi.Transform() + obj_start_pose.p = gymapi.Vec3(0.0, 0.0, 0.0) + obj_start_pose.r = gymapi.Quat(0., 0., 0., 1.) + self.obj_pose_list.append(obj_start_pose) + else: + print(f'load object asset into IsaacGym: {self.scene_id}') + # scene_urdf_path = f"{self.scene_id}/SceneDescription.urdf" + scene_urdf_path = f"{self.scene_id}/fine_urdf/fine_scene.urdf" + object_asset_options = gymapi.AssetOptions() + object_asset_options.density = self.cfg['object']['density'] + # update + object_asset_options.linear_damping = self.cfg['object']['damping']['linear'] + object_asset_options.angular_damping = self.cfg['object']['damping']['angular'] + + object_asset_options.fix_base_link = True + object_asset_options.disable_gravity = True + object_asset_options.armature = 0.01 + object_asset_options.use_mesh_materials = True + object_asset_options.mesh_normal_mode = gymapi.COMPUTE_PER_VERTEX + object_asset_options.override_com = True # recompute center of mesh + object_asset_options.override_inertia = True # recompute inertia + object_asset_options.vhacd_enabled = True + object_asset_options.vhacd_params = gymapi.VhacdParams() + object_asset_options.vhacd_params.resolution = 3000000 + + obj_asset = self.gym.load_asset( + self.sim, './envs/assets/scene_description', scene_urdf_path, object_asset_options) + self.obj_asset_list.append(obj_asset) + obj_start_pose = gymapi.Transform() + obj_start_pose.p = gymapi.Vec3(0.0, 0.0, 0.0) + obj_start_pose.r = gymapi.Quat(0., 0., 0., 1.) + self.obj_pose_list.append(obj_start_pose) + + def _load_obj(self, env_ptr, env_id): + + if self.obj_loaded == False: + self._load_obj_asset() + self.obj_loaded = True + if self.scene_id is None: + obj_type = env_id % self.num_scenes + subenv_id = env_id // self.num_scenes + else: + obj_type = env_id // self.env_per_object + subenv_id = env_id % self.env_per_object + obj_actor = self.gym.create_actor( + env_ptr, + self.obj_asset_list[obj_type], + self.obj_pose_list[obj_type], + "scene{}-{}".format(obj_type, subenv_id), + env_id, + 0, + 0) + # color of scene + + obj_shape_props = self.gym.get_actor_rigid_shape_properties(env_ptr, obj_actor) + self.gym.set_rigid_body_color(env_ptr, 1, 0, gymapi.MESH_VISUAL, gymapi.Vec3(0.75, 0.75, 0.75)) + for shape in obj_shape_props: + shape.friction = self.cfg['object']['shape']['friction'] + self.gym.set_actor_rigid_shape_properties(env_ptr, obj_actor, obj_shape_props) + + self.obj_actor_list.append(obj_actor) + assert(self.gym.get_actor_rigid_body_names(self.env_ptr_list[0], 1)[0] == 'table') + + def _place_agents(self, env_num, spacing): + + print("Simulator: creating agents") + + lower = gymapi.Vec3(-spacing, -spacing, 0.0) + upper = gymapi.Vec3(spacing, spacing, spacing) + self.space_middle = torch.zeros((env_num, 3), device=self.device) + self.space_range = torch.zeros((env_num, 3), device=self.device) + self.space_middle[:, 0] = self.space_middle[:, 1] = 0 + self.space_middle[:, 2] = spacing / 2 + self.space_range[:, 0] = self.space_range[:, 1] = spacing + self.space_middle[:, 2] = spacing / 2 + num_per_row = int(np.sqrt(env_num)) + + with tqdm(total=env_num) as pbar: + pbar.set_description('Enumerating envs:') + for env_id in range(env_num): + env_ptr = self.gym.create_env( + self.sim, lower, upper, num_per_row) + self.env_ptr_list.append(env_ptr) + self._load_agent(env_ptr, env_id) + self._load_obj(env_ptr, env_id) + pbar.update(1) + + def _create_ground_plane(self): + plane_params = gymapi.PlaneParams() + plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0) + plane_params.distance = 0.1 + plane_params.static_friction = 0.1 + plane_params.dynamic_friction = 0.1 + self.gym.add_ground(self.sim, plane_params) + + def _draw_line(self, src, dst): + line_vec = np.stack([src, dst]).flatten().astype(np.float32) + color = np.array([1, 0, 0], dtype=np.float32) + self.gym.clear_lines(self.viewer) + self.gym.add_lines( + self.viewer, + self.env_ptr_list[0], + self.env_num, + line_vec, + color + ) + + +def get_sim_param(): + # initialize sim + sim_params = gymapi.SimParams() + sim_params.dt = 1. / 60. + sim_params.num_client_threads = 0 + sim_params.physx.solver_type = 1 + sim_params.physx.num_position_iterations = 4 + sim_params.physx.num_velocity_iterations = 0 + sim_params.physx.num_threads = 4 + sim_params.physx.use_gpu = True + sim_params.physx.num_subscenes = 0 + sim_params.physx.max_gpu_contact_pairs = 8 * 1024 * 1024 + sim_params.use_gpu_pipeline = True + sim_params.physx.use_gpu = True + sim_params.physx.num_threads = 0 + return sim_params + + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--scene_idx', default=-1, type=int) + parser.add_argument('--num_envs', default=1024, type=int) + args = parser.parse_args() + + SCENE_ID = f'dthvl15jruz9i2fok6bsy3qamp8c4nex{str(args.scene_idx).zfill(3)}' + num_envs = args.num_envs + init_qpos = torch.tensor([[0., -0.785, 0, -2.356, 0., 1.571, 0.785]], device='cuda').repeat(num_envs, 1) + sim_params = get_sim_param() + path = "envs/tasks/franka_panda.yaml" + with open(path) as f: + config = yaml.safe_load(f) + scene_config = None + if args.scene_idx == -1: + env = FrankaMotion_Player(config=config, sim_params=sim_params, physics_engine=gymapi.SIM_PHYSX, + device_type="cuda", device_id=0, headless=False, + scene_id=None, init_qpos=init_qpos) + else: + env = FrankaMotion_Player(config=config, sim_params=sim_params, physics_engine=gymapi.SIM_PHYSX, + device_type="cuda", device_id=0, headless=False, + scene_id=f"dthvl15jruz9i2fok6bsy3qamp8c4nex000", init_qpos=init_qpos) + while not env.gym.query_viewer_has_closed(env.viewer): + env.step_qpos(init_qpos) + cur_qpos = env.get_cur_qpos() + print(f'cur pose: {cur_qpos.mean(0)}') \ No newline at end of file diff --git a/envs/tasks/franka_panda.yaml b/envs/tasks/franka_panda.yaml new file mode 100644 index 0000000..8cecbe6 --- /dev/null +++ b/envs/tasks/franka_panda.yaml @@ -0,0 +1,49 @@ +task: + useStage: False + useSlider: False + useGroundTruth: False + +env: + numScenes: 200 + asset: + assetRoot: "envs/assets" + env_name: "franka" + envSpacing: 1.5 + numEnvs: 6 + driveMode: "pos" + enableCameraSensors: False + +sim: + substeps: 2 + physx: + num_threads: 4 + solver_type: 1 # 0: pgs, 1: tgs + num_position_iterations: 8 + num_velocity_iterations: 0 + contact_offset: 0.002 + rest_offset: 0.0 + bounce_threshold_velocity: 0.2 + max_depenetration_velocity: 1000.0 + default_buffer_size_multiplier: 5.0 + flex: + num_outer_iterations: 5 + num_inner_iterations: 20 + warm_start: 0.8 + relaxation: 0.75 + +object: + density: 1000 + damping: + linear: 1 + angular: 1 + shape: + friction: 1. + +agent: + density: 1000 + dof_props: + stiffness: 100.0 + velocity: 1. + damping: 100.0 + shape: + friction: 1. diff --git a/models/__init__.py b/models/__init__.py index aa7a27e..13c138f 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -5,3 +5,4 @@ from .optimizer.motion_in_scene import MotionInSceneOptimizer from .optimizer.path_in_scene import PathInSceneOptimizer from .planner.pathplanning import GreedyPathPlanner +from .planner.fk2planning import GreedyFK2Planner diff --git a/models/environment.py b/models/environment.py index 135510f..526398e 100644 --- a/models/environment.py +++ b/models/environment.py @@ -9,11 +9,17 @@ from tqdm import tqdm from omegaconf import DictConfig from collections import defaultdict +import json +import pickle +import copy +import yaml from models.optimizer.utils import transform_verts from utils.registry import Registry from utils.misc import random_str from utils.visualize import create_trimesh_node, create_trimesh_nodes_path, render_scannet_path +from envs.tasks.franka_panda import FrankaMotion_Player +from isaacgym import gymapi, gymtorch, gymutil ENV = Registry('Env') @@ -304,6 +310,242 @@ def get_trajectory_length(self) -> torch.Tensor: """ return self.trajectory_length +class FK2PlanningEnvCore(): + def __init__(self, data: Dict, scene_id: str, sims_per_step: int = 10, + arrive_threshold: float = 0.3, max_trajectory_length: int = 500) -> None: + """ A Franka Forward kinematics planning environment + + Args: + data: dataloader provided data dict + scene_id: the scene id of data + max_trajectory_length: max length of the trajectory + """ + self.scene_id = scene_id + self.device = data['start'].device + self.batch = data['start'].shape[0] + self.max_trajectory_length = max_trajectory_length + self.arrive_threshold = arrive_threshold + self.sims_per_step = sims_per_step + # + # self.angle_normalize = angle_normalize + # self.angle_denormalize = angle_denormalize + + self.state = data['start'].squeeze(1) + self.target = data['target'] + + self.end = torch.zeros(self.batch, dtype=bool, device=self.device) + self.trajectory = [[self.state[i].clone().cpu()] for i in range(self.batch)] + self.franka_isaac_env = None + self.initial_qpos = self.state + self.isaac_config = yaml.safe_load(open("envs/tasks/franka_panda.yaml")) + self.trajectory_length = torch.ones(self.batch, dtype=torch.int32, device=self.device) * max_trajectory_length + self.init_isaac_env(initial_qpos=self.initial_qpos) + + def get_sim_param(self): + # initialize sim + sim_params = gymapi.SimParams() + sim_params.dt = 1. / 60. + sim_params.num_client_threads = 0 + sim_params.physx.solver_type = 1 + sim_params.physx.num_position_iterations = 4 + sim_params.physx.num_velocity_iterations = 0 + sim_params.physx.num_threads = 4 + sim_params.physx.use_gpu = True + sim_params.physx.num_subscenes = 0 + sim_params.physx.max_gpu_contact_pairs = 8 * 1024 * 1024 + sim_params.use_gpu_pipeline = True + sim_params.physx.use_gpu = True + sim_params.physx.num_threads = 0 + return sim_params + + def __del__(self): + del self.franka_isaac_env + + def init_isaac_env(self, initial_qpos: torch.Tensor): + self.initial_qpos = initial_qpos.detach().clone() + sim_params = self.get_sim_param() + self.franka_isaac_env = FrankaMotion_Player(config=self.isaac_config, sim_params=sim_params, + physics_engine=gymapi.SIM_PHYSX, + device_type=self.device, device_id=0, headless=False, + scene_id=self.scene_id, init_qpos=initial_qpos) + + def step(self, step: int, target_qpos: torch.Tensor): + ## 1. step target qpos + next_state = self.step_target_qpos(target_qpos) + ## 2. compute next state + self.state[~self.end, ...] = next_state[~self.end, ...] + for i in range(self.batch): + self.trajectory[i].append(self.target[i].clone().cpu() if self.end[i] else self.state[i].clone().cpu()) + ## 3. check arriving the target position + dist = torch.norm(self.state - self.target, dim=-1) + arrived = dist < self.arrive_threshold + + # print(self.end) + self.trajectory_length[torch.logical_and(arrived, ~self.end)] = step + self.end = torch.logical_or(self.end, arrived) + + def get_trajectories(self): + return self.trajectory + + def step_target_qpos(self, target_qpos: torch.Tensor): + for i in range(self.sims_per_step): + self.franka_isaac_env.step_qpos(qpos=target_qpos) + next_cur_qpos = self.get_current_qpos() + return next_cur_qpos + + def get_current_qpos(self): + cur_qpos = self.franka_isaac_env.get_cur_qpos() + return cur_qpos.squeeze(1) + + def all_end(self) -> bool: + """ If all case is end + Returns: + A bool + """ + return torch.all(self.end) + +@ENV.register() +@torch.no_grad() +class FK2PlanningEnvWrapper(): + def __init__(self, cfg: DictConfig) -> None: + """ FK2Planning environment class for path planning task. + + Args: + cfg: + """ + self.max_sample_each_step = cfg.max_sample_each_step + self.max_trajectory_length = cfg.max_trajectory_length + self.arrive_threshold = cfg.arrive_threshold + self.eval_case_num_per_scene = cfg.eval_case_num_per_scene + self.sims_per_step = cfg.sims_per_step + self.arrive_threshold = cfg.arrive_threshold + + self.angle_normalize = None + self.angle_denormalize = None + + def run(self, + model: torch.nn.Module, + dataloader: torch.utils.data.DataLoader, + save_dir: str + ) -> None: + """ Planning within the environment + + Args: + model: diffusion model + dataloader: test dataloader + save_dir: save_directory of result + """ + # save_record_path = os.path.join('/home/puhao/dev/SceneDiffuser', 'succ_record.json') + model.eval() + device = model.device + + self.angle_normalize = dataloader.dataset.angle_normalize + self.angle_denormalize = dataloader.dataset.angle_denormalize + # task_name = 'pla-0.1@demo' + # res_saver_path = os.path.join('outputs', task_name, 'metrics.json') + # replay_saver_path = os.path.join('outputs', task_name, 'replay.pkl') + + replay_res = {} + + # succ_scene_record = json.load(open(save_record_path, 'r')) + res = defaultdict(list) + res['succ'] = [] + res['eval_cnt'] = 0 + res['max_trajectory_length'] = self.max_trajectory_length + res['arrive_threshold'] = self.arrive_threshold + res['scene_list'] = [] + # if os.path.exists(res_saver_path): + # with open(res_saver_path, 'r') as fp: + # res = json.load(fp) + # else: + # with open(res_saver_path, 'w') as fp: + # json.dump(res, fp) + for i, data in enumerate(dataloader): + for key in data: + if torch.is_tensor(data[key]): + data[key] = data[key].to(device) + + i_scene_id = data['scene_id'][0] + i_target_qpos = np.array(data['target'].cpu()) + # if data['scene_id'][0] in succ_scene_record.keys(): + # continue + if data['scene_id'][0] in res['scene_list']: + continue + succ_record_list = [] + ## de-normalize for env input + if dataloader.dataset.normalize_x: + data['x'] = self.angle_denormalize(data['x'].cpu()).cuda() + data['target'] = self.angle_denormalize(data['target'].cpu()).cuda() + data['start'] = self.angle_denormalize(data['start'].cpu()).cuda() + + env = FK2PlanningEnvCore( + data=copy.deepcopy(data), + scene_id=data['scene_id'][0], + arrive_threshold=self.arrive_threshold, + sims_per_step=self.sims_per_step, + max_trajectory_length=self.max_trajectory_length + ) + + ## re-normalize for model input + if dataloader.dataset.normalize_x: + data['x'] = self.angle_normalize(data['x'].cpu()).cuda() + data['target'] = self.angle_normalize(data['target'].cpu()).cuda() + data['start'] = self.angle_normalize(data['start'].cpu()).cuda() + + for j in tqdm(range(self.max_trajectory_length)): + outputs = model.sample(data, k=self.max_sample_each_step) + # assert torch.sum(outputs[:, 0:1, -1, 0, :] - env.state) < 1e-6 + assert (outputs.shape[1] == 1) + + pred_next_qpos = outputs[:, 0, -1, 10, :] + # # todo: l2 loss + # pred_next_qpos = data['target'] + + if dataloader.dataset.normalize_x: + pred_next_qpos = self.angle_denormalize(pred_next_qpos.cpu()).cuda() + env.step(j + 1, pred_next_qpos) + + data['start'] = env.state.clone() + if dataloader.dataset.normalize_x: + data['start'] = self.angle_normalize(data['start'].cpu()).cuda() + data['start'] = data['start'].unsqueeze(1) + if env.all_end(): + break + print() + i_trajectories = env.get_trajectories() + i_trajectories = np.array([np.stack(case) for case in i_trajectories]) + # save replayer recorder + replay_res[i_scene_id] = {'sample_trajs': i_trajectories, + 'target_qpos': i_target_qpos} + + for j in range(env.batch): + succ_record_list.append(bool(env.end[j].cpu())) + res['scene_list'].append(data['scene_id'][0]) + res['succ'].append(bool(env.end[j].cpu())) + res['length'].append(float(env.trajectory_length[j])) + # print(f'[{i}/{200}]') + # succ_scene_record[data['scene_id'][0]] = succ_record_list + res['eval_cnt'] += env.batch + + # json.dump(succ_scene_record, open(save_record_path, 'w')) + + del env + # with open(res_saver_path, 'w') as fp: + # json.dump(res, fp) + + # with open(replay_saver_path, 'wb') as fp: + # pickle.dump(replay_res, fp) + for key in ['succ', 'length']: + res[key+'_average'] = sum(res[key]) / len(res[key]) + + ## save quantitative results + save_path = os.path.join(save_dir, 'metrics.json') + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, 'w') as fp: + json.dump(res, fp) + # with open(res_saver_path, 'w') as fp: + # json.dump(res, fp) + @ENV.register() @torch.no_grad() class PathPlanningEnvWrapper(): diff --git a/models/model/pointnet2/pointnet2_modules.py b/models/model/pointnet2/pointnet2_modules.py new file mode 100644 index 0000000..7cce31a --- /dev/null +++ b/models/model/pointnet2/pointnet2_modules.py @@ -0,0 +1,513 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +''' Pointnet2 layers. +Modified based on: https://github.com/erikwijmans/Pointnet2_PyTorch +Extended with the following: +1. Uniform sampling in each local region (sample_uniformly) +2. Return sampled points indices to support votenet. +''' +import torch +import torch.nn as nn +import torch.nn.functional as F + +from models.model.pointnet2 import pointnet2_utils +import models.model.pointnet2.pytorch_utils as pt_utils +from typing import List, Tuple + + +class _PointnetSAModuleBase(nn.Module): + + def __init__(self): + super().__init__() + self.npoint = None + self.groupers = None + self.mlps = None + + def forward(self, xyz: torch.Tensor, + features: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Parameters + ---------- + xyz : torch.Tensor + (B, N, 3) tensor of the xyz coordinates of the features + features : torch.Tensor + (B, N, C) tensor of the descriptors of the the features + + Returns + ------- + new_xyz : torch.Tensor + (B, npoint, 3) tensor of the new features' xyz + new_features : torch.Tensor + (B, npoint, \sum_k(mlps[k][-1])) tensor of the new_features descriptors + """ + + new_features_list = [] + + xyz_flipped = xyz.transpose(1, 2).contiguous() + new_xyz = pointnet2_utils.gather_operation( + xyz_flipped, + pointnet2_utils.furthest_point_sample(xyz, self.npoint) + ).transpose(1, 2).contiguous() if self.npoint is not None else None + + for i in range(len(self.groupers)): + new_features = self.groupers[i]( + xyz, new_xyz, features + ) # (B, C, npoint, nsample) + + new_features = self.mlps[i]( + new_features + ) # (B, mlp[-1], npoint, nsample) + new_features = F.max_pool2d( + new_features, kernel_size=[1, new_features.size(3)] + ) # (B, mlp[-1], npoint, 1) + new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) + + new_features_list.append(new_features) + + return new_xyz, torch.cat(new_features_list, dim=1) + + +class PointnetSAModuleMSG(_PointnetSAModuleBase): + r"""Pointnet set abstrction layer with multiscale grouping + + Parameters + ---------- + npoint : int + Number of features + radii : list of float32 + list of radii to group with + nsamples : list of int32 + Number of samples in each ball query + mlps : list of list of int32 + Spec of the pointnet before the global max_pool for each scale + bn : bool + Use batchnorm + """ + + def __init__( + self, + *, + npoint: int, + radii: List[float], + nsamples: List[int], + mlps: List[List[int]], + bn: bool = True, + use_xyz: bool = True, + sample_uniformly: bool = False + ): + super().__init__() + + assert len(radii) == len(nsamples) == len(mlps) + + self.npoint = npoint + self.groupers = nn.ModuleList() + self.mlps = nn.ModuleList() + for i in range(len(radii)): + radius = radii[i] + nsample = nsamples[i] + self.groupers.append( + pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, sample_uniformly=sample_uniformly) + if npoint is not None else pointnet2_utils.GroupAll(use_xyz) + ) + mlp_spec = mlps[i] + if use_xyz: + mlp_spec[0] += 3 + + self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) + + +class PointnetSAModule(PointnetSAModuleMSG): + r"""Pointnet set abstrction layer + + Parameters + ---------- + npoint : int + Number of features + radius : float + Radius of ball + nsample : int + Number of samples in the ball query + mlp : list + Spec of the pointnet before the global max_pool + bn : bool + Use batchnorm + """ + + def __init__( + self, + *, + mlp: List[int], + npoint: int = None, + radius: float = None, + nsample: int = None, + bn: bool = True, + use_xyz: bool = True + ): + super().__init__( + mlps=[mlp], + npoint=npoint, + radii=[radius], + nsamples=[nsample], + bn=bn, + use_xyz=use_xyz + ) + + +class PointnetSAModuleVotes(nn.Module): + ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG + with extra support for returning point indices for getting their GT votes ''' + + def __init__( + self, + *, + mlp: List[int], + npoint: int = None, + radius: float = None, + nsample: int = None, + bn: bool = True, + use_xyz: bool = True, + pooling: str = 'max', + sigma: float = None, # for RBF pooling + normalize_xyz: bool = False, # noramlize local XYZ with radius + sample_uniformly: bool = False, + ret_unique_cnt: bool = False + ): + super().__init__() + + self.npoint = npoint + self.radius = radius + self.nsample = nsample + self.pooling = pooling + self.mlp_module = None + self.use_xyz = use_xyz + self.sigma = sigma + if self.sigma is None: + self.sigma = self.radius/2 + self.normalize_xyz = normalize_xyz + self.ret_unique_cnt = ret_unique_cnt + + if npoint is not None: + self.grouper = pointnet2_utils.QueryAndGroup(radius, nsample, + use_xyz=use_xyz, ret_grouped_xyz=True, normalize_xyz=normalize_xyz, + sample_uniformly=sample_uniformly, ret_unique_cnt=ret_unique_cnt) + else: + self.grouper = pointnet2_utils.GroupAll(use_xyz, ret_grouped_xyz=True) + + mlp_spec = mlp + if use_xyz and len(mlp_spec)>0: + mlp_spec[0] += 3 + self.mlp_module = pt_utils.SharedMLP(mlp_spec, bn=bn) + + + def forward(self, xyz: torch.Tensor, + features: torch.Tensor = None, + inds: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Parameters + ---------- + xyz : torch.Tensor + (B, N, 3) tensor of the xyz coordinates of the features + features : torch.Tensor + (B, C, N) tensor of the descriptors of the the features + inds : torch.Tensor + (B, npoint) tensor that stores index to the xyz points (values in 0-N-1) + + Returns + ------- + new_xyz : torch.Tensor + (B, npoint, 3) tensor of the new features' xyz + new_features : torch.Tensor + (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors + inds: torch.Tensor + (B, npoint) tensor of the inds + """ + + xyz_flipped = xyz.transpose(1, 2).contiguous() + if inds is None: + inds = pointnet2_utils.furthest_point_sample(xyz, self.npoint) + else: + assert(inds.shape[1] == self.npoint) + new_xyz = pointnet2_utils.gather_operation( + xyz_flipped, inds + ).transpose(1, 2).contiguous() if self.npoint is not None else None + + if not self.ret_unique_cnt: + grouped_features, grouped_xyz = self.grouper( + xyz, new_xyz, features + ) # (B, C, npoint, nsample) + else: + grouped_features, grouped_xyz, unique_cnt = self.grouper( + xyz, new_xyz, features + ) # (B, C, npoint, nsample), (B,3,npoint,nsample), (B,npoint) + + new_features = self.mlp_module( + grouped_features + ) # (B, mlp[-1], npoint, nsample) + if self.pooling == 'max': + new_features = F.max_pool2d( + new_features, kernel_size=[1, new_features.size(3)] + ) # (B, mlp[-1], npoint, 1) + elif self.pooling == 'avg': + new_features = F.avg_pool2d( + new_features, kernel_size=[1, new_features.size(3)] + ) # (B, mlp[-1], npoint, 1) + elif self.pooling == 'rbf': + # Use radial basis function kernel for weighted sum of features (normalized by nsample and sigma) + # Ref: https://en.wikipedia.org/wiki/Radial_basis_function_kernel + rbf = torch.exp(-1 * grouped_xyz.pow(2).sum(1,keepdim=False) / (self.sigma**2) / 2) # (B, npoint, nsample) + new_features = torch.sum(new_features * rbf.unsqueeze(1), -1, keepdim=True) / float(self.nsample) # (B, mlp[-1], npoint, 1) + new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) + + if not self.ret_unique_cnt: + return new_xyz, new_features, inds + else: + return new_xyz, new_features, inds, unique_cnt + +class PointnetSAModuleMSGVotes(nn.Module): + ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG + with extra support for returning point indices for getting their GT votes ''' + + def __init__( + self, + *, + mlps: List[List[int]], + npoint: int, + radii: List[float], + nsamples: List[int], + bn: bool = True, + use_xyz: bool = True, + sample_uniformly: bool = False + ): + super().__init__() + + assert(len(mlps) == len(nsamples) == len(radii)) + + self.npoint = npoint + self.groupers = nn.ModuleList() + self.mlps = nn.ModuleList() + for i in range(len(radii)): + radius = radii[i] + nsample = nsamples[i] + self.groupers.append( + pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, sample_uniformly=sample_uniformly) + if npoint is not None else pointnet2_utils.GroupAll(use_xyz) + ) + mlp_spec = mlps[i] + if use_xyz: + mlp_spec[0] += 3 + + self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) + + def forward(self, xyz: torch.Tensor, + features: torch.Tensor = None, inds: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]: + r""" + Parameters + ---------- + xyz : torch.Tensor + (B, N, 3) tensor of the xyz coordinates of the features + features : torch.Tensor + (B, C, C) tensor of the descriptors of the the features + inds : torch.Tensor + (B, npoint) tensor that stores index to the xyz points (values in 0-N-1) + + Returns + ------- + new_xyz : torch.Tensor + (B, npoint, 3) tensor of the new features' xyz + new_features : torch.Tensor + (B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors + inds: torch.Tensor + (B, npoint) tensor of the inds + """ + new_features_list = [] + + xyz_flipped = xyz.transpose(1, 2).contiguous() + if inds is None: + inds = pointnet2_utils.furthest_point_sample(xyz, self.npoint) + new_xyz = pointnet2_utils.gather_operation( + xyz_flipped, inds + ).transpose(1, 2).contiguous() if self.npoint is not None else None + + for i in range(len(self.groupers)): + new_features = self.groupers[i]( + xyz, new_xyz, features + ) # (B, C, npoint, nsample) + new_features = self.mlps[i]( + new_features + ) # (B, mlp[-1], npoint, nsample) + new_features = F.max_pool2d( + new_features, kernel_size=[1, new_features.size(3)] + ) # (B, mlp[-1], npoint, 1) + new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint) + + new_features_list.append(new_features) + + return new_xyz, torch.cat(new_features_list, dim=1), inds + + +class PointnetFPModule(nn.Module): + r"""Propigates the features of one set to another + + Parameters + ---------- + mlp : list + Pointnet module parameters + bn : bool + Use batchnorm + """ + + def __init__(self, *, mlp: List[int], bn: bool = True): + super().__init__() + self.mlp = pt_utils.SharedMLP(mlp, bn=bn) + + def forward( + self, unknown: torch.Tensor, known: torch.Tensor, + unknow_feats: torch.Tensor, known_feats: torch.Tensor + ) -> torch.Tensor: + r""" + Parameters + ---------- + unknown : torch.Tensor + (B, n, 3) tensor of the xyz positions of the unknown features + known : torch.Tensor + (B, m, 3) tensor of the xyz positions of the known features + unknow_feats : torch.Tensor + (B, C1, n) tensor of the features to be propigated to + known_feats : torch.Tensor + (B, C2, m) tensor of features to be propigated + + Returns + ------- + new_features : torch.Tensor + (B, mlp[-1], n) tensor of the features of the unknown features + """ + + if known is not None: + dist, idx = pointnet2_utils.three_nn(unknown, known) + dist_recip = 1.0 / (dist + 1e-8) + norm = torch.sum(dist_recip, dim=2, keepdim=True) + weight = dist_recip / norm + + interpolated_feats = pointnet2_utils.three_interpolate( + known_feats, idx, weight + ) + else: + interpolated_feats = known_feats.expand( + *known_feats.size()[0:2], unknown.size(1) + ) + + if unknow_feats is not None: + new_features = torch.cat([interpolated_feats, unknow_feats], + dim=1) #(B, C2 + C1, n) + else: + new_features = interpolated_feats + + new_features = new_features.unsqueeze(-1) + new_features = self.mlp(new_features) + + return new_features.squeeze(-1) + +class PointnetLFPModuleMSG(nn.Module): + ''' Modified based on _PointnetSAModuleBase and PointnetSAModuleMSG + learnable feature propagation layer.''' + + def __init__( + self, + *, + mlps: List[List[int]], + radii: List[float], + nsamples: List[int], + post_mlp: List[int], + bn: bool = True, + use_xyz: bool = True, + sample_uniformly: bool = False + ): + super().__init__() + + assert(len(mlps) == len(nsamples) == len(radii)) + + self.post_mlp = pt_utils.SharedMLP(post_mlp, bn=bn) + + self.groupers = nn.ModuleList() + self.mlps = nn.ModuleList() + for i in range(len(radii)): + radius = radii[i] + nsample = nsamples[i] + self.groupers.append( + pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz, + sample_uniformly=sample_uniformly) + ) + mlp_spec = mlps[i] + if use_xyz: + mlp_spec[0] += 3 + + self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn)) + + def forward(self, xyz2: torch.Tensor, xyz1: torch.Tensor, + features2: torch.Tensor, features1: torch.Tensor) -> torch.Tensor: + r""" Propagate features from xyz1 to xyz2. + Parameters + ---------- + xyz2 : torch.Tensor + (B, N2, 3) tensor of the xyz coordinates of the features + xyz1 : torch.Tensor + (B, N1, 3) tensor of the xyz coordinates of the features + features2 : torch.Tensor + (B, C2, N2) tensor of the descriptors of the the features + features1 : torch.Tensor + (B, C1, N1) tensor of the descriptors of the the features + + Returns + ------- + new_features1 : torch.Tensor + (B, \sum_k(mlps[k][-1]), N1) tensor of the new_features descriptors + """ + new_features_list = [] + + for i in range(len(self.groupers)): + new_features = self.groupers[i]( + xyz1, xyz2, features1 + ) # (B, C1, N2, nsample) + new_features = self.mlps[i]( + new_features + ) # (B, mlp[-1], N2, nsample) + new_features = F.max_pool2d( + new_features, kernel_size=[1, new_features.size(3)] + ) # (B, mlp[-1], N2, 1) + new_features = new_features.squeeze(-1) # (B, mlp[-1], N2) + + if features2 is not None: + new_features = torch.cat([new_features, features2], + dim=1) #(B, mlp[-1] + C2, N2) + + new_features = new_features.unsqueeze(-1) + new_features = self.post_mlp(new_features) + + new_features_list.append(new_features) + + return torch.cat(new_features_list, dim=1).squeeze(-1) + + +if __name__ == "__main__": + from torch.autograd import Variable + torch.manual_seed(1) + torch.cuda.manual_seed_all(1) + xyz = Variable(torch.randn(2, 9, 3).cuda(), requires_grad=True) + xyz_feats = Variable(torch.randn(2, 9, 6).cuda(), requires_grad=True) + + test_module = PointnetSAModuleMSG( + npoint=2, radii=[5.0, 10.0], nsamples=[6, 3], mlps=[[9, 3], [9, 6]] + ) + test_module.cuda() + print(test_module(xyz, xyz_feats)) + + for _ in range(1): + _, new_features = test_module(xyz, xyz_feats) + new_features.backward( + torch.cuda.FloatTensor(*new_features.size()).fill_(1) + ) + print(new_features) + print(xyz.grad) \ No newline at end of file diff --git a/models/model/pointnet2/pointnet2_semseg.py b/models/model/pointnet2/pointnet2_semseg.py new file mode 100644 index 0000000..5dd3200 --- /dev/null +++ b/models/model/pointnet2/pointnet2_semseg.py @@ -0,0 +1,321 @@ +import torch +import torch.nn as nn +from models.model.pointnet2.pointnet2_modules import PointnetFPModule, PointnetSAModule, PointnetSAModuleMSG +import models.model.pointnet2.pytorch_utils as pt_utils + + +def get_model(num_classes, is_msg=True, input_channels=6, use_xyz=True, bn=True): + if is_msg: + model = Pointnet2MSG( + num_classes=num_classes, + input_channels=input_channels, + use_xyz=use_xyz, + bn=bn + ) + else: + model = Pointnet2SSG( + num_classes=num_classes, + input_channels=input_channels, + use_xyz=use_xyz, + bn=bn + ) + + return model + +class Pointnet2MSG(nn.Module): + def __init__(self, num_classes, input_channels=3, use_xyz=True, bn=True): + super().__init__() + + NPOINTS = [1024, 256, 64, 16] + RADIUS = [[0.05, 0.1], [0.1, 0.2], [0.2, 0.4], [0.4, 0.8]] + NSAMPLE = [[16, 32], [16, 32], [16, 32], [16, 32]] + MLPS = [[[16, 16, 32], [32, 32, 64]], [[64, 64, 128], [64, 96, 128]], + [[128, 196, 256], [128, 196, 256]], [[256, 256, 512], [256, 384, 512]]] + FP_MLPS = [[128, 128], [256, 256], [512, 512], [512, 512]] + CLS_FC = [128] + DP_RATIO = 0.5 + + self.SA_modules = nn.ModuleList() + channel_in = input_channels + + skip_channel_list = [input_channels] + for k in range(NPOINTS.__len__()): + mlps = MLPS[k].copy() + channel_out = 0 + for idx in range(mlps.__len__()): + mlps[idx] = [channel_in] + mlps[idx] + channel_out += mlps[idx][-1] + + self.SA_modules.append( + PointnetSAModuleMSG( + npoint=NPOINTS[k], + radii=RADIUS[k], + nsamples=NSAMPLE[k], + mlps=mlps, + use_xyz=use_xyz, + bn=bn + ) + ) + skip_channel_list.append(channel_out) + channel_in = channel_out + + self.FP_modules = nn.ModuleList() + + for k in range(FP_MLPS.__len__()): + pre_channel = FP_MLPS[k + 1][-1] if k + 1 < len(FP_MLPS) else channel_out + self.FP_modules.append( + PointnetFPModule( + mlp=[pre_channel + skip_channel_list[k]] + FP_MLPS[k], + bn=bn + ) + ) + + cls_layers = [] + pre_channel = FP_MLPS[0][-1] + for k in range(0, CLS_FC.__len__()): + cls_layers.append(pt_utils.Conv1d(pre_channel, CLS_FC[k], bn=bn)) + pre_channel = CLS_FC[k] + cls_layers.append(pt_utils.Conv1d(pre_channel, num_classes, activation=None, bn=bn)) + cls_layers.insert(1, nn.Dropout(DP_RATIO)) + self.cls_layer = nn.Sequential(*cls_layers) + + def _break_up_pc(self, pc): + xyz = pc[..., 0:3].contiguous() + features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None + + return xyz, features + + def forward(self, pointcloud: torch.cuda.FloatTensor): + xyz, features = self._break_up_pc(pointcloud) + + l_xyz, l_features = [xyz], [features] + for i in range(len(self.SA_modules)): + li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i]) + l_xyz.append(li_xyz) + l_features.append(li_features) + + for i in range(-1, -(len(self.FP_modules) + 1), -1): + l_features[i - 1] = self.FP_modules[i]( + l_xyz[i - 1], l_xyz[i], l_features[i - 1], l_features[i] + ) + + pred_cls = self.cls_layer(l_features[0]).transpose(1, 2).contiguous() # (B, N, 1) + return pred_cls + +class Pointnet2SSG(nn.Module): + def __init__(self, num_classes, input_channels=3, use_xyz=True, bn=True): + super().__init__() + + NPOINTS = [1024, 256, 64, 16] + RADIUS = [0.1, 0.2, 0.4, 0.8] + NSAMPLE = [32, 32, 32, 32] + MLPS = [[32, 32, 64], [64, 64, 128], + [128, 128, 256], [256, 256, 512]] + FP_MLPS = [[128, 128], [256, 128], [256, 256], [256, 256]] + CLS_FC = [128] + DP_RATIO = 0.5 + + self.SA_modules = nn.ModuleList() + channel_in = input_channels + + skip_channel_list = [input_channels] + for k in range(NPOINTS.__len__()): + mlps = MLPS[k].copy() + channel_out = 0 + mlps = [channel_in] + mlps + channel_out += mlps[-1] + + self.SA_modules.append( + PointnetSAModule( + npoint=NPOINTS[k], + radius=RADIUS[k], + nsample=NSAMPLE[k], + mlp=mlps, + use_xyz=use_xyz, + bn=bn + ) + ) + skip_channel_list.append(channel_out) + channel_in = channel_out + + self.FP_modules = nn.ModuleList() + + for k in range(FP_MLPS.__len__()): + pre_channel = FP_MLPS[k + 1][-1] if k + 1 < len(FP_MLPS) else channel_out + self.FP_modules.append( + PointnetFPModule( + mlp=[pre_channel + skip_channel_list[k]] + FP_MLPS[k], + bn=bn + ) + ) + + cls_layers = [] + pre_channel = FP_MLPS[0][-1] + for k in range(0, CLS_FC.__len__()): + cls_layers.append(pt_utils.Conv1d(pre_channel, CLS_FC[k], bn=bn)) + pre_channel = CLS_FC[k] + cls_layers.append(pt_utils.Conv1d(pre_channel, num_classes, activation=None, bn=bn)) + cls_layers.insert(1, nn.Dropout(DP_RATIO)) + self.cls_layer = nn.Sequential(*cls_layers) + + def _break_up_pc(self, pc): + xyz = pc[..., 0:3].contiguous() + features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None + + return xyz, features + + def forward(self, pointcloud: torch.cuda.FloatTensor): + xyz, features = self._break_up_pc(pointcloud) + + l_xyz, l_features = [xyz], [features] + for i in range(len(self.SA_modules)): + li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i]) + l_xyz.append(li_xyz) + l_features.append(li_features) + + # last layer, l_xyz <4, 16, 3>, l_feature <4, 512, 16> + for i in range(-1, -(len(self.FP_modules) + 1), -1): + l_features[i - 1] = self.FP_modules[i]( + l_xyz[i - 1], l_xyz[i], l_features[i - 1], l_features[i] + ) + + pred_cls = self.cls_layer(l_features[0]).transpose(1, 2).contiguous() # (B, N, categories) + return pred_cls + + +def get_feature_extractor(is_msg=True, input_channels=6, use_xyz=True, bn=True): + if is_msg: + model = Pointnet2MSG_Feature( + input_channels=input_channels, + use_xyz=use_xyz, + bn=bn + ) + else: + model = Pointnet2SSG_Feature( + input_channels=input_channels, + use_xyz=use_xyz, + bn=bn + ) + + return model + +class Pointnet2MSG_Feature(nn.Module): + def __init__(self, input_channels=3, use_xyz=True, bn=True): + super().__init__() + + NPOINTS = [1024, 256, 64, 16] + RADIUS = [[0.05, 0.1], [0.1, 0.2], [0.2, 0.4], [0.4, 0.8]] + NSAMPLE = [[16, 32], [16, 32], [16, 32], [16, 32]] + MLPS = [[[16, 16, 32], [32, 32, 64]], + [[64, 64, 128], [64, 96, 128]], + [[128, 196, 256], [128, 196, 256]], + [[256, 256, 512], [256, 384, 512]]] + + self.SA_modules = nn.ModuleList() + channel_in = input_channels + + skip_channel_list = [input_channels] + for k in range(NPOINTS.__len__()): + mlps = MLPS[k].copy() + channel_out = 0 + for idx in range(mlps.__len__()): + mlps[idx] = [channel_in] + mlps[idx] + channel_out += mlps[idx][-1] + + self.SA_modules.append( + PointnetSAModuleMSG( + npoint=NPOINTS[k], + radii=RADIUS[k], + nsamples=NSAMPLE[k], + mlps=mlps, + use_xyz=use_xyz, + bn=bn + ) + ) + skip_channel_list.append(channel_out) + channel_in = channel_out + + def _break_up_pc(self, pc): + xyz = pc[..., 0:3].contiguous() + features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None + + return xyz, features + + def forward(self, pointcloud: torch.cuda.FloatTensor): + xyz, features = self._break_up_pc(pointcloud) + + l_xyz, l_features = [xyz], [features] + for i in range(len(self.SA_modules)): + li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i]) + l_xyz.append(li_xyz) + l_features.append(li_features) + + return l_xyz, l_features + +class Pointnet2SSG_Feature(nn.Module): + def __init__(self, input_channels=3, use_xyz=True, bn=True): + super().__init__() + + # NPOINTS = [1024, 256, 64, 16] + NPOINTS = [2048, 512, 128, 16] + + RADIUS = [0.02, 0.04, 0.06, 0.08] + NSAMPLE = [32, 32, 16, 16] + MLPS = [[32, 32, 64], [64, 64, 128], + [128, 128, 256], [256, 256, 512]] + + self.SA_modules = nn.ModuleList() + channel_in = input_channels + + skip_channel_list = [input_channels] + for k in range(NPOINTS.__len__()): + mlps = MLPS[k].copy() + channel_out = 0 + mlps = [channel_in] + mlps + channel_out += mlps[-1] + + self.SA_modules.append( + PointnetSAModule( + npoint=NPOINTS[k], + radius=RADIUS[k], + nsample=NSAMPLE[k], + mlp=mlps, + use_xyz=use_xyz, + bn=bn + ) + ) + skip_channel_list.append(channel_out) + channel_in = channel_out + + def _break_up_pc(self, pc): + xyz = pc[..., 0:3].contiguous() + features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None + + return xyz, features + + def forward(self, pointcloud: torch.cuda.FloatTensor): + xyz, features = self._break_up_pc(pointcloud) + + l_xyz, l_features = [xyz], [features] + for i in range(len(self.SA_modules)): + li_xyz, li_features = self.SA_modules[i](l_xyz[i], l_features[i]) + l_xyz.append(li_xyz) + l_features.append(li_features) + + return l_xyz, l_features + + +def pointnet2_enc_repro(c=3, num_points=2048): + assert (num_points == 2048) ## cannot adapt num of points here + model = get_feature_extractor(is_msg=False, input_channels=c-3, use_xyz=True, bn=True) + return model + + +if __name__ == '__main__': + print(f'get pointnet2 semseg') + msg_model = get_feature_extractor(is_msg=False, input_channels=0, use_xyz=True, bn=True) + msg_model.cuda() + + dummy_pc_input = torch.randn(8, 1024, 3, device='cuda') + dummy_feat = msg_model(dummy_pc_input) + print(dummy_feat[-1][-1].shape) \ No newline at end of file diff --git a/models/model/pointnet2/pointnet2_utils.py b/models/model/pointnet2/pointnet2_utils.py new file mode 100644 index 0000000..8996dc8 --- /dev/null +++ b/models/model/pointnet2/pointnet2_utils.py @@ -0,0 +1,427 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +''' Modified based on: https://github.com/erikwijmans/Pointnet2_PyTorch ''' +from __future__ import ( + division, + absolute_import, + with_statement, + print_function, + unicode_literals, +) +from typing import Tuple, Any +import torch +from torch.autograd import Function +import torch.nn as nn +import models.model.pointnet2.pytorch_utils as pt_utils +import sys + +try: + import builtins +except: + raise NotImplementedError('module __builtin__ not support yet') + # import __builtin__ as builtins + +try: + import pointnet2._ext as _ext +except ImportError: + if not getattr(builtins, "__POINTNET2_SETUP__", False): + raise ImportError( + "Could not import _ext module.\n" + "Please see the setup instructions in the README: " + "https://github.com/erikwijmans/Pointnet2_PyTorch/blob/master/README.rst" + ) + +if False: + # Workaround for type hints without depending on the `typing` module + from typing import * + + +class RandomDropout(nn.Module): + def __init__(self, p=0.5, inplace=False): + super(RandomDropout, self).__init__() + self.p = p + self.inplace = inplace + + def forward(self, X): + theta = torch.Tensor(1).uniform_(0, self.p)[0] + return pt_utils.feature_dropout_no_scaling(X, theta, self.train, self.inplace) + + +class FurthestPointSampling(Function): + @staticmethod + def forward(ctx, xyz, npoint): + # type: (Any, torch.Tensor, int) -> torch.Tensor + r""" + Uses iterative furthest point sampling to select a set of npoint features that have the largest + minimum distance + + Parameters + ---------- + xyz : torch.Tensor + (B, N, 3) tensor where N > npoint + npoint : int32 + number of features in the sampled set + + Returns + ------- + torch.Tensor + (B, npoint) tensor containing the set + """ + fps_inds = _ext.furthest_point_sampling(xyz, npoint) + ctx.mark_non_differentiable(fps_inds) + return fps_inds + + @staticmethod + def backward(xyz, a=None): + return None, None + + +furthest_point_sample = FurthestPointSampling.apply + + +class GatherOperation(Function): + @staticmethod + def forward(ctx, features, idx): + # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor + r""" + + Parameters + ---------- + features : torch.Tensor + (B, C, N) tensor + + idx : torch.Tensor + (B, npoint) tensor of the features to gather + + Returns + ------- + torch.Tensor + (B, C, npoint) tensor + """ + + _, C, N = features.size() + + ctx.for_backwards = (idx, C, N) + + return _ext.gather_points(features, idx) + + @staticmethod + def backward(ctx, grad_out): + idx, C, N = ctx.for_backwards + + grad_features = _ext.gather_points_grad(grad_out.contiguous(), idx, N) + return grad_features, None + + +gather_operation = GatherOperation.apply + + +class ThreeNN(Function): + @staticmethod + def forward(ctx, unknown, known): + # type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] + r""" + Find the three nearest neighbors of unknown in known + Parameters + ---------- + unknown : torch.Tensor + (B, n, 3) tensor of known features + known : torch.Tensor + (B, m, 3) tensor of unknown features + + Returns + ------- + dist : torch.Tensor + (B, n, 3) l2 distance to the three nearest neighbors + idx : torch.Tensor + (B, n, 3) index of 3 nearest neighbors + """ + dist2, idx = _ext.three_nn(unknown, known) + + return torch.sqrt(dist2), idx + + @staticmethod + def backward(ctx, a=None, b=None): + return None, None + + +three_nn = ThreeNN.apply + + +class ThreeInterpolate(Function): + @staticmethod + def forward(ctx, features, idx, weight): + # type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor + r""" + Performs weight linear interpolation on 3 features + Parameters + ---------- + features : torch.Tensor + (B, c, m) Features descriptors to be interpolated from + idx : torch.Tensor + (B, n, 3) three nearest neighbors of the target features in features + weight : torch.Tensor + (B, n, 3) weights + + Returns + ------- + torch.Tensor + (B, c, n) tensor of the interpolated features + """ + B, c, m = features.size() + n = idx.size(1) + + ctx.three_interpolate_for_backward = (idx, weight, m) + + return _ext.three_interpolate(features, idx, weight) + + @staticmethod + def backward(ctx, grad_out): + # type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + r""" + Parameters + ---------- + grad_out : torch.Tensor + (B, c, n) tensor with gradients of ouputs + + Returns + ------- + grad_features : torch.Tensor + (B, c, m) tensor with gradients of features + + None + + None + """ + idx, weight, m = ctx.three_interpolate_for_backward + + grad_features = _ext.three_interpolate_grad( + grad_out.contiguous(), idx, weight, m + ) + + return grad_features, None, None + + +three_interpolate = ThreeInterpolate.apply + + +class GroupingOperation(Function): + @staticmethod + def forward(ctx, features, idx): + # type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor + r""" + + Parameters + ---------- + features : torch.Tensor + (B, C, N) tensor of features to group + idx : torch.Tensor + (B, npoint, nsample) tensor containing the indicies of features to group with + + Returns + ------- + torch.Tensor + (B, C, npoint, nsample) tensor + """ + B, nfeatures, nsample = idx.size() + _, C, N = features.size() + + ctx.for_backwards = (idx, N) + + return _ext.group_points(features, idx) + + @staticmethod + def backward(ctx, grad_out): + # type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor] + r""" + + Parameters + ---------- + grad_out : torch.Tensor + (B, C, npoint, nsample) tensor of the gradients of the output from forward + + Returns + ------- + torch.Tensor + (B, C, N) gradient of the features + None + """ + idx, N = ctx.for_backwards + + grad_features = _ext.group_points_grad(grad_out.contiguous(), idx, N) + + return grad_features, None + + +grouping_operation = GroupingOperation.apply + + +class BallQuery(Function): + @staticmethod + def forward(ctx, radius, nsample, xyz, new_xyz): + # type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor + r""" + + Parameters + ---------- + radius : float + radius of the balls + nsample : int + maximum number of features in the balls + xyz : torch.Tensor + (B, N, 3) xyz coordinates of the features + new_xyz : torch.Tensor + (B, npoint, 3) centers of the ball query + + Returns + ------- + torch.Tensor + (B, npoint, nsample) tensor with the indicies of the features that form the query balls + """ + inds = _ext.ball_query(new_xyz, xyz, radius, nsample) + ctx.mark_non_differentiable(inds) + return inds + + @staticmethod + def backward(ctx, a=None): + return None, None, None, None + + +ball_query = BallQuery.apply + + +class QueryAndGroup(nn.Module): + r""" + Groups with a ball query of radius + + Parameters + --------- + radius : float32 + Radius of ball + nsample : int32 + Maximum number of features to gather in the ball + """ + + def __init__(self, radius, nsample, use_xyz=True, ret_grouped_xyz=False, normalize_xyz=False, sample_uniformly=False, ret_unique_cnt=False): + ## type: (QueryAndGroup, float, int, bool) -> None + super(QueryAndGroup, self).__init__() + self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz + self.ret_grouped_xyz = ret_grouped_xyz + self.normalize_xyz = normalize_xyz + self.sample_uniformly = sample_uniformly + self.ret_unique_cnt = ret_unique_cnt + if self.ret_unique_cnt: + assert(self.sample_uniformly) + + def forward(self, xyz, new_xyz, features=None): + # type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor] + r""" + Parameters + ---------- + xyz : torch.Tensor + xyz coordinates of the features (B, N, 3) + new_xyz : torch.Tensor + centriods (B, npoint, 3) + features : torch.Tensor + Descriptors of the features (B, C, N) + + Returns + ------- + new_features : torch.Tensor + (B, 3 + C, npoint, nsample) tensor + """ + idx = ball_query(self.radius, self.nsample, xyz, new_xyz) + + if self.sample_uniformly: + unique_cnt = torch.zeros((idx.shape[0], idx.shape[1])) + for i_batch in range(idx.shape[0]): + for i_region in range(idx.shape[1]): + unique_ind = torch.unique(idx[i_batch, i_region, :]) + num_unique = unique_ind.shape[0] + unique_cnt[i_batch, i_region] = num_unique + sample_ind = torch.randint(0, num_unique, (self.nsample - num_unique,), dtype=torch.long) + all_ind = torch.cat((unique_ind, unique_ind[sample_ind])) + idx[i_batch, i_region, :] = all_ind + + + xyz_trans = xyz.transpose(1, 2).contiguous() + grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample) + grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1) + if self.normalize_xyz: + grouped_xyz /= self.radius + + if features is not None: + grouped_features = grouping_operation(features, idx) + if self.use_xyz: + new_features = torch.cat( + [grouped_xyz, grouped_features], dim=1 + ) # (B, C + 3, npoint, nsample) + else: + new_features = grouped_features + else: + assert ( + self.use_xyz + ), "Cannot have not features and not use xyz as a feature!" + new_features = grouped_xyz + + ret = [new_features] + if self.ret_grouped_xyz: + ret.append(grouped_xyz) + if self.ret_unique_cnt: + ret.append(unique_cnt) + if len(ret) == 1: + return ret[0] + else: + return tuple(ret) + + +class GroupAll(nn.Module): + r""" + Groups all features + + Parameters + --------- + """ + + def __init__(self, use_xyz=True, ret_grouped_xyz=False): + # type: (GroupAll, bool) -> None + super(GroupAll, self).__init__() + self.use_xyz = use_xyz + + def forward(self, xyz, new_xyz, features=None): + # type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor] + r""" + Parameters + ---------- + xyz : torch.Tensor + xyz coordinates of the features (B, N, 3) + new_xyz : torch.Tensor + Ignored + features : torch.Tensor + Descriptors of the features (B, C, N) + + Returns + ------- + new_features : torch.Tensor + (B, C + 3, 1, N) tensor + """ + + grouped_xyz = xyz.transpose(1, 2).unsqueeze(2) + if features is not None: + grouped_features = features.unsqueeze(2) + if self.use_xyz: + new_features = torch.cat( + [grouped_xyz, grouped_features], dim=1 + ) # (B, 3 + C, 1, N) + else: + new_features = grouped_features + else: + new_features = grouped_xyz + + if self.ret_grouped_xyz: + return new_features, grouped_xyz + else: + return new_features \ No newline at end of file diff --git a/models/model/pointnet2/pytorch_utils.py b/models/model/pointnet2/pytorch_utils.py new file mode 100644 index 0000000..6dc56b0 --- /dev/null +++ b/models/model/pointnet2/pytorch_utils.py @@ -0,0 +1,297 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +''' Modified based on Ref: https://github.com/erikwijmans/Pointnet2_PyTorch ''' +import torch +import torch.nn as nn +from typing import List, Tuple + +class SharedMLP(nn.Sequential): + + def __init__( + self, + args: List[int], + *, + bn: bool = False, + activation=nn.ReLU(inplace=True), + preact: bool = False, + first: bool = False, + name: str = "" + ): + super().__init__() + + for i in range(len(args) - 1): + self.add_module( + name + 'layer{}'.format(i), + Conv2d( + args[i], + args[i + 1], + bn=(not first or not preact or (i != 0)) and bn, + activation=activation + if (not first or not preact or (i != 0)) else None, + preact=preact + ) + ) + + +class _BNBase(nn.Sequential): + + def __init__(self, in_size, batch_norm=None, name=""): + super().__init__() + self.add_module(name + "bn", batch_norm(in_size)) + + nn.init.constant_(self[0].weight, 1.0) + nn.init.constant_(self[0].bias, 0) + + +class BatchNorm1d(_BNBase): + + def __init__(self, in_size: int, *, name: str = ""): + super().__init__(in_size, batch_norm=nn.BatchNorm1d, name=name) + + +class BatchNorm2d(_BNBase): + + def __init__(self, in_size: int, name: str = ""): + super().__init__(in_size, batch_norm=nn.BatchNorm2d, name=name) + + +class BatchNorm3d(_BNBase): + + def __init__(self, in_size: int, name: str = ""): + super().__init__(in_size, batch_norm=nn.BatchNorm3d, name=name) + + +class _ConvBase(nn.Sequential): + + def __init__( + self, + in_size, + out_size, + kernel_size, + stride, + padding, + activation, + bn, + init, + conv=None, + batch_norm=None, + bias=True, + preact=False, + name="" + ): + super().__init__() + + bias = bias and (not bn) + conv_unit = conv( + in_size, + out_size, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=bias + ) + init(conv_unit.weight) + if bias: + nn.init.constant_(conv_unit.bias, 0) + + if bn: + if not preact: + bn_unit = batch_norm(out_size) + else: + bn_unit = batch_norm(in_size) + + if preact: + if bn: + self.add_module(name + 'bn', bn_unit) + + if activation is not None: + self.add_module(name + 'activation', activation) + + self.add_module(name + 'conv', conv_unit) + + if not preact: + if bn: + self.add_module(name + 'bn', bn_unit) + + if activation is not None: + self.add_module(name + 'activation', activation) + + +class Conv1d(_ConvBase): + + def __init__( + self, + in_size: int, + out_size: int, + *, + kernel_size: int = 1, + stride: int = 1, + padding: int = 0, + activation=nn.ReLU(inplace=True), + bn: bool = False, + init=nn.init.kaiming_normal_, + bias: bool = True, + preact: bool = False, + name: str = "" + ): + super().__init__( + in_size, + out_size, + kernel_size, + stride, + padding, + activation, + bn, + init, + conv=nn.Conv1d, + batch_norm=BatchNorm1d, + bias=bias, + preact=preact, + name=name + ) + + +class Conv2d(_ConvBase): + + def __init__( + self, + in_size: int, + out_size: int, + *, + kernel_size: Tuple[int, int] = (1, 1), + stride: Tuple[int, int] = (1, 1), + padding: Tuple[int, int] = (0, 0), + activation=nn.ReLU(inplace=True), + bn: bool = False, + init=nn.init.kaiming_normal_, + bias: bool = True, + preact: bool = False, + name: str = "" + ): + super().__init__( + in_size, + out_size, + kernel_size, + stride, + padding, + activation, + bn, + init, + conv=nn.Conv2d, + batch_norm=BatchNorm2d, + bias=bias, + preact=preact, + name=name + ) + + +class Conv3d(_ConvBase): + + def __init__( + self, + in_size: int, + out_size: int, + *, + kernel_size: Tuple[int, int, int] = (1, 1, 1), + stride: Tuple[int, int, int] = (1, 1, 1), + padding: Tuple[int, int, int] = (0, 0, 0), + activation=nn.ReLU(inplace=True), + bn: bool = False, + init=nn.init.kaiming_normal_, + bias: bool = True, + preact: bool = False, + name: str = "" + ): + super().__init__( + in_size, + out_size, + kernel_size, + stride, + padding, + activation, + bn, + init, + conv=nn.Conv3d, + batch_norm=BatchNorm3d, + bias=bias, + preact=preact, + name=name + ) + + +class FC(nn.Sequential): + + def __init__( + self, + in_size: int, + out_size: int, + *, + activation=nn.ReLU(inplace=True), + bn: bool = False, + init=None, + preact: bool = False, + name: str = "" + ): + super().__init__() + + fc = nn.Linear(in_size, out_size, bias=not bn) + if init is not None: + init(fc.weight) + if not bn: + nn.init.constant_(fc.bias, 0) + + if preact: + if bn: + self.add_module(name + 'bn', BatchNorm1d(in_size)) + + if activation is not None: + self.add_module(name + 'activation', activation) + + self.add_module(name + 'fc', fc) + + if not preact: + if bn: + self.add_module(name + 'bn', BatchNorm1d(out_size)) + + if activation is not None: + self.add_module(name + 'activation', activation) + +def set_bn_momentum_default(bn_momentum): + + def fn(m): + if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)): + m.momentum = bn_momentum + + return fn + + +class BNMomentumScheduler(object): + + def __init__( + self, model, bn_lambda, last_epoch=-1, + setter=set_bn_momentum_default + ): + if not isinstance(model, nn.Module): + raise RuntimeError( + "Class '{}' is not a PyTorch nn Module".format( + type(model).__name__ + ) + ) + + self.model = model + self.setter = setter + self.lmbd = bn_lambda + + self.step(last_epoch + 1) + self.last_epoch = last_epoch + + def step(self, epoch=None): + if epoch is None: + epoch = self.last_epoch + 1 + + self.last_epoch = epoch + self.model.apply(self.setter(self.lmbd(epoch))) + diff --git a/models/model/scene_model.py b/models/model/scene_model.py index 8edda9d..5ca2ee0 100644 --- a/models/model/scene_model.py +++ b/models/model/scene_model.py @@ -4,6 +4,7 @@ from utils.registry import Registry from models.model.pointtransformer import pointtransformer_enc_repro from models.model.pointnet import pointnet_enc_repro +from models.model.pointnet2.pointnet2_semseg import pointnet2_enc_repro SCENEMODEL = Registry('SceneModel') @@ -26,3 +27,7 @@ def PointTransformer(**kwargs): @SCENEMODEL.register() def PointNet(**kwargs): return pointnet_enc_repro(**kwargs) + +@SCENEMODEL.register() +def PointNet2(**kwargs): + return pointnet2_enc_repro(**kwargs) \ No newline at end of file diff --git a/models/model/unet.py b/models/model/unet.py index c62a93d..9e8fc7a 100644 --- a/models/model/unet.py +++ b/models/model/unet.py @@ -38,7 +38,8 @@ def __init__(self, cfg: DictConfig, slurm: bool, *args, **kwargs) -> None: self.scene_model = create_scene_model(cfg.scene_model.name, **scene_model_args) ## load pretrained weights weight_path = cfg.scene_model.pretrained_weights_slurm if slurm else cfg.scene_model.pretrained_weights - self.scene_model.load_pretrained_weight(weigth_path=weight_path) + if weight_path is not None: + self.scene_model.load_pretrained_weight(weigth_path=weight_path) if cfg.freeze_scene_model: for p in self.scene_model.parameters(): p.requires_grad_(False) @@ -143,6 +144,11 @@ def condition(self, data: Dict) -> torch.Tensor: b = data['pos'].shape[0] pos = data['pos'].to(torch.float32) scene_feat = self.scene_model(pos).reshape(b, self.scene_model.num_groups, -1) + elif self.scene_model_name == 'PointNet2': + b = data['pos'].shape[0] + pos = data['pos'].to(torch.float32) + _, scene_feat_list = self.scene_model(pos) + scene_feat = scene_feat_list[-1].transpose(1, 2) else: raise Exception('Unexcepted scene model.') diff --git a/models/optimizer/grasp_with_object.py b/models/optimizer/grasp_with_object.py index 45e9ed1..2103c7c 100644 --- a/models/optimizer/grasp_with_object.py +++ b/models/optimizer/grasp_with_object.py @@ -1,14 +1,32 @@ from typing import Dict +import os +import numpy as np import torch from omegaconf import DictConfig from utils.handmodel import get_handmodel from models.optimizer.optimizer import Optimizer from models.base import OPTIMIZER - +import pickle +import torch.functional as F @OPTIMIZER.register() class GraspWithObject(Optimizer): + _BATCH_SIZE = 16 + _N_OBJ = 4096 + _joint_angle_lower = torch.tensor([-0.5235988, -0.7853982, -0.43633232, 0., 0., 0., -0.43633232, 0., 0., 0., + -0.43633232, 0., 0., 0., 0., -0.43633232, 0., 0., 0., -1.047, 0., -0.2618, + -0.5237, 0.], device='cuda') + _joint_angle_upper = torch.tensor([0.17453292, 0.61086524, 0.43633232, 1.5707964, 1.5707964, 1.5707964, 0.43633232, + 1.5707964, 1.5707964, 1.5707964, 0.43633232, 1.5707964, 1.5707964, 1.5707964, + 0.6981317, 0.43633232, 1.5707964, 1.5707964, 1.5707964, 1.047, 1.309, 0.2618, + 0.5237, 1.], device='cuda') + + _global_trans_lower = torch.tensor([-0.13128923, -0.10665303, -0.45753425], device='cuda') + _global_trans_upper = torch.tensor([0.12772022, 0.22954416, -0.21764427], device='cuda') + + _NORMALIZE_LOWER = -1. + _NORMALIZE_UPPER = 1. def __init__(self, cfg: DictConfig, slurm: bool, *args, **kwargs) -> None: if 'device' in kwargs: self.device = kwargs['device'] @@ -16,36 +34,71 @@ def __init__(self, cfg: DictConfig, slurm: bool, *args, **kwargs) -> None: self.device = 'cpu' self.slurm = slurm self.scale = cfg.scale - self.palm_alignment = cfg.palm_alignment - self.palm_alignment_weight = cfg.palm_alignment_weight + self.collision = cfg.collision + self.collision_weight = cfg.collision_weight self.clip_grad_by_value = cfg.clip_grad_by_value self.modeling_keys = cfg.modeling_keys - self.batch_size = cfg.batch_size - self.hand_model = get_handmodel(batch_size=self.batch_size, device=self.device) + self.normalize_x = cfg.normalize_x + self.normalize_x_trans = cfg.normalize_x_trans + + self.asset_dir = cfg.asset_dir_slrum if self.slurm else cfg.asset_dir + self.obj_pcds_nors_dict = pickle.load(open(os.path.join(self.asset_dir, 'object_pcds_nors.pkl'), 'rb')) + self.hand_model = get_handmodel(batch_size=self._BATCH_SIZE, device=self.device) - def optimize(self, x: torch.Tensor, data: Dict) -> torch.Tensor: + self.relu = torch.nn.ReLU() + + def optimize(self, x: torch.Tensor, data: Dict, t: int) -> torch.Tensor: """ Compute gradient for optimizer constraint Args: x: the denosied signal at current step, which is detached and is required grad data: data dict that provides original data + t: sample time Return: The optimizer objective value of current step """ loss = 0. + self.hand_model.update_kinematics(q=x) + hand_pcd = self.hand_model.get_surface_points() + n_hand = hand_pcd.shape[1] + n_obj = self._N_OBJ + obj_pcd_nor_list = [] + for object_name in data['scene_id'][t*self._BATCH_SIZE:(t+1)*self._BATCH_SIZE]: + obj_pcd_nor_list.append(self.obj_pcds_nors_dict[object_name][:self._N_OBJ, :]) + obj_pcd_nor = np.stack(obj_pcd_nor_list, axis=0) + obj_pcd_nor = torch.tensor(obj_pcd_nor, device='cuda') + ## compute palm alignment optimization - if self.palm_alignment: - # palm_alignment = 1. - toward_of_palm(qpos) (dot prodoct) toward_of_object(qpos) - # loss += self.palm_alignment_weight * palm_alignment - pass + if self.collision: + obj_pcd = obj_pcd_nor[..., :3] + obj_nor = obj_pcd_nor[..., 3:6] + # batch the obj pcd and hand pcd + batch_obj_pcd = obj_pcd_nor[:, :, :3].view(self._BATCH_SIZE, 1, n_obj, 3).repeat(1, n_hand, 1, 1) + batch_hand_pcd = hand_pcd.view(self._BATCH_SIZE, n_hand, 1, 3).repeat(1, 1, n_obj, 1) + # compute the pair wise dist + hand_obj_dist = (batch_obj_pcd - batch_hand_pcd).norm(dim=3) + hand_obj_dist, hand_obj_indices = hand_obj_dist.min(dim=2) + # gather the obj points and normals w.r.t. hand points + hand_obj_points = torch.stack([obj_pcd[i, x, :] for i, x in enumerate(hand_obj_indices)], dim=0) + hand_obj_normals = torch.stack([obj_nor[i, x, :] for i, x in enumerate(hand_obj_indices)], dim=0) + # compute the signs + hand_obj_signs = ((hand_obj_points - hand_pcd) * hand_obj_normals).sum(dim=2) + hand_obj_signs = (hand_obj_signs > 0.).float() + # signs dot dist to compute collision value + # collision_value = (hand_obj_signs * hand_obj_dist).max(dim=1).values + collision_value = (hand_obj_signs * hand_obj_dist).sum(dim=1) + # collision_value = self.relu(collision_value - 0.1) + # collision_value = torch.abs(collision_value - 0.005) + loss += self.collision_weight * collision_value.mean() - raise NotImplementedError + return (-1.0) * loss - def gradient(self, x: torch.Tensor, data: Dict) -> torch.Tensor: + def gradient(self, x: torch.Tensor, data: Dict, variance: torch.Tensor) -> torch.Tensor: + # print(f'compute gradient...') """ Compute gradient for optimizer constraint Args: @@ -55,10 +108,50 @@ def gradient(self, x: torch.Tensor, data: Dict) -> torch.Tensor: Return: Commputed gradient """ + assert (x.shape[0] % self._BATCH_SIZE == 0) with torch.enable_grad(): - x_in = x.detach().requires_grad_(True) - obj = self.optimize(x_in, data) - grad = torch.autograd.grad(obj, x_in)[0] + # concatenate the id rot to x_in + id_6d_rot = torch.tensor([1., 0., 0., 0., 1., 0.], device='cuda').view(1, 6).repeat(x.shape[0], 1) + x = torch.cat([x[:, :3], id_6d_rot, x[:, 3:]], dim=-1) + x_in = x.detach().requires_grad_(False) + grad_list = [] + obj_list = [] + for i in range(x.shape[0] // self._BATCH_SIZE): + i_x_in = x_in[i*self._BATCH_SIZE:(i+1)*self._BATCH_SIZE, :].detach().requires_grad_(True) + if self.normalize_x_trans: + i_x_in_denorm_trans = self.trans_denormalize(i_x_in[:, :3]) + else: + i_x_in_denorm_trans = i_x_in[:, :3] + if self.normalize_x: + i_x_in_denorm_angle = self.angle_denormalize(i_x_in[:, 9:]) + else: + i_x_in_denorm_angle = i_x_in[:, 9:] + i_x_in_denorm = torch.cat([i_x_in_denorm_trans, i_x_in[:, 3:9], i_x_in_denorm_angle], dim=-1) + obj = self.optimize(i_x_in_denorm, data, t=i) + i_grad = torch.autograd.grad(obj, i_x_in)[0] + obj_list.append(obj.abs().mean().detach().cpu()) + grad_list.append(i_grad) + # print(f'loss: {np.mean(obj_list)}') + grad = torch.cat(grad_list, dim=0) ## clip gradient by value + # print(f'grad norm: {grad.abs().mean()}') + grad = grad * self.scale grad = torch.clip(grad, **self.clip_grad_by_value) - raise NotImplementedError + # grad = torch.cat([grad[:, :3], grad[:, 9:]], dim=-1) + # grad = torch.cat([torch.zeros_like(grad[:, :3], device=self.device), grad[:, 9:]], dim=-1) + grad = torch.cat([torch.zeros_like(grad[:, :3], device=self.device), + torch.zeros_like(grad[:, 9:11], device=self.device), + grad[:, 11:]], dim=-1) + return grad + + def angle_denormalize(self, joint_angle: torch.Tensor): + joint_angle_denorm = joint_angle + (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) / 2 + joint_angle_denorm /= (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) + joint_angle_denorm = joint_angle_denorm * (self._joint_angle_upper - self._joint_angle_lower) + self._joint_angle_lower + return joint_angle_denorm + + def trans_denormalize(self, global_trans: torch.Tensor): + global_trans_denorm = global_trans + (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) / 2 + global_trans_denorm /= (self._NORMALIZE_UPPER - self._NORMALIZE_LOWER) + global_trans_denorm = global_trans_denorm * (self._global_trans_upper - self._global_trans_lower) + self._global_trans_lower + return global_trans_denorm \ No newline at end of file diff --git a/models/planner/fk2planning.py b/models/planner/fk2planning.py new file mode 100644 index 0000000..7969c43 --- /dev/null +++ b/models/planner/fk2planning.py @@ -0,0 +1,69 @@ +from typing import Dict +import torch +import torch.nn.functional as F +import numpy as np +from omegaconf import DictConfig + +from models.optimizer.optimizer import Optimizer +from models.base import PLANNER + +@PLANNER.register() +class GreedyFK2Planner(Optimizer): + + def __init__(self, cfg: DictConfig, *args, **kwargs) -> None: + self.scale = cfg.scale + + self.greedy_type = cfg.greedy_type + + self.clip_grad_by_value = cfg.clip_grad_by_value + + def objective(self, x: torch.Tensor, data: Dict): + """ Compute gradient for planner guidance + + Args: + x: the denosied signal at current step, which is detached and is required grad + data: data dict that provides original data + + Return: + The optimizer objective value of current step + """ + loss = 0. + + ## important!!! scale x with region_area + ## convert x to reality psoition + target = data['target'] # + + if self.greedy_type == 'last_frame': + loss += F.l1_loss(x[:, -1, :], target, reduction='mean') + elif self.greedy_type == 'all_frame': + loss += F.l1_loss(x, target.unsqueeze(1), reduction='mean') + elif self.greedy_type == 'all_frame_exp': + traj_dist = torch.norm(x - target.unsqueeze(1), dim=-1, p=1) + loss += (-1.0) * torch.exp(1 / traj_dist.clamp(min=0.01)).sum() + else: + raise Exception('Unsupported greedy type') + + return (-1.0) * loss + + def gradient(self, x: torch.Tensor, data: Dict, variance: torch.Tensor) -> torch.Tensor: + """ Compute gradient for planner guidance + Args: + x: the denosied signal at current step + data: data dict that provides original data + + Return: + Commputed gradient + """ + with torch.enable_grad(): + x_in = x.detach().requires_grad_(True) + obj = self.objective(x_in, data) + grad = torch.autograd.grad(obj, x_in)[0] + + # print(f'obj: {-obj.detach().cpu()}') + + ## clip gradient by value + grad = grad * self.scale + grad = torch.clip(grad, **self.clip_grad_by_value) + ## TODO clip gradient by norm + + return grad \ No newline at end of file diff --git a/models/visualizer.py b/models/visualizer.py index f266bea..56ae62a 100644 --- a/models/visualizer.py +++ b/models/visualizer.py @@ -14,6 +14,8 @@ from utils.visualize import create_trimesh_nodes_path, create_trimesh_node from utils.handmodel import get_handmodel from utils.plotly_utils import plot_mesh +from utils.rot6d import rot_to_orthod6d, robust_compute_rotation_matrix_from_ortho6d, random_rot +from tqdm import tqdm VISUALIZER = Registry('Visualizer') @@ -404,6 +406,91 @@ def visualize( if cnt >= vis_cnt: break +@VISUALIZER.register() +@torch.no_grad() +class GraspGenURVisualizer(): + def __init__(self, cfg: DictConfig) -> None: + """ Visual evaluation class for pose generation task. + Args: + cfg: visuzalizer configuration + """ + self.ksample = cfg.ksample + self.hand_model = get_handmodel(batch_size=1, device='cuda') + + def visualize( + self, + model: torch.nn.Module, + dataloader: torch.utils.data.DataLoader, + save_dir: str + ) -> None: + """ Visualize method + Args: + model: diffusion model + dataloader: test dataloader + save_dir: save directory of rendering images + """ + model.eval() + device = model.device + + os.makedirs(save_dir, exist_ok=True) + os.makedirs(os.path.join(save_dir, 'html'), exist_ok=True) + # save + pbar = tqdm(total=len(dataloader.dataset.split) * self.ksample) + object_pcds_dict = dataloader.dataset.scene_pcds + res = {'method': 'diffuser@w/o-opt', + 'desc': 'w/o optimizer grasp pose generation', + 'sample_qpos': {}} + + # for syn training dataset + # train_obj_list_ds = ["contactdb+alarm_clock", "contactdb+banana", "contactdb+binoculars", + # "contactdb+cube_medium", "contactdb+mouse", "contactdb+piggy_bank", + # "contactdb+ps_controller", "contactdb+stapler", "contactdb+train", + # "ycb+toy_airplane"] + + for object_name in dataloader.dataset._test_split: + # for object_name in train_obj_list_ds: + obj_pcd_can = torch.tensor(object_pcds_dict[object_name], device=device).unsqueeze(0).repeat(self.ksample, 1, 1) + obj_pcd_can = obj_pcd_can[:, :dataloader.dataset.num_points, :] + i_rot_list = [] + for k_rot in range(self.ksample): + i_rot_list.append(random_rot(device)) + i_rot = torch.stack(i_rot_list).to(torch.float64) + obj_pcd_rot = torch.matmul(i_rot, obj_pcd_can.transpose(1, 2)).transpose(1, 2) + + # construct data + data = {'x': torch.randn(self.ksample, 27, device=device), + 'pos': obj_pcd_rot.to(device), + 'scene_rot_mat': i_rot, + 'scene_id': [object_name for i in range(self.ksample)], + 'cam_trans': [None for i in range(self.ksample)]} + outputs = model.sample(data, k=1).squeeze(1)[:, -1, :].to(torch.float64) + ## denormalization + if dataloader.dataset.normalize_x: + outputs[:, 3:] = dataloader.dataset.angle_denormalize(joint_angle=outputs[:, 3:].cpu()).cuda() + if dataloader.dataset.normalize_x_trans: + outputs[:, :3] = dataloader.dataset.trans_denormalize(global_trans=outputs[:, :3].cpu()).cuda() + + id_6d_rot = torch.tensor([1., 0., 0., 0., 1., 0.], device=device).view(1, 6).repeat(self.ksample, 1).to(torch.float64) + outputs_3d_rot = rot_to_orthod6d(torch.bmm(i_rot, robust_compute_rotation_matrix_from_ortho6d(id_6d_rot))) + outputs[:, :3] = torch.bmm(i_rot, outputs[:, :3].unsqueeze(-1)).squeeze(-1) + outputs = torch.cat([outputs[:, :3], outputs_3d_rot, outputs[:, 3:]], dim=-1) + + # visualization for checking + scene_id = data['scene_id'][0] + scene_dataset, scene_object = scene_id.split('+') + mesh_path = os.path.join('assets/object', scene_dataset, scene_object, f'{scene_object}.stl') + obj_mesh = trimesh.load(mesh_path) + for i in range(outputs.shape[0]): + self.hand_model.update_kinematics(q=outputs[i:i+1, :]) + vis_data = [plot_mesh(obj_mesh, color='lightblue')] + vis_data += self.hand_model.get_plotly_data(opacity=1.0, color='pink') + save_path = os.path.join(save_dir, 'html', f'{scene_id}+sample-{i}.html') + fig = go.Figure(data=vis_data) + fig.write_html(save_path) + pbar.update(1) + res['sample_qpos'][object_name] = np.array(outputs.cpu().detach()) + pickle.dump(res, open(os.path.join(save_dir, 'res_diffuser.pkl'), 'wb')) + @VISUALIZER.register() @torch.no_grad() class PoseGenVisualizerHF(): diff --git a/plan.py b/plan.py index b1a58be..e9fe5fc 100644 --- a/plan.py +++ b/plan.py @@ -1,3 +1,5 @@ +from isaacgym import gymapi, gymutil, gymtorch + import os import hydra import torch diff --git a/scripts/franka_planning/plan.sh b/scripts/franka_planning/plan.sh new file mode 100644 index 0000000..c66ac15 --- /dev/null +++ b/scripts/franka_planning/plan.sh @@ -0,0 +1,12 @@ +CKPT=$1 + +python plan.py hydra/job_logging=none hydra/hydra_logging=none \ + exp_dir=${CKPT} \ + diffuser=ddpm \ + diffuser.steps=30 \ + model=unet_fk2 \ + model.use_position_embedding=true \ + task=franka_planning \ + task.dataset.normalize_x=true \ + planner=greedy_fk2_planning \ + planner.scale=0.2 \ No newline at end of file diff --git a/scripts/franka_planning/train.sh b/scripts/franka_planning/train.sh new file mode 100644 index 0000000..fb60a87 --- /dev/null +++ b/scripts/franka_planning/train.sh @@ -0,0 +1,11 @@ +EXP_NAME=$1 + +python train.py hydra/job_logging=none hydra/hydra_logging=none \ + exp_name=${EXP_NAME} \ + diffuser=ddpm \ + diffuser.loss_type=l1 \ + diffuser.steps=30 \ + model=unet_fk2 \ + model.use_position_embedding=true \ + task=franka_planning \ + task.dataset.normalize_x=true \ No newline at end of file diff --git a/scripts/franka_planning/train_ddm.sh b/scripts/franka_planning/train_ddm.sh new file mode 100644 index 0000000..7061309 --- /dev/null +++ b/scripts/franka_planning/train_ddm.sh @@ -0,0 +1,12 @@ +EXP_NAME=$1 + +CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 --use_env train_ddm.py \ + hydra/job_logging=none hydra/hydra_logging=none \ + exp_name=${EXP_NAME} \ + diffuser=ddpm \ + diffuser.loss_type=l1 \ + diffuser.steps=30 \ + model=unet_fk2 \ + model.use_position_embedding=true \ + task=franka_planning \ + task.dataset.normalize_x=true \ No newline at end of file diff --git a/scripts/grasp_gen/sample.sh b/scripts/grasp_gen/sample.sh deleted file mode 100755 index d1a7138..0000000 --- a/scripts/grasp_gen/sample.sh +++ /dev/null @@ -1,8 +0,0 @@ -CKPT=$1 - -python sample.py hydra/job_logging=none hydra/hydra_logging=none \ - exp_dir=${CKPT} \ - diffuser=ddpm_l1 \ - model=unet_grasp \ - task=grasp_gen \ - task.visualizer.ksample=10 diff --git a/scripts/grasp_gen/train.sh b/scripts/grasp_gen/train.sh deleted file mode 100755 index 2ccdfec..0000000 --- a/scripts/grasp_gen/train.sh +++ /dev/null @@ -1,7 +0,0 @@ -EXP_NAME=$1 - -python train.py hydra/job_logging=none hydra/hydra_logging=none \ - exp_name=${EXP_NAME} \ - diffuser=ddpm \ - model=unet_grasp \ - task=grasp_gen diff --git a/scripts/grasp_gen/train.slurm b/scripts/grasp_gen/train.slurm deleted file mode 100755 index 9275e57..0000000 --- a/scripts/grasp_gen/train.slurm +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash -#SBATCH -J gpg_diffuser -#SBATCH --comment "3D-conditioned diffusion" - -#SBATCH --qos singlegpu -#SBATCH -p gpu -#SBATCH -N 1 -#SBATCH --gpus=1 -#SBATCH --cpus-per-task=8 -#SBATCH --ntasks=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --time=3-00:00:00 - -#SBATCH -o "/home/lipuhao/dev/SceneDiffuser/3d_diffusion/outputs/slurm/%x.%j.out" -#SBATCH -e "/home/lipuhao/dev/SceneDiffuser/3d_diffusion/outputs/slurm/%x.%j.err" - -EXP_NAME=$1 -BATCH_SIZE=$2 -LR=$3 - -SLURM=1 python train.py hydra/job_logging=none hydra/hydra_logging=none \ - output_dir=outputs/exps/ \ - exp_name=${EXP_NAME} \ - diffuser=ddpm_l1 \ - model=unet_grasp \ - task=grasp_gen \ - task.visualizer.visualize=false \ - task.train.batch_size=${BATCH_SIZE} \ - task.lr=${LR} \ No newline at end of file diff --git a/scripts/grasp_gen_ur/sample.sh b/scripts/grasp_gen_ur/sample.sh new file mode 100755 index 0000000..e954d8b --- /dev/null +++ b/scripts/grasp_gen_ur/sample.sh @@ -0,0 +1,35 @@ +CKPT=$1 +OPT=$2 + +if [ -z ${CKPT} ] +then + echo "No ckpt input." + exit +fi + +if [ -z ${OPT} ] || [ ${OPT} != "OPT" ] +then + echo "Without optimizer guidance." + python sample.py hydra/job_logging=none hydra/hydra_logging=none \ + exp_dir=${CKPT} \ + diffuser=ddpm \ + diffuser.loss_type=l1 \ + diffuser.steps=100 \ + model=unet_grasp \ + task=grasp_gen_ur \ + task.dataset.normalize_x=true \ + task.dataset.normalize_x_trans=true +else + echo "With optimizer guidance." + python sample.py hydra/job_logging=none hydra/hydra_logging=none \ + exp_dir=${CKPT} \ + diffuser=ddpm \ + diffuser.loss_type=l1 \ + diffuser.steps=100 \ + model=unet_grasp \ + task=grasp_gen_ur \ + task.dataset.normalize_x=true \ + task.dataset.normalize_x_trans=true \ + optimizer=grasp_with_object \ + optimizer.scale=0.1 +fi \ No newline at end of file diff --git a/scripts/grasp_gen_ur/test.sh b/scripts/grasp_gen_ur/test.sh new file mode 100755 index 0000000..e69de29 diff --git a/scripts/grasp_gen_ur/train.sh b/scripts/grasp_gen_ur/train.sh new file mode 100755 index 0000000..0a93edf --- /dev/null +++ b/scripts/grasp_gen_ur/train.sh @@ -0,0 +1,11 @@ +EXP_NAME=$1 + +python train.py hydra/job_logging=none hydra/hydra_logging=none \ + exp_name=${EXP_NAME} \ + diffuser=ddpm \ + diffuser.loss_type=l1 \ + diffuser.steps=100 \ + model=unet_grasp \ + task=grasp_gen_ur \ + task.dataset.normalize_x=true \ + task.dataset.normalize_x_trans=true \ No newline at end of file diff --git a/utils/misc.py b/utils/misc.py index 46bc55c..9ba60e3 100644 --- a/utils/misc.py +++ b/utils/misc.py @@ -31,6 +31,10 @@ def compute_model_dim(cfg: DictConfig) -> int: return 2 elif cfg.name == 'grasp_gen': return 3 + 6 + 24 + elif cfg.name == 'grasp_gen_ur': + return 3 + 24 + elif cfg.name == 'franka_planning': + return 7 else: raise Exception('Unsupported task.') diff --git a/utils/rot6d.py b/utils/rot6d.py index a11ab2e..5f60b1f 100644 --- a/utils/rot6d.py +++ b/utils/rot6d.py @@ -3,6 +3,19 @@ import transforms3d +def random_rot(device='cuda'): + rot_angles = np.random.random(3) * np.pi * 2 + theta_x, theta_y, theta_z = rot_angles[0], rot_angles[1], rot_angles[2] + Rx = torch.tensor([[1, 0, 0], [0, np.cos(theta_x), -np.sin(theta_x)], [0, np.sin(theta_x), np.cos(theta_x)]]).to(device) + Ry = torch.tensor([[np.cos(theta_y), 0, np.sin(theta_y)], [0, 1, 0], [-np.sin(theta_y), 0, np.cos(theta_y)]]).to(device) + Rz = torch.tensor([[np.cos(theta_z), -np.sin(theta_z), 0], [np.sin(theta_z), np.cos(theta_z), 0], [0, 0, 1]]).to(device) + return (Rx @ Ry @ Rz).clone().detach() # [3, 3] + + +def rot_to_orthod6d(rot): + return rot.transpose(1, 2)[:, :2].reshape([-1, 6]) + + def get_rot6d_from_rot3d(rot3d): global_rotation = np.array(transforms3d.euler.euler2mat(rot3d[0], rot3d[1], rot3d[2])) return global_rotation.T.reshape(9)[:6] @@ -76,4 +89,4 @@ def cross_product(u, v): out = torch.cat((i.view(batch, 1), j.view(batch, 1), k.view(batch, 1)), 1) - return out + return out \ No newline at end of file From 0513d0e7d065f093b5cc5d05f24c0dd7433be130 Mon Sep 17 00:00:00 2001 From: Xiaoyao-Li <1617410526@qq.com> Date: Fri, 21 Apr 2023 11:38:14 +0800 Subject: [PATCH 2/7] update git ignore file --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index e358e11..318702d 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,5 @@ assets/object/ assets/urdf/ envs/assets/franka_description/ -envs/assets/scene_description/ \ No newline at end of file +envs/assets/scene_description/ +envs/assets/movable_hand_urdf/ \ No newline at end of file From 29f85cccd1e61ad865cec24a1cb5ff48480ed5da Mon Sep 17 00:00:00 2001 From: Xiaoyao-Li <1617410526@qq.com> Date: Fri, 21 Apr 2023 11:39:57 +0800 Subject: [PATCH 3/7] update code for grasp pose evaluation --- envs/__init__.py | 0 envs/tasks/__init__.py | 0 envs/tasks/base_task_grasp.py | 230 ++++++ envs/tasks/grasp_test_force.yaml | 92 +++ envs/tasks/grasp_test_force_shadowhand.py | 867 ++++++++++++++++++++++ envs/tasks/utils/angle.py | 255 +++++++ 6 files changed, 1444 insertions(+) create mode 100644 envs/__init__.py create mode 100644 envs/tasks/__init__.py create mode 100644 envs/tasks/base_task_grasp.py create mode 100644 envs/tasks/grasp_test_force.yaml create mode 100644 envs/tasks/grasp_test_force_shadowhand.py create mode 100644 envs/tasks/utils/angle.py diff --git a/envs/__init__.py b/envs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/envs/tasks/__init__.py b/envs/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/envs/tasks/base_task_grasp.py b/envs/tasks/base_task_grasp.py new file mode 100644 index 0000000..2f67b27 --- /dev/null +++ b/envs/tasks/base_task_grasp.py @@ -0,0 +1,230 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. + +import sys +import os +import operator +from copy import deepcopy +import random + +from isaacgym import gymapi +from isaacgym.gymutil import get_property_setter_map, get_property_getter_map, get_default_setter_args, apply_random_samples, check_buckets, generate_random_samples +from isaacgym import gymtorch +import numpy as np +import torch +import ipdb + + +# Base class for RL tasks +class BaseTask(): + def __init__(self, cfg, enable_camera_sensors=False, + cam_pos=(10, 10, 3.0), cam_target=(0, 0, 0.0)): + self.gym = gymapi.acquire_gym() + + self.device_type = cfg.get("self.device_type", "cuda") + self.device_id = cfg.get("self.device_id", 0) + + self.device = "cpu" + if self.device_type == "cuda" or self.device_type == "GPU": + self.device = "cuda" + ":" + str(self.device_id) + + self.headless = cfg["headless"] + + # double check! + self.graphics_device_id = self.device_id + if enable_camera_sensors == False and self.headless == True: + self.graphics_device_id = -1 + + self.num_envs = cfg["env"]["numTrain"] + self.num_obs = cfg["env"]["numObservations"] + self.num_states = cfg["env"].get("numStates", 0) + self.num_actions = cfg["env"]["numActions"] + + self.control_freq_inv = cfg["env"].get("controlFrequencyInv", 1) + + # optimization flags for pytorch JIT + torch._C._jit_set_profiling_mode(False) + torch._C._jit_set_profiling_executor(False) + + # allocate buffers + self.obs_buf = torch.zeros( + (self.num_envs, self.num_obs), device=self.device, dtype=torch.float) + self.states_buf = torch.zeros( + (self.num_envs, self.num_states), device=self.device, dtype=torch.float) + self.rew_buf = torch.zeros( + self.num_envs, device=self.device, dtype=torch.float) + self.reset_buf = torch.ones( + self.num_envs, device=self.device, dtype=torch.long) + self.progress_buf = torch.zeros( + self.num_envs, device=self.device, dtype=torch.long) + self.randomize_buf = torch.zeros( + self.num_envs, device=self.device, dtype=torch.long) + self.extras = {} + + self.original_props = {} + self.dr_randomizations = {} + self.first_randomization = True + self.actor_params_generator = None + self.extern_actor_params = {} + for env_id in range(self.num_envs): + self.extern_actor_params[env_id] = None + + self.last_step = -1 + self.last_rand_step = -1 + + # create envs, sim and viewer + + self.create_sim() + self.gym.prepare_sim(self.sim) + + # todo: read from config + self.enable_viewer_sync = True + self.viewer = None + self.need_update = False + + # if running with a viewer, set up keyboard shortcuts and camera + if self.headless == False: + # subscribe to keyboard shortcuts + self.viewer = self.gym.create_viewer( + self.sim, gymapi.CameraProperties()) + self.gym.subscribe_viewer_keyboard_event( + self.viewer, gymapi.KEY_ESCAPE, "QUIT") + self.gym.subscribe_viewer_keyboard_event( + self.viewer, gymapi.KEY_V, "toggle_viewer_sync") + + # set the camera position based on up axis + sim_params = self.gym.get_sim_params(self.sim) + if sim_params.up_axis == gymapi.UP_AXIS_Z: + cam_pos = gymapi.Vec3(cam_pos[0], cam_pos[1], cam_pos[2]) + cam_target = gymapi.Vec3(cam_target[0], cam_target[1], cam_target[2]) + else: + cam_pos = gymapi.Vec3(cam_pos[0], cam_pos[1], cam_pos[2]) + cam_target = gymapi.Vec3(cam_target[0], cam_target[1], cam_target[2]) + + self.gym.viewer_camera_look_at( + self.viewer, None, cam_pos, cam_target) + + + + # set gravity based on up axis and return axis index + def set_sim_params_up_axis(self, sim_params, axis): + if axis == 'z': + sim_params.up_axis = gymapi.UP_AXIS_Z + sim_params.gravity.x = 0 + sim_params.gravity.y = 0 + sim_params.gravity.z = -9.81 + return 2 + return 1 + + def create_sim(self, compute_device, graphics_device, physics_engine, sim_params): + sim = self.gym.create_sim(compute_device, graphics_device, physics_engine, sim_params) + if sim is None: + print("*** Failed to create sim") + quit() + + return sim + + def step(self, actions): + if self.dr_randomizations.get('actions', None): + actions = self.dr_randomizations['actions']['noise_lambda'](actions) + + # apply actions + self.pre_physics_step(actions) + + # step physics and render each frame + for _ in range(self.control_freq_inv): + self.render() + self.gym.simulate(self.sim) + + # to fix! + if self.device == 'cpu': + self.gym.fetch_results(self.sim, True) + + # compute observations, rewards, resets, ... + self.post_physics_step() + + if self.dr_randomizations.get('observations', None): + self.obs_buf = self.dr_randomizations['observations']['noise_lambda'](self.obs_buf) + + + + # def control_ik(self, dpose): + # j_eef, self.num_envs + # # solve damped least squares + # j_eef_T = torch.transpose(j_eef, 1, 2) + # lmbda = torch.eye(6, self.device=self.device) * (self.damping ** 2) + # u = (j_eef_T @ torch.inverse(j_eef @ j_eef_T + lmbda) @ dpose).view(self.num_envs, 7) + # return u + + def render(self, sync_frame_time=False): + if self.viewer: + # check for window closed + if self.gym.query_viewer_has_closed(self.viewer): + sys.exit() + + # check for keyboard events + for evt in self.gym.query_viewer_action_events(self.viewer): + if evt.action == "QUIT" and evt.value > 0: + sys.exit() + elif evt.action == "toggle_viewer_sync" and evt.value > 0: + self.enable_viewer_sync = not self.enable_viewer_sync + + # fetch results + if self.device != 'cpu': + self.gym.fetch_results(self.sim, True) + + # step graphics + if self.enable_viewer_sync: + self.gym.step_graphics(self.sim) + self.gym.draw_viewer(self.viewer, self.sim, True) + else: + self.gym.poll_viewer_events(self.viewer) + + def get_actor_params_info(self, dr_params, env): + """Returns a flat array of actor params, their names and ranges.""" + if "actor_params" not in dr_params: + return None + params = [] + names = [] + lows = [] + highs = [] + param_getters_map = get_property_getter_map(self.gym) + for actor, actor_properties in dr_params["actor_params"].items(): + handle = self.gym.find_actor_handle(env, actor) + for prop_name, prop_attrs in actor_properties.items(): + if prop_name == 'color': + continue # this is set randomly + props = param_getters_map[prop_name](env, handle) + if not isinstance(props, list): + props = [props] + for prop_idx, prop in enumerate(props): + for attr, attr_randomization_params in prop_attrs.items(): + name = prop_name+'_'+str(prop_idx)+'_'+attr + lo_hi = attr_randomization_params['range'] + distr = attr_randomization_params['distribution'] + if 'uniform' not in distr: + lo_hi = (-1.0*float('Inf'), float('Inf')) + if isinstance(prop, np.ndarray): + for attr_idx in range(prop[attr].shape[0]): + params.append(prop[attr][attr_idx]) + names.append(name+'_'+str(attr_idx)) + lows.append(lo_hi[0]) + highs.append(lo_hi[1]) + else: + params.append(getattr(prop, attr)) + names.append(name) + lows.append(lo_hi[0]) + highs.append(lo_hi[1]) + return params, names, lows, highs + + def pre_physics_step(self, actions): + raise NotImplementedError + + def post_physics_step(self): + raise NotImplementedError + + diff --git a/envs/tasks/grasp_test_force.yaml b/envs/tasks/grasp_test_force.yaml new file mode 100644 index 0000000..33cffb0 --- /dev/null +++ b/envs/tasks/grasp_test_force.yaml @@ -0,0 +1,92 @@ +task: + useStage: False + useSlider: False + useGroundTruth: False + +env: + env_name: "grasp" + numTrain: 6 + numObservations: 50 + numActions: 9 + envSpacing: 1.0 + maxEpisodeLength: 256 + actionSpeedScale: 20 + enableDebugVis: False + + contactBufferSize: 9999 + contactMovingThreshold: 0.1 + + pointFeatureDim: 512 + objPointDownsampleNum: 2048 + handPointDownsampleNum: 64 + pointNetLR: 0.0001 + visualizePointcloud: False + + enableCameraSensors: False + depth_bar: 10 + map_dis_bar: 0.1 + moving_pc_mode: False + + driveMode: "pos" + clipObservations: 5.0 + clipActions: 1.0 + + asset: + assetRoot: "envs/assets" + AssetNumTrain: 1 + trainObjAssets: + 0: + name: 0 + path: object_to_grasp/apple/apple.urdf + +sim: + substeps: 2 + physx: + num_threads: 4 + solver_type: 1 # 0: pgs, 1: tgs + num_position_iterations: 8 + num_velocity_iterations: 0 + contact_offset: 0.002 + rest_offset: 0.0 + bounce_threshold_velocity: 0.2 + max_depenetration_velocity: 1000.0 + default_buffer_size_multiplier: 5.0 + flex: + num_outer_iterations: 5 + num_inner_iterations: 20 + warm_start: 0.8 + relaxation: 0.75 + +object: + density: 1000 + damping: + linear: 10 + angular: 100 + shape: + friction: 2. + +agent: + density: 1000 + dof_props: + stiffness: 400.0 + velocity: 0.8 + damping: 400.0 + shape: + friction: 2. + +eval_policy: + init: + steps: 200 + dynamic: + directions: + 0: [1., 0., 0.] + 1: [-1., 0., 0.] + 2: [0., 1., 0.] + 3: [0., -1., 0.] + 4: [0., 0., 1.] + 5: [0., 0., -1.] + num_steps: 50 + magnitude_per_volume: 500. + error: + distance: 0.02 + diff --git a/envs/tasks/grasp_test_force_shadowhand.py b/envs/tasks/grasp_test_force_shadowhand.py new file mode 100644 index 0000000..e25e496 --- /dev/null +++ b/envs/tasks/grasp_test_force_shadowhand.py @@ -0,0 +1,867 @@ +import os +from requests import get +from isaacgym import gymutil, gymtorch, gymapi +from isaacgym.torch_utils import * +import numpy as np +from random import shuffle, randint +import yaml +from envs.tasks.base_task_grasp import BaseTask +from tqdm import tqdm +import ipdb +from scipy.spatial.transform import Rotation as R +import copy +import numpy as np +import math +import functools +from typing import Optional +from envs.tasks.utils.angle import r2euler +import torch +import transforms3d +from utils.handmodel import get_handmodel +import torch.nn.functional as F +import trimesh as tm + + +def quat_axis(q, axis=0): + """ ?? """ + basis_vec = torch.zeros(q.shape[0], 3, device=q.device) + basis_vec[:, axis] = 1 + return quat_rotate(q, basis_vec) + + +class IsaacGraspTestForce_shadowhand(BaseTask): + + def __init__(self, cfg, sim_params, physics_engine, device_type, device_id, headless, + init_opt_q: torch.tensor, object_name: str, object_volume: float, + agent_index=[[[0, 1, 2, 3, 4, 5]], [[0, 1, 2, 3, 4, 5]]], + fix_object=False): + self.gym = None + self.viewer = None + self.num_sim = 0 + self.cfg = cfg + self.sim_params = sim_params + self.device = device_type + self.init_opt_q = init_opt_q.clone() + self.init_q = self.q_transfer_o2s(init_opt_q) + self.object_name = object_name + self.object_volume = object_volume + self.physics_engine = physics_engine + self.agent_index = agent_index + self.up_axis = 'z' + self.cfg["device_type"] = device_type + self.cfg["device_id"] = device_id + self.cfg["headless"] = headless + self.device_type = device_type + self.device_id = device_id + self.headless = headless + self.fix_object = fix_object + self.device = "cuda" + if self.device_type == "cuda" or self.device_type == "GPU": + self.device = "cuda" + ":" + str(self.device_id) + + self.max_episode_length = self.cfg["env"]["maxEpisodeLength"] + + cfg["env"]["numTrain"] = init_opt_q.shape[0] + self.env_num_train = cfg["env"]["numTrain"] + self.env_num = self.env_num_train + self.asset_root = cfg["env"]["asset"]["assetRoot"] + self.num_train = cfg["env"]["asset"]["AssetNumTrain"] + self.tot_num = self.num_train + train_list_len = len(cfg["env"]["asset"]["trainObjAssets"]) + self.train_name_list = [] + self.exp_name = cfg['env']["env_name"] + + print("Simulator: number of objects", self.tot_num) + print("Simulator: number of environments", self.env_num) + if self.num_train: + assert (self.env_num_train % self.num_train == 0) + + # the number of used length must less than real length + assert (self.num_train <= train_list_len) + # the number of used length must less than real length + # each object should have equal number envs + assert (self.env_num % self.tot_num == 0) + self.env_per_object = self.env_num // self.tot_num + + for name in cfg["env"]["asset"]["trainObjAssets"]: + self.train_name_list.append( + cfg["env"]["asset"]["trainObjAssets"][name]["name"]) + + self.env_ptr_list = [] + self.obj_loaded = False + self.dexterous_loaded = False + self.hand_idxs = [] + + super().__init__(cfg=self.cfg, enable_camera_sensors=cfg["env"]["enableCameraSensors"], + cam_pos=(0.5, 0., 0.0), cam_target=(0, 0, 0.0)) + + # acquire tensors + self.root_tensor = gymtorch.wrap_tensor( + self.gym.acquire_actor_root_state_tensor(self.sim)).to(self.device) + self.dof_state_tensor = gymtorch.wrap_tensor( + self.gym.acquire_dof_state_tensor(self.sim)).to(self.device) + self.rigid_body_tensor = gymtorch.wrap_tensor( + self.gym.acquire_rigid_body_state_tensor(self.sim)).to(self.device).reshape(self.num_envs, -1, 13) + + self.gym.refresh_actor_root_state_tensor(self.sim) + self.gym.refresh_dof_state_tensor(self.sim) + self.gym.refresh_rigid_body_state_tensor(self.sim) + + self.root_tensor = self.root_tensor.view(self.num_envs, -1, 13) + self.dof_state_tensor = self.dof_state_tensor.view( + self.num_envs, -1, 2) + self.initial_dof_states = self.dof_state_tensor.clone() + self.initial_root_states = self.root_tensor.clone() + self.initial_obj_state = self.initial_root_states[:, 1, :3] + + self.dexterous_dof_tensor = self.dof_state_tensor[:, :self.dexterous_num_dofs, :] + + self.dexterous_root_tensor = self.root_tensor[:, 0, :] + self.object_root_tensor = self.root_tensor[:, 1, :] + + self.dof_dim = self.dexterous_num_dofs + self.pos_act = self.initial_dof_states[:, :self.dexterous_num_dofs, 0].clone() + self.eff_act = torch.zeros( + (self.num_envs, self.dof_dim), device=self.device) + # Real + # joint_name = self.gym.get_actor_joint_names(env_ptr, dexterous_actor) + # ['move_x', 'move_y', 'move_z', 'rot_r', 'rot_p', 'rot_y', + # 'WRJ2', 'WRJ1', + # 'FFJ4', 'FFJ3', 'FFJ2', 'FFJ1', + # 'LFJ5', 'LFJ4', 'LFJ3', 'LFJ2', 'LFJ1', + # 'MFJ4', 'MFJ3', 'MFJ2', 'MFJ1', + # 'RFJ4', 'RFJ3', 'RFJ2', 'RFJ1', + # 'THJ5', 'THJ4', 'THJ3', 'THJ2', 'THJ1'] + + # Puhao: + # WRJ2 WRJ1 + # FFJ4 FFJ3 FFJ2 FFJ1 + # MFJ4 MFJ3 MFJ2 MFJ1 + # RFJ4 RFJ3 RFJ2 RFJ1 + # LFJ5 LFJ4 LFJ3 LFJ2 LFJ1 + # THJ5 THJ4 THJ3 THJ2 THJ1 + # rigid props + + # ik parameters + self.damping = 0.05 + self.prepare_ik() + + # flags for switching between training and evaluation mode + # self.simulate() + + def __del__(self): + if self.gym is not None: + self.gym.destroy_sim(self.sim) + if self.viewer is not None: + self.gym.destroy_viewer(self.viewer) + del self.gym + + def _load_meta_data(self): + import warnings + warnings.filterwarnings("ignore") + metadata_basedir = self.cfg["metadata"]["basedir"] + tra_dir = os.path.join(metadata_basedir, 'tra_dir') + data_listdir = os.listdir(tra_dir) + data_listdir.sort(key=lambda x: int(x.split('-')[2].split('.pt')[0])) + self.metadata_dict = { + 'contactdb+apple': [], 'contactdb+camera': [], 'contactdb+cylinder_medium': [], + 'contactdb+door_knob': [], 'contactdb+flashlight': [], 'contactdb+stanford_bunny': [] + } + for data_path in data_listdir: + i_data = torch.load(os.path.join(tra_dir, data_path)) + i_q_tra = i_data['q_tra'] + i_energy = i_data['energy'] + i_object_name = i_data['object_name'] + i_sample = i_data['i_sample'] + # compute penetration + opt_q = i_q_tra[i_energy.min(0)[1], -1, :].unsqueeze(0).clone() + penetration = 0. + # penetration = compute_penetration(opt_q, i_object_name) + if penetration > self.cfg['metadata']['pen_th']: + print(f'remove data: {data_path} | pen: {penetration}') + continue + print(f'load data: {data_path} | pen: {penetration}') + # TODO: for debug + if type(penetration) != float: + penetration = penetration.cpu().item() + self.metadata_dict[i_object_name].append( + (self.q_transfer_o2s(opt_q).to(self.device), + i_energy.min(0)[0], + i_sample, + penetration)) + self.len_per_object = {x: len(self.metadata_dict[x]) for x in self.metadata_dict.keys()} + print('#######################data size#########################') + print(self.len_per_object) + print('#######################data size#########################') + num_envs = self._get_num_envs() + self.sim_q_final_dict = {} + for object_name in self.metadata_dict.keys(): + self.sim_q_final_dict[object_name] = [self.metadata_dict[object_name][x][0] for x in range(self.len_per_object[object_name])] + zero_pose = torch.zeros(self.metadata_dict[object_name][0][0]) + zero_pose[:, :3] = 1. + self.sim_q_final_dict[object_name] += [zero_pose.clone() for i in range(num_envs - self.len_per_object[object_name])] + self.sim_q_final_dict[object_name] = torch.cat(self.sim_q_final_dict[object_name], dim=0) + + def _get_num_envs(self): + return max(self.len_per_object.values()) + + def is_obj_stable(self, start_pos, end_pos): + """ + :param start_pos: [B x 3] + :param end_pos: [B x 3] + :return: + """ + is_obj_stable = torch.ones(start_pos.shape[0]).bool().to(self.device) + error_distance = self.cfg['eval_policy']['error']['distance'] + + obj_distance = torch.norm(start_pos - end_pos, dim=1) + is_obj_stable *= (abs(obj_distance) < error_distance) + return is_obj_stable + + def step_sim_q(self, sim_q): + sim_q = sim_q.to(self.device) + self.gym.set_dof_position_target_tensor(self.sim, gymtorch.unwrap_tensor(sim_q)) + self.gym.set_dof_actuation_force_tensor(self.sim, gymtorch.unwrap_tensor(self.effort_action)) + + self.gym.simulate(self.sim) + self.gym.fetch_results(self.sim, True) + + # refresh tensors + self.gym.refresh_rigid_body_state_tensor(self.sim) + self.gym.refresh_dof_state_tensor(self.sim) + self.gym.refresh_jacobian_tensors(self.sim) + self.gym.refresh_mass_matrix_tensors(self.sim) + + # update viewer + self.gym.step_graphics(self.sim) + self.gym.draw_viewer(self.viewer, self.sim, True) + self.gym.sync_frame_time(self.sim) + + def _set_init_pose(self): + # todo: for debug | data Aug. 28 + self.closure_q = self.init_q.clone() + self.closure_q[:, 8:] += 0.3 + for i_iter in range(self.cfg['eval_policy']['init']['steps']): + self.step_sim_q(self.closure_q.clone()) + pass + + def _set_normal_force_pose(self): + step_size = 0.2 + learning_rate = 0.5 + contact_threshold = 0.005 + num_object_pts = 2048 + object_pts = self.object_mesh.sample(num_object_pts) + object_pts = torch.tensor(object_pts, device=self.device) + + opt_q_global = self.init_opt_q[:, :11].clone().detach().to(self.device) + opt_q_angle = self.init_opt_q[:, 11:].clone().detach().to(self.device) + opt_q_angle.requires_grad = True + opt_q = torch.cat([opt_q_global, opt_q_angle], dim=1) + optimizer = torch.optim.Adam([opt_q_angle], lr=learning_rate) + hand_model = get_handmodel(robot='shadowhand', batch_size=self.env_num, device=self.device, hand_scale=1.) + surface_points, surface_normal = hand_model.get_surface_points_and_normals(q=opt_q) + surface_points = surface_points.reshape(-1, 3) + surface_normal = surface_normal.reshape(-1, 3) + with torch.no_grad(): + surface_points_distance2mesh = surface_points.clone().reshape(-1, 1, 3).repeat(1, num_object_pts, 1) + + surface_points_distance2mesh -= object_pts + + surface_points_distance2mesh = surface_points_distance2mesh.norm(dim=2) + surface_points_distance2mesh_values = torch.min(surface_points_distance2mesh, dim=1).values + normal_mask = torch.zeros_like(surface_points_distance2mesh_values, device=self.device) + normal_mask[surface_points_distance2mesh_values < contact_threshold] = 1. + normal_mask = normal_mask.reshape(-1, 1).repeat(1, 3) + + surface_points_target = (surface_points.clone() + step_size * normal_mask * surface_normal.clone()).detach() + loss = (surface_points - surface_points_target).norm(dim=1).sum() + optimizer.zero_grad() + loss.backward() + optimizer.step() + + self.closure_q = torch.cat([opt_q_global, opt_q_angle.clone().detach()], dim=1) + self.closure_q.requires_grad = False + # self.closure_q = self.init_opt_q + torch.clamp(self.closure_q - self.init_opt_q, min=-0.3, max=0.3) + self.closure_q = self.q_transfer_o2s(self.closure_q) + for i_iter in range(self.cfg['eval_policy']['init']['steps']): + self.step_sim_q(self.closure_q.clone()) + + def push_object(self): + # self._set_init_pose() + self._set_normal_force_pose() + achieve_6dir = torch.ones(self.num_envs).bool().to(self.device) + self.force_num_steps = self.cfg['eval_policy']['dynamic']['num_steps'] + self.object_force_directions = self.cfg['eval_policy']['dynamic']['directions'].values() + with tqdm(range(self.force_num_steps * 6), desc='Test Push Object:') as pbar: + for i_direction in self.object_force_directions: + i_achieve = self._push_object_with_direction(i_direction, pbar) + # print(f'i_direction: {i_achieve}') + achieve_6dir *= i_achieve + return achieve_6dir + + def _push_object_with_direction(self, i_direction, pbar): + object_force_magnitude = self.cfg['eval_policy']['dynamic']['magnitude_per_volume'] * self.object_volume + # print(f'force: {object_force_magnitude}') + object_pos_start = self.get_obj_pos() + for i_iter in range(self.force_num_steps): + + self.gym.refresh_actor_root_state_tensor(self.sim) + self.gym.refresh_dof_state_tensor(self.sim) + self.gym.refresh_rigid_body_state_tensor(self.sim) + force_position = self.rigid_body_tensor[:, :, :3].clone() + object_force = torch.zeros_like(force_position, device='cuda') + object_force[:, -1, :] = object_force_magnitude * torch.tensor(i_direction, device='cuda') + self.gym.apply_rigid_body_force_at_pos_tensors(self.sim, gymtorch.unwrap_tensor(object_force), + gymtorch.unwrap_tensor(force_position), gymapi.ENV_SPACE) + self.gym.simulate(self.sim) + self.gym.fetch_results(self.sim, True) + # # todo: for debug + # print(f'contact len: {len(self.gym.get_rigid_contacts(self.sim))}') + if not self.headless: + self.render() + if self.cfg["env"]["enableCameraSensors"] is True: + self.gym.step_graphics(self.sim) + self.gym.draw_viewer(self.viewer, self.sim, True) + self.gym.sync_frame_time(self.sim) + pbar.update() + object_pos_terminal = self.get_obj_pos() + return self.is_obj_stable(object_pos_start, object_pos_terminal) + + def create_sim(self): + self.dt = self.sim_params.dt + self.up_axis_idx = self.set_sim_params_up_axis( + self.sim_params, self.up_axis) + + self.sim = super().create_sim(self.device_id, self.graphics_device_id, + self.physics_engine, self.sim_params) + self._create_ground_plane() + self._place_agents( + self.cfg["env"]["numTrain"], self.cfg["env"]["envSpacing"]) + + def _load_agent(self, env_ptr, env_id): + + if self.dexterous_loaded == False: + self.dexterous_actor_list = [] + asset_root = self.asset_root + dexterous_asset_file = "movable_hand_urdf/shadowhand/robots/movable_shadowhand.urdf" + asset_options = gymapi.AssetOptions() + asset_options.density = self.cfg['agent']['density'] + asset_options.fix_base_link = True + asset_options.disable_gravity = True + # Switch Meshes from Z-up left-handed system to Y-up Right-handed coordinate system. + asset_options.flip_visual_attachments = False + asset_options.armature = 0.01 + asset_options.use_mesh_materials = True + asset_options.mesh_normal_mode = gymapi.COMPUTE_PER_VERTEX + asset_options.override_com = True # recompute center of mesh + asset_options.override_inertia = True # recompute inertia + asset_options.vhacd_enabled = True + asset_options.vhacd_params = gymapi.VhacdParams() + # asset_options.vhacd_params.resolution = 3000000 + asset_options.vhacd_params.resolution = 1000000 + + self.dexterous_asset = self.gym.load_asset( + self.sim, asset_root, dexterous_asset_file, asset_options) + self.dexterous_loaded = True + + dexterous_dof_max_torque, self.dexterous_dof_lower_limits, self.dexterous_dof_upper_limits = self._get_dof_property( + self.dexterous_asset) + + dof_props = self.gym.get_asset_dof_properties(self.dexterous_asset) + if self.cfg["env"]["driveMode"] in ["pos", "ik"]: + dof_props["driveMode"][:].fill(gymapi.DOF_MODE_POS) + dof_props["stiffness"][:].fill(self.cfg['agent']['dof_props']['stiffness']) + dof_props["velocity"][:].fill(self.cfg['agent']['dof_props']['velocity']) + dof_props["damping"][:].fill(self.cfg['agent']['dof_props']['damping']) + else: # osc + dof_props["driveMode"][:].fill(gymapi.DOF_MODE_EFFORT) + dof_props["stiffness"][:].fill(0.0) + dof_props["damping"][:].fill(0.0) + # root pose + initial_dexterous_pose = gymapi.Transform() + initial_dexterous_pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0) + initial_dexterous_pose.p = gymapi.Vec3(0.0, 0.0, 0.0) + + # set start dof + self.dexterous_num_dofs = self.gym.get_asset_dof_count(self.dexterous_asset) + + # default_dof_pos = np.zeros(self.dexterous_num_dofs, dtype=np.float32) + default_dof_pos = self.init_q[env_id, :].cpu() + + # initialize for the pose and rotation + # default_dof_pos = self.compute_pos_action(self.translation, self.rpy, self.new_joint_angle, env_id) + + dexterous_dof_state = np.zeros_like( + dexterous_dof_max_torque, gymapi.DofState.dtype) + dexterous_dof_state["pos"] = default_dof_pos + + dexterous_actor = self.gym.create_actor( + env_ptr, + self.dexterous_asset, + initial_dexterous_pose, + "dexterous", + env_id, + 0, + 0) + + dexterous_shape_props = self.gym.get_actor_rigid_shape_properties(env_ptr, dexterous_actor) + for shape in dexterous_shape_props: + shape.friction = self.cfg['agent']['shape']['friction'] + + self.gym.set_actor_rigid_shape_properties(env_ptr, dexterous_actor, dexterous_shape_props) + self.gym.set_actor_dof_properties(env_ptr, dexterous_actor, dof_props) + self.gym.set_actor_dof_states( + env_ptr, dexterous_actor, dexterous_dof_state, gymapi.STATE_ALL) + self.dexterous_actor_list.append(dexterous_actor) + + self.dexterous_link_dict = self.gym.get_asset_rigid_body_dict(self.dexterous_asset) + + self.forearm_index = self.dexterous_link_dict["forearm"] + # get global index of hand in rigid body state tensor + hand_idx = self.gym.find_actor_rigid_body_index(env_ptr, dexterous_actor, "forearm", gymapi.DOMAIN_SIM) + self.hand_idxs.append(hand_idx) + + def prepare_ik(self): + + # get dof state tensor + self.dof_pos = self.dexterous_dof_tensor[:, :, 0].view(self.num_envs, self.dexterous_num_dofs, 1) + self.dof_vel = self.dexterous_dof_tensor[:, :, 1].view(self.num_envs, self.dexterous_num_dofs, 1) + + # Set action tensors + self.pos_action = torch.zeros_like(self.dof_pos).squeeze(-1) + self.effort_action = torch.zeros_like(self.pos_action) + + def q_transfer_o2s(self, opt_q): + """ + :param opt_q: + :return: + """ + opt_q = opt_q.detach().clone() + opt_translation = opt_q[:, :3].float().clone() + sim_translation = opt_translation.clone() + opt_rotation = opt_q[:, 3:9].clone() + opt_joint_angle = opt_q[:, 9:].clone() + + rot_matrix = robust_compute_rotation_matrix_from_ortho6d(opt_rotation).cpu() + sim_rpy = torch.zeros(rot_matrix.shape[0], 3, device=self.device) + for i in range(rot_matrix.shape[0]): + # print('------**------') + # print(rot_matrix[i]) + sim_rpy[i] = torch.tensor(r2euler(rot_matrix[i], type="XYZ"), device=self.device) + # # TODO: for debug + # print(get_rot6d_from_rpy(sim_rpy[i])) + # quit() + sim_joint_angle = torch.zeros_like(opt_joint_angle, device=self.device) + sim_joint_angle[:, :6] = opt_joint_angle[:, :6] + sim_joint_angle[:, 6:11] = opt_joint_angle[:, 14:19] + sim_joint_angle[:, 11:15] = opt_joint_angle[:, 6:10] + sim_joint_angle[:, 15:19] = opt_joint_angle[:, 10:14] + sim_joint_angle[:, 19:24] = opt_joint_angle[:, 19:24] + + pos_action = torch.cat([sim_translation, sim_rpy, sim_joint_angle], dim=1) + return pos_action + + def step_opt_q(self, opt_q): + # opt_q = opt_q.detach().cpu() + opt_q = opt_q.to(self.device) + self.gym.simulate(self.sim) + self.gym.fetch_results(self.sim, True) + + self.pos_action = self.q_transfer_o2s(opt_q) + self.gym.set_dof_position_target_tensor(self.sim, gymtorch.unwrap_tensor(self.pos_action)) + self.gym.set_dof_actuation_force_tensor(self.sim, gymtorch.unwrap_tensor(self.effort_action)) + + # refresh tensors + self.gym.refresh_rigid_body_state_tensor(self.sim) + self.gym.refresh_dof_state_tensor(self.sim) + self.gym.refresh_jacobian_tensors(self.sim) + self.gym.refresh_mass_matrix_tensors(self.sim) + + self.num_sim += 1 + # update viewer + self.gym.step_graphics(self.sim) + self.gym.draw_viewer(self.viewer, self.sim, True) + self.gym.sync_frame_time(self.sim) + + def get_q_as_opt(self): + # print(self.dof_state_tensor[:, :, 0].size()) + # refresh tensors + sim_q_state = self.dof_state_tensor[:, :, 0].clone() + sim_translation = sim_q_state[:, :3] + sim_rpy = sim_q_state[:, 3:6] + sim_q_angle = sim_q_state[:, 6:] + + opt_translation = sim_translation + opt_rotation = torch.zeros(sim_q_state.shape[0], 6, device=self.device) + for i in range(sim_q_state.shape[0]): + opt_rotation[i, :] = torch.tensor(get_rot6d_from_rpy(sim_rpy[i, :]), device=self.device) + opt_q_angle = torch.zeros_like(sim_q_angle, device=self.device) + opt_q_angle[:, :6] = sim_q_angle[:, :6] + opt_q_angle[:, 6:10] = sim_q_angle[:, 11:15] + opt_q_angle[:, 10:14] = sim_q_angle[:, 15:19] + opt_q_angle[:, 14:19] = sim_q_angle[:, 6:11] + opt_q_angle[:, 19:24] = sim_q_angle[:, 19:24] + return torch.cat([opt_translation, opt_rotation, opt_q_angle], dim=1) + + def get_obj_pos(self): + return self.rigid_body_tensor[:, -1, :3].detach().clone() + + def _get_dof_property(self, asset): + dof_props = self.gym.get_asset_dof_properties(asset) + dof_num = self.gym.get_asset_dof_count(asset) + dof_lower_limits = [] + dof_upper_limits = [] + dof_max_torque = [] + for i in range(dof_num): + dof_max_torque.append(dof_props['effort'][i]) + dof_lower_limits.append(dof_props['lower'][i]) + dof_upper_limits.append(dof_props['upper'][i]) + dof_max_torque = np.array(dof_max_torque) + dof_lower_limits = np.array(dof_lower_limits) + dof_upper_limits = np.array(dof_upper_limits) + return dof_max_torque, dof_lower_limits, dof_upper_limits + + def _load_obj_asset(self): + + self.obj_name_list = [] + self.obj_asset_list = [] + self.table_asset_list = [] + self.obj_pose_list = [] + self.table_pose_list = [] + self.obj_actor_list = [] + self.table_actor_list = [] + + train_len = len(self.cfg["env"]["asset"]["trainObjAssets"].items()) + train_len = min(train_len, self.num_train) + total_len = train_len + used_len = min(total_len, self.tot_num) + + select_train_asset = [i for i in range(train_len)] + + with tqdm(total=used_len) as pbar: + pbar.set_description('Loading assets:') + cur = 0 + + obj_asset_list = [] + # prepare the assets to be used + if self.object_name is None: + raise NotImplementedError + # for id, (name, val) in enumerate(self.cfg["env"]["asset"]["trainObjAssets"].items()): + # if id in select_train_asset: + # obj_asset_list.append((id, (name, val))) + # for id, (name, val) in obj_asset_list: + # self.obj_name_list.append(name) + # # load object + else: + print(f'load object asset into IsaacGym: {self.object_name}') + object_asset_options = gymapi.AssetOptions() + object_asset_options.density = self.cfg['object']['density'] + # update + object_asset_options.linear_damping = self.cfg['object']['damping']['linear'] + object_asset_options.angular_damping = self.cfg['object']['damping']['angular'] + + object_asset_options.fix_base_link = self.fix_object + object_asset_options.disable_gravity = True + object_asset_options.use_mesh_materials = True + object_asset_options.mesh_normal_mode = gymapi.COMPUTE_PER_VERTEX + object_asset_options.override_com = True # recompute center of mesh + object_asset_options.override_inertia = True # recompute inertia + object_asset_options.vhacd_enabled = True + object_asset_options.vhacd_params = gymapi.VhacdParams() + # object_asset_options.vhacd_params.resolution = 3000000 + object_asset_options.vhacd_params.resolution = 1000000 + + dataset_name = self.object_name.split('+')[0] + object_name = self.object_name.split('+')[1] + object_urdf_path = f'object/{dataset_name}/{object_name}/{object_name}.urdf' + object_mesh_path = f'assets/object/{dataset_name}/{object_name}/{object_name}.stl' + self.object_mesh = tm.load(object_mesh_path) + + obj_asset = self.gym.load_asset( + self.sim, 'assets', object_urdf_path, object_asset_options) + self.obj_asset_list.append(obj_asset) + rig_dict = self.gym.get_asset_rigid_body_dict(obj_asset) + self.obj_rig_name = list(rig_dict.keys())[0] + obj_start_pose = gymapi.Transform() + obj_start_pose.p = gymapi.Vec3(0.0, 0.0, 0.0) + obj_start_pose.r = gymapi.Quat(0., 0., 0., 1.) + self.obj_pose_list.append(obj_start_pose) + + def _load_obj(self, env_ptr, env_id): + + if self.obj_loaded == False: + self._load_obj_asset() + self.obj_loaded = True + + obj_type = env_id // self.env_per_object + subenv_id = env_id % self.env_per_object + obj_actor = self.gym.create_actor( + env_ptr, + self.obj_asset_list[obj_type], + self.obj_pose_list[obj_type], + "obj{}-{}".format(obj_type, subenv_id), + env_id, + 0, + 0) + + obj_shape_props = self.gym.get_actor_rigid_shape_properties(env_ptr, obj_actor) + for shape in obj_shape_props: + shape.friction = self.cfg['object']['shape']['friction'] + self.gym.set_actor_rigid_shape_properties(env_ptr, obj_actor, obj_shape_props) + + self.obj_actor_list.append(obj_actor) + assert(self.gym.get_actor_rigid_body_names(self.env_ptr_list[0], 1)[0] == 'object') + + def _place_agents(self, env_num, spacing): + + print("Simulator: creating agents") + + lower = gymapi.Vec3(-spacing, -spacing, 0.0) + upper = gymapi.Vec3(spacing, spacing, spacing) + self.space_middle = torch.zeros((env_num, 3), device=self.device) + self.space_range = torch.zeros((env_num, 3), device=self.device) + self.space_middle[:, 0] = self.space_middle[:, 1] = 0 + self.space_middle[:, 2] = spacing / 2 + self.space_range[:, 0] = self.space_range[:, 1] = spacing + self.space_middle[:, 2] = spacing / 2 + num_per_row = int(np.sqrt(env_num)) + + with tqdm(total=env_num) as pbar: + pbar.set_description('Enumerating envs:') + for env_id in range(env_num): + env_ptr = self.gym.create_env( + self.sim, lower, upper, num_per_row) + self.env_ptr_list.append(env_ptr) + self._load_agent(env_ptr, env_id) + self._load_obj(env_ptr, env_id) + pbar.update(1) + + def _create_ground_plane(self): + plane_params = gymapi.PlaneParams() + plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0) + plane_params.distance = 1. + plane_params.static_friction = 0.1 + plane_params.dynamic_friction = 0.1 + self.gym.add_ground(self.sim, plane_params) + + def _draw_line(self, src, dst): + line_vec = np.stack([src, dst]).flatten().astype(np.float32) + color = np.array([1, 0, 0], dtype=np.float32) + self.gym.clear_lines(self.viewer) + self.gym.add_lines( + self.viewer, + self.env_ptr_list[0], + self.env_num, + line_vec, + color + ) + + # def __del__(self): + # self.gym.destroy_sim(self.sim) + # self.gym.destroy_viewer(self.viewer) + + # def simulate(self): + # self.num_sim = 0 + # while not self.gym.query_viewer_has_closed(self.viewer): + # # step the physics + # self.gym.simulate(self.sim) + # self.gym.fetch_results(self.sim, True) + # + # self.move_base(self.translation, self.rpy, self.new_joint_angle, self.num_sim) + # + # # refresh tensors + # self.gym.refresh_rigid_body_state_tensor(self.sim) + # self.gym.refresh_dof_state_tensor(self.sim) + # self.gym.refresh_jacobian_tensors(self.sim) + # self.gym.refresh_mass_matrix_tensors(self.sim) + # + # self.num_sim += 1 + # # update viewer + # self.gym.step_graphics(self.sim) + # self.gym.draw_viewer(self.viewer, self.sim, True) + # self.gym.sync_frame_time(self.sim) + + +def orientation_error(desired, current): + cc = quat_conjugate(current) + q_r = quat_mul(desired, cc) + return q_r[:, 0:3] * torch.sign(q_r[:, 3]).unsqueeze(-1) + + +def control_ik(j_eef, device, dpose, num_envs): + # Set controller parameters + # IK params + damping = 0.05 + # solve damped least squares + j_eef_T = torch.transpose(j_eef, 1, 2) + lmbda = torch.eye(6, device=device) * (damping ** 2) + u = (j_eef_T @ torch.inverse(j_eef @ j_eef_T + lmbda) + @ dpose).view(num_envs, -1) + return u + + +def relative_pose(src, dst): + shape = dst.shape + p = dst.view(-1, shape[-1])[:, :3] - src.view(-1, src.shape[-1])[:, :3] + ip = dst.view(-1, shape[-1])[:, 3:] + ret = torch.cat((p, ip), dim=1) + return ret.view(*shape) + + +def get_sim_param(): + # initialize sim + sim_params = gymapi.SimParams() + sim_params.dt = 1. / 60. + sim_params.num_client_threads = 0 + sim_params.physx.solver_type = 1 + sim_params.physx.num_position_iterations = 4 + sim_params.physx.num_velocity_iterations = 0 + sim_params.physx.num_threads = 4 + sim_params.physx.use_gpu = True + sim_params.physx.num_subscenes = 0 + sim_params.physx.max_gpu_contact_pairs = 8 * 1024 * 1024 + sim_params.use_gpu_pipeline = True + sim_params.physx.use_gpu = True + sim_params.physx.num_threads = 0 + return sim_params + + +def compute_rotation_matrix_from_ortho6d(poses): + """ + Code from + https://github.com/papagina/RotationContinuity + On the Continuity of Rotation Representations in Neural Networks + Zhou et al. CVPR19 + https://zhouyisjtu.github.io/project_rotation/rotation.html + """ + x_raw = poses[:, 0:3] # batch*3 + y_raw = poses[:, 3:6] # batch*3 + + x = normalize_vector(x_raw) # batch*3 + z = cross_product(x, y_raw) # batch*3 + z = normalize_vector(z) # batch*3 + y = cross_product(z, x) # batch*3 + + x = x.view(-1, 3, 1) + y = y.view(-1, 3, 1) + z = z.view(-1, 3, 1) + matrix = torch.cat((x, y, z), 2) # batch*3*3 + return matrix + + +def robust_compute_rotation_matrix_from_ortho6d(poses): + """ + Instead of making 2nd vector orthogonal to first + create a base that takes into account the two predicted + directions equally + """ + x_raw = poses[:, 0:3] # batch*3 + y_raw = poses[:, 3:6] # batch*3 + + x = normalize_vector(x_raw) # batch*3 + y = normalize_vector(y_raw) # batch*3 + middle = normalize_vector(x + y) + orthmid = normalize_vector(x - y) + x = normalize_vector(middle + orthmid) + y = normalize_vector(middle - orthmid) + # Their scalar product should be small ! + # assert torch.einsum("ij,ij->i", [x, y]).abs().max() < 0.00001 + z = normalize_vector(cross_product(x, y)) + + x = x.view(-1, 3, 1) + y = y.view(-1, 3, 1) + z = z.view(-1, 3, 1) + matrix = torch.cat((x, y, z), 2) # batch*3*3 + # Check for reflection in matrix ! If found, flip last vector TODO + # assert (torch.stack([torch.det(mat) for mat in matrix ])< 0).sum() == 0 + return matrix + + +def cross_product(u, v): + batch = u.shape[0] + i = u[:, 1] * v[:, 2] - u[:, 2] * v[:, 1] + j = u[:, 2] * v[:, 0] - u[:, 0] * v[:, 2] + k = u[:, 0] * v[:, 1] - u[:, 1] * v[:, 0] + + out = torch.cat((i.view(batch, 1), j.view(batch, 1), k.view(batch, 1)), 1) + + return out + + +def normalize_vector(v): + batch = v.shape[0] + v_mag = torch.sqrt(v.pow(2).sum(1)) # batch + v_mag = torch.max(v_mag, v.new([1e-8])) + v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1]) + v = v / v_mag + return v + + +def quaternion_to_axis_angle(quaternions): + """ + Convert rotations given as quaternions to axis/angle. + + Args: + quaternions: quaternions with real part first, + as tensor of shape (..., 4). + + Returns: + Rotations given as a vector in axis angle form, as a tensor + of shape (..., 3), where the magnitude is the angle + turned anticlockwise in radians around the vector's + direction. + """ + norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True) + half_angles = torch.atan2(norms, quaternions[..., :1]) + angles = 2 * half_angles + eps = 1e-6 + small_angles = angles.abs() < eps + sin_half_angles_over_angles = torch.empty_like(angles) + sin_half_angles_over_angles[~small_angles] = ( + torch.sin(half_angles[~small_angles]) / angles[~small_angles] + ) + # for x small, sin(x/2) is about x/2 - (x/2)^3/6 + # so sin(x/2)/x is about 1/2 - (x*x)/48 + sin_half_angles_over_angles[small_angles] = ( + 0.5 - (angles[small_angles] * angles[small_angles]) / 48 + ) + return quaternions[..., 1:] / sin_half_angles_over_angles + + +# def get_rot6d_from_rot3d(rot3d): +# global_rotation = np.array(transforms3d.euler.euler2mat(rot3d[0], rot3d[1], rot3d[2], axes='sxyz')) +# return global_rotation.T.reshape(9)[:6] +def get_rot6d_from_rpy(rpy): + rpy_inverse = - rpy + global_rotation = np.array( + transforms3d.euler.euler2mat(rpy_inverse[0], rpy_inverse[1], rpy_inverse[2], axes='sxyz')) + return global_rotation.reshape(9)[:6] + + +if __name__ == "__main__": + sim_params = get_sim_param() + path = "envs/tasks/grasp_test_force.yaml" + with open(path) as f: + cfg = yaml.safe_load(f) + from utils.get_models import get_handmodel + + init_opt_q = torch.zeros(16, 9 + 24, device='cuda') + init_opt_q[:, :3] = torch.tensor([0.2, 0.02, -0.3], device='cuda') + init_opt_q[:, 3:9] = torch.tensor([1., 0., 0., 0., 1., 0.], device='cuda') + hand_model = get_handmodel('shadowhand', 1, 'cuda', 1.) + q_opt_joint_pose = hand_model.revolute_joints_q_lower + hand_model.revolute_joints_q_upper + q_opt_joint_pose = (q_opt_joint_pose / 2).repeat(16, 1) + print(f'upper: {hand_model.revolute_joints_q_upper}') + print(f'lower: {hand_model.revolute_joints_q_lower}') + print(f'upper - lower: {hand_model.revolute_joints_q_upper - hand_model.revolute_joints_q_lower}') + init_opt_q[:, 9:].copy_(hand_model.revolute_joints_q_lower.repeat(16, 1)) + env = IsaacGraspTestForce_shadowhand(cfg, sim_params, gymapi.SIM_PHYSX, "cuda", 0, False, + init_opt_q=init_opt_q, object_name='contactdb+apple', object_volume=0.0002) + init_opt_q[:, 9:].copy_(hand_model.revolute_joints_q_upper.repeat(16, 1)) + while not env.gym.query_viewer_has_closed(env.viewer): + env.step_opt_q(init_opt_q) + + + + + diff --git a/envs/tasks/utils/angle.py b/envs/tasks/utils/angle.py new file mode 100644 index 0000000..b02f6d6 --- /dev/null +++ b/envs/tasks/utils/angle.py @@ -0,0 +1,255 @@ + +import math +import numpy as np +from scipy.spatial.transform import Rotation as R + + +def r2euler(R, type): + R = np.array(R) + type = str(type).upper() + err = float(0.001) + + if type == "XYZ": + # R[0,2]/sqrt((R[1,2])**2 + (R[2,2])**2) == sin(beta)/|cos(beta)| + # ==> beta (-pi/2, pi/2) + beta = math.atan2(R[0,2], math.sqrt((R[1,2])**2 + (R[2,2])**2)) + + if beta >= math.pi/2-err and beta <= math.pi/2+err: + beta = math.pi/2 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[1,0], R[1,1]) + elif beta >= -(math.pi/2)-err and beta <= -(math.pi/2)+err: + beta = -math.pi/2 + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[1,0], R[1,1]) + else: + alpha = math.atan2(-(R[1,2])/(math.cos(beta)),(R[2,2])/(math.cos(beta))) + gamma = math.atan2(-(R[0,1])/(math.cos(beta)), (R[0,0])/(math.cos(beta))) + + elif type == "XZY": + # -R[0,1]/sqrt((R[1,1])**2 + (R[2,1])**2) == sin(beta)/|cos(beta)| + # ==> beta (-pi/2, pi/2) + beta = math.atan2(-R[0,1], math.sqrt((R[1,1])**2 + (R[2,1])**2)) + + if beta >= math.pi/2-err and beta <= math.pi/2+err: + beta = math.pi/2 + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[1,2], R[1,0]) + elif beta >= -(math.pi/2)-err and beta <= -(math.pi/2)+err: + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[1,2], -R[1,0]) + else: + alpha = math.atan2((R[2,1])/(math.cos(beta)), (R[1,1])/(math.cos(beta))) + gamma = math.atan2((R[0,2])/(math.cos(beta)), (R[0,0])/(math.cos(beta))) + + elif type == "YXZ": + # -R[1,2]/sqrt(R[0,2]**2 + R[2,2]**2) == sin(beta)/|cos(beta)| + # ==> beta (-pi/2, pi/2) + beta = math.atan2(-R[1,2], math.sqrt((R[0,2])**2 + (R[2,2])**2)) + + if beta >= math.pi/2-err and beta <= math.pi/2+err: + beta = math.pi/2 + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[0,1], R[0,0]) + elif beta >= -(math.pi/2)-err and beta <= -(math.pi/2)+err: + beta = -math.pi/2 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[0,1], R[0,0]) + else: + alpha = math.atan2((R[0,2])/(math.cos(beta)), (R[2,2])/(math.cos(beta))) + gamma = math.atan2((R[1,0])/(math.cos(beta)), (R[1,1])/(math.cos(beta))) + + elif type == "YZX": + # R[1,0]/sqrt(R[0,0]**2 + R[2,0]**2) == sin(beta)/|cos(beta)| + # ==> beta (-pi/2, pi/2) + beta = math.atan2(R[1,0], math.sqrt((R[0,0])**2 + (R[2,0])**2)) + + if beta >= math.pi/2-err and beta <= math.pi/2+err: + beta = math.pi/2 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[0,2], -R[0,1]) + elif beta >= -(math.pi/2)-err and beta <= -(math.pi/2)+err: + beta = -math.pi/2 + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[0,2], R[0,1]) + else: + alpha = math.atan2(-(R[2,0])/(math.cos(beta)), (R[0,0])/(math.cos(beta))) + gamma = math.atan2(-(R[1,2])/(math.cos(beta)), (R[1,1])/(math.cos(beta))) + + elif type == "ZXY": + # R[2,1]/sqrt(R[0,1]**2 + R[1,1]**2) == sin(beta)/|cos(beta)| + # ==> beta (-pi/2, pi/2) + beta = math.atan2(R[2,1], math.sqrt((R[0,1])**2 + (R[1,1])**2)) + + if beta >= math.pi/2-err and beta <= math.pi/2+err: + beta = math.pi/2 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[0,2], R[0,0]) + elif beta >= -(math.pi/2)-err and beta <= -(math.pi/2)+err: + beta = -math.pi/2 + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[0,2], R[0,0]) + else: + alpha = math.atan2(-(R[0,1])/(math.cos(beta)), (R[1,1])/(math.cos(beta))) + gamma = math.atan2(-(R[2,0])/(math.cos(beta)), (R[2,2])/(math.cos(beta))) + + elif type == "ZYX": + # -R[2,0]/sqrt(R[0,0]**2 + R[1,0]**2) == sin(beta)/|cos(beta)| + # ==> beta (-pi/2, pi/2) + beta = math.atan2(-R[2,0], math.sqrt((R[0,0])**2 + (R[1,0])**2)) + + if beta >= math.pi/2-err and beta <= math.pi/2+err: + beta = math.pi/2 + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[0,1], R[1,2]) + elif beta >= -(math.pi/2)-err and beta <= -(math.pi/2)+err: + beta = -math.pi/2 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[0,1], -R[1,2]) + else: + alpha = math.atan2((R[1,0])/(math.cos(beta)), (R[0,0])/(math.cos(beta))) + gamma = math.atan2((R[2,1])/(math.cos(beta)), (R[2,2])/(math.cos(beta))) + + elif type == "XYX": + # sqrt(R[0,1]**2 + R[0,2]**2)/R[0,0] == |sin(beta)|/cos(beta) + # ==> beta (0, pi) + beta = math.atan2(math.sqrt((R[0,1])**2 + (R[0,2])**2), R[0,0]) + if beta >= 0.0-err and beta <= 0.0+err: + beta = 0.0 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[1,2], R[1,1]) + elif beta >= math.pi-err and beta <= math.pi+err: + beta = math.pi + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[1,2], R[1,1]) + else: + alpha = math.atan2((R[1,0])/(math.sin(beta)), -(R[2,0])/(math.sin(beta))) + gamma = math.atan2((R[0,1])/(math.sin(beta)), (R[0,2])/(math.sin(beta))) + + elif type == "XZX": + # sqrt(R[1,0]**2 + R[2,0]**2)/R[0,0] == |sin(beta)|/cos(beta) + # ==> beta (0, pi) + beta = math.atan2(math.sqrt((R[1,0])**2 + (R[2,0])**2), R[0,0]) + if beta >= 0.0-err and beta <= 0.0+err: + beta = 0.0 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[1,2], R[1,1]) + elif beta >= math.pi-err and beta <= math.pi+err: + beta = math.pi + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[1,2], -R[1,1]) + else: + alpha = math.atan2((R[2,0])/(math.sin(beta)), (R[1,0])/(math.sin(beta))) + gamma = math.atan2((R[0,2])/(math.sin(beta)), -(R[0,1])/(math.sin(beta))) + + elif type == "YXY": + # sqrt(R[0,1]**2 + R[2,1]**2)/R[1,1] == |sin(beta)|/cos(beta) + # ==> beta(0, pi) + beta = math.atan2(math.sqrt((R[0,1])**2 + (R[2,1])**2), R[1,1]) + if beta >= 0.0-err and beta <= 0.0+err: + beta = 0.0 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[0,2], R[0,0]) + elif beta >= math.pi-err and beta <= math.pi+err: + beta = math.pi + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[0,2], R[0,0]) + else: + alpha = math.atan2((R[0,1])/(math.sin(beta)), (R[2,1])/(math.sin(beta))) + gamma = math.atan2((R[1,0])/(math.sin(beta)), -(R[1,2])/(math.sin(beta))) + + elif type == "YZY": + # sqrt(R[0,1]**2 + R[2,1]**2)/R[1,1] == |sin(beta)|/cos(beta) + # ==> beta(0, pi) + beta = math.atan2(math.sqrt((R[0,1])**2 + (R[2,1])**2), R[1,1]) + if beta >= 0.0-err and beta <= 0.0+err: + beta = 0.0 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[0,2], R[0,0]) + elif beta >= math.pi-err and beta <= math.pi+err: + beta = math.pi + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[0,2], -R[0,0]) + else: + alpha = math.atan2((R[2,1])/(math.sin(beta)), -(R[0,1])/(math.sin(beta))) + gamma = math.atan2((R[1,2])/(math.sin(beta)), (R[1,0])/(math.sin(beta))) + + elif type == "ZXZ": + # sqrt(R[0,2]**2 + R[1,2]**2)/R[2,2] == |sin(beta)|/cos(beta) + # ==> beta(0, pi) + beta = math.atan2(math.sqrt((R[0,2])**2 + (R[1,2])**2), R[2,2]) + if beta >= 0.0-err and beta <= 0.0+err: + beta = 0.0 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[0,1], R[0,0]) + elif beta >= math.pi-err and beta <= math.pi+err: + beta = math.pi + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[0,1], R[0,0]) + else: + alpha = math.atan2((R[0,2])/(math.sin(beta)), -(R[1,2])/(math.sin(beta))) + gamma = math.atan2((R[2,0])/(math.sin(beta)), (R[2,1])/(math.sin(beta))) + + elif type == "ZYZ": + # sqrt(R[0,2]**2 + R[1,2]**2)/R[2,2] == |sin(beta)|/cos(beta) + # ==> beta(0, pi) + beta = math.atan2(math.sqrt((R[0,2])**2 + (R[1,2])**2), R[2,2]) + if beta >= 0.0-err and beta <= 0.0+err: + beta = 0.0 + # alpha + gamma is fixed + alpha = 0.0 + gamma = math.atan2(-R[0,1], R[0,0]) + elif beta >= math.pi+err and beta <= math.pi+err: + beta = math.pi + # alpha - gamma is fixed + alpha = 0.0 + gamma = math.atan2(R[0,1], -R[0,0]) + else: + alpha = math.atan2((R[1,2])/(math.sin(beta)), (R[0,2])/(math.sin(beta))) + gamma = math.atan2((R[2,1])/(math.sin(beta)), -(R[2,0])/(math.sin(beta))) + + return alpha, beta, gamma + + +if __name__ == "__main__": + RM = np.loadtxt ('RotationMatrix.txt') + print("Upper -- intrinsic\nLower -- extrinsic") + + euler_type = str(input("Please input the type of euler angle...\n")) + + if euler_type.isupper(): + angle_0, angle_1, angle_2 = r2euler(RM, euler_type) + print("Intrinsic") + print("Angle about {} is {}".format(euler_type[0], angle_0)) + print("Angle about {} is {}".format(euler_type[1], angle_1)) + print("Angle about {} is {}".format(euler_type[2], angle_2)) + elif euler_type.islower(): + angle_0, angle_1, angle_2 = r2euler(RM, euler_type.upper()[::-1])[::-1] + print("extrinsic") + print("Angle about {} is {}".format(euler_type[0], angle_0)) + print("Angle about {} is {}".format(euler_type[1], angle_1)) + print("Angle about {} is {}".format(euler_type[2], angle_2)) + else: + pass From 2c3894fd1d4ccbe55ecd96e6040ec2278cbf59e0 Mon Sep 17 00:00:00 2001 From: Xiaoyao-Li <1617410526@qq.com> Date: Fri, 21 Apr 2023 11:40:24 +0800 Subject: [PATCH 4/7] update scripts for grasp pose evaluation --- scripts/grasp_gen_ur/test.py | 169 +++++++++++++++++++++++++++++++++++ scripts/grasp_gen_ur/test.sh | 11 +++ 2 files changed, 180 insertions(+) create mode 100644 scripts/grasp_gen_ur/test.py diff --git a/scripts/grasp_gen_ur/test.py b/scripts/grasp_gen_ur/test.py new file mode 100644 index 0000000..ba1eacf --- /dev/null +++ b/scripts/grasp_gen_ur/test.py @@ -0,0 +1,169 @@ +import os +import sys +sys.path.append(os.getcwd()) + +import gc +import yaml +import pickle +import argparse +from loguru import logger + +from isaacgym import gymapi, gymutil, gymtorch +import torch +import random +import numpy as np + +import trimesh as tm +from utils.handmodel import get_handmodel, compute_collision +from envs.tasks.grasp_test_force_shadowhand import IsaacGraspTestForce_shadowhand as IsaacGraspTestForce + + +def set_global_seed(seed: int) -> None: + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description='Test Scripts of Grasp Generation') + parser.add_argument('--stability_config', type=str, + default='envs/tasks/grasp_test_force.yaml', + help='stability config file path') + parser.add_argument('--eval_dir', type=str, required=True, + help='evaluation directory path (e.g.,\ + "outputs/2022-11-15_18-07-50_GPUR_l1_pn2_T100/eval/final/2023-04-20_13-06-44")') + parser.add_argument('--dataset_dir', type=str, required=True, + help='dataset directory path (e.g.,\ + ("/path/to/MultiDex_UR/")') + parser.add_argument('--seed', type=int, default=42, + help='random seed') + parser.add_argument('--cpu', action='store_true', default=False, + help='run all on cpu') + parser.add_argument('--onscreen', action='store_true', default=False, + help='run simulator onscreen') + + return parser.parse_args() + + +def get_sim_param(): + # initialize sim + sim_params = gymapi.SimParams() + sim_params.dt = 1./60. + sim_params.num_client_threads = 0 + sim_params.physx.solver_type = 1 + sim_params.physx.num_position_iterations = 4 + sim_params.physx.num_velocity_iterations = 0 + sim_params.physx.num_threads = 4 + sim_params.physx.use_gpu = True + sim_params.physx.num_subscenes = 0 + sim_params.physx.max_gpu_contact_pairs = 8 * 1024 * 1024 + sim_params.use_gpu_pipeline = True + sim_params.physx.use_gpu = True + sim_params.physx.num_threads = 0 + return sim_params + + +def stability_tester(args: argparse.Namespace) -> dict: + with open(args.stability_config) as f: + stability_config = yaml.safe_load(f) + sim_params = get_sim_param() + sim_headless = not args.onscreen + + # load generated grasp results here + grasps = pickle.load(open(os.path.join(args.eval_dir, 'res_diffuser.pkl'), 'rb')) + isaac_env = None + results = {} + across_all_cases = 0 + across_all_succ = 0 + + for object_name in grasps['sample_qpos'].keys(): + logger.info(f'Stability test for [{object_name}]') + q_generated = grasps['sample_qpos'][object_name] + q_generated = torch.tensor(q_generated, device=args.device).to(torch.float32) + + # load object mesh + object_mesh_path = f'./assets/object/{object_name.split("+")[0]}/{object_name.split("+")[1]}/{object_name.split("+")[1]}.stl' + object_mesh = tm.load(object_mesh_path) + object_volume = object_mesh.volume + isaac_env = IsaacGraspTestForce(stability_config, sim_params, gymapi.SIM_PHYSX, + args.device, 0, headless=sim_headless, init_opt_q=q_generated, + object_name=object_name, object_volume=object_volume, fix_object=False) + succ_grasp_object = isaac_env.push_object() + results[object_name] = {'total': int(succ_grasp_object.shape[0]), + 'succ': int(succ_grasp_object.sum()), + 'case_list': succ_grasp_object.tolist()} + logger.info(f'Success rate of [{object_name}]: {int(succ_grasp_object.sum())} / {int(succ_grasp_object.shape[0])}') + across_all_succ += int(succ_grasp_object.sum()) + across_all_cases += int(succ_grasp_object.shape[0]) + if isaac_env is not None: + del isaac_env + gc.collect() + logger.info(f'**Success Rate** across all objects: {across_all_succ} / {across_all_cases}') + + return results + + +def diversity_tester(args: argparse.Namespace, stability_results: dict) -> None: + grasps = pickle.load(open(os.path.join(args.eval_dir, 'res_diffuser.pkl'), 'rb')) + + qpos_std = [] + for object_name in grasps['sample_qpos'].keys(): + i_qpos = grasps['sample_qpos'][object_name][:, 9:] + i_qpos = i_qpos[stability_results[object_name]['case_list'], :] + if i_qpos.shape[0]: + i_qpos = np.sqrt(i_qpos.var(axis=0)) + qpos_std.append(i_qpos) + + qpos_std = np.stack(qpos_std, axis=0) + qpos_std = qpos_std.mean(axis=0).mean() + logger.info(f'**Diversity** (std: rad.) across all success grasps: {qpos_std}') + + +def collision_tester(args: argparse.Namespace, stability_results: dict) -> None: + _BATCHSIZE = 16 #NOTE: adjust this batchsize to fit your GPU memory && need to be divided by generated grasps per object + _NPOINTS = 4096 #NOTE: number of surface points sampled from a object + + grasps = pickle.load(open(os.path.join(args.eval_dir, 'res_diffuser.pkl'), 'rb')) + obj_pcds_nors_dict = pickle.load(open('/home/puhao/data/MultiDex_UR/object_pcds_nors.pkl', 'rb')) + hand_model = get_handmodel(batch_size=_BATCHSIZE, device=args.device) + + collisions_dict = {obj: [] for obj in grasps['sample_qpos'].keys()} + for object_name in grasps['sample_qpos'].keys(): + qpos = grasps['sample_qpos'][object_name] + obj_pcd_nor = obj_pcds_nors_dict[object_name][:_NPOINTS, :] + + for i in range(qpos.shape[0] // _BATCHSIZE): + i_qpos = qpos[i * _BATCHSIZE: (i + 1) * _BATCHSIZE, :] + hand_model.update_kinematics(q=torch.tensor(i_qpos, device=args.device)) + hand_surface_points = hand_model.get_surface_points() + #TODO: needed to be checked + depth_collision = compute_collision(torch.tensor(obj_pcd_nor, device=args.device), hand_surface_points) + collisions_dict[object_name].append(np.array(depth_collision.cpu()[stability_results[object_name]['case_list'][i * _BATCHSIZE : (i + 1) * _BATCHSIZE]])) + collisions_dict[object_name] = np.concatenate(collisions_dict[object_name], axis=0) + + collision_values = np.concatenate([collisions_dict[object_name] for object_name in grasps['sample_qpos'].keys()], axis=0) + logger.info(f'**Collision** (depth: mm.) across all grasps: {collision_values.mean() * 1e3}') + +def main() -> None: + args = parse_args() + set_global_seed(args.seed) + args.device = 'cpu' if args.cpu else 'cuda' + + logger.add(args.eval_dir + '/evaluation.log') + logger.info(f'Evaluation directory: {args.eval_dir}') + + logger.info('Start evaluating..') + + stability_results = stability_tester(args) + diversity_tester(args, stability_results) + collision_tester(args, stability_results) + + logger.info('End evaluating..') + + +if __name__ == '__main__': + main() diff --git a/scripts/grasp_gen_ur/test.sh b/scripts/grasp_gen_ur/test.sh index e69de29..a63c5fc 100755 --- a/scripts/grasp_gen_ur/test.sh +++ b/scripts/grasp_gen_ur/test.sh @@ -0,0 +1,11 @@ +EVAL_DIR=$1 +DATASET_DIR=$2 + +# eval_dir: path to the directory where the evaluation results be sampled +# (e.g., "outputs/2022-11-15_18-07-50_GPUR_l1_pn2_T100/eval/final/2023-04-20_13-06-44") +# dataset_dir: path to the directory where the MultiDex_UR dataset is stored +# (e.g., "/home/puhao/data/MultiDex_UR") +python ./scripts/grasp_gen_ur/test.py --eval_dir=${EVAL_DIR} \ + --dataset_dir=${DATASET_DIR} \ + --stability_config='envs/tasks/grasp_test_force.yaml' \ + --seed=42 \ From 7cc7b5dcf6bc1542ec80da9ca2336f0fa6318566 Mon Sep 17 00:00:00 2001 From: Xiaoyao-Li <1617410526@qq.com> Date: Fri, 21 Apr 2023 11:40:47 +0800 Subject: [PATCH 5/7] update handmodel for mask print output --- utils/handmodel.py | 37 +++++++++++++++++++++++++++++++++++-- 1 file changed, 35 insertions(+), 2 deletions(-) diff --git a/utils/handmodel.py b/utils/handmodel.py index c822b45..365da22 100644 --- a/utils/handmodel.py +++ b/utils/handmodel.py @@ -48,7 +48,7 @@ def __init__(self, robot_name, urdf_filename, mesh_path, raise NotImplementedError for i_link, link in enumerate(visual.links): - print(f"Processing link #{i_link}: {link.name}") + # print(f"Processing link #{i_link}: {link.name}") # load mesh if len(link.visuals) == 0: continue @@ -184,7 +184,7 @@ def get_surface_points(self, q=None, downsample=True): surface_points.append( torch.matmul(trans_matrix, self.surface_points[link_name].transpose(1, 2)).transpose(1, 2)[..., :3]) surface_points = torch.cat(surface_points, 1) - surface_points = torch.matmul(self.global_rotation, surface_points.transpose(1, 2)).transpose(1, + surface_points = torch.matmul(self.global_rotation.float(), surface_points.transpose(1, 2)).transpose(1, 2) + self.global_translation.unsqueeze( 1) # if downsample: @@ -296,6 +296,39 @@ def get_handmodel(batch_size, device, hand_scale=1., robot='shadowhand'): return hand_model +def compute_collision(obj_pcd_nor: torch.Tensor, hand_pcd: torch.Tensor): + """ + :param obj_pcd_nor: N_obj x 6 + :param hand_surface_points: B x N_hand x 3 + :return: + """ + b = hand_pcd.shape[0] + n_obj = obj_pcd_nor.shape[0] + n_hand = hand_pcd.shape[1] + + obj_pcd = obj_pcd_nor[:, :3] + obj_nor = obj_pcd_nor[:, 3:6] + + # batch the obj pcd + batch_obj_pcd = obj_pcd.unsqueeze(0).repeat(b, 1, 1).view(b, 1, n_obj, 3) + batch_obj_pcd = batch_obj_pcd.repeat(1, n_hand, 1, 1) + # batch the hand pcd + batch_hand_pcd = hand_pcd.view(b, n_hand, 1, 3).repeat(1, 1, n_obj, 1) + # compute the pair wise dist + hand_obj_dist = (batch_obj_pcd - batch_hand_pcd).norm(dim=3) + hand_obj_dist, hand_obj_indices = hand_obj_dist.min(dim=2) + # gather the obj points and normals w.r.t. hand points + hand_obj_points = torch.stack([obj_pcd[x, :] for x in hand_obj_indices], dim=0) + hand_obj_normals = torch.stack([obj_nor[x, :] for x in hand_obj_indices], dim=0) + # compute the signs + hand_obj_signs = ((hand_obj_points - hand_pcd) * hand_obj_normals).sum(dim=2) + hand_obj_signs = (hand_obj_signs > 0.).float() + # signs dot dist to compute collision value + collision_value = (hand_obj_signs * hand_obj_dist).max(dim=1).values + # collision_value = (hand_obj_signs * hand_obj_dist).mean(dim=1) + return collision_value + + if __name__ == '__main__': from plotly_utils import plot_point_cloud seed = 0 From 02c7fee542c66d1dc6d0f2f8702149668e93c915 Mon Sep 17 00:00:00 2001 From: Silverster98 <1539168414@qq.com> Date: Fri, 21 Apr 2023 15:05:45 +0800 Subject: [PATCH 6/7] Update readme for obj branch --- README.md | 87 +++++++++++++++++++++++++++++++++--- scripts/grasp_gen_ur/test.sh | 2 +- 2 files changed, 83 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 3fad51e..7d908ec 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,9 @@ # Diffusion-based Generation, Optimization, and Planning in 3D Scenes

+ + Paper PDF + Paper arXiv @@ -28,6 +31,7 @@ This repository is the official implementation of paper "Diffusion-based Generat We introduce SceneDiffuser, a conditional generative model for 3D scene understanding. SceneDiffuser provides a unified model for solving scene-conditioned generation, optimization, and planning. In contrast to prior work, SceneDiffuser is intrinsically scene-aware, physics-based, and goal-oriented. +[Paper](https://scenediffuser.github.io/paper.pdf) | [arXiv](https://arxiv.org/abs/2301.06015) | [Project](https://scenediffuser.github.io/) | [HuggingFace Demo](https://huggingface.co/spaces/SceneDiffuser/SceneDiffuserDemo) | @@ -41,23 +45,33 @@ We introduce SceneDiffuser, a conditional generative model for 3D scene understa We introduce SceneDiffuser, a conditional generative model for 3D scene understanding. SceneDiffuser provides a unified model for solving scene-conditioned generation, optimization, and planning. In contrast to prior works, SceneDiffuser is intrinsically scene-aware, physics-based, and goal-oriented. With an iterative sampling strategy, SceneDiffuser jointly formulates the scene-aware generation, physics-based optimization, and goal-oriented planning via a diffusion-based denoising process in a fully differentiable fashion. Such a design alleviates the discrepancies among different modules and the posterior collapse of previous scene-conditioned generative models. We evaluate SceneDiffuser with various 3D scene understanding tasks, including human pose and motion generation, dexterous grasp generation, path planning for 3D navigation, and motion planning for robot arms. The results show significant improvements compared with previous models, demonstrating the tremendous potential of SceneDiffuser for the broad community of 3D scene understanding. +## News + +- [ 2023.04 ] We release the code for grasp generation and arm motion planning! + ## Setup -1. Create a new `conda` environemnt and activate it +1. Create a new `conda` environemnt and activate it. ```bash conda create -n 3d python=3.8 conda activate 3d ``` -2. Install dependent libraries with `pip` +2. Install dependent libraries with `pip`. ```bash pip install -r pre-requirements.txt pip install -r requirements.txt ``` - - We use `pytorch1.11` and `cuda11.3`, modify `pre-requirements.txt` to install [other versions](https://pytorch.org/get-started/previous-versions/) of `pytorch` + - We use `pytorch1.11` and `cuda11.3`, modify `pre-requirements.txt` to install [other versions](https://pytorch.org/get-started/previous-versions/) of `pytorch`. + +3. Install [Isaac Gym](https://developer.nvidia.com/isaac-gym) and install [pointnet2](https://github.com/daveredrum/Pointnet2.ScanNet) by executing the following command (optional for grasp generation and arm motion planning). + + ```bash + pip install git+https://github.com/daveredrum/Pointnet2.ScanNet.git#subdirectory=pointnet2 + ``` ## Data & Checkpoints @@ -77,10 +91,13 @@ Download our [pre-trained model](https://drive.google.com/drive/folders/1CKJER3C task|checkpoints|desc -|-|- +Pretrained Point Transformer|2022-04-13_18-29-56_POINTTRANS_C_32768| Pose Generation|2022-11-09_11-22-52_PoseGen_ddm4_lr1e-4_ep100| Motion Generation|2022-11-09_12-54-50_MotionGen_ddm_T200_lr1e-4_ep300|w/o start position Motion Generation|2022-11-09_14-28-12_MotionGen_ddm_T200_lr1e-4_ep300_obser|w/ start position Path Planning|2022-11-25_20-57-28_Path_ddm4_LR1e-4_E100_REL| +Grasp Generation|2022-11-15_18-07-50_GPUR_l1_pn2_T100| +Arm Motion Planning|2022-11-11_14-28-30_FK2Plan_ptr_T30_4|denoising step is 30 ## Task-1: Human Pose Generation in 3D Scenes @@ -155,7 +172,39 @@ bash scripts/motion_gen/sample.sh ${CKPT} [OPT] ## Task-3: Dexterous Grasp Generation for 3D Objects -coming soon. +To run this code, you first need to change the git branch to `obj` by execute + +```bash +git checkout obj +``` + +Make sure you have installed [Isaac Gym](https://developer.nvidia.com/isaac-gym) and [pointnet2](https://github.com/daveredrum/Pointnet2.ScanNet). See [Setup](#setup) section. + +### Train + +- Train with single gpu (one gpu is enough) + + ```bash + bash scripts/grasp_gen_ur/train.sh ${EXP_NAME} + ``` + +### Sample (Qualitative Visualization) + +```bash +bash scripts/grasp_gen_ur/sample.sh ${CKPT} [OPT] +# e.g., bash scripts/grasp_gen_ur/sample.sh ./outputs/2022-11-15_18-07-50_GPUR_l1_pn2_T100/ OPT +``` + +- `[OPT]` is optional for optimization-guided sampling. + +### Test (Quantitative Evaluation) + +You first need to run `scripts/grasp_gen_ur/sample.sh` to sample some results. Then we will compute quantitative metrics with these sampled results. + +```bash +bash scripts/grasp_gen_ur/test.sh ${EVAL_DIR} ${DATASET_DIR} +# e.g., bash scripts/grasp_gen_ur/test.sh outputs/2022-11-15_18-07-50_GPUR_l1_pn2_T100/eval/final/2023-04-20_13-06-44 YOUR_PATH/data/MultiDex_UR +``` ## Task-4: Path Planning in 3D Scenes @@ -177,6 +226,7 @@ coming soon. ```bash bash scripts/path_planning/plan.sh ${CKPT} +# e.g., bash scripts/path_planning/plan.sh ./outputs/2022-11-25_20-57-28_Path_ddm4_LR1e-4_E100_REL/ ``` ### Sample (Qualitative Visualization) @@ -192,7 +242,34 @@ bash scripts/path_planning/sample.sh ${CKPT} [OPT] [PLA] ## Task-5: Motion Planning for Robot Arms -coming soon. +To run this code, you first need to change the git branch to `obj` by execute + +```bash +git checkout obj +``` + +Make sure you have installed [Isaac Gym](https://developer.nvidia.com/isaac-gym) and [pointnet2](https://github.com/daveredrum/Pointnet2.ScanNet). See [Setup](#setup) section. + +### Train + +- Train with single gpu + + ```bash + bash scripts/franka_planning/train.sh ${EXP_NAME} + ``` + +- Train with 4 GPUs (modify `scripts/path_planning/train_ddm.sh` to specify the visible GPUs) + + ```bash + bash scripts/franka_planning/train_ddm.sh ${EXP_NAME} + ``` + +### Test (Quantitative Evaluation) + +```bash +bash scripts/franka_planning/plan.sh ${CKPT} +# e.g., bash scripts/franka_planning/plan.sh outputs/2022-11-11_14-28-30_FK2Plan_ptr_T30_4/ +``` ## Citation diff --git a/scripts/grasp_gen_ur/test.sh b/scripts/grasp_gen_ur/test.sh index a63c5fc..c3aea1a 100755 --- a/scripts/grasp_gen_ur/test.sh +++ b/scripts/grasp_gen_ur/test.sh @@ -8,4 +8,4 @@ DATASET_DIR=$2 python ./scripts/grasp_gen_ur/test.py --eval_dir=${EVAL_DIR} \ --dataset_dir=${DATASET_DIR} \ --stability_config='envs/tasks/grasp_test_force.yaml' \ - --seed=42 \ + --seed=42 From eccee61d4b4d2c16358de967813d3d99ebdddce0 Mon Sep 17 00:00:00 2001 From: Silverster98 <1539168414@qq.com> Date: Fri, 21 Apr 2023 15:08:46 +0800 Subject: [PATCH 7/7] Update readme --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 7d908ec..41309b3 100644 --- a/README.md +++ b/README.md @@ -172,7 +172,7 @@ bash scripts/motion_gen/sample.sh ${CKPT} [OPT] ## Task-3: Dexterous Grasp Generation for 3D Objects -To run this code, you first need to change the git branch to `obj` by execute +To run this code, you first need to change the git branch to `obj` by executing ```bash git checkout obj @@ -242,7 +242,7 @@ bash scripts/path_planning/sample.sh ${CKPT} [OPT] [PLA] ## Task-5: Motion Planning for Robot Arms -To run this code, you first need to change the git branch to `obj` by execute +To run this code, you first need to change the git branch to `obj` by executing ```bash git checkout obj