Isaac Sim 4.0 support, Kinematics API doc, Windows support

src/curobo/cuda_robot_model/__init__.py

@@ -10,34 +10,75 @@
 #
 
 """
-This module contains code for cuda accelerated kinematics.
+This module contains GPU accelerated kinematics leveraging CUDA. Kinematics computations enable
+mapping from a robot's joint configuration to the pose of the robot's links in Cartesian space
+(with reference to the robot's base link). In cuRobo, robot's geometry is approximated with spheres
+and their positions are also computed as part of kinematics. This mapping is differentiable,
+enabling their use in optimization problems and as part of neural networks.
+
+
+..  figure:: ../images/robot_representation.png
+    :width: 400px
+    :align: center
+
+    Robot representation in cuRobo is shown for the Franka Panda robot.
+
+
 
 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.
+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.KinematicsParser` 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. An
+example workflow for setting up a robot from URDF is shown below:
 
-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`.
+.. graphviz::
 
-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
+   digraph {
+    rankdir=LR;
+   bgcolor="#808080";
+   edge [color = "#FFFFFF"; fontsize=10];
+   node [shape="box", style="rounded, filled", fontsize=12, color="#76b900", fontcolor="#FFFFFF"];
+   "CudaRobotGenerator" [color="#FFFFFF", fontcolor="#000000"]
+   "UrdfKinematicsParser" [fillcolor="#FFFFFF", fontcolor="#000000", style="box, filled", color="#000000"]
+   "CudaRobotGenerator" [color="#FFFFFF", fontcolor="#000000"]
+    "URDF" [fillcolor="#FFFFFF", fontcolor="#000000", style="box, filled, dashed", color="#000000"]
+    "XRDF" [fillcolor="#FFFFFF", fontcolor="#000000", style="box, filled, dashed", color="#000000"]
+    "cuRobo YML" [fillcolor="#FFFFFF", fontcolor="#000000", style="box, filled", color="#000000"]
+
+   "CudaRobotGeneratorConfig" -> "CudaRobotGenerator";
+   "CudaRobotGenerator" -> "UrdfKinematicsParser" [dir="both"];
+   "CudaRobotGenerator" -> "CudaRobotModelConfig";
+   "URDF" -> "cuRobo YML";
+   "XRDF" -> "cuRobo YML" [style="dashed",label="Optional", fontcolor="#FFFFFF"];
+   "cuRobo YML" -> "CudaRobotGeneratorConfig";
+
+   }
+
+
+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`. cuRobo also supports using
+`XRDF <https://nvidia-isaac-ros.github.io/concepts/manipulation/xrdf.html>`_ for representing
+the additional parameters of the robot that are not available in URDF.
+
+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`.
+    :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

@@ -8,16 +8,23 @@
 # without an express license agreement from NVIDIA CORPORATION or
 # its affiliates is strictly prohibited.
 #
+"""
+Generates a Tensor representation of kinematics for use in
+:class:`~curobo.cuda_robot_model.CudaRobotModel`. This module reads the robot from a
+:class:`~curobo.cuda_robot_model.kinematics_parser.KinematicsParser` and
+generates the necessary tensors for kinematics computation.
+
+"""
+
 from __future__ import annotations
 
 # Standard Library
 import copy
 import os
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 # Third Party
-import importlib_resources
 import torch
 import torch.autograd.profiler as profiler
 
@@ -45,13 +52,13 @@ try:
 except ImportError:
     log_info(
         "USDParser failed to import, install curobo with pip install .[usd] "
-        + "or pip install usd-core, NOTE: Do not install this if using with ISAAC SIM."
+        + "or pip install usd-core, NOTE: Do not install this if using with Isaac Sim."
     )
 
 
 @dataclass
 class CudaRobotGeneratorConfig:
-    """Create Cuda Robot Model Configuration."""
+    """Robot representation generator configuration, loads from a dictionary."""
 
     #: Name of base link for kinematic tree.
     base_link: str
@@ -69,7 +76,7 @@ class CudaRobotGeneratorConfig:
     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 can be generated for robot using `Isaac Sim Robot Description Editor <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.
@@ -79,13 +86,17 @@ class CudaRobotGeneratorConfig:
     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)
+    #: for self collision avoidance (e.g., `MoveIt Panda Issue <https://github.com/ros-planning/panda_moveit_config/pull/35#issuecomment-671333863>`_).
     self_collision_buffer: Optional[Dict[str, float]] = None
 
-    #: Self-collision ignore
+    #: Dictionary with each key as a link name and value as a list of link names to ignore self
+    #: collision. E.g., {"link1": ["link2", "link3"], "link2": ["link3", "link4"]} will
+    #: ignore self collision between link1 and link2, link1 and link3, link2 and link3, link2 and
+    #: link4. The mapping is bidirectional so it's sufficient to mention the mapping in one
+    #: direction (i.e., not necessary to mention "link1" in ignore list for "link2").
     self_collision_ignore: Optional[Dict[str, List[str]]] = None
 
-    #: debug config
+    #: Debugging information to pass to kinematics module.
     debug: Optional[Dict[str, Any]] = None
 
     #: Enabling this flag writes out the cumulative transformation matrix to global memory. This
@@ -129,33 +140,46 @@ class CudaRobotGeneratorConfig:
     #: joint angle given from this dictionary.
     lock_joints: Optional[Dict[str, float]] = None
 
+    #: Additional links to add to parsed kinematics tree. This is useful for adding fixed links
+    #: that are not present in the URDF or USD.
     extra_links: Optional[Dict[str, LinkParams]] = None
 
-    #: this is deprecated
+    #: Deprecated way to add a fixed link.
     add_object_link: bool = False
 
+    #: Deprecated flag to load assets from external module. Now, pass absolute path to
+    #: asset_root_path or use :class:`~curobo.util.file_path.ContentPath`.
     use_external_assets: bool = False
 
+    #: Deprecated path to load assets from external module. Use
+    #: :class:`~curobo.util.file_path.ContentPath` instead.
     external_asset_path: Optional[str] = None
+
+    #: Deprecated path to load robot configs from external module. Use
+    #: :class:`~curobo.util.file_path.ContentPath` instead.
     external_robot_configs_path: Optional[str] = None
 
     #: Create n collision spheres for links with name
     extra_collision_spheres: Optional[Dict[str, int]] = None
 
-    #: cspace config
+    #: Configuration space parameters for robot (e.g, acceleration, jerk limits).
     cspace: Union[None, CSpaceConfig, Dict[str, List[Any]]] = None
 
     #: Enable loading meshes from kinematics parser.
     load_meshes: bool = False
 
     def __post_init__(self):
+        """Post initialization adds absolute paths, converts dictionaries to objects."""
+
         # add root path:
         # Check if an external asset path is provided:
         asset_path = get_assets_path()
         robot_path = get_robot_configs_path()
         if self.external_asset_path is not None:
+            log_warn("Deprecated: external_asset_path is deprecated, use ContentPath")
             asset_path = self.external_asset_path
         if self.external_robot_configs_path is not None:
+            log_warn("Deprecated: external_robot_configs_path is deprecated, use ContentPath")
             robot_path = self.external_robot_configs_path
 
         if self.urdf_path is not None:
@@ -214,30 +238,45 @@ class CudaRobotGeneratorConfig:
 
 
 class CudaRobotGenerator(CudaRobotGeneratorConfig):
+    """Robot Kinematics Representation Generator.
+
+    The word "Chain" is used interchangeably with "Tree" in this class.
+
+    """
+
     def __init__(self, config: CudaRobotGeneratorConfig) -> None:
+        """Initialize the robot generator.
+
+        Args:
+            config: Parameters to initialize the robot generator.
+        """
         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._kinematics_config = None
         self.initialize_tensors()
 
     @property
-    def kinematics_config(self):
+    def kinematics_config(self) -> KinematicsTensorConfig:
+        """Kinematics representation as Tensors."""
         return self._kinematics_config
 
     @property
-    def self_collision_config(self):
+    def self_collision_config(self) -> SelfCollisionKinematicsConfig:
+        """Self collision configuration for robot."""
         return self._self_collision_data
 
     @property
     def kinematics_parser(self):
+        """Kinematics parser used to generate robot parameters."""
         return self._kinematics_parser
 
     @profiler.record_function("robot_generator/initialize_tensors")
     def initialize_tensors(self):
+        """Initialize tensors for kinematics representatiobn."""
         self._joint_limits = None
         self._self_collision_data = None
         self.lock_jointstate = None
@@ -335,10 +374,15 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
             lock_jointstate=self.lock_jointstate,
             mimic_joints=self._mimic_joint_data,
         )
-        if self.asset_root_path != "":
+        if self.asset_root_path is not None and self.asset_root_path != "":
             self._kinematics_parser.add_absolute_path_to_link_meshes(self.asset_root_path)
 
     def add_link(self, link_params: LinkParams):
+        """Add an extra link to the robot kinematics tree.
+
+        Args:
+            link_params: Parameters of the link to add.
+        """
         self.extra_links[link_params.link_name] = link_params
 
     def add_fixed_link(
@@ -348,6 +392,14 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         joint_name: Optional[str] = None,
         transform: Optional[Pose] = None,
     ):
+        """Add a fixed link to the robot kinematics tree.
+
+        Args:
+            link_name: Name of the link to add.
+            parent_link_name: Parent link to add the fixed link to.
+            joint_name: Name of fixed to joint to create.
+            transform: Offset transform of the fixed link from the joint.
+        """
         if transform is None:
             transform = (
                 Pose.from_list([0, 0, 0, 1, 0, 0, 0], self.tensor_args)
@@ -368,7 +420,22 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         self.add_link(link_params)
 
     @profiler.record_function("robot_generator/build_chain")
-    def _build_chain(self, base_link, ee_link, other_links, link_names):
+    def _build_chain(
+        self,
+        base_link: str,
+        ee_link: str,
+        other_links: List[str],
+    ) -> List[str]:
+        """Build kinematic tree of the robot.
+
+        Args:
+            base_link: Name of base link for the chain.
+            ee_link: Name of end-effector link for the chain.
+            other_links: List of other links to add to the chain.
+
+        Returns:
+            List[str]: List of link names in the chain.
+        """
         self._n_dofs = 0
         self._controlled_links = []
         self._bodies = []
@@ -404,7 +471,13 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         self.non_fixed_joint_names = self.joint_names.copy()
         return chain_link_names
 
-    def _get_mimic_joint_data(self):
+    def _get_mimic_joint_data(self) -> Dict[str, List[int]]:
+        """Get joints that are mimicked from actuated joints joints.
+
+        Returns:
+            Dict[str, List[int]]: Dictionary containing name of actuated joint and list of mimic
+                joint indices.
+        """
         # get joint types:
         mimic_joint_data = {}
         for i in range(1, len(self._bodies)):
@@ -417,7 +490,16 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         return mimic_joint_data
 
     @profiler.record_function("robot_generator/build_kinematics_tensors")
-    def _build_kinematics_tensors(self, base_link, ee_link, link_names, chain_link_names):
+    def _build_kinematics_tensors(self, base_link, link_names, chain_link_names):
+        """Create kinematic tensors for robot given kinematic tree.
+
+        Args:
+            base_link: Name of base link for the tree.
+            link_names: Namer of links to compute kinematics for. This is used to determine link
+                indices to store pose during forward kinematics.
+            chain_link_names: List of link names in the kinematic tree. Used to traverse the
+                kinematic tree.
+        """
         self._active_joints = []
         self._mimic_joint_data = {}
         link_map = [0 for i in range(len(self._bodies))]
@@ -493,9 +575,19 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         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)
+    def _build_kinematics(
+        self, base_link: str, ee_link: str, other_links: List[str], link_names: List[str]
+    ):
+        """Build kinematics tensors given base link, end-effector link and other links.
+
+        Args:
+            base_link: Name of base link for the kinematic tree.
+            ee_link: Name of end-effector link for the kinematic tree.
+            other_links: List of other links to add to the kinematic tree.
+            link_names: List of link names to store poses after kinematics computation.
+        """
+        chain_link_names = self._build_chain(base_link, ee_link, other_links)
+        self._build_kinematics_tensors(base_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
@@ -504,13 +596,26 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
     @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,
+        base_link: str,
+        ee_link: str,
+        other_links: List[str],
+        link_names: List[str],
         lock_joints: Dict[str, float],
     ):
-        chain_link_names = self._build_chain(base_link, ee_link, other_links, link_names)
+        """Build kinematics with locked joints.
+
+        This function will first build the chain with no locked joints, find the transforms
+        when the locked joints are set to the given values, and then use these transforms as
+        fixed transforms for the locked joints.
+
+        Args:
+            base_link: Base link of the kinematic tree.
+            ee_link: End-effector link of the kinematic tree.
+            other_links: Other links to add to the kinematic tree.
+            link_names: List of link names to store poses after kinematics computation.
+            lock_joints: Joints to lock in the kinematic tree with value to lock at.
+        """
+        chain_link_names = self._build_chain(base_link, ee_link, other_links)
         # find links attached to lock joints:
         lock_joint_names = list(lock_joints.keys())
 
@@ -531,7 +636,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         new_link_names = list(set(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)
+        self._build_kinematics_tensors(base_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
@@ -628,12 +733,12 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         collision_link_names: List[str],
         collision_sphere_buffer: Union[float, Dict[str, float]] = 0.0,
     ):
-        """
+        """Build collision model for robot.
 
         Args:
-            collision_spheres (_type_): _description_
-            collision_link_names (_type_): _description_
-            collision_sphere_buffer (float, optional): _description_. Defaults to 0.0.
+            collision_spheres: Spheres for each link of the robot.
+            collision_link_names: Name of links to load spheres for.
+            collision_sphere_buffer: Additional padding to add to collision spheres.
         """
 
         # We create all tensors on cpu and then finally move them to gpu
