Modifies behavior of debug visualization for better UI experience (#208)

# Description

Earlier the UI was disabling buttons for which the config object set
`debug_vis` flag as False. This wasn't an elegant behavior as users may
want to enable the visualization from the GUI at runtime. This MR
modifies the behavior of all classes that provide debug visualization to
support this behavior.

## Type of change

- Bug fix (non-breaking change which fixes an issue)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file

source/extensions/omni.isaac.orbit/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.9.23"
+version = "0.9.24"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

@@ -1,6 +1,18 @@
 Changelog
 ---------
 
+0.9.24 (2023-10-27)
+~~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+* Changed the behavior of setting up debug visualization for assets, sensors and command generators.
+  Earlier it was raising an error if debug visualization was not enabled in the configuration object.
+  Now it checks whether debug visualization is implemented and only sets up the callback if it is
+  implemented.
+
+
 0.9.23 (2023-10-27)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/asset_base.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import inspect
 import re
 import weakref
 from abc import ABC, abstractmethod
@@ -68,14 +69,10 @@ def __init__(self, cfg: AssetBaseCfg):
             lambda event, obj=weakref.proxy(self): obj._invalidate_initialize_callback(event),
             order=10,
         )
-        # add callback for debug visualization
-        if self.cfg.debug_vis:
-            app_interface = omni.kit.app.get_app_interface()
-            self._debug_visualization_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop(
-                lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event),
-            )
-        else:
-            self._debug_visualization_handle = None
+        # add handle for debug visualization (this is set to a valid handle inside set_debug_vis)
+        self._debug_vis_handle = None
+        # set initial state of debug visualization
+        self.set_debug_vis(self.cfg.debug_vis)
 
     def __del__(self):
         """Unsubscribe from the callbacks."""
@@ -87,9 +84,9 @@ def __del__(self):
             self._invalidate_initialize_handle.unsubscribe()
             self._invalidate_initialize_handle = None
         # clear debug visualization
-        if self._debug_visualization_handle:
-            self._debug_visualization_handle.unsubscribe()
-            self._debug_visualization_handle = None
+        if self._debug_vis_handle:
+            self._debug_vis_handle.unsubscribe()
+            self._debug_vis_handle = None
 
     """
     Properties
@@ -107,21 +104,47 @@ def data(self) -> Any:
         """Data related to the asset."""
         return NotImplementedError
 
+    @property
+    def has_debug_vis_implementation(self) -> bool:
+        """Whether the asset has a debug visualization implemented."""
+        # check if function raises NotImplementedError
+        source_code = inspect.getsource(self._debug_vis_callback)
+        return "NotImplementedError" not in source_code
+
     """
     Operations.
     """
 
-    def set_debug_vis(self, debug_vis: bool):
+    def set_debug_vis(self, debug_vis: bool) -> bool:
         """Sets whether to visualize the asset data.
 
         Args:
             debug_vis: Whether to visualize the asset data.
 
-        Raises:
-            RuntimeError: If the asset debug visualization is not enabled.
+        Returns:
+            Whether the debug visualization was successfully set. False if the asset
+            does not support debug visualization.
         """
-        if not self.cfg.debug_vis:
-            raise RuntimeError("Debug visualization is not enabled for this sensor.")
+        # check if debug visualization is supported
+        if not self.has_debug_vis_implementation:
+            return False
+        # toggle debug visualization objects
+        self._set_debug_vis_impl(debug_vis)
+        # toggle debug visualization handles
+        if debug_vis:
+            # create a subscriber for the post update event if it doesn't exist
+            if self._debug_vis_handle is None:
+                app_interface = omni.kit.app.get_app_interface()
+                self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop(
+                    lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event)
+                )
+        else:
+            # remove the subscriber if it exists
+            if self._debug_vis_handle is not None:
+                self._debug_vis_handle.unsubscribe()
+                self._debug_vis_handle = None
+        # return success
+        return True
 
     @abstractmethod
     def reset(self, env_ids: Sequence[int] | None = None):
@@ -158,12 +181,24 @@ def _initialize_impl(self):
         """Initializes the PhysX handles and internal buffers."""
         raise NotImplementedError
 
-    def _debug_vis_impl(self):
-        """Perform debug visualization of the asset."""
-        pass
+    def _set_debug_vis_impl(self, debug_vis: bool):
+        """Set debug visualization into visualization objects.
+
+        This function is responsible for creating the visualization objects if they don't exist
+        and input ``debug_vis`` is True. If the visualization objects exist, the function should
+        set their visibility into the stage.
+        """
+        raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.")
+
+    def _debug_vis_callback(self, event):
+        """Callback for debug visualization.
+
+        This function calls the visualization objects and sets the data to visualize into them.
+        """
+        raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.")
 
     """
