release repository

src/curobo/geom/__init__.py

@@ -0,0 +1,17 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+"""
+This module contains code for geometric processing such as pointcloud processing, analytic signed 
+distance computation for the environment, and also signed distance computation between robot and the
+environment. These functions can be used for robot self collision checking and also for robot 
+environment collision checking. 
+"""

src/curobo/geom/sdf/__init__.py

@@ -0,0 +1,18 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+"""
+This module contains code for geometric processing such as pointcloud processing, analytic signed 
+distance computation for the environment, and also signed distance computation between robot and 
+the environment. These functions can be used for robot self collision checking and also for robot 
+environment collision checking. The module also contains neural networks for learned signed distance
+fields.
+"""

src/curobo/geom/sdf/sdf_grid.py

@@ -0,0 +1,73 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+# Third Party
+import torch
+
+
+# @torch.jit.script
+def lookup_distance(pt, dist_matrix_flat, num_voxels):
+    # flatten:
+    ind_pt = (
+        (pt[..., 0]) * (num_voxels[1] * num_voxels[2]) + pt[..., 1] * num_voxels[2] + pt[..., 2]
+    )
+    dist = dist_matrix_flat[ind_pt]
+    return dist
+
+
+# @torch.jit.script
+def compute_sdf_gradient(pt, dist_matrix_flat, num_voxels, dist):
+    grad_l = []
+    for i in range(3):  # x,y,z
+        pt_n = pt.clone()
+        pt_p = pt.clone()
+        pt_n[..., i] -= 1
+        pt_p[..., i] += 1
+        # get distance from minus 1 and plus 1 idx:
+        pt_n[pt_n < 0] = 0
+        # pt_n[pt_n>nu]
+        # pt_p[pt_p > num_voxels] = num_voxels[i]
+        d_n = lookup_distance(pt_n, dist_matrix_flat, num_voxels)
+        d_p = lookup_distance(pt_p, dist_matrix_flat, num_voxels)
+        mask = d_n < d_p
+        dx = torch.where(mask, -1, 1)
+        ds = torch.where(mask, d_n - dist, d_p - dist)
+        g_d = ds / dx
+        grad_l.append(g_d)
+        # print(i, dist,  pt)
+    g_pt = torch.stack(grad_l, dim=-1)
+    # g_pt = g_pt/torch.linalg.norm(g_pt, dim=-1, keepdim=True)
+    return g_pt
+
+
+class SDFGrid(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, pt, dist_matrix_flat, num_voxels):
+        # input = x_id,y_id,z_id
+        pt = (pt).to(dtype=torch.int64)
+
+        dist = lookup_distance(pt, dist_matrix_flat, num_voxels)
+        ctx.save_for_backward(pt, dist_matrix_flat, num_voxels, dist)
+        return dist.unsqueeze(-1)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        pt, dist_matrix_flat, num_voxels, dist = ctx.saved_tensors
+        grad_pt = grad_voxels = grad_matrix_flat = None
+        if ctx.needs_input_grad[0]:
+            pt = pt.to(dtype=torch.int64)
+            g_pt = compute_sdf_gradient(pt, dist_matrix_flat, num_voxels, dist)
+            # print(g_pt)
+            grad_pt = grad_output * g_pt
+        if ctx.needs_input_grad[1]:
+            raise NotImplementedError("SDFGrid: Can't get gradient w.r.t. dist_matrix")
+        if ctx.needs_input_grad[2]:
+            raise NotImplementedError("SDFGrid: Can't get gradient w.r.t. num_voxels")
+        return grad_pt, grad_matrix_flat, grad_voxels

src/curobo/geom/sdf/utils.py

@@ -0,0 +1,35 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+# CuRobo
+from curobo.geom.sdf.world import CollisionCheckerType, WorldCollisionConfig
+from curobo.util.logger import log_error
+
+
+def create_collision_checker(config: WorldCollisionConfig):
+    if config.checker_type == CollisionCheckerType.PRIMITIVE:
+        # CuRobo
+        from curobo.geom.sdf.world import WorldPrimitiveCollision
+
+        return WorldPrimitiveCollision(config)
+    elif config.checker_type == CollisionCheckerType.BLOX:
+        # CuRobo
+        from curobo.geom.sdf.world_blox import WorldBloxCollision
+
+        return WorldBloxCollision(config)
+    elif config.checker_type == CollisionCheckerType.MESH:
+        # CuRobo
+        from curobo.geom.sdf.world_mesh import WorldMeshCollision
+
+        return WorldMeshCollision(config)
+    else:
+        log_error("Not implemented", exc_info=True)
+        raise NotImplementedError

src/curobo/geom/sdf/world.py