@@ -680,7 +785,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         )
 
         # build self collision distance tensor:
-        self.self_collision_distance = (
+        self_collision_distance = (
             torch.zeros(
                 (self.total_spheres, self.total_spheres),
                 dtype=cpu_tensor_args.dtype,
@@ -691,7 +796,6 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         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):
@@ -724,44 +828,66 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
                         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
+                            self_collision_distance[sp1, sp2] = rad1[k1] + rad2[k2] + c1 + c2
 
+        self_collision_distance = self_collision_distance.to(device=self.tensor_args.device)
         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))
+            d_mat = self_collision_distance
+            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
-        )
-
+        ) = self._create_self_collision_thread_data(self_collision_distance)
         use_experimental_kernel = True
+
+        if (
+            self.debug is not None
+            and "self_collision_experimental" in self.debug
+            and self.debug["self_collision_experimental"] is not None
+        ):
+            use_experimental_kernel = self.debug["self_collision_experimental"]
+
+        if not valid_data:
+            use_experimental_kernel = False
+            log_warn("Self Collision checks are greater than 32 * 512, using slower kernel")
+        if use_experimental_kernel:
+            self_coll_matrix = torch.zeros((2), device=self.tensor_args.device, dtype=torch.uint8)
+        else:
+            self_coll_matrix = (self_collision_distance != -(torch.inf)).to(dtype=torch.uint8)
+        # self_coll_matrix = (self_collision_distance != -(torch.inf)).to(dtype=torch.uint8)
+
         # 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,
+            collision_matrix=self_coll_matrix,
+            experimental_kernel=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):
+    def _create_self_collision_thread_data(
+        self, collision_threshold: torch.Tensor
+    ) -> Tuple[torch.Tensor, int, bool, int]:
+        """Create thread data for self collision checks.
+
+        Args:
+            collision_threshold: Collision distance between spheres of the robot. Used to
+                skip self collision checks when distance is -inf.
+
+        Returns:
+            Tuple[torch.Tensor, int, bool, int]: Thread location for self collision checks,
+                number of self collision checks, if thread calculation was successful,
+                and number of checks per thread.
+
+        """
         coll_cpu = collision_threshold.cpu()
         max_checks_per_thread = 512
         thread_loc = torch.zeros((2 * 32 * max_checks_per_thread), dtype=torch.int16) - 1
@@ -815,7 +941,10 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         elif val < 512:
             max_checks_per_thread = 512
         else:
-            raise ValueError("Self Collision not supported")
+            log_error(
+                "Self Collision not supported as checks are greater than 32 * 512, \
+                      reduce number of spheres used to approximate the robot."
+            )
 
         if max_checks_per_thread < 2:
             max_checks_per_thread = 2
@@ -829,7 +958,13 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         )
 
     @profiler.record_function("robot_generator/add_body_to_tree")
-    def _add_body_to_tree(self, link_name, base=False):
+    def _add_body_to_tree(self, link_name: str, base=False):
+        """Add link to kinematic tree.
+
+        Args:
+            link_name: Name of the link to add.
+            base: Is this the base link of the kinematic tree?
+        """
         body_idx = len(self._bodies)
 
         rigid_body_params = self._kinematics_parser.get_link_parameters(link_name, base=base)
@@ -848,7 +983,17 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
         )
         self._name_to_idx_map[rigid_body_params.link_name] = body_idx
 
-    def _get_joint_links(self, joint_names: List[str]):
+    def _get_joint_links(self, joint_names: List[str]) -> Dict[str, Dict[str, Union[str, int]]]:
+        """Get data (parent link, child link, mimic, link_index) for joints given in the list.
+
+        Args:
+            joint_names: Names of joints to get data for.
+
+        Returns:
+            Dict[str, Dict[str, Union[str, int]]]: Dictionary containing joint name as key and
+                dictionary containing parent link, child link, and link index as
+                values. Also includes mimic joint data if present.
+        """
         j_data = {}
 
         for j in joint_names:
@@ -878,6 +1023,22 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
     def _get_link_poses(
         self, q: torch.Tensor, link_names: List[str], kinematics_config: KinematicsTensorConfig
     ) -> Pose:
+        """Get Pose of links at given joint angles using forward kinematics.
+
+        This is implemented here to avoid circular dependencies with
+        :class:`~curobo.cuda_robot_model.cuda_robot_model.CudaRobotModel` module. This is used
+        to calculate fixed transforms for locked joints in this class. This implementation
+        does not compute position of robot spheres.
+
+        Args:
+            q: Joint angles to compute forward kinematics for.
+            link_names: Name of links to return pose.
+            kinematics_config: Tensor Configuration for kinematics computation.
+
+        Returns:
+            Pose: Pose of links at given joint angles.
+        """
+        q = q.view(1, -1)
         link_pos_seq = torch.zeros(
             (1, len(self.link_names), 3),
             device=self.tensor_args.device,
@@ -889,7 +1050,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
             dtype=self.tensor_args.dtype,
         )
         batch_robot_spheres = torch.zeros(
-            (1, self.total_spheres, 4),
+            (1, 0, 4),
             device=self.tensor_args.device,
             dtype=self.tensor_args.dtype,
         )
@@ -903,23 +1064,24 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
             device=self.tensor_args.device,
             dtype=self.tensor_args.dtype,
         )
+
         link_pos_seq, link_quat_seq, _ = get_cuda_kinematics(
             link_pos_seq,
             link_quat_seq,
-            batch_robot_spheres,
+            batch_robot_spheres.contiguous(),
             global_cumul_mat,
             q,
-            kinematics_config.fixed_transforms,
-            kinematics_config.link_spheres,
+            kinematics_config.fixed_transforms.contiguous(),
+            kinematics_config.link_spheres.contiguous(),
             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_sphere_idx_map.contiguous(),  # sphere idx map
             kinematics_config.link_chain_map,
             kinematics_config.joint_offset_map,
             grad_out_q,
-            self.use_global_cumul,
+            False,
         )
         position = torch.zeros(
             (q.shape[0], len(link_names), 3),
@@ -936,14 +1098,21 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
             i = self.link_names.index(l)
             position[:, li, :] = link_pos_seq[:, i, :]
             quaternion[:, li, :] = link_quat_seq[:, i, :]
-        return Pose(position=position, quaternion=quaternion)
+        return Pose(position=position.clone(), quaternion=quaternion.clone())
 
-    @property
-    def get_joint_limits(self):
+    def get_joint_limits(self) -> JointLimits:
+        """Get joint limits for the robot."""
         return self._joint_limits
 
     @profiler.record_function("robot_generator/get_joint_limits")
-    def _get_joint_position_velocity_limits(self):
+    def _get_joint_position_velocity_limits(self) -> Dict[str, torch.Tensor]:
+        """Compute joint position and velocity limits for the robot.
+
+        Returns:
+            Dict[str, torch.Tensor]: Dictionary containing position and velocity limits for the
+                robot. Each value is a tensor of shape (2, n_joints) with first row containing
+                minimum limits and second row containing maximum limits.
+        """
         joint_limits = {"position": [[], []], "velocity": [[], []]}
 
         for idx in self._active_joints:
@@ -959,6 +1128,7 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
 
     @profiler.record_function("robot_generator/update_joint_limits")
     def _update_joint_limits(self):
+        """Update limits from CSpaceConfig (acceleration, jerk limits and position clips)."""
         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)]