-    Simulation callbacks.
+    Internal simulation callbacks.
     """
 
     def _initialize_callback(self, event):
@@ -180,7 +215,3 @@ def _initialize_callback(self, event):
     def _invalidate_initialize_callback(self, event):
         """Invalidates the scene elements."""
         self._is_initialized = False
-
-    def _debug_vis_callback(self, event):
-        """Visualizes the asset data."""
-        self._debug_vis_impl()

source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/command_generator_base.py

@@ -12,6 +12,7 @@
 
 from __future__ import annotations
 
+import inspect
 import torch
 import weakref
 from abc import ABC, abstractmethod
@@ -48,6 +49,7 @@ def __init__(self, cfg: CommandGeneratorBaseCfg, env: BaseEnv):
         # store the inputs
         self.cfg = cfg
         self._env = env
+
         # create buffers to store the command
         # -- metrics that can be used for logging
         self.metrics = dict()
@@ -56,21 +58,16 @@ def __init__(self, cfg: CommandGeneratorBaseCfg, env: BaseEnv):
         # -- counter for the number of times the command has been resampled within the current episode
         self.command_counter = torch.zeros(self.num_envs, device=self.device, dtype=torch.long)
 
-        # add callback for debug visualization
-        if self.cfg.debug_vis:
-            app_interface = omni.kit.app.get_app_interface()
-            # NOTE: Use weakref on callback to ensure that this object can be deleted when its destructor is called.
-            self._debug_visualization_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop(
-                lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event),
-            )
-        else:
-            self._debug_visualization_handle = None
+        # add handle for debug visualization (this is set to a valid handle inside set_debug_vis)
+        self._debug_vis_handle = None
+        # set initial state of debug visualization
+        self.set_debug_vis(self.cfg.debug_vis)
 
     def __del__(self):
         """Unsubscribe from the callbacks."""
-        if self._debug_visualization_handle is not None:
-            self._debug_visualization_handle.unsubscribe()
-            self._debug_visualization_handle = None
+        if self._debug_vis_handle:
+            self._debug_vis_handle.unsubscribe()
+            self._debug_vis_handle = None
 
     """
     Properties
@@ -92,21 +89,47 @@ def command(self) -> torch.Tensor:
         """The command tensor. Shape is (num_envs, command_dim)."""
         raise NotImplementedError
 
+    @property
+    def has_debug_vis_implementation(self) -> bool:
+        """Whether the command generator has a debug visualization implemented."""
+        # check if function raises NotImplementedError
+        source_code = inspect.getsource(self._debug_vis_callback)
+        return "NotImplementedError" not in source_code
+
     """
     Operations.
     """
 
-    def set_debug_vis(self, debug_vis: bool):
+    def set_debug_vis(self, debug_vis: bool) -> bool:
         """Sets whether to visualize the command data.
 
         Args:
             debug_vis: Whether to visualize the command data.
 
-        Raises:
-            RuntimeError: If the command debug visualization is not enabled.
+        Returns:
+            Whether the debug visualization was successfully set. False if the command
+            generator does not support debug visualization.
         """
-        if not self.cfg.debug_vis:
-            raise RuntimeError("Debug visualization is not enabled for this sensor.")
+        # check if debug visualization is supported
+        if not self.has_debug_vis_implementation:
+            return False
+        # toggle debug visualization objects
+        self._set_debug_vis_impl(debug_vis)
+        # toggle debug visualization handles
+        if debug_vis:
+            # create a subscriber for the post update event if it doesn't exist
+            if self._debug_vis_handle is None:
+                app_interface = omni.kit.app.get_app_interface()
+                self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop(
+                    lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event)
+                )
+        else:
+            # remove the subscriber if it exists
+            if self._debug_vis_handle is not None:
+                self._debug_vis_handle.unsubscribe()
+                self._debug_vis_handle = None
+        # return success
+        return True
 
     def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]:
         """Reset the command generator and log metrics.
@@ -173,14 +196,6 @@ def _resample(self, env_ids: Sequence[int]):
         # resample the command
         self._resample_command(env_ids)
 
-    """
-    Simulation callbacks.
-    """
-
-    def _debug_vis_callback(self, event):
-        """Visualizes the sensor data."""
-        self._debug_vis_impl()
-
     """
     Implementation specific functions.
     """
@@ -200,9 +215,18 @@ def _update_metrics(self):
         """Update the metrics based on the current state."""
         raise NotImplementedError
 
-    def _debug_vis_impl(self):
-        """Visualize the command in the simulator.
+    def _set_debug_vis_impl(self, debug_vis: bool):
+        """Set debug visualization into visualization objects.
+
+        This function is responsible for creating the visualization objects if they don't exist
+        and input ``debug_vis`` is True. If the visualization objects exist, the function should
+        set their visibility into the stage.
+        """
+        raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.")
+
+    def _debug_vis_callback(self, event):
+        """Callback for debug visualization.
 
-        This is an optional function that can be used to visualize the command in the simulator.
+        This function calls the visualization objects and sets the data to visualize into them.
         """
