release repository

src/curobo/cuda_robot_model/__init__.py

@@ -0,0 +1,43 @@
+#
+# 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 cuda accelerated kinematics.
+
+Kinematics in CuRobo currently supports single axis actuated joints, where the joint can be actuated
+as prismatic or revolute joints. Continuous joints are approximated to revolute joints with limits
+at [-6, +6] radians. Mimic joints are not supported, so convert mimic joints to independent joints.
+
+CuRobo loads a robot's kinematic tree from :class:`types.KinematicsTensorConfig`. This config is
+generated using :class:`cuda_robot_generator.CudaRobotGenerator`. A parser base class 
+:class:`kinematics_parser.KinematicsParer` is provided to help with parsing kinematics from 
+standard formats. Kinematics parsing from URDF is implemented in 
+:class:`urdf_kinematics_parser.UrdfKinematicsParser`. An experimental USD kinematics parser is 
+provided in :class:`usd_kinematics_parser.UsdKinematicsParser`, which is missing an additional 
+transform between the joint origin and link origin, so this might not work for all robots.
+
+In addition to parsing data from a kinematics file (urdf, usd), CuRobo also needs a sphere 
+representation of the robot that approximates the volume of the robot's links with spheres. 
+Several other parameters are also needed to represent kinematics in CuRobo. A tutorial on setting up a 
+robot is provided in :ref:`tut_robot_configuration`.
+
+Once a robot configuration file is setup, you can pass this to 
+:class:`cuda_robot_model.CudaRobotModelConfig` to generate an instance of kinematics configuraiton.
+:class:`cuda_robot_model.CudaRobotModel` takes this configuration and provides access to kinematics
+computations.
+
+.. note::
+    :class:`cuda_robot_model.CudaRobotModel` creates memory tensors that are used by CUDA kernels 
+    while :class:`cuda_robot_model.CudaRobotModelConfig` contains only the robot kinematics 
+    configuration. To reduce memory overhead, you can pass one instance of 
+    :class:`cuda_robot_model.CudaRobotModelConfig` to many instances of 
+    :class:`cuda_robot_model.CudaRobotModel`.
+"""

src/curobo/cuda_robot_model/cuda_robot_generator.py

@@ -0,0 +1,888 @@
+#
+# 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
+import copy
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+# Third Party
+import importlib_resources
+import torch
+import torch.autograd.profiler as profiler
+
+# CuRobo
+from curobo.cuda_robot_model.kinematics_parser import LinkParams
+from curobo.cuda_robot_model.types import (
+    CSpaceConfig,
+    JointLimits,
+    JointType,
+    KinematicsTensorConfig,
+    SelfCollisionKinematicsConfig,
+)
+from curobo.cuda_robot_model.urdf_kinematics_parser import UrdfKinematicsParser
+from curobo.curobolib.kinematics import get_cuda_kinematics
+from curobo.geom.types import tensor_sphere
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.state import JointState
+from curobo.util.logger import log_error, log_info, log_warn
+from curobo.util_file import get_assets_path, get_robot_configs_path, join_path, load_yaml
+
+try:
+    # CuRobo
+    from curobo.cuda_robot_model.usd_kinematics_parser import UsdKinematicsParser
+except ImportError:
+    log_warn(
+        "USDParser failed to import, install curobo with pip install .[usd] "
+        + "or pip install usd-core"
+    )
+
+
+@dataclass
+class CudaRobotGeneratorConfig:
+    """Create Cuda Robot Model Configuration."""
+
+    #: Name of base link for kinematic tree.
+    base_link: str
+
+    #: Name of end-effector link to compute pose.
+    ee_link: str
+
+    #: Device to load cuda robot model.
+    tensor_args: TensorDeviceType = TensorDeviceType()
+
+    #: Name of link names to compute pose in addition to ee_link.
+    link_names: Optional[List[str]] = None
+
+    #: Name of links to compute sphere positions for use in collision checking.
+    collision_link_names: Optional[List[str]] = None
+
+    #: Collision spheres that fill the volume occupied by the links of the robot.
+    #: Collision spheres can be generated for robot using https://docs.omniverse.nvidia.com/app_isaacsim/app_isaacsim/advanced_tutorials/tutorial_motion_generation_robot_description_editor.html#collision-spheres
+    collision_spheres: Union[None, str, Dict[str, Any]] = None
+
+    #: Radius buffer to add to collision spheres as padding.
+    collision_sphere_buffer: float = 0.0
+
+    #: Compute jacobian of link poses. Currently not supported.
+    compute_jacobian: bool = False
+
+    #: Padding to add for self collision between links. Some robots use a large padding
+    #: for self collision avoidance (https://github.com/ros-planning/panda_moveit_config/pull/35#issuecomment-671333863)
+    self_collision_buffer: Optional[Dict[str, float]] = None
+
+    #: Self-collision ignore
+    self_collision_ignore: Optional[Dict[str, List[str]]] = None
+
+    #: debug config
+    debug: Optional[Dict[str, Any]] = None
+
+    #: Enabling this flag writes out the cumulative transformation matrix to global memory. This
+    #: allows for reusing the cumulative matrix during backward of kinematics (15% speedup over
+    #: recomputing cumul in backward).
+    use_global_cumul: bool = True
+
+    #: Path of meshes of robot links. Currently not used as we represent robot link geometry with
+    #: collision spheres.
+    asset_root_path: str = ""
+
+    #: Names of links to load meshes for visualization. This is only used for exporting
+    #: visualizations.
+    mesh_link_names: Optional[List[str]] = None
+
+    #: Set this to true to add mesh_link_names to link_names when computing kinematics.
+    load_link_names_with_mesh: bool = False
+
+    #: Path to load robot urdf.
+    urdf_path: Optional[str] = None
+
+    #: Path to load robot usd.
+    usd_path: Optional[str] = None
+
+    #: Root prim of robot in usd.
+    usd_robot_root: Optional[str] = None
+
+    #: Path of robot in Isaac server.
+    isaac_usd_path: Optional[str] = None
+
+    #: Load Kinematics chain from usd.
+    use_usd_kinematics: bool = False
+
+    #: Joints to flip axis when loading from USD
+    usd_flip_joints: Optional[List[str]] = None
+
+    #: Flip joint limits in USD.
+    usd_flip_joint_limits: Optional[List[str]] = None
+
+    #: Lock active joints in the kinematic tree. This will convert the joint to a fixed joint with
+    #: joint angle given from this dictionary.
+    lock_joints: Optional[Dict[str, float]] = None
+
+    extra_links: Optional[Dict[str, LinkParams]] = None
+
+    #: this is deprecated
+    add_object_link: bool = False
+
+    use_external_assets: bool = False
+
+    #: Create n collision spheres for links with name
+    extra_collision_spheres: Optional[Dict[str, int]] = None
+
+    #: cspace config
+    cspace: Union[None, CSpaceConfig, Dict[str, List[Any]]] = None
+
+    def __post_init__(self):
+        # add root path:
+        if self.urdf_path is not None:
+            self.urdf_path = join_path(get_assets_path(), self.urdf_path)
+        if self.usd_path is not None:
+            self.usd_path = join_path(get_assets_path(), self.usd_path)
+        if self.asset_root_path != "":
+            if self.use_external_assets:
+                with importlib_resources.files("content") as path:
+                    content_dir_posix = path
+                self.asset_root_path = join_path(content_dir_posix, self.asset_root_path)
+
+            else:
+                self.asset_root_path = join_path(get_assets_path(), self.asset_root_path)
+        elif self.urdf_path is not None:
+            self.asset_root_path = os.path.dirname(self.urdf_path)
+
+        if self.collision_spheres is None and (
+            self.collision_link_names is not None and len(self.collision_link_names) > 0
+        ):
+            log_error("collision link names are provided without robot collision spheres")
+        if self.load_link_names_with_mesh:
+            if self.link_names is None:
+                self.link_names = copy.deepcopy(self.mesh_link_names)
+            else:
+                for i in self.mesh_link_names:
+                    if i not in self.link_names:
+                        self.link_names.append(i)
+        if self.link_names is None:
+            self.link_names = [self.ee_link]
+        if self.collision_link_names is None:
+            self.collision_link_names = []
+        if self.ee_link not in self.link_names:
+            self.link_names.append(self.ee_link)
+        if self.collision_spheres is not None:
+            if isinstance(self.collision_spheres, str):
+                coll_yml = join_path(get_robot_configs_path(), self.collision_spheres)
+                coll_params = load_yaml(coll_yml)
+
+                self.collision_spheres = coll_params["collision_spheres"]
+            if self.extra_collision_spheres is not None:
+                for k in self.extra_collision_spheres.keys():
+                    new_spheres = [
+                        {"center": [0.0, 0.0, 0.0], "radius": -10.0}
+                        for n in range(self.extra_collision_spheres[k])
+                    ]
+                    self.collision_spheres[k] = new_spheres
+        if self.use_usd_kinematics and self.usd_path is None:
+            log_error("usd_path is required to load kinematics from usd")
+        if self.usd_flip_joints is None:
+            self.usd_flip_joints = {}
+        if self.usd_flip_joint_limits is None:
+            self.usd_flip_joint_limits = []
+        if self.extra_links is None:
+            self.extra_links = {}
+        else:
+            for k in self.extra_links.keys():
+                if isinstance(self.extra_links[k], dict):
+                    self.extra_links[k] = LinkParams.from_dict(self.extra_links[k])
+        if isinstance(self.cspace, Dict):
+            self.cspace = CSpaceConfig(**self.cspace, tensor_args=self.tensor_args)
+
+
+class CudaRobotGenerator(CudaRobotGeneratorConfig):
+    def __init__(self, config: CudaRobotGeneratorConfig) -> None:
+        super().__init__(**vars(config))
+        self.cpu_tensor_args = self.tensor_args.cpu()
+
+        self._self_collision_data = None
+        self.non_fixed_joint_names = []
+        self._n_dofs = 1
+
+        self.initialize_tensors()
+
+    @property
+    def kinematics_config(self):
+        return self._kinematics_config
+
+    @property
+    def self_collision_config(self):
+        return self._self_collision_data
+
+    @property
+    def kinematics_parser(self):
+        return self._kinematics_parser
+
+    @profiler.record_function("robot_generator/initialize_tensors")
+    def initialize_tensors(self):
+        self._joint_limits = None
+        self._self_collision_data = None
+        self.lock_jointstate = None
+        self.lin_jac, self.ang_jac = None, None
+
+        self._link_spheres_tensor = torch.empty(
+            (0, 4), device=self.tensor_args.device, dtype=self.tensor_args.dtype
+        )
+        self._link_sphere_idx_map = torch.empty(
+            (0), dtype=torch.int16, device=self.tensor_args.device
+        )
+        self.total_spheres = 0
+        self.self_collision_distance = (
+            torch.zeros(
+                (self.total_spheres, self.total_spheres),
+                dtype=self.tensor_args.dtype,
+                device=self.tensor_args.device,
+            )
+            - torch.inf
+        )
+        self.self_collision_offset = torch.zeros(
+            (self.total_spheres), dtype=self.tensor_args.dtype, device=self.tensor_args.device
+        )
+        # create a mega list of all links that we need:
+        other_links = copy.deepcopy(self.link_names)
+
+        for i in self.collision_link_names:
+            if i not in self.link_names:
+                other_links.append(i)
+        for i in self.extra_links:
+            p_name = self.extra_links[i].parent_link_name
+            if p_name not in self.link_names and p_name not in other_links:
+                other_links.append(p_name)
+
+        # other_links = list(set(self.link_names + self.collision_link_names))
+
+        # load kinematics parser based on file type:
+        # NOTE: Also add option to load from data buffers.
+        if self.use_usd_kinematics:
+            self._kinematics_parser = UsdKinematicsParser(
+                self.usd_path,
+                flip_joints=self.usd_flip_joints,
+                flip_joint_limits=self.usd_flip_joint_limits,
+                extra_links=self.extra_links,
+                usd_robot_root=self.usd_robot_root,
+            )
+        else:
+            self._kinematics_parser = UrdfKinematicsParser(
+                self.urdf_path, mesh_root=self.asset_root_path, extra_links=self.extra_links
+            )
+
+        if self.lock_joints is None:
+            self._build_kinematics(self.base_link, self.ee_link, other_links, self.link_names)
+        else:
+            self._build_kinematics_with_lock_joints(
+                self.base_link, self.ee_link, other_links, self.link_names, self.lock_joints
+            )
+        if self.cspace is None:
+            jpv = self._get_joint_position_velocity_limits()
+            self.cspace = CSpaceConfig.load_from_joint_limits(
+                jpv["position"][1, :], jpv["position"][0, :], self.joint_names, self.tensor_args
+            )
+
+        self.cspace.inplace_reindex(self.joint_names)
+        self._update_joint_limits()
+        self._ee_idx = self.link_names.index(self.ee_link)
+
+        # create kinematics tensor:
+        self._kinematics_config = KinematicsTensorConfig(
+            fixed_transforms=self._fixed_transform,
+            link_map=self._link_map,
+            joint_map=self._joint_map,
+            joint_map_type=self._joint_map_type,
+            store_link_map=self._store_link_map,
+            link_chain_map=self._link_chain_map,
+            link_names=self.link_names,
+            link_spheres=self._link_spheres_tensor,
+            link_sphere_idx_map=self._link_sphere_idx_map,
+            n_dof=self._n_dofs,
+            joint_limits=self._joint_limits,
+            non_fixed_joint_names=self.non_fixed_joint_names,
+            total_spheres=self.total_spheres,
+            link_name_to_idx_map=self._name_to_idx_map,
+            joint_names=self.joint_names,
+            debug=self.debug,
+            ee_idx=self._ee_idx,
+            mesh_link_names=self.mesh_link_names,
+            cspace=self.cspace,
+            base_link=self.base_link,
+            ee_link=self.ee_link,
+            lock_jointstate=self.lock_jointstate,
+        )
+        if self.asset_root_path != "":
+            self._kinematics_parser.add_absolute_path_to_link_meshes(self.asset_root_path)
+
+    def add_link(self, link_params: LinkParams):
+        self.extra_links[link_params.link_name] = link_params
+
+    def add_fixed_link(
+        self,
+        link_name: str,
+        parent_link_name: str,
+        joint_name: Optional[str] = None,
+        transform: Optional[Pose] = None,
+    ):
+        if transform is None:
+            transform = (
+                Pose.from_list([0, 0, 0, 1, 0, 0, 0], self.tensor_args)
+                .get_matrix()
+                .view(4, 4)
+                .cpu()
+                .numpy()
+            )
+        if joint_name is None:
+            joint_name = link_name + "_j_" + parent_link_name
+        link_params = LinkParams(
+            link_name=link_name,
+            parent_link_name=parent_link_name,
+            joint_name=joint_name,
+            fixed_transform=transform,
+            joint_type=JointType.FIXED,
+        )
+        self.add_link(link_params)
+
+    @profiler.record_function("robot_generator/build_chain")
+    def _build_chain(self, base_link, ee_link, other_links, link_names):
+        self._n_dofs = 0
+        self._controlled_joints = []
+        self._bodies = []
+
+        self._name_to_idx_map = dict()
+        self.base_link = base_link
+        self.ee_link = ee_link
+        self.joint_names = []
+        self._fixed_transform = []
+        chain_link_names = self._kinematics_parser.get_chain(base_link, ee_link)
+        self._add_body_to_tree(chain_link_names[0], base=True)
+        for i, l_name in enumerate(chain_link_names[1:]):
+            self._add_body_to_tree(l_name)
+        # check if all links are in the built tree:
+
+        for i in other_links:
+            if i in self._name_to_idx_map:
+                continue
+            if i not in self.extra_links.keys():
+                chain_l_names = self._kinematics_parser.get_chain(base_link, i)
+
+                for k in chain_l_names:
+                    if k in chain_link_names:
+                        continue
+                    # if link name is not in chain, add to chain
+                    chain_link_names.append(k)
+                    # add to tree:
+                    self._add_body_to_tree(k, base=False)
+        for i in self.extra_links.keys():
+            if i not in chain_link_names:
+                self._add_body_to_tree(i, base=False)
+                chain_link_names.append(i)
+
+        self.non_fixed_joint_names = self.joint_names.copy()
+        return chain_link_names
+
+    @profiler.record_function("robot_generator/build_kinematics_tensors")
+    def _build_kinematics_tensors(self, base_link, ee_link, link_names, chain_link_names):
+        link_map = [0 for i in range(len(self._bodies))]
+        store_link_map = []  # [-1 for i in range(len(self._bodies))]
+
+        joint_map = [
+            -1 if i not in self._controlled_joints else i for i in range(len(self._bodies))
+        ]
+        joint_map_type = [
+            -1 if i not in self._controlled_joints else i for i in range(len(self._bodies))
+        ]
+        all_joint_names = []
+        j_count = 0
+        ordered_link_names = []
+        # add body 0 details:
+        if self._bodies[0].link_name in link_names:
+            store_link_map.append(chain_link_names.index(self._bodies[0].link_name))
+            ordered_link_names.append(self._bodies[0].link_name)
+        # get joint types:
+        for i in range(1, len(self._bodies)):
+            body = self._bodies[i]
+            parent_name = body.parent_link_name
+            link_map[i] = self._name_to_idx_map[parent_name]
+            all_joint_names.append(body.joint_name)
+            if body.link_name in link_names:
+                store_link_map.append(chain_link_names.index(body.link_name))
+                ordered_link_names.append(body.link_name)
+            if i in self._controlled_joints:
+                joint_map[i] = j_count
+                joint_map_type[i] = body.joint_type.value
+                j_count += 1
+        self.link_names = ordered_link_names
+        # do a for loop to get link matrix:
+        link_chain_map = torch.eye(
+            len(chain_link_names), dtype=torch.int16, device=self.cpu_tensor_args.device
+        )
+
+        # iterate and set true:
+        for i in range(len(chain_link_names)):
+            chain_l_names = self._kinematics_parser.get_chain(base_link, chain_link_names[i])
+            for k in chain_l_names:
+                link_chain_map[i, chain_link_names.index(k)] = 1.0
+
+        self._link_map = torch.as_tensor(
+            link_map, device=self.tensor_args.device, dtype=torch.int16
+        )
+        self._joint_map = torch.as_tensor(
+            joint_map, device=self.tensor_args.device, dtype=torch.int16
+        )
+        self._joint_map_type = torch.as_tensor(
+            joint_map_type, device=self.tensor_args.device, dtype=torch.int8
+        )
+        self._store_link_map = torch.as_tensor(
+            store_link_map, device=self.tensor_args.device, dtype=torch.int16
+        )
+        self._link_chain_map = link_chain_map.to(device=self.tensor_args.device)
+        self._fixed_transform = torch.cat((self._fixed_transform), dim=0).to(
+            device=self.tensor_args.device
+        )
+        self._all_joint_names = all_joint_names
+
+    @profiler.record_function("robot_generator/build_kinematics")
+    def _build_kinematics(self, base_link, ee_link, other_links, link_names):
+        chain_link_names = self._build_chain(base_link, ee_link, other_links, link_names)
+        self._build_kinematics_tensors(base_link, ee_link, link_names, chain_link_names)
+        if self.collision_spheres is not None and len(self.collision_link_names) > 0:
+            self._build_collision_model(
+                self.collision_spheres, self.collision_link_names, self.collision_sphere_buffer
+            )
+
+    @profiler.record_function("robot_generator/build_kinematics_with_lock_joints")
+    def _build_kinematics_with_lock_joints(
+        self,
+        base_link,
+        ee_link,
+        other_links,
+        link_names,
+        lock_joints: Dict[str, float],
+    ):
+        chain_link_names = self._build_chain(base_link, ee_link, other_links, link_names)
+        # find links attached to lock joints:
+        lock_joint_names = list(lock_joints.keys())
+
+        joint_data = self._get_joint_links(lock_joint_names)
+
+        lock_links = list(
+            set(
+                [joint_data[j]["parent"] for j in joint_data.keys()]
+                + [joint_data[j]["child"] for j in joint_data.keys()]
+            )
+        )
+        new_link_names = link_names + lock_links
+
+        # rebuild kinematic tree with link names added to link pose computation:
+        self._build_kinematics_tensors(base_link, ee_link, new_link_names, chain_link_names)
+        if self.collision_spheres is not None and len(self.collision_link_names) > 0:
+            self._build_collision_model(
+                self.collision_spheres, self.collision_link_names, self.collision_sphere_buffer
+            )
+        # do forward kinematics and get transform for locked joints:
+        q = torch.zeros(
+            (1, self._n_dofs), device=self.tensor_args.device, dtype=self.tensor_args.dtype
+        )
+        # set lock joints in the joint angles:
+        l_idx = torch.as_tensor(
+            [self.joint_names.index(l) for l in lock_joints.keys()],
+            dtype=torch.long,
+            device=self.tensor_args.device,
+        )
+        l_val = self.tensor_args.to_device([lock_joints[l] for l in lock_joints.keys()])
+
+        q[0, l_idx] = l_val
+        kinematics_config = KinematicsTensorConfig(
+            fixed_transforms=self._fixed_transform,
+            link_map=self._link_map,
+            joint_map=self._joint_map,
+            joint_map_type=self._joint_map_type,
+            store_link_map=self._store_link_map,
+            link_chain_map=self._link_chain_map,
+            link_names=self.link_names,
+            link_spheres=self._link_spheres_tensor,
+            link_sphere_idx_map=self._link_sphere_idx_map,
+            n_dof=self._n_dofs,
+            joint_limits=self._joint_limits,
+            non_fixed_joint_names=self.non_fixed_joint_names,
+            total_spheres=self.total_spheres,
+        )
+        link_poses = self._get_link_poses(q, lock_links, kinematics_config)
+
+        # remove lock links from store map:
+        store_link_map = [chain_link_names.index(l) for l in link_names]
+        self._store_link_map = torch.as_tensor(
+            store_link_map, device=self.tensor_args.device, dtype=torch.int16
+        )
+        self.link_names = link_names
+
+        # compute a fixed transform for fixing joints:
+        with profiler.record_function("cuda_robot_generator/fix_locked_joints"):
+            # convert tensors to cpu:
+            self._joint_map_type = self._joint_map_type.to(device=self.cpu_tensor_args.device)
+            self._joint_map = self._joint_map.to(device=self.cpu_tensor_args.device)
+
+            for j in lock_joint_names:
+                w_parent = lock_links.index(joint_data[j]["parent"])
+                w_child = lock_links.index(joint_data[j]["child"])
+                parent_t_child = (
+                    link_poses.get_index(0, w_parent)
+                    .inverse()
+                    .multiply(link_poses.get_index(0, w_child))
+                )
+                # Make this joint as fixed
+                i = self._all_joint_names.index(j) + 1
+                self._joint_map_type[i] = -1
+                self._joint_map[i:] -= 1
+                self._joint_map[i] = -1
+                self._n_dofs -= 1
+                self._fixed_transform[i] = parent_t_child.get_matrix()
+                self._controlled_joints.remove(i)
+                self.joint_names.remove(j)
+            self._joint_map[self._joint_map < -1] = -1
+            self._joint_map = self._joint_map.to(device=self.tensor_args.device)
+            self._joint_map_type = self._joint_map_type.to(device=self.tensor_args.device)
+        if len(self.lock_joints.keys()) > 0:
+            self.lock_jointstate = JointState(
+                position=l_val, joint_names=list(self.lock_joints.keys())
+            )
+
+    @profiler.record_function("robot_generator/build_collision_model")
+    def _build_collision_model(
+        self,
+        collision_spheres: Dict,
+        collision_link_names: List[str],
+        collision_sphere_buffer: float = 0.0,
+    ):
+        """
+
+        Args:
+            collision_spheres (_type_): _description_
+            collision_link_names (_type_): _description_
+            collision_sphere_buffer (float, optional): _description_. Defaults to 0.0.
+        """
+
+        # We create all tensors on cpu and then finally move them to gpu
+        coll_link_spheres = []
+        # we store as [n_link, 7]
+        link_sphere_idx_map = []
+        cpu_tensor_args = self.tensor_args.cpu()
+        with profiler.record_function("robot_generator/build_collision_spheres"):
+            for j_idx, j in enumerate(collision_link_names):
+                # print(j_idx)
+                n_spheres = len(collision_spheres[j])
+                link_spheres = torch.zeros(
+                    (n_spheres, 4), dtype=cpu_tensor_args.dtype, device=cpu_tensor_args.device
+                )
+                # find link index in global map:
+                l_idx = self._name_to_idx_map[j]
+
+                for i in range(n_spheres):
+                    link_spheres[i, :] = tensor_sphere(
+                        collision_spheres[j][i]["center"],
+                        collision_spheres[j][i]["radius"],
+                        tensor_args=cpu_tensor_args,
+                        tensor=link_spheres[i, :],
+                    )
+                    link_sphere_idx_map.append(l_idx)
+                coll_link_spheres.append(link_spheres)
+                self.total_spheres += n_spheres
+
+        self._link_spheres_tensor = torch.cat(coll_link_spheres, dim=0)
+        new_radius = self._link_spheres_tensor[..., 3] + collision_sphere_buffer
+        flag = torch.logical_and(new_radius > -1.0, new_radius <= 0.0)
+        new_radius[flag] = 0.001
+        self._link_spheres_tensor[:, 3] = new_radius
+        self._link_sphere_idx_map = torch.as_tensor(
+            link_sphere_idx_map, dtype=torch.int16, device=cpu_tensor_args.device
+        )
+
+        # build self collision distance tensor:
+        self.self_collision_distance = (
+            torch.zeros(
+                (self.total_spheres, self.total_spheres),
+                dtype=cpu_tensor_args.dtype,
+                device=cpu_tensor_args.device,
+            )
+            - torch.inf
+        )
+        self.self_collision_offset = torch.zeros(
+            (self.total_spheres), dtype=cpu_tensor_args.dtype, device=cpu_tensor_args.device
+        )
+
+        with profiler.record_function("robot_generator/self_collision_distance"):
+            # iterate through each link:
+            for j_idx, j in enumerate(collision_link_names):
+                ignore_links = []
+                if j in self.self_collision_ignore.keys():
+                    ignore_links = self.self_collision_ignore[j]
+                link1_idx = self._name_to_idx_map[j]
+                link1_spheres_idx = torch.nonzero(self._link_sphere_idx_map == link1_idx)
+
+                rad1 = self._link_spheres_tensor[link1_spheres_idx, 3]
+                if j not in self.self_collision_buffer.keys():
+                    self.self_collision_buffer[j] = 0.0
+                c1 = self.self_collision_buffer[j]
+                self.self_collision_offset[link1_spheres_idx] = c1
+                for _, i_name in enumerate(collision_link_names):
+                    if i_name == j or i_name in ignore_links:
+                        continue
+                    if i_name not in collision_link_names:
+                        log_error("Self Collision Link name not found in collision_link_names")
+                    # find index of this link name:
+                    if i_name not in self.self_collision_buffer.keys():
+                        self.self_collision_buffer[i_name] = 0.0
+                    c2 = self.self_collision_buffer[i_name]
+                    link2_idx = self._name_to_idx_map[i_name]
+                    # update collision distance between spheres from these two links:
+                    link2_spheres_idx = torch.nonzero(self._link_sphere_idx_map == link2_idx)
+                    rad2 = self._link_spheres_tensor[link2_spheres_idx, 3]
+
+                    for k1 in range(len(rad1)):
+                        sp1 = link1_spheres_idx[k1]
+                        for k2 in range(len(rad2)):
+                            sp2 = link2_spheres_idx[k2]
+                            self.self_collision_distance[sp1, sp2] = rad1[k1] + rad2[k2] + c1 + c2
+
+        with profiler.record_function("robot_generator/self_collision_min"):
+            d_mat = self.self_collision_distance
+            self.self_collision_distance = torch.minimum(d_mat, d_mat.transpose(0, 1))
+
+        if self.debug is not None and "self_collision_experimental" in self.debug:
+            use_experimental_kernel = self.debug["self_collision_experimental"]
+        self.self_collision_distance = self.self_collision_distance.to(
+            device=self.tensor_args.device
+        )
+        (
+            self._self_coll_thread_locations,
+            self._self_coll_idx,
+            valid_data,
+            checks_per_thread,
+        ) = self._create_self_collision_thread_data(self.self_collision_distance)
+        self._self_coll_matrix = (self.self_collision_distance != -(torch.inf)).to(
+            dtype=torch.uint8
+        )
+
+        use_experimental_kernel = True
+        # convert all tensors to gpu:
+        self._link_sphere_idx_map = self._link_sphere_idx_map.to(device=self.tensor_args.device)
+        self._link_spheres_tensor = self._link_spheres_tensor.to(device=self.tensor_args.device)
+        self.self_collision_offset = self.self_collision_offset.to(device=self.tensor_args.device)
+        self._self_collision_data = SelfCollisionKinematicsConfig(
+            offset=self.self_collision_offset,
+            distance_threshold=self.self_collision_distance,
+            thread_location=self._self_coll_thread_locations,
+            thread_max=self._self_coll_idx,
+            collision_matrix=self._self_coll_matrix,
+            experimental_kernel=valid_data and use_experimental_kernel,
+            checks_per_thread=checks_per_thread,
+        )
+
+    @profiler.record_function("robot_generator/create_self_collision_thread_data")
+    def _create_self_collision_thread_data(self, collision_threshold):
+        coll_cpu = collision_threshold.cpu()
+        max_checks_per_thread = 512
+        thread_loc = torch.zeros((2 * 32 * max_checks_per_thread), dtype=torch.int16) - 1
+        n_spheres = coll_cpu.shape[0]
+        sl_idx = 0
+        skip_count = 0
+        all_val = 0
+        valid_data = True
+        for i in range(n_spheres):
+            if not valid_data:
+                break
+            if torch.max(coll_cpu[i]) == -torch.inf:
+                print("skip", i)
+            for j in range(i + 1, n_spheres):
+                if sl_idx > thread_loc.shape[0] - 1:
+                    valid_data = False
+                    log_warn(
+                        "Self Collision checks are greater than "
+                        + str(32 * max_checks_per_thread)
+                        + ", using slower kernel"
+                    )
+                    break
+                if coll_cpu[i, j] != -torch.inf:
+                    thread_loc[sl_idx] = i
+                    sl_idx += 1
+                    thread_loc[sl_idx] = j
+                    sl_idx += 1
+                else:
+                    skip_count += 1
+                all_val += 1
+        log_info("Self Collision threads, skipped %: " + str(100 * float(skip_count) / all_val))
+        log_info("Self Collision count: " + str(sl_idx / (2)))
+        log_info("Self Collision per thread: " + str(sl_idx / (2 * 1024)))
+
+        max_checks_per_thread = 512
+        val = sl_idx / (2 * 1024)
+        if val < 1:
+            max_checks_per_thread = 1
+        elif val < 2:
+            max_checks_per_thread = 2
+        elif val < 4:
+            max_checks_per_thread = 4
+        elif val < 8:
+            max_checks_per_thread = 8
+        elif val < 32:
+            max_checks_per_thread = 32
+        elif val < 64:
+            max_checks_per_thread = 64
+        elif val < 128:
+            max_checks_per_thread = 128
+        elif val < 512:
+            max_checks_per_thread = 512
+        else:
+            raise ValueError("Self Collision not supported")
+
+        if max_checks_per_thread < 2:
+            max_checks_per_thread = 2
+        log_info("Self Collision using: " + str(max_checks_per_thread))
+
+        return (
+            thread_loc.to(device=collision_threshold.device),
+            sl_idx,
+            valid_data,
+            max_checks_per_thread,
+        )
+
+    @profiler.record_function("robot_generator/add_body_to_tree")
+    def _add_body_to_tree(self, link_name, base=False):
+        body_idx = len(self._bodies)
+
+        rigid_body_params = self._kinematics_parser.get_link_parameters(link_name, base=base)
+        self._bodies.append(rigid_body_params)
+        if rigid_body_params.joint_type != JointType.FIXED:
+            self._controlled_joints.append(body_idx)
+            self.joint_names.append(rigid_body_params.joint_name)
+            self._n_dofs = self._n_dofs + 1
+        self._fixed_transform.append(
+            torch.as_tensor(
+                rigid_body_params.fixed_transform,
+                device=self.cpu_tensor_args.device,
+                dtype=self.cpu_tensor_args.dtype,
+            ).unsqueeze(0)
+        )
+        self._name_to_idx_map[rigid_body_params.link_name] = body_idx
+
+    def _get_joint_links(self, joint_names: List[str]):
+        j_data = {}
+        for j in joint_names:
+            for b in self._bodies:
+                if b.joint_name == j:
+                    j_data[j] = {"parent": b.parent_link_name, "child": b.link_name}
+        return j_data
+
+    @profiler.record_function("robot_generator/get_link_poses")
+    def _get_link_poses(
+        self, q: torch.Tensor, link_names: List[str], kinematics_config: KinematicsTensorConfig
+    ) -> Pose:
+        if q.is_contiguous():
+            q_in = q.view(-1)
+        else:
+            q_in = q.reshape(-1)  # .reshape(-1)
+        # q_in = q.reshape(-1)
+        link_pos_seq = torch.zeros(
+            (1, len(self.link_names), 3),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+        link_quat_seq = torch.zeros(
+            (1, len(self.link_names), 4),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+        batch_robot_spheres = torch.zeros(
+            (1, self.total_spheres, 4),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+        grad_out_q = torch.zeros(
+            (1 * q.shape[-1]),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+        global_cumul_mat = torch.zeros(
+            (1, self._link_map.shape[0], 4, 4),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+        link_pos_seq, link_quat_seq, _ = get_cuda_kinematics(
+            # self._link_mat_seq,  # data will be stored here
+            link_pos_seq,
+            link_quat_seq,
+            batch_robot_spheres,
+            global_cumul_mat,
+            q_in,
+            kinematics_config.fixed_transforms,
+            kinematics_config.link_spheres,
+            kinematics_config.link_map,  # tells which link is attached to which link i
+            kinematics_config.joint_map,  # tells which joint is attached to a link i
+            kinematics_config.joint_map_type,  # joint type
+            kinematics_config.store_link_map,
+            kinematics_config.link_sphere_idx_map,  # sphere idx map
+            kinematics_config.link_chain_map,
+            grad_out_q,
+            self.use_global_cumul,
+        )
+        position = torch.zeros(
+            (q.shape[0], len(link_names), 3),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+        quaternion = torch.zeros(
+            (q.shape[0], len(link_names), 4),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+
+        for li, l in enumerate(link_names):
+            i = self.link_names.index(l)
+            position[:, li, :] = link_pos_seq[:, i, :]
+            quaternion[:, li, :] = link_quat_seq[:, i, :]
+        return Pose(position=position, quaternion=quaternion)
+
+    @property
+    def get_joint_limits(self):
+        return self._joint_limits
+
+    @profiler.record_function("robot_generator/get_joint_limits")
+    def _get_joint_position_velocity_limits(self):
+        joint_limits = {"position": [[], []], "velocity": [[], []]}
+
+        for idx in self._controlled_joints:
+            joint_limits["position"][0].append(self._bodies[idx].joint_limits[0])
+            joint_limits["position"][1].append(self._bodies[idx].joint_limits[1])
+            joint_limits["velocity"][0].append(self._bodies[idx].joint_velocity_limits[0])
+            joint_limits["velocity"][1].append(self._bodies[idx].joint_velocity_limits[1])
+        for k in joint_limits:
+            joint_limits[k] = torch.as_tensor(
+                joint_limits[k], device=self.tensor_args.device, dtype=self.tensor_args.dtype
+            )
+        return joint_limits
+
+    @profiler.record_function("robot_generator/update_joint_limits")
+    def _update_joint_limits(self):
+        joint_limits = self._get_joint_position_velocity_limits()
+        joint_limits["jerk"] = torch.cat(
+            [-1.0 * self.cspace.max_jerk.unsqueeze(0), self.cspace.max_jerk.unsqueeze(0)]
+        )
+        joint_limits["acceleration"] = torch.cat(
+            [
+                -1.0 * self.cspace.max_acceleration.unsqueeze(0),
+                self.cspace.max_acceleration.unsqueeze(0),
+            ]
+        )
+        self._joint_limits = JointLimits(joint_names=self.joint_names, **joint_limits)

src/curobo/cuda_robot_model/cuda_robot_model.py

@@ -0,0 +1,386 @@
+#
+# 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
+from typing import Any, Dict, List, Optional
+
+# Third Party
+import torch
+import torch.autograd.profiler as profiler
+
+# CuRobo
+from curobo.cuda_robot_model.cuda_robot_generator import (
+    CudaRobotGenerator,
+    CudaRobotGeneratorConfig,
+)
+from curobo.cuda_robot_model.kinematics_parser import KinematicsParser
+from curobo.cuda_robot_model.types import (
+    CudaRobotModelState,
+    KinematicsTensorConfig,
+    SelfCollisionKinematicsConfig,
+)
+from curobo.curobolib.kinematics import get_cuda_kinematics
+from curobo.geom.types import Sphere
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util.logger import log_error
+from curobo.util_file import get_robot_path, join_path, load_yaml
+
+
+@dataclass
+class CudaRobotModelConfig:
+    tensor_args: TensorDeviceType
+    link_names: List[str]
+    kinematics_config: KinematicsTensorConfig
+    self_collision_config: Optional[SelfCollisionKinematicsConfig] = None
+    kinematics_parser: Optional[KinematicsParser] = None
+    compute_jacobian: bool = False
+    use_global_cumul: bool = False
+    generator_config: Optional[CudaRobotGeneratorConfig] = None
+
+    def get_joint_limits(self):
+        return self.kinematics_config.joint_limits
+
+    @staticmethod
+    def from_basic_urdf(
+        urdf_path: str,
+        base_link: str,
+        ee_link: str,
+        tensor_args: TensorDeviceType = TensorDeviceType(),
+    ) -> CudaRobotModelConfig:
+        """Load a cuda robot model from only urdf. This does not support collision queries.
+
+        Args:
+            urdf_path : Path of urdf file.
+            base_link : Name of base link.
+            ee_link : Name of end-effector link.
+            tensor_args : Device to load robot model. Defaults to TensorDeviceType().
+
+        Returns:
+            CudaRobotModelConfig: cuda robot model configuration.
+        """
+        config = CudaRobotGeneratorConfig(base_link, ee_link, tensor_args, urdf_path=urdf_path)
+        return CudaRobotModelConfig.from_config(config)
+
+    @staticmethod
+    def from_basic_usd(
+        usd_path: str,
+        usd_robot_root: str,
+        base_link: str,
+        ee_link: str,
+        tensor_args: TensorDeviceType = TensorDeviceType(),
+    ) -> CudaRobotModelConfig:
+        """Load a cuda robot model from only urdf. This does not support collision queries.
+
+        Args:
+            urdf_path : Path of urdf file.
+            base_link : Name of base link.
+            ee_link : Name of end-effector link.
+            tensor_args : Device to load robot model. Defaults to TensorDeviceType().
+
+        Returns:
+            CudaRobotModelConfig: cuda robot model configuration.
+        """
+        config = CudaRobotGeneratorConfig(
+            tensor_args,
+            base_link,
+            ee_link,
+            usd_path=usd_path,
+            usd_robot_root=usd_robot_root,
+            use_usd_kinematics=True,
+        )
+        return CudaRobotModelConfig.from_config(config)
+
+    @staticmethod
+    def from_robot_yaml_file(
+        file_path: str,
+        ee_link: Optional[str] = None,
+        tensor_args: TensorDeviceType = TensorDeviceType(),
+    ):
+        config_file = load_yaml(join_path(get_robot_path(), file_path))["robot_cfg"]["kinematics"]
+        if ee_link is not None:
+            config_file["ee_link"] = ee_link
+        return CudaRobotModelConfig.from_config(
+            CudaRobotGeneratorConfig(**config_file, tensor_args=tensor_args)
+        )
+
+    @staticmethod
+    def from_data_dict(
+        data_dict: Dict[str, Any], tensor_args: TensorDeviceType = TensorDeviceType()
+    ):
+        return CudaRobotModelConfig.from_config(
+            CudaRobotGeneratorConfig(**data_dict, tensor_args=tensor_args)
+        )
+
+    @staticmethod
+    def from_config(config: CudaRobotGeneratorConfig):
+        # create a config generator and load all values
+        generator = CudaRobotGenerator(config)
+        return CudaRobotModelConfig(
+            tensor_args=generator.tensor_args,
+            link_names=generator.link_names,
+            kinematics_config=generator.kinematics_config,
+            self_collision_config=generator.self_collision_config,
+            kinematics_parser=generator.kinematics_parser,
+            use_global_cumul=generator.use_global_cumul,
+            compute_jacobian=generator.compute_jacobian,
+            generator_config=config,
+        )
+
+    @property
+    def cspace(self):
+        return self.kinematics_config.cspace
+
+
+class CudaRobotModel(CudaRobotModelConfig):
+    """
+    CUDA Accelerated Robot Model
+
+    NOTE: Currently dof is created only for links that we need to compute kinematics.
+    E.g., for robots with many serial chains, add all links of the robot to get the correct dof.
+    This is not an issue if you are loading collision spheres as that will cover the full geometry
+    of the robot.
+    """
+
+    def __init__(self, config: CudaRobotModelConfig):
+        super().__init__(**vars(config))
+        self._batch_size = 0
+        self.update_batch_size(1, reset_buffers=True)
+
+    def update_batch_size(self, batch_size, force_update=False, reset_buffers=False):
+        if batch_size == 0:
+            log_error("batch size is zero")
+        if force_update and self._batch_size == batch_size and self.compute_jacobian:
+            self.lin_jac = self.lin_jac.detach()  # .requires_grad_(True)
+            self.ang_jac = self.ang_jac.detach()  # .requires_grad_(True)
+        elif self._batch_size != batch_size or reset_buffers:
+            self._batch_size = batch_size
+            self._link_pos_seq = torch.zeros(
+                (self._batch_size, len(self.link_names), 3),
+                device=self.tensor_args.device,
+                dtype=self.tensor_args.dtype,
+            )
+            self._link_quat_seq = torch.zeros(
+                (self._batch_size, len(self.link_names), 4),
+                device=self.tensor_args.device,
+                dtype=self.tensor_args.dtype,
+            )
+
+            self._batch_robot_spheres = torch.zeros(
+                (self._batch_size, self.kinematics_config.total_spheres, 4),
+                device=self.tensor_args.device,
+                dtype=self.tensor_args.dtype,
+            )
+            self._grad_out_q = torch.zeros(
+                (self._batch_size, self.get_dof()),
+                device=self.tensor_args.device,
+                dtype=self.tensor_args.dtype,
+            )
+            self._global_cumul_mat = torch.zeros(
+                (self._batch_size, self.kinematics_config.link_map.shape[0], 4, 4),
+                device=self.tensor_args.device,
+                dtype=self.tensor_args.dtype,
+            )
+            if self.compute_jacobian:
+                self.lin_jac = torch.zeros(
+                    [batch_size, 3, self.kinematics_config.n_dofs],
+                    device=self.tensor_args.device,
+                    dtype=self.tensor_args.dtype,
+                )
+                self.ang_jac = torch.zeros(
+                    [batch_size, 3, self.kinematics_config.n_dofs],
+                    device=self.tensor_args.device,
+                    dtype=self.tensor_args.dtype,
+                )
+
+    @profiler.record_function("cuda_robot_model/forward_kinematics")
+    def forward(self, q, link_name=None, calculate_jacobian=False):
+        # pos, rot = self.compute_forward_kinematics(q, qd, link_name)
+        if len(q.shape) > 2:
+            raise ValueError("q shape should be [batch_size, dof]")
+        batch_size = q.shape[0]
+        self.update_batch_size(batch_size, force_update=q.requires_grad)
+
+        # do fused forward:
+        link_pos_seq, link_quat_seq, link_spheres_tensor = self._cuda_forward(q)
+
+        if len(self.link_names) == 1:
+            ee_pos = link_pos_seq.squeeze(1)
+            ee_quat = link_quat_seq.squeeze(1)
+        else:
+            link_idx = self.kinematics_config.ee_idx
+            if link_name is not None:
+                link_idx = self.link_names.index(link_name)
+            ee_pos = link_pos_seq.contiguous()[..., link_idx, :]
+            ee_quat = link_quat_seq.contiguous()[..., link_idx, :]
+        lin_jac = ang_jac = None
+
+        # compute jacobians?
+        if calculate_jacobian:
+            raise NotImplementedError
+        return (
+            ee_pos,
+            ee_quat,
+            lin_jac,
+            ang_jac,
+            link_pos_seq,
+            link_quat_seq,
+            link_spheres_tensor,
+        )
+
+    def get_state(self, q, link_name=None, calculate_jacobian=False) -> CudaRobotModelState:
+        out = self.forward(q, link_name, calculate_jacobian)
+        state = CudaRobotModelState(
+            out[0],
+            out[1],
+            None,
+            None,
+            out[4],
+            out[5],
+            out[6],
+            self.link_names,
+        )
+        return state
+
+    def get_robot_as_mesh(self, q: torch.Tensor):
+        # get all link meshes:
+        m_list = [self.get_link_mesh(l) for l in self.mesh_link_names]
+        pose = self.get_link_poses(q, self.mesh_link_names)
+        for li, l in enumerate(self.mesh_link_names):
+            m_list[li].pose = pose.get_index(0, li).tolist()
+
+        return m_list
+
+    def get_robot_as_spheres(self, q: torch.Tensor, filter_valid: bool = True):
+        state = self.get_state(q)
+
+        # state has sphere position and radius
+
+        sph_all = state.get_link_spheres().cpu().numpy()
+
+        sph_traj = []
+        for j in range(sph_all.shape[0]):
+            sph = sph_all[j, :, :]
+            if filter_valid:
+                sph_list = [
+                    Sphere(
+                        name="robot_curobo_sphere_" + str(i),
+                        pose=[sph[i, 0], sph[i, 1], sph[i, 2], 1, 0, 0, 0],
+                        radius=sph[i, 3],
+                    )
+                    for i in range(sph.shape[0])
+                    if (sph[i, 3] > 0.0)
+                ]
+            else:
+                sph_list = [
+                    Sphere(
+                        name="robot_curobo_sphere_" + str(i),
+                        pose=[sph[i, 0], sph[i, 1], sph[i, 2], 1, 0, 0, 0],
+                        radius=sph[i, 3],
+                    )
+                    for i in range(sph.shape[0])
+                ]
+            sph_traj.append(sph_list)
+        return sph_traj
+
+    def get_link_poses(self, q: torch.Tensor, link_names: List[str]) -> Pose:
+        state = self.get_state(q)
+        position = torch.zeros(
+            (q.shape[0], len(link_names), 3),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+        quaternion = torch.zeros(
+            (q.shape[0], len(link_names), 4),
+            device=self.tensor_args.device,
+            dtype=self.tensor_args.dtype,
+        )
+
+        for li, l in enumerate(link_names):
+            i = self.link_names.index(l)
+            position[:, li, :] = state.links_position[:, i, :]
+            quaternion[:, li, :] = state.links_quaternion[:, i, :]
+        return Pose(position=position, quaternion=quaternion)
+
+    def _cuda_forward(self, q):
+        link_pos, link_quat, robot_spheres = get_cuda_kinematics(
+            # self._link_mat_seq,  # data will be stored here
+            self._link_pos_seq,
+            self._link_quat_seq,
+            self._batch_robot_spheres,
+            self._global_cumul_mat,
+            q,
+            self.kinematics_config.fixed_transforms,
+            self.kinematics_config.link_spheres,
+            self.kinematics_config.link_map,  # tells which link is attached to which link i
+            self.kinematics_config.joint_map,  # tells which joint is attached to a link i
+            self.kinematics_config.joint_map_type,  # joint type
+            self.kinematics_config.store_link_map,
+            self.kinematics_config.link_sphere_idx_map,  # sphere idx map
+            self.kinematics_config.link_chain_map,
+            self._grad_out_q,
+            self.use_global_cumul,
+        )
+        # if(robot_spheres.shape[0]<10):
+        #    print(robot_spheres)
+        return link_pos, link_quat, robot_spheres
+
+    @property
+    def all_articulated_joint_names(self):
+        return self.kinematics_config.non_fixed_joint_names
+
+    def get_self_collision_config(self) -> SelfCollisionKinematicsConfig:
+        return self.self_collision_config
+
+    def get_link_mesh(self, link_name: str) -> Mesh:
+        mesh = self.kinematics_parser.get_link_mesh(link_name)
+        mesh.pose = [0, 0, 0, 1, 0, 0, 0]
+        return mesh
+
+    def get_link_transform(self, link_name: str) -> Pose:
+        mat = self._kinematics_config.fixed_transforms[self._name_to_idx_map[link_name]]
+        pose = Pose(position=mat[:3, 3], rotation=mat[:3, :3])
+        return pose
+
+    def get_all_link_transforms(self) -> Pose:
+        pose = Pose(
+            self.kinematics_config.fixed_transforms[:, :3, 3],
+            rotation=self.kinematics_config.fixed_transforms[:, :3, :3],
+        )
+        return pose
+
+    def get_dof(self) -> int:
+        return self.kinematics_config.n_dof
+
+    @property
+    def joint_names(self) -> List[str]:
+        return self.kinematics_config.joint_names
+
+    @property
+    def total_spheres(self) -> int:
+        return self.kinematics_config.total_spheres
+
+    @property
+    def lock_jointstate(self):
+        return self.kinematics_config.lock_jointstate
+
+    @property
+    def ee_link(self):
+        return self.kinematics_config.ee_link
+
+    @property
+    def base_link(self):
+        return self.kinematics_config.base_link
+
+    def update_kinematics_config(self, new_kin_config: KinematicsTensorConfig):
+        self.kinematics_config.copy(new_kin_config)

src/curobo/cuda_robot_model/kinematics_parser.py

@@ -0,0 +1,101 @@
+#
+# 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, field
+from typing import Dict, List, Optional, Union
+
+# Third Party
+import numpy as np
+
+# CuRobo
+from curobo.cuda_robot_model.types import JointType
+from curobo.geom.types import Mesh
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util.logger import log_error
+
+
+@dataclass
+class LinkParams:
+    link_name: str
+    joint_name: str
+    joint_type: JointType
+    fixed_transform: np.ndarray
+    parent_link_name: Optional[str] = None
+    joint_limits: Optional[List[float]] = None
+    joint_axis: Optional[np.ndarray] = None
+    joint_id: Optional[int] = None
+    joint_velocity_limits: List[float] = field(default_factory=lambda: [-2.0, 2.0])
+
+    @staticmethod
+    def from_dict(dict_data):
+        dict_data["joint_type"] = JointType[dict_data["joint_type"]]
+        dict_data["fixed_transform"] = (
+            Pose.from_list(dict_data["fixed_transform"], tensor_args=TensorDeviceType())
+            .get_numpy_matrix()
+            .reshape(4, 4)
+        )
+
+        return LinkParams(**dict_data)
+
+
+class KinematicsParser:
+    def __init__(self, extra_links: Optional[Dict[str, LinkParams]] = None) -> None:
+        self._parent_map = {}
+        self.extra_links = extra_links
+        self.build_link_parent()
+        # add extra links to parent:
+        if self.extra_links is not None and len(list(self.extra_links.keys())) > 0:
+            for i in self.extra_links:
+                self._parent_map[i] = self.extra_links[i].parent_link_name
+
+    def build_link_parent(self):
+        """Build a map of parent links to each link in the kinematic tree.
+
+        NOTE: Use this function to fill self._parent_map.
+        """
+        pass
+
+    def get_chain(self, base_link: str, ee_link: str) -> List[str]:
+        """Get list of links attaching ee_link to base_link.
+
+        Args:
+            base_link (str): Name of base link.
+            ee_link (str): Name of end-effector link.
+
+        Returns:
+            List[str]: List of link names starting from base_link to ee_link.
+        """
+        chain_links = [ee_link]
+        link = ee_link
+        while link != base_link:
+            link = self._parent_map[link]
+            # add link to chain:
+            chain_links.append(link)
+        chain_links.reverse()
+        return chain_links
+
+    def _get_from_extra_links(self, link_name: str) -> LinkParams:
+        if self.extra_links is None:
+            return None
+        if link_name in self.extra_links.keys():
+            return self.extra_links[link_name]
+        return None
+
+    def get_link_parameters(self, link_name: str, base: bool = False) -> LinkParams:
+        pass
+
+    def add_absolute_path_to_link_meshes(self, mesh_dir: str = ""):
+        pass
+
+    def get_link_mesh(self, link_name: str) -> Mesh:
+        pass

src/curobo/cuda_robot_model/types.py

@@ -0,0 +1,388 @@
+#
+# 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
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+# Third Party
+import torch
+
+# CuRobo
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.state import JointState
+from curobo.types.tensor import T_DOF
+from curobo.util.tensor_util import copy_if_not_none
+
+
+class JointType(Enum):
+    FIXED = -1
+    X_PRISM = 0
+    Y_PRISM = 1
+    Z_PRISM = 2
+    X_ROT = 3
+    Y_ROT = 4
+    Z_ROT = 5
+
+
+@dataclass
+class JointLimits:
+    joint_names: List[str]
+    position: torch.Tensor
+    velocity: torch.Tensor
+    acceleration: torch.Tensor
+    jerk: torch.Tensor
+    effort: Optional[torch.Tensor] = None
+    tensor_args: TensorDeviceType = TensorDeviceType()
+
+    @staticmethod
+    def from_data_dict(data: Dict, tensor_args: TensorDeviceType = TensorDeviceType()):
+        p = tensor_args.to_device(data["position"])
+        v = tensor_args.to_device(data["velocity"])
+        a = tensor_args.to_device(data["acceleration"])
+        j = tensor_args.to_device(data["jerk"])
+        e = None
+        if "effort" in data and data["effort"] is not None:
+            e = tensor_args.to_device(data["effort"])
+
+        return JointLimits(data["joint_names"], p, v, a, j, e)
+
+    def clone(self) -> JointLimits:
+        return JointLimits(
+            self.joint_names.copy(),
+            self.position.clone(),
+            self.velocity.clone(),
+            self.acceleration.clone(),
+            self.jerk.clone(),
+            self.effort.clone() if self.effort is not None else None,
+            self.tensor_args,
+        )
+
+
+@dataclass
+class CSpaceConfig:
+    joint_names: List[str]
+    retract_config: Optional[T_DOF] = None
+    cspace_distance_weight: Optional[T_DOF] = None
+    null_space_weight: Optional[T_DOF] = None  # List[float]
+    tensor_args: TensorDeviceType = TensorDeviceType()
+    max_acceleration: Union[float, List[float]] = 10.0
+    max_jerk: Union[float, List[float]] = 500.0
+    velocity_scale: Union[float, List[float]] = 1.0
+    acceleration_scale: Union[float, List[float]] = 1.0
+    jerk_scale: Union[float, List[float]] = 1.0
+
+    def __post_init__(self):
+        if self.retract_config is not None:
+            self.retract_config = self.tensor_args.to_device(self.retract_config)
+        if self.cspace_distance_weight is not None:
+            self.cspace_distance_weight = self.tensor_args.to_device(self.cspace_distance_weight)
+        if self.null_space_weight is not None:
+            self.null_space_weight = self.tensor_args.to_device(self.null_space_weight)
+        if isinstance(self.max_acceleration, float):
+            self.max_acceleration = self.tensor_args.to_device(
+                [self.max_acceleration for _ in self.joint_names]
+            )
+
+        if isinstance(self.velocity_scale, float) or len(self.velocity_scale) == 1:
+            self.velocity_scale = self.tensor_args.to_device(
+                [self.velocity_scale for _ in self.joint_names]
+            ).view(-1)
+
+        if isinstance(self.acceleration_scale, float) or len(self.acceleration_scale) == 1:
+            self.acceleration_scale = self.tensor_args.to_device(
+                [self.acceleration_scale for _ in self.joint_names]
+            ).view(-1)
+
+        if isinstance(self.jerk_scale, float) or len(self.jerk_scale) == 1:
+            self.jerk_scale = self.tensor_args.to_device(
+                [self.jerk_scale for _ in self.joint_names]
+            ).view(-1)
+
+        if isinstance(self.max_acceleration, List):
+            self.max_acceleration = self.tensor_args.to_device(self.max_acceleration)
+        if isinstance(self.max_jerk, float):
+            self.max_jerk = [self.max_jerk for _ in self.joint_names]
+        if isinstance(self.max_jerk, List):
+            self.max_jerk = self.tensor_args.to_device(self.max_jerk)
+        if isinstance(self.velocity_scale, List):
+            self.velocity_scale = self.tensor_args.to_device(self.velocity_scale)
+
+        if isinstance(self.acceleration_scale, List):
+            self.acceleration_scale = self.tensor_args.to_device(self.acceleration_scale)
+        if isinstance(self.jerk_scale, List):
+            self.jerk_scale = self.tensor_args.to_device(self.jerk_scale)
+
+    def inplace_reindex(self, joint_names: List[str]):
+        new_index = [self.joint_names.index(j) for j in joint_names]
+        if self.retract_config is not None:
+            self.retract_config = self.retract_config[new_index].clone()
+        if self.cspace_distance_weight is not None:
+            self.cspace_distance_weight = self.cspace_distance_weight[new_index].clone()
+        if self.null_space_weight is not None:
+            self.null_space_weight = self.null_space_weight[new_index].clone()
+        self.max_acceleration = self.max_acceleration[new_index].clone()
+        self.max_jerk = self.max_jerk[new_index].clone()
+        self.velocity_scale = self.velocity_scale[new_index].clone()
+        self.acceleration_scale = self.acceleration_scale[new_index].clone()
+        self.jerk_scale = self.jerk_scale[new_index].clone()
+        joint_names = [self.joint_names[n] for n in new_index]
+        self.joint_names = joint_names
+
+    def clone(self) -> CSpaceConfig:
+        return CSpaceConfig(
+            joint_names=self.joint_names.copy(),
+            retract_config=copy_if_not_none(self.retract_config),
+            null_space_weight=copy_if_not_none(self.null_space_weight),
+            cspace_distance_weight=copy_if_not_none(self.cspace_distance_weight),
+            tensor_args=self.tensor_args,
+            max_jerk=self.max_jerk.clone(),
+            max_acceleration=self.max_acceleration.clone(),
+            velocity_scale=self.velocity_scale.clone(),
+            acceleration_scale=self.acceleration_scale.clone(),
+            jerk_scale=self.jerk_scale.clone(),
+        )
+
+    def scale_joint_limits(self, joint_limits: JointLimits):
+        if self.velocity_scale is not None:
+            joint_limits.velocity = joint_limits.velocity * self.velocity_scale
+        if self.acceleration_scale is not None:
+            joint_limits.acceleration = joint_limits.acceleration * self.acceleration_scale
+        if self.jerk_scale is not None:
+            joint_limits.jerk = joint_limits.jerk * self.jerk_scale
+
+        return joint_limits
+
+    @staticmethod
+    def load_from_joint_limits(
+        joint_position_upper: torch.Tensor,
+        joint_position_lower: torch.Tensor,
+        joint_names: List[str],
+        tensor_args: TensorDeviceType = TensorDeviceType(),
+    ):
+        retract_config = ((joint_position_upper + joint_position_lower) / 2).flatten()
+        n_dof = retract_config.shape[-1]
+        null_space_weight = torch.ones(n_dof, **vars(tensor_args))
+        cspace_distance_weight = torch.ones(n_dof, **vars(tensor_args))
+        return CSpaceConfig(
+            joint_names,
+            retract_config,
+            cspace_distance_weight,
+            null_space_weight,
+            tensor_args=tensor_args,
+        )
+
+
+@dataclass
+class KinematicsTensorConfig:
+    fixed_transforms: torch.Tensor
+    link_map: torch.Tensor
+    joint_map: torch.Tensor
+    joint_map_type: torch.Tensor
+    store_link_map: torch.Tensor
+    link_chain_map: torch.Tensor
+    link_names: List[str]
+    joint_limits: JointLimits
+    non_fixed_joint_names: List[str]
+    n_dof: int
+    mesh_link_names: Optional[List[str]] = None
+    joint_names: Optional[List[str]] = None
+    lock_jointstate: Optional[JointState] = None
+    link_spheres: Optional[torch.Tensor] = None
+    link_sphere_idx_map: Optional[torch.Tensor] = None
+    link_name_to_idx_map: Optional[Dict[str, int]] = None
+    total_spheres: int = 0
+    debug: Optional[Any] = None
+    ee_idx: int = 0
+    cspace: Optional[CSpaceConfig] = None
+    base_link: str = "base_link"
+    ee_link: str = "ee_link"
+
+    def __post_init__(self):
+        if self.cspace is None and self.joint_limits is not None:
+            self.load_cspace_cfg_from_kinematics()
+        if self.joint_limits is not None and self.cspace is not None:
+            self.joint_limits = self.cspace.scale_joint_limits(self.joint_limits)
+
+    def load_cspace_cfg_from_kinematics(self):
+        retract_config = (
+            (self.joint_limits.position[1] + self.joint_limits.position[0]) / 2
+        ).flatten()
+        null_space_weight = torch.ones(self.n_dof, **vars(self.tensor_args))
+        cspace_distance_weight = torch.ones(self.n_dof, **vars(self.tensor_args))
+        joint_names = self.joint_names
+        self.cspace = CSpaceConfig(
+            joint_names,
+            retract_config,
+            cspace_distance_weight,
+            null_space_weight,
+            tensor_args=self.tensor_args,
+            max_acceleration=self.joint_limits.acceleration[1],
+            max_jerk=self.joint_limits.max_jerk[1],
+        )
+
+    def get_sphere_index_from_link_name(self, link_name: str) -> torch.Tensor:
+        link_idx = self.link_name_to_idx_map[link_name]
+        # get sphere index for this link:
+        link_spheres_idx = torch.nonzero(self.link_sphere_idx_map == link_idx).view(-1)
+        return link_spheres_idx
+
+    def update_link_spheres(
+        self, link_name: str, sphere_position_radius: torch.Tensor, start_sph_idx: int = 0
+    ):
+        """Update sphere parameters
+
+        #NOTE: This currently does not update self collision distances.
+        Args:
+            link_name: _description_
+            sphere_position_radius: _description_
+            start_sph_idx: _description_. Defaults to 0.
+        """
+        # get sphere indices from link name:
+        link_sphere_index = self.get_sphere_index_from_link_name(link_name)[
+            start_sph_idx : start_sph_idx + sphere_position_radius.shape[0]
+        ]
+        # update sphere data:
+        self.link_spheres[link_sphere_index, :] = sphere_position_radius
+
+    def get_link_spheres(
+        self,
+        link_name: str,
+    ):
+        link_sphere_index = self.get_sphere_index_from_link_name(link_name)
+        return self.link_spheres[link_sphere_index, :]
+
+    def attach_object(
+        self,
+        sphere_radius: Optional[float] = None,
+        sphere_tensor: Optional[torch.Tensor] = None,
+        link_name: str = "attached_object",
+    ) -> bool:
+        """_summary_
+
+        Args:
+            sphere_radius: _description_. Defaults to None.
+            sphere_tensor: _description_. Defaults to None.
+            link_name: _description_. Defaults to "attached_object".
+
+        Raises:
+            ValueError: _description_
+            ValueError: _description_
+
+        Returns:
+            _description_
+        """
+        if link_name not in self.link_name_to_idx_map.keys():
+            raise ValueError(link_name + " not found in spheres")
+        curr_spheres = self.get_link_spheres(link_name)
+
+        if sphere_radius is not None:
+            curr_spheres[:, 3] = sphere_radius
+        if sphere_tensor is not None:
+            if sphere_tensor.shape != curr_spheres.shape and sphere_tensor.shape[0] != 1:
+                raise ValueError("sphere_tensor shape does not match current spheres")
+            curr_spheres[:, :] = sphere_tensor
+        self.update_link_spheres(link_name, curr_spheres)
+        return True
+
+    def detach_object(self, link_name: str = "attached_object") -> bool:
+        """Detach object from robot end-effector
+
+        Args:
+            link_name: _description_. Defaults to "attached_object".
+
+        Raises:
+            ValueError: _description_
+
+        Returns:
+            _description_
+        """
+        if link_name not in self.link_name_to_idx_map.keys():
+            raise ValueError(link_name + " not found in spheres")
+        curr_spheres = self.get_link_spheres(link_name)
+        curr_spheres[:, 3] = -100.0
+        curr_spheres[:, :3] = 0.0
+        self.update_link_spheres(link_name, curr_spheres)
+
+        return True
+
+    def get_number_of_spheres(self, link_name: str) -> int:
+        """Get number of spheres for a link
+
+        Args:
+            link_name: name of link
+        """
+        return self.get_link_spheres(link_name).shape[0]
+
+
+@dataclass
+class SelfCollisionKinematicsConfig:
+    """Dataclass that stores self collision attributes to pass to cuda kernel."""
+
+    offset: Optional[torch.Tensor] = None
+    distance_threshold: Optional[torch.Tensor] = None
+    thread_location: Optional[torch.Tensor] = None
+    thread_max: Optional[int] = None
+    collision_matrix: Optional[torch.Tensor] = None
+    experimental_kernel: bool = True
+    checks_per_thread: int = 32
+
+
+@dataclass(frozen=True)
+class CudaRobotModelState:
+    """Dataclass that stores kinematics information."""
+
+    #: End-effector position stored as x,y,z in meters [b, 3]. End-effector is defined by
+    #: :py:attr:`curobo.cuda_robot_model.cuda_robot_generator.CudaRobotGeneratorConfig.ee_link`.
+    ee_position: torch.Tensor
+
+    #: End-effector orientaiton stored as quaternion qw, qx, qy, qz [b,4]. End-effector is defined
+    # by :py:attr:`CudaRobotModelConfig.ee_link`.
+    ee_quaternion: torch.Tensor
+
+    #: Linear Jacobian. Currently not supported.
+    lin_jacobian: Optional[torch.Tensor] = None
+
+    #: Angular Jacobian. Currently not supported.
+    ang_jacobian: Optional[torch.Tensor] = None
+
+    #: Position of links specified by link_names  (:py:attr:`CudaRobotModelConfig.link_names`).
+    links_position: Optional[torch.Tensor] = None
+
+    #: Quaternions of links specified by link names (:py:attr:`CudaRobotModelConfig.link_names`).
+    links_quaternion: Optional[torch.Tensor] = None
+
+    #: Position of spheres specified by collision spheres (:py:attr:`CudaRobotModelConfig.collision_spheres`)
+    #: in x, y, z, r format [b,n,4].
+    link_spheres_tensor: Optional[torch.Tensor] = None
+
+    link_names: Optional[str] = None
+
+    @property
+    def ee_pose(self):
+        return Pose(self.ee_position, self.ee_quaternion)
+
+    def get_link_spheres(self):
+        return self.link_spheres_tensor
+
+    @property
+    def link_pose(self):
+        link_poses = None
+        if self.link_names is not None:
+            link_poses = {}
+            link_pos = self.links_position.contiguous()
+            link_quat = self.links_quaternion.contiguous()
+            for i, v in enumerate(self.link_names):
+                link_poses[v] = Pose(link_pos[..., i, :], link_quat[..., i, :])
+        return link_poses

src/curobo/cuda_robot_model/urdf_kinematics_parser.py

@@ -0,0 +1,167 @@
+#
+# 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 Dict, Optional
+
+# Third Party
+import numpy as np
+import yourdfpy
+from lxml import etree
+
+# CuRobo
+from curobo.cuda_robot_model.kinematics_parser import KinematicsParser, LinkParams
+from curobo.cuda_robot_model.types import JointType
+from curobo.geom.types import Mesh as CuroboMesh
+from curobo.types.base import TensorDeviceType
+from curobo.util.logger import log_error, log_warn
+from curobo.util_file import join_path
+
+
+class UrdfKinematicsParser(KinematicsParser):
+    def __init__(
+        self,
+        urdf_path,
+        load_meshes: bool = False,
+        mesh_root: str = "",
+        extra_links: Optional[Dict[str, LinkParams]] = None,
+    ) -> None:
+        # load robot from urdf:
+        self._robot = yourdfpy.URDF.load(
+            urdf_path,
+            load_meshes=load_meshes,
+            build_scene_graph=False,
+            mesh_dir=mesh_root,
+            filename_handler=yourdfpy.filename_handler_null,
+        )
+        super().__init__(extra_links)
+
+    def build_link_parent(self):
+        self._parent_map = {}
+        for j in self._robot.joint_map:
+            self._parent_map[self._robot.joint_map[j].child] = self._robot.joint_map[j].parent
+
+    def _find_parent_joint_of_link(self, link_name):
+        for j_idx, j in enumerate(self._robot.joint_map):
+            if self._robot.joint_map[j].child == link_name:
+                return j_idx, j
+        log_error("Link is not attached to any joint")
+
+    def _get_joint_name(self, idx):
+        joint = self._robot.joint_names[idx]
+        return joint
+
+    def get_link_parameters(self, link_name: str, base=False) -> LinkParams:
+        link_params = self._get_from_extra_links(link_name)
+        if link_params is not None:
+            return link_params
+        body_params = {}
+        body_params["link_name"] = link_name
+
+        if base:
+            body_params["parent_link_name"] = None
+            joint_transform = np.eye(4)
+            joint_name = "base_joint"
+            joint_type = "fixed"
+            joint_limits = None
+            joint_axis = None
+            body_params["joint_id"] = 0
+        else:
+            body_params["parent_link_name"] = self._parent_map[link_name]
+
+            jid, joint_name = self._find_parent_joint_of_link(link_name)
+            body_params["joint_id"] = jid
+            joint = self._robot.joint_map[joint_name]
+            joint_transform = joint.origin
+            if joint_transform is None:
+                joint_transform = np.eye(4)
+            joint_type = joint.type
+            joint_limits = None
+            joint_axis = None
+            if joint_type != "fixed":
+                if joint_type != "continuous":
+                    joint_limits = {
+                        "effort": joint.limit.effort,
+                        "lower": joint.limit.lower,
+                        "upper": joint.limit.upper,
+                        "velocity": joint.limit.velocity,
+                    }
+                else:
+                    # log_warn("Converting continuous joint to revolute")
+                    joint_type = "revolute"
+                    joint_limits = {
+                        "effort": joint.limit.effort,
+                        "lower": -3.14 * 2,
+                        "upper": 3.14 * 2,
+                        "velocity": joint.limit.velocity,
+                    }
+
+                joint_axis = joint.axis
+                if (joint_axis < 0).any():
+                    log_warn("Joint axis has negative value (-1). This is not supported.")
+                joint_axis = np.abs(joint_axis)
+                body_params["joint_limits"] = [joint_limits["lower"], joint_limits["upper"]]
+                body_params["joint_velocity_limits"] = [
+                    -1.0 * joint_limits["velocity"],
+                    joint_limits["velocity"],
+                ]
+
+        body_params["fixed_transform"] = joint_transform
+        body_params["joint_name"] = joint_name
+
+        body_params["joint_axis"] = joint_axis
+
+        if joint_type == "fixed":
+            joint_type = JointType.FIXED
+        elif joint_type == "prismatic":
+            if joint_axis[0] == 1:
+                joint_type = JointType.X_PRISM
+            if joint_axis[1] == 1:
+                joint_type = JointType.Y_PRISM
+            if joint_axis[2] == 1:
+                joint_type = JointType.Z_PRISM
+        elif joint_type == "revolute":
+            if joint_axis[0] == 1:
+                joint_type = JointType.X_ROT
+            if joint_axis[1] == 1:
+                joint_type = JointType.Y_ROT
+            if joint_axis[2] == 1:
+                joint_type = JointType.Z_ROT
+        else:
+            log_error("Joint type not supported")
+
+        body_params["joint_type"] = joint_type
+        link_params = LinkParams(**body_params)
+
+        return link_params
+
+    def add_absolute_path_to_link_meshes(self, mesh_dir: str = ""):
+        # read all link meshes and update their mesh paths by prepending mesh_dir
+        links = self._robot.link_map
+        for k in links.keys():
+            # read visual and collision
+            vis = links[k].visuals
+            for i in range(len(vis)):
+                m = vis[i].geometry.mesh
+                if m is not None:
+                    m.filename = join_path(mesh_dir, m.filename)
+            col = links[k].collisions
+            for i in range(len(col)):
+                m = col[i].geometry.mesh
+                if m is not None:
+                    m.filename = join_path(mesh_dir, m.filename)
+
+    def get_urdf_string(self):
+        txt = etree.tostring(self._robot.write_xml(), method="xml", encoding="unicode")
+        return txt
+
+    def get_link_mesh(self, link_name):
+        m = self._robot.link_map[link_name].visuals[0].geometry.mesh
+        return CuroboMesh(name=link_name, pose=None, scale=m.scale, file_path=m.filename)

src/curobo/cuda_robot_model/usd_kinematics_parser.py

@@ -0,0 +1,213 @@
+#
+# 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 Dict, List, Optional, Tuple
+
+# Third Party
+import numpy as np
+from pxr import Usd, UsdPhysics
+
+# CuRobo
+from curobo.cuda_robot_model.kinematics_parser import KinematicsParser, LinkParams
+from curobo.cuda_robot_model.types import JointType
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.util.logger import log_error
+
+
+class UsdKinematicsParser(KinematicsParser):
+    """An experimental kinematics parser from USD.
+    NOTE: A more complete solution will be available in a future release. Current implementation
+    does not account for link geometry transformations after a joints.
+    """
+
+    def __init__(
+        self,
+        usd_path: str,
+        flip_joints: List[str] = [],
+        flip_joint_limits: List[str] = [],
+        usd_robot_root: str = "robot",
+        tensor_args: TensorDeviceType = TensorDeviceType(),
+        extra_links: Optional[Dict[str, LinkParams]] = None,
+    ) -> None:
+        # load usd file:
+
+        # create a usd stage
+        self._flip_joints = flip_joints
+        self._flip_joint_limits = flip_joint_limits
+        self._stage = Usd.Stage.Open(usd_path)
+        self._usd_robot_root = usd_robot_root
+        self._parent_joint_map = {}
+        self.tensor_args = tensor_args
+        super().__init__(extra_links)
+
+    @property
+    def robot_prim_root(self):
+        return self._usd_robot_root
+
+    def build_link_parent(self):
+        self._parent_map = {}
+        all_joints = [
+            x
+            for x in self._stage.Traverse()
+            if (x.IsA(UsdPhysics.Joint) and str(x.GetPath()).startswith(self._usd_robot_root))
+        ]
+        for l in all_joints:
+            parent, child = get_links_for_joint(l)
+            if child is not None and parent is not None:
+                self._parent_map[child.GetName()] = parent.GetName()
+                self._parent_joint_map[child.GetName()] = l  # store joint prim
+
+    def get_link_parameters(self, link_name: str, base: bool = False) -> LinkParams:
+        """Get Link parameters from usd stage.
+
+        NOTE: USD kinematics "X" axis joints map to "Z" in URDF. Specifically,
+        uniform token physics:axis = "X" value only matches "Z" in URDF. This is because of usd
+        files assuming Y axis as up while urdf files assume Z axis as up.
+
+        Args:
+            link_name (str): Name of link.
+            base (bool, optional): flag to specify base link. Defaults to False.
+
+        Returns:
+            LinkParams: obtained link parameters.
+        """
+        link_params = self._get_from_extra_links(link_name)
+        if link_params is not None:
+            return link_params
+        joint_limits = None
+        joint_axis = None
+        if base:
+            parent_link_name = None
+            joint_transform = np.eye(4)
+            joint_name = "base_joint"
+            joint_type = JointType.FIXED
+
+        else:
+            parent_link_name = self._parent_map[link_name]
+            joint_prim = self._parent_joint_map[link_name]  # joint prim connects link
+            joint_transform = self._get_joint_transform(joint_prim)
+            joint_axis = None
+            joint_name = joint_prim.GetName()
+            if joint_prim.IsA(UsdPhysics.FixedJoint):
+                joint_type = JointType.FIXED
+            elif joint_prim.IsA(UsdPhysics.RevoluteJoint):
+                j_prim = UsdPhysics.RevoluteJoint(joint_prim)
+                joint_axis = j_prim.GetAxisAttr().Get()
+                joint_limits = np.radians(
+                    np.ravel([j_prim.GetLowerLimitAttr().Get(), j_prim.GetUpperLimitAttr().Get()])
+                )
+                if joint_name in self._flip_joints.keys():
+                    joint_axis = self._flip_joints[joint_name]
+                if joint_axis == "X":
+                    joint_type = JointType.X_ROT
+                elif joint_axis == "Y":
+                    joint_type = JointType.Y_ROT
+                elif joint_axis == "Z":
+                    joint_type = JointType.Z_ROT
+                else:
+                    log_error("Joint axis not supported" + str(joint_axis))
+
+            elif joint_prim.IsA(UsdPhysics.PrismaticJoint):
+                j_prim = UsdPhysics.PrismaticJoint(joint_prim)
+
+                joint_axis = j_prim.GetAxisAttr().Get()
+                joint_limits = np.ravel(
+                    [j_prim.GetLowerLimitAttr().Get(), j_prim.GetUpperLimitAttr().Get()]
+                )
+                if joint_name in self._flip_joints.keys():
+                    joint_axis = self._flip_joints[joint_name]
+                if joint_name in self._flip_joint_limits:
+                    joint_limits = np.ravel(
+                        [-1.0 * j_prim.GetUpperLimitAttr().Get(), j_prim.GetLowerLimitAttr().Get()]
+                    )
+                if joint_axis == "X":
+                    joint_type = JointType.X_PRISM
+                elif joint_axis == "Y":
+                    joint_type = JointType.Y_PRISM
+                elif joint_axis == "Z":
+                    joint_type = JointType.Z_PRISM
+                else:
+                    log_error("Joint axis not supported" + str(joint_axis))
+            else:
+                log_error("Joint type not supported")
+        link_params = LinkParams(
+            link_name=link_name,
+            joint_name=joint_name,
+            joint_type=joint_type,
+            fixed_transform=joint_transform,
+            parent_link_name=parent_link_name,
+            joint_limits=joint_limits,
+        )
+        return link_params
+
+    def _get_joint_transform(self, prim: Usd.Prim):
+        j_prim = UsdPhysics.Joint(prim)
+        position = np.ravel(j_prim.GetLocalPos0Attr().Get())
+        quatf = j_prim.GetLocalRot0Attr().Get()
+        quat = np.zeros(4)
+        quat[0] = quatf.GetReal()
+        quat[1:] = quatf.GetImaginary()
+
+        # create a homogenous transformation matrix:
+        transform_0 = Pose(self.tensor_args.to_device(position), self.tensor_args.to_device(quat))
+
+        position = np.ravel(j_prim.GetLocalPos1Attr().Get())
+        quatf = j_prim.GetLocalRot1Attr().Get()
+        quat = np.zeros(4)
+        quat[0] = quatf.GetReal()
+        quat[1:] = quatf.GetImaginary()
+
+        # create a homogenous transformation matrix:
+        transform_1 = Pose(self.tensor_args.to_device(position), self.tensor_args.to_device(quat))
+        transform = (
+            transform_0.multiply(transform_1.inverse()).get_matrix().cpu().view(4, 4).numpy()
+        )
+
+        # get attached link transform:
+
+        return transform
+
+
+def get_links_for_joint(prim: Usd.Prim) -> Tuple[Optional[Usd.Prim], Optional[Usd.Prim]]:
+    """Get all link prims from the given joint prim.
+
+    Note:
+        This assumes that the `body0_rel_targets` and `body1_rel_targets` are configured such
+        that the parent link is specified in `body0_rel_targets` and the child links is specified
+        in `body1_rel_targets`.
+    """
+    stage = prim.GetStage()
+    joint_api = UsdPhysics.Joint(prim)
+
+    rel0_targets = joint_api.GetBody0Rel().GetTargets()
+    if len(rel0_targets) > 1:
+        raise NotImplementedError(
+            "`get_links_for_joint` does not currently handle more than one relative"
+            f" body target in the joint. joint_prim: {prim}, body0_rel_targets:"
+            f" {rel0_targets}"
+        )
+    link0_prim = None
+    if len(rel0_targets) != 0:
+        link0_prim = stage.GetPrimAtPath(rel0_targets[0])
+
+    rel1_targets = joint_api.GetBody1Rel().GetTargets()
+    if len(rel1_targets) > 1:
+        raise NotImplementedError(
+            "`get_links_for_joint` does not currently handle more than one relative"
+            f" body target in the joint. joint_prim: {prim}, body1_rel_targets:"
+            f" {rel0_targets}"
+        )
+    link1_prim = None
+    if len(rel1_targets) != 0:
+        link1_prim = stage.GetPrimAtPath(rel1_targets[0])
+
+    return (link0_prim, link1_prim)