@@ -970,7 +1140,6 @@ class CudaRobotGenerator(CudaRobotGeneratorConfig):
             ]
         )
         # clip joint position:
-        # NOTE: change this to be per joint
         joint_limits["position"][0] += self.cspace.position_limit_clip
         joint_limits["position"][1] -= self.cspace.position_limit_clip
         joint_limits["velocity"][0] *= self.cspace.velocity_scale

src/curobo/cuda_robot_model/cuda_robot_model.py

@@ -8,6 +8,14 @@
 # without an express license agreement from NVIDIA CORPORATION or
 # its affiliates is strictly prohibited.
 #
+
+"""
+This module builds a kinematic representation of a robot on the GPU and provides
+differentiable mapping from it's joint configuration to Cartesian pose of it's links
+(forward kinematics). This module also computes the position of the spheres of the robot as part
+of the forward kinematics function.
+"""
+
 from __future__ import annotations
 
 # Standard Library
@@ -24,28 +32,75 @@ from curobo.cuda_robot_model.cuda_robot_generator import (
     CudaRobotGeneratorConfig,
 )
 from curobo.cuda_robot_model.kinematics_parser import KinematicsParser
-from curobo.cuda_robot_model.types import KinematicsTensorConfig, SelfCollisionKinematicsConfig
+from curobo.cuda_robot_model.types import (
+    CSpaceConfig,
+    JointLimits,
+    KinematicsTensorConfig,
+    SelfCollisionKinematicsConfig,
+)
+from curobo.cuda_robot_model.util import load_robot_yaml
 from curobo.curobolib.kinematics import get_cuda_kinematics
-from curobo.geom.types import Mesh, Sphere
+from curobo.geom.sphere_fit import SphereFitType
+from curobo.geom.types import Mesh, Obstacle, Sphere
 from curobo.types.base import TensorDeviceType
+from curobo.types.file_path import ContentPath
 from curobo.types.math import Pose
 from curobo.types.state import JointState
-from curobo.util.logger import log_error
-from curobo.util_file import get_robot_path, join_path, load_yaml
+from curobo.util.logger import log_error, log_info, log_warn
+from curobo.util_file import is_file_xrdf
 
 
 @dataclass
 class CudaRobotModelConfig:
+    """
+    Configuration for robot kinematics on GPU.
+
+    Helper functions are provided to load this configuration from an URDF file or from
+    a cuRobo robot configuration file (:ref:`tut_robot_configuration`). To create from a XRDF, use
+    :ref:`curobo.util.xrdf_utils.convert_xrdf_to_curobo`.
+    """
+
+    #: Device and floating point precision to use for kinematics.
     tensor_args: TensorDeviceType
+
+    #: Names of links to compute poses with forward kinematics.
     link_names: List[str]
+
+    #: Tensors representing kinematics of the robot. This can be created using
+    #: :class:`~curobo.cuda_robot_model.cuda_robot_generator.CudaRobotGenerator`.
     kinematics_config: KinematicsTensorConfig
+
+    #: Collision pairs to ignore when computing self collision between spheres across all
+    #: robot links. This also contains distance threshold between spheres pairs and which thread
+    #: indices for calculating the distances. More details on computing these parameters is in
+    #: :func:`~curobo.cuda_robot_model.cuda_robot_generator.CudaRobotGenerator._build_collision_model`.
     self_collision_config: Optional[SelfCollisionKinematicsConfig] = None
+
+    #: Parser to load kinematics from URDF or USD files. This is used to load kinematics
+    #: representation of the robot. This is created using
+    #: :class:`~curobo.cuda_robot_model.kinematics_parser.KinematicsParser`.
+    #: USD is an experimental feature and might not work for all robots.
     kinematics_parser: Optional[KinematicsParser] = None
+
+    #: Output jacobian during forward kinematics. This is not implemented. The forward kinematics
+    #: function does use Jacobian during backward pass. What's not supported is
     compute_jacobian: bool = False
-    use_global_cumul: bool = False
+
+    #: Store transformation matrix of every link during forward kinematics call in global memory.
+    #: This helps speed up backward pass as we don't need to recompute the transformation matrices.
+    #: However, this increases memory usage and also slightly slows down forward kinematics.
+    #: Enabling this is recommended for getting the best performance.
+    use_global_cumul: bool = True
+
+    #: Generator config used to create this robot kinematics model.
     generator_config: Optional[CudaRobotGeneratorConfig] = None
 
-    def get_joint_limits(self):
+    def get_joint_limits(self) -> JointLimits:
+        """Get limits of actuated joints of the robot.
+
+        Returns:
+            JointLimits: Joint limits of the robot's actuated joints.
+        """
         return self.kinematics_config.joint_limits
 
     @staticmethod
@@ -99,29 +154,96 @@ class CudaRobotModelConfig:
         return CudaRobotModelConfig.from_config(config)
 
     @staticmethod
-    def from_robot_yaml_file(
-        file_path: str,
+    def from_content_path(
+        content_path: ContentPath,
         ee_link: Optional[str] = None,
         tensor_args: TensorDeviceType = TensorDeviceType(),
-    ):
-        config_file = load_yaml(join_path(get_robot_path(), file_path))["robot_cfg"]["kinematics"]
+    ) -> CudaRobotModelConfig:
+        """Load robot from Contentpath containing paths to robot description files.
+
+        Args:
+            content_path: Path to robot configuration files.
+            ee_link: End-effector link name. If None, it is read from the file.
+            tensor_args: Device to load robot model, defaults to cuda:0.
+
+        Returns:
+            CudaRobotModelConfig: cuda robot model configuration.
+        """
+
+        config_file = load_robot_yaml(content_path)
+        if "robot_cfg" in config_file:
+            config_file = config_file["robot_cfg"]
+        if "kinematics" in config_file:
+            config_file = config_file["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_robot_yaml_file(
+        file_path: Union[str, Dict],
+        ee_link: Optional[str] = None,
+        tensor_args: TensorDeviceType = TensorDeviceType(),
+        urdf_path: Optional[str] = None,
+    ) -> CudaRobotModelConfig:
+        """Load robot from a yaml file that is in cuRobo's format (:ref:`tut_robot_configuration`).
+
+        Args:
+            file_path: Path to robot configuration file (yml or xrdf).
+            ee_link: End-effector link name. If None, it is read from the file.
+            tensor_args: Device to load robot model, defaults to cuda:0.
+            urdf_path: Path to urdf file. This is required when loading a xrdf file.
+
+        Returns:
+            CudaRobotModelConfig: cuda robot model configuration.
+        """
+        if isinstance(file_path, dict):
+            content_path = ContentPath(robot_urdf_file=urdf_path, robot_config_file=file_path)
+        else:
+            if is_file_xrdf(file_path):
+                content_path = ContentPath(robot_urdf_file=urdf_path, robot_xrdf_file=file_path)
+            else:
+                content_path = ContentPath(robot_urdf_file=urdf_path, robot_config_file=file_path)
+
+        return CudaRobotModelConfig.from_content_path(content_path, ee_link, tensor_args)
+
     @staticmethod
     def from_data_dict(
         data_dict: Dict[str, Any],
         tensor_args: TensorDeviceType = TensorDeviceType(),
-    ):
+    ) -> CudaRobotModelConfig:
+        """Load robot from a dictionary containing data for :class:`~curobo.cuda_robot_model.cuda_robot_generator.CudaRobotGeneratorConfig`.
+
+        :tut_robot_configuration discusses the data required to load a robot.
+
+        Args:
+            data_dict: Input dictionary containing robot configuration.
+            tensor_args: Device to load robot model, defaults to cuda:0.
+
+        Returns:
+            CudaRobotModelConfig: cuda robot model configuration.
+        """
+        if "robot_cfg" in data_dict:
+            data_dict = data_dict["robot_cfg"]
+        if "kinematics" in data_dict:
+            data_dict = data_dict["kinematics"]
         return CudaRobotModelConfig.from_config(
             CudaRobotGeneratorConfig(**data_dict, tensor_args=tensor_args)
         )
 
     @staticmethod
-    def from_config(config: CudaRobotGeneratorConfig):
+    def from_config(config: CudaRobotGeneratorConfig) -> CudaRobotModelConfig:
+        """Create a robot model configuration from a generator configuration.
+
+        Args:
+            config: Input robot generator configuration.
+
+        Returns:
+            CudaRobotModelConfig: robot model configuration.
+        """
         # create a config generator and load all values
         generator = CudaRobotGenerator(config)
         return CudaRobotModelConfig(
@@ -136,17 +258,19 @@ class CudaRobotModelConfig:
         )
 
     @property
-    def cspace(self):
+    def cspace(self) -> CSpaceConfig:
+        """Get cspace parameters of the robot."""
         return self.kinematics_config.cspace
 
     @property
     def dof(self) -> int:
+        """Get the number of actuated joints (degrees of freedom) of the robot"""
         return self.kinematics_config.n_dof
 
 
 @dataclass
 class CudaRobotModelState:
-    """Dataclass that stores kinematics information."""
+    """Kinematic state of robot."""
 
     #: End-effector position stored as x,y,z in meters [b, 3]. End-effector is defined by
     #: :attr:`CudaRobotModel.ee_link`.
