Draws connection lines for FrameTransformer visualization

source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/frame_transformer/frame_transformer.py

@@ -22,6 +22,7 @@
     is_identity_pose,
     subtract_frame_transforms,
 )
+import omni.isaac.lab.utils.string as string_utils
 
 from ..sensor_base import SensorBase
 from .frame_transformer_data import FrameTransformerData
@@ -88,6 +89,26 @@ def data(self) -> FrameTransformerData:
         # return the data
         return self._data
 
+    @property
+    def num_bodies(self) -> int:
+        """Returns the number of target bodies being tracked.
+
+        Note:
+            This is an alias used for consistency with other sensors. Otherwise, we recommend using
+            :attr:`len(data.target_frame_names)` to access the number of target frames.
+        """
+        return len(self._target_frame_body_names)
+
+    @property
+    def body_names(self) -> list[str]:
+        """Returns the names of the target bodies being tracked.
+
+        Note:
+            This is an alias used for consistency with other sensors. Otherwise, we recommend using
+            :attr:`data.target_frame_names` to access the target frame names.
+        """
+        return self._target_frame_body_names
+
     """
     Operations
     """
@@ -99,6 +120,18 @@ def reset(self, env_ids: Sequence[int] | None = None):
         if env_ids is None:
             env_ids = ...
 
+    def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]:
+        """Find bodies in the articulation based on the name keys.
+
+        Args:
+            name_keys: A regular expression or a list of regular expressions to match the body names.
+            preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False.
+
+        Returns:
+            A tuple of lists containing the body indices and names.
+        """
+        return string_utils.resolve_matching_names(name_keys, self._target_frame_names, preserve_order)
+
     """
     Implementation.
     """
@@ -395,16 +428,33 @@ def _set_debug_vis_impl(self, debug_vis: bool):
         if debug_vis:
             if not hasattr(self, "frame_visualizer"):
                 self.frame_visualizer = VisualizationMarkers(self.cfg.visualizer_cfg)
+
+                import omni.isaac.debug_draw._debug_draw as isaac_debug_draw
+
+                self.debug_draw = isaac_debug_draw.acquire_debug_draw_interface()
+
             # set their visibility to true
             self.frame_visualizer.set_visibility(True)
         else:
             if hasattr(self, "frame_visualizer"):
                 self.frame_visualizer.set_visibility(False)
+                # clear the lines
+                self.debug_draw.clear_lines()
 
     def _debug_vis_callback(self, event):
         # Update the visualized markers
-        if self.frame_visualizer is not None:
-            self.frame_visualizer.visualize(self._data.target_pos_w.view(-1, 3), self._data.target_quat_w.view(-1, 4))
+        all_pos = torch.cat([self._data.source_pos_w, self._data.target_pos_w.view(-1, 3)], dim=0)
+        all_quat = torch.cat([self._data.source_quat_w, self._data.target_quat_w.view(-1, 4)], dim=0)
+        self.frame_visualizer.visualize(all_pos, all_quat)
+
+        # Draw lines connecting the source frame to the target frames
+        self.debug_draw.clear_lines()
+        # make the lines color yellow
+        source_pos = self._data.source_pos_w.cpu().tolist()
+        colors = [[1, 1, 0, 1]] * self._num_envs
+        for frame_index in range(len(self._target_frame_names)):
+            target_pos = self._data.target_pos_w[:, frame_index].cpu().tolist()
+            self.debug_draw.draw_lines(source_pos, target_pos, colors, [1.5] * self._num_envs)
 
     """
     Internal simulation callbacks.