Improvements of MDP terms in nav_tasks and bugfix in trajectory sampler (#5)

# Description

### ✨ Additions & Enhancements
* `NavigationSE2Action` momentum support.
* **Height-scan observations**: occlusion-aware utilities incl. door
recognition.
* **Curriculum**: new helper `change_reward_param`.
* **Terrain generation**: `StairsRampTerrainCfg` now accepts
`random_state_file` for deterministic builds.

### 🐛 Fixes
* **nav_suite 0.2.4**
* `TrajectorySampling`: removed unnecessary `torch.int64` cast that
caused dtype mismatch.
* **nav_tasks 0.3.10**
  * Height-scan clipping now respects min/max bounds.
  * Correct usage of `quat_rotate` after Isaac Lab API change.
  * `TrajectorySampling` index dtype cast to `torch.int64`.
  * Third-party list in `pyproject.toml` updated to new package names.
  * Corrected `GoalCommand.command` docstring (returns `(num_envs, 4)`).
* `FixGoalCommand.command` once again returns **position + heading**
(shape `(num_envs, 4)`) instead of position-only.
 

## Type of change

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./formatter.sh`
- [x] I have made corresponding changes to the documentation (will come
shortly)
- [x] My changes generate no new warnings
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

Author: pascal-roth

exts/nav_suite/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.2.3"
+version = "0.2.4"
 
 # Description
 title = "IsaacLab Navigation Suite"

exts/nav_suite/docs/CHANGELOG.rst

@@ -1,6 +1,16 @@
 Changelog
 ---------
 
+0.2.4 (2025-08-08)
+~~~~~~~~~~~~~~~~~~
+
+Fixes
+^^^^^
+
+- Fixed dtype mismatch in :class:`nav_suite.collectors.TrajectorySampling` by removing the casting of trajectory lengths to ``torch.int64``.
+- Renamed configuration option ``attach_yaw_only`` to  ``ray_alignment`` according to IsaacLab API changes in 2.1.1
+
+
 0.2.3 (2025-06-11)
 ~~~~~~~~~~~~~~~~~~
 

exts/nav_suite/nav_suite/collectors/trajectory_sampling.py

@@ -232,7 +232,7 @@ def sample_paths_by_terrain(
                 clipped = randomized_samples[mask][:samples_per_terrain_level]
                 samples_by_terrain[row, :, :3] = self.terrain_analyser.points[clipped[:, 0].type(torch.int64)]
                 samples_by_terrain[row, :, 3:6] = self.terrain_analyser.points[clipped[:, 1].type(torch.int64)]
-                samples_by_terrain[row, :, 6] = clipped[:, 2].type(torch.int64)
+                samples_by_terrain[row, :, 6] = clipped[:, 2]
         else:
             subterrain_idx_origins = (
                 randomized_samples_subterrains_origins[:, 0] * num_cols + randomized_samples_subterrains_origins[:, 1]
@@ -255,7 +255,7 @@ def sample_paths_by_terrain(
                 clipped = randomized_samples[mask][:samples_per_terrain]
                 samples_by_terrain[row, col, :, :3] = self.terrain_analyser.points[clipped[:, 0].type(torch.int64)]
                 samples_by_terrain[row, col, :, 3:6] = self.terrain_analyser.points[clipped[:, 1].type(torch.int64)]
-                samples_by_terrain[row, col, :, 6] = clipped[:, 2].type(torch.int64)
+                samples_by_terrain[row, col, :, 6] = clipped[:, 2]
 
         # save curr_data as pickle
         if self.cfg.enable_saved_paths_loading:

exts/nav_suite/nav_suite/sensors/matterport_raycaster.py

@@ -145,7 +145,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]):
         self._data.quat_w[env_ids] = quat_w
 
         # ray cast based on the sensor poses
-        if self.cfg.attach_yaw_only:
+        if self.cfg.ray_alignment == "yaw":
             # 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)

exts/nav_tasks/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.3.9"
+version = "0.3.10"
 
 # Description
 title = "IsaacLab Navigation RL Tasks"

exts/nav_tasks/docs/CHANGELOG.rst

@@ -1,6 +1,33 @@
 Changelog
 ---------
 
+0.3.10 (2025-08-08)
+~~~~~~~~~~~~~~~~~~
+
+Changed
+^^^^^^^
+
+- Updated configuration option ``attach_yaw_only`` to ``ray_alignment`` across sensor and environment configurations according to IsaacLab API changes in 2.1.1
+- Updated math calls with :func:`isaaclab.utils.math.quat_rotate` according to IsaacLab API changes in 2.1.1.
+- Updated :meth:`nav_tasks.mdp.commands.FixGoalCommand.command` to return only position plus heading instead of just position.
+
+Added
+^^^^^
+
+- Added ``momentum`` feature and corresponding configuration options to :class:`nav_tasks.mdp.actions.NavigationSE2Action` and :class:`nav_tasks.mdp.actions.NavigationSE2ActionCfg`.
+- Added occlusion-aware height scan utilities that better handle door recognition and terrain occlusions.
+- Added curriculum :func:`nav_tasks.mdp.curriculums.change_reward_param`.
+- Added ``random_state_file`` option from :class:`nav_tasks.terrains.StairsRampTerrainCfg`.
+
+Fixes
+^^^^^
+
+- Fixed height scan clipping in :func:`nav_tasks.mdp.observations.height_scan_clipped` to correctly apply limits.
+- Fixed dtype mismatch in sample selection inside :class:`nav_tasks.collectors.TrajectorySampling` by casting indices to ``torch.int64``.
+- Fixed third-party module list in ``pyproject.toml`` to new Isaac Lab package names.
+- Fixed docstring of :meth:`nav_tasks.mdp.commands.GoalCommand.command` to return pose with shape ``(num_envs, 4)`` instead of ``(num_envs, 7)``.
+
+
 0.3.9 (2025-08-04)
 ~~~~~~~~~~~~~~~~~~
 

exts/nav_tasks/nav_tasks/configs/env_cfg_base.py

@@ -76,7 +76,7 @@ class NavTasksDepthNavSceneCfg(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"],

exts/nav_tasks/nav_tasks/mdp/actions/navigation_actions.py

@@ -88,6 +88,19 @@ def __init__(self, cfg: NavigationSE2ActionCfg, env: ManagerBasedRLEnv):
                     f"Unsupported clip mode: {self.cfg.clip_mode}. Supported modes are 'minmax', 'tanh' and 'none'."
                 )
 
+        # parse momentum
+        if self.cfg.momentum is not None:
+            if isinstance(self.cfg.momentum, (float, int)):
+                self._momentum = float(self.cfg.momentum)
+            elif isinstance(self.cfg.momentum, list):
+                self._momentum = torch.tensor([self.cfg.momentum], device=self.device).repeat(self.num_envs, 1)
+            else:
+                raise ValueError(
+                    f"Unsupported momentum type: {type(self.cfg.momentum)}. Supported types are float, int, and list."
+                )
+        else:
+            self._momentum = 0
+
         # set up buffers
         self._init_buffers()
 
@@ -136,6 +149,11 @@ def process_actions(self, actions):
         self._processed_navigation_velocity_actions *= self._scale
         self._processed_navigation_velocity_actions += self._offset
 
+        # apply a momentum offset
+        self._processed_navigation_velocity_actions += (
+            torch.norm(self._env.scene.articulations["robot"].data.root_lin_vel_b, dim=1, keepdim=True) * self._momentum
+        )
+
     def apply_actions(self):
         """Apply low-level actions for the simulator to the physics engine. This functions is called with the
         simulation frequency of 200Hz. Since low-level locomotion runs at 50Hz, we need to decimate the actions."""

exts/nav_tasks/nav_tasks/mdp/actions/navigation_actions_cfg.py

@@ -64,3 +64,15 @@ class NavigationSE2ActionCfg(ActionTermCfg):
     .. note::
         Offset is applied after scaling. If a list is provided, it must be of the same length as the number of action dimensions.
     """
+
+    momentum: float | list[float] | None = None
+    """Momentum of the action space. Default is None.
+
+    Momentum is computed as the norm of the robot's base linear velocity in the body frame times the momentum factor.
+
+    .. math::
+        momentum = \text{norm}(v_{base}) \times \text{momentum}
+
+    .. note::
+        Momentum is applied after scaling and offset. If a list is provided, it must be of the same length as the number of action dimensions.
+    """

exts/nav_tasks/nav_tasks/mdp/commands/fix_goal_command.py

@@ -12,7 +12,7 @@
 from typing import TYPE_CHECKING
 
 from isaaclab.assets import Articulation
-from isaaclab.utils.math import quat_apply_inverse, sample_uniform, yaw_quat
+from isaaclab.utils.math import quat_apply_inverse, sample_uniform, wrap_to_pi, yaw_quat
 
 from nav_suite.terrain_analysis import TerrainAnalysis, TerrainAnalysisSingleton
 
@@ -41,6 +41,9 @@ def __init__(self, cfg: FixGoalCommandCfg, env: ManagerBasedRLEnv):
         Args:
             cfg: The configuration parameters for the command.
             env: The environment object.
+
+        Raises:
+            AssertionError: If ``trajectory_num_samples`` is not divisible by ``num_envs``.
         """
         super().__init__(cfg, env)
         # -- robot
@@ -94,6 +97,10 @@ def __init__(self, cfg: FixGoalCommandCfg, env: ManagerBasedRLEnv):
             torch.zeros((self.num_envs, 1), device=self.device),
         ])
 
+        # heading command is just straight forward
+        self.heading_command_w = torch.zeros(self.num_envs, device=self.device)
+        self.heading_command_b = torch.zeros_like(self.heading_command_w)
+
         # -- spawn locations (x, y, z, heading)
         self.pos_spawn_w = env.scene.env_origins.clone()
         self.heading_spawn_w = torch.zeros(self.num_envs, device=self.device)
@@ -114,6 +121,11 @@ def __init__(self, cfg: FixGoalCommandCfg, env: ManagerBasedRLEnv):
 
         # EVAL case that maximum number of samples is set
         if self.cfg.trajectory_num_samples is not None:
+            assert self.cfg.trajectory_num_samples % env.num_envs == 0, (
+                "[FixGoalCommand] ``trajectory_num_samples`` must be divisible by ``num_envs`` to allow for equal"
+                " sampling of paths per environment."
+            )
+
             self.all_goals = self.pos_command_w.unsqueeze(1).expand(
                 -1, self.cfg.trajectory_num_samples // env.num_envs, -1
             )
@@ -187,8 +199,8 @@ def __str__(self) -> str:
 
     @property
     def command(self) -> torch.Tensor:
-        """The desired base position in base frame. Shape is (num_envs, 3)."""
-        return self.pos_command_b
+        """The desired base pose in base frame. Shape is (num_envs, 4)."""
+        return torch.cat((self.pos_command_b, self.heading_command_b.unsqueeze(-1)), dim=1)
 
     @property
     def analysis(self) -> TerrainAnalysis | TerrainAnalysisSingleton:
@@ -216,6 +228,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]
             self.nb_sampled_paths = 0
             self.not_updated_envs.fill_(False)
             self.prev_not_updated_envs.fill_(False)
+            self.path_idx_per_env.fill_(0)
 
         return super().reset(env_ids=env_ids)
 
@@ -252,12 +265,19 @@ def _resample_command(self, env_ids: Sequence[int]):
             perturbation[:, 2] = 0.1
             self.pos_spawn_w[env_ids] = self._env.scene.env_origins[env_ids] + perturbation
 
+        # add robot height offset to the spawn position
+        self.pos_spawn_w[env_ids, 2] += self.robot.data.default_root_state[env_ids, 2]
+
     def _update_command(self):
         """Re-target the position command to the current root position."""
         target_vec = self.pos_command_w - self.robot.data.root_pos_w[:, :3]
         target_vec[:, 2] = 0.0  # ignore z component
         self.pos_command_b[:] = quat_apply_inverse(yaw_quat(self.robot.data.root_quat_w), target_vec)
 
+        # update the heading command in the base frame
+        # heading_w is angle world x axis to robot base x axis
+        self.heading_command_b[:] = wrap_to_pi(self.heading_command_w - self.robot.data.heading_w)
+
     def _update_metrics(self):
         """Update metrics."""
         self.metrics["error_pos"] = torch.norm(self.pos_command_w - self.robot.data.root_pos_w[:, :3], dim=1)