@@ -207,20 +331,38 @@ class CudaRobotModel(CudaRobotModelConfig):
     """
     CUDA Accelerated Robot Model
 
+    Load basic kinematics from an URDF with :func:`~CudaRobotModelConfig.from_basic_urdf`.
+    Check :ref:`tut_robot_configuration` for details on how to also create a geometric
+    representation of the robot.
     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):
+        """Initialize kinematics instance with a robot model configuration.
+
+        Args:
+            config: Input robot model configuration.
+        """
         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):
+    def update_batch_size(
+        self, batch_size: int, force_update: bool = False, reset_buffers: bool = False
+    ):
+        """Update batch size of the robot model.
+
+        Args:
+            batch_size: Batch size to update the robot model.
+            force_update: Detach gradients of tensors. This is not supported.
+            reset_buffers: Recreate the tensors even if the batch size is same.
+        """
         if batch_size == 0:
             log_error("batch size is zero")
         if force_update and self._batch_size == batch_size and self.compute_jacobian:
+            log_error("Outputting jacobian is not supported")
             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:
@@ -252,6 +394,7 @@ class CudaRobotModel(CudaRobotModelConfig):
                 dtype=self.tensor_args.dtype,
             )
             if self.compute_jacobian:
+                log_error("Outputting jacobian is not supported")
                 self.lin_jac = torch.zeros(
                     [batch_size, 3, self.kinematics_config.n_dofs],
                     device=self.tensor_args.device,
@@ -264,10 +407,27 @@ class CudaRobotModel(CudaRobotModelConfig):
                 )
 
     @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)
+    def forward(
+        self, q, link_name=None, calculate_jacobian=False
+    ) -> Tuple[Tensor, Tensor, None, None, Tensor, Tensor, Tensor]:
+        """Compute forward kinematics of the robot.
+
+        Use :func:`~get_state` to get a structured output.
+
+        Args:
+            q: Joint configuration of the robot. Shape should be [batch_size, dof].
+            link_name: Name of link to return pose of. If None, returns end-effector pose.
+            calculate_jacobian: Calculate jacobian of the robot. Not supported.
+
+        Returns:
+            Tuple[Tensor, Tensor, None, None, Tensor, Tensor, Tensor]: End-effector position,
+            end-effector quaternion (wxyz), linear jacobian(None), angular jacobian(None),
+            link positions, link quaternion (wxyz), link spheres.
+        """
         if len(q.shape) > 2:
-            raise ValueError("q shape should be [batch_size, dof]")
+            log_error("q shape should be [batch_size, dof]")
+        if len(q.shape) == 1:
+            q = q.unsqueeze(0)
         batch_size = q.shape[0]
         self.update_batch_size(batch_size, force_update=q.requires_grad)
 
@@ -287,7 +447,7 @@ class CudaRobotModel(CudaRobotModelConfig):
 
         # compute jacobians?
         if calculate_jacobian:
-            raise NotImplementedError
+            log_error("Outputting jacobian is not supported")
         return (
             ee_pos,
             ee_quat,
@@ -298,7 +458,19 @@ class CudaRobotModel(CudaRobotModelConfig):
             link_spheres_tensor,
         )
 
-    def get_state(self, q, link_name=None, calculate_jacobian=False) -> CudaRobotModelState:
+    def get_state(
+        self, q: torch.Tensor, link_name: str = None, calculate_jacobian: bool = False
+    ) -> CudaRobotModelState:
+        """Get kinematic state of the robot by computing forward kinematics.
+
+        Args:
+            q: Joint configuration of the robot. Shape should be [batch_size, dof].
+            link_name: Name of link to return pose of. If None, returns end-effector pose.
+            calculate_jacobian: Calculate jacobian of the robot. Not supported.
+
+        Returns:
+            CudaRobotModelState: Kinematic state of the robot.
+        """
         out = self.forward(q, link_name, calculate_jacobian)
         state = CudaRobotModelState(
             out[0],
@@ -312,12 +484,46 @@ class CudaRobotModel(CudaRobotModelConfig):
         )
         return state
 
-    def get_robot_link_meshes(self):
+    def compute_kinematics(
+        self, js: JointState, link_name: Optional[str] = None, calculate_jacobian: bool = False
+    ) -> CudaRobotModelState:
+        """Compute forward kinematics of the robot.
+
+        Args:
+            js: Joint state of robot.
+            link_name: Name of link to return pose of. If None, returns end-effector pose.
+            calculate_jacobian: Calculate jacobian of the robot. Not supported.
+
+
+        Returns:
+            CudaRobotModelState: Kinematic state of the robot.
+
+        """
+        if js.joint_names is not None:
+            if js.joint_names != self.kinematics_config.joint_names:
+                log_error("Joint names do not match, reoder joints before forward kinematics")
+
+        return self.get_state(js.position, link_name, calculate_jacobian)
+
+    def get_robot_link_meshes(self) -> List[Mesh]:
+        """Get meshes of all links of the robot.
+
+        Returns:
+            List[Mesh]: List of all link meshes.
+        """
         m_list = [self.get_link_mesh(l) for l in self.kinematics_config.mesh_link_names]
 
         return m_list
 
-    def get_robot_as_mesh(self, q: torch.Tensor):
+    def get_robot_as_mesh(self, q: torch.Tensor) -> List[Mesh]:
+        """Transform robot links to Cartesian poses using forward kinematics and return as meshes.
+
+        Args:
+            q: Joint configuration of the robot, shape should be [1, dof].
+
+        Returns:
+            List[Mesh]: List of all link meshes.
+        """
         # get all link meshes:
         m_list = self.get_robot_link_meshes()
         pose = self.get_link_poses(q, self.kinematics_config.mesh_link_names)
@@ -328,7 +534,16 @@ class CudaRobotModel(CudaRobotModelConfig):
 
         return m_list
 
-    def get_robot_as_spheres(self, q: torch.Tensor, filter_valid: bool = True):
+    def get_robot_as_spheres(self, q: torch.Tensor, filter_valid: bool = True) -> List[Sphere]:
+        """Get robot spheres using forward kinematics on given joint configuration q.
+
+        Args:
+            q: Joint configuration of the robot, shape should be [1, dof].
+            filter_valid: Filter out spheres with radius <= 0.
+
+        Returns:
+            List[Sphere]: List of all robot spheres.
+        """
         state = self.get_state(q)
 
         # state has sphere position and radius
@@ -361,6 +576,18 @@ class CudaRobotModel(CudaRobotModelConfig):
         return sph_traj
 
     def get_link_poses(self, q: torch.Tensor, link_names: List[str]) -> Pose:
+        """Get Pose of links at given joint configuration q using forward kinematics.
+
+        Note that only the links specified in :var:`~link_names` are returned.
+
+        Args:
+            q: Joint configuration of the robot, shape should be [batch_size, dof].
+            link_names: Names of links to get pose of. This should be a subset of
+                :var:`~link_names`.
+
+        Returns:
+            Pose: Poses of links at given joint configuration.
+        """
         state = self.get_state(q)
         position = torch.zeros(
             (q.shape[0], len(link_names), 3),
@@ -379,7 +606,15 @@ class CudaRobotModel(CudaRobotModelConfig):
             quaternion[:, li, :] = state.links_quaternion[:, i, :]
         return Pose(position=position, quaternion=quaternion)
 
-    def _cuda_forward(self, q):
+    def _cuda_forward(self, q: torch.Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        """Compute forward kinematics on GPU. Use :func:`~get_state` or :func:`~forward` instead.
+
+        Args:
+            q: Joint configuration of the robot, shape should be [batch_size, dof].
+
+        Returns:
+            Tuple[Tensor, Tensor, Tensor]: Link positions, link quaternions, link
+        """
         link_pos, link_quat, robot_spheres = get_cuda_kinematics(
             self._link_pos_seq,
             self._link_quat_seq,
@@ -401,22 +636,36 @@ class CudaRobotModel(CudaRobotModelConfig):
         return link_pos, link_quat, robot_spheres
 
     @property
-    def all_articulated_joint_names(self):
+    def all_articulated_joint_names(self) -> List[str]:
+        """Names of all articulated joints of the robot."""
         return self.kinematics_config.non_fixed_joint_names
 
     def get_self_collision_config(self) -> SelfCollisionKinematicsConfig:
+        """Get self collision configuration parameters of the robot."""
         return self.self_collision_config
 
     def get_link_mesh(self, link_name: str) -> Mesh:
+        """Get mesh of a link of the robot."""
         mesh = self.kinematics_parser.get_link_mesh(link_name)
         return mesh
 
     def get_link_transform(self, link_name: str) -> Pose:
-        mat = self.kinematics_config.fixed_transforms[self.kinematics_config.link_name_to_idx_map[link_name]]
+        """Get pose offset of a link from it's parent joint.
+
+        Args:
+            link_name: Name of link to get pose of.
+
+        Returns:
+            Pose: Pose of the link.
+        """
+        mat = self.kinematics_config.fixed_transforms[
+            self.kinematics_config.link_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:
+        """Get offset pose of all links with respect to their parent joint."""
         pose = Pose(
             self.kinematics_config.fixed_transforms[:, :3, 3],
             rotation=self.kinematics_config.fixed_transforms[:, :3, :3],
@@ -424,32 +673,59 @@ class CudaRobotModel(CudaRobotModelConfig):
         return pose
 
     def get_dof(self) -> int:
+        """Get degrees of freedom of the robot."""
         return self.kinematics_config.n_dof
 
     @property
     def dof(self) -> int:
+        """Degrees of freedom of the robot."""
         return self.kinematics_config.n_dof
 
     @property
     def joint_names(self) -> List[str]:
+        """Names of actuated joints."""
         return self.kinematics_config.joint_names
 
     @property
     def total_spheres(self) -> int:
+        """Number of spheres used to approximate robot geometry."""
         return self.kinematics_config.total_spheres
 
     @property
-    def lock_jointstate(self):
+    def lock_jointstate(self) -> JointState:
+        """State of joints that are locked in the kinematic representation."""
         return self.kinematics_config.lock_jointstate
 
-    def get_full_js(self, js: JointState):
+    def get_full_js(self, js: JointState) -> JointState:
+        """Get state of all joints, including locked joints.
+
+        This function will not provide state of mimic joints. If you need mimic joints, use
+        :func:`~get_mimic_js`.
+
+        Args:
+            js: State containing articulated joints.
+
+        Returns:
+            JointState: State of all joints.
+        """
         all_joint_names = self.all_articulated_joint_names
         lock_joint_state = self.lock_jointstate
 
         new_js = js.get_augmented_joint_state(all_joint_names, lock_joint_state)
         return new_js
 
-    def get_mimic_js(self, js: JointState):
+    def get_mimic_js(self, js: JointState) -> JointState:
+        """Get state of mimic joints from active joints.
+
+        Current implementation uses a for loop over joints to calculate the state. This can be
+        optimized by using a custom CUDA kernel or a matrix multiplication.
+
+        Args:
+            js: State containing articulated joints.
+
+        Returns:
+            JointState: State of active, locked, and mimic joints.
+        """
         if self.kinematics_config.mimic_joints is None:
             return None
         extra_joints = {"position": [], "joint_names": []}
@@ -467,21 +743,118 @@ class CudaRobotModel(CudaRobotModelConfig):
         new_js = js.get_augmented_joint_state(js.joint_names + extra_js.joint_names, extra_js)
         return new_js
 
+    def update_kinematics_config(self, new_kin_config: KinematicsTensorConfig):
+        """Update kinematics representation of the robot.
+
+        A kinematics representation can be updated with new parameters. Some parameters that could
+        require updating are state of locked joints, when a robot grasps an object. Another instance
+        is when using different planners for different parts of the robot, example updating the
+        state of robot base or another arm. Updations should result in the same tensor dimensions,
+        if not then the instance of this class requires reinitialization.
+
+        Args:
+            new_kin_config: New kinematics representation of the robot.
+        """
+
+        self.kinematics_config.copy_(new_kin_config)
+
+    def attach_external_objects_to_robot(
+        self,
+        joint_state: JointState,
+        external_objects: List[Obstacle],
+        surface_sphere_radius: float = 0.001,
+        link_name: str = "attached_object",
+        sphere_fit_type: SphereFitType = SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
+        voxelize_method: str = "ray",
+        world_objects_pose_offset: Optional[Pose] = None,
+    ) -> bool:
+        """Attach external objects to a robot's link. See :ref:`attach_object_note` for details.
+
+        Args:
+            joint_state: Joint state of the robot.
+            external_objects: List of external objects to attach to the robot.
+            surface_sphere_radius: Radius (in meters) to use for points sampled on surface of the
+                object. A smaller radius will allow for generating motions very close to obstacles.
+            link_name: Name of the link (frame) to attach the objects to. The assumption is that
+                this link does not have any geometry and all spheres of this link represent
+                attached objects.
+            sphere_fit_type: Sphere fit algorithm to use. See :ref:`attach_object_note` for more
+                details. The default method :attr:`SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE`
+                voxelizes the volume of the objects and adds spheres representing the voxels, then
+                samples points on the surface of the object, adds :attr:`surface_sphere_radius` to
+                these points. This should be used for most cases.
+            voxelize_method: Method to use for voxelization, passed to
+                :py:func:`trimesh.voxel.creation.voxelize`.
+            world_objects_pose_offset: Offset to apply to the object poses before attaching to the
+                robot. This is useful when attaching an object that's in contact with the world.
+                The offset is applied in the world frame before attaching to the robot.
+        """
+        log_info("Attach objects to robot")
+        if len(external_objects) == 0:
+            log_error("no object in external_objects")
+        kin_state = self.compute_kinematics(joint_state)
+        ee_pose = kin_state.ee_pose  # w_T_ee
+        if world_objects_pose_offset is not None:
+            # add offset from ee:
+            ee_pose = world_objects_pose_offset.inverse().multiply(ee_pose)
+            # new ee_pose:
+            # w_T_ee = offset_T_w * w_T_ee
+            # ee_T_w
+        ee_pose = ee_pose.inverse()  # ee_T_w to multiply all objects later
+        max_spheres = self.kinematics_config.get_number_of_spheres(link_name)
+        object_names = [x.name for x in external_objects]
+        n_spheres = int(max_spheres / len(object_names))
+        sphere_tensor = torch.zeros((max_spheres, 4))
+        sphere_tensor[:, 3] = -10.0
+        sph_list = []
+        if n_spheres == 0:
+            log_warn(
+                "No spheres found, max_spheres: "
+                + str(max_spheres)
+                + " n_objects: "
+                + str(len(object_names))
+            )
+            return False
+        for i, x in enumerate(object_names):
+            obs = external_objects[i]
+            sph = obs.get_bounding_spheres(
+                n_spheres,
+                surface_sphere_radius,
+                pre_transform_pose=ee_pose,
+                tensor_args=self.tensor_args,
+                fit_type=sphere_fit_type,
+                voxelize_method=voxelize_method,
+            )
+            sph_list += [s.position + [s.radius] for s in sph]
+
+        log_info("MG: Computed spheres for attach objects to robot")
+
+        spheres = self.tensor_args.to_device(torch.as_tensor(sph_list))
+
+        if spheres.shape[0] > max_spheres:
+            spheres = spheres[: spheres.shape[0]]
+        sphere_tensor[: spheres.shape[0], :] = spheres.contiguous()
+
+        self.kinematics_config.attach_object(sphere_tensor=sphere_tensor, link_name=link_name)
+
+        return True
+
     @property
-    def ee_link(self):
+    def ee_link(self) -> str:
+        """End-effector link of the robot. Changing requires reinitializing this class."""
         return self.kinematics_config.ee_link
 
     @property
-    def base_link(self):
+    def base_link(self) -> str:
+        """Base link of the robot. Changing requires reinitializing this class."""
         return self.kinematics_config.base_link
 
     @property
     def robot_spheres(self):
+        """Spheres representing robot geometry."""
         return self.kinematics_config.link_spheres
 
-    def update_kinematics_config(self, new_kin_config: KinematicsTensorConfig):
-        self.kinematics_config.copy_(new_kin_config)
-
     @property
-    def retract_config(self):
+    def retract_config(self) -> torch.Tensor:
+        """Retract configuration of the robot. Use :func:`~joint_names` to get joint names."""
         return self.kinematics_config.cspace.retract_config

src/curobo/cuda_robot_model/kinematics_parser.py

@@ -8,10 +8,21 @@
 # without an express license agreement from NVIDIA CORPORATION or
 # its affiliates is strictly prohibited.
 #
+"""
+Base module for parsing kinematics from different representations.
+
+cuRobo provides kinematics parsing from an URDF and a partial implementation for parsing from
+a USD. To parse from other representations, an user can extend the :class:`~KinematicsParser`
+class and implement only the abstract methods. Optionally, user can also provide functions for
+reading meshes, useful for debugging and visualization.
+
+"""
+from __future__ import annotations
 
 # Standard Library
+from abc import abstractmethod
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 # Third Party
 import numpy as np
@@ -25,6 +36,8 @@ from curobo.types.math import Pose
 
 @dataclass
 class LinkParams:
+    """Parameters of a link in the kinematic tree."""
+
     link_name: str
     joint_name: str
     joint_type: JointType
@@ -38,7 +51,15 @@ class LinkParams:
     mimic_joint_name: Optional[str] = None
 
     @staticmethod
-    def from_dict(dict_data):
+    def from_dict(dict_data: Dict[str, Any]) -> LinkParams:
+        """Create a LinkParams object from a dictionary.
+
+        Args:
+            dict_data: Dictionary containing link parameters.
+
+        Returns:
+            LinkParams: Link parameters object.
+        """
         dict_data["joint_type"] = JointType[dict_data["joint_type"]]
         dict_data["fixed_transform"] = (
             Pose.from_list(dict_data["fixed_transform"], tensor_args=TensorDeviceType())
@@ -50,7 +71,21 @@ class LinkParams:
 
 
 class KinematicsParser:
+    """
+    Base class for parsing kinematics.
+
+    Implement abstractmethods to parse kinematics from any representation. Optionally, implement
+    methods for reading meshes for visualization and debugging.
+    """
+
     def __init__(self, extra_links: Optional[Dict[str, LinkParams]] = None) -> None:
+        """Initialize the KinematicsParser.
+
+        Args:
+            extra_links: Additional links to be added to the kinematic tree.
+        """
+
+        #: Parent link for all link in the kinematic tree.
         self._parent_map = {}
         self.extra_links = extra_links
         self.build_link_parent()
@@ -59,10 +94,45 @@ class KinematicsParser:
             for i in self.extra_links:
                 self._parent_map[i] = {"parent": self.extra_links[i].parent_link_name}
 
+    @abstractmethod
     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.
+        Use this function to fill ``_parent_map``. Check
+        :meth:`curobo.cuda_robot_model.urdf_kinematics_parser.UrdfKinematicsParser.build_link_parent`
+        for an example implementation.
+        """
+        pass
+
+    @abstractmethod
+    def get_link_parameters(self, link_name: str, base: bool = False) -> LinkParams:
+        """Get parameters of a link in the kinematic tree.
+
+        Args:
+            link_name: Name of the link.
+            base: Is this the base link of the robot?
+
+        Returns:
+            LinkParams: Parameters of the link.
+        """
+        pass
+
+    def add_absolute_path_to_link_meshes(self, mesh_dir: str = ""):
+        """Add absolute path to link meshes.
+
+        Args:
+            mesh_dir: Absolute path to the directory containing link meshes.
+        """
+        pass
+
+    def get_link_mesh(self, link_name: str) -> Mesh:
+        """Get mesh of a link.
+
+        Args:
+            link_name: Name of the link.
+
+        Returns:
+            Mesh: Mesh of the link.
         """
         pass
 
