removed: origins from kinematics, inheritance and ROS from Robot

Author: mhubii

cli/rr_cam_swarm.py

@@ -6,7 +6,7 @@
 import torch
 
 from roboreg import differentiable as rrd
-from roboreg.io import find_files, parse_camera_info, parse_mono_data
+from roboreg.io import URDFParser, find_files, parse_camera_info, parse_mono_data
 from roboreg.losses import soft_dice_loss
 from roboreg.optim import LinearParticleSwarm, ParticleSwarmOptimizer
 from roboreg.util import (
@@ -301,9 +301,9 @@ def main() -> None:
         device=device,
     )
 
-    urdf_parser = rrd.URDFParser()
+    urdf_parser = URDFParser()
     urdf_parser.from_ros_xacro(ros_package=args.ros_package, xacro_path=args.xacro_path)
-    robot = rrd.Robot(
+    robot = rrd.Robot.from_urdf_parser(
         urdf_parser=urdf_parser,
         root_link_name=args.root_link_name,
         end_link_name=args.end_link_name,

cli/rr_hydra.py

@@ -183,7 +183,7 @@ def main():
 
     # instantiate robot
     batch_size = len(joint_states)
-    robot = Robot(
+    robot = Robot.from_urdf_parser(
         urdf_parser=urdf_parser,
         root_link_name=root_link_name,
         end_link_name=end_link_name,

roboreg/differentiable/kinematics.py

@@ -3,43 +3,31 @@
 import pytorch_kinematics as pk
 import torch
 
-from roboreg.io import URDFParser
-
 
 class TorchKinematics:
     __slots__ = [
         "_root_link_name",
         "_end_link_name",
         "_chain",
-        "_mesh_origins_lookup",
         "_device",
     ]
 
     def __init__(
         self,
-        urdf_parser: URDFParser,
+        urdf: str,
         root_link_name: str,
         end_link_name: str,
-        device: torch.device = "cuda",
+        device: torch.device = torch.device("cuda"),
     ) -> None:
         self._root_link_name = root_link_name
         self._end_link_name = end_link_name
         self._chain = self._build_serial_chain_from_urdf(
-            urdf_parser.urdf,
+            urdf=urdf,
             root_link_name=self._root_link_name,
             end_link_name=self._end_link_name,
         )
-
-        self._mesh_origins_lookup = urdf_parser.mesh_origins(
-            root_link_name=root_link_name, end_link_name=end_link_name
-        )
-        self._mesh_origins_lookup = {
-            key: torch.from_numpy(value).to(device=device, dtype=torch.float32)
-            for key, value in self._mesh_origins_lookup.items()
-        }
-
-        # default move to device
-        self.to(device=device)
+        self._device = device
+        self.to(device=self._device)
 
     def _build_serial_chain_from_urdf(
         self, urdf: str, root_link_name: str, end_link_name: str
@@ -50,18 +38,11 @@ def _build_serial_chain_from_urdf(
 
     def to(self, device: torch.device) -> None:
         self._chain.to(device=device)
-        for link_name in self._mesh_origins_lookup:
-            self._mesh_origins_lookup[link_name] = self._mesh_origins_lookup[
-                link_name
-            ].to(device=device)
         self._device = device
 
-    def mesh_forward_kinematics(self, q: torch.Tensor) -> Dict[str, torch.Tensor]:
-        r"""Computes forward kinematics and returns corresponding homogeneous transformations.
-        Corrects for mesh offsets. Meshes that are tranformed by the returned transformation appear physically correct.
-        """
+    def forward_kinematics(self, q: torch.Tensor) -> Dict[str, torch.Tensor]:
         ht_lookup = {
-            key: value.get_matrix() @ self._mesh_origins_lookup[key]
+            key: value.get_matrix()
             for key, value in self._chain.forward_kinematics(q, end_only=False).items()
         }
         return ht_lookup

roboreg/differentiable/rendering.py

@@ -13,7 +13,7 @@ class NVDiffRastRenderer:
 
     __slots__ = ["_device", "_ctx"]
 
-    def __init__(self, device: torch.device = "cuda") -> None:
+    def __init__(self, device: torch.device = torch.device("cuda")) -> None:
         super().__init__()
         if not torch.cuda.is_available():
             raise ValueError("CUDA is not available.")

roboreg/differentiable/robot.py

@@ -6,36 +6,69 @@
 from .structs import TorchMeshContainer
 
 
-class Robot(TorchMeshContainer):
-    __slots__ = ["_kinematics", "_configured_vertices"]
+class Robot:
+    __slots__ = ["_mesh_container", "_kinematics", "_configured_vertices", "_device"]
 
     def __init__(
         self,
+        mesh_container: TorchMeshContainer,
+        kinematics: TorchKinematics,
+        device: torch.device = torch.device("cuda"),
+    ) -> None:
+        self._mesh_container = mesh_container
+        self._kinematics = kinematics
+        self._configured_vertices = self.mesh_container.vertices.clone()
+        self._device = device
+        self.to(device=self._device)
+
+    @classmethod
+    def from_urdf_parser(
+        cls,
         urdf_parser: URDFParser,
         root_link_name: str,
         end_link_name: str,
         collision: bool = False,
         batch_size: int = 1,
-        device: torch.device = "cuda",
+        device: torch.device = torch.device("cuda"),
         target_reduction: float = 0.0,
-    ) -> None:
-        super().__init__(
-            mesh_paths=urdf_parser.ros_package_mesh_paths(
-                root_link_name=root_link_name,
-                end_link_name=end_link_name,
-                collision=collision,
-            ),
+    ) -> "Robot":
+        from roboreg.io import apply_mesh_origins, load_meshes, simplify_meshes
+
+        # parse data from URDF
+        mesh_paths = urdf_parser.ros_package_mesh_paths(
+            root_link_name=root_link_name,
+            end_link_name=end_link_name,
+            collision=collision,
+        )
+        mesh_origins = urdf_parser.mesh_origins(
+            root_link_name=root_link_name,
+            end_link_name=end_link_name,
+            collision=collision,
+        )
+
+        # load and preprocess meshes
+        meshes = load_meshes(paths=mesh_paths)
+        meshes = simplify_meshes(
+            meshes=meshes,
+            target_reduction=target_reduction,
+        )
+        meshes = apply_mesh_origins(meshes=meshes, origins=mesh_origins)
+
+        # configure this robot
+        mesh_container = TorchMeshContainer(
+            meshes=meshes,
             batch_size=batch_size,
             device=device,
-            target_reduction=target_reduction,
         )
-        self._kinematics = TorchKinematics(
-            urdf_parser=urdf_parser,
+
+        kinematics = TorchKinematics(
+            urdf=urdf_parser.urdf,
             root_link_name=root_link_name,
             end_link_name=end_link_name,
             device=device,
         )
-        self._configured_vertices = self.vertices.clone()
+
+        return cls(mesh_container=mesh_container, kinematics=kinematics, device=device)
 
     def configure(
         self, q: torch.FloatTensor, ht_root: torch.FloatTensor = None
@@ -44,37 +77,51 @@ def configure(
             raise ValueError(
                 f"Expected joint states of shape {self._kinematics.chain.n_joints}, got {q.shape[-1]}."
             )
-        if q.shape[0] != self._batch_size:
+        if q.shape[0] != self._mesh_container.batch_size:
             raise ValueError(
-                f"Batch size mismatch. Meshes: {self._batch_size}, joint states: {q.shape[0]}."
+                f"Batch size mismatch. Meshes: {self._mesh_container.batch_size}, joint states: {q.shape[0]}."
             )
         if ht_root is None:
             ht_root = torch.eye(4, device=self._device).unsqueeze(0)
-        ht_target_lookup = self._kinematics.mesh_forward_kinematics(q)
-        self._configured_vertices = self.vertices.clone()
+        ht_target_lookup = self._kinematics.forward_kinematics(q)
+        self._configured_vertices = self.mesh_container.vertices.clone()
         for link_name, ht in ht_target_lookup.items():
             self._configured_vertices[
                 :,
-                self.lower_vertex_index_lookup[
+                self.mesh_container.lower_vertex_index_lookup[
                     link_name
-                ] : self.upper_vertex_index_lookup[link_name],
+                ] : self.mesh_container.upper_vertex_index_lookup[link_name],
             ] = torch.matmul(
                 torch.matmul(
                     self._configured_vertices[
                         :,
-                        self.lower_vertex_index_lookup[
+                        self.mesh_container.lower_vertex_index_lookup[
                             link_name
-                        ] : self.upper_vertex_index_lookup[link_name],
+                        ] : self.mesh_container.upper_vertex_index_lookup[link_name],
                     ],
                     ht.transpose(-1, -2),
                 ),
                 ht_root.transpose(-1, -2),
             )
 
+    def to(self, device: torch.device) -> None:
+        self._mesh_container.to(device=device)
+        self._kinematics.to(device=device)
+        self._configured_vertices = self._configured_vertices.to(device=device)
+        self._device = device
+
     @property
     def kinematics(self) -> TorchKinematics:
         return self._kinematics
 
+    @property
+    def mesh_container(self) -> TorchMeshContainer:
+        return self._mesh_container
+
     @property
     def configured_vertices(self) -> torch.FloatTensor:
         return self._configured_vertices
+
+    @property
+    def device(self) -> torch.device:
+        return self._device

roboreg/differentiable/scene.py

@@ -66,7 +66,7 @@ def observe_from(
         )
         return self._renderer.constant_color(
             observed_vertices,
-            self._robot.faces,
+            self._robot.mesh_container.faces,
             self._cameras[camera_name].resolution,
         )
 

roboreg/differentiable/structs.py

@@ -34,7 +34,7 @@ def __init__(
         self,
         meshes: Dict[str, Mesh],
         batch_size: int = 1,
-        device: torch.device = "cuda",
+        device: torch.device = torch.device("cuda"),
     ) -> None:
         self._names = []
         self._vertices = []
@@ -62,7 +62,7 @@ def __init__(
     def _populate_container(
         self,
         meshes: Dict[str, Mesh],
-        device: torch.device = "cuda",
+        device: torch.device("cuda"),
     ) -> None:
         offset = 0
         for name, mesh in meshes.items():
@@ -184,7 +184,7 @@ def __init__(
         resolution: Tuple[int, int],
         intrinsics: Optional[Union[torch.FloatTensor, np.ndarray]] = None,
         extrinsics: Optional[Union[torch.FloatTensor, np.ndarray]] = None,
-        device: torch.device = "cuda",
+        device: torch.device = torch.device("cuda"),
         name: str = "camera",
     ) -> None:
         if intrinsics is None:
@@ -312,7 +312,7 @@ def __init__(
         extrinsics: Optional[Union[torch.FloatTensor, np.ndarray]] = None,
         zmin: float = 0.1,
         zmax: float = 100.0,
-        device: torch.device = "cuda",
+        device: torch.device = torch.device("cuda"),
     ) -> None:
         super().__init__(resolution, intrinsics, extrinsics, device)
 

roboreg/io/meshes.py

@@ -64,6 +64,8 @@ def simplify_mesh(mesh: Mesh, target_reduction: float = 0.0) -> Mesh:
     Returns:
         Mesh: The simplified mesh.
     """
+    if target_reduction == 0.0:
+        return mesh
     if 0.0 > target_reduction > 1.0:
         raise ValueError(
             f"Expected target reduction in [0, 1], got {target_reduction}."
@@ -88,7 +90,41 @@ def simplify_meshes(
     Returns:
         Dict[str, Mesh]: The simplified meshes. 
     """
+    if target_reduction == 0.0:
+        return meshes
     return {
         name: simplify_mesh(mesh, target_reduction=target_reduction)
         for name, mesh in meshes.items()
     }
+
+
+def apply_mesh_origin(vertices: np.ndarray, origin: np.ndarray) -> np.ndarray:
+    r"""Apply a homogeneous transformation to mesh vertices.
+
+    Args:
+        vertices (np.ndarray): The mesh vertices of shape Nx3.
+        origin (np.ndarray): The mesh origin as a homogeneous transform of shape 4x4.
+    Returns:
+        np.ndarray: The transformed mesh vertices of shape Nx3.
+    """
+    return vertices @ origin[:3, :3].T + origin[:3, 3].T
+
+
+def apply_mesh_origins(
+    meshes: Dict[str, Mesh], origins: Dict[str, np.ndarray]
+) -> Dict[str, Mesh]:
+    r"""Apply mesh origins to multiple meshes.
+
+    Args:
+        meshes (Dict[str, Mesh]): The meshes to apply origins to.
+        origins (Dict[str, np.ndarray]): The mesh origins.
+
+    Returns:
+        Dict[str, Mesh]: The meshes with applied origins.
+    """
+    for name in meshes.keys():
+        if name in origins:
+            meshes[name].vertices = apply_mesh_origin(
+                meshes[name].vertices, origins[name]
+            )
+    return meshes

roboreg/util/factories.py

@@ -56,7 +56,7 @@ def create_robot_scene(
         )
 
     # instantiate robot
-    robot = Robot(
+    robot = Robot.from_urdf_parser(
         urdf_parser=urdf_parser,
         root_link_name=root_link_name,
         end_link_name=end_link_name,