@@ -0,0 +1,900 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+# Standard Library
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Optional, Union
+
+# Third Party
+import torch
+
+# CuRobo
+from curobo.curobolib.geom import SdfSphereOBB, SdfSweptSphereOBB
+from curobo.geom.types import Cuboid, Mesh, Obstacle, WorldConfig, batch_tensor_cube
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util.logger import log_error, log_info, log_warn
+
+
+@dataclass
+class CollisionBuffer:
+    distance_buffer: torch.Tensor
+    grad_distance_buffer: torch.Tensor
+    sparsity_index_buffer: torch.Tensor
+    shape: Optional[torch.Size] = None
+
+    def __post_init__(self):
+        self.shape = self.distance_buffer.shape
+
+    @classmethod
+    def initialize_from_shape(cls, shape: torch.Size, tensor_args: TensorDeviceType):
+        batch, horizon, n_spheres, _ = shape
+        distance_buffer = torch.zeros(
+            (batch, horizon, n_spheres), device=tensor_args.device, dtype=tensor_args.dtype
+        )
+        grad_distance_buffer = torch.zeros(
+            (batch, horizon, n_spheres, 4), device=tensor_args.device, dtype=tensor_args.dtype
+        )
+        sparsity_idx = torch.zeros(
+            (batch, horizon, n_spheres),
+            device=tensor_args.device,
+            dtype=torch.uint8,
+        )
+        return CollisionBuffer(distance_buffer, grad_distance_buffer, sparsity_idx)
+
+    def _update_from_shape(self, shape: torch.Size, tensor_args: TensorDeviceType):
+        batch, horizon, n_spheres, _ = shape
+        self.distance_buffer = torch.zeros(
+            (batch, horizon, n_spheres), device=tensor_args.device, dtype=tensor_args.dtype
+        )
+        self.grad_distance_buffer = torch.zeros(
+            (batch, horizon, n_spheres, 4), device=tensor_args.device, dtype=tensor_args.dtype
+        )
+        self.sparsity_index_buffer = torch.zeros(
+            (batch, horizon, n_spheres),
+            device=tensor_args.device,
+            dtype=torch.uint8,
+        )
+        self.shape = shape[:3]
+
+    def update_buffer_shape(self, shape: torch.Size, tensor_args: TensorDeviceType):
+        if self.shape != shape[:3]:
+            # print("Recreating PRIM: ",self.shape, shape)
+
+            # self = CollisionBuffer.initialize_from_shape(
+            #    shape,
+            #    tensor_args,
+            # )
+            self._update_from_shape(shape, tensor_args)
+            # print("New shape:",self.shape)
+
+    def clone(self):
+        dist_buffer = self.distance_buffer.clone()
+        grad_buffer = self.grad_distance_buffer.clone()
+        sparse_buffer = self.sparsity_index_buffer.clone()
+        return CollisionBuffer(dist_buffer, grad_buffer, sparse_buffer)
+
+    def __mul__(self, scalar: float):
+        self.distance_buffer *= scalar
+        self.grad_distance_buffer *= scalar
+        self.sparsity_index_buffer *= int(scalar)
+        return self
+
+
+@dataclass
+class CollisionQueryBuffer:
+    """Class stores all buffers required to query collisions
+    This class currently has three main buffers. We initialize the required
+    buffers based on ?
+    """
+
+    primitive_collision_buffer: Optional[CollisionBuffer] = None
+    mesh_collision_buffer: Optional[CollisionBuffer] = None
+    blox_collision_buffer: Optional[CollisionBuffer] = None
+    shape: Optional[torch.Size] = None
+
+    def __post_init__(self):
+        if self.shape is None:
+            if self.primitive_collision_buffer is not None:
+                self.shape = self.primitive_collision_buffer.shape
+            elif self.mesh_collision_buffer is not None:
+                self.shape = self.mesh_collision_buffer.shape
+            elif self.blox_collision_buffer is not None:
+                self.shape = self.blox_collision_buffer.shape
+
+    def __mul__(self, scalar: float):
+        if self.primitive_collision_buffer is not None:
+            self.primitive_collision_buffer = self.primitive_collision_buffer * scalar
+        if self.mesh_collision_buffer is not None:
+            self.mesh_collision_buffer = self.mesh_collision_buffer * scalar
+        if self.blox_collision_buffer is not None:
+            self.blox_collision_buffer = self.blox_collision_buffer * scalar
+        return self
+
+    def clone(self):
+        prim_buffer = mesh_buffer = blox_buffer = None
+        if self.primitive_collision_buffer is not None:
+            prim_buffer = self.primitive_collision_buffer.clone()
+        if self.mesh_collision_buffer is not None:
+            mesh_buffer = self.mesh_collision_buffer.clone()
+        if self.blox_collision_buffer is not None:
+            blox_buffer = self.blox_collision_buffer.clone()
+        return CollisionQueryBuffer(prim_buffer, mesh_buffer, blox_buffer, self.shape)
+
+    @classmethod
+    def initialize_from_shape(
+        cls,
+        shape: torch.Size,
+        tensor_args: TensorDeviceType,
+        collision_types: Dict[str, bool],
+    ):
+        primitive_buffer = mesh_buffer = blox_buffer = None
+        if "primitive" in collision_types and collision_types["primitive"]:
+            primitive_buffer = CollisionBuffer.initialize_from_shape(shape, tensor_args)
+        if "mesh" in collision_types and collision_types["mesh"]:
+            mesh_buffer = CollisionBuffer.initialize_from_shape(shape, tensor_args)
+        if "blox" in collision_types and collision_types["blox"]:
+            blox_buffer = CollisionBuffer.initialize_from_shape(shape, tensor_args)
+        return CollisionQueryBuffer(primitive_buffer, mesh_buffer, blox_buffer)
+
+    def create_from_shape(
+        self,
+        shape: torch.Size,
+        tensor_args: TensorDeviceType,
+        collision_types: Dict[str, bool],
+    ):
+        if "primitive" in collision_types and collision_types["primitive"]:
+            self.primitive_collision_buffer = CollisionBuffer.initialize_from_shape(
+                shape, tensor_args
+            )
+        if "mesh" in collision_types and collision_types["mesh"]:
+            self.mesh_collision_buffer = CollisionBuffer.initialize_from_shape(shape, tensor_args)
+        if "blox" in collision_types and collision_types["blox"]:
+            self.blox_collision_buffer = CollisionBuffer.initialize_from_shape(shape, tensor_args)
+        self.shape = shape
+        # return self
+
+    def update_buffer_shape(
+        self,
+        shape: torch.Size,
+        tensor_args: TensorDeviceType,
+        collision_types: Optional[Dict[str, bool]],
+    ):
+        # print(shape, self.shape)
+        # update buffers:
+        assert len(shape) == 4  # shape is: batch, horizon, n_spheres, 4
+        if self.shape is None:  # buffers not initialized:
+            self.create_from_shape(shape, tensor_args, collision_types)
+            # print("Creating new memory", self.shape)
+        else:
+            # update buffers if shape doesn't match:
+            # TODO: allow for dynamic change of collision_types
+            if self.primitive_collision_buffer is not None:
+                # print(self.primitive_collision_buffer.shape, shape)
+                self.primitive_collision_buffer.update_buffer_shape(shape, tensor_args)
+            if self.mesh_collision_buffer is not None:
+                self.mesh_collision_buffer.update_buffer_shape(shape, tensor_args)
+            if self.blox_collision_buffer is not None:
+                self.blox_collision_buffer.update_buffer_shape(shape, tensor_args)
+            self.shape = shape
+
+    def get_gradient_buffer(
+        self,
+    ):
+        prim_buffer = mesh_buffer = blox_buffer = None
+        current_buffer = None
+        if self.primitive_collision_buffer is not None:
+            prim_buffer = self.primitive_collision_buffer.grad_distance_buffer
+            current_buffer = prim_buffer.clone()
+
+        if self.mesh_collision_buffer is not None:
+            mesh_buffer = self.mesh_collision_buffer.grad_distance_buffer
+            if current_buffer is None:
+                current_buffer = mesh_buffer.clone()
+            else:
+                current_buffer += mesh_buffer
+        if self.blox_collision_buffer is not None:
+            blox_buffer = self.blox_collision_buffer.grad_distance_buffer
+            if current_buffer is None:
+                current_buffer = blox_buffer.clone()
+            else:
+                current_buffer += blox_buffer
+
+        return current_buffer
+
+
+class CollisionCheckerType(Enum):
+    """Type of collision checker to use.
+    Args:
+        Enum (_type_): _description_
+    """
+
+    PRIMITIVE = "PRIMITIVE"
+    BLOX = "BLOX"
+    MESH = "MESH"
+
+
+@dataclass
+class WorldCollisionConfig:
+    tensor_args: TensorDeviceType
+    world_model: Optional[Union[List[WorldConfig], WorldConfig]] = None
+    cache: Optional[Dict[Obstacle, int]] = None
+    n_envs: int = 1
+    checker_type: CollisionCheckerType = CollisionCheckerType.PRIMITIVE
+    max_distance: float = 0.01
+
+    def __post_init__(self):
+        if self.world_model is not None and isinstance(self.world_model, list):
+            self.n_envs = len(self.world_model)
+
+    @staticmethod
+    def load_from_dict(
+        world_coll_checker_dict: Dict,
+        world_model_dict: Union[WorldConfig, Dict, List[WorldConfig]] = None,
+        tensor_args: TensorDeviceType = TensorDeviceType(),
+    ):
+        world_cfg = world_model_dict
+        if world_model_dict is not None:
+            if isinstance(world_model_dict, list) and isinstance(world_model_dict[0], dict):
+                world_cfg = [WorldConfig.from_dict(x) for x in world_model_dict]
+            elif isinstance(world_model_dict, dict):
+                world_cfg = WorldConfig.from_dict(world_model_dict)
+        world_coll_checker_dict["checker_type"] = CollisionCheckerType(
+            world_coll_checker_dict["checker_type"]
+        )
+        return WorldCollisionConfig(
+            tensor_args=tensor_args, world_model=world_cfg, **world_coll_checker_dict
+        )
+
+
+class WorldCollision(WorldCollisionConfig):
+    def __init__(self, config: Optional[WorldCollisionConfig] = None):
+        if config is not None:
+            WorldCollisionConfig.__init__(self, **vars(config))
+        self.collision_types = {}  # Use this dictionary to store collision types
+
+    def load_collision_model(self, world_model: WorldConfig):
+        raise NotImplementedError
+
+    def get_sphere_distance(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        """
+        Computes the signed distance via analytic function
+        Args:
+        tensor_sphere: b, n, 4
+        """
+        raise NotImplementedError
+
+    def get_sphere_collision(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        """
+        Computes the signed distance via analytic function
+        Args:
+        tensor_sphere: b, n, 4
+        we assume we don't need gradient for this function. If you need gradient, use get_sphere_distance
+        """
+
+        raise NotImplementedError
+
+    def get_swept_sphere_distance(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        sweep_steps: int,
+        enable_speed_metric=False,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        raise NotImplementedError
+
+    def get_swept_sphere_collision(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        sweep_steps: int,
+        enable_speed_metric=False,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        raise NotImplementedError
+
+    def get_sphere_trace(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        sweep_steps: int,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        raise NotImplementedError
+
+
+class WorldPrimitiveCollision(WorldCollision):
+    """World Oriented Bounding Box representation object
+
+    We represent the world with oriented bounding boxes. For speed, we assume there is a
+    maximum number of obbs that can be instantiated. This number is read from the WorldCollisionConfig.
+    If no cache is setup, we use the number from the first call of load_collision_model.
+    """
+
+    def __init__(self, config: WorldCollisionConfig):
+        super().__init__(config)
+        self._world_cubes = None
+        self._cube_tensor_list = None
+        self._env_n_obbs = None
+        self._env_obbs_names = None
+        self._init_cache()
+        if self.world_model is not None:
+            if isinstance(self.world_model, list):
+                self.load_batch_collision_model(self.world_model)
+            else:
+                self.load_collision_model(self.world_model)
+
+    def _init_cache(self):
+        if self.cache is not None and "obb" in self.cache and self.cache["obb"] not in [None, 0]:
+            self._create_obb_cache(self.cache["obb"])
+
+    def load_collision_model(self, world_config: WorldConfig, env_idx=0):
+        self._load_collision_model_in_cache(world_config, env_idx)
+
+    def load_batch_collision_model(self, world_config_list: List[WorldConfig]):
+        """Load a batch of collision environments from a list of world configs.
+
+        Args:
+            world_config_list: list of world configs to load from.
+        """
+        # First find largest number of cuboid:
+        c_len = []
+        pose_batch = []
+        dims_batch = []
+        names_batch = []
+        for i in world_config_list:
+            c = i.cuboid
+            if c is not None:
+                c_len.append(len(c))
+                pose_batch.extend([i.pose for i in c])
+                dims_batch.extend([i.dims for i in c])
+                names_batch.extend([i.name for i in c])
+            else:
+                c_len.append(0)
+
+        max_obb = max(c_len)
+        if max_obb < 1:
+            log_warn("No obbs found")
+            return
+        # check if number of environments is same as config:
+        reset_buffers = False
+        if self._env_n_obbs is not None and len(world_config_list) != len(self._env_n_obbs):
+            log_warn(
+                "env_n_obbs is not same as world_config_list, reloading collision buffers (breaks CG)"
+            )
+            reset_buffers = True
+            self.n_envs = len(world_config_list)
+            self._env_n_obbs = torch.zeros(
+                (self.n_envs), device=self.tensor_args.device, dtype=torch.int32
+            )
+
+        if self._cube_tensor_list is not None and self._cube_tensor_list[0].shape[1] < max_obb:
+            log_warn(
+                "number of obbs is greater than buffer, reloading collision buffers (breaks CG)"
+            )
+            reset_buffers = True
+        # create cache if does not exist:
+        if self._cube_tensor_list is None or reset_buffers:
+            log_info("Creating Obb cache" + str(max_obb))
+            self._create_obb_cache(max_obb)
+
+        # load obstacles:
+        ## load data into gpu:
+        cube_batch = batch_tensor_cube(pose_batch, dims_batch, self.tensor_args)
+        c_start = 0
+        for i in range(len(self._env_n_obbs)):
+            if c_len[i] > 0:
+                # load obb:
+                self._cube_tensor_list[0][i, : c_len[i], :3] = cube_batch[0][
+                    c_start : c_start + c_len[i]
+                ]
+                self._cube_tensor_list[1][i, : c_len[i], :7] = cube_batch[1][
+                    c_start : c_start + c_len[i]
+                ]
+                self._cube_tensor_list[2][i, : c_len[i]] = 1
+                self._env_obbs_names[i][: c_len[i]] = names_batch[c_start : c_start + c_len[i]]
+                self._cube_tensor_list[2][i, c_len[i] :] = 0
+                c_start += c_len[i]
+        self._env_n_obbs[:] = torch.as_tensor(
+            c_len, dtype=torch.int32, device=self.tensor_args.device
+        )
+        self.collision_types["primitive"] = True
+
+    def _load_collision_model_in_cache(self, world_config: WorldConfig, env_idx: int = 0):
+        cube_objs = world_config.cuboid
+        max_obb = len(cube_objs)
+        self.world_model = world_config
+        if max_obb < 1:
+            log_info("No OBB objs")
+            return
+        if self._cube_tensor_list is None or self._cube_tensor_list[0].shape[1] < max_obb:
+            log_info("Creating Obb cache" + str(max_obb))
+            self._create_obb_cache(max_obb)
+
+        # load as a batch:
+        pose_batch = [c.pose for c in cube_objs]
+        dims_batch = [c.dims for c in cube_objs]
+        names_batch = [c.name for c in cube_objs]
+        cube_batch = batch_tensor_cube(pose_batch, dims_batch, self.tensor_args)
+
+        self._cube_tensor_list[0][env_idx, :max_obb, :3] = cube_batch[0]
+        self._cube_tensor_list[1][env_idx, :max_obb, :7] = cube_batch[1]
+
+        self._cube_tensor_list[2][env_idx, :max_obb] = 1  # enabling obstacle
+
+        self._cube_tensor_list[2][env_idx, max_obb:] = 0  # disabling obstacle
+        # self._cube_tensor_list[1][env_idx, max_obb:, 0] = 1000.0  # Not needed. TODO: remove
+
+        self._env_n_obbs[env_idx] = max_obb
+        self._env_obbs_names[env_idx][:max_obb] = names_batch
+        self.collision_types["primitive"] = True
+
+    def _create_obb_cache(self, obb_cache):
+        box_dims = (
+            torch.zeros(
+                (self.n_envs, obb_cache, 4),
+                dtype=self.tensor_args.dtype,
+                device=self.tensor_args.device,
+            )
+            + 0.01
+        )
+        box_pose = torch.zeros(
+            (self.n_envs, obb_cache, 8),
+            dtype=self.tensor_args.dtype,
+            device=self.tensor_args.device,
+        )
+        box_pose[..., 3] = 1.0
+        obs_enable = torch.zeros(
+            (self.n_envs, obb_cache), dtype=torch.uint8, device=self.tensor_args.device
+        )
+        self._env_n_obbs = torch.zeros(
+            (self.n_envs), device=self.tensor_args.device, dtype=torch.int32
+        )
+        self._cube_tensor_list = [box_dims, box_pose, obs_enable]
+        self.collision_types["primitive"] = True
+        self._env_obbs_names = [[None for _ in range(obb_cache)] for _ in range(self.n_envs)]
+
+    def add_obb_from_raw(
+        self,
+        name: str,
+        dims: torch.Tensor,
+        env_idx: int,
+        w_obj_pose: Optional[Pose] = None,
+        obj_w_pose: Optional[Pose] = None,
+    ):
+        """
+        Args:
+        dims: lenght, width, height
+        position: x,y,z
+        rotation: matrix (3x3)
+        """
+        assert w_obj_pose is not None or obj_w_pose is not None
+        if name in self._env_obbs_names[env_idx]:
+            log_error("Obstacle already exists with name: " + name, exc_info=True)
+        if w_obj_pose is not None:
+            obj_w_pose = w_obj_pose.inverse()
+        # cube = tensor_cube(w_obj_pose, dims, tensor_args=self.tensor_args)
+
+        self._cube_tensor_list[0][env_idx, self._env_n_obbs[env_idx], :3] = dims
+        self._cube_tensor_list[1][
+            env_idx, self._env_n_obbs[env_idx], :7
+        ] = obj_w_pose.get_pose_vector()
+        self._cube_tensor_list[2][env_idx, self._env_n_obbs[env_idx]] = 1
+        self._env_obbs_names[env_idx][self._env_n_obbs[env_idx]] = name
+        self._env_n_obbs[env_idx] += 1
+        return self._env_n_obbs[env_idx] - 1
+
+    def add_obb(
+        self,
+        cuboid: Cuboid,
+        env_idx: int = 0,
+    ):
+        return self.add_obb_from_raw(
+            cuboid.name,
+            self.tensor_args.to_device(cuboid.dims),
+            env_idx,
+            Pose.from_list(cuboid.pose, self.tensor_args),
+        )
+
+    def update_obb_dims(
+        self,
+        obj_dims: torch.Tensor,
+        name: Optional[str] = None,
+        env_obj_idx: Optional[torch.Tensor] = None,
+        env_idx: int = 0,
+    ):
+        """Update obstacle dimensions
+
+        Args:
+            obj_dims (torch.Tensor): [dim.x,dim.y, dim.z], give as [b,3]
+            obj_idx (torch.Tensor or int):
+
+        """
+        if env_obj_idx is not None:
+            self._cube_tensor_list[0][env_obj_idx, :3] = obj_dims
+        else:
+            # find index of given name:
+            obs_idx = self.get_obb_idx(name, env_idx)
+
+            self._cube_tensor_list[0][env_idx, obs_idx, :3] = obj_dims
+
+    def enable_obstacle(
+        self,
+        name: str,
+        enable: bool = True,
+        env_idx: int = 0,
+    ):
+        return self.enable_obb(enable, name, None, env_idx)
+
+    def enable_obb(
+        self,
+        enable: bool = True,
+        name: Optional[str] = None,
+        env_obj_idx: Optional[torch.Tensor] = None,
+        env_idx: int = 0,
+    ):
+        """Update obstacle dimensions
+
+        Args:
+            obj_dims (torch.Tensor): [dim.x,dim.y, dim.z], give as [b,3]
+            obj_idx (torch.Tensor or int):
+
+        """
+        if env_obj_idx is not None:
+            self._cube_tensor_list[2][env_obj_idx] = int(enable)  # enable == 1
+        else:
+            # find index of given name:
+            obs_idx = self.get_obb_idx(name, env_idx)
+
+            self._cube_tensor_list[2][env_idx, obs_idx] = int(enable)
+
+    def update_obstacle_pose(
+        self,
+        name: str,
+        w_obj_pose: Pose,
+        env_idx: int = 0,
+    ):
+        if self._env_obbs_names is not None and name in self._env_obbs_names[env_idx]:
+            self.update_obb_pose(name=name, w_obj_pose=w_obj_pose, env_idx=env_idx)
+        else:
+            log_error("obstacle not found in OBB world model: " + name)
+
+    def update_obb_pose(
+        self,
+        w_obj_pose: Optional[Pose] = None,
+        obj_w_pose: Optional[Pose] = None,
+        name: Optional[str] = None,
+        env_obj_idx: Optional[torch.Tensor] = None,
+        env_idx: int = 0,
+    ):
+        """Update pose of a specific objects.
+        This also updates the signed distance grid to account for the updated object pose.
+        Args:
+        obj_w_pose: Pose
+        obj_idx:
+        """
+        obj_w_pose = self._get_obstacle_poses(w_obj_pose, obj_w_pose)
+        if env_obj_idx is not None:
+            self._cube_tensor_list[1][env_obj_idx, :7] = obj_w_pose.get_pose_vector()
+        else:
+            obs_idx = self.get_obb_idx(name, env_idx)
+            self._cube_tensor_list[1][env_idx, obs_idx, :7] = obj_w_pose.get_pose_vector()
+
+    @classmethod
+    def _get_obstacle_poses(
+        cls,
+        w_obj_pose: Optional[Pose] = None,
+        obj_w_pose: Optional[Pose] = None,
+    ):
+        if w_obj_pose is not None:
+            w_inv_pose = w_obj_pose.inverse()
+        elif obj_w_pose is not None:
+            w_inv_pose = obj_w_pose
+        else:
+            raise ValueError("Object pose is not given")
+        return w_inv_pose
+
+    def get_obb_idx(
+        self,
+        name: str,
+        env_idx: int = 0,
+    ) -> int:
+        if name not in self._env_obbs_names[env_idx]:
+            log_error("Obstacle with name: " + name + " not found in current world", exc_info=True)
+        return self._env_obbs_names[env_idx].index(name)
+
+    def get_sphere_distance(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss=False,
+    ):
+        if "primitive" not in self.collision_types or not self.collision_types["primitive"]:
+            raise ValueError("Primitive Collision has no obstacles")
+
+        b, h, n, _ = query_sphere.shape  # This can be read from collision query buffer
+        use_batch_env = True
+        if env_query_idx is None:
+            use_batch_env = False
+            env_query_idx = self._env_n_obbs
+
+        dist = SdfSphereOBB.apply(
+            query_sphere,
+            collision_query_buffer.primitive_collision_buffer.distance_buffer,
+            collision_query_buffer.primitive_collision_buffer.grad_distance_buffer,
+            collision_query_buffer.primitive_collision_buffer.sparsity_index_buffer,
+            weight,
+            activation_distance,
+            self._cube_tensor_list[0],
+            self._cube_tensor_list[0],
+            self._cube_tensor_list[1],
+            self._cube_tensor_list[2],
+            self._env_n_obbs,
+            env_query_idx,
+            self._cube_tensor_list[0].shape[1],
+            b,
+            h,
+            n,
+            query_sphere.requires_grad,
+            True,
+            use_batch_env,
+            return_loss,
+        )
+
+        return dist
+
+    def get_sphere_collision(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss=False,
+    ):
+        if "primitive" not in self.collision_types or not self.collision_types["primitive"]:
+            raise ValueError("Primitive Collision has no obstacles")
+        if return_loss:
+            raise ValueError("cannot return loss for classification, use get_sphere_distance")
+        b, h, n, _ = query_sphere.shape
+        use_batch_env = True
+        if env_query_idx is None:
+            use_batch_env = False
+            env_query_idx = self._env_n_obbs
+
+        dist = SdfSphereOBB.apply(
+            query_sphere,
+            collision_query_buffer.primitive_collision_buffer.distance_buffer,
+            collision_query_buffer.primitive_collision_buffer.grad_distance_buffer,
+            collision_query_buffer.primitive_collision_buffer.sparsity_index_buffer,
+            weight,
+            activation_distance,
+            self._cube_tensor_list[0],
+            self._cube_tensor_list[0],
+            self._cube_tensor_list[1],
+            self._cube_tensor_list[2],
+            self._env_n_obbs,
+            env_query_idx,
+            self._cube_tensor_list[0].shape[1],
+            b,
+            h,
+            n,
+            query_sphere.requires_grad,
+            False,
+            use_batch_env,
+            return_loss,
+        )
+        return dist
+
+    def get_swept_sphere_distance(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        sweep_steps: int,
+        enable_speed_metric=False,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss=False,
+    ):
+        """
+        Computes the signed distance via analytic function
+        Args:
+        tensor_sphere: b, n, 4
+        """
+        if "primitive" not in self.collision_types or not self.collision_types["primitive"]:
+            raise ValueError("Primitive Collision has no obstacles")
+
+        b, h, n, _ = query_sphere.shape
+        use_batch_env = True
+        if env_query_idx is None:
+            use_batch_env = False
+            env_query_idx = self._env_n_obbs
+
+        dist = SdfSweptSphereOBB.apply(
+            query_sphere,
+            collision_query_buffer.primitive_collision_buffer.distance_buffer,
+            collision_query_buffer.primitive_collision_buffer.grad_distance_buffer,
+            collision_query_buffer.primitive_collision_buffer.sparsity_index_buffer,
+            weight,
+            activation_distance,
+            speed_dt,
+            self._cube_tensor_list[0],
+            self._cube_tensor_list[0],
+            self._cube_tensor_list[1],
+            self._cube_tensor_list[2],
+            self._env_n_obbs,
+            env_query_idx,
+            self._cube_tensor_list[0].shape[1],
+            b,
+            h,
+            n,
+            sweep_steps,
+            enable_speed_metric,
+            query_sphere.requires_grad,
+            True,
+            use_batch_env,
+            return_loss,
+        )
+
+        return dist
+
+    def get_swept_sphere_collision(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        sweep_steps: int,
+        enable_speed_metric=False,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss=False,
+    ):
+        """
+        Computes the signed distance via analytic function
+        Args:
+        tensor_sphere: b, n, 4
+        """
+        if "primitive" not in self.collision_types or not self.collision_types["primitive"]:
+            raise ValueError("Primitive Collision has no obstacles")
+        if return_loss:
+            raise ValueError("cannot return loss for classify, use get_swept_sphere_distance")
+        b, h, n, _ = query_sphere.shape
+
+        use_batch_env = True
+        if env_query_idx is None:
+            use_batch_env = False
+            env_query_idx = self._env_n_obbs
+        dist = SdfSweptSphereOBB.apply(
+            query_sphere,
+            collision_query_buffer.primitive_collision_buffer.distance_buffer,
+            collision_query_buffer.primitive_collision_buffer.grad_distance_buffer,
+            collision_query_buffer.primitive_collision_buffer.sparsity_index_buffer,
+            weight,
+            activation_distance,
+            speed_dt,
+            self._cube_tensor_list[0],
+            self._cube_tensor_list[0],
+            self._cube_tensor_list[1],
+            self._cube_tensor_list[2],
+            self._env_n_obbs,
+            env_query_idx,
+            self._cube_tensor_list[0].shape[1],
+            b,
+            h,
+            n,
+            sweep_steps,
+            enable_speed_metric,
+            query_sphere.requires_grad,
+            False,
+            use_batch_env,
+            return_loss,
+        )
+
+        return dist
+
+    def clear_cache(self):
+        pass
+
+    def get_voxels_in_bounding_box(
+        self,
+        cuboid: Cuboid,
+        voxel_size: float = 0.02,
+    ) -> Union[List[Cuboid], torch.Tensor]:
+        bounds = cuboid.dims
+        low = [-bounds[0], -bounds[1], -bounds[2]]
+        high = [bounds[0], bounds[1], bounds[2]]
+        trange = [h - l for l, h in zip(low, high)]
+
+        x = torch.linspace(
+            -bounds[0], bounds[0], int(trange[0] // voxel_size), device=self.tensor_args.device
+        )
+        y = torch.linspace(
+            -bounds[1], bounds[1], int(trange[1] // voxel_size), device=self.tensor_args.device
+        )
+        z = torch.linspace(
+            -bounds[2], bounds[2], int(trange[2] // voxel_size), device=self.tensor_args.device
+        )
+        w, l, h = x.shape[0], y.shape[0], z.shape[0]
+        xyz = (
+            torch.stack(torch.meshgrid(x, y, z, indexing="ij")).permute((1, 2, 3, 0)).reshape(-1, 3)
+        )
+
+        pose = Pose.from_list(cuboid.pose, tensor_args=self.tensor_args)
+        xyz = pose.transform_points(xyz.contiguous())
+        r = torch.zeros_like(xyz[:, 0:1]) + voxel_size
+        xyzr = torch.cat([xyz, r], dim=1)
+        xyzr = xyzr.reshape(-1, 1, 1, 4)
+        collision_buffer = CollisionQueryBuffer()
+        collision_buffer.update_buffer_shape(
+            xyzr.shape,
+            self.tensor_args,
+            self.collision_types,
+        )
+        weight = self.tensor_args.to_device([1.0])
+        act_distance = self.tensor_args.to_device([0.0])
+
+        d_sph = self.get_sphere_collision(
+            xyzr,
+            collision_buffer,
+            weight,
+            act_distance,
+        )
+        d_sph = d_sph.reshape(-1)
+        xyzr = xyzr.reshape(-1, 4)
+        # get occupied voxels:
+        occupied = xyzr[d_sph > 0.0]
+        return occupied
+
+    def get_mesh_in_bounding_box(
+        self,
+        cuboid: Cuboid,
+        voxel_size: float = 0.02,
+    ) -> Mesh:
+        voxels = self.get_voxels_in_bounding_box(cuboid, voxel_size)
+        mesh = Mesh.from_pointcloud(
+            voxels[:, :3].detach().cpu().numpy(),
+            pitch=voxel_size,
+        )
+        return mesh

src/curobo/geom/sdf/world_blox.py

@@ -0,0 +1,478 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+# Standard Library
+from typing import List, Optional
+
+# Third Party
+import torch
+import torch.autograd.profiler as profiler
+
+# CuRobo
+from curobo.geom.sdf.world import CollisionQueryBuffer, WorldCollisionConfig
+from curobo.geom.sdf.world_mesh import WorldMeshCollision
+from curobo.geom.types import Cuboid, Mesh, Sphere, SphereFitType, WorldConfig
+from curobo.types.camera import CameraObservation
+from curobo.types.math import Pose
+from curobo.util.logger import log_error, log_info, log_warn
+
+try:
+    # Third Party
+    from nvblox_torch.mapper import Mapper
+except ImportError:
+    log_warn("nvblox torch wrapper is not installed, loading abstract class")
+    # Standard Library
+    from abc import ABC as Mapper
+
+
+class WorldBloxCollision(WorldMeshCollision):
+    """World Collision Representaiton using Nvidia's nvblox library.
+
+    This class depends on pytorch wrapper for nvblox.
+    Additionally, this representation does not support batched environments as we only store
+    one world via nvblox.
+
+    There are two ways to use nvblox, one is by loading maps from disk and the other is by
+    creating maps online. In both these instances, we might load more than one map and need to check
+    collisions against all maps.
+
+    To facilitate online map creation and updation, we build apis in this
+    class that provide.
+    """
+
+    def __init__(self, config: WorldCollisionConfig):
+        self._pose_offset = None
+        self._blox_mapper = None
+        self._blox_tensor_list = None
+        self._blox_voxel_sizes = [0.02]
+        super().__init__(config)
+
+    def load_collision_model(self, world_model: WorldConfig):
+        # load nvblox mesh
+        if len(world_model.blox) > 0:
+            # check if there is a mapper instance:
+            for k in world_model.blox:
+                if k.mapper_instance is not None:
+                    if self._blox_mapper is None:
+                        self._blox_mapper = k.mapper_instance
+                    else:
+                        log_error("Only one blox mapper instance is supported")
+            if self._blox_mapper is None:
+                voxel_sizes = []
+                integrator_types = []
+                # get voxel sizes and integration types:
+                for k in world_model.blox:
+                    voxel_sizes.append(k.voxel_size)
+                    integrator_types.append(k.integrator_type)
+                # create a mapper instance:
+                self._blox_mapper = Mapper(
+                    voxel_sizes=voxel_sizes,
+                    integrator_types=integrator_types,
+                    cuda_device_id=self.tensor_args.device.index,
+                    free_on_destruction=False,
+                )
+                self._blox_voxel_sizes = voxel_sizes
+            # load map from file if it exists:
+
+            names = []
+            pose = []
+
+            for i, k in enumerate(world_model.blox):
+                names.append(k.name)
+                pose.append(k.pose)
+                if k.map_path is not None:
+                    log_info("loading nvblox map for layer: " + str(i) + " from " + k.map_path)
+                    self._blox_mapper.load_from_file(k.map_path, i)
+
+            # load buffers:
+            pose_tensor = torch.zeros(
+                (len(names), 8), device=self.tensor_args.device, dtype=self.tensor_args.dtype
+            )
+            enable_tensor = torch.ones(
+                (len(names)), device=self.tensor_args.device, dtype=torch.uint8
+            )
+
+            pose_tensor[:, 3] = 1.0
+
+            pose_tensor[:, :7] = (
+                Pose.from_batch_list(pose, tensor_args=self.tensor_args).inverse().get_pose_vector()
+            )
+            self._blox_tensor_list = [pose_tensor, enable_tensor]
+            self._blox_names = names
+            self.collision_types["blox"] = True
+
+        return super().load_collision_model(world_model)
+
+    def clear_cache(self):
+        self._blox_mapper.clear()
+        super().clear_cache()
+
+    def clear_blox_layer(self, layer_name: str):
+        index = self._blox_names.index(layer_name)
+        self._blox_mapper.clear(index)
+
+    def _get_blox_sdf(
+        self,
+        query_spheres,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight,
+        activation_distance,
+    ):
+        d = self._blox_mapper.query_sphere_sdf_cost(
+            query_spheres,
+            collision_query_buffer.blox_collision_buffer.distance_buffer,
+            collision_query_buffer.blox_collision_buffer.grad_distance_buffer,
+            collision_query_buffer.blox_collision_buffer.sparsity_index_buffer,
+            weight,
+            activation_distance,
+            self._blox_tensor_list[0],
+            self._blox_tensor_list[1],
+        )
+        return d
+
+    def _get_blox_swept_sdf(
+        self,
+        query_spheres,
+        collision_query_buffer,
+        weight,
+        activation_distance,
+        speed_dt,
+        sweep_steps,
+        enable_speed_metric,
+    ):
+        d = self._blox_mapper.query_sphere_trajectory_sdf_cost(
+            query_spheres,
+            collision_query_buffer.blox_collision_buffer.distance_buffer,
+            collision_query_buffer.blox_collision_buffer.grad_distance_buffer,
+            collision_query_buffer.blox_collision_buffer.sparsity_index_buffer,
+            weight,
+            activation_distance,
+            speed_dt,
+            self._blox_tensor_list[0],
+            self._blox_tensor_list[1],
+            sweep_steps,
+            enable_speed_metric,
+        )
+        return d
+
+    def get_sphere_distance(
+        self,
+        query_sphere: torch.Tensor,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        if "blox" not in self.collision_types or not self.collision_types["blox"]:
+            return super().get_sphere_distance(
+                query_sphere,
+                collision_query_buffer,
+                weight,
+                activation_distance,
+                env_query_idx,
+                return_loss,
+            )
+
+        d = self._get_blox_sdf(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            activation_distance=activation_distance,
+        )
+
+        if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
+            "mesh" not in self.collision_types or not self.collision_types["mesh"]
+        ):
+            return d
+        d_base = super().get_sphere_distance(
+            query_sphere,
+            collision_query_buffer,
+            weight,
+            activation_distance,
+            env_query_idx,
+            return_loss,
+        )
+        d = d + d_base
+
+        return d
+
+    def get_sphere_collision(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx=None,
+        return_loss: bool = False,
+    ):
+        if "blox" not in self.collision_types or not self.collision_types["blox"]:
+            return super().get_sphere_collision(
+                query_sphere,
+                collision_query_buffer,
+                weight,
+                activation_distance,
+                env_query_idx,
+                return_loss,
+            )
+
+        d = self._get_blox_sdf(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            activation_distance=activation_distance,
+        )
+        if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
+            "mesh" not in self.collision_types or not self.collision_types["mesh"]
+        ):
+            return d
+        d_base = super().get_sphere_collision(
+            query_sphere,
+            collision_query_buffer,
+            weight,
+            activation_distance,
+            env_query_idx,
+            return_loss,
+        )
+        d = d + d_base
+
+        return d
+
+    def get_swept_sphere_distance(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        sweep_steps: int,
+        enable_speed_metric=False,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        if "blox" not in self.collision_types or not self.collision_types["blox"]:
+            return super().get_swept_sphere_distance(
+                query_sphere,
+                collision_query_buffer,
+                weight,
+                activation_distance,
+                speed_dt,
+                sweep_steps,
+                enable_speed_metric,
+                env_query_idx,
+                return_loss=return_loss,
+            )
+
+        d = self._get_blox_swept_sdf(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            activation_distance=activation_distance,
+            speed_dt=speed_dt,
+            sweep_steps=sweep_steps,
+            enable_speed_metric=enable_speed_metric,
+        )
+
+        if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
+            "mesh" not in self.collision_types or not self.collision_types["mesh"]
+        ):
+            return d
+        d_base = super().get_swept_sphere_distance(
+            query_sphere,
+            collision_query_buffer,
+            weight,
+            activation_distance,
+            speed_dt,
+            sweep_steps,
+            enable_speed_metric,
+            env_query_idx,
+            return_loss=return_loss,
+        )
+        d = d + d_base
+
+        return d
+
+    def get_swept_sphere_collision(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        sweep_steps,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        enable_speed_metric=False,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        if "blox" not in self.collision_types or not self.collision_types["blox"]:
+            return super().get_swept_sphere_collision(
+                query_sphere,
+                collision_query_buffer,
+                weight,
+                sweep_steps,
+                activation_distance,
+                speed_dt,
+                enable_speed_metric,
+                env_query_idx,
+                return_loss=return_loss,
+            )
+        d = self._get_blox_swept_sdf(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            activation_distance=activation_distance,
+            speed_dt=speed_dt,
+            sweep_steps=sweep_steps,
+            enable_speed_metric=enable_speed_metric,
+        )
+
+        if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
+            "mesh" not in self.collision_types or not self.collision_types["mesh"]
+        ):
+            return d
+        d_base = super().get_swept_sphere_collision(
+            query_sphere,
+            collision_query_buffer,
+            weight,
+            activation_distance,
+            speed_dt,
+            sweep_steps,
+            enable_speed_metric,
+            env_query_idx,
+            return_loss=return_loss,
+        )
+        d = d + d_base
+        return d
+
+    def enable_obstacle(
+        self,
+        name: str,
+        enable: bool = True,
+        env_idx: int = 0,
+    ):
+        if self._blox_names is not None and name in self._blox_names:
+            self.enable_blox(enable, name)
+        else:
+            super().enable_obstacle(name, enable, env_idx)
+
+    def enable_blox(self, enable: bool = True, name: Optional[str] = None):
+        index = self._blox_names.index(name)
+        self._blox_tensor_list[1][index] = int(enable)
+
+    def update_blox_pose(
+        self,
+        w_obj_pose: Optional[Pose] = None,
+        obj_w_pose: Optional[Pose] = None,
+        name: Optional[str] = None,
+    ):
+        obj_w_pose = self._get_obstacle_poses(w_obj_pose, obj_w_pose)
+        index = self._blox_names.index(name)
+        self._blox_tensor_list[0][index][:7] = obj_w_pose.get_pose_vector()
+
+    def clear_bounding_box(
+        self,
+        cuboid: Cuboid,
+        layer_name: Optional[str] = None,
+    ):
+        raise NotImplementedError
+
+    def get_bounding_spheres(
+        self,
+        bounding_box: Cuboid,
+        obstacle_name: Optional[str] = None,
+        n_spheres: int = 1,
+        surface_sphere_radius: float = 0.002,
+        fit_type: SphereFitType = SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
+        voxelize_method: str = "ray",
+        pre_transform_pose: Optional[Pose] = None,
+        clear_region: bool = False,
+    ) -> List[Sphere]:
+        mesh = self.get_mesh_in_bounding_box(bounding_box, obstacle_name, clear_region=clear_region)
+
+        spheres = mesh.get_bounding_spheres(
+            n_spheres=n_spheres,
+            surface_sphere_radius=surface_sphere_radius,
+            fit_type=fit_type,
+            voxelize_method=voxelize_method,
+            pre_transform_pose=pre_transform_pose,
+            tensor_args=self.tensor_args,
+        )
+        return spheres
+
+    @profiler.record_function("world_blox/add_camera_frame")
+    def add_camera_frame(self, camera_observation: CameraObservation, layer_name: str):
+        index = self._blox_names.index(layer_name)
+        pose_mat = camera_observation.pose.get_matrix().view(4, 4)
+        if camera_observation.rgb_image is not None:
+            with profiler.record_function("world_blox/add_color_frame"):
+                self._blox_mapper.add_color_frame(
+                    camera_observation.rgb_image,
+                    pose_mat,
+                    camera_observation.intrinsics,
+                    mapper_id=index,
+                )
+
+        if camera_observation.depth_image is not None:
+            with profiler.record_function("world_blox/add_depth_frame"):
+                self._blox_mapper.add_depth_frame(
+                    camera_observation.depth_image,
+                    pose_mat,
+                    camera_observation.intrinsics,
+                    mapper_id=index,
+                )
+
+    def process_camera_frames(self, layer_name: Optional[str] = None, process_aux: bool = False):
+        self.update_blox_esdf(layer_name)
+        if process_aux:
+            self.update_blox_mesh(layer_name)
+
+    @profiler.record_function("world_blox/update_hashes")
+    def update_blox_hashes(self):
+        self._blox_mapper.update_hashmaps()
+
+    @profiler.record_function("world_blox/update_esdf")
+    def update_blox_esdf(self, layer_name: Optional[str] = None):
+        index = -1
+        if layer_name is not None:
+            index = self._blox_names.index(layer_name)
+        self._blox_mapper.update_esdf(index)
+
+    @profiler.record_function("world_blox/update_mesh")
+    def update_blox_mesh(self, layer_name: Optional[str] = None):
+        index = -1
+        if layer_name is not None:
+            index = self._blox_names.index(layer_name)
+        self._blox_mapper.update_mesh(index)
+
+    @profiler.record_function("world_blox/get_mesh")
+    def get_mesh_from_blox_layer(self, layer_name: str, mode: str = "nvblox") -> Mesh:
+        index = self._blox_names.index(layer_name)
+        if mode == "nvblox":
+            mesh_data = self._blox_mapper.get_mesh(index)
+            mesh = Mesh(
+                name=self._blox_names[index],
+                pose=self._blox_tensor_list[0][index, :7].squeeze().cpu().tolist(),
+                vertices=mesh_data["vertices"].tolist(),
+                faces=mesh_data["triangles"].tolist(),
+                vertex_colors=mesh_data["colors"],
+                vertex_normals=mesh_data["normals"],
+            )
+        elif mode == "voxel":
+            # query points using collision checker and use trimesh to create a mesh:
+            mesh = self.get_mesh_in_bounding_box(
+                Cuboid("test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1.5, 1.5, 1]),
+                voxel_size=0.03,
+            )
+        return mesh
+
+    def decay_layer(self, layer_name: str):
+        index = self._blox_names.index(layer_name)
+        self._blox_mapper.decay_occupancy(mapper_id=index)

src/curobo/geom/sdf/world_mesh.py

@@ -0,0 +1,599 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+# Standard Library
+from dataclasses import dataclass
+from typing import List, Optional
+
+# Third Party
+import numpy as np
+import torch
+import warp as wp
+
+# CuRobo
+from curobo.geom.sdf.warp_primitives import SdfMeshWarpPy, SweptSdfMeshWarpPy
+from curobo.geom.sdf.world import (
+    CollisionQueryBuffer,
+    WorldCollisionConfig,
+    WorldPrimitiveCollision,
+)
+from curobo.geom.types import Mesh, WorldConfig
+from curobo.types.math import Pose
+from curobo.util.logger import log_error, log_info, log_warn
+from curobo.util.warp import init_warp
+
+
+@dataclass(frozen=True)
+class WarpMeshData:
+    name: str
+    m_id: int
+    vertices: wp.array
+    faces: wp.array
+    mesh: wp.Mesh
+
+
+class WorldMeshCollision(WorldPrimitiveCollision):
+    """World Mesh Collision using Nvidia's warp library
+
+    This currently requires passing int64 array from torch to warp which is only
+    available when compiled from source.
+    """
+
+    def __init__(self, config: WorldCollisionConfig):
+        # WorldCollision.(self)
+        init_warp()
+
+        self.tensor_args = config.tensor_args
+
+        self._env_n_mesh = None
+        self._mesh_tensor_list = None
+        self._env_mesh_names = None
+        self._wp_device = wp.torch.device_from_torch(self.tensor_args.device)
+        self._wp_mesh_cache = {}  # stores warp meshes across environments
+
+        super().__init__(config)
+
+    def _init_cache(self):
+        if (
+            self.cache is not None
+            and "mesh" in self.cache
+            and (not self.cache["mesh"] in [None, 0])
+        ):
+            self._create_mesh_cache(self.cache["mesh"])
+
+        return super()._init_cache()
+
+    def load_collision_model(
+        self, world_model: WorldConfig, env_idx: int = 0, load_obb_obs: bool = True
+    ):
+        max_nmesh = len(world_model.mesh)
+        if max_nmesh > 0:
+            if self._mesh_tensor_list is None or self._mesh_tensor_list[0].shape[1] < max_nmesh:
+                log_warn("Creating new Mesh cache: " + str(max_nmesh))
+                self._create_mesh_cache(max_nmesh)
+
+            # load all meshes as a batch:
+            name_list, w_mid, w_inv_pose = self._load_batch_mesh_to_warp(world_model.mesh)
+            self._mesh_tensor_list[0][env_idx, :max_nmesh] = w_mid
+            self._mesh_tensor_list[1][env_idx, :max_nmesh, :7] = w_inv_pose
+            self._mesh_tensor_list[2][env_idx, :max_nmesh] = 1
+
+            self._env_mesh_names[env_idx][:max_nmesh] = name_list
+            self._env_n_mesh[env_idx] = max_nmesh
+
+            self.collision_types["mesh"] = True
+        if load_obb_obs:
+            super().load_collision_model(world_model, env_idx)
+        else:
+            self.world_model = world_model
+
+    def load_batch_collision_model(self, world_config_list: List[WorldConfig]):
+        max_nmesh = max([len(x.mesh) for x in world_config_list])
+        if self._mesh_tensor_list is None or self._mesh_tensor_list[0].shape[1] < max_nmesh:
+            log_info("Creating new Mesh cache: " + str(max_nmesh))
+            self._create_mesh_cache(max_nmesh)
+
+        for env_idx, world_model in enumerate(world_config_list):
+            self.load_collision_model(world_model, env_idx=env_idx, load_obb_obs=False)
+        super().load_batch_collision_model(world_config_list)
+
+    def _load_mesh_to_warp(self, mesh: Mesh):
+        verts, faces = mesh.get_mesh_data()
+        v = wp.array(verts, dtype=wp.vec3, device=self._wp_device)
+        f = wp.array(np.ravel(faces), dtype=int, device=self._wp_device)
+        new_mesh = wp.Mesh(points=v, indices=f)
+        return WarpMeshData(mesh.name, new_mesh.id, v, f, new_mesh)
+
+    def _load_mesh_into_cache(self, mesh: Mesh) -> WarpMeshData:
+        #
+        if mesh.name not in self._wp_mesh_cache:
+            # load mesh into cache:
+            self._wp_mesh_cache[mesh.name] = self._load_mesh_to_warp(mesh)
+            # return self._wp_mesh_cache[mesh.name]
+        else:
+            log_info("Object already in warp cache, using existing instance: " + mesh.name)
+        return self._wp_mesh_cache[mesh.name]
+
+    def _load_batch_mesh_to_warp(self, mesh_list: List[Mesh]):
+        # First load all verts and faces:
+        name_list = []
+        pose_list = []
+        id_list = torch.zeros((len(mesh_list)), device=self.tensor_args.device, dtype=torch.int64)
+        for i, m_idx in enumerate(mesh_list):
+            m_data = self._load_mesh_into_cache(m_idx)
+            pose_list.append(m_idx.pose)
+
+            id_list[i] = m_data.m_id
+            name_list.append(m_data.name)
+        pose_buffer = Pose.from_batch_list(pose_list, self.tensor_args)
+        inv_pose_buffer = pose_buffer.inverse()
+        return name_list, id_list, inv_pose_buffer.get_pose_vector()
+
+    def add_mesh(self, new_mesh: Mesh, env_idx: int = 0):
+        if self._env_n_mesh[env_idx] >= self._mesh_tensor_list[0].shape[1]:
+            log_error(
+                "Cannot add new mesh as we are at mesh cache limit, increase cache limit in WorldMeshCollision"
+            )
+            return
+
+        wp_mesh_data = self._load_mesh_into_cache(new_mesh)
+
+        # get mesh pose:
+        w_obj_pose = Pose.from_list(new_mesh.pose, self.tensor_args)
+        # add loaded mesh into scene:
+
+        curr_idx = self._env_n_mesh[env_idx]
+        self._mesh_tensor_list[0][env_idx, curr_idx] = wp_mesh_data.m_id
+        self._mesh_tensor_list[1][env_idx, curr_idx, :7] = w_obj_pose.inverse().get_pose_vector()
+        self._mesh_tensor_list[2][env_idx, curr_idx] = 1
+        self._env_mesh_names[env_idx][curr_idx] = wp_mesh_data.name
+        self._env_n_mesh[env_idx] = curr_idx + 1
+
+    def get_mesh_idx(
+        self,
+        name: str,
+        env_idx: int = 0,
+    ) -> int:
+        if name not in self._env_mesh_names[env_idx]:
+            log_error("Obstacle with name: " + name + " not found in current world", exc_info=True)
+        return self._env_mesh_names[env_idx].index(name)
+
+    def create_collision_cache(self, mesh_cache=None, obb_cache=None, n_envs=None):
+        if n_envs is not None:
+            self.n_envs = n_envs
+        if mesh_cache is not None:
+            self._create_mesh_cache(mesh_cache)
+        if obb_cache is not None:
+            self._create_obb_cache(obb_cache)
+
+    def _create_mesh_cache(self, mesh_cache):
+        # create cache to store meshes, mesh poses and inverse poses
+
+        self._env_n_mesh = torch.zeros(
+            (self.n_envs), device=self.tensor_args.device, dtype=torch.int32
+        )
+
+        obs_enable = torch.zeros(
+            (self.n_envs, mesh_cache), dtype=torch.uint8, device=self.tensor_args.device
+        )
+        obs_inverse_pose = torch.zeros(
+            (self.n_envs, mesh_cache, 8),
+            dtype=self.tensor_args.dtype,
+            device=self.tensor_args.device,
+        )
+        obs_ids = torch.zeros(
+            (self.n_envs, mesh_cache), device=self.tensor_args.device, dtype=torch.int64
+        )
+        # v_empty = [[None for _ in range(mesh_cache)] for _ in range(self.n_envs)]
+        # @f_empty = [[None for _ in range(mesh_cache)] for _ in range(self.n_envs)]
+        # wp_m_empty = [[None for _ in range(mesh_cache)] for _ in range(self.n_envs)]
+        # warp requires uint64 for mesh indices, supports conversion from int64 to uint64
+        self._mesh_tensor_list = [
+            obs_ids,
+            obs_inverse_pose,
+            obs_enable,
+        ]  # 0=mesh idx, 1=pose, 2=mesh enable
+        self.collision_types["mesh"] = True  # TODO: enable this after loading first mesh
+        self._env_mesh_names = [[None for _ in range(mesh_cache)] for _ in range(self.n_envs)]
+
+        self._wp_mesh_cache = {}
+
+    def update_mesh_pose(
+        self,
+        w_obj_pose: Optional[Pose] = None,
+        obj_w_pose: Optional[Pose] = None,
+        name: Optional[str] = None,
+        env_obj_idx: Optional[torch.Tensor] = None,
+        env_idx: int = 0,
+    ):
+        w_inv_pose = self._get_obstacle_poses(w_obj_pose, obj_w_pose)
+
+        if name is not None:
+            obs_idx = self.get_mesh_idx(name, env_idx)
+            self._mesh_tensor_list[1][env_idx, obs_idx, :7] = w_inv_pose.get_pose_vector()
+        elif env_obj_idx is not None:
+            self._mesh_tensor_list[1][env_idx, env_obj_idx, :7] = w_inv_pose.get_pose_vector()
+        else:
+            raise ValueError("name or env_obj_idx needs to be given to update mesh pose")
+
+    def update_all_mesh_pose(
+        self,
+        w_obj_pose: Optional[Pose] = None,
+        obj_w_pose: Optional[Pose] = None,
+        name: Optional[List[str]] = None,
+        env_obj_idx: Optional[torch.Tensor] = None,
+        env_idx: int = 0,
+    ):
+        """Update poses for a list of meshes in the same environment
+
+        Args:
+            w_obj_pose (Optional[Pose], optional): _description_. Defaults to None.
+            obj_w_pose (Optional[Pose], optional): _description_. Defaults to None.
+            name (Optional[List[str]], optional): _description_. Defaults to None.
+            env_obj_idx (Optional[torch.Tensor], optional): _description_. Defaults to None.
+            env_idx (int, optional): _description_. Defaults to 0.
+        """
+        w_inv_pose = self._get_obstacle_poses(w_obj_pose, obj_w_pose)
+        raise NotImplementedError
+
+    def update_mesh_pose_env(
+        self,
+        w_obj_pose: Optional[Pose] = None,
+        obj_w_pose: Optional[Pose] = None,
+        name: Optional[str] = None,
+        env_obj_idx: Optional[torch.Tensor] = None,
+        env_idx: List[int] = [0],
+    ):
+        """Update pose of a single object in a list of environments
+
+        Args:
+            w_obj_pose (Optional[Pose], optional): _description_. Defaults to None.
+            obj_w_pose (Optional[Pose], optional): _description_. Defaults to None.
+            name (Optional[List[str]], optional): _description_. Defaults to None.
+            env_obj_idx (Optional[torch.Tensor], optional): _description_. Defaults to None.
+            env_idx (List[int], optional): _description_. Defaults to [0].
+        """
+        w_inv_pose = self._get_obstacle_poses(w_obj_pose, obj_w_pose)
+        # collect index of mesh across environments:
+        # index_tensor = torch.zeros((1, len(env_idx)), dtype=torch.long, device=self.tensor_args.device)
+
+        # for i, e in enumerate[env_idx]:
+        #    index_tensor[0,i] = self.get_mesh_idx(name, e)
+        raise NotImplementedError
+        # self._mesh_tensor_list[1][env_idx, obj_idx]
+
+    def update_mesh_from_warp(
+        self,
+        warp_mesh_idx: int,
+        w_obj_pose: Optional[Pose] = None,
+        obj_w_pose: Optional[Pose] = None,
+        obj_idx: int = 0,
+        env_idx: int = 0,
+        name: Optional[str] = None,
+    ):
+        if name is not None:
+            obj_idx = self.get_mesh_idx(name, env_idx)
+
+        if obj_idx >= self._mesh_tensor_list[0][env_idx].shape[0]:
+            log_error("Out of cache memory")
+        w_inv_pose = self._get_obstacle_poses(w_obj_pose, obj_w_pose)
+
+        self._mesh_tensor_list[0][env_idx, obj_idx] = warp_mesh_idx
+        self._mesh_tensor_list[1][env_idx, obj_idx] = w_inv_pose
+        self._mesh_tensor_list[2][env_idx, obj_idx] = 1
+        self._env_mesh_names[env_idx][obj_idx] = name
+        if self._env_n_mesh[env_idx] <= obj_idx:
+            self._env_n_mesh[env_idx] = obj_idx + 1
+
+    def update_obstacle_pose(
+        self,
+        name: str,
+        w_obj_pose: Pose,
+        env_idx: int = 0,
+    ):
+        if self._env_mesh_names is not None and name in self._env_mesh_names[env_idx]:
+            self.update_mesh_pose(name=name, w_obj_pose=w_obj_pose, env_idx=env_idx)
+        elif self._env_obbs_names is not None and name in self._env_obbs_names[env_idx]:
+            self.update_obb_pose(name=name, w_obj_pose=w_obj_pose, env_idx=env_idx)
+        else:
+            log_error("obstacle not found in OBB world model: " + name)
+
+    def enable_obstacle(
+        self,
+        name: str,
+        enable: bool = True,
+        env_idx: int = 0,
+    ):
+        if self._env_mesh_names is not None and name in self._env_mesh_names[env_idx]:
+            self.enable_mesh(enable, name, None, env_idx)
+        elif self._env_obbs_names is not None and name in self._env_obbs_names[env_idx]:
+            self.enable_obb(enable, name, None, env_idx)
+        else:
+            log_error("Obstacle not found in world model: " + name)
+        self.world_model.objects
+
+    def enable_mesh(
+        self,
+        enable: bool = True,
+        name: Optional[str] = None,
+        env_mesh_idx: Optional[torch.Tensor] = None,
+        env_idx: int = 0,
+    ):
+        """Update obstacle dimensions
+
+        Args:
+            obj_dims (torch.Tensor): [dim.x,dim.y, dim.z], give as [b,3]
+            obj_idx (torch.Tensor or int):
+
+        """
+        if env_mesh_idx is not None:
+            self._mesh_tensor_list[2][env_mesh_idx] = int(enable)  # enable == 1
+        else:
+            # find index of given name:
+            obs_idx = self.get_mesh_idx(name, env_idx)
+            self._mesh_tensor_list[2][env_idx, obs_idx] = int(enable)
+
+    def _get_sdf(
+        self,
+        query_spheres,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx=None,
+        return_loss=False,
+    ):
+        d = SdfMeshWarpPy.apply(
+            query_spheres,
+            collision_query_buffer.mesh_collision_buffer.distance_buffer,
+            collision_query_buffer.mesh_collision_buffer.grad_distance_buffer,
+            collision_query_buffer.mesh_collision_buffer.sparsity_index_buffer,
+            weight,
+            activation_distance,
+            self._mesh_tensor_list[0],
+            self._mesh_tensor_list[1],
+            self._mesh_tensor_list[2],
+            self._env_n_mesh,
+            self.max_distance,
+            env_query_idx,
+            return_loss,
+        )
+        return d
+
+    def _get_swept_sdf(
+        self,
+        query_spheres,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        sweep_steps: int,
+        enable_speed_metric=False,
+        env_query_idx=None,
+        return_loss: bool = False,
+    ):
+        d = SweptSdfMeshWarpPy.apply(
+            query_spheres,
+            collision_query_buffer.mesh_collision_buffer.distance_buffer,
+            collision_query_buffer.mesh_collision_buffer.grad_distance_buffer,
+            collision_query_buffer.mesh_collision_buffer.sparsity_index_buffer,
+            weight,
+            activation_distance,
+            speed_dt,
+            self._mesh_tensor_list[0],
+            self._mesh_tensor_list[1],
+            self._mesh_tensor_list[2],
+            self._env_n_mesh,
+            sweep_steps,
+            enable_speed_metric,
+            self.max_distance,
+            env_query_idx,
+            return_loss,
+        )
+        return d
+
+    def get_sphere_distance(
+        self,
+        query_sphere: torch.Tensor,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        # TODO: if no mesh object exist, call primitive
+        if "mesh" not in self.collision_types or not self.collision_types["mesh"]:
+            return super().get_sphere_distance(
+                query_sphere,
+                collision_query_buffer,
+                weight=weight,
+                activation_distance=activation_distance,
+                env_query_idx=env_query_idx,
+                return_loss=return_loss,
+            )
+
+        d = self._get_sdf(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            activation_distance=activation_distance,
+            env_query_idx=env_query_idx,
+            return_loss=return_loss,
+        )
+
+        if "primitive" not in self.collision_types or not self.collision_types["primitive"]:
+            return d
+        d_prim = super().get_sphere_distance(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            activation_distance=activation_distance,
+            env_query_idx=env_query_idx,
+            return_loss=return_loss,
+        )
+        d_val = d.view(d_prim.shape) + d_prim
+        return d_val
+
+    def get_sphere_collision(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        env_query_idx=None,
+        return_loss=False,
+    ):
+        if "mesh" not in self.collision_types or not self.collision_types["mesh"]:
+            return super().get_sphere_collision(
+                query_sphere,
+                collision_query_buffer,
+                weight=weight,
+                activation_distance=activation_distance,
+                env_query_idx=env_query_idx,
+                return_loss=return_loss,
+            )
+
+        d = self._get_sdf(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            activation_distance=activation_distance,
+            env_query_idx=env_query_idx,
+            return_loss=return_loss,
+        )
+
+        if "primitive" not in self.collision_types or not self.collision_types["primitive"]:
+            return d
+
+        d_prim = super().get_sphere_collision(
+            query_sphere,
+            collision_query_buffer,
+            weight,
+            activation_distance=activation_distance,
+            env_query_idx=env_query_idx,
+            return_loss=return_loss,
+        )
+        d_val = d.view(d_prim.shape) + d_prim
+
+        return d_val
+
+    def get_swept_sphere_distance(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        sweep_steps: int,
+        enable_speed_metric=False,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        # log_warn("Swept: Mesh + Primitive Collision Checking is experimental")
+        if "mesh" not in self.collision_types or not self.collision_types["mesh"]:
+            return super().get_swept_sphere_distance(
+                query_sphere,
+                collision_query_buffer,
+                weight=weight,
+                env_query_idx=env_query_idx,
+                sweep_steps=sweep_steps,
+                activation_distance=activation_distance,
+                speed_dt=speed_dt,
+                enable_speed_metric=enable_speed_metric,
+                return_loss=return_loss,
+            )
+
+        d = self._get_swept_sdf(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            env_query_idx=env_query_idx,
+            sweep_steps=sweep_steps,
+            activation_distance=activation_distance,
+            speed_dt=speed_dt,
+            enable_speed_metric=enable_speed_metric,
+            return_loss=return_loss,
+        )
+        if "primitive" not in self.collision_types or not self.collision_types["primitive"]:
+            return d
+
+        d_prim = super().get_swept_sphere_distance(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            env_query_idx=env_query_idx,
+            sweep_steps=sweep_steps,
+            activation_distance=activation_distance,
+            speed_dt=speed_dt,
+            enable_speed_metric=enable_speed_metric,
+            return_loss=return_loss,
+        )
+        d_val = d.view(d_prim.shape) + d_prim
+
+        return d_val
+
+    def get_swept_sphere_collision(
+        self,
+        query_sphere,
+        collision_query_buffer: CollisionQueryBuffer,
+        weight: torch.Tensor,
+        sweep_steps,
+        activation_distance: torch.Tensor,
+        speed_dt: torch.Tensor,
+        enable_speed_metric=False,
+        env_query_idx: Optional[torch.Tensor] = None,
+        return_loss: bool = False,
+    ):
+        if "mesh" not in self.collision_types or not self.collision_types["mesh"]:
+            return super().get_swept_sphere_collision(
+                query_sphere,
+                collision_query_buffer,
+                weight=weight,
+                env_query_idx=env_query_idx,
+                sweep_steps=sweep_steps,
+                activation_distance=activation_distance,
+                speed_dt=speed_dt,
+                enable_speed_metric=enable_speed_metric,
+                return_loss=return_loss,
+            )
+        d = self._get_swept_sdf(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            activation_distance=activation_distance,
+            speed_dt=speed_dt,
+            env_query_idx=env_query_idx,
+            sweep_steps=sweep_steps,
+            enable_speed_metric=enable_speed_metric,
+            return_loss=return_loss,
+        )
+        if "primitive" not in self.collision_types or not self.collision_types["primitive"]:
+            return d
+
+        d_prim = super().get_swept_sphere_collision(
+            query_sphere,
+            collision_query_buffer,
+            weight=weight,
+            env_query_idx=env_query_idx,
+            sweep_steps=sweep_steps,
+            activation_distance=activation_distance,
+            speed_dt=speed_dt,
+            enable_speed_metric=enable_speed_metric,
+            return_loss=return_loss,
+        )
+        d_val = d.view(d_prim.shape) + d_prim
+
+        return d_val
+
+    def clear_cache(self):
+        self._wp_mesh_cache = {}
+        if self._mesh_tensor_list is not None:
+            self._mesh_tensor_list[2][:] = 0
+        super().clear_cache()

src/curobo/geom/sphere_fit.py

@@ -0,0 +1,239 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+# Standard Library
+from enum import Enum
+from typing import List, Tuple
+
+# Third Party
+import numpy as np
+import torch
+import trimesh
+from trimesh.voxel.creation import voxelize
+
+# CuRobo
+from curobo.util.logger import log_warn
+
+
+class SphereFitType(Enum):
+    """Supported sphere fit types are listed here. VOXEL_VOLUME_SAMPLE_SURFACE works best.
+
+    See :ref:`attach_object_note` for more details.
+    """
+
+    #:  samples the surface of the mesh to approximate geometry
+    SAMPLE_SURFACE = "sample_surface"
+    #: voxelizes the volume and returns voxel positions that are intersecting with surface.
+    VOXEL_SURFACE = "voxel_surface"
+    #: voxelizes the volume and returns all ocupioed voxel positions.
+    VOXEL_VOLUME = "voxel_volume"
+    #: voxelizes the volume and returns voxel positions that are inside the surface of the geometry
+    VOXEL_VOLUME_INSIDE = "voxel_volume_inside"
+    #: voxelizes the volume and returns voxel positions that are inside the surface,
+    #: along with surface sampled positions
+    VOXEL_VOLUME_SAMPLE_SURFACE = "voxel_volume_sample_surface"
+
+
+def sample_even_fit_mesh(mesh: trimesh.Trimesh, n_spheres: int, sphere_radius: float):
+    n_pts = trimesh.sample.sample_surface_even(mesh, n_spheres)[0]
+    n_radius = [sphere_radius for _ in range(len(n_pts))]
+    return n_pts, n_radius
+
+
+def get_voxel_pitch(mesh: trimesh.Trimesh, n_cubes: int):
+    d = mesh.extents
+    cube_volume = d[0] * d[1] * d[2]
+    v = mesh.volume
+    if v > 0:
+        occupancy = 1.0 - ((cube_volume - v) / cube_volume)
+    else:
+        log_warn("sphere_fit: occupancy test failed, assuming cuboid volume")
+        occupancy = 1.0
+
+    # given the extents, find the radius to get number of spheres:
+    pitch = (occupancy * cube_volume / n_cubes) ** (1 / 3)
+    return pitch
+
+
+def voxel_fit_surface_mesh(
+    mesh: trimesh.Trimesh,
+    n_spheres: int,
+    sphere_radius: float,
+    voxelize_method: str = "ray",
+):
+    pts, rad = get_voxelgrid_from_mesh(mesh, n_spheres, voxelize_method)
+    if pts is None:
+        return pts, rad
+    pr = trimesh.proximity.ProximityQuery(mesh)
+    if len(pts) <= n_spheres:
+        surface_pts, _, _ = pr.on_surface(pts)
+        n_radius = [sphere_radius for _ in range(len(surface_pts))]
+        return surface_pts, n_radius
+
+    # compute signed distance:
+    dist = pr.signed_distance(pts)
+
+    # calculate distance to boundary:
+    dist = np.abs(dist - rad)
+    # get the first n points closest to boundary:
+    _, idx = torch.topk(torch.as_tensor(dist), k=n_spheres, largest=False)
+
+    n_pts = pts[idx]
+    n_radius = [sphere_radius for _ in range(len(n_pts))]
+    return n_pts, n_radius
+
+
+def get_voxelgrid_from_mesh(mesh: trimesh.Trimesh, n_spheres: int, voxelize_method: str = "ray"):
+    pitch = get_voxel_pitch(mesh, n_spheres)
+    radius = pitch / 2.0
+    try:
+        voxel = voxelize(mesh, pitch, voxelize_method)
+        voxel = voxel.fill("base")
+        pts = voxel.points
+        rad = np.ravel([radius for _ in range(len(pts))])
+    except:
+        log_warn("voxelization failed")
+        pts = rad = None
+    return pts, rad
+
+
+def voxel_fit_mesh(
+    mesh: trimesh.Trimesh,
+    n_spheres: int,
+    surface_sphere_radius: float,
+    voxelize_method: str = "ray",
+):
+    pts, rad = get_voxelgrid_from_mesh(mesh, n_spheres, voxelize_method)
+    if pts is None:
+        return pts, rad
+    # compute signed distance:
+    pr = trimesh.proximity.ProximityQuery(mesh)
+    dist = pr.signed_distance(pts)
+
+    # calculate distance to boundary:
+    dist = dist - rad
+
+    # all negative values are outside the mesh:
+    idx = dist < 0.0
+    surface_pts, _, _ = pr.on_surface(pts[idx])
+    pts[idx] = surface_pts
+    rad[idx] = surface_sphere_radius
+    if len(pts) > n_spheres:
+        # remove some surface pts:
+        inside_idx = dist >= 0.0
+        inside_pts = pts[inside_idx]
+        if len(inside_pts) < n_spheres:
+            new_pts = np.zeros((n_spheres, 3))
+            new_pts[: len(inside_pts)] = inside_pts
+            new_radius = np.zeros(n_spheres)
+            new_radius[: len(inside_pts)] = rad[inside_idx]
+
+            new_pts[len(inside_pts) :] = surface_pts[: n_spheres - len(inside_pts)]
+            new_radius[len(inside_pts) :] = surface_sphere_radius
+            pts = new_pts
+            rad = new_radius
+
+    n_pts = pts
+    n_radius = rad.tolist()
+
+    return n_pts, n_radius
+
+
+def voxel_fit_volume_sample_surface_mesh(
+    mesh: trimesh.Trimesh,
+    n_spheres: int,
+    surface_sphere_radius: float,
+    voxelize_method: str = "ray",
+):
+    pts, rad = voxel_fit_volume_inside_mesh(mesh, 0.75 * n_spheres, voxelize_method)
+    if pts is None:
+        return pts, rad
+    # compute surface points:
+    if len(pts) >= n_spheres:
+        return pts, rad
+
+    sample_count = n_spheres - (len(pts))
+
+    surface_sample_pts, sample_radius = sample_even_fit_mesh(
+        mesh, sample_count, surface_sphere_radius
+    )
+    pts = np.concatenate([pts, surface_sample_pts])
+    rad = np.concatenate([rad, sample_radius])
+    return pts, rad
+
+
+def voxel_fit_volume_inside_mesh(
+    mesh: trimesh.Trimesh,
+    n_spheres: int,
+    voxelize_method: str = "ray",
+):
+    pts, rad = get_voxelgrid_from_mesh(mesh, 2 * n_spheres, voxelize_method)
+    if pts is None:
+        return pts, rad
+    # compute signed distance:
+    pr = trimesh.proximity.ProximityQuery(mesh)
+    dist = pr.signed_distance(pts)
+
+    # calculate distance to boundary:
+    dist = dist - rad
+    # all negative values are outside the mesh:
+    idx = dist > 0.0
+    n_pts = pts[idx]
+    n_radius = rad[idx].tolist()
+    return n_pts, n_radius
+
+
+def fit_spheres_to_mesh(
+    mesh: trimesh.Trimesh,
+    n_spheres: int,
+    surface_sphere_radius: float = 0.01,
+    fit_type: SphereFitType = SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
+    voxelize_method: str = "ray",
+) -> Tuple[np.ndarray, List[float]]:
+    """Approximate a mesh with spheres. See :ref:`attach_object_note` for more details.
+
+
+    Args:
+        mesh: Input trimesh
+        n_spheres: _description_
+        surface_sphere_radius: _description_. Defaults to 0.01.
+        fit_type: _description_. Defaults to SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE.
+        voxelize_method: _description_. Defaults to "ray".
+
+    Returns:
+        _description_
+    """
+    n_pts = n_radius = None
+    if fit_type == SphereFitType.SAMPLE_SURFACE:
+        n_pts, n_radius = sample_even_fit_mesh(mesh, n_spheres, surface_sphere_radius)
+    elif fit_type == SphereFitType.VOXEL_SURFACE:
+        n_pts, n_radius = voxel_fit_surface_mesh(
+            mesh, n_spheres, surface_sphere_radius, voxelize_method
+        )
+    elif fit_type == SphereFitType.VOXEL_VOLUME_INSIDE:
+        n_pts, n_radius = voxel_fit_volume_inside_mesh(mesh, n_spheres, voxelize_method)
+    elif fit_type == SphereFitType.VOXEL_VOLUME:
+        n_pts, n_radius = voxel_fit_mesh(mesh, n_spheres, surface_sphere_radius, voxelize_method)
+    elif fit_type == SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE:
+        n_pts, n_radius = voxel_fit_volume_sample_surface_mesh(
+            mesh, n_spheres, surface_sphere_radius, voxelize_method
+        )
+    if n_pts is None or len(n_pts) < 1:
+        log_warn("sphere_fit: falling back to sampling surface")
+        n_pts, n_radius = sample_even_fit_mesh(mesh, n_spheres, surface_sphere_radius)
+
+    if n_pts is not None and len(n_pts) > n_spheres:
+        samples = torch.as_tensor(n_pts)
+        dist = torch.linalg.norm(samples - torch.mean(samples, dim=-1).unsqueeze(1), dim=-1)
+        _, knn_i = dist.topk(n_spheres, largest=True)
+        n_pts = samples[knn_i].cpu().numpy()
+        n_radius = np.ravel(n_radius)[knn_i.cpu().flatten().tolist()].tolist()
+    return n_pts, n_radius

src/curobo/geom/types.py

@@ -0,0 +1,813 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+from __future__ import annotations
+
+# Standard Library
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Sequence
+
+# Third Party
+import numpy as np
+import torch
+import trimesh
+
+# CuRobo
+from curobo.geom.sphere_fit import SphereFitType, fit_spheres_to_mesh
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util.logger import log_error, log_warn
+from curobo.util_file import get_assets_path, join_path
+
+
+@dataclass
+class Material:
+    metallic: float = 0.0
+    roughness: float = 0.4
+
+
+@dataclass
+class Obstacle:
+    """Base class for all obstacles."""
+
+    #: Unique name of obstacle.
+    name: str
+
+    #: Pose of obstacle as a list with format [x y z qw qx qy qz]
+    pose: Optional[List[float]] = None
+
+    #: NOTE: implement scaling for all obstacle types.
+    scale: Optional[List[float]] = None
+
+    #: Color of obstacle to use in visualization.
+    color: Optional[List[float]] = None
+
+    #: texture to apply to obstacle in visualization.
+    texture_id: Optional[str] = None
+
+    #: texture to apply to obstacle in visualization.
+    texture: Optional[str] = None
+
+    #: material properties to apply in visualization.
+    material: Material = Material()
+
+    tensor_args: TensorDeviceType = TensorDeviceType()
+
+    def get_trimesh_mesh(self, process: bool = True) -> trimesh.Trimesh:
+        """Create a trimesh instance from the obstacle representation.
+
+        Args:
+            process (bool, optional): process when loading from file. Defaults to True.
+
+        Raises:
+            NotImplementedError: requires implementation in derived class.
+
+        Returns:
+            trimesh.Trimesh: instance of obstacle as a trimesh.
+        """
+        raise NotImplementedError
+
+    def save_as_mesh(self, file_path: str):
+        mesh_scene = self.get_trimesh_mesh()
+        mesh_scene.export(file_path)
+
+    def get_cuboid(self) -> Cuboid:
+        """Get oriented bounding box of obstacle (OBB).
+
+        Returns:
+            Cuboid: returns obstacle as a cuboid.
+        """
+        # create a trimesh object:
+        m = self.get_trimesh_mesh()
+        # compute bounding box:
+        dims = m.bounding_box_oriented.primitive.extents
+        dims = dims
+        offset = m.bounding_box_oriented.primitive.transform
+        new_pose = self.pose
+
+        base_pose = Pose.from_list(self.pose)
+        offset_pose = Pose.from_matrix(self.tensor_args.to_device(offset))
+        new_base_pose = base_pose.multiply(offset_pose)
+        new_pose = new_base_pose.tolist()
+
+        if self.color is None:
+            # get average color:
+            if isinstance(m.visual, trimesh.visual.texture.TextureVisuals):
+                m.visual = m.visual.to_color()
+            if isinstance(m.visual, trimesh.visual.color.ColorVisuals):
+                self.color = (np.ravel(m.visual.main_color) / 255.0).tolist()
+        return Cuboid(
+            name=self.name,
+            pose=new_pose,
+            dims=dims.tolist(),
+            color=self.color,
+            texture=self.texture,
+            material=self.material,
+            tensor_args=self.tensor_args,
+        )
+
+    def get_mesh(self, process: bool = True) -> Mesh:
+        """Get obstacle as a mesh.
+
+        Args:
+            process (bool, optional): process mesh from file. Defaults to True.
+
+        Returns:
+            Mesh: obstacle as a mesh.
+        """
+        # load sphere as a mesh in trimesh:
+        m = self.get_trimesh_mesh(process=process)
+        if isinstance(m.visual, trimesh.visual.texture.TextureVisuals):
+            m.visual = m.visual.to_color()
+
+        return Mesh(
+            name=self.name,
+            vertices=m.vertices,
+            faces=m.faces,
+            pose=self.pose,
+            color=self.color,
+            texture_id=self.texture_id,
+            texture=self.texture,
+            material=self.material,
+            vertex_colors=m.visual.vertex_colors,
+            face_colors=m.visual.face_colors,
+            tensor_args=self.tensor_args,
+        )
+
+    def get_transform_matrix(self) -> np.ndarray:
+        """Get homogenous transformation matrix from pose.
+
+        Returns:
+           np.ndarray : transformation matrix.
+        """
+        # convert pose to matrix:
+        mat = trimesh.scene.transforms.kwargs_to_matrix(
+            translation=self.pose[:3], quaternion=self.pose[3:]
+        )
+        return mat
+
+    def get_sphere(self, n: int = 1) -> Sphere:
+        """Compute a sphere that fits in the volume of the object.
+
+        Args:
+            n: number of spheres
+        Returns:
+            spheres
+        """
+        obb = self.get_cuboid()
+
+        # from obb, compute the number of spheres
+
+        # fit a sphere of radius equal to the smallest dim
+        r = min(obb.dims)
+        sph = Sphere(
+            name="m_sphere",
+            pose=obb.pose,
+            position=obb.pose[:3],
+            radius=r,
+            tensor_args=self.tensor_args,
+        )
+
+        return sph
+
+    def get_bounding_spheres(
+        self,
+        n_spheres: int = 1,
+        surface_sphere_radius: float = 0.002,
+        fit_type: SphereFitType = SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
+        voxelize_method: str = "ray",
+        pre_transform_pose: Optional[Pose] = None,
+        tensor_args: TensorDeviceType = TensorDeviceType(),
+    ) -> List[Sphere]:
+        """Compute n spheres that fits in the volume of the object.
+
+        Args:
+            n: number of spheres
+        Returns:
+            spheres
+        """
+        mesh = self.get_trimesh_mesh()
+        pts, n_radius = fit_spheres_to_mesh(
+            mesh, n_spheres, surface_sphere_radius, fit_type, voxelize_method=voxelize_method
+        )
+
+        obj_pose = Pose.from_list(self.pose, tensor_args)
+        # transform object:
+
+        # transform points:
+        if pre_transform_pose is not None:
+            obj_pose = pre_transform_pose.multiply(obj_pose)  # convert object pose to another frame
+
+        if pts is None or len(pts) == 0:
+            log_warn("spheres could not be fit!, adding one sphere at origin")
+            pts = np.zeros((1, 3))
+            pts[0, :] = mesh.centroid
+            n_radius = [0.02]
+            obj_pose = Pose.from_list([0, 0, 0, 1, 0, 0, 0], tensor_args)
+
+        points_cuda = tensor_args.to_device(pts)
+        pts = obj_pose.transform_points(points_cuda).cpu().view(-1, 3).numpy()
+
+        new_spheres = [
+            Sphere(
+                name="sph_" + str(i),
+                pose=[pts[i, 0], pts[i, 1], pts[i, 2], 1, 0, 0, 0],
+                radius=n_radius[i],
+            )
+            for i in range(pts.shape[0])
+        ]
+
+        return new_spheres
+
+
+@dataclass
+class Cuboid(Obstacle):
+    """Represent obstacle as a cuboid."""
+
+    #: Dimensions of cuboid in meters [x_length, y_length, z_length].
+    dims: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
+
+    def __post_init__(self):
+        if self.pose is None:
+            log_error("Cuboid Obstacle requires Pose")
+
+    def get_trimesh_mesh(self, process: bool = True):
+        m = trimesh.creation.box(extents=self.dims)
+        if self.color is not None:
+            color_visual = trimesh.visual.color.ColorVisuals(
+                mesh=m, face_colors=self.color, vertex_colors=self.color
+            )
+            m.visual = color_visual
+        return m
+
+
+@dataclass
+class Capsule(Obstacle):
+    radius: float = 0.0
+    base: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
+    tip: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
+
+    def get_trimesh_mesh(self, process: bool = True):
+        height = self.tip[2] - self.base[2]
+        if (
+            height < 0
+            or self.tip[0] != 0
+            or self.tip[1] != 0
+            or self.base[0] != 0
+            or self.base[1] != 0
+        ):
+            log_error(
+                "Capsule to Mesh is only supported when base and tip are at xy origin with a positive height"
+            )
+        m = trimesh.creation.capsule(radius=self.radius, height=self.tip[2] - self.base[2])
+        if self.color is not None:
+            color_visual = trimesh.visual.color.ColorVisuals(
+                mesh=m, face_colors=self.color, vertex_colors=self.color
+            )
+            m.visual = color_visual
+        return m
+
+
+@dataclass
+class Cylinder(Obstacle):
+    radius: float = 0.0
+    height: float = 0.0
+
+    def get_trimesh_mesh(self, process: bool = True):
+        m = trimesh.creation.cylinder(radius=self.radius, height=self.height)
+        if self.color is not None:
+            color_visual = trimesh.visual.color.ColorVisuals(
+                mesh=m, face_colors=self.color, vertex_colors=self.color
+            )
+            m.visual = color_visual
+        return m
+
+
+@dataclass
+class Sphere(Obstacle):
+    radius: float = 0.0
+
+    #: position is deprecated, use pose instead
+    position: Optional[List[float]] = None
+
+    def __post_init__(self):
+        if self.position is not None:
+            self.pose = self.position + [1, 0, 0, 0]
+            log_warn("Sphere.position is deprecated, use Sphere.pose instead")
+        if self.pose is not None:
+            self.position = self.pose[:3]
+
+    def get_trimesh_mesh(self, process: bool = True):
+        m = trimesh.creation.icosphere(radius=self.radius)
+        if self.color is not None:
+            color_visual = trimesh.visual.color.ColorVisuals(
+                mesh=m, face_colors=self.color, vertex_colors=self.color
+            )
+            m.visual = color_visual
+        return m
+
+    def get_cuboid(self) -> Cuboid:
+        """Get oriented bounding box of obstacle (OBB).
+
+        Returns:
+            Cuboid: returns obstacle as a cuboid.
+        """
+
+        # create a trimesh object:
+        m = self.get_trimesh_mesh()
+        # compute bounding box:
+        dims = m.bounding_box_oriented.primitive.extents
+        dims = dims
+        offset = m.bounding_box_oriented.primitive.transform
+        new_pose = self.pose
+
+        base_pose = Pose.from_list(self.pose)
+        offset_pose = Pose.from_matrix(self.tensor_args.to_device(offset))
+        new_base_pose = base_pose.multiply(offset_pose)
+        new_pose = new_base_pose.tolist()
+
+        # since a sphere is symmetrical, we set the orientation to identity
+        new_pose[3:] = [1, 0, 0, 0]
+
+        if self.color is None:
+            # get average color:
+            if isinstance(m.visual, trimesh.visual.texture.TextureVisuals):
+                m.visual = m.visual.to_color()
+            if isinstance(m.visual, trimesh.visual.color.ColorVisuals):
+                self.color = (np.ravel(m.visual.main_color) / 255.0).tolist()
+        return Cuboid(
+            name=self.name,
+            pose=new_pose,
+            dims=dims.tolist(),
+            color=self.color,
+            texture=self.texture,
+            material=self.material,
+            tensor_args=self.tensor_args,
+        )
+
+
+@dataclass
+class Mesh(Obstacle):
+    file_path: Optional[str] = None
+    file_string: Optional[
+        str
+    ] = None  # storing full mesh as a string, loading from this is not implemented yet.
+    urdf_path: Optional[str] = None  # useful for visualization in isaac gym.
+    vertices: Optional[List[List[float]]] = None
+    faces: Optional[List[int]] = None
+    vertex_colors: Optional[List[List[float]]] = None
+    vertex_normals: Optional[List[List[float]]] = None
+    face_colors: Optional[List[List[float]]] = None
+
+    def __post_init__(self):
+        if self.file_path is not None:
+            self.file_path = join_path(get_assets_path(), self.file_path)
+        if self.urdf_path is not None:
+            self.urdf_path = join_path(get_assets_path(), self.urdf_path)
+        if self.scale is not None and self.vertices is not None:
+            self.vertices = np.ravel(self.scale) * self.vertices
+            self.scale = None
+
+    def get_trimesh_mesh(self, process: bool = True):
+        # load mesh from filepath or verts and faces:
+        if self.file_path is not None:
+            m = trimesh.load(self.file_path, process=process, force="mesh")
+            if isinstance(m, trimesh.Scene):
+                m = m.dump(concatenate=True)
+
+            if isinstance(m.visual, trimesh.visual.texture.TextureVisuals):
+                m.visual = m.visual.to_color()
+        else:
+            m = trimesh.Trimesh(
+                self.vertices,
+                self.faces,
+                vertex_colors=self.vertex_colors,
+                vertex_normals=self.vertex_normals,
+                face_colors=self.face_colors,
+            )
+        # if self.scale is not None:
+        #    m.vertices = np.ravel(self.scale) * m.vertices
+
+        return m
+
+    def get_mesh_data(self, process: bool = True):
+        verts = faces = None
+        if self.file_path is not None:
+            m = self.get_trimesh_mesh(process=process)
+            verts = m.vertices.view(np.ndarray)
+            faces = m.faces
+        elif self.vertices is not None and self.faces is not None:
+            verts = self.vertices
+            faces = self.faces
+        else:
+            ValueError("No Mesh object found")
+
+        return verts, faces
+
+    @staticmethod
+    def from_pointcloud(
+        pointcloud: np.ndarray,
+        pitch: float = 0.02,
+        name="world_pc",
+        pose: List[float] = [0, 0, 0, 1, 0, 0, 0],
+        filter_close_points: float = 0.0,
+    ):
+        if filter_close_points > 0.0:
+            # remove points that are closer than given threshold
+            dist = np.linalg.norm(pointcloud, axis=-1)
+            pointcloud = pointcloud[dist > filter_close_points]
+
+        pc = trimesh.PointCloud(pointcloud)
+        scene_mesh = trimesh.voxel.ops.points_to_marching_cubes(pc.vertices, pitch=pitch)
+        return Mesh(name, pose=pose, vertices=scene_mesh.vertices, faces=scene_mesh.faces)
+
+
+@dataclass
+class BloxMap(Obstacle):
+    map_path: Optional[str] = None
+    scale: List[float] = field(default_factory=lambda: [1.0, 1.0, 1.0])
+    voxel_size: float = 0.02
+    #: integrator type to use in nvblox. Options: ["tsdf", "occupancy"]
+    integrator_type: str = "tsdf"
+    mesh_file_path: Optional[str] = None
+    mapper_instance: Any = None
+    mesh: Optional[Mesh] = None
+
+    def __post_init__(self):
+        if self.map_path is not None:
+            self.map_path = join_path(get_assets_path(), self.map_path)
+        if self.mesh_file_path is not None:
+            self.mesh = Mesh(
+                name=self.name + "_mesh", file_path=self.mesh_file_path, pose=self.pose
+            )
+
+    def get_trimesh_mesh(self, process: bool = True):
+        if self.mesh is not None:
+            return self.mesh.get_trimesh_mesh(process)
+        else:
+            log_warn("no mesh found for obstacle: " + self.name)
+            return None
+
+
+@dataclass
+class WorldConfig(Sequence):
+    """Representation of World for use in CuRobo."""
+
+    #: List of Sphere obstacles.
+    sphere: Optional[List[Sphere]] = None
+
+    #: List of Cuboid obstacles.
+    cuboid: Optional[List[Cuboid]] = None
+
+    #: List of Capsule obstacles.
+    capsule: Optional[List[Capsule]] = None
+
+    #: List of Cylinder obstacles.
+    cylinder: Optional[List[Cylinder]] = None
+
+    #: List of Mesh obstacles.
+    mesh: Optional[List[Mesh]] = None
+
+    #: BloxMap obstacle.
+    blox: Optional[List[BloxMap]] = None
+
+    #: List of all obstacles in world.
+    objects: Optional[List[Obstacle]] = None
+
+    def __post_init__(self):
+        # create objects list:
+        if self.objects is None:
+            self.objects = []
+            if self.sphere is not None:
+                self.objects += self.sphere
+            if self.cuboid is not None:
+                self.objects += self.cuboid
+            if self.capsule is not None:
+                self.objects += self.capsule
+            if self.mesh is not None:
+                self.objects += self.mesh
+            if self.blox is not None:
+                self.objects += self.blox
+            if self.cylinder is not None:
+                self.objects += self.cylinder
+        if self.sphere is None:
+            self.sphere = []
+        if self.cuboid is None:
+            self.cuboid = []
+        if self.capsule is None:
+            self.capsule = []
+        if self.mesh is None:
+            self.mesh = []
+        if self.cylinder is None:
+            self.cylinder = []
+        if self.blox is None:
+            self.blox = []
+
+    def __len__(self):
+        return len(self.objects)
+
+    def __getitem__(self, idx):
+        return self.objects[idx]
+
+    def clone(self):
+        return WorldConfig(
+            cuboid=self.cuboid.copy() if self.cuboid is not None else None,
+            sphere=self.sphere.copy() if self.sphere is not None else None,
+            mesh=self.mesh.copy() if self.mesh is not None else None,
+            capsule=self.capsule.copy() if self.capsule is not None else None,
+            cylinder=self.cylinder.copy() if self.cylinder is not None else None,
+            blox=self.blox.copy() if self.blox is not None else None,
+        )
+
+    @staticmethod
+    def from_dict(data_dict: Dict[str, Any]) -> WorldConfig:
+        cuboid = None
+        sphere = None
+        capsule = None
+        mesh = None
+        blox = None
+        cylinder = None
+        # load yaml:
+        if "cuboid" in data_dict.keys():
+            cuboid = [Cuboid(name=x, **data_dict["cuboid"][x]) for x in data_dict["cuboid"]]
+        if "sphere" in data_dict.keys():
+            sphere = [Sphere(name=x, **data_dict["sphere"][x]) for x in data_dict["sphere"]]
+        if "mesh" in data_dict.keys():
+            mesh = [Mesh(name=x, **data_dict["mesh"][x]) for x in data_dict["mesh"]]
+        if "capsule" in data_dict.keys():
+            capsule = [Capsule(name=x, **data_dict["capsule"][x]) for x in data_dict["capsule"]]
+        if "cylinder" in data_dict.keys():
+            cylinder = [Cylinder(name=x, **data_dict["cylinder"][x]) for x in data_dict["cylinder"]]
+        if "blox" in data_dict.keys():
+            blox = [BloxMap(name=x, **data_dict["blox"][x]) for x in data_dict["blox"]]
+
+        return WorldConfig(
+            cuboid=cuboid,
+            sphere=sphere,
+            capsule=capsule,
+            cylinder=cylinder,
+            mesh=mesh,
+            blox=blox,
+        )
+
+    # load world config as obbs: convert all types to obbs
+    @staticmethod
+    def create_obb_world(current_world: WorldConfig):
+        sphere_obb = []
+        capsule_obb = []
+        cylinder_obb = []
+        mesh_obb = []
+        blox_obb = []
+        cuboid_obb = current_world.cuboid
+
+        if current_world.capsule is not None and len(current_world.capsule) > 0:
+            capsule_obb = [x.get_cuboid() for x in current_world.capsule]
+        if current_world.sphere is not None and len(current_world.sphere) > 0:
+            sphere_obb = [x.get_cuboid() for x in current_world.sphere]
+        if current_world.cylinder is not None and len(current_world.cylinder) > 0:
+            cylinder_obb = [x.get_cuboid() for x in current_world.cylinder]
+        if current_world.blox is not None and len(current_world.blox) > 0:
+            for i in range(len(current_world.blox)):
+                if current_world.blox[i].mesh is not None:
+                    blox_obb.append(current_world.blox[i].get_cuboid())
+
+        if current_world.mesh is not None and len(current_world.mesh) > 0:
+            mesh_obb = [x.get_cuboid() for x in current_world.mesh]
+        return WorldConfig(
+            cuboid=cuboid_obb + sphere_obb + capsule_obb + cylinder_obb + mesh_obb + blox_obb
+        )
+
+    @staticmethod
+    def create_mesh_world(current_world: WorldConfig, process: bool = False):
+        sphere_obb = []
+        capsule_obb = []
+        cuboid_obb = []
+        cylinder_obb = []
+        blox_obb = []
+        if current_world.capsule is not None and len(current_world.capsule) > 0:
+            capsule_obb = [x.get_mesh(process=process) for x in current_world.capsule]
+
+        if current_world.sphere is not None and len(current_world.sphere) > 0:
+            sphere_obb = [x.get_mesh(process=process) for x in current_world.sphere]
+
+        if current_world.cuboid is not None and len(current_world.cuboid) > 0:
+            cuboid_obb = [x.get_mesh(process=process) for x in current_world.cuboid]
+
+        if current_world.cylinder is not None and len(current_world.cylinder) > 0:
+            cylinder_obb = [x.get_mesh(process=process) for x in current_world.cylinder]
+        if current_world.blox is not None and len(current_world.blox) > 0:
+            for i in range(len(current_world.blox)):
+                if current_world.blox[i].mesh is not None:
+                    blox_obb.append(current_world.blox[i].get_mesh(process=process))
+
+        return WorldConfig(
+            mesh=current_world.mesh
+            + sphere_obb
+            + capsule_obb
+            + cuboid_obb
+            + cylinder_obb
+            + blox_obb
+        )
+
+    @staticmethod
+    def create_collision_support_world(current_world: WorldConfig, process: bool = True):
+        sphere_obb = []
+        capsule_obb = []
+        cuboid_obb = []
+        cylinder_obb = []
+        blox_obb = []
+        if current_world.capsule is not None and len(current_world.capsule) > 0:
+            capsule_obb = [x.get_mesh(process=process) for x in current_world.capsule]
+
+        if current_world.sphere is not None and len(current_world.sphere) > 0:
+            sphere_obb = [x.get_mesh(process=process) for x in current_world.sphere]
+
+        if current_world.cuboid is not None and len(current_world.cuboid) > 0:
+            cuboid_obb = current_world.cuboid
+
+        if current_world.cylinder is not None and len(current_world.cylinder) > 0:
+            cylinder_obb = [x.get_mesh(process=process) for x in current_world.cylinder]
+        if current_world.blox is not None and len(current_world.blox) > 0:
+            for i in range(len(current_world.blox)):
+                if current_world.blox[i].mesh is not None:
+                    blox_obb.append(current_world.blox[i].get_mesh(process=process))
+
+        return WorldConfig(
+            mesh=current_world.mesh + sphere_obb + capsule_obb + cylinder_obb + blox_obb,
+            cuboid=cuboid_obb,
+        )
+
+    @staticmethod
+    def get_scene_graph(current_world: WorldConfig):
+        m_world = WorldConfig.create_mesh_world(current_world)
+        mesh_scene = trimesh.scene.scene.Scene(base_frame="world")
+        mesh_list = m_world
+        for m in mesh_list:
+            mesh_scene.add_geometry(
+                m.get_trimesh_mesh(),
+                geom_name=m.name,
+                parent_node_name="world",
+                transform=m.get_transform_matrix(),
+            )
+
+        return mesh_scene
+
+    @staticmethod
+    def create_merged_mesh_world(current_world: WorldConfig, process: bool = True):
+        mesh_scene = WorldConfig.get_scene_graph(current_world)
+        mesh_scene = mesh_scene.dump(concatenate=True)
+        new_mesh = Mesh(
+            vertices=mesh_scene.vertices.view(np.ndarray),
+            faces=mesh_scene.faces,
+            name="merged_mesh",
+            pose=[0, 0, 0, 1, 0, 0, 0],
+        )
+        return WorldConfig(mesh=[new_mesh])
+
+    def get_obb_world(self):
+        return WorldConfig.create_obb_world(self)
+
+    def get_mesh_world(self, merge_meshes: bool = False, process: bool = False):
+        if merge_meshes:
+            return WorldConfig.create_merged_mesh_world(self, process=process)
+        else:
+            return WorldConfig.create_mesh_world(self, process=process)
+
+    def get_collision_check_world(self, mesh_process: bool = False):
+        return WorldConfig.create_collision_support_world(self, process=mesh_process)
+
+    def save_world_as_mesh(self, file_path: str, save_as_scene_graph=False):
+        mesh_scene = WorldConfig.get_scene_graph(self)
+        if save_as_scene_graph:
+            mesh_scene = mesh_scene.dump(concatenate=True)
+
+        mesh_scene.export(file_path)
+
+    def get_cache_dict(self) -> Dict[str, int]:
+        """Computes the number of obstacles in each type
+
+        Returns:
+            _description_
+        """
+        cache = {"obb": len(self.cuboid), "mesh": len(self.mesh)}
+        return cache
+
+    def add_obstacle(self, obstacle: Obstacle):
+        if isinstance(obstacle, Mesh):
+            self.mesh.append(obstacle)
+        elif isinstance(obstacle, Cuboid):
+            self.cuboid.append(obstacle)
+        elif isinstance(obstacle, Sphere):
+            self.sphere.append(obstacle)
+        elif isinstance(obstacle, Cylinder):
+            self.cylinder.append(obstacle)
+        elif isinstance(obstacle, Capsule):
+            self.capsule.append(obstacle)
+        else:
+            ValueError("Obstacle type not supported")
+        self.objects.append(obstacle)
+
+    def randomize_color(self, r=[0, 1], g=[0, 1], b=[0, 1]):
+        """Randomize color of objects within the given range
+
+        Args:
+            r: _description_. Defaults to [0,1].
+            g: _description_. Defaults to [0,1].
+            b: _description_. Defaults to [0,1].
+
+        Returns:
+            _description_
+        """
+        n = len(self.objects)
+        r_l = np.random.uniform(r[0], r[1], n)
+        g_l = np.random.uniform(g[0], g[1], n)
+        b_l = np.random.uniform(b[0], b[1], n)
+        for i, i_val in enumerate(self.objects):
+            i_val.color = [r_l[i], g_l[i], b_l[i], 1.0]
+
+    def add_color(self, rgba=[0.0, 0.0, 0.0, 1.0]):
+        for i, i_val in enumerate(self.objects):
+            i_val.color = rgba
+
+    def add_material(self, material=Material()):
+        for i, i_val in enumerate(self.objects):
+            i_val.material = material
+
+    def get_obstacle(self, name: str) -> Union[None, Obstacle]:
+        for i in self.objects:
+            if i.name == name:
+                return i
+        return None
+
+    def remove_obstacle(self, name: str):
+        for i in range(len(self.objects)):
+            if self.objects[i].name == name:
+                del self.objects[i]
+                return
+
+    def remove_absolute_paths(self) -> WorldConfig:
+        for obj in self.objects:
+            if obj.name.startswith("/"):
+                obj.name = obj.name[1:]
+
+
+def tensor_sphere(pt, radius, tensor=None, tensor_args=TensorDeviceType()):
+    if tensor is None:
+        tensor = torch.empty(4, device=tensor_args.device, dtype=tensor_args.dtype)
+    tensor[:3] = torch.as_tensor(pt, device=tensor_args.device, dtype=tensor_args.dtype)
+    tensor[3] = radius
+    return tensor
+
+
+def tensor_capsule(base, tip, radius, tensor=None, tensor_args=TensorDeviceType()):
+    if tensor is None:
+        tensor = torch.empty(7, device=tensor_args.device, dtype=tensor_args.dtype)
+    tensor[:3] = torch.as_tensor(base, device=tensor_args.device, dtype=tensor_args.dtype)
+    tensor[3:6] = torch.as_tensor(tip, device=tensor_args.device, dtype=tensor_args.dtype)
+    tensor[6] = radius
+    return tensor
+
+
+def tensor_cube(pose, dims, tensor_args=TensorDeviceType()):
+    """
+
+    Args:
+        pose (_type_): x,y,z, qw,qx,qy,qz
+        dims (_type_): _description_
+        tensor_args (_type_, optional): _description_. Defaults to TensorDeviceType().
+
+    Returns:
+        _type_: _description_
+    """
+    w_T_b = Pose.from_list(pose, tensor_args=tensor_args)
+    b_T_w = w_T_b.inverse()
+    dims_t = torch.tensor(
+        [dims[0], dims[1], dims[2]], device=tensor_args.device, dtype=tensor_args.dtype
+    )
+    cube = [dims_t, b_T_w.get_pose_vector()]
+    return cube
+
+
+def batch_tensor_cube(pose, dims, tensor_args=TensorDeviceType()):
+    """
+
+    Args:
+        pose (_type_): x,y,z, qw,qx,qy,qz
+        dims (_type_): _description_
+        tensor_args (_type_, optional): _description_. Defaults to TensorDeviceType().
+
+    Returns:
+        _type_: _description_
+    """
+    w_T_b = Pose.from_batch_list(pose, tensor_args=tensor_args)
+    b_T_w = w_T_b.inverse()
+    dims_t = torch.as_tensor(np.array(dims), device=tensor_args.device, dtype=tensor_args.dtype)
+    cube = [dims_t, b_T_w.get_pose_vector()]
+    return cube