@@ -85,18 +155,31 @@ class KinematicsParser:
         chain_links.reverse()
         return chain_links
 
+    def get_controlled_joint_names(self) -> List[str]:
+        """Get names of all controlled joints in the robot.
+
+        Returns:
+            Names of all controlled joints in the robot.
+        """
+        j_list = []
+        for k in self._parent_map.keys():
+            joint_name = self._parent_map[k]["joint_name"]
+            joint = self._robot.joint_map[joint_name]
+            if joint.type != "fixed" and joint.mimic is None:
+                j_list.append(joint_name)
+        return j_list
+
     def _get_from_extra_links(self, link_name: str) -> LinkParams:
+        """Get link parameters for extra links.
+
+        Args:
+            link_name: Name of the link.
+
+        Returns:
+            LinkParams: Link parameters if found, else None.
+        """
         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

@@ -8,6 +8,8 @@
 # without an express license agreement from NVIDIA CORPORATION or
 # its affiliates is strictly prohibited.
 #
+"""Common structures used for Kinematics are defined in this module."""
+
 from __future__ import annotations
 
 # Standard Library
@@ -28,33 +30,87 @@ from curobo.util.tensor_util import clone_if_not_none, copy_if_not_none
 
 
 class JointType(Enum):
+    """Type of Joint. Arbitrary axis of change is not supported."""
+
+    #: Fixed joint.
     FIXED = -1
+
+    #: Prismatic joint along x-axis.
     X_PRISM = 0
+
+    #: Prismatic joint along y-axis.
     Y_PRISM = 1
+
+    #: Prismatic joint along z-axis.
     Z_PRISM = 2
+
+    #: Revolute joint along x-axis.
     X_ROT = 3
+
+    #: Revolute joint along y-axis.
     Y_ROT = 4
+
+    #: Revolute joint along z-axis.
     Z_ROT = 5
