Use masked_fill_ for boolean-mask scalar assignments

Replace boolean-mask index assignments of the form ``tensor[mask] = scalar``
with ``Tensor.masked_fill_`` across camera/depth observation processing, the
linear interpolation utility, the Shadow Hand feature extractor, the drone
navigation observations, and the TacSL visuotactile renderer.

``masked_fill_`` performs the fill in a single fused kernel and avoids the
``nonzero``-based advanced-indexing path (and its device sync), which makes
these GPU tensor operations more performant while preserving behavior.

Author: ooctipus

source/isaaclab/changelog.d/masked-fill-perf.rst

@@ -0,0 +1,7 @@
+Changed
+^^^^^^^
+
+* Changed boolean-mask scalar assignments to use :meth:`torch.Tensor.masked_fill_`
+  in :class:`~isaaclab.utils.interpolation.LinearInterpolation` and the camera
+  image observation in :mod:`isaaclab.envs.mdp.observations` to avoid the
+  ``nonzero``-based advanced-indexing path and improve performance.

source/isaaclab/isaaclab/envs/mdp/observations.py

@@ -425,7 +425,7 @@ def image(
             mean_tensor = torch.mean(images, dim=(1, 2), keepdim=True)
             images -= mean_tensor
         elif "distance_to" in data_type or "depth" in data_type:
-            images[images == float("inf")] = 0
+            images.masked_fill_(images == float("inf"), 0)
         elif "normals" in data_type:
             images = (images + 1.0) * 0.5
 

source/isaaclab/isaaclab/utils/interpolation/linear_interpolation.py

@@ -75,7 +75,7 @@ def compute(self, q: torch.Tensor) -> torch.Tensor:
         # compute the weight as: (q_i - x_lb) / (x_ub - x_lb)
         weight = (q_1d - self._x[lower_bound]) / (self._x[upper_bound] - self._x[lower_bound])
         # If a point is out of bounds assign weight 0.0
-        weight[upper_bound == lower_bound] = 0.0
+        weight.masked_fill_(upper_bound == lower_bound, 0.0)
 
         # Perform linear interpolation
         fq = self._y[lower_bound] + weight * (self._y[upper_bound] - self._y[lower_bound])

source/isaaclab_contrib/changelog.d/masked-fill-perf.rst

@@ -0,0 +1,6 @@
+Changed
+^^^^^^^
+
+* Changed boolean-mask scalar assignments to use :meth:`torch.Tensor.masked_fill_`
+  in the TacSL visuotactile renderer to avoid the ``nonzero``-based
+  advanced-indexing path and improve performance.

source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_render.py

@@ -175,7 +175,7 @@ def render(self, height_map: torch.Tensor) -> torch.Tensor:
             Rendered image tensor. Shape is (N, H, W, 3).
         """
         height_map = height_map.clone()
-        height_map[torch.abs(height_map) < 1e-6] = 0  # remove minor artifact
+        height_map.masked_fill_(torch.abs(height_map) < 1e-6, 0)  # remove minor artifact
         height_map = height_map * -1000.0
         height_map /= self.cfg.mm_per_pixel
 
@@ -251,7 +251,7 @@ def _generate_normals(self, img: torch.Tensor) -> tuple[torch.Tensor, torch.Tens
         grad_mag_orig = torch.sqrt(dzdx**2 + dzdy**2)
         grad_mag = torch.arctan(grad_mag_orig)  # seems that arctan is used as a squashing function
         grad_dir = torch.arctan2(dzdx, dzdy)
-        grad_dir[grad_mag_orig == 0] = 0
+        grad_dir.masked_fill_(grad_mag_orig == 0, 0)
 
         # handle edges
         grad_mag = torch.nn.functional.pad(grad_mag[:, 1:-1, 1:-1], pad=(1, 1, 1, 1))

source/isaaclab_tasks/changelog.d/masked-fill-perf.rst

@@ -0,0 +1,7 @@
+Changed
+^^^^^^^
+
+* Changed boolean-mask scalar assignments to use :meth:`torch.Tensor.masked_fill_`
+  in the Shadow Hand feature extractor, the Cartpole camera environments, and the
+  drone navigation observations to avoid the ``nonzero``-based advanced-indexing
+  path and improve performance.

source/isaaclab_tasks/isaaclab_tasks/contrib/cartpole_showcase/cartpole_camera/cartpole_camera_env.py

@@ -52,7 +52,7 @@ def _get_observations(self) -> dict:
             camera_data -= mean_tensor
         elif "depth" in self.cfg.tiled_camera.data_types:
             camera_data = self._tiled_camera.data.output[data_type]
-            camera_data[camera_data == float("inf")] = 0
+            camera_data.masked_fill_(camera_data == float("inf"), 0)
 
         # fundamental spaces
         # - Box

source/isaaclab_tasks/isaaclab_tasks/contrib/drone_arl/mdp/observations.py

@@ -196,11 +196,11 @@ def __call__(self, env: ManagerBasedEnv, sensor_cfg: SceneEntityCfg, data_type:
 
         if self.normalize:
             if self.data_type == "distance_to_image_plane":
-                images[images == float("inf")] = 10.0
-                images[images == -float("inf")] = 10.0
-                images[images > 10.0] = 10.0
+                images.masked_fill_(images == float("inf"), 10.0)
+                images.masked_fill_(images == -float("inf"), 10.0)
+                images.masked_fill_(images > 10.0, 10.0)
                 images = images / 10.0
-                images[images < 0.02] = -1.0
+                images.masked_fill_(images < 0.02, -1.0)
             else:
                 raise ValueError(f"Image data type: {self.data_type} not supported")
 

source/isaaclab_tasks/isaaclab_tasks/core/cartpole/cartpole_direct_camera_env.py

@@ -113,7 +113,7 @@ def _get_observations(self) -> dict:
             camera_data -= mean_tensor
         elif "depth" in self.cfg.tiled_camera.data_types:
             camera_data = self._tiled_camera.data.output[data_type]
-            camera_data[camera_data == float("inf")] = 0
+            camera_data.masked_fill_(camera_data == float("inf"), 0)
         elif "semantic_segmentation" in self.cfg.tiled_camera.data_types:
             camera_data = self._tiled_camera.data.output[data_type]
 

source/isaaclab_tasks/isaaclab_tasks/core/shadow_hand/feature_extractor.py

@@ -231,7 +231,7 @@ def _preprocess_images(self, camera_output: dict[str, torch.Tensor]) -> torch.Te
             if dt == "rgb":
                 img = img / 255.0
             elif dt == "depth":
-                img[img == float("inf")] = 0
+                img.masked_fill_(img == float("inf"), 0)
                 img /= 5.0
                 max_val = img.max()
                 if max_val > 0:
@@ -262,7 +262,7 @@ def _save_images(self, camera_output: dict[str, torch.Tensor]):
             img = camera_output[dt].float()
             if dt == "depth":
                 img = img.clone()
-                img[img == float("inf")] = 0
+                img.masked_fill_(img == float("inf"), 0)
                 max_val = img.max()
                 if max_val > 0:
                     img = img / max_val