Fix missed rays in base height reward

source/isaaclab/changelog.d/antoiner-fix-base-height-ray-misses.rst

@@ -0,0 +1,5 @@
+Fixed
+^^^^^
+
+* Fixed :func:`~isaaclab.envs.mdp.base_height_l2` returning non-finite penalties when ray-caster rays missed the
+  terrain.

source/isaaclab/isaaclab/envs/mdp/rewards.py

@@ -105,18 +105,29 @@ def base_height_l2(
     asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
     sensor_cfg: SceneEntityCfg | None = None,
 ) -> torch.Tensor:
-    """Penalize asset height from its target using L2 squared kernel.
+    """Penalize the asset height error using a squared L2 kernel.
 
-    Note:
-        For flat terrain, target height is in the world frame. For rough terrain,
-        sensor readings can adjust the target height to account for the terrain.
+    When :paramref:`sensor_cfg` is provided, the target height is relative to the mean height of the valid ray hits.
+    Non-finite ray hits are ignored. If all rays miss the terrain, the penalty is zero for that environment.
+
+    Args:
+        env: The environment.
+        target_height: Target asset height [m] relative to the terrain, or to the world origin when
+            :paramref:`sensor_cfg` is not provided.
+        asset_cfg: The asset whose height is penalized.
+        sensor_cfg: The ray-caster sensor used to estimate the terrain height.
+
+    Returns:
+        The squared height error [m²], shape (num_envs,).
     """
     # extract the used quantities (to enable type-hinting)
     asset: RigidObject = env.scene[asset_cfg.name]
     if sensor_cfg is not None:
         sensor: RayCaster = env.scene[sensor_cfg.name]
-        # Adjust the target height using the sensor data
-        adjusted_target_height = target_height + torch.mean(sensor.data.ray_hits_w.torch[..., 2], dim=1)
+        ray_hits_z = sensor.data.ray_hits_w.torch[..., 2]
+        terrain_height = torch.nanmean(ray_hits_z.masked_fill(torch.isinf(ray_hits_z), torch.nan), dim=1)
+        height_error = asset.data.root_pos_w.torch[:, 2] - target_height - terrain_height
+        return torch.square(height_error).masked_fill_(torch.isnan(terrain_height), 0.0)
     else:
         # Use the provided target height directly for flat terrain
         adjusted_target_height = target_height

source/isaaclab/test/envs/test_mdp_rewards.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from types import SimpleNamespace
+
+import torch
+
+from isaaclab.envs.mdp.rewards import base_height_l2
+from isaaclab.managers import SceneEntityCfg
+
+
+def _make_env(root_heights: list[float], ray_hit_heights: list[list[float]]) -> SimpleNamespace:
+    root_pos_w = torch.zeros((len(root_heights), 3))
+    root_pos_w[:, 2] = torch.tensor(root_heights)
+
+    ray_hits_w = torch.zeros((len(ray_hit_heights), len(ray_hit_heights[0]), 3))
+    ray_hits_w[..., 2] = torch.tensor(ray_hit_heights)
+
+    scene = {
+        "robot": SimpleNamespace(data=SimpleNamespace(root_pos_w=SimpleNamespace(torch=root_pos_w))),
+        "height_scanner": SimpleNamespace(data=SimpleNamespace(ray_hits_w=SimpleNamespace(torch=ray_hits_w))),
+    }
+    return SimpleNamespace(scene=scene)
+
+
+def test_base_height_l2_handles_missed_ray_hits():
+    """The reward ignores partial misses and is zero when all rays miss."""
+    env = _make_env(
+        root_heights=[1.0, 1.0],
+        ray_hit_heights=[[0.0, 0.2, float("inf")], [float("inf"), float("inf"), float("inf")]],
+    )
+
+    penalty = base_height_l2(
+        env,
+        target_height=0.5,
+        sensor_cfg=SceneEntityCfg("height_scanner"),
+    )
+
+    torch.testing.assert_close(penalty, torch.tensor([0.16, 0.0]))