+
+    #: Prismatic joint along negative x-axis.
     X_PRISM_NEG = 6
+
+    #: Prismatic joint along negative y-axis.
     Y_PRISM_NEG = 7
+
+    #: Prismatic joint along negative z-axis.
     Z_PRISM_NEG = 8
+
+    #: Revolute joint along negative x-axis.
     X_ROT_NEG = 9
+
+    #: Revolute joint along negative y-axis.
     Y_ROT_NEG = 10
+
+    #: Revolute joint along negative z-axis.
     Z_ROT_NEG = 11
 
 
 @dataclass
 class JointLimits:
+    """Joint limits for a robot."""
+
+    #: Names of the joints. All tensors are indexed by joint names.
     joint_names: List[str]
+
+    #: Position limits for each joint. Shape [n_joints, 2] with columns having [min, max] values.
     position: torch.Tensor
+
+    #: Velocity limits for each joint. Shape [n_joints, 2] with columns having [min, max] values.
     velocity: torch.Tensor
+
+    #: Acceleration limits for each joint. Shape [n_joints, 2] with columns having [min, max]
+    #: values.
     acceleration: torch.Tensor
+
+    #: Jerk limits for each joint. Shape [n_joints, 2] with columns having [min, max] values.
     jerk: torch.Tensor
+
+    #: Effort limits for each joint. This is not used.
     effort: Optional[torch.Tensor] = None
+
+    #: Device and floating point precision for tensors.
     tensor_args: TensorDeviceType = TensorDeviceType()
 
     @staticmethod
-    def from_data_dict(data: Dict, tensor_args: TensorDeviceType = TensorDeviceType()):
+    def from_data_dict(
+        data: Dict, tensor_args: TensorDeviceType = TensorDeviceType()
+    ) -> JointLimits:
+        """Create JointLimits from a dictionary.
+
+        Args:
+            data: Dictionary containing joint limits. E.g., {"position": [0, 1], ...}.
+            tensor_args: Device and floating point precision for tensors.
+
+        Returns:
+            JointLimits: Joint limits instance.
+        """
         p = tensor_args.to_device(data["position"])
         v = tensor_args.to_device(data["velocity"])
         a = tensor_args.to_device(data["acceleration"])
@@ -66,6 +122,7 @@ class JointLimits:
         return JointLimits(data["joint_names"], p, v, a, j, e)
 
     def clone(self) -> JointLimits:
+        """Clone joint limits."""
         return JointLimits(
             self.joint_names.copy(),
             self.position.clone(),
@@ -76,7 +133,15 @@ class JointLimits:
             self.tensor_args,
         )
 
-    def copy_(self, new_jl: JointLimits):
+    def copy_(self, new_jl: JointLimits) -> JointLimits:
+        """Copy joint limits from another instance. This maintains reference and copies the data.
+
+        Args:
+            new_jl: JointLimits instance to copy from.
+
+        Returns:
+            JointLimits: Data copied joint limits.
+        """
         self.joint_names = new_jl.joint_names.copy()
         self.position.copy_(new_jl.position)
         self.velocity.copy_(new_jl.velocity)
@@ -87,19 +152,53 @@ class JointLimits:
 
 @dataclass
 class CSpaceConfig:
+    """Configuration space parameters of the robot."""
+
+    #: Names of the joints.
     joint_names: List[str]
+
+    #: Retract configuration for the robot. This is the configuration used to bias graph search
+    #: and also regularize inverse kinematics. This configuration is also used to initialize
+    #: the robot during warmup phase of an optimizer. Set this to a collision-free configuration
+    #: for good performance. When this configuration is in collision, it's not used to bias
+    #: graph search.
     retract_config: Optional[T_DOF] = None
+
+    #: Weight for each joint in configuration space. Used to measure distance between nodes in
+    #: graph search-based planning.
     cspace_distance_weight: Optional[T_DOF] = None
+
+    #: Weight for each joint, used in regularization cost term for inverse kinematics.
     null_space_weight: Optional[T_DOF] = None
+
+    #: Device and floating point precision for tensors.
     tensor_args: TensorDeviceType = TensorDeviceType()
+
+    #: Maximum acceleration for each joint. Accepts a scalar or a list of values for each joint.
     max_acceleration: Union[float, List[float]] = 10.0
+
+    #: Maximum jerk for each joint. Accepts a scalar or a list of values for each joint.
     max_jerk: Union[float, List[float]] = 500.0
+
+    #: Velocity scale for each joint. Accepts a scalar or a list of values for each joint.
+    #: This is used to scale the velocity limits for each joint.
     velocity_scale: Union[float, List[float]] = 1.0
+
+    #: Acceleration scale for each joint. Accepts a scalar or a list of values for each joint.
+    #: This is used to scale the acceleration limits for each joint.
     acceleration_scale: Union[float, List[float]] = 1.0
+
+    #: Jerk scale for each joint. Accepts a scalar or a list of values for each joint.
+    #: This is used to scale the jerk limits for each joint.
     jerk_scale: Union[float, List[float]] = 1.0
+
+    #: Position limit clip value. This is used to clip the position limits for each joint.
+    #: Accepts a scalar or a list of values for each joint. This is useful to truncate limits
+    #: to account for any safety margins imposed by real robot controllers.
     position_limit_clip: Union[float, List[float]] = 0.0
 
     def __post_init__(self):
+        """Post initialization checks and data transfer to device tensors."""
         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:
@@ -139,6 +238,8 @@ class CSpaceConfig:
             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)
+        if isinstance(self.position_limit_clip, List):
+            self.position_limit_clip = self.tensor_args.to_device(self.position_limit_clip)
         # check shapes:
         if self.retract_config is not None:
             dof = self.retract_config.shape
@@ -148,6 +249,12 @@ class CSpaceConfig:
                 log_error("null_space_weight shape does not match retract_config")
 
     def inplace_reindex(self, joint_names: List[str]):
+        """Change order of joints in configuration space tensors to match given order of names.
+
+        Args:
+            joint_names: New order of joint names.
+
+        """
         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()
@@ -163,7 +270,18 @@ class CSpaceConfig:
         joint_names = [self.joint_names[n] for n in new_index]
         self.joint_names = joint_names
 
-    def copy_(self, new_config: CSpaceConfig):
+    def copy_(self, new_config: CSpaceConfig) -> CSpaceConfig:
+        """Copy parameters from another instance.
+
+        This maintains reference and copies the data. Assumes that the new instance has the same
+        number of joints as the current instance and also same shape of tensors.
+
+        Args:
+            new_config: New parameters to copy into current instance.
+
+        Returns:
+            CSpaceConfig: Same instance of cspace configuration which has updated parameters.
+        """
         self.joint_names = new_config.joint_names.copy()
         self.retract_config = copy_if_not_none(new_config.retract_config, self.retract_config)
         self.null_space_weight = copy_if_not_none(
@@ -183,6 +301,8 @@ class CSpaceConfig:
         return self
 
     def clone(self) -> CSpaceConfig:
+        """Clone configuration space parameters."""
+
         return CSpaceConfig(
             joint_names=self.joint_names.copy(),
             retract_config=clone_if_not_none(self.retract_config),
@@ -194,9 +314,22 @@ class CSpaceConfig:
             velocity_scale=self.velocity_scale.clone(),
             acceleration_scale=self.acceleration_scale.clone(),
             jerk_scale=self.jerk_scale.clone(),
+            position_limit_clip=(
+                self.position_limit_clip.clone()
+                if isinstance(self.position_limit_clip, torch.Tensor)
+                else self.position_limit_clip
+            ),
         )
 
-    def scale_joint_limits(self, joint_limits: JointLimits):
+    def scale_joint_limits(self, joint_limits: JointLimits) -> JointLimits:
+        """Scale joint limits by the given scale factors.
+
+        Args:
+            joint_limits: Joint limits to scale.
+
+        Returns:
+            JointLimits: Scaled joint limits.
+        """
         if self.velocity_scale is not None:
             joint_limits.velocity = joint_limits.velocity * self.velocity_scale
         if self.acceleration_scale is not None:
@@ -212,7 +345,20 @@ class CSpaceConfig:
         joint_position_lower: torch.Tensor,
         joint_names: List[str],
         tensor_args: TensorDeviceType = TensorDeviceType(),
-    ):
+    ) -> CSpaceConfig:
+        """Load CSpace configuration from joint limits.
+
+        Args:
+            joint_position_upper: Upper position limits for each joint.
+            joint_position_lower: Lower position limits for each joint.
+            joint_names: Names of the joints. This should match the order of joints in the upper
+                and lower limits.
+            tensor_args: Device and floating point precision for tensors.
+
+        Returns:
+            CSpaceConfig: CSpace configuration with retract configuration set to the middle of the
+                joint limits and all weights set to 1.
+        """
         retract_config = ((joint_position_upper + joint_position_lower) / 2).flatten()
         n_dof = retract_config.shape[-1]
         null_space_weight = torch.ones(n_dof, **(tensor_args.as_torch_dict()))
@@ -238,25 +384,26 @@ class KinematicsTensorConfig:
     #: Static Homogenous Transform from parent link to child link for all links [n_links,4,4].
     fixed_transforms: torch.Tensor
 
-    #: index of fixed_transform given link index [n_links].
+    #: Index of fixed_transform given link index [n_links].
     link_map: torch.Tensor
 
-    #: joint index given link index [n_links].
+    #: Joint index given link index [n_links].
     joint_map: torch.Tensor
 
-    #: type of joint given link index [n_links].
+    #: Type of joint given link index [n_links].
     joint_map_type: torch.Tensor
 
+    #: Joint offset to store scalars for mimic joints and negative axis joints.
     joint_offset_map: torch.Tensor
 
-    #: index of link to write out pose [n_store_links].
+    #: Index of link to write out pose [n_store_links].
     store_link_map: torch.Tensor
 
     #: Mapping between each link to every other link, this is used to check
     #: if a link is part of a serial chain formed by another link [n_links, n_links].
     link_chain_map: torch.Tensor
 
-    #: Name of links whose pose will be stored [n_store_links].
+    #: Name of links to compute pose during kinematics computation [n_store_links].
     link_names: List[str]
 
     #: Joint limits
@@ -274,31 +421,32 @@ class KinematicsTensorConfig:
     #: Name of all actuated joints.
     joint_names: Optional[List[str]] = None
 