-        pass
+        raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.")

source/extensions/omni.isaac.orbit/omni/isaac/orbit/command_generators/position_command_generator.py

@@ -41,11 +41,16 @@ def __init__(self, cfg: TerrainBasedPositionCommandGeneratorCfg, env: BaseEnv):
             cfg: The configuration parameters for the command generator.
             env: The environment object.
         """
+        # initialize the base class
         super().__init__(cfg, env)
+
+        # obtain the robot and terrain assets
         # -- robot
         self.robot: Articulation = env.scene[cfg.asset_name]
         # -- terrain
         self.terrain: TerrainImporter = env.scene.terrain
+
+        # crete buffers to store the command
         # -- commands: (x, y, z, heading)
         self.pos_command_w = torch.zeros(self.num_envs, 3, device=self.device)
         self.heading_command_w = torch.zeros(self.num_envs, device=self.device)
@@ -54,8 +59,6 @@ def __init__(self, cfg: TerrainBasedPositionCommandGeneratorCfg, env: BaseEnv):
         # -- metrics
         self.metrics["error_pos"] = torch.zeros(self.num_envs, device=self.device)
         self.metrics["error_heading"] = torch.zeros(self.num_envs, device=self.device)
-        # -- debug vis
-        self.box_goal_visualizer = None
 
     def __str__(self) -> str:
         msg = "TerrainBasedPositionCommandGenerator:\n"
@@ -73,15 +76,6 @@ def command(self) -> torch.Tensor:
         """The desired base position in base frame. Shape is (num_envs, 3)."""
         return self.pos_command_b
 
-    """
-    Operations.
-    """
-
-    def set_debug_vis(self, debug_vis: bool):
-        super().set_debug_vis(debug_vis)
-        if self.box_goal_visualizer is not None:
-            self.box_goal_visualizer.set_visibility(debug_vis)
-
     """
     Implementation specific functions.
     """
@@ -120,12 +114,20 @@ def _update_metrics(self):
         self.metrics["error_pos"] = torch.norm(self.pos_command_w - self.robot.data.root_pos_w[:, :3], dim=1)
         self.metrics["error_heading"] = torch.abs(wrap_to_pi(self.heading_command_w - self.robot.heading_w))
 
-    def _debug_vis_impl(self):
-        # create the box marker if necessary
-        if self.box_goal_visualizer is None:
-            marker_cfg = CUBOID_MARKER_CFG.copy()
-            marker_cfg.prim_path = "/Visuals/Command/position_goal"
-            marker_cfg.markers["cuboid"].scale = (0.1, 0.1, 0.1)
-            self.box_goal_visualizer = VisualizationMarkers(marker_cfg)
+    def _set_debug_vis_impl(self, debug_vis: bool):
+        # create markers if necessary for the first tome
+        if debug_vis:
+            if not hasattr(self, "box_goal_visualizer"):
+                marker_cfg = CUBOID_MARKER_CFG.copy()
+                marker_cfg.prim_path = "/Visuals/Command/position_goal"
+                marker_cfg.markers["cuboid"].scale = (0.1, 0.1, 0.1)
+                self.box_goal_visualizer = VisualizationMarkers(marker_cfg)
+            # set their visibility to true
+            self.box_goal_visualizer.set_visibility(True)
+        else:
+            if hasattr(self, "box_goal_visualizer"):
+                self.box_goal_visualizer.set_visibility(False)
+
+    def _debug_vis_callback(self, event):
         # update the box marker
         self.box_goal_visualizer.visualize(self.pos_command_w)