├── scenecollisionnet ├── __init__.py ├── version.py ├── policy │ ├── __init__.py │ ├── constants.py │ ├── utils.py │ ├── robot.py │ ├── urdf.py │ └── collision_checker.py ├── collision_models │ ├── __init__.py │ ├── fcl.py │ ├── sdf.py │ ├── pointnet_grid.py │ └── collision_nets.py ├── trainer.py ├── utils.py ├── scene.py └── datasets.py ├── .gitattributes ├── LICENSE.pdf ├── data └── panda │ ├── meshes │ └── collision │ │ ├── hand.stl │ │ ├── finger.stl │ │ ├── link0.stl │ │ ├── link1.stl │ │ ├── link2.stl │ │ ├── link3.stl │ │ ├── link4.stl │ │ ├── link5.stl │ │ ├── link6.stl │ │ └── link7.stl │ └── panda.urdf ├── .gitignore ├── scripts ├── install_manifold.sh ├── train_robotcollisionnet_docker.sh ├── benchmark_baseline_docker.sh ├── benchmark_robotcollisionnet_docker.sh ├── train_scenecollisionnet_docker.sh ├── benchmark_scenecollisionnet_docker.sh └── download_weights.sh ├── .dockerignore ├── .flake8 ├── cfg ├── benchmark_robotcollisionnet.yaml ├── benchmark_scenecollisionnet.yaml ├── train_robotcollisionnet.yaml ├── benchmark_baseline.yaml └── train_scenecollisionnet.yaml ├── .gitmodules ├── resources └── urdf.template ├── Dockerfile ├── tools ├── generate_acronym_dataset.py ├── rollout_policy.py ├── loss_plots.py ├── generate_mesh_dataset.py ├── train_robotcollisionnet.py ├── benchmark_robotcollisionnet.py ├── train_scenecollisionnet.py ├── benchmark_baseline.py └── benchmark_scenecollisionnet.py ├── setup.py └── README.md /scenecollisionnet/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /scenecollisionnet/version.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.0.1" 2 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | *.pth.tar filter=lfs diff=lfs merge=lfs -text 2 | -------------------------------------------------------------------------------- /LICENSE.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/LICENSE.pdf -------------------------------------------------------------------------------- /scenecollisionnet/policy/__init__.py: -------------------------------------------------------------------------------- 1 | from .robot_environment import RobotEnvironment 2 | from .mppi import MPPIPolicy 3 | -------------------------------------------------------------------------------- /data/panda/meshes/collision/hand.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/hand.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/finger.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/finger.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/link0.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/link0.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/link1.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/link1.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/link2.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/link2.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/link3.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/link3.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/link4.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/link4.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/link5.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/link5.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/link6.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/link6.stl -------------------------------------------------------------------------------- /data/panda/meshes/collision/link7.stl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NVlabs/SceneCollisionNet/HEAD/data/panda/meshes/collision/link7.stl -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | *.egg-info 3 | *.pyc 4 | 5 | cached_dataset 6 | datasets 7 | models 8 | benchmark 9 | weights 10 | 11 | .vscode 12 | 13 | extern/isaacgym -------------------------------------------------------------------------------- /scripts/install_manifold.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Install Manifold 4 | cd extern/Manifold 5 | mkdir build && cd build 6 | cmake .. -DCMAKE_BUILD_TYPE=Release 7 | make -j4 8 | -------------------------------------------------------------------------------- /scripts/train_robotcollisionnet_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | cd /SceneCollisionNet 4 | python3 tools/train_robotcollisionnet.py --cfg /cfg/train_robotcollisionnet.yaml $@ 5 | -------------------------------------------------------------------------------- /scripts/benchmark_baseline_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | cd /SceneCollisionNet 4 | PYOPENGL_PLATFORM=egl python3 tools/benchmark_baseline.py --cfg /cfg/benchmark_baseline.yaml $@ 5 | -------------------------------------------------------------------------------- /scripts/benchmark_robotcollisionnet_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | cd /SceneCollisionNet 4 | python3 tools/benchmark_robotcollisionnet.py --cfg /cfg/benchmark_robotcollisionnet.yaml $@ 5 | -------------------------------------------------------------------------------- /scripts/train_scenecollisionnet_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | cd /SceneCollisionNet 4 | PYOPENGL_PLATFORM=egl python3 tools/train_scenecollisionnet.py --cfg /cfg/train_scenecollisionnet.yaml $@ 5 | -------------------------------------------------------------------------------- /scripts/benchmark_scenecollisionnet_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | cd /SceneCollisionNet 4 | PYOPENGL_PLATFORM=egl python3 tools/benchmark_scenecollisionnet.py --cfg /cfg/benchmark_scenecollisionnet.yaml $@ 5 | -------------------------------------------------------------------------------- /.dockerignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | *.egg-info 3 | *.pyc 4 | **/build 5 | **/*.so 6 | 7 | cfg 8 | cached_dataset 9 | datasets 10 | models 11 | benchmark 12 | weights 13 | 14 | .vscode 15 | 16 | README.md 17 | Dockerfile 18 | .dockerignore -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | max-line-length = 119 3 | ignore = 4 | W503, E203, E231 # conflicts with `black` formatting 5 | E402 # from headless matplotlib use 6 | per-file-ignores = __init__.py: F401 7 | exclude = .git, cfg/, .eggs/, extern/ -------------------------------------------------------------------------------- /cfg/benchmark_robotcollisionnet.yaml: -------------------------------------------------------------------------------- 1 | output: /benchmark 2 | device: cuda:0 3 | # num_workers: 4 # Defaults to os.cpu_count() if not specified. 4 | iterations: 250 5 | 6 | model: 7 | name: self_coll_nn 8 | path: /models 9 | 10 | dataset: 11 | batch_size: 512 -------------------------------------------------------------------------------- /scenecollisionnet/collision_models/__init__.py: -------------------------------------------------------------------------------- 1 | from .collision_nets import RobotCollisionNet, SceneCollisionNet 2 | from .fcl import FCLMultiSceneCollisionChecker 3 | from .pointnet_grid import PointNetGrid 4 | from .sdf import SDFObjCollisionChecker, SDFSceneCollisionChecker 5 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "extern/pointnet2"] 2 | path = extern/pointnet2 3 | url = ../../mjd3/pointnet2 4 | [submodule "extern/Manifold"] 5 | path = extern/Manifold 6 | url = ../../hjwdzh/Manifold 7 | [submodule "extern/kaolin"] 8 | path = extern/kaolin 9 | url = ../../mjd3/kaolin 10 | [submodule "extern/tracikpy"] 11 | path = extern/tracikpy 12 | url = ../../mjd3/tracikpy 13 | -------------------------------------------------------------------------------- /cfg/benchmark_scenecollisionnet.yaml: -------------------------------------------------------------------------------- 1 | output: /benchmark 2 | device: cuda:0 3 | # num_workers: 4 # Defaults to os.cpu_count() if not specified. Note: set to 0 if vis=1 4 | vis: 1 # output visualizations 5 | iterations: 1000 6 | 7 | model: 8 | name: scene_coll_net 9 | path: /models 10 | 11 | dataset: 12 | meshes: /dataset 13 | batch_size: 1 14 | query_size: 64 15 | trajectories: 1 -------------------------------------------------------------------------------- /cfg/train_robotcollisionnet.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: self_coll_nn 3 | path: /models # For docker container 4 | 5 | trainer: 6 | device: cuda:0 7 | # num_workers: 4 # defaults to os.cpu_count() if not specified 8 | train_iterations: 100 9 | val_iterations: 25 10 | epochs: 100 11 | lr: 0.001 12 | momentum: 0.9 13 | 14 | dataset: 15 | robot_urdf: data/panda/panda.urdf 16 | batch_size: 256 -------------------------------------------------------------------------------- /scripts/download_weights.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | export fileid=1lz8N2kMsfgrZf2PVmbpq5mxBzqmQHUdL 4 | export filename=weights.zip 5 | 6 | wget --save-cookies cookies.txt 'https://docs.google.com/uc?export=download&id='$fileid -O- \ 7 | | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1/p' > confirm.txt 8 | 9 | wget --load-cookies cookies.txt -O $filename \ 10 | 'https://docs.google.com/uc?export=download&id='$fileid'&confirm='$( 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | -------------------------------------------------------------------------------- /cfg/benchmark_baseline.yaml: -------------------------------------------------------------------------------- 1 | output: /benchmark 2 | vis: 0 3 | # num_workers: 4 # Defaults to os.cpu_count() if not specified. 4 | iterations: 1000 5 | 6 | model: 7 | type: sdf_scene 8 | device: cuda:0 9 | 10 | dataset: 11 | meshes: /dataset 12 | batch_size: 1 13 | query_size: 2048 14 | trajectories: 16 15 | bounds: 16 | - [-0.5, -0.8, -0.06] 17 | - [0.5, 0.8, 0.3] 18 | intrinsics: 19 | frame: camera 20 | fx: 616.36529541 21 | fy: 616.20294189 22 | cx: 310.25881958 23 | cy: 236.59980774 24 | skew: 0.0 25 | width: 640 26 | height: 480 27 | 28 | extrinsics: 29 | azimuth: [-0.2, 0.2] 30 | elevation: [0.6, 1.0] 31 | radius: [1.5, 2.0] -------------------------------------------------------------------------------- /scenecollisionnet/policy/constants.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import trimesh.transformations as tra 3 | 4 | GRASP_FRAME_INIT_OFFSET = tra.quaternion_matrix([0.707, 0, 0, -0.707]) 5 | GRASP_FRAME_INIT_OFFSET[:3, 3] = [0, 0, 0.0] 6 | GRASP_FRAME_FINAL_OFFSET = tra.quaternion_matrix([0.707, 0, 0, -0.707]) 7 | GRASP_FRAME_FINAL_OFFSET[:3, 3] = [0, 0, 0.1] 8 | 9 | TABLE_LABEL = 0 10 | ROBOT_LABEL = np.iinfo(np.uint8).max 11 | DEPTH_CLIP_RANGE = 3.0 12 | 13 | ROBOT_Q_INIT = np.array( 14 | [ 15 | -1.22151887, 16 | -1.54163973, 17 | -0.3665906, 18 | -2.23575787, 19 | 0.5335327, 20 | 1.04913162, 21 | -0.14688508, 22 | 0.04, 23 | 0.04, 24 | ] 25 | ) 26 | -------------------------------------------------------------------------------- /cfg/train_scenecollisionnet.yaml: -------------------------------------------------------------------------------- 1 | model: 2 | name: scene_coll_net 3 | type: SceneCollisionNet # Options: SceneCollisionNet or PointNetGrid 4 | path: /models 5 | bounds: 6 | - [-0.5, -0.8, -0.06] 7 | - [0.5, 0.8, 0.3] 8 | vox_size: [0.125, 0.1, 0.09] 9 | 10 | trainer: 11 | device: cuda:0 12 | num_workers: 4 # defaults to os.cpu_count() if not specified 13 | train_iterations: 1000 14 | val_iterations: 250 15 | epochs: 1000 16 | lr: 0.001 17 | momentum: 0.9 18 | loss_pct: 0.1 19 | 20 | dataset: 21 | meshes: /dataset 22 | batch_size: 1 23 | query_size: 2048 24 | n_obj_points: 1024 25 | n_scene_points: 32768 26 | rotations: 1 27 | trajectories: 64 28 | 29 | camera: 30 | intrinsics: 31 | frame: camera 32 | fx: 616.36529541 33 | fy: 616.20294189 34 | cx: 310.25881958 35 | cy: 236.59980774 36 | skew: 0.0 37 | width: 640 38 | height: 480 39 | 40 | extrinsics: 41 | azimuth: [-0.2, 0.2] 42 | elevation: [0.6, 1.0] 43 | radius: [1.5, 2.0] 44 | -------------------------------------------------------------------------------- /scenecollisionnet/policy/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import trimesh.transformations as tra 5 | from isaacgym import gymapi 6 | 7 | 8 | def from_rpy(a, e, th): 9 | rt = tra.euler_matrix(a, e, th) 10 | q = tra.quaternion_from_matrix(rt) 11 | q = np.roll(q, -1) 12 | return gymapi.Quat(q[0], q[1], q[2], q[3]) 13 | 14 | 15 | def gym_pose_to_matrix(pose): 16 | q = [pose["r"][3], pose["r"][0], pose["r"][1], pose["r"][2]] 17 | trans = tra.quaternion_matrix(q) 18 | trans[:3, 3] = [pose["p"][i] for i in range(3)] 19 | 20 | return trans 21 | 22 | 23 | def write_urdf( 24 | obj_name, 25 | obj_path, 26 | output_folder, 27 | ): 28 | content = open("resources/urdf.template").read() 29 | content = content.replace("NAME", obj_name) 30 | content = content.replace("MEAN_X", "0.0") 31 | content = content.replace("MEAN_Y", "0.0") 32 | content = content.replace("MEAN_Z", "0.0") 33 | content = content.replace("SCALE", "1.0") 34 | content = content.replace("COLLISION_OBJ", obj_path) 35 | content = content.replace("GEOMETRY_OBJ", obj_path) 36 | urdf_path = os.path.abspath( 37 | os.path.join(output_folder, obj_name + ".urdf") 38 | ) 39 | if not os.path.exists(output_folder): 40 | os.makedirs(output_folder) 41 | open(urdf_path, "w").write(content) 42 | return urdf_path 43 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM nvidia/cudagl:10.2-devel-ubuntu18.04 2 | 3 | # env variables for tzdata install 4 | ARG DEBIAN_FRONTEND=noninteractive 5 | ENV TZ=America/Vancouver 6 | 7 | RUN apt-get update -y && \ 8 | apt-get install -y \ 9 | python3-pip \ 10 | python3-dev \ 11 | ninja-build \ 12 | libcudnn8=8.1.1.33-1+cuda10.2 \ 13 | libcudnn8-dev=8.1.1.33-1+cuda10.2 \ 14 | libsm6 libxext6 libxrender-dev \ 15 | freeglut3-dev \ 16 | liboctomap-dev libfcl-dev \ 17 | gifsicle libfreetype6-dev libpng-dev && \ 18 | rm -rf /var/lib/apt/lists/* 19 | 20 | # RUN add-apt-repository ppa:deadsnakes/ppa 21 | # RUN apt-get install -y python3.8 python3.8-dev python3-pip 22 | 23 | # Install pytorch and pytorch scatter deps (need to match CUDA version) 24 | RUN pip3 install --no-cache-dir \ 25 | torch==1.8.1+cu102 -f https://download.pytorch.org/whl/torch_stable.html \ 26 | torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cu102.html 27 | 28 | # Build and install pointnet2 dep 29 | RUN mkdir -p SceneCollisionNet 30 | COPY [ "./extern", "SceneCollisionNet/extern" ] 31 | RUN TORCH_CUDA_ARCH_LIST="6.0 6.1 6.2 7.0 7.5 8.0 8.6+PTX" python3 -m pip install --no-cache-dir SceneCollisionNet/extern/pointnet2 32 | 33 | # Build and install kaolin dep (for SDF baselines) 34 | RUN IGNORE_TORCH_VER=1 FORCE_CUDA=1 TORCH_CUDA_ARCH_LIST="6.0 6.1 6.2 7.0 7.5 8.0 8.6+PTX" python3 -m pip install --no-cache-dir SceneCollisionNet/extern/kaolin 35 | 36 | # Install all other python deps 37 | COPY [ "./setup.py", "SceneCollisionNet/setup.py" ] 38 | RUN python3 -m pip install --no-cache-dir `python3 SceneCollisionNet/setup.py --list-bench` 39 | 40 | # Install repo 41 | COPY [ ".", "SceneCollisionNet/" ] 42 | RUN python3 -m pip install --no-cache-dir SceneCollisionNet/ 43 | -------------------------------------------------------------------------------- /scenecollisionnet/policy/robot.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from .urdf import TorchURDF 5 | 6 | 7 | # Thin wrapper around urdfpy class that adds cfg sampling 8 | class Robot: 9 | def __init__(self, urdf_path, ee_name, device=None): 10 | if device is None: 11 | device = torch.device("cpu") 12 | self._robot = TorchURDF.load(urdf_path, device=device) 13 | self.dof = len(self.joint_names) 14 | self._ee_name = ee_name 15 | self.device = device 16 | self.min_joints = torch.tensor( 17 | [ 18 | j.limit.lower 19 | for j in self._robot.joints 20 | if j.joint_type != "fixed" 21 | ], 22 | device=self.device, 23 | ).unsqueeze_(0) 24 | self.max_joints = torch.tensor( 25 | [ 26 | j.limit.upper 27 | for j in self._robot.joints 28 | if j.joint_type != "fixed" 29 | ], 30 | device=self.device, 31 | ).unsqueeze_(0) 32 | 33 | @property 34 | def joint_names(self): 35 | return [j.name for j in self._robot.joints if j.joint_type != "fixed"] 36 | 37 | @property 38 | def links(self): 39 | return self._robot.links 40 | 41 | @property 42 | def mesh_links(self): 43 | return [ 44 | link 45 | for link in self._robot.links 46 | if link.collision_mesh is not None 47 | ] 48 | 49 | @property 50 | def link_map(self): 51 | return self._robot.link_map 52 | 53 | @property 54 | def link_poses(self): 55 | return self._link_poses 56 | 57 | @property 58 | def ee_pose(self): 59 | return self.link_poses[self.link_map[self._ee_name]] 60 | 61 | def set_joint_cfg(self, q): 62 | if q is not None: 63 | if isinstance(q, np.ndarray): 64 | q = torch.from_numpy(q).float().to(self.device) 65 | if q.device != self.device: 66 | q = q.to(self.device) 67 | if q.ndim == 1: 68 | q = q.reshape(1, -1) 69 | if q.ndim > 2: 70 | raise ValueError("Tensor is wrong shape, must have 2 dims") 71 | link_poses = self._robot.link_fk_batch(q[:, : self.dof]) 72 | self._link_poses = link_poses 73 | else: 74 | self._link_poses = None 75 | 76 | def sample_cfg(self, num=1): 77 | alpha = torch.rand(num, self.dof, device=self.device) 78 | return alpha * self.min_joints + (1 - alpha) * self.max_joints 79 | -------------------------------------------------------------------------------- /tools/generate_acronym_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | import h5py 5 | from tqdm import tqdm 6 | 7 | from scenecollisionnet.utils import ( 8 | MeshLoader, 9 | ProcessKillingExecutor, 10 | process_mesh, 11 | ) 12 | 13 | 14 | def process_mesh_timeout(*args, **kwargs): 15 | pass 16 | 17 | 18 | if __name__ == "__main__": 19 | 20 | parser = argparse.ArgumentParser(description="Generate mesh dataset") 21 | parser.add_argument( 22 | "meshes_dir", 23 | type=str, 24 | help="path to the Shapenet mesh dataset directory", 25 | ) 26 | parser.add_argument( 27 | "grasps_dir", 28 | type=str, 29 | help="path to ACRONYM grasp dataset directory of hdf5 files", 30 | ) 31 | parser.add_argument( 32 | "output_dir", 33 | type=str, 34 | help="output dataset path", 35 | ) 36 | 37 | args = parser.parse_args() 38 | mesh_dir = args.meshes_dir 39 | grasps_dir = args.grasps_dir 40 | out_dir = args.output_dir 41 | 42 | if not os.path.exists(mesh_dir): 43 | print("Input directory does not exist!") 44 | mesh_loader = MeshLoader(mesh_dir) 45 | 46 | if not os.path.exists(out_dir): 47 | os.makedirs(out_dir) 48 | out_ds_file = os.path.join(out_dir, "object_info.hdf5") 49 | 50 | inputs = [] 51 | for f in tqdm(os.listdir(grasps_dir), desc="Generating Inputs"): 52 | mc, mk, ms = os.path.splitext(f)[0].split("_") 53 | try: 54 | in_path = mesh_loader.get_path(mk) 55 | except ValueError: 56 | continue 57 | out_path = os.path.join( 58 | out_dir, 59 | mc, 60 | os.path.splitext(os.path.basename(in_path))[0] + ".obj", 61 | ) 62 | inputs.append( 63 | (in_path, out_path, float(ms), os.path.join(grasps_dir, f)) 64 | ) 65 | 66 | executor = ProcessKillingExecutor(max_workers=8) 67 | generator = executor.map( 68 | process_mesh, 69 | inputs, 70 | timeout=120, 71 | callback_timeout=process_mesh_timeout, 72 | ) 73 | 74 | categories = {} 75 | with h5py.File(out_ds_file, "a") as f: 76 | f.create_group("meshes") 77 | f.create_group("categories") 78 | for mesh_info in tqdm( 79 | generator, total=len(inputs), desc="Processing Meshes" 80 | ): 81 | if mesh_info is not None: 82 | mk, minfo = mesh_info 83 | f["meshes"].create_group(mk) 84 | for key in minfo: 85 | f["meshes"][mk][key] = minfo[key] 86 | 87 | if minfo["category"] not in categories: 88 | categories[minfo["category"]] = [mk] 89 | else: 90 | categories[minfo["category"]].append(mk) 91 | 92 | for c in categories: 93 | f["categories"][c] = categories[c] 94 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | """ 2 | Setup of codebase 3 | Author: Michael Danielczuk 4 | """ 5 | import os 6 | import sys 7 | 8 | from setuptools import find_packages, setup 9 | 10 | root_dir = os.path.dirname(os.path.realpath(__file__)) 11 | 12 | requirements_default = set( 13 | [ 14 | "numpy", # Basic math/array utilities 15 | "tqdm", # Progress bars 16 | "h5py", # Reading/writing dataset info 17 | "trimesh", # Mesh loading/utilities 18 | ] 19 | ) 20 | 21 | requirements_train = set( 22 | [ 23 | "pyrender", # Scene rendering 24 | "python-fcl", # Scene GT collisions 25 | "urdfpy", # Robot FK 26 | "torch>=1.5", # Training/Benchmarking 27 | "torch-scatter", # For collision models 28 | f"pointnet2 @ file://localhost{root_dir}/extern/pointnet2", # Network modules 29 | "autolab_core", # used for loading cfg files and Image classes 30 | ] 31 | ) 32 | 33 | requirements_bench = requirements_train.union( 34 | [ 35 | "matplotlib", # Plotting 36 | "seaborn", # Fancy plotting 37 | "pandas", # Loss plotting made easy 38 | "natsort", # Loss plotting sort files 39 | "pygifsicle", # Benchmarking GIFs 40 | "imageio", # Benchmarking GIFs 41 | ] 42 | ) 43 | 44 | requirements_policy = requirements_train.union( 45 | [ 46 | f"tracikpy @ file://localhost{root_dir}/extern/tracikpy", # Fast IK solutions 47 | "KNN-CUDA @ https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl", # CUDA impl of KNN 48 | f"isaacgym @ file://localhost{root_dir}/extern/isaacgym/python", # Isaac Gym Simulation 49 | ] 50 | ) 51 | 52 | # if someone wants to output a requirements file 53 | # `python setup.py --list-train > requirements.txt` 54 | if "--list-train" in sys.argv: 55 | print("\n".join(requirements_train.union(requirements_default))) 56 | exit() 57 | elif "--list-bench" in sys.argv: 58 | print("\n".join(requirements_bench.union(requirements_default))) 59 | exit() 60 | 61 | # load __version__ without importing anything 62 | version_file = os.path.join( 63 | os.path.dirname(__file__), "scenecollisionnet/version.py" 64 | ) 65 | with open(version_file, "r") as f: 66 | # use eval to get a clean string of version from file 67 | __version__ = eval(f.read().strip().split("=")[-1]) 68 | 69 | # load README.md as long_description 70 | long_description = "" 71 | if os.path.exists("README.md"): 72 | with open("README.md", "r") as f: 73 | long_description = f.read() 74 | 75 | setup( 76 | name="scenecollisionnet", 77 | version=__version__, 78 | description="SceneCollisionNet ICRA21 Paper Code", 79 | long_description=long_description, 80 | author="Michael Danielczuk", 81 | author_email="mdanielczuk@berkeley.edu", 82 | license="MIT Software License", 83 | url="https://github.com/mjd3/SceneCollisionNet", 84 | keywords="robotics computer vision", 85 | classifiers=[ 86 | "License :: OSI Approved :: MIT Software License", 87 | "Programming Language :: Python", 88 | "Programming Language :: Python :: 3.6", 89 | "Programming Language :: Python :: 3.7", 90 | "Programming Language :: Python :: 3.8", 91 | "Natural Language :: English", 92 | "Topic :: Scientific/Engineering", 93 | ], 94 | packages=find_packages(), 95 | install_requires=list(requirements_default), 96 | extras_require={ 97 | "train": list(requirements_train), 98 | "bench": list(requirements_bench), 99 | "policy": list(requirements_policy), 100 | }, 101 | ) 102 | -------------------------------------------------------------------------------- /tools/rollout_policy.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import random 3 | 4 | import numpy as np 5 | from autolab_core import Logger 6 | 7 | from scenecollisionnet.policy import MPPIPolicy, RobotEnvironment 8 | 9 | 10 | def rollout(args): 11 | logger = Logger.get_logger("tools/rollout_policy.py") 12 | 13 | if args.seed is not None: 14 | np.random.seed(args.seed) 15 | random.seed(args.seed) 16 | 17 | env = RobotEnvironment(args) 18 | policy = MPPIPolicy( 19 | num_proc=args.mppi_nproc, 20 | num_path=args.mppi_num_rollouts, 21 | horizon=args.mppi_horizon, 22 | collision_steps=args.mppi_collision_steps, 23 | max_step=args.mppi_max_step, 24 | lift_height=args.mppi_lift_height, 25 | lift_steps=args.mppi_lift_steps, 26 | transition_threshold=args.mppi_transition_threshold, 27 | noise_std=args.mppi_q_noise_std, 28 | self_coll_nn=args.self_coll_nn, 29 | scene_coll_nn=args.scene_coll_nn, 30 | cam_type=args.mppi_cam, 31 | device=args.compute_device_id, 32 | log_file=args.log_file, 33 | ) 34 | 35 | input("press any key to start") 36 | o, info = env.step() 37 | while True: 38 | a = policy.get_action(o) 39 | next_o, info = env.step(a) 40 | if info["done"]: 41 | break 42 | 43 | o = next_o 44 | 45 | logger.info(f"Finished {2 * args.num_objects} attempts. Exiting...") 46 | 47 | 48 | if __name__ == "__main__": 49 | parser = argparse.ArgumentParser( 50 | description="", 51 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, 52 | ) 53 | parser.add_argument("--compute-device-id", type=int, default=0) 54 | parser.add_argument("--graphics-device-id", type=int, default=1) 55 | parser.add_argument( 56 | "--robot-asset-root", 57 | type=str, 58 | default="extern/isaacgym/assets/", 59 | ) 60 | parser.add_argument( 61 | "--robot-asset-file", 62 | type=str, 63 | default="urdf/franka_description/robots/franka_panda.urdf", 64 | ) 65 | parser.add_argument("--ee-link-name", type=str, default="panda_hand") 66 | parser.add_argument("--headless", action="store_true", default=False) 67 | parser.add_argument( 68 | "--dataset-root", 69 | type=str, 70 | default="datasets/shapenet", 71 | ) 72 | parser.add_argument( 73 | "--urdf-cache", type=str, default="datasets/scene_cache" 74 | ) 75 | parser.add_argument("--num-objects", type=int, default=10) 76 | parser.add_argument("--camera-height", type=int, default=480) 77 | parser.add_argument("--camera-width", type=int, default=640) 78 | 79 | parser.add_argument("--seed", type=int) 80 | parser.add_argument("--control-frequency", type=float, default=20) 81 | parser.add_argument("--npoints", type=int, default=8192) 82 | parser.add_argument("--mppi-nproc", type=int, default=20) 83 | parser.add_argument("--mppi-horizon", type=int, default=40) 84 | parser.add_argument("--mppi-num-rollouts", type=int, default=200) 85 | parser.add_argument("--mppi-q-noise-std", type=float, default=0.08) 86 | parser.add_argument("--mppi-max-step", type=float, default=0.05) 87 | parser.add_argument("--mppi-lift-height", type=float, default=0.6) 88 | parser.add_argument("--mppi-lift-steps", type=int, default=20) 89 | parser.add_argument("--mppi-collision-steps", type=int, default=10) 90 | parser.add_argument( 91 | "--mppi-transition-threshold", type=float, default=0.05 92 | ) 93 | parser.add_argument("--mppi-cam", type=str, default="ws") 94 | parser.add_argument("--scene-coll-nn", type=str) 95 | parser.add_argument("--self-coll-nn", type=str) 96 | parser.add_argument("--log-file", type=str) 97 | 98 | args = parser.parse_args() 99 | rollout(args) 100 | -------------------------------------------------------------------------------- /tools/loss_plots.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | 4 | import matplotlib 5 | import pandas as pd 6 | from natsort import natsorted 7 | 8 | matplotlib.use("agg") 9 | import matplotlib.pyplot as plt 10 | import seaborn as sns 11 | 12 | sns.set() 13 | 14 | if __name__ == "__main__": 15 | 16 | parser = argparse.ArgumentParser(description="Make loss plots for model") 17 | parser.add_argument( 18 | "csv_files", 19 | type=str, 20 | nargs="+", 21 | help="path to the csv file with losses", 22 | ) 23 | parser.add_argument( 24 | "-o", "--output_path", type=str, help="path to store output images" 25 | ) 26 | parser.add_argument( 27 | "-s", "--sample", type=int, default=1, help="downsample plotting rate" 28 | ) 29 | args = parser.parse_args() 30 | 31 | output_path = ( 32 | os.path.dirname(args.csv_files[0]) 33 | if args.output_path is None 34 | else args.output_path 35 | ) 36 | 37 | legend = len(args.csv_files) > 1 38 | 39 | f1, axes1 = plt.subplots(2, 2, figsize=(9, 6)) 40 | f2, axes2 = plt.subplots(2, 2, figsize=(9, 6)) 41 | for csvf in natsorted(args.csv_files): 42 | 43 | data = pd.read_csv(csvf) 44 | data = data[data["iteration"] % args.sample == 0] 45 | 46 | # Make and save training plots 47 | label = os.path.dirname(csvf).split("/")[-1] if legend else None 48 | sns.lineplot( 49 | x="epoch", 50 | y="train/loss", 51 | data=data, 52 | ax=axes1[0, 0], 53 | label=label, 54 | legend=False, 55 | ) 56 | sns.lineplot( 57 | x="epoch", 58 | y="train/acc", 59 | data=data, 60 | ax=axes1[0, 1], 61 | label=label, 62 | legend=False, 63 | ) 64 | sns.lineplot( 65 | x="epoch", 66 | y="train/f1", 67 | data=data, 68 | ax=axes1[1, 0], 69 | label=label, 70 | legend=False, 71 | ) 72 | sns.lineplot( 73 | x="epoch", 74 | y="train/ap", 75 | data=data, 76 | ax=axes1[1, 1], 77 | label=label, 78 | legend=False, 79 | ) 80 | 81 | # Make and save validation plots 82 | sns.lineplot( 83 | x="epoch", 84 | y="valid/loss", 85 | data=data, 86 | ax=axes2[0, 0], 87 | label=label, 88 | legend=False, 89 | ) 90 | sns.lineplot( 91 | x="epoch", 92 | y="valid/acc", 93 | data=data, 94 | ax=axes2[0, 1], 95 | label=label, 96 | legend=False, 97 | ) 98 | sns.lineplot( 99 | x="epoch", 100 | y="valid/f1", 101 | data=data, 102 | ax=axes2[1, 0], 103 | label=label, 104 | legend=False, 105 | ) 106 | sns.lineplot( 107 | x="epoch", 108 | y="valid/ap", 109 | data=data, 110 | ax=axes2[1, 1], 111 | label=label, 112 | legend=False, 113 | ) 114 | 115 | if legend: 116 | f1.legend( 117 | *axes1[0, 0].get_legend_handles_labels(), 118 | loc="center left", 119 | bbox_to_anchor=(1.0, 0.5) 120 | ) 121 | f2.legend( 122 | *axes2[0, 0].get_legend_handles_labels(), 123 | loc="center left", 124 | bbox_to_anchor=(1.0, 0.5) 125 | ) 126 | 127 | f1.tight_layout() 128 | f2.tight_layout() 129 | f1.savefig( 130 | os.path.join(output_path, "training_plots.png"), bbox_inches="tight" 131 | ) 132 | f2.savefig( 133 | os.path.join(output_path, "validation_plots.png"), bbox_inches="tight" 134 | ) 135 | -------------------------------------------------------------------------------- /tools/generate_mesh_dataset.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | from shutil import copyfile 4 | 5 | import h5py 6 | import numpy as np 7 | from tqdm import tqdm 8 | 9 | from scenecollisionnet.utils import ( 10 | MeshLoader, 11 | ProcessKillingExecutor, 12 | process_mesh, 13 | ) 14 | 15 | 16 | def process_mesh_timeout(*args, **kwargs): 17 | pass 18 | 19 | 20 | if __name__ == "__main__": 21 | parser = argparse.ArgumentParser(description="Generate mesh dataset") 22 | parser.add_argument( 23 | "meshes_dir", 24 | type=str, 25 | help="path to the mesh dataset directory", 26 | ) 27 | parser.add_argument( 28 | "--dataset_file", type=str, help="path to existing dataset hdf5" 29 | ) 30 | parser.add_argument( 31 | "output_dir", 32 | type=str, 33 | help="output dataset path", 34 | ) 35 | 36 | args = parser.parse_args() 37 | mesh_dir = args.meshes_dir 38 | dataset_file = args.dataset_file 39 | out_dir = args.output_dir 40 | 41 | if not os.path.exists(mesh_dir): 42 | print("Input directory does not exist!") 43 | mesh_loader = MeshLoader(mesh_dir) 44 | 45 | if not os.path.exists(out_dir): 46 | os.makedirs(out_dir) 47 | 48 | out_ds_file = os.path.join(out_dir, "object_info.hdf5") 49 | if dataset_file is not None: 50 | copyfile(dataset_file, out_ds_file) 51 | 52 | mesh_keys, mesh_scales, mesh_cats = [], [], [] 53 | with h5py.File(out_ds_file, "r") as f: 54 | for mk in f["meshes"]: 55 | mesh_keys.append("_".join(mk.split("_")[:-1])) 56 | mesh_scales.append(f["meshes"][mk]["scale"][()]) 57 | mesh_cats.append(f["meshes"][mk]["category"].asstr()[()]) 58 | mesh_stps = np.zeros(len(mesh_keys), dtype=bool) 59 | else: 60 | mesh_keys = mesh_loader.meshes() 61 | mesh_scales = np.ones(len(mesh_keys)) 62 | mesh_cats = [ 63 | mk.split("~")[0] if len(mk.split("~")) > 1 else "" 64 | for mk in mesh_keys 65 | ] 66 | mesh_stps = np.ones(len(mesh_keys), dtype=bool) 67 | 68 | inputs = [] 69 | for mk, ms, mc, mp in tqdm( 70 | zip(mesh_keys, mesh_scales, mesh_cats, mesh_stps), 71 | total=len(mesh_keys), 72 | desc="Generating Inputs", 73 | ): 74 | try: 75 | in_path = mesh_loader.get_path(mk) 76 | except ValueError: 77 | continue 78 | out_path = os.path.abspath( 79 | os.path.join( 80 | out_dir, 81 | mc, 82 | os.path.splitext(os.path.basename(in_path))[0] + ".obj", 83 | ) 84 | ) 85 | inputs.append((in_path, out_path, ms, None, mp)) 86 | 87 | executor = ProcessKillingExecutor(max_workers=8) 88 | generator = executor.map( 89 | process_mesh, 90 | inputs, 91 | timeout=120, 92 | callback_timeout=process_mesh_timeout, 93 | ) 94 | 95 | with h5py.File(out_ds_file, "a") as f: 96 | if "meshes" not in f: 97 | f.create_group("meshes") 98 | 99 | categories = {} 100 | if "categories" not in f: 101 | f.create_group("categories") 102 | 103 | for mesh_info in tqdm( 104 | generator, total=len(inputs), desc="Processing Meshes" 105 | ): 106 | if mesh_info is not None: 107 | mk, minfo = mesh_info 108 | if mk not in f["meshes"]: 109 | f["meshes"].create_group(mk) 110 | for key in minfo: 111 | if key in f["meshes"][mk]: 112 | del f["meshes"][mk][key] 113 | f["meshes"][mk][key] = minfo[key] 114 | 115 | if minfo["category"] not in categories: 116 | categories[minfo["category"]] = [mk] 117 | elif mk not in categories[minfo["category"]]: 118 | categories[minfo["category"]].append(mk) 119 | 120 | for c in categories: 121 | if c in f["categories"]: 122 | del f["categories"][c] 123 | f["categories"][c] = categories[c] 124 | -------------------------------------------------------------------------------- /tools/train_robotcollisionnet.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import os.path as osp 4 | import sys 5 | 6 | import numpy as np 7 | import torch 8 | import torch.nn as nn 9 | from autolab_core import YamlConfig, keyboard_input 10 | from torch.cuda.amp import GradScaler 11 | from torch.utils.data import DataLoader 12 | 13 | from scenecollisionnet.collision_models import RobotCollisionNet 14 | from scenecollisionnet.datasets import IterableRobotCollisionDataset 15 | from scenecollisionnet.trainer import Trainer 16 | 17 | 18 | class RobotCollisionTrainer(Trainer): 19 | def to_binary(self, pred, true): 20 | return torch.clamp(pred, 0.0, 1.0), true > 1.0 21 | 22 | 23 | if __name__ == "__main__": 24 | 25 | parser = argparse.ArgumentParser(description="Train a SelfCollisionNet") 26 | parser.add_argument( 27 | "--cfg", 28 | type=str, 29 | default="cfg/train_robotcollisionnet.yaml", 30 | help="config file with training params", 31 | ) 32 | parser.add_argument( 33 | "--resume", action="store_true", help="resume training" 34 | ) 35 | parser.add_argument("-b", "--batch_size", type=int) 36 | parser.add_argument("-n", "--model_name", type=str) 37 | parser.add_argument( 38 | "-d", "--device", type=int, default=0, help="GPU index" 39 | ) 40 | args = parser.parse_args() 41 | 42 | # Replace config with args 43 | config = YamlConfig(args.cfg) 44 | if args.batch_size is not None: 45 | config["dataset"]["batch_size"] = args.batch_size 46 | if args.model_name is not None: 47 | config["model"]["name"] = args.model_name 48 | resume = args.resume 49 | 50 | # Create output directory for model 51 | out = osp.join(config["model"]["path"], config["model"]["name"]) 52 | if osp.exists(out) and not resume: 53 | response = keyboard_input( 54 | "A model exists at {}. Would you like to overwrite?".format(out), 55 | yesno=True, 56 | ) 57 | if response.lower() == "n": 58 | sys.exit(0) 59 | elif osp.exists(out) and resume: 60 | resume = resume and osp.exists(osp.join(out, "checkpoint.pth.tar")) 61 | else: 62 | resume = False 63 | os.makedirs(out) 64 | 65 | # Check whether GPU available 66 | cuda = torch.cuda.is_available() 67 | if not cuda: 68 | print("No CUDA available!") 69 | sys.exit(0) 70 | 71 | # 1. Dataset 72 | kwargs = { 73 | "num_workers": config["trainer"]["num_workers"] 74 | if "num_workers" in config["trainer"] 75 | else os.cpu_count(), 76 | "pin_memory": True, 77 | "worker_init_fn": lambda _: np.random.seed(), 78 | } 79 | train_set = IterableRobotCollisionDataset( 80 | **config["dataset"], 81 | ) 82 | train_loader = DataLoader(train_set, batch_size=None, **kwargs) 83 | 84 | # 2. Model 85 | start_epoch = 0 86 | start_iteration = 0 87 | config["model"]["num_joints"] = len(train_set.robot.actuated_joints) 88 | model = RobotCollisionNet(num_joints=config["model"]["num_joints"]) 89 | if resume: 90 | checkpoint = torch.load(osp.join(out, "checkpoint.pth.tar")) 91 | model.load_state_dict(checkpoint["model_state_dict"]) 92 | start_epoch = checkpoint["epoch"] 93 | start_iteration = checkpoint["iteration"] 94 | 95 | # 3. Optimizer and loss 96 | criterion = nn.MSELoss() 97 | optimizer = torch.optim.SGD( 98 | model.parameters(), 99 | lr=config["trainer"]["lr"], 100 | momentum=config["trainer"]["momentum"], 101 | ) 102 | scaler = GradScaler() 103 | if resume: 104 | optimizer.load_state_dict(checkpoint["optim_state_dict"]) 105 | scaler.load_state_dict(checkpoint["amp"]) 106 | 107 | # Save out training config 108 | config.save(osp.join(out, "train.yaml")) 109 | 110 | trainer = RobotCollisionTrainer( 111 | device=torch.device(config["trainer"]["device"]), 112 | model=model, 113 | optimizer=optimizer, 114 | scaler=scaler, 115 | criterion=criterion, 116 | train_loader=train_loader, 117 | train_iterations=config["trainer"]["train_iterations"], 118 | val_iterations=config["trainer"]["val_iterations"], 119 | epochs=config["trainer"]["epochs"], 120 | out=out, 121 | ) 122 | trainer.epoch = start_epoch 123 | trainer.iteration = start_iteration 124 | trainer.train() 125 | -------------------------------------------------------------------------------- /scenecollisionnet/collision_models/fcl.py: -------------------------------------------------------------------------------- 1 | import multiprocessing as mp 2 | import queue as Queue 3 | 4 | import numpy as np 5 | import trimesh 6 | 7 | 8 | class FCLProc(mp.Process): 9 | """ 10 | Used for finding collisions in parallel using FCL. 11 | """ 12 | 13 | def __init__(self, output_queue): 14 | """ 15 | Args: 16 | output_queue: mp.Queue, the queue that all the output data 17 | that is computed is added to. 18 | """ 19 | super().__init__() 20 | self.output_queue = output_queue 21 | self.input_queue = mp.Queue() 22 | 23 | def _collides(self, poses, inds): 24 | """computes collisions.""" 25 | coll = np.zeros(len(inds), dtype=np.bool) 26 | for k, pose in enumerate(poses[inds]): 27 | self.obj_manager.set_transform("obj", pose) 28 | coll[k] = self.scene_manager.in_collision_other(self.obj_manager) 29 | return coll 30 | 31 | def _set_scene(self, scene): 32 | self.scene_manager = trimesh.collision.CollisionManager() 33 | self.scene_manager.add_object("scene", scene) 34 | 35 | def _set_object(self, obj): 36 | self.obj_manager = trimesh.collision.CollisionManager() 37 | self.obj_manager.add_object("obj", obj) 38 | 39 | def run(self): 40 | """ 41 | the main function of each FCL process. 42 | """ 43 | while True: 44 | try: 45 | request = self.input_queue.get(timeout=1) 46 | except Queue.Empty: 47 | continue 48 | if request[0] == "set_scene": 49 | self._set_scene(request[1]), 50 | elif request[0] == "set_object": 51 | self._set_object(request[1]), 52 | elif request[0] == "collides": 53 | self.output_queue.put( 54 | ( 55 | request[3], 56 | self._collides(*request[1:3]), 57 | ) 58 | ) 59 | 60 | def set_scene(self, scene): 61 | self.input_queue.put(("set_scene", scene)) 62 | 63 | def set_object(self, obj): 64 | self.input_queue.put(("set_object", obj)) 65 | 66 | def collides(self, poses, inds, pind=None): 67 | self.input_queue.put(("collides", poses, inds, pind)) 68 | 69 | 70 | class FCLMultiSceneCollisionChecker: 71 | def __init__(self, n_proc=10): 72 | self._n_proc = n_proc 73 | 74 | self.output_queue = mp.Queue() 75 | self.coll_procs = [] 76 | for _ in range(self._n_proc): 77 | self.coll_procs.append( 78 | FCLProc( 79 | self.output_queue, 80 | ) 81 | ) 82 | self.coll_procs[-1].daemon = True 83 | self.coll_procs[-1].start() 84 | 85 | def set_scene(self, scene_pc): 86 | pc = trimesh.PointCloud(scene_pc) 87 | scene_mesh = trimesh.voxel.ops.points_to_marching_cubes( 88 | pc.vertices, pitch=0.01 89 | ) 90 | for proc in self.coll_procs: 91 | proc.set_scene(scene_mesh) 92 | 93 | def set_object(self, obj_pc): 94 | pc = trimesh.PointCloud(obj_pc) 95 | obj_mesh = trimesh.voxel.ops.points_to_marching_cubes( 96 | pc.vertices, pitch=0.01 97 | ) 98 | for proc in self.coll_procs: 99 | proc.set_object(obj_mesh) 100 | 101 | def __call__(self, scene, obj, trans, rots): 102 | self.set_scene(scene.squeeze().numpy()) 103 | self.set_object(obj.squeeze().numpy()) 104 | 105 | poses = np.repeat(np.eye(4)[None, ...], len(trans), axis=0) 106 | poses[:, :3, :2] = rots.numpy().reshape(-1, 3, 2) 107 | poses[:, :3, 2] = np.cross(rots[:, :3].numpy(), rots[:, 3:].numpy()) 108 | poses[:, :3, 3] = trans 109 | coll = np.zeros(len(trans), dtype=np.bool) 110 | for i in range(self._n_proc): 111 | self.coll_procs[i].collides( 112 | poses, 113 | np.arange( 114 | i * len(trans) // self._n_proc, 115 | (i + 1) * len(trans) // self._n_proc, 116 | ), 117 | pind=i, 118 | ) 119 | 120 | # collect computed collisions 121 | for _ in range(self._n_proc): 122 | i, proc_coll = self.output_queue.get(True) 123 | coll[ 124 | i 125 | * len(trans) 126 | // self._n_proc : (i + 1) 127 | * len(trans) 128 | // self._n_proc 129 | ] = proc_coll 130 | 131 | return coll 132 | -------------------------------------------------------------------------------- /tools/benchmark_robotcollisionnet.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import os.path as osp 4 | import sys 5 | from timeit import default_timer as timer 6 | 7 | import matplotlib as mpl 8 | import numpy as np 9 | import torch 10 | import torch.nn as nn 11 | import tqdm 12 | from torch.autograd import Variable 13 | from torch.utils.data import DataLoader 14 | 15 | mpl.use("agg") 16 | import matplotlib.pyplot as plt 17 | import seaborn as sns 18 | from autolab_core import BinaryClassificationResult, YamlConfig, keyboard_input 19 | 20 | from scenecollisionnet.collision_models import RobotCollisionNet 21 | from scenecollisionnet.datasets import IterableRobotCollisionDataset 22 | 23 | sns.set() 24 | 25 | 26 | def benchmark(device, model, criterion, test_loader, iterations, out_path): 27 | model.eval() 28 | time_start = timer() 29 | 30 | test_loss = 0 31 | preds, trues = np.array([]), np.array([]) 32 | data_loader = iter(test_loader) 33 | passes = [] 34 | for _ in tqdm.trange( 35 | iterations, 36 | desc="Benchmarking", 37 | ): 38 | data = next(data_loader) 39 | data = [Variable(d.float().to(device)) for d in data] 40 | centers, colls = data 41 | fw_pass_start = timer() 42 | with torch.no_grad(): 43 | out = model(centers).squeeze(dim=-1) 44 | passes.append(timer() - fw_pass_start) 45 | 46 | loss = criterion(out, colls.float()) 47 | pred_batch = out.cpu().numpy().flatten() 48 | true_batch = colls.cpu().numpy().flatten() 49 | preds = np.append(preds, pred_batch) 50 | trues = np.append(trues, true_batch) 51 | test_loss += loss.item() 52 | 53 | test_loss /= iterations 54 | 55 | # sort by probs 56 | accs = [] 57 | f1s = [] 58 | tprs = [] 59 | taus = np.linspace(0, 10, 51) 60 | for t in taus: 61 | b = BinaryClassificationResult(preds >= t, trues >= t) 62 | accs.append(b.accuracy) 63 | f1s.append(b.f1_score) 64 | tprs.append(b.tpr) 65 | 66 | np.savez_compressed( 67 | osp.join(out_path, "accuracy_curve.npz"), 68 | taus=taus, 69 | accs=accs, 70 | f1s=f1s, 71 | tprs=tprs, 72 | ) 73 | 74 | ax = sns.lineplot(x=taus, y=accs, label="Accuracy") 75 | ax = sns.lineplot(x=taus, y=f1s, label="F1") 76 | ax = sns.lineplot(x=taus, y=tprs, label="TPR") 77 | ax.legend() 78 | plt.xlabel("Threshold") 79 | plt.ylabel("Score") 80 | plt.ylim([0.49, 1.01]) 81 | plt.xlim([taus.min() - 0.01, taus.max() + 0.01]) 82 | plt.savefig(osp.join(out_path, "accuracy_curve.png")) 83 | 84 | bcr = BinaryClassificationResult( 85 | preds >= taus.mean(), trues >= taus.mean() 86 | ) 87 | with open(osp.join(out_path, "results.txt"), "w") as f: 88 | elapsed_time = timer() - time_start 89 | log = ( 90 | ["Images: {:d}".format(iterations)] 91 | + ["Queries: {:d}".format(iterations * len(colls.flatten()))] 92 | + ["Loss: {:.5f}".format(test_loss)] 93 | + ["Accuracy: {:.3f}".format(bcr.accuracy)] 94 | + ["F1 Score: {:.3f}".format(bcr.f1_score)] 95 | + ["TPR: {:.3f}".format(bcr.tpr)] 96 | + ["AP Score: {:.3f}".format(bcr.ap_score)] 97 | + [ 98 | "FW Pass Time: {:.4f} +- {:.4f} s".format( 99 | np.mean(passes), np.std(passes) 100 | ) 101 | ] 102 | + ["Time: {:.2f} s".format(elapsed_time)] 103 | ) 104 | log = map(str, log) 105 | f.write("\n".join(log)) 106 | 107 | 108 | if __name__ == "__main__": 109 | 110 | parser = argparse.ArgumentParser( 111 | description="Benchmark a SelfCollisionNet" 112 | ) 113 | parser.add_argument( 114 | "--cfg", 115 | type=str, 116 | default="cfg/benchmark_robotcollisionnet.yaml", 117 | help="config file with benchmarking params", 118 | ) 119 | args = parser.parse_args() 120 | 121 | config = YamlConfig(args.cfg) 122 | 123 | # Create output directory for model 124 | out = osp.join(config["output"], config["model"]["name"]) 125 | if osp.exists(out): 126 | response = keyboard_input( 127 | "A benchmark folder exists at {}. Overwrite?".format(out), 128 | yesno=True, 129 | ) 130 | if response.lower() == "n": 131 | os.exit() 132 | else: 133 | os.makedirs(out) 134 | config.save(osp.join(out, "benchmark.yaml")) 135 | 136 | # Check whether GPU available 137 | if not torch.cuda.is_available(): 138 | print("No CUDA available!") 139 | sys.exit(0) 140 | device = torch.device(config["device"]) 141 | 142 | # 1. Model cfg 143 | model_path = osp.join(config["model"]["path"], config["model"]["name"]) 144 | train_cfg = YamlConfig(osp.join(model_path, "train.yaml")) 145 | 146 | # 2. Dataset 147 | kwargs = { 148 | "num_workers": config["num_workers"] 149 | if "num_workers" in config 150 | else os.cpu_count(), 151 | "pin_memory": True, 152 | "worker_init_fn": lambda _: np.random.seed(), 153 | } 154 | train_cfg["dataset"]["batch_size"] = config["dataset"]["batch_size"] 155 | test_set = IterableRobotCollisionDataset( 156 | **train_cfg["dataset"], 157 | ) 158 | test_loader = DataLoader(test_set, batch_size=None, **kwargs) 159 | 160 | # 3. Model 161 | model = RobotCollisionNet(num_joints=len(test_set.robot.actuated_joints)) 162 | checkpoint = torch.load(osp.join(model_path, "model.pth.tar")) 163 | model.load_state_dict(checkpoint["model_state_dict"], strict=False) 164 | model = model.to(device) 165 | 166 | benchmark( 167 | device=device, 168 | model=model, 169 | criterion=nn.MSELoss(), 170 | test_loader=test_loader, 171 | iterations=config["iterations"], 172 | out_path=out, 173 | ) 174 | -------------------------------------------------------------------------------- /tools/train_scenecollisionnet.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import os.path as osp 4 | import sys 5 | 6 | import numpy as np 7 | import torch 8 | import torch.nn as nn 9 | from autolab_core import YamlConfig, keyboard_input 10 | from torch.cuda.amp import GradScaler 11 | from torch.utils.data import DataLoader 12 | 13 | from scenecollisionnet.collision_models import PointNetGrid, SceneCollisionNet 14 | from scenecollisionnet.datasets import IterableSceneCollisionDataset 15 | from scenecollisionnet.trainer import Trainer 16 | 17 | 18 | class SceneCollisionNetTrainer(Trainer): 19 | def to_binary(self, pred, true): 20 | return torch.sigmoid(pred), true 21 | 22 | 23 | if __name__ == "__main__": 24 | 25 | parser = argparse.ArgumentParser(description="Train a SceneCollisionNet") 26 | parser.add_argument( 27 | "--cfg", 28 | type=str, 29 | default="cfg/train_scenecollisionnet.yaml", 30 | help="config file with training params", 31 | ) 32 | parser.add_argument( 33 | "--resume", action="store_true", help="resume training" 34 | ) 35 | parser.add_argument("-b", "--batch_size", type=int) 36 | parser.add_argument("-q", "--query_size", type=int) 37 | parser.add_argument("-o", "--obj_points", type=int) 38 | parser.add_argument("-s", "--scene_points", type=int) 39 | parser.add_argument("-t", "--trajectories", type=int) 40 | parser.add_argument("-r", "--rotations", type=int) 41 | parser.add_argument("-n", "--model_name", type=str) 42 | parser.add_argument( 43 | "-l", "--loss_pct", type=float, help="top k hard negative loss" 44 | ) 45 | parser.add_argument("--dataset_path", type=str) 46 | parser.add_argument("--device", type=int, default=0, help="GPU index") 47 | args = parser.parse_args() 48 | 49 | # Replace config with args 50 | config = YamlConfig(args.cfg) 51 | if args.batch_size is not None: 52 | config["dataset"]["batch_size"] = args.batch_size 53 | if args.query_size is not None: 54 | config["dataset"]["query_size"] = args.query_size 55 | if args.obj_points is not None: 56 | config["dataset"]["n_obj_points"] = args.obj_points 57 | if args.scene_points is not None: 58 | config["dataset"]["n_scene_points"] = args.scene_points 59 | if args.trajectories is not None: 60 | config["dataset"]["trajectories"] = args.trajectories 61 | if args.rotations is not None: 62 | config["dataset"]["rotations"] = args.rotations 63 | if args.model_name is not None: 64 | config["model"]["name"] = args.model_name 65 | if args.dataset_path is not None: 66 | config["dataset"]["meshes"] = args.dataset_path 67 | if args.device is not None: 68 | config["trainer"]["device"] = args.device 69 | if args.loss_pct is not None: 70 | config["trainer"]["loss_pct"] = args.loss_pct 71 | resume = args.resume 72 | 73 | # Create output directory for model 74 | out = osp.join(config["model"]["path"], config["model"]["name"]) 75 | if osp.exists(out) and not resume: 76 | response = keyboard_input( 77 | "A model exists at {}. Would you like to overwrite?".format(out), 78 | yesno=True, 79 | ) 80 | if response.lower() == "n": 81 | sys.exit(0) 82 | elif osp.exists(out) and resume: 83 | resume = resume and osp.exists(osp.join(out, "checkpoint.pth.tar")) 84 | else: 85 | resume = False 86 | os.makedirs(out) 87 | config.save(osp.join(out, "train.yaml")) 88 | 89 | # Check whether GPU available 90 | cuda = torch.cuda.is_available() 91 | if not cuda: 92 | print("No CUDA available!") 93 | sys.exit(0) 94 | 95 | # 1. Model 96 | start_epoch = 0 97 | start_iteration = 0 98 | model_type = ( 99 | SceneCollisionNet 100 | if config["model"]["type"] == "SceneCollisionNet" 101 | else PointNetGrid 102 | ) 103 | model = model_type( 104 | bounds=config["model"]["bounds"], 105 | vox_size=config["model"]["vox_size"], 106 | ) 107 | if resume: 108 | checkpoint = torch.load(osp.join(out, "checkpoint.pth.tar")) 109 | model.load_state_dict(checkpoint["model_state_dict"]) 110 | start_epoch = checkpoint["epoch"] 111 | start_iteration = checkpoint["iteration"] 112 | 113 | # 2. Dataset 114 | kwargs = { 115 | "num_workers": config["trainer"]["num_workers"] 116 | if "num_workers" in config["trainer"] 117 | else os.cpu_count(), 118 | "pin_memory": True, 119 | "worker_init_fn": lambda _: np.random.seed(), 120 | } 121 | train_set = IterableSceneCollisionDataset( 122 | **config["dataset"], 123 | **config["camera"], 124 | bounds=config["model"]["bounds"] 125 | ) 126 | train_loader = DataLoader(train_set, batch_size=None, **kwargs) 127 | 128 | # 3. Optimizer and loss 129 | criterion = nn.BCEWithLogitsLoss(reduction="none") 130 | optimizer = torch.optim.SGD( 131 | model.parameters(), 132 | lr=config["trainer"]["lr"], 133 | momentum=config["trainer"]["momentum"], 134 | ) 135 | scaler = GradScaler() 136 | if resume: 137 | optimizer.load_state_dict(checkpoint["optim_state_dict"]) 138 | scaler.load_state_dict(checkpoint["amp"]) 139 | 140 | trainer = SceneCollisionNetTrainer( 141 | device=torch.device(config["trainer"]["device"]), 142 | model=model, 143 | optimizer=optimizer, 144 | scaler=scaler, 145 | criterion=criterion, 146 | train_loader=train_loader, 147 | train_iterations=config["trainer"]["train_iterations"], 148 | val_iterations=config["trainer"]["val_iterations"], 149 | epochs=config["trainer"]["epochs"], 150 | out=out, 151 | loss_pct=config["trainer"]["loss_pct"], 152 | ) 153 | trainer.epoch = start_epoch 154 | trainer.iteration = start_iteration 155 | trainer.train() 156 | -------------------------------------------------------------------------------- /scenecollisionnet/collision_models/sdf.py: -------------------------------------------------------------------------------- 1 | import multiprocessing as mp 2 | import queue as Queue 3 | 4 | import numpy as np 5 | import torch 6 | import trimesh 7 | 8 | try: 9 | import kaolin as kal 10 | except ModuleNotFoundError: 11 | print("Kaolin not installed, SDF baselines not available!") 12 | 13 | 14 | class SDFProc(mp.Process): 15 | """ 16 | Used for finding collisions in parallel using SDFs. 17 | """ 18 | 19 | def __init__(self, output_queue, device): 20 | """ 21 | Args: 22 | output_queue: mp.Queue, the queue that all the output data 23 | that is computed is added to. 24 | """ 25 | super().__init__() 26 | self.output_queue = output_queue 27 | self.device = device 28 | self.input_queue = mp.Queue() 29 | 30 | def _set_sdf(self, verts, faces): 31 | self.v = torch.from_numpy(verts).float().to(self.device) 32 | self.f = torch.from_numpy(faces).long().to(self.device) 33 | 34 | def _collides(self, poses, inds): 35 | """computes collisions.""" 36 | ind_poses = torch.from_numpy(poses[inds]).float().to(self.device) 37 | tf_pts = torch.cat( 38 | ( 39 | self.points, 40 | torch.ones((len(self.points), 1), device=self.device), 41 | ), 42 | dim=1, 43 | ) 44 | tf_pts = torch.matmul(ind_poses, tf_pts.T).transpose(2, 1)[..., :3] 45 | sdf_vals = kal.ops.mesh.check_sign( 46 | self.v.expand(len(inds), -1, -1), self.f, tf_pts 47 | ) 48 | return sdf_vals.any(dim=-1).cpu().numpy() 49 | 50 | def run(self): 51 | """ 52 | the main function of each process. 53 | """ 54 | while True: 55 | try: 56 | request = self.input_queue.get(timeout=1) 57 | except Queue.Empty: 58 | continue 59 | if request[0] == "set_sdf": 60 | self._set_sdf(*request[1:]) 61 | elif request[0] == "set_points": 62 | self.points = request[1].float().to(self.device) 63 | elif request[0] == "collides": 64 | self.output_queue.put( 65 | ( 66 | request[3], 67 | self._collides(*request[1:3]), 68 | ) 69 | ) 70 | 71 | def set_sdf(self, verts, faces): 72 | self.input_queue.put(("set_sdf", verts, faces)) 73 | 74 | def set_points(self, pts): 75 | self.input_queue.put(("set_points", pts)) 76 | 77 | def collides(self, poses, inds, pind=None): 78 | self.input_queue.put(("collides", poses, inds, pind)) 79 | 80 | 81 | class SDFMultiSceneCollisionChecker: 82 | def __init__(self, device="cuda:0", n_proc=2): 83 | self.device = torch.device(device) 84 | self._n_proc = n_proc 85 | 86 | if self._n_proc > 0: 87 | self.output_queue = mp.Queue() 88 | self.coll_procs = [] 89 | for _ in range(self._n_proc): 90 | self.coll_procs.append(SDFProc(self.output_queue, self.device)) 91 | self.coll_procs[-1].daemon = True 92 | self.coll_procs[-1].start() 93 | 94 | def set_points(self, pts): 95 | if self._n_proc == 0: 96 | self.points = pts.float().to(self.device) 97 | else: 98 | for proc in self.coll_procs: 99 | proc.set_points(pts) 100 | 101 | def set_sdf(self, verts, faces): 102 | if self._n_proc == 0: 103 | self.v = torch.from_numpy(verts).float().to(self.device) 104 | self.f = torch.from_numpy(faces).long().to(self.device) 105 | else: 106 | for proc in self.coll_procs: 107 | proc.set_sdf(verts, faces) 108 | 109 | def _points_to_mesh(self, pts): 110 | pc = trimesh.PointCloud(pts) 111 | mesh = trimesh.voxel.ops.points_to_marching_cubes( 112 | pc.vertices, pitch=0.01 113 | ) 114 | return mesh.vertices, mesh.faces 115 | 116 | def _to_poses(self, trans, rots): 117 | poses = np.repeat(np.eye(4)[None, ...], len(trans), axis=0) 118 | poses[:, :3, :2] = rots.numpy().reshape(-1, 3, 2) 119 | poses[:, :3, 2] = np.cross(rots[:, :3].numpy(), rots[:, 3:].numpy()) 120 | poses[:, :3, 3] = trans 121 | return poses 122 | 123 | def __call__(self, scene, obj, trans, rots): 124 | self.set_scene(scene.squeeze()) 125 | self.set_object(obj.squeeze()) 126 | poses = self._to_poses(trans, rots) 127 | 128 | if self._n_proc > 0: 129 | coll = np.zeros(len(trans), dtype=np.bool) 130 | for i in range(self._n_proc): 131 | self.coll_procs[i].collides( 132 | poses, 133 | np.arange( 134 | i * len(trans) // self._n_proc, 135 | (i + 1) * len(trans) // self._n_proc, 136 | ), 137 | pind=i, 138 | ) 139 | 140 | # collect computed collisions 141 | for _ in range(self._n_proc): 142 | i, proc_coll = self.output_queue.get(True) 143 | coll[ 144 | i 145 | * len(trans) 146 | // self._n_proc : (i + 1) 147 | * len(trans) 148 | // self._n_proc 149 | ] = proc_coll 150 | 151 | else: 152 | poses = torch.from_numpy(poses).float().to(self.device) 153 | tf_pts = torch.cat( 154 | ( 155 | self.points, 156 | torch.ones((len(self.points), 1), device=self.device), 157 | ), 158 | dim=1, 159 | ) 160 | tf_pts = torch.matmul(poses, tf_pts.T).transpose(2, 1)[..., :3] 161 | sdf_vals = kal.ops.mesh.check_sign( 162 | self.v.expand(len(poses), -1, -1), self.f, tf_pts 163 | ) 164 | coll = sdf_vals.any(dim=-1).cpu().numpy() 165 | 166 | return coll 167 | 168 | 169 | class SDFSceneCollisionChecker(SDFMultiSceneCollisionChecker): 170 | def set_scene(self, scene_pts): 171 | self.set_sdf(*self._points_to_mesh(scene_pts)) 172 | 173 | def set_object(self, obj_pts): 174 | self.set_points(obj_pts) 175 | 176 | 177 | class SDFObjCollisionChecker(SDFMultiSceneCollisionChecker): 178 | def set_scene(self, scene_pts): 179 | self.set_points(scene_pts) 180 | 181 | def set_object(self, obj_pts): 182 | self.set_sdf(*self._points_to_mesh(obj_pts)) 183 | 184 | def _to_poses(self, trans, rots): 185 | poses = super()._to_poses(trans, rots) 186 | return np.linalg.inv(poses) 187 | -------------------------------------------------------------------------------- /scenecollisionnet/collision_models/pointnet_grid.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from pointnet2.pointnet2_modules import PointnetSAModuleMSG 4 | from pointnet2.pytorch_utils import Conv1d 5 | 6 | OBJ_NPOINTS = [512, 256, 128, 1] 7 | OBJ_MLPS = [ 8 | [[16, 16, 32], [32, 32, 64]], 9 | [[64, 64, 128], [64, 96, 128]], 10 | [[128, 196, 256], [128, 196, 256]], 11 | [[256, 256, 512], [256, 384, 512]], 12 | ] 13 | OBJ_RADIUS = [[0.005, 0.01], [0.01, 0.02], [0.02, 0.04], [0.04, 0.08]] 14 | OBJ_NSAMPLES = [[32, 32], [32, 32], [32, 32], [32, 32]] 15 | 16 | SCENE_NPOINTS = [8192, 2048, 512, 128, 32, 1] 17 | SCENE_MLPS = [ 18 | [[16, 16, 32], [32, 32, 64]], 19 | [[64, 64, 128], [64, 96, 128]], 20 | [[128, 196, 256], [128, 196, 256]], 21 | [[256, 256, 512], [256, 384, 512]], 22 | [[512, 768, 1024], [512, 768, 1024]], 23 | [[1024, 1024, 2048], [1024, 1536, 2048]], 24 | ] 25 | SCENE_RADIUS = [ 26 | [0.005, 0.01], 27 | [0.01, 0.05], 28 | [0.1, 0.2], 29 | [0.2, 0.4], 30 | [0.4, 0.8], 31 | [0.8, 1.6], 32 | ] 33 | SCENE_NSAMPLE = [[32, 32], [32, 32], [32, 32], [32, 32], [32, 32], [32, 32]] 34 | CLS_FC = [2048, 1024, 512, 512] 35 | 36 | 37 | class PointNetGrid(nn.Module): 38 | def __init__(self, bounds, vox_size): 39 | super().__init__() 40 | 41 | self.bounds = nn.Parameter( 42 | torch.FloatTensor(bounds), requires_grad=False 43 | ) 44 | self.vox_size = nn.Parameter( 45 | torch.FloatTensor(vox_size), requires_grad=False 46 | ) 47 | if torch.allclose(self.vox_size, self.bounds[1] - self.bounds[0]): 48 | self.num_voxels = nn.Parameter( 49 | torch.ones(3, dtype=torch.long), requires_grad=False 50 | ) 51 | else: 52 | self.num_voxels = nn.Parameter( 53 | (self.bounds[1] - self.bounds[0]) / (self.vox_size / 2) + 1, 54 | requires_grad=False, 55 | ).long() 56 | self.vox_centers = torch.meshgrid( 57 | [ 58 | torch.linspace(self.bounds[0, i], self.bounds[1, i], n) 59 | for i, n in enumerate(self.num_voxels) 60 | ] 61 | ) 62 | self.vox_centers = nn.Parameter( 63 | torch.stack([vc.flatten() for vc in self.vox_centers], dim=1), 64 | requires_grad=False, 65 | ) 66 | 67 | # Scene pointnet layers 68 | channel_in = 0 69 | self.scene_SA_modules = nn.ModuleList() 70 | for k in range(SCENE_NPOINTS.__len__()): 71 | mlps = SCENE_MLPS[k].copy() 72 | channel_out = 0 73 | for idx in range(mlps.__len__()): 74 | mlps[idx] = [channel_in] + mlps[idx] 75 | channel_out += mlps[idx][-1] 76 | 77 | self.scene_SA_modules.append( 78 | PointnetSAModuleMSG( 79 | npoint=SCENE_NPOINTS[k], 80 | radii=SCENE_RADIUS[k], 81 | nsamples=SCENE_NSAMPLE[k], 82 | mlps=mlps, 83 | use_xyz=True, 84 | bn=False, 85 | ) 86 | ) 87 | channel_in = channel_out 88 | 89 | # Obj pointnet layers 90 | channel_in = 0 91 | self.obj_SA_modules = nn.ModuleList() 92 | for k in range(OBJ_NPOINTS.__len__()): 93 | mlps = OBJ_MLPS[k].copy() 94 | channel_out = 0 95 | for idx in range(mlps.__len__()): 96 | mlps[idx] = [channel_in] + mlps[idx] 97 | channel_out += mlps[idx][-1] 98 | 99 | self.obj_SA_modules.append( 100 | PointnetSAModuleMSG( 101 | npoint=OBJ_NPOINTS[k], 102 | radii=OBJ_RADIUS[k], 103 | nsamples=OBJ_NSAMPLES[k], 104 | mlps=mlps, 105 | use_xyz=True, 106 | bn=False, 107 | ) 108 | ) 109 | channel_in = channel_out 110 | 111 | cls_layers = [] 112 | pre_channel = 4096 + 1024 + 3 + 6 113 | for k in range(0, CLS_FC.__len__()): 114 | cls_layers.append(Conv1d(pre_channel, CLS_FC[k])) 115 | pre_channel = CLS_FC[k] 116 | cls_layers.append(Conv1d(pre_channel, 1, activation=None)) 117 | cls_layers.insert(1, nn.Dropout(0.5)) 118 | self.cls_layer = nn.Sequential(*cls_layers) 119 | 120 | def _break_up_pc(self, pc): 121 | xyz = pc[..., 0:3].contiguous() 122 | features = ( 123 | pc[..., 3:].transpose(1, 2).contiguous() 124 | if pc.size(-1) > 3 125 | else None 126 | ) 127 | 128 | return xyz, features 129 | 130 | def get_obj_features(self, obj): 131 | xyz, features = self._break_up_pc(obj) 132 | for i in range(len(self.obj_SA_modules)): 133 | xyz, features = self.obj_SA_modules[i](xyz, features) 134 | return features.squeeze(dim=1) 135 | 136 | def get_scene_features(self, scene): 137 | xyz, features = self._break_up_pc(scene) 138 | for i in range(len(self.scene_SA_modules)): 139 | xyz, features = self.scene_SA_modules[i](xyz, features) 140 | return features.squeeze(dim=1) 141 | 142 | def forward(self, scene, obj, tr, rot): 143 | 144 | # Voxelize scene points 145 | scene_vox_mask = ( 146 | torch.abs(scene.transpose(1, 0) - self.vox_centers) 147 | < self.vox_size / 2 148 | ).all(dim=-1) 149 | scene_vox_counts = scene_vox_mask.sum(dim=0) 150 | nonzero_inds = torch.where(scene_vox_counts > 0)[0] 151 | val_vox_counts = scene_vox_counts[nonzero_inds] 152 | val_vox_centers = self.vox_centers[nonzero_inds] 153 | 154 | # Sample points and normalize per voxel - super janky but works 155 | scene_smpl_inds = ( 156 | torch.randint( 157 | scene.size(1), 158 | size=(1024, len(nonzero_inds)), 159 | device=scene.device, 160 | ) 161 | % val_vox_counts 162 | ) 163 | offsets = torch.roll(val_vox_counts.cumsum(dim=0), 1, 0) 164 | offsets[0] = 0 165 | scene_smpl_inds += offsets 166 | scene_pts = scene[ 167 | :, torch.where(scene_vox_mask.T)[1][scene_smpl_inds.flatten()] 168 | ].reshape(*scene_smpl_inds.shape, -1) 169 | norm_scene_pts = (scene_pts - val_vox_centers).transpose(1, 0) 170 | 171 | # Normalize translations 172 | tr_vox_mask = ( 173 | torch.abs(tr.unsqueeze(1) - val_vox_centers) < self.vox_size / 2 174 | ).all(dim=-1) 175 | tr_val_inds, tr_vox_inds = torch.where(tr_vox_mask) 176 | norm_trans = tr[tr_val_inds] - val_vox_centers[tr_vox_inds] 177 | 178 | q = len(norm_trans) 179 | obj_feats = self.get_obj_features(obj) 180 | scene_feats = self.get_scene_features(norm_scene_pts) 181 | cls_in = torch.cat( 182 | ( 183 | obj_feats.expand(q, -1, -1), 184 | scene_feats[tr_vox_inds], 185 | norm_trans.unsqueeze(2), 186 | rot[tr_val_inds].unsqueeze(2), 187 | ), 188 | dim=1, 189 | ) 190 | pred_cls = self.cls_layer(cls_in).squeeze(dim=-1).squeeze(dim=-1) 191 | return ( 192 | torch.zeros(q, device=tr.device).scatter_add( 193 | 0, tr_val_inds, pred_cls 194 | ) 195 | / 8.0 196 | ) 197 | -------------------------------------------------------------------------------- /scenecollisionnet/trainer.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | import os 3 | import os.path as osp 4 | import shutil 5 | import sys 6 | from abc import abstractmethod 7 | 8 | import numpy as np 9 | import torch 10 | import tqdm 11 | from autolab_core import BinaryClassificationResult 12 | from torch.autograd import Variable 13 | from torch.cuda.amp import autocast 14 | 15 | 16 | class Trainer(object): 17 | def __init__( 18 | self, 19 | device, 20 | model, 21 | optimizer, 22 | scaler, 23 | criterion, 24 | train_loader, 25 | train_iterations, 26 | val_iterations, 27 | epochs, 28 | out, 29 | loss_pct=1.0, 30 | ): 31 | self.device = device 32 | self.model = model 33 | self.optim = optimizer 34 | self.scaler = scaler 35 | self.criterion = criterion 36 | 37 | self.data_loader = iter(train_loader) 38 | self.train_iterations = train_iterations 39 | self.val_iterations = val_iterations 40 | self.loss_pct = loss_pct 41 | 42 | self.timestamp_start = datetime.datetime.now() 43 | 44 | self.out = out 45 | if not osp.exists(self.out): 46 | os.makedirs(self.out) 47 | 48 | self.log_headers = [ 49 | "epoch", 50 | "iteration", 51 | "train/loss", 52 | "train/acc", 53 | "train/f1", 54 | "train/tpr", 55 | "train/ap", 56 | "valid/loss", 57 | "valid/acc", 58 | "valid/f1", 59 | "valid/tpr", 60 | "valid/ap", 61 | "elapsed_time", 62 | ] 63 | if not osp.exists(osp.join(self.out, "log.csv")): 64 | with open(osp.join(self.out, "log.csv"), "w") as f: 65 | f.write(",".join(self.log_headers) + "\n") 66 | 67 | self.epoch = 0 68 | self.iteration = 0 69 | self.max_epoch = epochs 70 | self.best_ap = 0.0 71 | 72 | def validate(self): 73 | self.model.eval() 74 | 75 | val_loss = 0 76 | val_preds, val_trues = np.array([]), np.array([]) 77 | for _ in tqdm.trange( 78 | self.val_iterations, 79 | desc="Valid iteration={:d}".format(self.iteration), 80 | leave=False, 81 | ): 82 | data = next(self.data_loader) 83 | data = [Variable(d.float().to(self.device)) for d in data] 84 | with torch.no_grad(): 85 | out = self.model(*data[:-1]).squeeze(dim=-1) 86 | 87 | coll = data[-1].type(out.type()) 88 | loss = self.criterion(out, coll).mean() 89 | 90 | out, coll = self.to_binary(out, coll) 91 | val_preds = np.append(val_preds, out.cpu().numpy().flatten()) 92 | val_trues = np.append( 93 | val_trues, coll.float().cpu().numpy().flatten() 94 | ) 95 | val_loss += loss.item() 96 | 97 | val_loss /= self.val_iterations 98 | bcr = BinaryClassificationResult(val_preds, val_trues) 99 | 100 | with open(osp.join(self.out, "log.csv"), "a") as f: 101 | elapsed_time = ( 102 | datetime.datetime.now() - self.timestamp_start 103 | ).total_seconds() 104 | log = ( 105 | [self.epoch, self.iteration] 106 | + [""] * 5 107 | + [val_loss] 108 | + [bcr.accuracy, bcr.f1_score, bcr.tpr, bcr.ap_score] 109 | + [elapsed_time] 110 | ) 111 | log = map(str, log) 112 | f.write(",".join(log) + "\n") 113 | 114 | is_best = bcr.ap_score > self.best_ap 115 | save_dict = { 116 | "epoch": self.epoch + 1, 117 | "iteration": self.iteration, 118 | "arch": self.model.__class__.__name__, 119 | "optim_state_dict": self.optim.state_dict(), 120 | "model_state_dict": self.model.state_dict(), 121 | "best_ap": self.best_ap, 122 | } 123 | save_dict.update({"amp": self.scaler.state_dict()}) 124 | torch.save(save_dict, osp.join(self.out, "checkpoint.pth.tar")) 125 | 126 | if is_best: 127 | shutil.copy( 128 | osp.join(self.out, "checkpoint.pth.tar"), 129 | osp.join(self.out, "model.pth.tar"), 130 | ) 131 | 132 | def train_epoch(self): 133 | self.model.train() 134 | 135 | train_bar = tqdm.trange( 136 | self.train_iterations, 137 | desc="Train epoch={:d}".format(self.epoch), 138 | leave=False, 139 | ) 140 | 141 | for batch_idx in train_bar: 142 | 143 | iteration = batch_idx + self.epoch * self.train_iterations 144 | if self.iteration != 0 and (iteration - 1) != self.iteration: 145 | continue # for resuming 146 | self.iteration = iteration 147 | 148 | assert self.model.training 149 | 150 | try: 151 | data = next(self.data_loader) 152 | except StopIteration: 153 | print("BAD NEWS BEARS") 154 | sys.exit(1) 155 | data = [Variable(d.float().to(self.device)) for d in data] 156 | self.optim.zero_grad() 157 | 158 | with autocast(): 159 | out = self.model(*data[:-1]).squeeze(dim=-1) 160 | 161 | coll = data[-1].type(out.type()) 162 | losses = self.criterion(out, coll) 163 | if self.loss_pct < 1.0: 164 | top_losses, _ = torch.topk( 165 | losses, int(losses.size(1) * self.loss_pct), sorted=False 166 | ) 167 | rand_losses = losses[ 168 | :, 169 | torch.randint( 170 | losses.size(1), (int(losses.size(1) * self.loss_pct),) 171 | ), 172 | ] 173 | loss = 0.5 * (top_losses.mean() + rand_losses.mean()) 174 | else: 175 | loss = losses.mean() 176 | loss_data = loss.item() 177 | if torch.isnan(loss.data): 178 | raise ValueError("loss is nan while training") 179 | self.scaler.scale(loss).backward() 180 | self.scaler.step(self.optim) 181 | self.scaler.update() 182 | 183 | with torch.no_grad(): 184 | out, coll = self.to_binary(out, coll) 185 | bcr = BinaryClassificationResult( 186 | out.data.cpu().numpy().flatten(), 187 | coll.data.cpu().numpy().flatten(), 188 | ) 189 | train_bar.set_postfix_str("Loss: {:.5f}".format(loss_data)) 190 | 191 | with open(osp.join(self.out, "log.csv"), "a") as f: 192 | elapsed_time = ( 193 | datetime.datetime.now() - self.timestamp_start 194 | ).total_seconds() 195 | log = ( 196 | [self.epoch, self.iteration] 197 | + [loss_data] 198 | + [bcr.accuracy, bcr.f1_score, bcr.tpr, bcr.ap_score] 199 | + [""] * 5 200 | + [elapsed_time] 201 | ) 202 | log = map(str, log) 203 | f.write(",".join(log) + "\n") 204 | 205 | @abstractmethod 206 | def to_binary(self, pred, true): 207 | pass 208 | 209 | def train(self): 210 | self.model.to(self.device) 211 | for epoch in tqdm.trange(self.epoch, self.max_epoch, desc="Train"): 212 | self.epoch = epoch 213 | self.train_epoch() 214 | self.validate() 215 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SceneCollisionNet 2 | This repo contains the code for "[Object Rearrangement Using Learned Implicit Collision Functions](https://arxiv.org/abs/2011.10726)", an ICRA 2021 paper. For more information, please visit the [project website](https://research.nvidia.com/publication/2021-03_Object-Rearrangement-Using). 3 | 4 | ## License 5 | This repo is released under [NVIDIA source code license](LICENSE.pdf). For business inquiries, please contact researchinquiries@nvidia.com. For press and other inquiries, please contact Hector Marinez at hmarinez@nvidia.com 6 | 7 | ## Install and Setup 8 | Clone and install the repo (we recommend a virtual environment, especially if training or benchmarking, to avoid dependency conflicts): 9 | ```shell 10 | git clone --recursive https://github.com/mjd3/SceneCollisionNet.git 11 | cd SceneCollisionNet 12 | pip install -e . 13 | ``` 14 | These commands install the minimum dependencies needed for generating a mesh dataset and then training/benchmarking using Docker. If you instead wish to train or benchmark without using Docker, please first install an appropriate version of [PyTorch](https://pytorch.org/get-started/locally/) and corresponding version of [PyTorch Scatter](https://github.com/rusty1s/pytorch_scatter) for your system. Then, execute these commands: 15 | ```shell 16 | git clone --recursive https://github.com/mjd3/SceneCollisionNet.git 17 | cd SceneCollisionNet 18 | pip install -e .[train] 19 | ``` 20 | If benchmarking, replace `train` in the last command with `bench`. 21 | 22 | To rollout the object rearrangement MPPI policy in a simulated tabletop environment, first download [Isaac Gym](https://developer.nvidia.com/isaac-gym) and place it in the `extern` folder within this repo. Next, follow the previous installation instructions for training, but replace the `train` option with `policy`. 23 | 24 | To download the pretrained weights for benchmarking or policy rollout, run `bash scripts/download_weights.sh`. 25 | 26 | ## Generating a Mesh Dataset 27 | To save time during training/benchmarking, meshes are preprocessed and mesh stable poses are calculated offline. SceneCollisionNet was trained using the [ACRONYM dataset](https://sites.google.com/nvidia.com/graspdataset). To use this dataset for training or benchmarking, download the ShapeNetSem meshes [here](https://shapenet.org/) (note: you must first register for an account) and the ACRONYM grasps [here](https://sites.google.com/nvidia.com/graspdataset). Next, build Manifold (an external library included as a submodule): 28 | ```shell 29 | ./scripts/install_manifold.sh 30 | ``` 31 | 32 | Then, use the following script to generate a preprocessed version of the ACRONYM dataset: 33 | ```shell 34 | python tools/generate_acronym_dataset.py /path/to/shapenetsem/meshes /path/to/acronym datasets/shapenet 35 | ``` 36 | 37 | If you have your own set of meshes, run: 38 | ```shell 39 | python tools/generate_mesh_dataset.py /path/to/meshes datasets/your_dataset_name 40 | ``` 41 | Note that this dataset will not include grasp data, which is not needed for training or benchmarking SceneCollisionNet, but is be used for rolling out the MPPI policy. 42 | 43 | ## Training/Benchmarking with Docker 44 | First, install Docker and `nvidia-docker2` following the instructions [here](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#installing-on-ubuntu-and-debian). Pull the SceneCollisionNet docker image from DockerHub (tag `scenecollisionnet`) or build locally using the provided Dockerfile (`docker build -t scenecollisionnet .`). Then, use the appropriate configuration `.yaml` file in `cfg` to set training or benchmarking parameters (note that cfg file paths are relative to the Docker container, not the local machine) and run one of the commands below (replacing paths with your local paths as needed; `-v` requires absolute paths). 45 | 46 | ### Train a SceneCollisionNet 47 | Edit `cfg/train_scenecollisionnet.yaml`, then run: 48 | ```shell 49 | docker run --gpus all --rm -it -v /path/to/dataset:/dataset:ro -v /path/to/models:/models:rw -v /path/to/cfg:/cfg:ro scenecollisionnet /SceneCollisionNet/scripts/train_scenecollisionnet_docker.sh 50 | ``` 51 | 52 | ### Train a RobotCollisionNet 53 | Edit `cfg/train_robotcollisionnet.yaml`, then run: 54 | ```shell 55 | docker run --gpus all --rm -it -v /path/to/models:/models:rw -v /path/to/cfg:/cfg:ro scenecollisionnet /SceneCollisionNet/scripts/train_robotcollisionnet_docker.sh 56 | ``` 57 | 58 | ### Benchmark a SceneCollisionNet 59 | Edit `cfg/benchmark_scenecollisionnet.yaml`, then run: 60 | ```shell 61 | docker run --gpus all --rm -it -v /path/to/dataset:/dataset:ro -v /path/to/models:/models:ro -v /path/to/cfg:/cfg:ro -v /path/to/benchmark_results:/benchmark:rw scenecollisionnet /SceneCollisionNet/scripts/benchmark_scenecollisionnet_docker.sh 62 | ``` 63 | 64 | 65 | ### Benchmark a RobotCollisionNet 66 | Edit `cfg/benchmark_robotcollisionnet.yaml`, then run: 67 | ```shell 68 | docker run --gpus all --rm -it -v /path/to/models:/models:rw -v /path/to/cfg:/cfg:ro -v /path/to/benchmark_results:/benchmark:rw scenecollisionnet /SceneCollisionNet/scripts/train_robotcollisionnet_docker.sh 69 | ``` 70 | 71 | ### Loss Plots 72 | To get loss plots while training, run: 73 | ```shell 74 | docker exec -d python3 tools/loss_plots.py /models//log.csv 75 | ``` 76 | 77 | ### Benchmark FCL or SDF Baselines 78 | Edit `cfg/benchmark_baseline.yaml`, then run: 79 | ```shell 80 | docker run --gpus all --rm -it -v /path/to/dataset:/dataset:ro -v /path/to/benchmark_results:/benchmark:rw -v /path/to/cfg:/cfg:ro scenecollisionnet /SceneCollisionNet/scripts/benchmark_baseline_docker.sh 81 | ``` 82 | 83 | ## Training/Benchmarking without Docker 84 | First, install system dependencies. The system dependencies listed assume an Ubuntu 18.04 install with NVIDIA drivers >= 450.80.02 and CUDA 10.2. You can adjust the dependencies accordingly for different driver/CUDA versions. Note that the NVIDIA drivers come packaged with EGL, which is used during training and benchmarking for headless rendering on the GPU. 85 | 86 | ### System Dependencies 87 | See Dockerfile for a full list. For training/benchmarking, you will need: 88 | ``` 89 | python3-dev 90 | python3-pip 91 | ninja-build 92 | libcudnn8=8.1.1.33-1+cuda10.2 93 | libcudnn8-dev=8.1.1.33-1+cuda10.2 94 | libsm6 95 | libxext6 96 | libxrender-dev 97 | freeglut3-dev 98 | liboctomap-dev 99 | libfcl-dev 100 | gifsicle 101 | libfreetype6-dev 102 | libpng-dev 103 | ``` 104 | 105 | ### Python Dependencies 106 | Follow the instructions above to install the necessary dependencies for your use case (either the `train`, `bench`, or `policy` options). 107 | 108 | ### Train a SceneCollisionNet 109 | Edit `cfg/train_scenecollisionnet.yaml`, then run: 110 | ```shell 111 | PYOPENGL_PLATFORM=egl python tools/train_scenecollisionnet.py 112 | ``` 113 | 114 | ### Train a RobotCollisionNet 115 | Edit `cfg/train_robotcollisionnet.yaml`, then run: 116 | ```shell 117 | python tools/train_robotcollisionnet.py 118 | ``` 119 | 120 | ### Benchmark a SceneCollisionNet 121 | Edit `cfg/benchmark_scenecollisionnet.yaml`, then run: 122 | ```shell 123 | PYOPENGL_PLATFORM=egl python tools/benchmark_scenecollisionnet.py 124 | ``` 125 | 126 | ### Benchmark a RobotCollisionNet 127 | Edit `cfg/benchmark_robotcollisionnet.yaml`, then run: 128 | ```shell 129 | python tools/benchmark_robotcollisionnet.py 130 | ``` 131 | 132 | ### Benchmark FCL or SDF Baselines 133 | Edit `cfg/benchmark_baseline.yaml`, then run: 134 | ```shell 135 | PYOPENGL_PLATFORM=egl python tools/benchmark_baseline.py 136 | ``` 137 | 138 | ## Policy Rollout 139 | To view a rearrangement MPPI policy rollout in a simulated Isaac Gym tabletop environment, run the following command (note that this requires a local machine with an available GPU and display): 140 | ```shell 141 | python tools/rollout_policy.py --self-coll-nn weights/self_coll_nn --scene-coll-nn weights/scene_coll_nn --control-frequency 1 142 | ``` 143 | There are many possible options for this command that can be viewed using the `--help` command line argument and set with the appropriate argument. If you get `RuntimeError: CUDA out of memory`, try reducing the horizon (`--mppi-horizon`, default 40), number of trajectories (`--mppi-num-rollouts`, default 200) or collision steps (`--mppi-collision-steps`, default 10). Note that this may affect policy performance. 144 | 145 | ## Citation 146 | If you use this code in your own research, please consider citing: 147 | ``` 148 | @inproceedings{danielczuk2021object, 149 | title={Object Rearrangement Using Learned Implicit Collision Functions}, 150 | author={Danielczuk, Michael and Mousavian, Arsalan and Eppner, Clemens and Fox, Dieter}, 151 | booktitle={Proc. IEEE Int. Conf. Robotics and Automation (ICRA)}, 152 | year={2021} 153 | } 154 | ``` 155 | -------------------------------------------------------------------------------- /data/panda/panda.urdf: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | 85 | 86 | 87 | 88 | 89 | 90 | 91 | 92 | 93 | 94 | 95 | 96 | 97 | 98 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | 134 | 135 | 136 | 137 | 138 | 139 | 140 | 141 | 142 | 143 | 145 | 146 | 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | 160 | 161 | 162 | 163 | 164 | 165 | 166 | 167 | 168 | 169 | 170 | 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | 211 | 212 | 213 | 214 | 215 | 216 | 217 | 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 | 234 | 235 | 236 | 237 | 238 | 239 | 240 | 241 | 242 | -------------------------------------------------------------------------------- /scenecollisionnet/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | from concurrent.futures import ProcessPoolExecutor 4 | from multiprocessing import Process, Queue 5 | from queue import Empty 6 | from typing import Any, Callable, Dict, Iterable 7 | 8 | import h5py 9 | import numpy as np 10 | import trimesh 11 | import trimesh.transformations as tra 12 | 13 | 14 | class MeshLoader(object): 15 | """A tool for loading meshes from a base directory. 16 | Attributes 17 | ---------- 18 | basedir : str 19 | basedir containing mesh files 20 | """ 21 | 22 | def __init__(self, basedir, only_ext=None): 23 | self.basedir = basedir 24 | self._map = {} 25 | for root, _, fns in os.walk(basedir): 26 | for fn in fns: 27 | full_fn = os.path.join(root, fn) 28 | f, ext = os.path.splitext(fn) 29 | if only_ext is not None and ext != only_ext: 30 | continue 31 | if basedir != root: 32 | f = os.path.basename(root) + "~" + f 33 | if ext[1:] not in trimesh.available_formats(): 34 | continue 35 | if f in self._map: 36 | continue 37 | # raise ValueError('Duplicate file named {}'.format(f)) 38 | self._map[f] = full_fn 39 | 40 | def meshes(self): 41 | return self._map.keys() 42 | 43 | def get_path(self, name): 44 | if name in self._map: 45 | return self._map[name] 46 | raise ValueError( 47 | "Could not find mesh with name {} in directory {}".format( 48 | name, self.basedir 49 | ) 50 | ) 51 | 52 | def load(self, name): 53 | m = trimesh.load(self.get_path(name)) 54 | m.metadata["name"] = name 55 | return m 56 | 57 | 58 | class ProcessKillingExecutor: 59 | """ 60 | The ProcessKillingExecutor works like an `Executor 61 | `_ 62 | in that it uses a bunch of processes to execute calls to a function with 63 | different arguments asynchronously. 64 | 65 | But other than the `ProcessPoolExecutor 66 | `_, 67 | the ProcessKillingExecutor forks a new Process for each function call that 68 | terminates after the function returns or if a timeout occurs. 69 | 70 | This means that contrary to the Executors and similar classes provided by 71 | the Python Standard Library, you can rely on the fact that a process will 72 | get killed if a timeout occurs and that absolutely no side can occur 73 | between function calls. 74 | 75 | Note that descendant processes of each process will not be terminated – 76 | they will simply become orphaned. 77 | """ 78 | 79 | def __init__(self, max_workers: int = None): 80 | self.processes = max_workers or os.cpu_count() 81 | 82 | def map( 83 | self, 84 | func: Callable, 85 | iterable: Iterable, 86 | timeout: float = None, 87 | callback_timeout: Callable = None, 88 | daemon: bool = True, 89 | ) -> Iterable: 90 | """ 91 | :param func: the function to execute 92 | :param iterable: an iterable of function arguments 93 | :param timeout: after this time, the process executing the function 94 | will be killed if it did not finish 95 | :param callback_timeout: this function will be called, if the task 96 | times out. It gets the same arguments as the original function 97 | :param daemon: define the child process as daemon 98 | """ 99 | executor = ProcessPoolExecutor(max_workers=self.processes) 100 | params = ( 101 | { 102 | "func": func, 103 | "fn_args": p_args, 104 | "p_kwargs": {}, 105 | "timeout": timeout, 106 | "callback_timeout": callback_timeout, 107 | "daemon": daemon, 108 | } 109 | for p_args in iterable 110 | ) 111 | return executor.map(self._submit_unpack_kwargs, params) 112 | 113 | def _submit_unpack_kwargs(self, params): 114 | """unpack the kwargs and call submit""" 115 | 116 | return self.submit(**params) 117 | 118 | def submit( 119 | self, 120 | func: Callable, 121 | fn_args: Any, 122 | p_kwargs: Dict, 123 | timeout: float, 124 | callback_timeout: Callable[[Any], Any], 125 | daemon: bool, 126 | ): 127 | """ 128 | Submits a callable to be executed with the given arguments. 129 | Schedules the callable to be executed as func(*args, **kwargs) in a new 130 | process. 131 | :param func: the function to execute 132 | :param fn_args: the arguments to pass to the function. Can be one argument 133 | or a tuple of multiple args. 134 | :param p_kwargs: the kwargs to pass to the function 135 | :param timeout: after this time, the process executing the function 136 | will be killed if it did not finish 137 | :param callback_timeout: this function will be called with the same 138 | arguments, if the task times out. 139 | :param daemon: run the child process as daemon 140 | :return: the result of the function, or None if the process failed or 141 | timed out 142 | """ 143 | p_args = fn_args if isinstance(fn_args, tuple) else (fn_args,) 144 | queue = Queue() 145 | p = Process( 146 | target=self._process_run, 147 | args=( 148 | queue, 149 | func, 150 | *p_args, 151 | ), 152 | kwargs=p_kwargs, 153 | ) 154 | 155 | if daemon: 156 | p.daemon = True 157 | 158 | p.start() 159 | try: 160 | ret = queue.get(block=True, timeout=timeout) 161 | if ret is None: 162 | callback_timeout(*p_args, **p_kwargs) 163 | return ret 164 | except Empty: 165 | callback_timeout(*p_args, **p_kwargs) 166 | if p.is_alive(): 167 | p.terminate() 168 | p.join() 169 | 170 | @staticmethod 171 | def _process_run( 172 | queue: Queue, func: Callable[[Any], Any] = None, *args, **kwargs 173 | ): 174 | """ 175 | Executes the specified function as func(*args, **kwargs). 176 | The result will be stored in the shared dictionary 177 | :param func: the function to execute 178 | :param queue: a Queue 179 | """ 180 | queue.put(func(*args, **kwargs)) 181 | 182 | 183 | def compute_camera_pose(distance, azimuth, elevation): 184 | cam_tf = tra.euler_matrix(np.pi / 2, 0, 0).dot( 185 | tra.euler_matrix(0, np.pi / 2, 0) 186 | ) 187 | 188 | extrinsics = np.eye(4) 189 | extrinsics[0, 3] += distance 190 | extrinsics = tra.euler_matrix(0, -elevation, azimuth).dot(extrinsics) 191 | 192 | cam_pose = extrinsics.dot(cam_tf) 193 | frame_pose = cam_pose.copy() 194 | frame_pose[:, 1:3] *= -1.0 195 | return cam_pose, frame_pose 196 | 197 | 198 | def process_mesh(in_path, out_path, scale, grasps, return_stps=True): 199 | mesh = trimesh.load(in_path, force="mesh", skip_materials=True) 200 | cat = ( 201 | "" 202 | if os.path.basename(os.path.dirname(out_path)) == "meshes" 203 | else os.path.basename(os.path.dirname(out_path)) 204 | ) 205 | if not os.path.exists(os.path.dirname(out_path)): 206 | os.makedirs(os.path.dirname(out_path)) 207 | 208 | if not mesh.is_watertight or len(mesh.faces) > 1000: 209 | obj_path = os.path.splitext(in_path)[0] + ".obj" 210 | is_obj = os.path.exists(obj_path) 211 | if not is_obj: 212 | mesh.export(obj_path) 213 | 214 | simplify_path = "{}_{:d}.obj".format( 215 | os.path.splitext(out_path)[0], 216 | np.random.RandomState().randint(2 ** 16), 217 | ) 218 | manifold_cmd = [ 219 | os.path.join( 220 | os.path.dirname(os.path.abspath(__file__)), 221 | "../extern/Manifold/build/manifold", 222 | ), 223 | obj_path, 224 | simplify_path, 225 | ] 226 | simplify_cmd = [ 227 | os.path.join( 228 | os.path.dirname(os.path.abspath(__file__)), 229 | "../extern/Manifold/build/simplify", 230 | ), 231 | "-i", 232 | simplify_path, 233 | "-o", 234 | simplify_path, 235 | "-m", 236 | "-f 1000", 237 | ] 238 | try: 239 | subprocess.check_output(manifold_cmd) 240 | except subprocess.CalledProcessError: 241 | if not is_obj: 242 | os.remove(obj_path) 243 | if os.path.exists(simplify_path): 244 | os.remove(simplify_path) 245 | return None 246 | if not is_obj: 247 | os.remove(obj_path) 248 | try: 249 | subprocess.check_output(simplify_cmd) 250 | except subprocess.CalledProcessError: 251 | if os.path.exists(simplify_path): 252 | os.remove(simplify_path) 253 | return None 254 | 255 | mesh = trimesh.load(simplify_path) 256 | os.remove(simplify_path) 257 | 258 | # Create final scaled and transformed mesh 259 | if not mesh.is_watertight: 260 | mesh.center_mass = mesh.centroid 261 | 262 | mesh.apply_scale(scale) 263 | mesh_offset = mesh.center_mass 264 | mesh.apply_transform( 265 | trimesh.transformations.translation_matrix(-mesh_offset) 266 | ) 267 | m_scale = "{}_{}.obj".format( 268 | os.path.splitext(os.path.basename(out_path))[0], scale 269 | ) 270 | s_out_path = os.path.join(os.path.dirname(out_path), m_scale) 271 | mesh.export(s_out_path) 272 | 273 | m_info = { 274 | "path": os.path.join(cat, m_scale), 275 | "scale": scale, 276 | "category": cat, 277 | } 278 | 279 | # Calculate stable poses and add grasps 280 | if return_stps: 281 | try: 282 | stps, probs = mesh.compute_stable_poses() 283 | if not probs.any(): 284 | os.remove(s_out_path) 285 | return None 286 | m_info.update({"stps": stps, "probs": probs / probs.sum()}) 287 | except Exception: 288 | os.remove(s_out_path) 289 | return None 290 | 291 | if grasps is not None: 292 | grasp_data = h5py.File(grasps, "r")["grasps"] 293 | positive_grasps = grasp_data["transforms"][:][ 294 | grasp_data["qualities/flex/object_in_gripper"][:] > 0 295 | ] 296 | positive_grasps[:, :3, 3] -= mesh_offset 297 | m_info["grasps"] = positive_grasps 298 | 299 | return os.path.splitext(m_scale)[0], m_info 300 | -------------------------------------------------------------------------------- /scenecollisionnet/collision_models/collision_nets.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch_scatter 5 | from pointnet2.pointnet2_modules import PointnetSAModule 6 | from pointnet2.pytorch_utils import FC, Conv1d, Conv3d 7 | from torch.cuda.amp import autocast 8 | 9 | OBJ_NPOINTS = [256, 64, None] 10 | OBJ_RADII = [0.02, 0.04, None] 11 | OBJ_NSAMPLES = [64, 128, None] 12 | OBJ_MLPS = [[0, 64, 128], [128, 128, 256], [256, 256, 512]] 13 | SCENE_PT_MLP = [3, 128, 256] 14 | SCENE_VOX_MLP = [256, 512, 1024, 512] 15 | CLS_FC = [2057, 1024, 256] 16 | 17 | 18 | class SceneCollisionNet(nn.Module): 19 | def __init__(self, bounds, vox_size): 20 | super().__init__() 21 | self.bounds = nn.Parameter( 22 | torch.from_numpy(np.asarray(bounds)).float(), requires_grad=False 23 | ) 24 | self.vox_size = nn.Parameter( 25 | torch.from_numpy(np.asarray(vox_size)).float(), requires_grad=False 26 | ) 27 | self.num_voxels = nn.Parameter( 28 | ((self.bounds[1] - self.bounds[0]) / self.vox_size).long(), 29 | requires_grad=False, 30 | ) 31 | 32 | self.scene_pt_mlp = nn.Sequential() 33 | for i in range(len(SCENE_PT_MLP) - 1): 34 | self.scene_pt_mlp.add_module( 35 | "pt_layer{}".format(i), 36 | Conv1d(SCENE_PT_MLP[i], SCENE_PT_MLP[i + 1], first=(i == 0)), 37 | ) 38 | 39 | self.scene_vox_mlp = nn.ModuleList() 40 | for i in range(len(SCENE_VOX_MLP) - 1): 41 | scene_conv = nn.Sequential() 42 | if SCENE_VOX_MLP[i + 1] > SCENE_VOX_MLP[i]: 43 | scene_conv.add_module( 44 | "3d_conv_layer{}".format(i), 45 | Conv3d( 46 | SCENE_VOX_MLP[i], 47 | SCENE_VOX_MLP[i + 1], 48 | kernel_size=3, 49 | padding=1, 50 | ), 51 | ) 52 | scene_conv.add_module( 53 | "3d_max_layer{}".format(i), nn.MaxPool3d(2, stride=2) 54 | ) 55 | else: 56 | scene_conv.add_module( 57 | "3d_convt_layer{}".format(i), 58 | nn.ConvTranspose3d( 59 | SCENE_VOX_MLP[i], 60 | SCENE_VOX_MLP[i + 1], 61 | kernel_size=2, 62 | stride=2, 63 | ), 64 | ) 65 | self.scene_vox_mlp.append(scene_conv) 66 | 67 | self.obj_SA_modules = nn.ModuleList() 68 | for k in range(OBJ_NPOINTS.__len__()): 69 | self.obj_SA_modules.append( 70 | PointnetSAModule( 71 | npoint=OBJ_NPOINTS[k], 72 | radius=OBJ_RADII[k], 73 | nsample=OBJ_NSAMPLES[k], 74 | mlp=OBJ_MLPS[k], 75 | bn=False, 76 | first=(i == 0), 77 | ) 78 | ) 79 | 80 | self.obj_FCs = nn.ModuleList( 81 | [ 82 | FC(OBJ_MLPS[-1][-1], 1024), 83 | FC(1024, 1024), 84 | ] 85 | ) 86 | 87 | self.classifier = nn.Sequential( 88 | FC(CLS_FC[0], CLS_FC[1], first=True), 89 | FC(CLS_FC[1], CLS_FC[2]), 90 | FC(CLS_FC[2], 1, activation=None), 91 | ) 92 | 93 | def _break_up_pc(self, pc): 94 | xyz = pc[..., 0:3].contiguous() 95 | features = ( 96 | pc[..., 3:].transpose(1, 2).contiguous() 97 | if pc.size(-1) > 3 98 | else None 99 | ) 100 | 101 | return xyz, features 102 | 103 | def _inds_to_flat(self, inds, scale=1): 104 | flat_inds = inds * torch.cuda.IntTensor( 105 | [ 106 | self.num_voxels[1:].prod() // (scale ** 2), 107 | self.num_voxels[2] // scale, 108 | 1, 109 | ], 110 | device=self.num_voxels.device, 111 | ) 112 | return flat_inds.sum(axis=-1) 113 | 114 | def _inds_from_flat(self, flat_inds, scale=1): 115 | ind0 = flat_inds // (self.num_voxels[1:].prod() // (scale ** 2)) 116 | ind1 = (flat_inds % (self.num_voxels[1:].prod() // (scale ** 2))) // ( 117 | self.num_voxels[2] // scale 118 | ) 119 | ind2 = (flat_inds % (self.num_voxels[1:].prod() // (scale ** 2))) % ( 120 | self.num_voxels[2] // scale 121 | ) 122 | return torch.stack((ind0, ind1, ind2), dim=-1) 123 | 124 | def voxel_inds(self, xyz, scale=1): 125 | inds = ((xyz - self.bounds[0]) // (scale * self.vox_size)).int() 126 | return self._inds_to_flat(inds, scale=scale) 127 | 128 | def get_obj_features(self, obj_pc): 129 | obj_xyz, obj_features = self._break_up_pc(obj_pc) 130 | 131 | # Featurize obj 132 | for i in range(len(self.obj_SA_modules)): 133 | obj_xyz, obj_features = self.obj_SA_modules[i]( 134 | obj_xyz, obj_features 135 | ) 136 | for i in range(len(self.obj_FCs)): 137 | obj_features = self.obj_FCs[i](obj_features.squeeze(axis=-1)) 138 | 139 | return obj_features 140 | 141 | def get_scene_features(self, scene_pc): 142 | scene_xyz, scene_features = self._break_up_pc(scene_pc) 143 | scene_inds = self.voxel_inds(scene_xyz) 144 | 145 | # Featurize scene points and max pool over voxels 146 | scene_vox_centers = ( 147 | self._inds_from_flat(scene_inds) * self.vox_size 148 | + self.vox_size / 2 149 | + self.bounds[0] 150 | ) 151 | scene_xyz_centered = (scene_pc[..., :3] - scene_vox_centers).transpose( 152 | 2, 1 153 | ) 154 | if scene_features is not None: 155 | scene_features = self.scene_pt_mlp( 156 | torch.cat((scene_xyz_centered, scene_features), dim=1) 157 | ) 158 | else: 159 | scene_features = self.scene_pt_mlp(scene_xyz_centered) 160 | max_vox_features = torch.zeros( 161 | (*scene_features.shape[:2], self.num_voxels.prod()) 162 | ).to(scene_pc.device) 163 | if scene_inds.max() >= self.num_voxels.prod(): 164 | print( 165 | scene_xyz[range(len(scene_pc)), scene_inds.max(axis=-1)[1]], 166 | scene_inds.max(), 167 | ) 168 | assert scene_inds.max() < self.num_voxels.prod() 169 | assert scene_inds.min() >= 0 170 | 171 | with autocast(enabled=False): 172 | max_vox_features[ 173 | ..., : scene_inds.max() + 1 174 | ] = torch_scatter.scatter_max( 175 | scene_features.float(), scene_inds[:, None, :] 176 | )[ 177 | 0 178 | ] 179 | max_vox_features = max_vox_features.reshape( 180 | *max_vox_features.shape[:2], *self.num_voxels.int() 181 | ) 182 | 183 | # 3D conv over voxels 184 | l_vox_features = [max_vox_features] 185 | for i in range(len(self.scene_vox_mlp)): 186 | li_vox_features = self.scene_vox_mlp[i](l_vox_features[i]) 187 | l_vox_features.append(li_vox_features) 188 | 189 | # Stack features from different levels 190 | stack_vox_features = torch.cat( 191 | (l_vox_features[1], l_vox_features[-1]), dim=1 192 | ) 193 | stack_vox_features = stack_vox_features.reshape( 194 | *stack_vox_features.shape[:2], -1 195 | ) 196 | return stack_vox_features 197 | 198 | # scene_pc: (b, n_scene_fts); obj_pc: (b, n_obj_fts); 199 | # trans: (q, 3); rots: (q, 6) 200 | def classify_tfs(self, obj_features, scene_features, trans, rots): 201 | b = len(scene_features) 202 | q = len(trans) 203 | 204 | # Get voxel indices for translations 205 | trans_inds = self.voxel_inds(trans, scale=2).long() 206 | if trans_inds.max() >= scene_features.shape[2]: 207 | print(trans[trans_inds.argmax()], trans_inds.max()) 208 | assert trans_inds.max() < scene_features.shape[2] 209 | assert trans_inds.min() >= 0 210 | vox_trans_features = scene_features[..., trans_inds].transpose(2, 1) 211 | 212 | # Calculate translation offsets from centers of voxels 213 | tr_vox_centers = ( 214 | self._inds_from_flat(trans_inds, scale=2) * self.vox_size * 2 215 | + self.vox_size / 2 216 | + self.bounds[0] 217 | ) 218 | trans_offsets = trans - tr_vox_centers.float() 219 | 220 | # Send concatenated features to classifier 221 | class_in = torch.cat( 222 | ( 223 | obj_features.unsqueeze(1).expand(b, q, obj_features.shape[-1]), 224 | vox_trans_features, 225 | trans_offsets.unsqueeze(0).expand(b, q, 3), 226 | rots.unsqueeze(0).expand(b, q, 6), 227 | ), 228 | dim=-1, 229 | ) 230 | 231 | return self.classifier(class_in) 232 | 233 | # scene_pc: (n_scene_fts); obj_pc: (b, n_obj_fts); 234 | # trans: (b, q, 3); rots: (b, q, 6) 235 | def classify_multi_obj_tfs( 236 | self, obj_features, scene_features, trans, rots 237 | ): 238 | b, q = trans.shape[:2] 239 | 240 | # Get voxel indices for translations 241 | trans_inds = self.voxel_inds(trans, scale=2).long() 242 | if trans_inds.max() >= scene_features.shape[-1]: 243 | print(trans[trans_inds.argmax()], trans_inds.max()) 244 | assert trans_inds.max() < scene_features.shape[-1] 245 | assert trans_inds.min() >= 0 246 | vox_trans_features = scene_features[..., trans_inds].permute(1, 2, 0) 247 | 248 | # Calculate translation offsets from centers of voxels 249 | tr_vox_centers = ( 250 | self._inds_from_flat(trans_inds, scale=2) * self.vox_size * 2 251 | + self.vox_size / 2 252 | + self.bounds[0] 253 | ) 254 | trans_offsets = trans - tr_vox_centers.float() 255 | 256 | # Send concatenated features to classifier 257 | class_in = torch.cat( 258 | ( 259 | obj_features.unsqueeze(1).expand(-1, q, -1), 260 | vox_trans_features, 261 | trans_offsets, 262 | rots, 263 | ), 264 | dim=-1, 265 | ) 266 | 267 | return self.classifier(class_in) 268 | 269 | # scene_pc: (b, n_scene_pts, 6); obj_pc: (b, n_obj_pts, 6); 270 | # trans: (b, q, 3); rots: (b, q, 6) 271 | def forward(self, scene_pc, obj_pc, trans, rots): 272 | 273 | obj_features = self.get_obj_features(obj_pc) 274 | scene_features = self.get_scene_features(scene_pc) 275 | return self.classify_tfs(obj_features, scene_features, trans, rots) 276 | 277 | 278 | class RobotCollisionNet(nn.Module): 279 | def __init__(self, num_joints): 280 | super().__init__() 281 | self.classifier = nn.Sequential( 282 | FC(num_joints, 64), 283 | FC(64, 64), 284 | FC(64, 64), 285 | FC(64, 64), 286 | FC(64, 64), 287 | FC(64, 1, activation=None), 288 | ) 289 | 290 | def forward(self, centers): 291 | return self.classifier(centers) 292 | -------------------------------------------------------------------------------- /tools/benchmark_baseline.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import os.path as osp 4 | from timeit import default_timer as timer 5 | 6 | import imageio 7 | import matplotlib as mpl 8 | import numpy as np 9 | import pygifsicle 10 | import tqdm 11 | from torch.utils.data import DataLoader 12 | 13 | mpl.use("agg") 14 | import matplotlib.pyplot as plt 15 | import seaborn as sns 16 | from autolab_core import BinaryClassificationResult, YamlConfig, keyboard_input 17 | 18 | from scenecollisionnet.collision_models import ( 19 | FCLMultiSceneCollisionChecker, 20 | SDFObjCollisionChecker, 21 | SDFSceneCollisionChecker, 22 | ) 23 | from scenecollisionnet.datasets import BenchmarkSceneCollisionDataset 24 | 25 | sns.set() 26 | 27 | 28 | def benchmark(model, test_loader, iterations, out_path, vis=True): 29 | time_start = timer() 30 | 31 | if vis: 32 | vis_path = osp.join(out_path, "vis") 33 | if not osp.exists(vis_path): 34 | os.mkdir(vis_path) 35 | 36 | preds, trues = np.array([]), np.array([]) 37 | data_loader = iter(test_loader) 38 | passes = [] 39 | coll_times = [] 40 | for batch_idx in tqdm.trange( 41 | iterations, 42 | desc="Benchmarking", 43 | ): 44 | data = next(data_loader) 45 | coll_time = data[-1] 46 | coll_times.append(coll_time) 47 | if vis: 48 | scene_manager, obj_mesh, obj_pose = data[:3] 49 | scene, obj, trans, rots, coll = data[3:-1] 50 | else: 51 | scene, obj, trans, rots, coll = data[:-1] 52 | fw_pass_start = timer() 53 | out = model(scene, obj, trans, rots) 54 | passes.append(timer() - fw_pass_start) 55 | 56 | preds = np.append(preds, out) 57 | trues = np.append(trues, coll.numpy()) 58 | 59 | if vis: 60 | obj_pose = obj_pose.numpy() 61 | trans = trans.numpy() 62 | rots = rots.numpy() 63 | scene = scene[0, :, :3].numpy() 64 | obj = obj[0, :, :3].numpy() 65 | obj_centroid = obj.mean(axis=0) 66 | with imageio.get_writer( 67 | osp.join(vis_path, "vis_{:d}.gif".format(batch_idx)), 68 | mode="I", 69 | duration=0.25, 70 | ) as writer: 71 | for i, tr, rot, pr, gt in zip( 72 | np.arange(len(trans)), trans, rots, out, coll.numpy() 73 | ): 74 | mesh_trans = tr - (obj_centroid - obj_pose[:3, 3]) 75 | mesh_tf = np.eye(4) 76 | mesh_tf[:3, 3] = mesh_trans 77 | if (rot != 0).any(): 78 | b1 = rot[:3] 79 | b2 = rot[3:] - b1.dot(rot[3:]) * b1 80 | b2 /= np.linalg.norm(b2) 81 | b3 = np.cross(b1, b2) 82 | mesh_tf[:3, :3] = np.stack((b1, b2, b3), axis=-1) 83 | new_pose = mesh_tf @ obj_pose 84 | 85 | # Render images with full meshes 86 | pred_color = plt.get_cmap("hsv")(0.3 * (1.0 - pr))[:-1] 87 | scene_manager.add_object( 88 | "pred_coll_obj", 89 | obj_mesh, 90 | pose=new_pose, 91 | color=pred_color, 92 | ) 93 | pred_im, _ = scene_manager._renderer.render_rgbd() 94 | scene_manager.remove_object("pred_coll_obj") 95 | gt_color = plt.get_cmap("hsv")(0.3 * (1.0 - gt))[:-1] 96 | scene_manager.add_object( 97 | "gt_coll_obj", obj_mesh, pose=new_pose, color=gt_color 98 | ) 99 | gt_im, _ = scene_manager._renderer.render_rgbd() 100 | scene_manager.remove_object("gt_coll_obj") 101 | 102 | # Render images with wireframe meshes and point clouds 103 | scene_dists = np.linalg.norm( 104 | scene_manager._renderer._camera_node.matrix[:3, 3] 105 | - scene, 106 | axis=1, 107 | ) 108 | scene_colors = (scene_dists - scene_dists.min()) / ( 109 | scene_dists.max() - scene_dists.min() 110 | ) 111 | for n in scene_manager._renderer._node_dict: 112 | scene_manager._renderer._node_dict[n].mesh.primitives[ 113 | 0 114 | ].material.wireframe = True 115 | scene_manager._renderer.add_points( 116 | scene, 117 | "scene_points", 118 | color=plt.get_cmap("hsv")(scene_colors), 119 | radius=0.005, 120 | ) 121 | 122 | scene_manager.add_object( 123 | "pred_coll_obj", 124 | obj_mesh, 125 | pose=new_pose, 126 | color=pred_color, 127 | ) 128 | pred_pts_im, _ = scene_manager._renderer.render_rgbd() 129 | scene_manager.remove_object("pred_coll_obj") 130 | scene_manager.add_object( 131 | "gt_coll_obj", obj_mesh, pose=new_pose, color=gt_color 132 | ) 133 | gt_pts_im, _ = scene_manager._renderer.render_rgbd() 134 | scene_manager.remove_object("gt_coll_obj") 135 | scene_manager._renderer.remove_object("scene_points") 136 | for n in scene_manager._renderer._node_dict: 137 | scene_manager._renderer._node_dict[n].mesh.primitives[ 138 | 0 139 | ].material.wireframe = False 140 | 141 | top_row = np.concatenate( 142 | ( 143 | np.pad( 144 | gt_im.data, 145 | ((0, 10), (0, 10), (0, 0)), 146 | constant_values=np.iinfo(np.uint8).max, 147 | ), 148 | np.pad( 149 | pred_im.data, 150 | ((0, 10), (0, 0), (0, 0)), 151 | constant_values=np.iinfo(np.uint8).max, 152 | ), 153 | ), 154 | axis=1, 155 | ) 156 | bot_row = np.concatenate( 157 | ( 158 | np.pad( 159 | gt_pts_im.data, 160 | ((0, 0), (0, 10), (0, 0)), 161 | constant_values=np.iinfo(np.uint8).max, 162 | ), 163 | np.pad( 164 | pred_pts_im.data, 165 | ((0, 0), (0, 0), (0, 0)), 166 | constant_values=np.iinfo(np.uint8).max, 167 | ), 168 | ), 169 | axis=1, 170 | ) 171 | full_gif_im = np.concatenate((top_row, bot_row), axis=0) 172 | writer.append_data(full_gif_im) 173 | scene_manager._renderer._renderer.delete() 174 | del scene_manager 175 | pygifsicle.optimize( 176 | osp.join(vis_path, "vis_{:d}.gif".format(batch_idx)) 177 | ) 178 | 179 | if not vis: 180 | # sort by probs 181 | accs = [] 182 | tprs = [] 183 | f1s = [] 184 | taus = np.linspace(0, 1, 21, endpoint=False) 185 | for t in taus: 186 | b = BinaryClassificationResult(preds, trues, threshold=t) 187 | accs.append(b.accuracy) 188 | f1s.append(b.f1_score) 189 | tprs.append(b.tpr) 190 | 191 | np.savez_compressed( 192 | osp.join(out_path, "accuracy_curve.npz"), 193 | taus=taus, 194 | accs=accs, 195 | f1s=f1s, 196 | tprs=tprs, 197 | ) 198 | 199 | ax = sns.lineplot(x=taus, y=accs, label="Accuracy") 200 | ax = sns.lineplot(x=taus, y=f1s, label="F1") 201 | ax = sns.lineplot(x=taus, y=tprs, label="TPR") 202 | ax.legend() 203 | plt.xlabel("Threshold") 204 | plt.ylabel("Score") 205 | plt.ylim([0.5, 1.0]) 206 | plt.xlim([-0.01, 1.01]) 207 | plt.savefig(osp.join(out_path, "accuracy_curve.png")) 208 | 209 | bcr = BinaryClassificationResult(preds, trues) 210 | with open(osp.join(out_path, "results.txt"), "w") as f: 211 | elapsed_time = timer() - time_start 212 | log = ( 213 | ["Images: {:d}".format(iterations)] 214 | + ["Queries: {:d}".format(iterations * len(coll.flatten()))] 215 | + ["Accuracy: {:.3f}".format(bcr.accuracy)] 216 | + ["F1 Score: {:.3f}".format(bcr.f1_score)] 217 | + ["AP Score: {:.3f}".format(bcr.ap_score)] 218 | + [ 219 | "FW Pass Time: {:.4f} +- {:.4f} s".format( 220 | np.mean(passes), np.std(passes) 221 | ) 222 | ] 223 | + ["Time: {:.2f} s".format(elapsed_time)] 224 | + [ 225 | "FCL Time: {:.3f} +- {:.3f} s".format( 226 | np.mean(coll_times), np.std(coll_times) 227 | ) 228 | ] 229 | ) 230 | log = map(str, log) 231 | f.write("\n".join(log)) 232 | print("\n".join(log)) 233 | 234 | 235 | if __name__ == "__main__": 236 | 237 | parser = argparse.ArgumentParser( 238 | description="Benchmark a CollisionNet Baseline" 239 | ) 240 | parser.add_argument("--model", type=str) 241 | parser.add_argument("--processes", type=int, default=4) 242 | parser.add_argument("--vis", action="store_true") 243 | parser.add_argument( 244 | "--cfg", 245 | type=str, 246 | default="cfg/benchmark_baseline.yaml", 247 | help="config file with benchmarking params", 248 | ) 249 | args = parser.parse_args() 250 | 251 | config = YamlConfig(args.cfg) 252 | model_type = args.model 253 | if model_type is not None: 254 | config["model"] = model_type 255 | n_proc = args.processes 256 | config["vis"] = args.vis 257 | config["iterations"] = 25 if config["vis"] else 1000 258 | 259 | # Create output directory for model 260 | out = osp.join(config["output"], config["model"]["type"]) 261 | if osp.exists(out): 262 | response = keyboard_input( 263 | "A benchmark folder exists at {}. Overwrite?".format(out), 264 | yesno=True, 265 | ) 266 | if response.lower() == "n": 267 | os.exit() 268 | else: 269 | os.makedirs(out) 270 | config.save(osp.join(out, "benchmark.yaml")) 271 | 272 | # 1. Model 273 | if config["model"]["type"] == "fcl": 274 | model = FCLMultiSceneCollisionChecker(n_proc=n_proc) 275 | elif config["model"]["type"] == "sdf_obj": 276 | model = SDFObjCollisionChecker( 277 | device=config["model"]["device"], n_proc=n_proc 278 | ) 279 | elif config["model"]["type"] == "sdf_scene": 280 | model = SDFSceneCollisionChecker( 281 | device=config["model"]["device"], n_proc=n_proc 282 | ) 283 | else: 284 | raise ValueError("Model not supported") 285 | 286 | # 2. Dataset 287 | kwargs = { 288 | "num_workers": 0 289 | if config["vis"] 290 | else ( 291 | config["num_workers"] 292 | if "num_workers" in config 293 | else os.cpu_count() 294 | ), 295 | "pin_memory": True, 296 | "worker_init_fn": lambda _: np.random.seed(), 297 | } 298 | test_set = BenchmarkSceneCollisionDataset( 299 | **config["dataset"], vis=config["vis"] 300 | ) 301 | test_loader = DataLoader(test_set, batch_size=None, **kwargs) 302 | 303 | benchmark( 304 | model=model, 305 | test_loader=test_loader, 306 | iterations=config["iterations"], 307 | out_path=out, 308 | vis=config["vis"], 309 | ) 310 | -------------------------------------------------------------------------------- /tools/benchmark_scenecollisionnet.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import os.path as osp 4 | import sys 5 | from timeit import default_timer as timer 6 | 7 | import imageio 8 | import matplotlib as mpl 9 | import numpy as np 10 | import pygifsicle 11 | import torch 12 | import torch.nn as nn 13 | import tqdm 14 | from torch.autograd import Variable 15 | from torch.utils.data import DataLoader 16 | 17 | mpl.use("agg") 18 | import matplotlib.pyplot as plt 19 | import seaborn as sns 20 | from autolab_core import BinaryClassificationResult, YamlConfig, keyboard_input 21 | 22 | from scenecollisionnet.collision_models import SceneCollisionNet 23 | from scenecollisionnet.datasets import BenchmarkSceneCollisionDataset 24 | 25 | sns.set() 26 | 27 | 28 | def benchmark( 29 | device, model, criterion, test_loader, iterations, out_path, vis=True 30 | ): 31 | model.eval() 32 | time_start = timer() 33 | 34 | if vis: 35 | vis_path = osp.join(out_path, "vis") 36 | if not osp.exists(vis_path): 37 | os.mkdir(vis_path) 38 | 39 | test_loss = 0 40 | preds, trues = np.array([]), np.array([]) 41 | data_loader = iter(test_loader) 42 | passes = [] 43 | coll_times = [] 44 | for batch_idx in tqdm.trange( 45 | iterations, 46 | desc="Benchmarking", 47 | ): 48 | data = next(data_loader) 49 | coll_time = data[-1] 50 | coll_times.append(coll_time) 51 | if vis: 52 | scene_manager, obj_mesh, obj_pose = data[:3] 53 | data = [Variable(d.float().to(device)) for d in data[3:-1]] 54 | else: 55 | data = [Variable(d.float().to(device)) for d in data[:-1]] 56 | scene, obj, trans, rots, coll = data 57 | fw_pass_start = timer() 58 | with torch.no_grad(): 59 | out = model(scene, obj, trans, rots).squeeze(dim=-1) 60 | passes.append(timer() - fw_pass_start) 61 | 62 | coll = data[-1].type(out.type()) 63 | loss = criterion(out, coll) 64 | 65 | pred_batch = torch.sigmoid(out).cpu().numpy().flatten() 66 | true_batch = coll.cpu().numpy().flatten() 67 | preds = np.append(preds, pred_batch) 68 | trues = np.append(trues, true_batch) 69 | test_loss += loss.item() 70 | 71 | if vis: 72 | obj_pose = obj_pose.numpy() 73 | trans = trans.cpu().numpy() 74 | rots = rots.cpu().numpy() 75 | scene = scene[0, :, :3].cpu().numpy() 76 | obj = obj[0, :, :3].cpu().numpy() 77 | obj_centroid = obj.mean(axis=0) 78 | with imageio.get_writer( 79 | osp.join(vis_path, "vis_{:d}.gif".format(batch_idx)), 80 | mode="I", 81 | duration=0.25, 82 | ) as writer: 83 | for i, tr, rot, pr, gt in zip( 84 | np.arange(len(trans)), trans, rots, pred_batch, true_batch 85 | ): 86 | mesh_trans = tr - (obj_centroid - obj_pose[:3, 3]) 87 | mesh_tf = np.eye(4) 88 | mesh_tf[:3, 3] = mesh_trans 89 | if (rot != 0).any(): 90 | b1 = rot[:3] 91 | b2 = rot[3:] - b1.dot(rot[3:]) * b1 92 | b2 /= np.linalg.norm(b2) 93 | b3 = np.cross(b1, b2) 94 | mesh_tf[:3, :3] = np.stack((b1, b2, b3), axis=-1) 95 | new_pose = mesh_tf @ obj_pose 96 | 97 | # Render images with full meshes 98 | pred_color = plt.get_cmap("hsv")(0.3 * (1.0 - pr))[:-1] 99 | scene_manager.add_object( 100 | "pred_coll_obj", 101 | obj_mesh, 102 | pose=new_pose, 103 | color=pred_color, 104 | ) 105 | pred_im, _ = scene_manager._renderer.render_rgbd() 106 | scene_manager.remove_object("pred_coll_obj") 107 | gt_color = plt.get_cmap("hsv")(0.3 * (1.0 - gt))[:-1] 108 | scene_manager.add_object( 109 | "gt_coll_obj", obj_mesh, pose=new_pose, color=gt_color 110 | ) 111 | gt_im, _ = scene_manager._renderer.render_rgbd() 112 | scene_manager.remove_object("gt_coll_obj") 113 | 114 | # Render images with wireframe meshes and point clouds 115 | scene_dists = np.linalg.norm( 116 | scene_manager._renderer._camera_node.matrix[:3, 3] 117 | - scene, 118 | axis=1, 119 | ) 120 | scene_colors = (scene_dists - scene_dists.min()) / ( 121 | scene_dists.max() - scene_dists.min() 122 | ) 123 | for n in scene_manager._renderer._node_dict: 124 | scene_manager._renderer._node_dict[n].mesh.primitives[ 125 | 0 126 | ].material.wireframe = True 127 | scene_manager._renderer.add_points( 128 | scene, 129 | "scene_points", 130 | color=plt.get_cmap("hsv")(scene_colors), 131 | radius=0.005, 132 | ) 133 | 134 | scene_manager.add_object( 135 | "pred_coll_obj", 136 | obj_mesh, 137 | pose=new_pose, 138 | color=pred_color, 139 | ) 140 | pred_pts_im, _ = scene_manager._renderer.render_rgbd() 141 | scene_manager.remove_object("pred_coll_obj") 142 | scene_manager.add_object( 143 | "gt_coll_obj", obj_mesh, pose=new_pose, color=gt_color 144 | ) 145 | gt_pts_im, _ = scene_manager._renderer.render_rgbd() 146 | scene_manager.remove_object("gt_coll_obj") 147 | scene_manager._renderer.remove_object("scene_points") 148 | for n in scene_manager._renderer._node_dict: 149 | scene_manager._renderer._node_dict[n].mesh.primitives[ 150 | 0 151 | ].material.wireframe = False 152 | 153 | top_row = np.concatenate( 154 | ( 155 | np.pad( 156 | gt_im.data, 157 | ((0, 10), (0, 10), (0, 0)), 158 | constant_values=np.iinfo(np.uint8).max, 159 | ), 160 | np.pad( 161 | pred_im.data, 162 | ((0, 10), (0, 0), (0, 0)), 163 | constant_values=np.iinfo(np.uint8).max, 164 | ), 165 | ), 166 | axis=1, 167 | ) 168 | bot_row = np.concatenate( 169 | ( 170 | np.pad( 171 | gt_pts_im.data, 172 | ((0, 0), (0, 10), (0, 0)), 173 | constant_values=np.iinfo(np.uint8).max, 174 | ), 175 | np.pad( 176 | pred_pts_im.data, 177 | ((0, 0), (0, 0), (0, 0)), 178 | constant_values=np.iinfo(np.uint8).max, 179 | ), 180 | ), 181 | axis=1, 182 | ) 183 | full_gif_im = np.concatenate((top_row, bot_row), axis=0) 184 | writer.append_data(full_gif_im) 185 | scene_manager._renderer._renderer.delete() 186 | del scene_manager 187 | pygifsicle.optimize( 188 | osp.join(vis_path, "vis_{:d}.gif".format(batch_idx)) 189 | ) 190 | test_loss /= iterations 191 | 192 | # sort by probs 193 | accs = [] 194 | tprs = [] 195 | f1s = [] 196 | taus = np.linspace(0, 1, 21, endpoint=False) 197 | for t in taus: 198 | b = BinaryClassificationResult(preds, trues, threshold=t) 199 | accs.append(b.accuracy) 200 | f1s.append(b.f1_score) 201 | tprs.append(b.tpr) 202 | 203 | np.savez_compressed( 204 | osp.join(out_path, "accuracy_curve.npz"), 205 | taus=taus, 206 | accs=accs, 207 | f1s=f1s, 208 | tprs=tprs, 209 | ) 210 | 211 | ax = sns.lineplot(x=taus, y=accs, label="Accuracy") 212 | ax = sns.lineplot(x=taus, y=f1s, label="F1") 213 | ax = sns.lineplot(x=taus, y=tprs, label="TPR") 214 | ax.legend() 215 | plt.xlabel("Threshold") 216 | plt.ylabel("Score") 217 | plt.ylim([0.5, 1.0]) 218 | plt.xlim([-0.01, 1.01]) 219 | plt.savefig(osp.join(out_path, "accuracy_curve.png")) 220 | 221 | bcr = BinaryClassificationResult(preds, trues) 222 | with open(osp.join(out_path, "results.txt"), "w") as f: 223 | elapsed_time = timer() - time_start 224 | log = ( 225 | ["Images: {:d}".format(iterations)] 226 | + ["Queries: {:d}".format(iterations * len(coll.flatten()))] 227 | + ["Loss: {:.5f}".format(test_loss)] 228 | + ["Accuracy: {:.3f}".format(bcr.accuracy)] 229 | + ["F1 Score: {:.3f}".format(bcr.f1_score)] 230 | + ["AP Score: {:.3f}".format(bcr.ap_score)] 231 | + [ 232 | "FW Pass Time: {:.4f} +- {:.4f} s".format( 233 | np.mean(passes), np.std(passes) 234 | ) 235 | ] 236 | + ["Time: {:.2f} s".format(elapsed_time)] 237 | + [ 238 | "FCL Time: {:.3f} +- {:.3f} s".format( 239 | np.mean(coll_times), np.std(coll_times) 240 | ) 241 | ] 242 | ) 243 | log = map(str, log) 244 | f.write("\n".join(log)) 245 | 246 | 247 | if __name__ == "__main__": 248 | 249 | parser = argparse.ArgumentParser(description="Benchmark a 3D CollisionNet") 250 | parser.add_argument( 251 | "--cfg", 252 | type=str, 253 | default="cfg/benchmark_scenecollisionnet.yaml", 254 | help="config file with benchmarking params", 255 | ) 256 | args = parser.parse_args() 257 | 258 | config = YamlConfig(args.cfg) 259 | 260 | # Create output directory for model 261 | out = osp.join(config["output"], config["model"]["name"]) 262 | if osp.exists(out): 263 | response = keyboard_input( 264 | "A benchmark folder exists at {}. Overwrite?".format(out), 265 | yesno=True, 266 | ) 267 | if response.lower() == "n": 268 | os.exit() 269 | else: 270 | os.makedirs(out) 271 | config.save(osp.join(out, "benchmark.yaml")) 272 | 273 | # Check whether GPU available 274 | if not torch.cuda.is_available(): 275 | print("No CUDA available!") 276 | sys.exit(0) 277 | device = torch.device(config["device"]) 278 | 279 | # 1. Model 280 | model_path = osp.join(config["model"]["path"], config["model"]["name"]) 281 | train_cfg = YamlConfig(osp.join(model_path, "train.yaml")) 282 | model = SceneCollisionNet( 283 | bounds=train_cfg["model"]["bounds"], 284 | vox_size=train_cfg["model"]["vox_size"], 285 | ) 286 | checkpoint = torch.load(osp.join(model_path, "model.pth.tar")) 287 | model.load_state_dict(checkpoint["model_state_dict"], strict=False) 288 | model = model.to(device) 289 | 290 | # 2. Dataset 291 | kwargs = { 292 | "num_workers": 0 293 | if config["vis"] 294 | else ( 295 | config["num_workers"] 296 | if "num_workers" in config 297 | else os.cpu_count() 298 | ), 299 | "pin_memory": True, 300 | "worker_init_fn": lambda _: np.random.seed(), 301 | } 302 | for k, v in config["dataset"].items(): 303 | train_cfg["dataset"][k] = v 304 | test_set = BenchmarkSceneCollisionDataset( 305 | **train_cfg["dataset"], 306 | **train_cfg["camera"], 307 | bounds=train_cfg["model"]["bounds"], 308 | vis=config["vis"] 309 | ) 310 | test_loader = DataLoader(test_set, batch_size=None, **kwargs) 311 | 312 | benchmark( 313 | device=device, 314 | model=model, 315 | criterion=nn.BCEWithLogitsLoss(), 316 | test_loader=test_loader, 317 | iterations=config["iterations"], 318 | out_path=out, 319 | vis=config["vis"], 320 | ) 321 | -------------------------------------------------------------------------------- /scenecollisionnet/scene.py: -------------------------------------------------------------------------------- 1 | import os 2 | from collections.abc import Iterable 3 | 4 | import h5py 5 | import numpy as np 6 | import pyrender 7 | import trimesh 8 | import trimesh.transformations as tra 9 | from autolab_core import ColorImage, DepthImage 10 | 11 | 12 | class ObjectPlacementNotFound(Exception): 13 | pass 14 | 15 | 16 | class SceneManager: 17 | def __init__(self, dataset_folder, renderer=None): 18 | self._dataset_path = dataset_folder 19 | obj_info = h5py.File( 20 | os.path.join(self._dataset_path, "object_info.hdf5"), "r" 21 | ) 22 | self.mesh_info = obj_info["meshes"] 23 | self.categories = obj_info["categories"] 24 | 25 | self._collision_manager = trimesh.collision.CollisionManager() 26 | if renderer is not None and not isinstance(renderer, SceneRenderer): 27 | raise ValueError("renderer must be of type SceneRenderer") 28 | self._renderer = renderer 29 | 30 | self.objs = {} 31 | 32 | self._gravity_axis = 2 33 | self._table_dims = np.array([1.0, 1.6, 0.6]) 34 | self._table_pose = np.eye(4) 35 | 36 | @property 37 | def camera_pose(self): 38 | if self._renderer is None: 39 | raise ValueError("SceneManager does not contain a renderer!") 40 | return self._renderer.camera_pose 41 | 42 | @camera_pose.setter 43 | def camera_pose(self, cam_pose): 44 | if self._renderer is None: 45 | raise ValueError("SceneManager does not contain a renderer!") 46 | self._renderer.camera_pose = cam_pose 47 | 48 | @property 49 | def table_bounds(self): 50 | if not hasattr(self, "_table_bounds"): 51 | lbs = self._table_pose[:3, 3] - 0.5 * self._table_dims 52 | ubs = self._table_pose[:3, 3] + 0.5 * self._table_dims 53 | lbs[self._gravity_axis] = ubs[self._gravity_axis] 54 | ubs[self._gravity_axis] += 0.001 55 | lbs[self._gravity_axis] += 0.001 56 | self._table_bounds = (lbs, ubs) 57 | return self._table_bounds 58 | 59 | @table_bounds.setter 60 | def table_bounds(self, bounds): 61 | if not isinstance(bounds, np.ndarray): 62 | bounds = np.asarray(bounds) 63 | if bounds.shape != (2, 3): 64 | raise ValueError("Bounds is incorrect shape, should be (2, 3)") 65 | self._table_bounds = bounds 66 | 67 | def collides(self): 68 | return self._collision_manager.in_collision_internal() 69 | 70 | def min_distance(self, obj_manager): 71 | return self._collision_manager.min_distance_other(obj_manager) 72 | 73 | def add_object(self, name, mesh, info={}, pose=None, color=None): 74 | if name in self.objs: 75 | raise ValueError( 76 | "Duplicate name: object {} already exists".format(name) 77 | ) 78 | 79 | if pose is None: 80 | pose = np.eye(4, dtype=np.float32) 81 | 82 | color = ( 83 | np.asarray((0.7, 0.7, 0.7)) if color is None else np.asarray(color) 84 | ) 85 | mesh.visual.face_colors = np.tile( 86 | np.reshape(color, [1, 3]), [mesh.faces.shape[0], 1] 87 | ) 88 | self.objs[name] = {"mesh": mesh, "pose": pose} 89 | if "grasps" in info: 90 | self.objs[name]["grasps"] = info["grasps"][()] 91 | self._collision_manager.add_object( 92 | name, 93 | mesh, 94 | transform=pose, 95 | ) 96 | if self._renderer is not None: 97 | self._renderer.add_object(name, mesh, pose) 98 | 99 | return True 100 | 101 | def remove_object(self, name): 102 | if name not in self.objs: 103 | raise ValueError("object {} needs to be added first".format(name)) 104 | self._collision_manager.remove_object(name) 105 | if self._renderer is not None: 106 | self._renderer.remove_object(name) 107 | del self.objs[name] 108 | 109 | def sample_obj(self, cat=None, obj=None): 110 | if cat is None: 111 | cat = np.random.choice(list(self.categories.keys())) 112 | elif isinstance(cat, Iterable): 113 | cat = np.random.choice(cat) 114 | if obj is None: 115 | obj = np.random.choice(list(self.categories[cat])) 116 | try: 117 | mesh_path = os.path.join( 118 | self._dataset_path, self.mesh_info[obj]["path"].asstr()[()] 119 | ) 120 | mesh = trimesh.load(mesh_path, force="mesh") 121 | mesh.metadata["key"] = obj 122 | mesh.metadata["path"] = mesh_path 123 | info = self.mesh_info[obj] 124 | except (ValueError, TypeError): 125 | mesh = None 126 | info = None 127 | 128 | return mesh, info 129 | 130 | def place_obj(self, obj_id, mesh, info, max_attempts=10): 131 | self.add_object(obj_id, mesh, info=info) 132 | for _ in range(max_attempts): 133 | rand_stp = self._random_object_pose(info, *self.table_bounds) 134 | self.set_object_pose(obj_id, rand_stp) 135 | if not self.collides(): 136 | return True 137 | 138 | self.remove_object(obj_id) 139 | return False 140 | 141 | def sample_and_place_obj(self, obj_id, max_attempts=10): 142 | for _ in range(max_attempts): 143 | obj_mesh, obj_info = self.sample_obj(cat=["Mug", "Book"]) 144 | if not obj_mesh: 145 | continue 146 | if self.place_obj(obj_id, obj_mesh, obj_info): 147 | break 148 | else: 149 | continue 150 | 151 | def arrange_scene(self, num_objects, max_attempts=10): 152 | 153 | # Create and add table mesh 154 | table_mesh = trimesh.creation.box(self._table_dims) 155 | table_mesh.metadata["key"] = "table" 156 | self.add_object( 157 | name="table", 158 | mesh=table_mesh, 159 | pose=self._table_pose, 160 | ) 161 | 162 | # Sample and add objects 163 | for i in range(num_objects): 164 | obj_id = "obj_{:d}".format(i + 1) 165 | self.sample_and_place_obj(obj_id, max_attempts=max_attempts) 166 | 167 | def get_object_pose(self, name): 168 | if name not in self.objs: 169 | raise ValueError("object {} needs to be added first".format(name)) 170 | return self.objs[name]["pose"] 171 | 172 | def set_object_pose(self, name, pose): 173 | if name not in self.objs: 174 | raise ValueError("object {} needs to be added first".format(name)) 175 | self.objs[name]["pose"] = pose 176 | self._collision_manager.set_transform( 177 | name, 178 | pose, 179 | ) 180 | if self._renderer is not None: 181 | self._renderer.set_object_pose(name, pose) 182 | 183 | def render_points(self): 184 | if self._renderer is not None: 185 | return self._renderer.render_points() 186 | 187 | def reset(self): 188 | if self._renderer is not None: 189 | self._renderer.reset() 190 | 191 | for name in self.objs: 192 | self._collision_manager.remove_object(name) 193 | 194 | self.objs = {} 195 | 196 | def _random_object_pose(self, obj_info, lbs, ubs): 197 | stps, probs = obj_info["stps"][()], obj_info["probs"][()] 198 | pose = stps[np.random.choice(len(stps), p=probs)].copy() 199 | pose[:3, 3] += np.random.uniform(lbs, ubs) 200 | z_rot = tra.rotation_matrix( 201 | 2 * np.pi * np.random.rand(), [0, 0, 1], point=pose[:3, 3] 202 | ) 203 | return z_rot @ pose 204 | 205 | 206 | class SceneRenderer: 207 | def __init__(self): 208 | 209 | self._scene = pyrender.Scene() 210 | self._node_dict = {} 211 | self._camera_intr = None 212 | self._camera_node = None 213 | self._light_node = None 214 | self._renderer = None 215 | 216 | def create_camera(self, intr, znear, zfar): 217 | cam = pyrender.IntrinsicsCamera( 218 | intr.fx, intr.fy, intr.cx, intr.cy, znear, zfar 219 | ) 220 | self._camera_intr = intr 221 | self._camera_node = pyrender.Node(camera=cam, matrix=np.eye(4)) 222 | self._scene.add_node(self._camera_node) 223 | light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=4.0) 224 | self._light_node = pyrender.Node(light=light, matrix=np.eye(4)) 225 | self._scene.add_node(self._light_node) 226 | self._renderer = pyrender.OffscreenRenderer( 227 | viewport_width=intr.width, 228 | viewport_height=intr.height, 229 | point_size=5.0, 230 | ) 231 | 232 | @property 233 | def camera_pose(self): 234 | if self._camera_node is None: 235 | return None 236 | return self._camera_node.matrix 237 | 238 | @camera_pose.setter 239 | def camera_pose(self, cam_pose): 240 | if self._camera_node is None: 241 | raise ValueError("No camera in scene!") 242 | self._scene.set_pose(self._camera_node, cam_pose) 243 | self._scene.set_pose(self._light_node, cam_pose) 244 | 245 | def render_rgbd(self, depth_only=False): 246 | 247 | if depth_only: 248 | depth = self._renderer.render( 249 | self._scene, pyrender.RenderFlags.DEPTH_ONLY 250 | ) 251 | color = None 252 | depth = DepthImage(depth, frame="camera") 253 | else: 254 | color, depth = self._renderer.render(self._scene) 255 | color = ColorImage(color, frame="camera") 256 | depth = DepthImage(depth, frame="camera") 257 | 258 | return color, depth 259 | 260 | def render_segmentation(self, full_depth=None): 261 | if full_depth is None: 262 | _, full_depth = self.render_rgbd(depth_only=True) 263 | 264 | self.hide_objects() 265 | output = np.zeros(full_depth.data.shape, dtype=np.uint8) 266 | for i, obj_name in enumerate(self._node_dict): 267 | self._node_dict[obj_name].mesh.is_visible = True 268 | _, depth = self.render_rgbd(depth_only=True) 269 | mask = np.logical_and( 270 | (np.abs(depth.data - full_depth.data) < 1e-6), 271 | np.abs(full_depth.data) > 0, 272 | ) 273 | if np.any(output[mask] != 0): 274 | raise ValueError("wrong label") 275 | output[mask] = i + 1 276 | self._node_dict[obj_name].mesh.is_visible = False 277 | self.show_objects() 278 | 279 | return output, ["BACKGROUND"] + list(self._node_dict.keys()) 280 | 281 | def render_points(self): 282 | _, depth = self.render_rgbd(depth_only=True) 283 | point_norm_cloud = depth.point_normal_cloud(self._camera_intr) 284 | 285 | pts = point_norm_cloud.points.data.T.reshape( 286 | depth.height, depth.width, 3 287 | ) 288 | norms = point_norm_cloud.normals.data.T.reshape( 289 | depth.height, depth.width, 3 290 | ) 291 | cp = self.get_camera_pose() 292 | cp[:, 1:3] *= -1 293 | 294 | pt_mask = np.logical_and( 295 | np.linalg.norm(pts, axis=-1) != 0.0, 296 | np.linalg.norm(norms, axis=-1) != 0.0, 297 | ) 298 | pts = tra.transform_points(pts[pt_mask], cp) 299 | return pts.astype(np.float32) 300 | 301 | def add_object(self, name, mesh, pose=None): 302 | if pose is None: 303 | pose = np.eye(4, dtype=np.float32) 304 | 305 | node = pyrender.Node( 306 | name=name, 307 | mesh=pyrender.Mesh.from_trimesh(mesh, smooth=False), 308 | matrix=pose, 309 | ) 310 | self._node_dict[name] = node 311 | self._scene.add_node(node) 312 | 313 | def add_points(self, points, name, pose=None, color=None, radius=0.005): 314 | points = np.asanyarray(points) 315 | if points.ndim == 1: 316 | points = np.array([points]) 317 | 318 | if pose is None: 319 | pose = np.eye(4) 320 | else: 321 | pose = pose.matrix 322 | 323 | color = ( 324 | np.asanyarray(color, dtype=np.float) if color is not None else None 325 | ) 326 | 327 | # If color specified per point, use sprites 328 | if color is not None and color.ndim > 1: 329 | self._renderer.point_size = 1000 * radius 330 | m = pyrender.Mesh.from_points(points, colors=color) 331 | # otherwise, we can make pretty spheres 332 | else: 333 | mesh = trimesh.creation.uv_sphere(radius, [20, 20]) 334 | if color is not None: 335 | mesh.visual.vertex_colors = color 336 | poses = None 337 | poses = np.tile(np.eye(4), (len(points), 1)).reshape( 338 | len(points), 4, 4 339 | ) 340 | poses[:, :3, 3::4] = points[:, :, None] 341 | m = pyrender.Mesh.from_trimesh(mesh, poses=poses) 342 | 343 | node = pyrender.Node(mesh=m, name=name, matrix=pose) 344 | self._node_dict[name] = node 345 | self._scene.add_node(node) 346 | 347 | def set_object_pose(self, name, pose): 348 | self._scene.set_pose(self._node_dict[name], pose) 349 | 350 | def has_object(self, name): 351 | return name in self._node_dict 352 | 353 | def remove_object(self, name): 354 | self._scene.remove_node(self._node_dict[name]) 355 | del self._node_dict[name] 356 | 357 | def show_objects(self, names=None): 358 | for name, node in self._node_dict.items(): 359 | if names is None or name in names: 360 | node.mesh.is_visible = True 361 | 362 | def toggle_wireframe(self, names=None): 363 | for name, node in self._node_dict.items(): 364 | if names is None or name in names: 365 | node.mesh.primitives[0].material.wireframe ^= True 366 | 367 | def hide_objects(self, names=None): 368 | for name, node in self._node_dict.items(): 369 | if names is None or name in names: 370 | node.mesh.is_visible = False 371 | 372 | def reset(self): 373 | for name in self._node_dict: 374 | self._scene.remove_node(self._node_dict[name]) 375 | self._node_dict = {} 376 | -------------------------------------------------------------------------------- /scenecollisionnet/datasets.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | from timeit import default_timer as timer 3 | 4 | import numpy as np 5 | import trimesh.transformations as tra 6 | from autolab_core import CameraIntrinsics 7 | from torch.utils.data import IterableDataset 8 | from trimesh.collision import CollisionManager 9 | from urdfpy import URDF 10 | 11 | from .scene import SceneManager, SceneRenderer 12 | from .utils import compute_camera_pose 13 | 14 | 15 | class IterableSceneCollisionDataset(IterableDataset): 16 | def __init__( 17 | self, 18 | meshes, 19 | batch_size, 20 | query_size, 21 | intrinsics, 22 | extrinsics, 23 | bounds, 24 | n_obj_points, 25 | n_scene_points, 26 | rotations=True, 27 | trajectories=1, 28 | ): 29 | self.meshes = meshes 30 | self.batch_size = batch_size 31 | self.query_size = query_size 32 | self.cam_intr = CameraIntrinsics(**intrinsics) 33 | self.cam_pose = extrinsics 34 | self.bounds = np.array(bounds) 35 | 36 | self.n_obj_points = n_obj_points 37 | self.n_scene_points = n_scene_points 38 | self.rotations = rotations 39 | self.trajectories = trajectories 40 | 41 | # Generator that yields batches of training tuples 42 | def __iter__(self): 43 | """ 44 | Generator that yields batches of training tuples 45 | Outputs: 46 | scene_points (self.batch_size, self.n_scene_points, 3): scene point cloud batch 47 | obj_points (self.n_obj_points, 3): object point cloud batch 48 | trans (self.query_size, 3): translations of object in scene 49 | rots (self.query_size, 6): rotations of object in scene (first two cols of rotation matrix) 50 | colls (self.query_size,): boolean array of GT collisions along trajectories 51 | scene_manager (utils.SceneManager): underlying scene manager for GT collisions 52 | """ 53 | while True: 54 | ( 55 | obj_points, 56 | obj_centroid, 57 | obj_mesh, 58 | obj_pose, 59 | camera_pose, 60 | ) = self.get_obj() 61 | scene_points, scene_manager = self.get_scene( 62 | camera_pose=camera_pose 63 | ) 64 | np.random.seed() 65 | trans, rots, colls = self.get_colls( 66 | scene_manager, 67 | obj_mesh, 68 | obj_pose, 69 | obj_centroid, 70 | ) 71 | del scene_manager 72 | 73 | yield scene_points, obj_points, trans, rots, colls 74 | 75 | def get_scene(self, low=10, high=20, camera_pose=None): 76 | """ 77 | Generate a scene point cloud by placing meshes on a tabletop 78 | Inputs: 79 | low (int): minimum number of objects to place 80 | high (int): maximum number of objects to place 81 | camera_pose (4, 4): optional camera pose 82 | Outputs: 83 | points_batch (self.batch_size, self.n_scene_points, 3): scene point cloud batch 84 | scene_manager (utils.SceneManager): underlying scene manager for GT collisions 85 | """ 86 | # Create scene with random number of objects btw low and high 87 | num_objs = np.random.randint(low, high) 88 | scene_manager = self._create_scene() 89 | scene_manager.arrange_scene(num_objs) 90 | 91 | # Render points from batch_size different angles 92 | points_batch = np.zeros( 93 | (self.batch_size, self.n_scene_points, 3), dtype=np.float32 94 | ) 95 | for i in range(self.batch_size): 96 | if camera_pose is None: 97 | camera_pose = self.sample_camera_pose() 98 | scene_manager.camera_pose = camera_pose 99 | 100 | points = scene_manager.render_points() 101 | points = points[ 102 | (points[:, :3] > self.bounds[0] + 1e-4).all(axis=1) 103 | ] 104 | points = points[ 105 | (points[:, :3] < self.bounds[1] - 1e-4).all(axis=1) 106 | ] 107 | pt_inds = np.random.choice( 108 | points.shape[0], size=self.n_scene_points 109 | ) 110 | sample_points = points[pt_inds] 111 | points_batch[i] = sample_points 112 | 113 | return points_batch, scene_manager 114 | 115 | def get_obj(self): 116 | """ 117 | Generate an object point cloud 118 | Outputs: 119 | points_batch (self.n_obj_points, 3): object point cloud batch 120 | points_center (3,): centroid of object point cloud 121 | obj_mesh (trimesh.Trimesh): object underlying mesh 122 | pose (4, 4): GT object stable pose (w/ z rotation) 123 | camera_pose (4, 4): Pose matrix of the camera 124 | """ 125 | obj_scene = self._create_scene() 126 | camera_pose = self.sample_camera_pose() 127 | obj_scene.camera_pose = camera_pose 128 | points = np.array([]) 129 | while not points.any(): 130 | obj_scene.reset() 131 | obj_mesh, obj_info = obj_scene.sample_obj() 132 | if obj_mesh is None: 133 | continue 134 | stps = obj_info["stps"] 135 | probs = obj_info["probs"] 136 | pose = stps[np.random.choice(len(stps), p=probs)].copy() 137 | z_rot = tra.rotation_matrix( 138 | 2 * np.pi * np.random.rand(), [0, 0, 1], point=pose[:3, 3] 139 | ) 140 | pose = z_rot @ pose 141 | obj_scene.add_object("obj", obj_mesh, pose) 142 | points = obj_scene.render_points() 143 | 144 | pt_inds = np.random.choice( 145 | points.shape[0], size=self.n_obj_points, replace=True 146 | ) 147 | points_batch = np.repeat( 148 | points[None, pt_inds], self.batch_size, axis=0 149 | ) 150 | points_center = np.mean(points_batch[0, :, :3], axis=0) 151 | del obj_scene 152 | 153 | return points_batch, points_center, obj_mesh, pose, camera_pose 154 | 155 | def get_colls( 156 | self, 157 | scene_manager, 158 | obj, 159 | obj_pose, 160 | obj_centroid, 161 | ): 162 | """ 163 | Generate object/scene collision trajectories 164 | Inputs: 165 | scene_manager (utils.SceneManager): Underlying scene manager 166 | obj (trimesh.Trimesh): Underlying object mesh 167 | obj_pose (4, 4): Object GT pose matrix 168 | obj_centroid (3,): Centroid of object point cloud 169 | Outputs: 170 | trans (self.query_size, 3): translations of object in scene 171 | rots (self.query_size, 6): rotations of object in scene (first two cols of rotation matrix) 172 | colls (self.query_size,): boolean array of GT collisions along trajectories 173 | """ 174 | # Generate trajectory translations and clip to workspace bounds 175 | trans_start, trans_end = np.random.uniform( 176 | self.bounds[0], 177 | self.bounds[1], 178 | size=(2, self.trajectories, len(self.bounds[0])), 179 | ) 180 | trans = ( 181 | np.linspace( 182 | trans_start, trans_end, self.query_size // self.trajectories 183 | ) 184 | .transpose(1, 0, 2) 185 | .reshape(self.query_size, 3) 186 | ) 187 | trans[trans < self.bounds[0] + 1e-4] = ( 188 | self.bounds[0, np.where(trans < self.bounds[0] + 1e-4)[1]] + 1e-4 189 | ) 190 | trans[trans > self.bounds[1] - 1e-4] = ( 191 | self.bounds[1, np.where(trans > self.bounds[1] - 1e-4)[1]] - 1e-4 192 | ) 193 | 194 | # Sample trajectory rotations (if desired) and interpolate 195 | mesh_trans = trans - (obj_centroid - obj_pose[:3, 3]) 196 | mesh_tfs = np.repeat(np.eye(4)[None, ...], self.query_size, axis=0) 197 | mesh_tfs[:, :3, 3] = mesh_trans 198 | if self.rotations: 199 | rots = np.random.randn(2 * self.trajectories, 2, 3).astype( 200 | np.float32 201 | ) 202 | b1 = rots[:, 0] / np.linalg.norm( 203 | rots[:, 0], axis=-1, keepdims=True 204 | ) 205 | b2 = ( 206 | rots[:, 1] 207 | - np.einsum("ij,ij->i", b1, rots[:, 1])[:, None] * b1 208 | ) 209 | b2 /= np.linalg.norm(b2, axis=1, keepdims=True) 210 | b3 = np.cross(b1, b2) 211 | rot_mats = np.stack((b1, b2, b3), axis=-1) 212 | step = self.query_size // self.trajectories 213 | for i in range(self.trajectories): 214 | quats = [ 215 | tra.quaternion_from_matrix(rm) 216 | for rm in rot_mats[2 * i : 2 * (i + 1)] 217 | ] 218 | d = np.dot(*quats) 219 | if d < 0.0: 220 | d = -d 221 | np.negative(quats[1], quats[1]) 222 | ang = np.arccos(d) 223 | t = np.linspace(0, 1, step, endpoint=True) 224 | quats_slerp = quats[0][None, :] * np.sin((1.0 - t) * ang)[ 225 | :, None 226 | ] / np.sin(ang) + quats[1][None, :] * np.sin(t * ang)[ 227 | :, None 228 | ] / np.sin( 229 | ang 230 | ) 231 | mesh_tfs[ 232 | i * step : (i + 1) * step, :3, :3 233 | ] = tra.quaternion_matrix(quats_slerp)[:, :3, :3] 234 | rots = mesh_tfs[:, :3, :2].transpose(0, 2, 1) 235 | else: 236 | rots = np.repeat( 237 | np.eye(3, dtype=np.float32)[None, :, :2], 238 | self.query_size, 239 | axis=0, 240 | ).transpose(0, 2, 1) 241 | 242 | # Apply to object within scene and find GT collision status 243 | new_obj_poses = mesh_tfs @ obj_pose 244 | colls = np.zeros(self.query_size, dtype=np.bool) 245 | scene_manager._collision_manager.add_object("query_obj", obj) 246 | for i in range(self.query_size): 247 | scene_manager._collision_manager.set_transform( 248 | "query_obj", new_obj_poses[i] 249 | ) 250 | colls[i] = scene_manager.collides() 251 | scene_manager._collision_manager.remove_object("query_obj") 252 | return ( 253 | trans, 254 | rots.reshape(self.query_size, -1), 255 | np.repeat( 256 | colls.reshape(1, self.query_size), self.batch_size, axis=0 257 | ), 258 | ) 259 | 260 | # Creates scene manager and renderer 261 | def _create_scene(self): 262 | r = SceneRenderer() 263 | r.create_camera(self.cam_intr, znear=0.04, zfar=5) 264 | s = SceneManager(self.meshes, renderer=r) 265 | return s 266 | 267 | # Samples a camera pose that looks at the center of the scene 268 | def sample_camera_pose(self, mean=False): 269 | if mean: 270 | az = np.mean(self.cam_pose["azimuth"]) 271 | elev = np.mean(self.cam_pose["elevation"]) 272 | radius = np.mean(self.cam_pose["radius"]) 273 | else: 274 | az = np.random.uniform(*self.cam_pose["azimuth"]) 275 | elev = np.random.uniform(*self.cam_pose["elevation"]) 276 | radius = np.random.uniform(*self.cam_pose["radius"]) 277 | 278 | sample_pose, _ = compute_camera_pose(radius, az, elev) 279 | 280 | return sample_pose 281 | 282 | 283 | class BenchmarkSceneCollisionDataset(IterableSceneCollisionDataset): 284 | def __init__( 285 | self, 286 | meshes, 287 | batch_size, 288 | query_size, 289 | intrinsics, 290 | extrinsics, 291 | bounds, 292 | n_obj_points=1024, 293 | n_scene_points=4096, 294 | rotations=True, 295 | trajectories=0, 296 | vis=True, 297 | **kwargs 298 | ): 299 | super().__init__( 300 | meshes, 301 | batch_size, 302 | query_size, 303 | intrinsics, 304 | extrinsics, 305 | bounds, 306 | n_obj_points, 307 | n_scene_points, 308 | rotations=rotations, 309 | trajectories=trajectories, 310 | ) 311 | self.vis = vis 312 | 313 | # Generator that yields batches of training tuples 314 | def __iter__(self): 315 | while True: 316 | ( 317 | obj_points, 318 | obj_centroid, 319 | obj_mesh, 320 | obj_pose, 321 | camera_pose, 322 | ) = self.get_obj() 323 | scene_points, scene_manager = self.get_scene( 324 | camera_pose=camera_pose 325 | ) 326 | np.random.seed() 327 | coll_start = timer() 328 | trans, rots, colls = self.get_colls( 329 | scene_manager, 330 | obj_mesh, 331 | obj_pose, 332 | obj_centroid, 333 | ) 334 | coll_time = timer() - coll_start 335 | 336 | if not self.vis: 337 | del scene_manager 338 | 339 | if self.vis: 340 | yield scene_manager, obj_mesh, obj_pose, scene_points, obj_points, trans, rots, colls, coll_time 341 | else: 342 | yield scene_points, obj_points, trans, rots, colls, coll_time 343 | 344 | 345 | class IterableRobotCollisionDataset(IterableDataset): 346 | def __init__(self, robot_urdf, batch_size): 347 | self.robot = URDF.load(robot_urdf) 348 | self.batch_size = batch_size 349 | 350 | # Get link poses for a series of random configurations 351 | low_joint_limits, high_joint_limits = self.robot.joint_limit_cfgs 352 | self.low_joint_vals = np.fromiter( 353 | low_joint_limits.values(), dtype=float 354 | ) 355 | self.high_joint_vals = np.fromiter( 356 | high_joint_limits.values(), dtype=float 357 | ) 358 | 359 | meshes = self.robot.collision_trimesh_fk().keys() 360 | self.link_meshes = list(meshes) 361 | self.num_links = len(self.link_meshes) 362 | self.link_combos = list(itertools.combinations(range(len(meshes)), 2)) 363 | 364 | # Add the meshes to the collision managers 365 | self.collision_managers = [] 366 | for m in self.link_meshes: 367 | collision_manager = CollisionManager() 368 | collision_manager.add_object("link", m) 369 | self.collision_managers.append(collision_manager) 370 | 371 | # Preprocess - find some collision-free configs and set min distances 372 | rand_cfgs = ( 373 | np.random.rand(self.batch_size, len(self.low_joint_vals)) 374 | * (self.high_joint_vals - self.low_joint_vals) 375 | + self.low_joint_vals 376 | ) 377 | mesh_poses = self.robot.collision_trimesh_fk_batch(rand_cfgs) 378 | colls = np.zeros(self.batch_size, dtype=np.bool) 379 | for k in range(self.batch_size): 380 | colls[k] = self.check_pairwise_distances( 381 | mesh_poses, k, boolean=True 382 | ) 383 | 384 | self.original_dists = np.inf * np.ones( 385 | (self.num_links, self.num_links) 386 | ) 387 | for k, c in enumerate(colls): 388 | if not c: 389 | dists = self.check_pairwise_distances( 390 | mesh_poses, k, normalize=False 391 | ) 392 | self.original_dists = np.minimum(self.original_dists, dists) 393 | 394 | # Generator that yields batches of training tuples 395 | def __iter__(self): 396 | while True: 397 | rand_cfgs = ( 398 | np.random.rand(self.batch_size, len(self.low_joint_vals)) 399 | * (self.high_joint_vals - self.low_joint_vals) 400 | + self.low_joint_vals 401 | ) 402 | mesh_poses = self.robot.collision_trimesh_fk_batch(rand_cfgs) 403 | colls = np.zeros(self.batch_size) 404 | for k in range(self.batch_size): 405 | colls[k] = np.sum(self.check_pairwise_distances(mesh_poses, k)) 406 | 407 | yield rand_cfgs.reshape(self.batch_size, -1), colls 408 | 409 | def check_pairwise_distances( 410 | self, mesh_poses, ind, boolean=False, normalize=False 411 | ): 412 | if boolean: 413 | coll = False 414 | else: 415 | dists = np.zeros((self.num_links, self.num_links)) 416 | for _, (i, j) in enumerate(self.link_combos): 417 | if abs(i - j) < 2 or ((i, j) == (6, 8)) or ((i, j) == (8, 10)): 418 | continue 419 | i_tf = mesh_poses[self.link_meshes[i]][ind] 420 | self.collision_managers[i].set_transform("link", i_tf) 421 | j_tf = mesh_poses[self.link_meshes[j]][ind] 422 | self.collision_managers[j].set_transform("link", j_tf) 423 | if boolean: 424 | coll |= self.collision_managers[i].in_collision_other( 425 | self.collision_managers[j] 426 | ) 427 | if coll: 428 | return coll 429 | else: 430 | dists[i, j] = dists[j, i] = self.get_dist_pair( 431 | i, j, normalize=normalize 432 | ) 433 | if boolean: 434 | return coll 435 | return dists 436 | 437 | def get_dist_pair(self, i, j, normalize=False): 438 | min_dist = self.collision_managers[i].min_distance_other( 439 | self.collision_managers[j] 440 | ) 441 | if normalize: 442 | min_dist = np.e ** ( 443 | -2.0 * (min_dist / self.original_dists[i, j]) ** 2 444 | ) 445 | return min_dist 446 | -------------------------------------------------------------------------------- /scenecollisionnet/policy/urdf.py: -------------------------------------------------------------------------------- 1 | import os 2 | from collections import OrderedDict 3 | 4 | import numpy as np 5 | import torch 6 | from lxml import etree as ET 7 | from urdfpy import ( 8 | URDF, 9 | Collision, 10 | Inertial, 11 | Joint, 12 | Link, 13 | Material, 14 | Transmission, 15 | Visual, 16 | ) 17 | from urdfpy.utils import parse_origin 18 | 19 | 20 | def configure_origin(value, device=None): 21 | """Convert a value into a 4x4 transform matrix. 22 | Parameters 23 | ---------- 24 | value : None, (6,) float, or (4,4) float 25 | The value to turn into the matrix. 26 | If (6,), interpreted as xyzrpy coordinates. 27 | Returns 28 | ------- 29 | matrix : (4,4) float or None 30 | The created matrix. 31 | """ 32 | assert isinstance( 33 | value, torch.Tensor 34 | ), "Invalid type for origin, expect 4x4 torch tensor" 35 | assert value.shape == (4, 4) 36 | return value.float().to(device) 37 | 38 | 39 | class TorchVisual(Visual): 40 | def __init__( 41 | self, geometry, name=None, origin=None, material=None, device=None 42 | ): 43 | self.device = device 44 | super().__init__(geometry, name, origin, material) 45 | 46 | @property 47 | def origin(self): 48 | """(4,4) float : The pose of this element relative to the link frame.""" 49 | return self._origin 50 | 51 | @origin.setter 52 | def origin(self, value): 53 | self._origin = configure_origin(value, self.device) 54 | 55 | @classmethod 56 | def _from_xml(cls, node, path, device): 57 | kwargs = cls._parse(node, path) 58 | kwargs["origin"] = torch.tensor(parse_origin(node)) 59 | kwargs["device"] = device 60 | return TorchVisual(**kwargs) 61 | 62 | 63 | class TorchCollision(Collision): 64 | @classmethod 65 | def _from_xml(cls, node, path, device): 66 | kwargs = cls._parse(node, path) 67 | kwargs["origin"] = parse_origin(node) 68 | return TorchCollision(**kwargs) 69 | 70 | 71 | class TorchLink(Link): 72 | _ELEMENTS = { 73 | "inertial": (Inertial, False, False), 74 | "visuals": (TorchVisual, False, True), 75 | "collisions": (TorchCollision, False, True), 76 | } 77 | 78 | def __init__(self, name, inertial, visuals, collisions, device=None): 79 | self.device = device 80 | super().__init__(name, inertial, visuals, collisions) 81 | 82 | @classmethod 83 | def _parse_simple_elements(cls, node, path, device): 84 | """Parse all elements in the _ELEMENTS array from the children of 85 | this node. 86 | Parameters 87 | ---------- 88 | node : :class:`lxml.etree.Element` 89 | The node to parse children for. 90 | path : str 91 | The string path where the XML file is located (used for resolving 92 | the location of mesh or image files). 93 | Returns 94 | ------- 95 | kwargs : dict 96 | Map from element names to the :class:`URDFType` subclass (or list, 97 | if ``multiple`` was set) created for that element. 98 | """ 99 | kwargs = {} 100 | for a in cls._ELEMENTS: 101 | t, r, m = cls._ELEMENTS[a] 102 | if not m: 103 | v = node.find(t._TAG) 104 | if r or v is not None: 105 | v = t._from_xml(v, path) 106 | else: 107 | vs = node.findall(t._TAG) 108 | if len(vs) == 0 and r: 109 | raise ValueError( 110 | "Missing required subelement(s) of type {} when " 111 | "parsing an object of type {}".format( 112 | t.__name__, cls.__name__ 113 | ) 114 | ) 115 | v = [t._from_xml(n, path, device) for n in vs] 116 | kwargs[a] = v 117 | return kwargs 118 | 119 | @classmethod 120 | def _parse(cls, node, path, device): 121 | """Parse all elements and attributes in the _ELEMENTS and _ATTRIBS 122 | arrays for a node. 123 | Parameters 124 | ---------- 125 | node : :class:`lxml.etree.Element` 126 | The node to parse. 127 | path : str 128 | The string path where the XML file is located (used for resolving 129 | the location of mesh or image files). 130 | Returns 131 | ------- 132 | kwargs : dict 133 | Map from names to Python classes created from the attributes 134 | and elements in the class arrays. 135 | """ 136 | kwargs = cls._parse_simple_attribs(node) 137 | kwargs.update(cls._parse_simple_elements(node, path, device)) 138 | return kwargs 139 | 140 | @classmethod 141 | def _from_xml(cls, node, path, device): 142 | """Create an instance of this class from an XML node. 143 | Parameters 144 | ---------- 145 | node : :class:`lxml.etree.Element` 146 | The node to parse. 147 | path : str 148 | The string path where the XML file is located (used for resolving 149 | the location of mesh or image files). 150 | Returns 151 | ------- 152 | obj : :class:`URDFType` 153 | An instance of this class parsed from the node. 154 | """ 155 | return cls(**cls._parse(node, path, device)) 156 | 157 | 158 | class TorchJoint(Joint): 159 | def __init__( 160 | self, 161 | name, 162 | joint_type, 163 | parent, 164 | child, 165 | axis=None, 166 | origin=None, 167 | limit=None, 168 | dynamics=None, 169 | safety_controller=None, 170 | calibration=None, 171 | mimic=None, 172 | device=None, 173 | ): 174 | self.device = device 175 | super().__init__( 176 | name, 177 | joint_type, 178 | parent, 179 | child, 180 | axis, 181 | origin, 182 | limit, 183 | dynamics, 184 | safety_controller, 185 | calibration, 186 | mimic, 187 | ) 188 | 189 | @property 190 | def origin(self): 191 | """(4,4) float : The pose of this element relative to the link frame.""" 192 | return self._origin 193 | 194 | @origin.setter 195 | def origin(self, value): 196 | self._origin = configure_origin(value, device=self.device) 197 | 198 | @property 199 | def axis(self): 200 | """(3,) float : The joint axis in the joint frame.""" 201 | return self._axis 202 | 203 | @axis.setter 204 | def axis(self, value): 205 | if value is None: 206 | value = torch.tensor([1.0, 0.0, 0.0], device=self.device) 207 | elif isinstance(value, torch.Tensor): 208 | assert value.shape == ( 209 | 3, 210 | ), "Invalid shape for axis, should be (3,)" 211 | value = value.to(self.device) 212 | value = value / torch.norm(value) 213 | else: 214 | value = torch.as_tensor(value).to(self.device) 215 | if value.shape != (3,): 216 | raise ValueError("Invalid shape for axis, should be (3,)") 217 | value = value / torch.norm(value) 218 | self._axis = value.float() 219 | 220 | @classmethod 221 | def _from_xml(cls, node, path, device): 222 | kwargs = cls._parse(node, path) 223 | kwargs["joint_type"] = str(node.attrib["type"]) 224 | kwargs["parent"] = node.find("parent").attrib["link"] 225 | kwargs["child"] = node.find("child").attrib["link"] 226 | axis = node.find("axis") 227 | if axis is not None: 228 | axis = torch.tensor(np.fromstring(axis.attrib["xyz"], sep=" ")) 229 | kwargs["axis"] = axis 230 | kwargs["origin"] = torch.tensor(parse_origin(node)) 231 | kwargs["device"] = device 232 | 233 | return TorchJoint(**kwargs) 234 | 235 | def _rotation_matrices(self, angles, axis): 236 | """Compute rotation matrices from angle/axis representations. 237 | Parameters 238 | ---------- 239 | angles : (n,) float 240 | The angles. 241 | axis : (3,) float 242 | The axis. 243 | Returns 244 | ------- 245 | rots : (n,4,4) 246 | The rotation matrices 247 | """ 248 | axis = axis / torch.norm(axis) 249 | sina = torch.sin(angles) 250 | cosa = torch.cos(angles) 251 | M = torch.eye(4, device=self.device).repeat((len(angles), 1, 1)) 252 | M[:, 0, 0] = cosa 253 | M[:, 1, 1] = cosa 254 | M[:, 2, 2] = cosa 255 | M[:, :3, :3] += ( 256 | torch.ger(axis, axis).repeat((len(angles), 1, 1)) 257 | * (1.0 - cosa)[:, np.newaxis, np.newaxis] 258 | ) 259 | M[:, :3, :3] += ( 260 | torch.tensor( 261 | [ 262 | [0.0, -axis[2], axis[1]], 263 | [axis[2], 0.0, -axis[0]], 264 | [-axis[1], axis[0], 0.0], 265 | ], 266 | device=self.device, 267 | ).repeat((len(angles), 1, 1)) 268 | * sina[:, np.newaxis, np.newaxis] 269 | ) 270 | return M 271 | 272 | def get_child_poses(self, cfg, n_cfgs): 273 | """Computes the child pose relative to a parent pose for a given set of 274 | configuration values. 275 | Parameters 276 | ---------- 277 | cfg : (n,) float or None 278 | The configuration values for this joint. They are interpreted 279 | based on the joint type as follows: 280 | - ``fixed`` - not used. 281 | - ``prismatic`` - a translation along the axis in meters. 282 | - ``revolute`` - a rotation about the axis in radians. 283 | - ``continuous`` - a rotation about the axis in radians. 284 | - ``planar`` - Not implemented. 285 | - ``floating`` - Not implemented. 286 | If ``cfg`` is ``None``, then this just returns the joint pose. 287 | Returns 288 | ------- 289 | poses : (n,4,4) float 290 | The poses of the child relative to the parent. 291 | """ 292 | if cfg is None: 293 | return self.origin.repeat((n_cfgs, 1, 1)) 294 | elif self.joint_type == "fixed": 295 | return self.origin.repeat((n_cfgs, 1, 1)) 296 | elif self.joint_type in ["revolute", "continuous"]: 297 | if cfg is None: 298 | cfg = torch.zeros(n_cfgs) 299 | return torch.matmul( 300 | self.origin.repeat((n_cfgs, 1, 1)), 301 | self._rotation_matrices(cfg.float(), self.axis), 302 | ) 303 | elif self.joint_type == "prismatic": 304 | if cfg is None: 305 | cfg = torch.zeros(n_cfgs) 306 | translation = torch.eye(4, device=self.device).repeat( 307 | (n_cfgs, 1, 1) 308 | ) 309 | translation[:, :3, 3] = self.axis * cfg[:, np.newaxis] 310 | return torch.matmul(self.origin, translation) 311 | elif self.joint_type == "planar": 312 | raise NotImplementedError() 313 | elif self.joint_type == "floating": 314 | raise NotImplementedError() 315 | else: 316 | raise ValueError("Invalid configuration") 317 | 318 | 319 | class TorchURDF(URDF): 320 | 321 | _ELEMENTS = { 322 | "links": (TorchLink, True, True), 323 | "joints": (TorchJoint, False, True), 324 | "transmissions": (Transmission, False, True), 325 | "materials": (Material, False, True), 326 | } 327 | 328 | def __init__( 329 | self, 330 | name, 331 | links, 332 | joints=None, 333 | transmissions=None, 334 | materials=None, 335 | other_xml=None, 336 | device=None, 337 | ): 338 | self.device = device 339 | super().__init__( 340 | name, links, joints, transmissions, materials, other_xml 341 | ) 342 | 343 | @staticmethod 344 | def load(file_obj, device=None): 345 | """Load a URDF from a file. 346 | Parameters 347 | ---------- 348 | file_obj : str or file-like object 349 | The file to load the URDF from. Should be the path to the 350 | ``.urdf`` XML file. Any paths in the URDF should be specified 351 | as relative paths to the ``.urdf`` file instead of as ROS 352 | resources. 353 | Returns 354 | ------- 355 | urdf : :class:`.URDF` 356 | The parsed URDF. 357 | """ 358 | if isinstance(file_obj, str): 359 | if os.path.isfile(file_obj): 360 | parser = ET.XMLParser( 361 | remove_comments=True, remove_blank_text=True 362 | ) 363 | tree = ET.parse(file_obj, parser=parser) 364 | path, _ = os.path.split(file_obj) 365 | else: 366 | raise ValueError("{} is not a file".format(file_obj)) 367 | else: 368 | parser = ET.XMLParser(remove_comments=True, remove_blank_text=True) 369 | tree = ET.parse(file_obj, parser=parser) 370 | path, _ = os.path.split(file_obj.name) 371 | 372 | node = tree.getroot() 373 | return TorchURDF._from_xml(node, path, device) 374 | 375 | @classmethod 376 | def _parse_simple_elements(cls, node, path, device): 377 | """Parse all elements in the _ELEMENTS array from the children of 378 | this node. 379 | Parameters 380 | ---------- 381 | node : :class:`lxml.etree.Element` 382 | The node to parse children for. 383 | path : str 384 | The string path where the XML file is located (used for resolving 385 | the location of mesh or image files). 386 | Returns 387 | ------- 388 | kwargs : dict 389 | Map from element names to the :class:`URDFType` subclass (or list, 390 | if ``multiple`` was set) created for that element. 391 | """ 392 | kwargs = {} 393 | for a in cls._ELEMENTS: 394 | t, r, m = cls._ELEMENTS[a] 395 | if not m: 396 | v = node.find(t._TAG) 397 | if r or v is not None: 398 | v = t._from_xml(v, path) 399 | else: 400 | vs = node.findall(t._TAG) 401 | if len(vs) == 0 and r: 402 | raise ValueError( 403 | "Missing required subelement(s) of type {} when " 404 | "parsing an object of type {}".format( 405 | t.__name__, cls.__name__ 406 | ) 407 | ) 408 | v = [t._from_xml(n, path, device) for n in vs] 409 | kwargs[a] = v 410 | return kwargs 411 | 412 | @classmethod 413 | def _parse(cls, node, path, device): 414 | """Parse all elements and attributes in the _ELEMENTS and _ATTRIBS 415 | arrays for a node. 416 | Parameters 417 | ---------- 418 | node : :class:`lxml.etree.Element` 419 | The node to parse. 420 | path : str 421 | The string path where the XML file is located (used for resolving 422 | the location of mesh or image files). 423 | Returns 424 | ------- 425 | kwargs : dict 426 | Map from names to Python classes created from the attributes 427 | and elements in the class arrays. 428 | """ 429 | kwargs = cls._parse_simple_attribs(node) 430 | kwargs.update(cls._parse_simple_elements(node, path, device)) 431 | return kwargs 432 | 433 | @classmethod 434 | def _from_xml(cls, node, path, device): 435 | valid_tags = set(["joint", "link", "transmission", "material"]) 436 | kwargs = cls._parse(node, path, device) 437 | 438 | extra_xml_node = ET.Element("extra") 439 | for child in node: 440 | if child.tag not in valid_tags: 441 | extra_xml_node.append(child) 442 | 443 | data = ET.tostring(extra_xml_node) 444 | kwargs["other_xml"] = data 445 | kwargs["device"] = device 446 | return cls(**kwargs) 447 | 448 | def _process_cfgs(self, cfgs): 449 | """Process a list of joint configurations into a dictionary mapping joints to 450 | configuration values. 451 | This should result in a dict mapping each joint to a list of cfg values, one 452 | per joint. 453 | """ 454 | joint_cfg = {} 455 | assert isinstance( 456 | cfgs, torch.Tensor 457 | ), "Incorrectly formatted config array" 458 | n_cfgs = len(cfgs) 459 | for i, j in enumerate(self.actuated_joints): 460 | joint_cfg[j] = cfgs[:, i] 461 | 462 | return joint_cfg, n_cfgs 463 | 464 | def link_fk(self, cfg=None, link=None, links=None, use_names=False): 465 | raise NotImplementedError("Not implemented") 466 | 467 | def link_fk_batch(self, cfgs=None, use_names=False): 468 | """Computes the poses of the URDF's links via forward kinematics in a batch. 469 | Parameters 470 | ---------- 471 | cfgs : dict, list of dict, or (n,m), float 472 | One of the following: (A) a map from joints or joint names to vectors 473 | of joint configuration values, (B) a list of maps from joints or joint names 474 | to single configuration values, or (C) a list of ``n`` configuration vectors, 475 | each of which has a vector with an entry for each actuated joint. 476 | use_names : bool 477 | If True, the returned dictionary will have keys that are string 478 | link names rather than the links themselves. 479 | Returns 480 | ------- 481 | fk : dict or (n,4,4) float 482 | A map from links to a (n,4,4) vector of homogenous transform matrices that 483 | position the links relative to the base link's frame 484 | """ 485 | joint_cfgs, n_cfgs = self._process_cfgs(cfgs) 486 | 487 | # Process link set 488 | link_set = self.links 489 | 490 | # Compute FK mapping each link to a vector of matrices, one matrix per cfg 491 | fk = OrderedDict() 492 | for lnk in self._reverse_topo: 493 | if lnk not in link_set: 494 | continue 495 | poses = torch.eye(4, device=self.device).repeat((n_cfgs, 1, 1)) 496 | path = self._paths_to_base[lnk] 497 | for i in range(len(path) - 1): 498 | child = path[i] 499 | parent = path[i + 1] 500 | joint = self._G.get_edge_data(child, parent)["joint"] 501 | 502 | cfg_vals = None 503 | if joint.mimic is not None: 504 | mimic_joint = self._joint_map[joint.mimic.joint] 505 | if mimic_joint in joint_cfgs: 506 | cfg_vals = joint_cfgs[mimic_joint] 507 | cfg_vals = ( 508 | joint.mimic.multiplier * cfg_vals 509 | + joint.mimic.offset 510 | ) 511 | elif joint in joint_cfgs: 512 | cfg_vals = joint_cfgs[joint] 513 | poses = torch.matmul( 514 | joint.get_child_poses(cfg_vals, n_cfgs), poses 515 | ) 516 | 517 | if parent in fk: 518 | poses = torch.matmul(fk[parent], poses) 519 | break 520 | fk[lnk] = poses 521 | 522 | if use_names: 523 | return {ell.name: fk[ell] for ell in fk} 524 | return fk 525 | 526 | def visual_geometry_fk_batch(self, cfgs=None): 527 | """Computes the poses of the URDF's visual geometries using fk. 528 | Parameters 529 | ---------- 530 | cfgs : dict, list of dict, or (n,m), float 531 | One of the following: (A) a map from joints or joint names to vectors 532 | of joint configuration values, (B) a list of maps from joints or joint names 533 | to single configuration values, or (C) a list of ``n`` configuration vectors, 534 | each of which has a vector with an entry for each actuated joint. 535 | links : list of str or list of :class:`.Link` 536 | The links or names of links to perform forward kinematics on. 537 | Only geometries from these links will be in the returned map. 538 | If not specified, all links are returned. 539 | Returns 540 | ------- 541 | fk : dict 542 | A map from :class:`Geometry` objects that are part of the visual 543 | elements of the specified links to the 4x4 homogenous transform 544 | matrices that position them relative to the base link's frame. 545 | """ 546 | lfk = self.link_fk_batch(cfgs=cfgs) 547 | 548 | fk = OrderedDict() 549 | for link in lfk: 550 | for visual in link.visuals: 551 | fk[visual.geometry] = torch.matmul(lfk[link], visual.origin) 552 | return fk 553 | -------------------------------------------------------------------------------- /scenecollisionnet/policy/collision_checker.py: -------------------------------------------------------------------------------- 1 | import abc 2 | import itertools 3 | import os 4 | import os.path as osp 5 | 6 | import numpy as np 7 | import torch 8 | import torch_scatter 9 | import trimesh 10 | import trimesh.transformations as tra 11 | from knn_cuda import KNN 12 | 13 | try: 14 | import kaolin as kal 15 | except ImportError: 16 | print("skipping kaolin import, only needed for baselines") 17 | import multiprocessing as mp 18 | import queue as Queue 19 | 20 | import ruamel.yaml as yaml 21 | 22 | from ..collision_models import RobotCollisionNet, SceneCollisionNet 23 | 24 | 25 | class CollisionChecker(abc.ABC): 26 | def __init__(self, **kwargs): 27 | super().__init__(**kwargs) 28 | 29 | @abc.abstractmethod 30 | def __call__(self): 31 | pass 32 | 33 | 34 | class SceneCollisionChecker(CollisionChecker): 35 | def __init__(self, robot, use_knn=True, **kwargs): 36 | super().__init__(**kwargs) 37 | self.robot = robot 38 | self.device = self.robot.device 39 | if use_knn: 40 | self.knn = KNN(k=8, transpose_mode=True).to(self.device) 41 | else: 42 | self.knn = None 43 | self.cur_scene_pc = None 44 | self.robot_to_model = None 45 | self.model_to_robot = None 46 | 47 | def set_scene(self, obs): 48 | orig_scene_pc = obs["pc"] 49 | scene_labels = obs["pc_label"] 50 | label_map = obs["label_map"] 51 | 52 | # Remove robot points plus excluded 53 | self.scene_pc_mask = np.logical_and( 54 | scene_labels != label_map["robot"], 55 | scene_labels != label_map["target"], 56 | ) 57 | 58 | # Transform into robot frame (z up) 59 | self.camera_pose = obs["camera_pose"] 60 | self.scene_pc = tra.transform_points(orig_scene_pc, self.camera_pose) 61 | 62 | def set_object(self, obs): 63 | scene_pc = obs["pc"] 64 | scene_labels = obs["pc_label"] 65 | label_map = obs["label_map"] 66 | obj_pc = scene_pc[scene_labels == label_map["target"]] 67 | 68 | self.obj_pc = tra.transform_points( 69 | obj_pc, 70 | obs["camera_pose"], 71 | ) 72 | 73 | def _aggregate_pc(self, cur_pc, new_pc): 74 | 75 | # Filter tiny clusters of points and split into vis/occluding 76 | cam_model_tr = ( 77 | torch.from_numpy(self.camera_pose[:3, 3]).float().to(self.device) 78 | ) 79 | new_pc = torch.from_numpy(new_pc).float().to(self.device) 80 | if self.knn is not None: 81 | nearest = self.knn(new_pc[None, ...], new_pc[None, ...])[0][ 82 | 0, :, -1 83 | ] 84 | vis_mask = torch.from_numpy(self.scene_pc_mask).to(self.device) 85 | dists = torch.norm(new_pc - cam_model_tr, dim=1) ** 2 86 | dists /= dists.max() 87 | if self.knn is None: 88 | nearest = torch.zeros_like(vis_mask) 89 | occ_scene_pc = new_pc[~vis_mask & (nearest < 0.1 * dists)] 90 | scene_pc = new_pc[vis_mask & (nearest < 0.1 * dists)] 91 | 92 | if cur_pc is not None: 93 | if self.knn is None: 94 | pass 95 | else: 96 | dists = torch.norm(cur_pc - cam_model_tr, dim=1) ** 2 97 | dists /= dists.max() 98 | 99 | nearest = self.knn(cur_pc[None, ...], cur_pc[None, ...])[0][ 100 | 0, :, -1 101 | ] 102 | cur_pc = cur_pc[nearest < 0.1 * dists].float() 103 | 104 | # Group points by rays; get mapping from points to unique rays 105 | cur_pc_rays = cur_pc - cam_model_tr 106 | cur_pc_dists = torch.norm(cur_pc_rays, dim=1, keepdim=True) + 1e-12 107 | cur_pc_rays /= cur_pc_dists 108 | occ_pc_rays = occ_scene_pc - cam_model_tr 109 | occ_pc_dists = torch.norm(occ_pc_rays, dim=1, keepdim=True) + 1e-12 110 | occ_pc_rays /= occ_pc_dists 111 | occ_rays = ( 112 | (torch.cat((cur_pc_rays, occ_pc_rays), dim=0) * 50) 113 | .round() 114 | .long() 115 | ) 116 | _, occ_uniq_inv, occ_uniq_counts = torch.unique( 117 | occ_rays, dim=0, return_inverse=True, return_counts=True 118 | ) 119 | 120 | # Build new point cloud from previous now-occluded points and new pc 121 | cur_occ_inv = occ_uniq_inv[: len(cur_pc_rays)] 122 | cur_occ_counts = torch.bincount( 123 | cur_occ_inv, minlength=len(occ_uniq_counts) 124 | ) 125 | mean_occ_dists = torch_scatter.scatter_max( 126 | occ_pc_dists.squeeze(), 127 | occ_uniq_inv[-len(occ_pc_rays) :], 128 | dim_size=occ_uniq_inv.max() + 1, 129 | )[0] 130 | occ_mask = (occ_uniq_counts > cur_occ_counts) & ( 131 | cur_occ_counts > 0 132 | ) 133 | occ_pc = cur_pc[ 134 | occ_mask[cur_occ_inv] 135 | & (cur_pc_dists.squeeze() > mean_occ_dists[cur_occ_inv] + 0.01) 136 | ] 137 | return torch.cat((occ_pc, scene_pc), dim=0) 138 | else: 139 | return scene_pc 140 | 141 | def _compute_model_tfs(self, obs): 142 | if "robot_to_model" in obs: 143 | self.robot_to_model = obs["robot_to_model"] 144 | self.model_to_robot = obs["model_to_robot"] 145 | assert len(self.robot_to_model) > 0 146 | assert len(self.model_to_robot) > 0 147 | else: 148 | scene_labels = obs["pc_label"] 149 | label_map = obs["label_map"] 150 | 151 | # Extract table transform from points 152 | tab_pts = self.scene_pc[scene_labels == label_map["table"]] 153 | if len(tab_pts) == 0: 154 | tab_pts = self.scene_pc[scene_labels == label_map["objs"]] 155 | tab_height = tab_pts.mean(axis=0)[2] 156 | tab_tf_2d = trimesh.bounds.oriented_bounds_2D(tab_pts[:, :2])[0] 157 | tab_tf = np.eye(4) 158 | tab_tf[:3, :2] = tab_tf_2d[:, :2] 159 | tab_tf[:3, 3] = np.append(tab_tf_2d[:2, 2], 0.3 - tab_height) 160 | 161 | # Fix "long" side of table by rotating 162 | self.robot_to_model = tra.euler_matrix(0, 0, np.pi / 2) @ tab_tf 163 | if self.robot_to_model[0, 0] < 0: 164 | self.robot_to_model = ( 165 | tra.euler_matrix(0, 0, -np.pi) @ self.robot_to_model 166 | ) 167 | self.model_to_robot = np.linalg.inv(self.robot_to_model) 168 | 169 | self.robot_to_model = ( 170 | torch.from_numpy(self.robot_to_model).float().to(self.device) 171 | ) 172 | self.model_to_robot = ( 173 | torch.from_numpy(self.model_to_robot).float().to(self.device) 174 | ) 175 | 176 | 177 | class FCLSelfCollisionChecker(CollisionChecker): 178 | def __init__(self, robot): 179 | self.robot = robot 180 | self._link_combos = list( 181 | itertools.combinations(robot.link_map.keys(), 2) 182 | ) 183 | self._link_managers = {} 184 | for link in robot.links: 185 | if link.collision_mesh is not None: 186 | self._link_managers[ 187 | link.name 188 | ] = trimesh.collision.CollisionManager() 189 | self._link_managers[link.name].add_object( 190 | link.name, link.collision_mesh 191 | ) 192 | 193 | def set_allowed_collisions(self, l1, l2): 194 | if (l1, l2) in self._link_combos: 195 | self._link_combos.remove((l1, l2)) 196 | elif (l2, l1) in self._link_combos: 197 | self._link_combos.remove((l2, l1)) 198 | 199 | def __call__(self, q): 200 | self.robot.set_joint_cfg(q) 201 | for (i, j) in self._link_combos: 202 | if i not in self._link_managers or j not in self._link_managers: 203 | continue 204 | i_tf = ( 205 | self.robot.link_poses[self.robot.link_map[i]] 206 | .squeeze() 207 | .cpu() 208 | .numpy() 209 | ) 210 | self._link_managers[i].set_transform(i, i_tf) 211 | j_tf = ( 212 | self.robot.link_poses[self.robot.link_map[j]] 213 | .squeeze() 214 | .cpu() 215 | .numpy() 216 | ) 217 | self._link_managers[j].set_transform(j, j_tf) 218 | if self._link_managers[i].in_collision_other( 219 | self._link_managers[j] 220 | ): 221 | return True 222 | return False 223 | 224 | 225 | class FCLSceneCollisionChecker(SceneCollisionChecker): 226 | def __init__(self, robot, use_scene_pc=False): 227 | super().__init__(robot=robot) 228 | self._use_scene_pc = use_scene_pc 229 | self.robot_manager = trimesh.collision.CollisionManager() 230 | 231 | for link in robot.mesh_links: 232 | self.robot_manager.add_object(link.name, link.collision_mesh) 233 | 234 | def set_scene(self, obs, scene=None): 235 | if self._use_scene_pc: 236 | super().set_scene(obs) 237 | self.cur_scene_pc = self._aggregate_pc( 238 | self.cur_scene_pc, self.scene_pc 239 | ) 240 | pc = trimesh.PointCloud(self.scene_pc) 241 | scene_mesh = trimesh.voxel.ops.points_to_marching_cubes( 242 | pc.vertices, pitch=0.01 243 | ) 244 | self.scene_manager = trimesh.collision.CollisionManager() 245 | self.scene_manager.add_object("scene", scene_mesh) 246 | else: 247 | self.scene_manager = scene 248 | 249 | def set_object(self, obs, scene=None): 250 | if self._use_scene_pc: 251 | super().set_scene(obs) 252 | pc = trimesh.PointCloud(self.scene_pc) 253 | pitch = pc.extents.max() / 100 254 | scene_mesh = trimesh.voxel.ops.points_to_marching_cubes( 255 | pc.vertices, pitch=pitch 256 | ) 257 | self.scene_manager = trimesh.collision.CollisionManager() 258 | self.scene_manager.add_object("scene", scene_mesh) 259 | else: 260 | self.scene_manager = scene 261 | 262 | def __call__(self, q, by_link=False, threshold=None): 263 | if q is None or len(q) == 0: 264 | return np.zeros((len(self.robot.mesh_links), 0), dtype=np.bool) 265 | coll = ( 266 | np.zeros((len(self.robot.mesh_links), len(q)), dtype=np.bool) 267 | if by_link 268 | else np.zeros(len(q), dtype=np.bool) 269 | ) 270 | self.robot.set_joint_cfg(q) 271 | for i in range(len(q)): 272 | for link in self.robot.mesh_links: 273 | pose = self.robot.link_poses[link][i].squeeze().cpu().numpy() 274 | import pdb 275 | 276 | pdb.set_trace() 277 | self.robot_manager.set_transform(link.name, pose) 278 | 279 | coll_q = self.robot_manager.in_collision_other( 280 | self.scene_manager, return_names=by_link 281 | ) 282 | if by_link: 283 | for j, link in enumerate(self.robot.mesh_links): 284 | coll[j, i] = np.any( 285 | [link.name in pair for pair in coll_q[1]] 286 | ) 287 | else: 288 | coll[i] = coll_q 289 | return coll 290 | 291 | 292 | class FCLProc(mp.Process): 293 | """ 294 | Used for finding collisions in parallel using FCL. 295 | """ 296 | 297 | def __init__(self, links, output_queue, use_scene_pc=True): 298 | """ 299 | Args: 300 | output_queue: mp.Queue, the queue that all the output data 301 | that is computed is added to. 302 | """ 303 | super().__init__() 304 | self.links = links 305 | self.output_queue = output_queue 306 | self.input_queue = mp.Queue() 307 | self.use_scene_pc = use_scene_pc 308 | 309 | def _collides(self, link_poses, inds, by_link): 310 | """computes collisions.""" 311 | coll = ( 312 | np.zeros((len(self.links), len(inds)), dtype=np.bool) 313 | if by_link 314 | else np.zeros(len(inds), dtype=np.bool) 315 | ) 316 | for k, i in enumerate(inds): 317 | for link in self.links: 318 | pose = link_poses[link.name][i].squeeze() 319 | self.robot_manager.set_transform(link.name, pose) 320 | 321 | coll_q = self.robot_manager.in_collision_other( 322 | self.scene_manager, return_names=by_link 323 | ) 324 | if by_link: 325 | for j, link in enumerate(self.links): 326 | coll[j, k] = np.any( 327 | [link.name in pair for pair in coll_q[1]] 328 | ) 329 | else: 330 | coll[k] = coll_q 331 | return coll 332 | 333 | def _object_collides(self, poses, inds): 334 | """computes collisions.""" 335 | coll = np.zeros(len(inds), dtype=np.bool) 336 | for k, i in enumerate(inds): 337 | self.object_manager.set_transform("obj", poses[i]) 338 | coll[k] = self.object_manager.in_collision_other( 339 | self.scene_manager 340 | ) 341 | return coll 342 | 343 | def _set_scene(self, scene): 344 | if self.use_scene_pc: 345 | self.scene_manager = trimesh.collision.CollisionManager() 346 | self.scene_manager.add_object("scene", scene) 347 | else: 348 | self.scene_manager = scene 349 | 350 | def _set_object(self, obj): 351 | if self.use_scene_pc: 352 | self.object_manager = trimesh.collision.CollisionManager() 353 | self.object_manager.add_object("obj", obj) 354 | 355 | def run(self): 356 | """ 357 | the main function of each FCL process. 358 | """ 359 | self.robot_manager = trimesh.collision.CollisionManager() 360 | for link in self.links: 361 | self.robot_manager.add_object(link.name, link.collision_mesh) 362 | while True: 363 | try: 364 | request = self.input_queue.get(timeout=1) 365 | except Queue.Empty: 366 | continue 367 | if request[0] == "set_scene": 368 | self._set_scene(request[1]), 369 | elif request[0] == "set_object": 370 | self._set_object(request[1]), 371 | elif request[0] == "collides": 372 | self.output_queue.put( 373 | ( 374 | request[4], 375 | self._collides(*request[1:4]), 376 | ) 377 | ) 378 | elif request[0] == "obj_collides": 379 | self.output_queue.put( 380 | ( 381 | request[3], 382 | self._object_collides(*request[1:3]), 383 | ) 384 | ) 385 | 386 | def set_scene(self, scene): 387 | self.input_queue.put(("set_scene", scene)) 388 | 389 | def set_object(self, obj): 390 | self.input_queue.put(("set_object", obj)) 391 | 392 | def collides(self, link_poses, inds, by_link, pind=None): 393 | self.input_queue.put(("collides", link_poses, inds, by_link, pind)) 394 | 395 | def object_collides(self, poses, inds, pind=None): 396 | self.input_queue.put(("obj_collides", poses, inds, pind)) 397 | 398 | 399 | class FCLMultiSceneCollisionChecker(SceneCollisionChecker): 400 | def __init__(self, robot, n_proc=10, use_scene_pc=True): 401 | super().__init__(robot=robot) 402 | self._n_proc = n_proc 403 | self._use_scene_pc = use_scene_pc 404 | 405 | self.output_queue = mp.Queue() 406 | self.coll_procs = [] 407 | for _ in range(self._n_proc): 408 | self.coll_procs.append( 409 | FCLProc( 410 | self.robot.mesh_links, 411 | self.output_queue, 412 | use_scene_pc=self._use_scene_pc, 413 | ) 414 | ) 415 | self.coll_procs[-1].daemon = True 416 | self.coll_procs[-1].start() 417 | 418 | def set_scene(self, obs, scene=None): 419 | if self._use_scene_pc: 420 | super().set_scene(obs) 421 | self.cur_scene_pc = self._aggregate_pc( 422 | self.cur_scene_pc, self.scene_pc 423 | ) 424 | pc = trimesh.PointCloud(self.cur_scene_pc.cpu().numpy()) 425 | self.scene_mesh = trimesh.voxel.ops.points_to_marching_cubes( 426 | pc.vertices, pitch=0.01 427 | ) 428 | for proc in self.coll_procs: 429 | proc.set_scene(self.scene_mesh) 430 | else: 431 | for proc in self.coll_procs: 432 | proc.set_scene(scene) 433 | 434 | def set_object(self, obs): 435 | if self.robot_to_model is None: 436 | self._compute_model_tfs(obs) 437 | if self._use_scene_pc: 438 | super().set_object(obs) 439 | obj_pc = self.obj_pc - self.obj_pc.mean() 440 | pc = trimesh.PointCloud(obj_pc) 441 | self.obj_mesh = trimesh.voxel.ops.points_to_marching_cubes( 442 | pc.vertices, pitch=0.025 443 | ) 444 | self.obj_mesh.vertices -= self.obj_mesh.centroid 445 | for proc in self.coll_procs: 446 | proc.set_object(self.obj_mesh) 447 | 448 | def sample_in_bounds(self, num=20000, offset=0.0): 449 | return ( 450 | torch.rand((num, 3), dtype=torch.float32, device=self.device) 451 | * (torch.tensor([1.0, 1.6, 0.4], device=self.device) - 2 * offset) 452 | + torch.tensor([-0.5, -0.8, 0.2], device=self.device) 453 | + offset 454 | ) 455 | 456 | def check_object_collisions(self, translations, threshold=None): 457 | coll = torch.zeros( 458 | len(translations), dtype=torch.bool, device=self.device 459 | ) 460 | tr = tra.transform_points( 461 | translations.cpu().numpy(), self.model_to_robot.cpu().numpy() 462 | ) 463 | poses = np.repeat(np.eye(4)[None, ...], len(tr), axis=0) 464 | poses[:, :3, 3] = tr 465 | for i in range(self._n_proc): 466 | self.coll_procs[i].object_collides( 467 | poses, 468 | np.arange( 469 | i * len(tr) // self._n_proc, 470 | (i + 1) * len(tr) // self._n_proc, 471 | ), 472 | pind=i, 473 | ) 474 | 475 | # collect computed iks 476 | for _ in range(self._n_proc): 477 | i, proc_coll = self.output_queue.get(True) 478 | coll[ 479 | i * len(tr) // self._n_proc : (i + 1) * len(tr) // self._n_proc 480 | ] = torch.from_numpy(proc_coll).to(self.device) 481 | 482 | return coll 483 | 484 | def __call__(self, q, by_link=False, threshold=None): 485 | if q is None or len(q) == 0: 486 | return torch.zeros( 487 | (len(self.robot.mesh_links), 0), 488 | dtype=torch.bool, 489 | device=self.device, 490 | ) 491 | coll = ( 492 | torch.zeros( 493 | (len(self.robot.mesh_links), len(q)), 494 | dtype=torch.bool, 495 | device=self.device, 496 | ) 497 | if by_link 498 | else torch.zeros(len(q), dtype=np.bool, device=self.device) 499 | ) 500 | self.robot.set_joint_cfg(q) 501 | poses = { 502 | k.name: v.cpu().numpy() for k, v in self.robot.link_poses.items() 503 | } 504 | for i in range(self._n_proc): 505 | self.coll_procs[i].collides( 506 | poses, 507 | np.arange( 508 | i * len(q) // self._n_proc, 509 | (i + 1) * len(q) // self._n_proc, 510 | ), 511 | by_link, 512 | pind=i, 513 | ) 514 | 515 | # collect computed iks 516 | for _ in range(self._n_proc): 517 | i, proc_coll = self.output_queue.get(True) 518 | if by_link: 519 | coll[ 520 | :, 521 | i 522 | * len(q) 523 | // self._n_proc : (i + 1) 524 | * len(q) 525 | // self._n_proc, 526 | ] = torch.from_numpy(proc_coll).to(self.device) 527 | else: 528 | coll[ 529 | i 530 | * len(q) 531 | // self._n_proc : (i + 1) 532 | * len(q) 533 | // self._n_proc 534 | ] = torch.from_numpy(proc_coll).to(self.device) 535 | 536 | return coll 537 | 538 | 539 | class SDFSceneCollisionChecker(SceneCollisionChecker): 540 | def __init__(self, robot, n_pts=1000, device=torch.device("cuda:0")): 541 | self.robot = robot 542 | self.device = device 543 | 544 | self.link_pts = {} 545 | for link in self.robot.mesh_links: 546 | self.link_pts[link] = ( 547 | torch.from_numpy(link.collision_mesh.sample(n_pts)) 548 | .float() 549 | .to(self.device) 550 | ) 551 | 552 | self.scene_mesh = None 553 | 554 | def set_scene(self, obs): 555 | super().set_scene(obs) 556 | pc = trimesh.PointCloud(self.scene_pc) 557 | pitch = pc.extents.max() / 100 558 | scene_mesh = trimesh.voxel.ops.points_to_marching_cubes( 559 | pc.vertices, pitch=pitch 560 | ) 561 | verts = torch.from_numpy(scene_mesh.vertices).float().to(self.device) 562 | faces = torch.from_numpy(scene_mesh.faces).long().to(self.device) 563 | self.scene_sdf = kal.conversions.trianglemesh_to_sdf( 564 | kal.rep.TriangleMesh.from_tensors(verts, faces) 565 | ) 566 | 567 | def __call__(self, q, by_link=False, threshold=None): 568 | if q is None or len(q) == 0: 569 | return np.zeros((len(self.robot.mesh_links), 0), dtype=np.bool) 570 | self.robot.set_joint_cfg(q) 571 | coll = torch.zeros( 572 | (len(self.robot.mesh_links), len(q)), 573 | dtype=torch.bool, 574 | device=self.device, 575 | ) 576 | for i, link in enumerate(self.robot.mesh_links): 577 | poses = self.robot.link_poses[link] 578 | link_pts = torch.cat( 579 | ( 580 | self.link_pts[link], 581 | torch.ones( 582 | (len(self.link_pts[link]), 1), device=self.device 583 | ), 584 | ), 585 | dim=1, 586 | ) 587 | tf_link_pts = torch.matmul(poses, link_pts.T).transpose(2, 1)[ 588 | ..., :3 589 | ] 590 | sdf_vals = self.scene_sdf(tf_link_pts.reshape(-1, 3)).reshape( 591 | len(poses), -1 592 | ) 593 | coll[i] = (sdf_vals <= 0).any(dim=1) 594 | 595 | return coll if by_link else coll.any(dim=0) 596 | 597 | 598 | class SDFRobotCollisionChecker(SceneCollisionChecker): 599 | def __init__(self, robot, device=torch.device("cuda:0")): 600 | self.robot = robot 601 | self.device = device 602 | 603 | self.link_sdfs = {} 604 | for link in self.robot.mesh_links: 605 | link_mesh = link.collision_mesh 606 | verts = ( 607 | torch.from_numpy(link_mesh.vertices).float().to(self.device) 608 | ) 609 | faces = torch.from_numpy(link_mesh.faces).long().to(self.device) 610 | self.link_sdfs[link] = kal.conversions.trianglemesh_to_sdf( 611 | kal.rep.TriangleMesh.from_tensors(verts, faces) 612 | ) 613 | 614 | self.scene_pts = None 615 | 616 | def set_scene(self, obs): 617 | super().set_scene(obs) 618 | self.scene_pts = ( 619 | torch.from_numpy(self.scene_pc).float().to(self.device) 620 | ) 621 | 622 | def __call__(self, q, by_link=False, threshold=None): 623 | if q is None or len(q) == 0: 624 | return np.zeros((len(self.robot.mesh_links), 0), dtype=np.bool) 625 | self.robot.set_joint_cfg(q) 626 | coll = torch.zeros( 627 | (len(self.robot.mesh_links), len(q)), 628 | dtype=torch.bool, 629 | device=self.device, 630 | ) 631 | for i, link in enumerate(self.robot.mesh_links): 632 | poses = torch.inverse(self.robot.link_poses[link]) 633 | scene_pts = torch.cat( 634 | ( 635 | self.scene_pts, 636 | torch.ones((len(self.scene_pts), 1), device=self.device), 637 | ), 638 | dim=1, 639 | ) 640 | tf_scene_pts = torch.matmul(poses, scene_pts.T).transpose(2, 1)[ 641 | ..., :3 642 | ] 643 | sdf_vals = self.link_sdfs[link]( 644 | tf_scene_pts.reshape(-1, 3) 645 | ).reshape(len(poses), -1) 646 | coll[i] = (sdf_vals <= 0).any(dim=1) 647 | 648 | return coll if by_link else coll.any(dim=0) 649 | 650 | 651 | class NNCollisionChecker(CollisionChecker): 652 | def __init__(self, model_path, device=torch.device("cuda:0"), **kwargs): 653 | super().__init__(**kwargs) 654 | self.model_path = model_path 655 | self.device = device 656 | with open(osp.join(self.model_path, "train.yaml")) as f: 657 | self.cfg = yaml.safe_load(f) 658 | 659 | def _setup(self): 660 | chk = torch.load( 661 | os.path.join(self.model_path, "model.pth.tar"), 662 | map_location=self.device, 663 | ) 664 | # Bounds and vox_size parameters can be loaded earlier 665 | self.model.load_state_dict(chk["model_state_dict"], strict=False) 666 | self.model = self.model.to(self.device) 667 | self.model.eval() 668 | 669 | @abc.abstractmethod 670 | def __call__(self): 671 | pass 672 | 673 | 674 | class NNSelfCollisionChecker(NNCollisionChecker): 675 | def __init__(self, model_path, threshold=3.0, **kwargs): 676 | super().__init__(model_path, **kwargs) 677 | self.model = RobotCollisionNet(self.cfg["model"]["num_joints"]) 678 | self.threshold = threshold 679 | self._setup() 680 | 681 | def set_allowed_collisions(self, l1, l2): 682 | pass 683 | 684 | def __call__(self, qs): 685 | if len(qs) == 0: 686 | return torch.cuda.BoolTensor([], device=self.device) 687 | if isinstance(qs, np.ndarray): 688 | qs = torch.from_numpy(qs).float().to(self.device) 689 | with torch.no_grad(): 690 | out = self.model(qs[:, :8]) 691 | return out.squeeze(-1) > self.threshold 692 | 693 | 694 | class NNSceneCollisionChecker(NNCollisionChecker, SceneCollisionChecker): 695 | def __init__(self, model_path, robot, **kwargs): 696 | super().__init__(model_path=model_path, robot=robot, **kwargs) 697 | self.model = SceneCollisionNet( 698 | bounds=self.cfg["model"]["bounds"], 699 | vox_size=self.cfg["model"]["vox_size"], 700 | ) 701 | self._setup() 702 | 703 | # Get features for robot links 704 | mesh_links = self.robot.mesh_links[1:] 705 | n_pts = self.cfg["dataset"]["n_obj_points"] 706 | self.link_pts = np.zeros((len(mesh_links), n_pts, 3), dtype=np.float32) 707 | for i, link in enumerate(mesh_links): 708 | pts = link.collision_mesh.sample(n_pts) 709 | l_pts = pts[None, ...] 710 | self.link_pts[i] = l_pts 711 | with torch.no_grad(): 712 | self.link_features = self.model.get_obj_features( 713 | torch.from_numpy(self.link_pts).to(self.device) 714 | ) 715 | 716 | def set_scene(self, obs): 717 | super().set_scene(obs) 718 | 719 | if self.robot_to_model is None: 720 | self._compute_model_tfs(obs) 721 | 722 | if self.cur_scene_pc is not None: 723 | self.cur_scene_pc = self._aggregate_pc( 724 | self.cur_scene_pc, self.scene_pc 725 | ) 726 | else: 727 | self.cur_scene_pc = self._aggregate_pc(None, self.scene_pc) 728 | model_scene_pc = ( 729 | self.robot_to_model 730 | @ torch.cat( 731 | ( 732 | self.cur_scene_pc, 733 | torch.ones( 734 | (len(self.cur_scene_pc), 1), device=self.device 735 | ), 736 | ), 737 | dim=1, 738 | ).T 739 | ) 740 | model_scene_pc = model_scene_pc[:3].T 741 | 742 | if self.model.bounds[0].device != self.device: 743 | self.model.bounds = [b.to(self.device) for b in self.model.bounds] 744 | self.model.vox_size = self.model.vox_size.to(self.device) 745 | self.model.num_voxels = self.model.num_voxels.to(self.device) 746 | 747 | # Clip points to model bounds and feed in for features 748 | in_bounds = ( 749 | model_scene_pc[..., :3] > self.model.bounds[0] + 1e-5 750 | ).all(dim=-1) 751 | in_bounds &= ( 752 | model_scene_pc[..., :3] < self.model.bounds[1] - 1e-5 753 | ).all(dim=-1) 754 | self.scene_features = self.model.get_scene_features( 755 | model_scene_pc[in_bounds].unsqueeze(0) 756 | ).squeeze(0) 757 | 758 | def set_object(self, obs): 759 | super().set_object(obs) 760 | if self.robot_to_model is None: 761 | self._compute_model_tfs(obs) 762 | 763 | obj_pc = tra.transform_points( 764 | self.obj_pc, 765 | self.robot_to_model.cpu().numpy(), 766 | ) 767 | 768 | obj_tensor = torch.from_numpy(obj_pc.astype(np.float32)).to( 769 | self.device 770 | ) 771 | obj_tensor -= obj_tensor.mean(dim=0) 772 | self.obj_features = self.model.get_obj_features( 773 | obj_tensor.unsqueeze(0) 774 | ).squeeze(0) 775 | 776 | def sample_in_bounds(self, num=20000, offset=0.0): 777 | return ( 778 | torch.rand((num, 3), dtype=torch.float32, device=self.device) 779 | * (-torch.sub(*self.model.bounds) - 2 * offset) 780 | + self.model.bounds[0] 781 | + offset 782 | ) 783 | 784 | def check_object_collisions(self, translations, threshold=0.45): 785 | res = torch.zeros(len(translations), dtype=bool, device=self.device) 786 | in_bounds = (translations > self.model.bounds[0] + 1e-5).all(dim=-1) 787 | in_bounds &= (translations < self.model.bounds[1] - 1e-5).all(dim=-1) 788 | if in_bounds.any(): 789 | tr_in = translations[in_bounds] 790 | rots = np.repeat( 791 | np.eye(4)[:3, :2].flatten()[None, :], len(tr_in), axis=0 792 | ) 793 | with torch.no_grad(): 794 | out = self.model.classify_tfs( 795 | self.obj_features[None, :], 796 | self.scene_features[None, ...], 797 | tr_in, 798 | torch.from_numpy(rots).float().to(self.device), 799 | ) 800 | res[in_bounds] = torch.sigmoid(out).squeeze() > threshold 801 | return res 802 | 803 | # Objs is an (n, m) array of points, tfs is an (n, t, 4, 4) array of 804 | # object transforms 805 | def __call__(self, qs, by_link=False, thresholded=True, threshold=0.4): 806 | self.robot.set_joint_cfg(qs) 807 | colls = torch.zeros( 808 | (len(self.link_features), len(qs)), 809 | dtype=torch.bool if thresholded else torch.float32, 810 | device=self.device, 811 | ) 812 | trans = torch.empty( 813 | (len(self.link_features), len(qs), 3), 814 | dtype=torch.float32, 815 | device=self.device, 816 | ) 817 | rots = torch.empty( 818 | (len(self.link_features), len(qs), 6), 819 | dtype=torch.float32, 820 | device=self.device, 821 | ) 822 | 823 | for i, link in enumerate(self.robot.mesh_links[1:]): 824 | poses_tf = self.robot_to_model @ self.robot.link_poses[link] 825 | trans[i] = poses_tf[:, :3, 3] 826 | rots[i] = poses_tf[:, :3, :2].reshape(len(qs), -1) 827 | 828 | # filter translations that are out of bounds 829 | in_bounds = (trans > self.model.bounds[0] + 1e-5).all(dim=-1) 830 | in_bounds &= (trans < self.model.bounds[1] - 1e-5).all(dim=-1) 831 | if in_bounds.any(): 832 | trans[~in_bounds] = 0.0 # Make inputs safe 833 | with torch.no_grad(): 834 | out = self.model.classify_multi_obj_tfs( 835 | self.link_features, 836 | self.scene_features, 837 | trans, 838 | rots, 839 | ) 840 | res = torch.sigmoid(out).squeeze(-1) 841 | res = res > threshold if thresholded else res 842 | colls = res * in_bounds 843 | 844 | if thresholded: 845 | return colls if by_link else colls.any(dim=0) 846 | else: 847 | return colls if by_link else colls.max(dim=0)[0] 848 | --------------------------------------------------------------------------------