-    #:
+    #: Name of joints to lock to a fixed value along with the locked value
     lock_jointstate: Optional[JointState] = None
 
-    #:
+    #: Joints that mimic other joints. This will be populated by :class:~`CudaRobotGenerator`
+    #  when parsing the kinematics of the robot.
     mimic_joints: Optional[dict] = None
 
-    #:
+    #: Sphere representation of the robot's geometry. This is used for collision detection.
     link_spheres: Optional[torch.Tensor] = None
 
-    #:
+    #: Mapping of link index to sphere index. This is used to get spheres for a link.
     link_sphere_idx_map: Optional[torch.Tensor] = None
 
-    #:
+    #: Mapping of link name to link index. This is used to get link index from link name.
     link_name_to_idx_map: Optional[Dict[str, int]] = None
 
-    #: total number of spheres that represent the robot's geometry.
+    #: Total number of spheres that represent the robot's geometry.
     total_spheres: int = 0
 
     #: Additional debug parameters.
     debug: Optional[Any] = None
 
-    #: index of end-effector in stored link poses.
+    #: Index of end-effector in stored link poses.
     ee_idx: int = 0
 
-    #: Cspace configuration
+    #: Cspace parameters for the robot.
     cspace: Optional[CSpaceConfig] = None
 
     #: Name of base link. This is the root link from which all kinematic parameters were computed.
@@ -312,6 +460,7 @@ class KinematicsTensorConfig:
     reference_link_spheres: Optional[torch.Tensor] = None
 
     def __post_init__(self):
+        """Post initialization checks and data transfer to device tensors."""
         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:
@@ -320,6 +469,18 @@ class KinematicsTensorConfig:
             self.reference_link_spheres = self.link_spheres.clone()
 
     def copy_(self, new_config: KinematicsTensorConfig) -> KinematicsTensorConfig:
+        """Copy parameters from another instance into current instance.
+
+        This maintains reference and copies the data. Assumes that the new instance has the same
+        number of joints as the current instance and also same shape of tensors.
+
+        Args:
+            new_config: New parameters to copy into current instance.
+
+        Returns:
+            KinematicsTensorConfig: Same instance of kinematics configuration which has updated
+                parameters.
+        """
         self.fixed_transforms.copy_(new_config.fixed_transforms)
         self.link_map.copy_(new_config.link_map)
         self.joint_map.copy_(new_config.joint_map)
@@ -350,10 +511,17 @@ class KinematicsTensorConfig:
         self.link_names = new_config.link_names
         self.mesh_link_names = new_config.mesh_link_names
         self.total_spheres = new_config.total_spheres
+        if self.lock_jointstate is not None and new_config.lock_jointstate is not None:
+            self.lock_jointstate.copy_(new_config.lock_jointstate)
+        self.mimic_joints = new_config.mimic_joints
 
         return self
 
     def load_cspace_cfg_from_kinematics(self):
+        """Load CSpace configuration from joint limits.
+
+        This sets the retract configuration to the middle of the joint limits and all weights to 1.
+        """
         retract_config = (
             (self.joint_limits.position[1] + self.joint_limits.position[0]) / 2
         ).flatten()
@@ -371,22 +539,27 @@ class KinematicsTensorConfig:
         )
 
     def get_sphere_index_from_link_name(self, link_name: str) -> torch.Tensor:
+        """Get indices of spheres for a link.
+
+        Args:
+            link_name: Name of the link.
+
+        Returns:
+            torch.Tensor: Indices of spheres for the link.
+        """
         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.
+        """Update sphere parameters of a specific link given by name.
 
         Args:
-            link_name: _description_
-            sphere_position_radius: _description_
-            start_sph_idx: _description_. Defaults to 0.
+            link_name: Name of the link.
+            sphere_position_radius: Tensor of shape [n_spheres, 4] with columns [x, y, z, r].
+            start_sph_idx: If providing a subset of spheres, this is the starting index.
         """
         # get sphere indices from link name:
         link_sphere_index = self.get_sphere_index_from_link_name(link_name)[
@@ -398,14 +571,31 @@ class KinematicsTensorConfig:
     def get_link_spheres(
         self,
         link_name: str,
-    ):
+    ) -> torch.Tensor:
+        """Get spheres of a link.
+
+        Args:
+            link_name: Name of link.
+
+        Returns:
+            torch.Tensor: Spheres of the link with shape [n_spheres, 4] with columns [x, y, z, r].
+        """
         link_sphere_index = self.get_sphere_index_from_link_name(link_name)
         return self.link_spheres[link_sphere_index, :]
 
     def get_reference_link_spheres(
         self,
         link_name: str,
-    ):
+    ) -> torch.Tensor:
+        """Get link spheres from the original robot configuration data before any modifications.
+
+        Args:
+            link_name: Name of link.
+
+        Returns:
+            torch.Tensor: Spheres of the link with shape [n_spheres, 4] with columns [x, y, z, r].
+        """
+
         link_sphere_index = self.get_sphere_index_from_link_name(link_name)
         return self.reference_link_spheres[link_sphere_index, :]
 
@@ -415,47 +605,45 @@ class KinematicsTensorConfig:
         sphere_tensor: Optional[torch.Tensor] = None,
         link_name: str = "attached_object",
     ) -> bool:
-        """_summary_
+        """Attach object approximated by spheres to a link of the robot.
+
+        This function updates the sphere parameters of the link to represent the attached object.
 
         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_
+            sphere_radius: Radius to change for existing spheres. If changing position of spheres
+                as well, then set this to None.
+            sphere_tensor: New sphere tensor to replace existing spheres. Shape [n_spheres, 4] with
+                columns [x, y, z, r]. If changing only radius, set this to None and use
+                sphere_radius.
+            link_name: Name of the link to attach object to. Defaults to "attached_object".
 
         Returns:
-            _description_
+            bool: True if successful.
         """
         if link_name not in self.link_name_to_idx_map.keys():
-            raise ValueError(link_name + " not found in spheres")
+            log_error(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")
+                log_error("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
+        """Detach object spheres from a link by setting all spheres to zero with negative radius.
 
         Args:
-            link_name: _description_. Defaults to "attached_object".
-
-        Raises:
-            ValueError: _description_
+            link_name: Name of the link to detach object from.
 
         Returns:
-            _description_
+            bool: True if successful.
         """
         if link_name not in self.link_name_to_idx_map.keys():
-            raise ValueError(link_name + " not found in spheres")
+            log_error(link_name + " not found in spheres")
         curr_spheres = self.get_link_spheres(link_name)
         curr_spheres[:, 3] = -100.0
         curr_spheres[:, :3] = 0.0
@@ -472,15 +660,25 @@ class KinematicsTensorConfig:
         return self.get_link_spheres(link_name).shape[0]
 
     def disable_link_spheres(self, link_name: str):
+        """Disable spheres of a link by setting all spheres to zero with negative radius.
+
+        Args:
+            link_name: Name of the link to disable spheres.
+        """
         if link_name not in self.link_name_to_idx_map.keys():
-            raise ValueError(link_name + " not found in spheres")
+            log_error(link_name + " not found in spheres")
         curr_spheres = self.get_link_spheres(link_name)
         curr_spheres[:, 3] = -100.0
         self.update_link_spheres(link_name, curr_spheres)
 
     def enable_link_spheres(self, link_name: str):
+        """Enable spheres of a link by resetting to values from initial robot configuration data.
+
+        Args:
+            link_name: Name of the link to enable spheres.
+        """
         if link_name not in self.link_name_to_idx_map.keys():
-            raise ValueError(link_name + " not found in spheres")
+            log_error(link_name + " not found in spheres")
         curr_spheres = self.get_reference_link_spheres(link_name)
         self.update_link_spheres(link_name, curr_spheres)
 
@@ -489,10 +687,29 @@ class KinematicsTensorConfig:
 class SelfCollisionKinematicsConfig:
     """Dataclass that stores self collision attributes to pass to cuda kernel."""
 
+    #: Offset radii for each sphere. This is used to inflate the spheres for self collision
+    #: detection.
     offset: Optional[torch.Tensor] = None
-    distance_threshold: Optional[torch.Tensor] = None
+
+    #: Sphere index to use for a given thread.
     thread_location: Optional[torch.Tensor] = None
+
+    #: Maximum number of threads to launch for computing self collision between spheres.
     thread_max: Optional[int] = None
-    collision_matrix: Optional[torch.Tensor] = None
+
+    #: Distance threshold for self collision detection. This is currently not used.
+    distance_threshold: Optional[torch.Tensor] = None
+
+    #: Two kernel implementations are available. Set this to True to use the experimental kernel
+    #: which is faster. Set this to False to use the collision matrix based kernel which is slower.
     experimental_kernel: bool = True
+
+    #: Collision matrix containing information about which pair of spheres need to be checked for
+    #: collision. This is only used when experimental_kernel is set to False.
+    collision_matrix: Optional[torch.Tensor] = None
+
+    #: Number of collision checks to perform per thread. Each thread loads a sphere and is allowed
+    #: to check upto checks_per_thread other spheres for collision. Note that all checks have to
+    #: be performed within 1024 threads as shared memory is used. So,
+    # checks_per_thread * n_spheres <= 1024.
     checks_per_thread: int = 32

src/curobo/cuda_robot_model/urdf_kinematics_parser.py

@@ -8,8 +8,13 @@
 # without an express license agreement from NVIDIA CORPORATION or
 # its affiliates is strictly prohibited.
 #
+"""
+Parses Kinematics from an `URDF <https://wiki.ros.org/urdf>`__ file that describes the
+kinematic tree of a robot.
+"""
+
 # Standard Library
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 # Third Party
 import numpy as np
@@ -20,13 +25,14 @@ from lxml import etree
 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.types.math import Pose
 from curobo.util.logger import log_error, log_warn
 from curobo.util_file import join_path
 
 
 class UrdfKinematicsParser(KinematicsParser):
+    """Parses Kinematics from an URDF file and provides access to the kinematic tree."""
+
     def __init__(
         self,
         urdf_path,
@@ -35,6 +41,16 @@ class UrdfKinematicsParser(KinematicsParser):
         extra_links: Optional[Dict[str, LinkParams]] = None,
         build_scene_graph: bool = False,
     ) -> None:
+        """Initialize instance with URDF file path.
+
+        Args:
+            urdf_path: Path to the URDF file.
+            load_meshes: Load meshes for links from the URDF file.
+            mesh_root: Absolute path to the directory where link meshes are stored.
+            extra_links: Extra links to add to the kinematic tree.
+            build_scene_graph: Build scene graph for the robot. Set this to True if you want to
+                determine the root link of the robot.
+        """
         # load robot from urdf:
         self._robot = yourdfpy.URDF.load(
             urdf_path,
@@ -46,6 +62,7 @@ class UrdfKinematicsParser(KinematicsParser):
         super().__init__(extra_links)
 
     def build_link_parent(self):
+        """Build parent map for the robot."""
         self._parent_map = {}
         for jid, j in enumerate(self._robot.joint_map):
             self._parent_map[self._robot.joint_map[j].child] = {
@@ -54,11 +71,31 @@ class UrdfKinematicsParser(KinematicsParser):
                 "joint_name": j,
             }
 
-    def _get_joint_name(self, idx):
+    def _get_joint_name(self, idx) -> str:
+        """Get the name of the joint at the given index.
+
+        Args:
+            idx: Index of the joint.
+
+        Returns:
+            str: Name of the joint.
+        """
         joint = self._robot.joint_names[idx]
         return joint
 
-    def _get_joint_limits(self, joint):
+    def _get_joint_limits(self, joint: yourdfpy.Joint) -> Tuple[Dict[str, float], str]:
+        """Get the limits of a joint.
+
+        This function converts continuous joints to revolute joints with limits [-6.28, 6.28].
+
+        Args:
+            joint: Instance of the joint.
+
+        Returns:
+            Tuple[Dict[str, float], str]: Limits of the joint and the type of the joint
+                (revolute or prismatic).
+        """
+
         joint_type = joint.type
         if joint_type != "continuous":
             joint_limits = {
@@ -79,6 +116,16 @@ class UrdfKinematicsParser(KinematicsParser):
         return joint_limits, joint_type
 
     def get_link_parameters(self, link_name: str, base=False) -> LinkParams:
+        """Get parameters of a link in the kinematic tree.
+
+        Args:
+            link_name: Name of the link.
+            base: Is this the base link of the robot?
+
+        Returns:
+            LinkParams: Parameters of the link.
+        """
+
         link_params = self._get_from_extra_links(link_name)
         if link_params is not None:
             return link_params
@@ -122,7 +169,7 @@ class UrdfKinematicsParser(KinematicsParser):
                     mimic_joint_name = joint_name
                     active_joint_name = joint.mimic.joint
                     active_joint = self._robot.joint_map[active_joint_name]
-                    active_joint_limits, active_joint_type = self._get_joint_limits(active_joint)
+                    active_joint_limits, _ = self._get_joint_limits(active_joint)
                     # check to make sure mimic joint limits are not larger than active joint:
                     if (
                         active_joint_limits["lower"] * joint_offset[0] + joint_offset[1]
@@ -196,6 +243,11 @@ class UrdfKinematicsParser(KinematicsParser):
         return link_params
 
     def add_absolute_path_to_link_meshes(self, mesh_dir: str = ""):
+        """Add absolute path to link meshes.
+
+        Args:
+            mesh_dir: Absolute path to the directory containing link meshes.
+        """
         # read all link meshes and update their mesh paths by prepending mesh_dir
         links = self._robot.link_map
         for k in links.keys():
@@ -211,11 +263,21 @@ class UrdfKinematicsParser(KinematicsParser):
                 if m is not None:
                     m.filename = join_path(mesh_dir, m.filename)
 
-    def get_urdf_string(self):
+    def get_urdf_string(self) -> str:
+        """Get the contents of URDF as a string."""
         txt = etree.tostring(self._robot.write_xml(), method="xml", encoding="unicode")
         return txt
 
-    def get_link_mesh(self, link_name):
+    def get_link_mesh(self, link_name: str) -> CuroboMesh:
+        """Get mesh of a link.
+
+        Args:
+            link_name: Name of the link.
+
+        Returns:
+            Mesh: Mesh of the link.
+        """
+
         link_data = self._robot.link_map[link_name]
 
         if len(link_data.visuals) == 0:
@@ -237,14 +299,12 @@ class UrdfKinematicsParser(KinematicsParser):
         )
 
     @property
-    def root_link(self):
-        return self._robot.base_link
+    def root_link(self) -> str:
+        """Returns the name of the base link of the robot.
 
