Updates to ray caster ray alignment and more customizable drift sampling (#2556)

This PR pushes some changes from our internal RAI fork of Isaac Lab.

The 2 main changes are:

#### Adds a new `ray_alignment` parameter to replace the previous
`yaw_only` arg.

Now the ray alignment can be aligned to `world`, `yaw`, or `base.

#### Improves drift height sampling

At the moment, the `RayCasterCfg` hosts a parameter called
`drift_range`. It allows to randomize the position of the sensor in
world frame, which comes with a couple of downsides:
* If the projection is done along the gravity vector, the z-component
drift term cannot be visualized (it is implicitly given in the pose, and
added to the height scan in the observation term).
* The perturbation is applied in world frame, i.e., if the ray cast is
yaw aligned, then the drift varies with the yaw angle
* The drift is applied to all (x,y,z) components equally

This PR adds a new parameter ray_cast_drift_range. It gives more freedom
in choosing the drift ranges as x, y and z components can be treated
separately. It also applies the drift after projecting and rotating the
ray cast points, i.e. the drift is now invariant under the yaw angle and
the drift in z is can be visualized

I've added Fabian Jenelten and Jeonghwan Kim to this PR as they worked
on these changes within RAI - I'm simply helping to push them up 😄

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] 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
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Signed-off-by: ooctipus <zhengyuz@nvidia.com>
Co-authored-by: ooctipus <zhengyuz@nvidia.com>

CONTRIBUTORS.md

@@ -57,6 +57,7 @@ Guidelines for modifications:
 * David Yang
 * Dhananjay Shendre
 * Dorsa Rohani
+* Fabian Jenelten
 * Felipe Mohr
 * Felix Yu
 * Gary Lvov
@@ -69,6 +70,7 @@ Guidelines for modifications:
 * Jack Zeng
 * Jan Kerner
 * Jean Tampon
+* Jeonghwan Kim
 * Jia Lin Yuan
 * Jiakai Zhang
 * Jinghuan Shang

scripts/tutorials/03_envs/create_quadruped_base_env.py

@@ -105,7 +105,7 @@ class MySceneCfg(InteractiveSceneCfg):
     height_scanner = RayCasterCfg(
         prim_path="{ENV_REGEX_NS}/Robot/base",
         offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
-        attach_yaw_only=True,
+        ray_alignment="yaw",
         pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
         debug_vis=True,
         mesh_prim_paths=["/World/ground"],

scripts/tutorials/04_sensors/add_sensors_on_robot.py

@@ -84,7 +84,7 @@ class SensorsSceneCfg(InteractiveSceneCfg):
         prim_path="{ENV_REGEX_NS}/Robot/base",
         update_period=0.02,
         offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
-        attach_yaw_only=True,
+        ray_alignment="yaw",
         pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
         debug_vis=True,
         mesh_prim_paths=["/World/defaultGroundPlane"],

scripts/tutorials/04_sensors/run_ray_caster.py

@@ -49,7 +49,7 @@ def define_sensor() -> RayCaster:
         prim_path="/World/Origin.*/ball",
         mesh_prim_paths=["/World/ground"],
         pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=(2.0, 2.0)),
-        attach_yaw_only=True,
+        ray_alignment="yaw",
         debug_vis=not args_cli.headless,
     )
     ray_caster = RayCaster(cfg=ray_caster_cfg)

source/isaaclab/isaaclab/scene/interactive_scene_cfg.py

@@ -52,7 +52,7 @@ class InteractiveSceneCfg:
             height_scanner = RayCasterCfg(
                 prim_path="{ENV_REGEX_NS}/Robot_1/base",
                 offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
-                attach_yaw_only=True,
+                ray_alignment="yaw",
                 pattern_cfg=GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
                 debug_vis=True,
                 mesh_prim_paths=["/World/ground"],

source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py

@@ -111,7 +111,13 @@ class RayCaster(SensorBase):
         if env_ids is None:
             env_ids = slice(None)
         # resample the drift
-        self.drift[env_ids] = self.drift[env_ids].uniform_(*self.cfg.drift_range)
+        r = torch.empty(len(env_ids), 3, device=self.device)
+        self.drift[env_ids] = r.uniform_(*self.cfg.drift_range)
+        # resample the height drift
+        r = torch.empty(len(env_ids), device=self.device)
+        self.ray_cast_drift[env_ids, 0] = r.uniform_(*self.cfg.ray_cast_drift_range["x"])
+        self.ray_cast_drift[env_ids, 1] = r.uniform_(*self.cfg.ray_cast_drift_range["y"])
+        self.ray_cast_drift[env_ids, 2] = r.uniform_(*self.cfg.ray_cast_drift_range["z"])
 
     """
     Implementation.
@@ -212,6 +218,7 @@ class RayCaster(SensorBase):
         self.ray_directions = self.ray_directions.repeat(self._view.count, 1, 1)
         # prepare drift
         self.drift = torch.zeros(self._view.count, 3, device=self.device)
+        self.ray_cast_drift = torch.zeros(self._view.count, 3, device=self.device)
         # fill the data buffer
         self._data.pos_w = torch.zeros(self._view.count, 3, device=self._device)
         self._data.quat_w = torch.zeros(self._view.count, 4, device=self._device)
@@ -233,23 +240,44 @@ class RayCaster(SensorBase):
         # note: we clone here because we are read-only operations
         pos_w = pos_w.clone()
         quat_w = quat_w.clone()
-        # apply drift
+        # apply drift to ray starting position in world frame
         pos_w += self.drift[env_ids]
         # store the poses
         self._data.pos_w[env_ids] = pos_w
         self._data.quat_w[env_ids] = quat_w
 
         # ray cast based on the sensor poses
+        if self.cfg.ray_alignment == "world":
+            # apply horizontal drift to ray starting position in ray caster frame
+            pos_w[:, 0:2] += self.ray_cast_drift[env_ids, 0:2]
+            # no rotation is considered and directions are not rotated
+            ray_starts_w = self.ray_starts[env_ids]
+            ray_starts_w += pos_w.unsqueeze(1)
+            ray_directions_w = self.ray_directions[env_ids]
+        elif self.cfg.ray_alignment == "yaw" or self.cfg.attach_yaw_only:
             if self.cfg.attach_yaw_only:
+                self.cfg.ray_alignment = "yaw"
+                omni.log.warn(
+                    "The `attach_yaw_only` property will be deprecated in a future release. Please use"
+                    " `ray_alignment='yaw'` instead."
+                )
+
+            # apply horizontal drift to ray starting position in ray caster frame
+            pos_w[:, 0:2] += quat_apply_yaw(quat_w, self.ray_cast_drift[env_ids])[:, 0:2]
             # only yaw orientation is considered and directions are not rotated
             ray_starts_w = quat_apply_yaw(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids])
             ray_starts_w += pos_w.unsqueeze(1)
             ray_directions_w = self.ray_directions[env_ids]
-        else:
+        elif self.cfg.ray_alignment == "base":
+            # apply horizontal drift to ray starting position in ray caster frame
+            pos_w[:, 0:2] += quat_apply(quat_w, self.ray_cast_drift[env_ids])[:, 0:2]
             # full orientation is considered
             ray_starts_w = quat_apply(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids])
             ray_starts_w += pos_w.unsqueeze(1)
             ray_directions_w = quat_apply(quat_w.repeat(1, self.num_rays), self.ray_directions[env_ids])
+        else:
+            raise RuntimeError(f"Unsupported ray_alignment type: {self.cfg.ray_alignment}.")
+
         # ray cast and store the hits
         # TODO: Make this work for multiple meshes?
         self._data.ray_hits_w[env_ids] = raycast_mesh(
@@ -259,6 +287,9 @@ class RayCaster(SensorBase):
             mesh=self.meshes[self.cfg.mesh_prim_paths[0]],
         )[0]
 
+        # apply vertical drift to ray starting position in ray caster frame
+        self._data.ray_hits_w[env_ids, :, 2] += self.ray_cast_drift[env_ids, 2].unsqueeze(-1)
+
     def _set_debug_vis_impl(self, debug_vis: bool):
         # set visibility of markers
         # note: parent only deals with callbacks. not their visibility

source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera_cfg.py

@@ -60,4 +60,4 @@ class RayCasterCameraCfg(RayCasterCfg):
 
     def __post_init__(self):
         # for cameras, this quantity should be False always.
-        self.attach_yaw_only = False
+        self.ray_alignment = "base"

source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py

@@ -7,6 +7,7 @@
 
 
 from dataclasses import MISSING
+from typing import Literal
 
 from isaaclab.markers import VisualizationMarkersCfg
 from isaaclab.markers.config import RAY_CASTER_MARKER_CFG
@@ -43,10 +44,22 @@ class RayCasterCfg(SensorBaseCfg):
     offset: OffsetCfg = OffsetCfg()
     """The offset pose of the sensor's frame from the sensor's parent frame. Defaults to identity."""
 
-    attach_yaw_only: bool = MISSING
+    attach_yaw_only: bool = False
     """Whether the rays' starting positions and directions only track the yaw orientation.
 
     This is useful for ray-casting height maps, where only yaw rotation is needed.
+
+    .. warning::
+
+        This attribute is deprecated. Use :attr:`~isaaclab.sensors.ray_caster.ray_caster_cfg.ray_alignment` instead.
+        To get the same behavior, set `ray_alignment` to `"yaw"`.
+    """
+
+    ray_alignment: Literal["base", "yaw", "world"] = "yaw"
+    """Specify in what frame the rays are projected onto the ground. Default is `world`.
+        * `base` if the rays' starting positions and directions track the full root position and orientation.
+        * `yaw` if the rays' starting positions and directions track root position and only yaw component of orientation. This is useful for ray-casting height maps.
+        * `world` if rays' starting positions and directions are always fixed. This is useful in combination with the grid map package.
     """
 
     pattern_cfg: PatternBaseCfg = MISSING
@@ -56,7 +69,13 @@ class RayCasterCfg(SensorBaseCfg):
     """Maximum distance (in meters) from the sensor to ray cast to. Defaults to 1e6."""
 
     drift_range: tuple[float, float] = (0.0, 0.0)
-    """The range of drift (in meters) to add to the ray starting positions (xyz). Defaults to (0.0, 0.0).
+    """The range of drift (in meters) to add to the ray starting positions (xyz) in world frame. Defaults to (0.0, 0.0).
+
+    For floating base robots, this is useful for simulating drift in the robot's pose estimation.
+    """
+
+    ray_cast_drift_range: dict[str, tuple[float, float]] = {"x": (0.0, 0.0), "y": (0.0, 0.0), "z": (0.0, 0.0)}
+    """The range of drift (in meters) to add to the projected ray points in local projection frame. Defaults to (0.0, 0.0) for x, y, and z drift.
 
     For floating base robots, this is useful for simulating drift in the robot's pose estimation.
     """

source/isaaclab/test/envs/check_manager_based_env_anymal_locomotion.py

@@ -92,7 +92,7 @@ class MySceneCfg(InteractiveSceneCfg):
     height_scanner = RayCasterCfg(
         prim_path="{ENV_REGEX_NS}/Robot/base",
         offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
-        attach_yaw_only=True,
+        ray_alignment="yaw",
         pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
         debug_vis=True,
         mesh_prim_paths=["/World/ground"],

source/isaaclab/test/markers/check_markers_visibility.py

@@ -70,7 +70,7 @@ class SensorsSceneCfg(InteractiveSceneCfg):
         prim_path="{ENV_REGEX_NS}/Robot/base",
         update_period=0.02,
         offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
-        attach_yaw_only=True,
+        ray_alignment="yaw",
         pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
         debug_vis=True,
         mesh_prim_paths=["/World/defaultGroundPlane"],

source/isaaclab/test/scene/check_interactive_scene.py

@@ -62,7 +62,7 @@ class MySceneCfg(InteractiveSceneCfg):
     height_scanner = RayCasterCfg(
         prim_path="{ENV_REGEX_NS}/Robot_1/base",
         offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
-        attach_yaw_only=True,
+        ray_alignment="yaw",
         pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
         debug_vis=True,
         mesh_prim_paths=["/World/ground"],

source/isaaclab/test/sensors/check_ray_caster.py

@@ -123,7 +123,7 @@ def main():
         prim_path="/World/envs/env_.*/ball",
         mesh_prim_paths=["/World/ground"],
         pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=(1.6, 1.0)),
-        attach_yaw_only=True,
+        ray_alignment="yaw",
         debug_vis=not args_cli.headless,
     )
     ray_caster = RayCaster(cfg=ray_caster_cfg)

source/isaaclab_assets/isaaclab_assets/sensors/velodyne.py

@@ -13,7 +13,7 @@ from isaaclab.sensors import RayCasterCfg, patterns
 ##
 
 VELODYNE_VLP_16_RAYCASTER_CFG = RayCasterCfg(
-    attach_yaw_only=False,
+    ray_alignment="base",
     pattern_cfg=patterns.LidarPatternCfg(
         channels=16, vertical_fov_range=(-15.0, 15.0), horizontal_fov_range=(-180.0, 180.0), horizontal_res=0.2
     ),

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/velocity_env_cfg.py

@@ -66,7 +66,7 @@ class MySceneCfg(InteractiveSceneCfg):
     height_scanner = RayCasterCfg(
         prim_path="{ENV_REGEX_NS}/Robot/base",
         offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
-        attach_yaw_only=True,
+        ray_alignment="yaw",
         pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
         debug_vis=False,
         mesh_prim_paths=["/World/ground"],