-    def get_controlled_joint_names(self) -> List[str]:
-        j_list = []
-        for k in self._parent_map.keys():
-            joint_name = self._parent_map[k]["joint_name"]
-            joint = self._robot.joint_map[joint_name]
-            if joint.type != "fixed" and joint.mimic is None:
-                j_list.append(joint_name)
-        return j_list
+        Only available when the URDF is loaded with build_scene_graph=True.
+
+        Returns:
+            str: Name of the base link.
+        """
+        return self._robot.base_link

src/curobo/cuda_robot_model/usd_kinematics_parser.py

@@ -8,6 +8,15 @@
 # without an express license agreement from NVIDIA CORPORATION or
 # its affiliates is strictly prohibited.
 #
+"""
+An experimental kinematics parser that reads the robot from a USD file or stage.
+
+Basic loading of simple robots should work but more complex robots may not be supported. E.g.,
+mimic joints cannot be parsed correctly. Use the URDF parser (:class:`~UrdfKinematicsParser`)
+for more complex robots.
+"""
+
+
 # Standard Library
 from typing import Dict, List, Optional, Tuple
 
@@ -33,8 +42,10 @@ except ImportError:
 
 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.
+
+    Current implementation does not account for link geometry transformations after a joints.
+    Also, cannot read mimic joints.
+
     """
 
     def __init__(
@@ -46,7 +57,17 @@ class UsdKinematicsParser(KinematicsParser):
         tensor_args: TensorDeviceType = TensorDeviceType(),
         extra_links: Optional[Dict[str, LinkParams]] = None,
     ) -> None:
-        # load usd file:
+        """Initialize instance with USD file path.
+
+        Args:
+            usd_path: path to usd reference. This will opened as a Usd Stage.
+            flip_joints: list of joint names to flip axis. This is required as current
+                implementation does not read transformations from joint to link correctly.
+            flip_joint_limits: list of joint names to flip joint limits.
+            usd_robot_root: Root prim of the robot in the Usd Stage.
+            tensor_args: Device and floating point precision for tensors.
+            extra_links: Additional links to add to the robot kinematics structure.
+        """
 
         # create a usd stage
         self._flip_joints = flip_joints
@@ -59,9 +80,11 @@ class UsdKinematicsParser(KinematicsParser):
 
     @property
     def robot_prim_root(self):
+        """Root prim of the robot in the Usd Stage."""
         return self._usd_robot_root
 
     def build_link_parent(self):
+        """Build a dictionary containing parent link for each link in the robot."""
         self._parent_map = {}
         all_joints = [
             x
@@ -77,16 +100,16 @@ class UsdKinematicsParser(KinematicsParser):
     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,
+        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.
+            base (bool, optional): Is this the base link of the robot?
 
         Returns:
-            LinkParams: obtained link parameters.
+            LinkParams: Obtained link parameters.
         """
         link_params = self._get_from_extra_links(link_name)
         if link_params is not None:
@@ -157,7 +180,15 @@ class UsdKinematicsParser(KinematicsParser):
         )
         return link_params
 
-    def _get_joint_transform(self, prim: Usd.Prim):
+    def _get_joint_transform(self, prim: Usd.Prim) -> Pose:
+        """Get pose of link from joint prim.
+
+        Args:
+            prim: joint prim in the usd stage.
+
+        Returns:
+            Pose: pose of the link from joint origin.
+        """
         j_prim = UsdPhysics.Joint(prim)
         position = np.ravel(j_prim.GetLocalPos0Attr().Get())
         quatf = j_prim.GetLocalRot0Attr().Get()
@@ -188,10 +219,16 @@ class UsdKinematicsParser(KinematicsParser):
 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`.
+
+    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`.
+
+    Args:
+        prim: joint prim in the usd stage.
+
+    Returns:
+        Tuple[Optional[Usd.Prim], Optional[Usd.Prim]]: parent link prim and child link prim.
     """
     stage = prim.GetStage()
     joint_api = UsdPhysics.Joint(prim)

src/curobo/cuda_robot_model/util.py

@@ -0,0 +1,63 @@
+#
+# 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 a function to load robot representation from a yaml or xrdf file. """
+
+# Standard Library
+from typing import Optional, Union
+
+# CuRobo
+from curobo.types.file_path import ContentPath
+from curobo.util.logger import log_error
+from curobo.util.xrdf_utils import convert_xrdf_to_curobo
+from curobo.util_file import join_path, load_yaml
+
+
+def load_robot_yaml(content_path: ContentPath = ContentPath()) -> dict:
+    """Load robot representation from a yaml or xrdf file.
+
+    Args:
+        content_path: Path to the robot configuration files.
+
+    Returns:
+        dict: Robot representation as a dictionary.
+    """
+    if isinstance(content_path, str):
+        log_error("content_path should be of type ContentPath")
+
+    robot_data = load_yaml(content_path.get_robot_configuration_path())
+
+    if "format" in robot_data and robot_data["format"] == "xrdf":
+        robot_data = convert_xrdf_to_curobo(
+            content_path=content_path,
+        )
+        robot_data["robot_cfg"]["kinematics"][
+            "asset_root_path"
+        ] = content_path.robot_asset_absolute_path
+
+    if "robot_cfg" not in robot_data:
+        robot_data["robot_cfg"] = robot_data
+    if "kinematics" not in robot_data["robot_cfg"]:
+        robot_data["robot_cfg"]["kinematics"] = robot_data
+    if content_path.robot_urdf_absolute_path is not None:
+        robot_data["robot_cfg"]["kinematics"]["urdf_path"] = content_path.robot_urdf_absolute_path
+    if content_path.robot_usd_absolute_path is not None:
+        robot_data["robot_cfg"]["kinematics"]["usd_path"] = content_path.robot_usd_absolute_path
+    if content_path.robot_asset_absolute_path is not None:
+        robot_data["robot_cfg"]["kinematics"][
+            "asset_root_path"
+        ] = content_path.robot_asset_absolute_path
+    if isinstance(robot_data["robot_cfg"]["kinematics"]["collision_spheres"], str):
+        robot_data["robot_cfg"]["kinematics"]["collision_spheres"] = join_path(
+            content_path.robot_config_root_path,
+            robot_data["robot_cfg"]["kinematics"]["collision_spheres"],
+        )
+    return robot_data