Fix replay crash on trailing idle action in IK tasks

scripts/tools/replay_demos.py

@@ -215,6 +215,9 @@ def replay_episodes_loop(  # noqa: C901
                     else:
                         has_next_action = True
                     actions[env_id] = env_next_action
+                if not has_next_action:
+                    # Stop before stepping once every environment has exhausted its recorded actions.
+                    break
                 if first_loop:
                     first_loop = False
                 else:

source/isaaclab/changelog.d/rwiltz-diff-ik-degenerate-quat.rst

@@ -0,0 +1,7 @@
+Fixed
+^^^^^
+
+* Fixed demonstration replay stepping once after all episodes completed.
+* Fixed :meth:`~isaaclab.controllers.DifferentialIKController.set_command` handling of
+  unnormalizable absolute-pose quaternions, which produced a NaN target orientation. Such
+  commands now hold the current end-effector orientation, or identity when none is provided.

source/isaaclab/isaaclab/controllers/differential_ik.py

@@ -83,6 +83,8 @@ def __init__(self, cfg: DifferentialIKControllerCfg, num_envs: int, device: str)
         # -- optional joint position limits for null-space joint-limit avoidance (set externally)
         self._joint_pos_lower = None
         self._joint_pos_upper = None
+        # -- identity quaternion (x, y, z, w), the last-resort fallback for a degenerate command
+        self._identity_quat = torch.tensor([0.0, 0.0, 0.0, 1.0], device=self._device).repeat(self.num_envs, 1)
 
     """
     Properties.
@@ -119,12 +121,17 @@ def set_command(
         It is up to the user to ensure that the command is given in the correct frame. The method only
         applies the relative mode if the command type is ``position_rel`` or ``pose_rel``.
 
+        Absolute ``pose`` commands normalize finite quaternions; unnormalizable entries use
+        :paramref:`ee_quat`, or identity when :paramref:`ee_quat` is omitted.
+
         Args:
             command: The input command in shape (N, 3) or (N, 6) or (N, 7).
             ee_pos: The current end-effector position in shape (N, 3).
                 This is only needed if the command type is ``position_rel`` or ``pose_rel``.
             ee_quat: The current end-effector orientation (x, y, z, w) in shape (N, 4).
-                This is only needed if the command type is ``position_*`` or ``pose_rel``.
+                This is needed if the command type is ``position_*`` or ``pose_rel``. For absolute
+                ``pose`` commands it is optional and used only as the fallback orientation for an
+                unnormalizable commanded quaternion.
 
         Raises:
             ValueError: If the command type is ``position_*`` and :attr:`ee_quat` is None.
@@ -158,9 +165,12 @@ def set_command(
                 self.ee_pos_des, self.ee_quat_des = apply_delta_pose(ee_pos, ee_quat, self._command)
             else:
                 self.ee_pos_des = self._command[:, 0:3]
-                # renormalize the commanded quaternion (callers may pass a slightly non-unit quat)
+                # normalize valid quaternions and use the fallback for non-finite results
                 quat = self._command[:, 3:7]
-                self.ee_quat_des = quat / torch.linalg.norm(quat, dim=-1, keepdim=True)
+                normalized_quat = quat / torch.linalg.norm(quat, dim=-1, keepdim=True)
+                is_valid = torch.isfinite(normalized_quat).all(dim=-1, keepdim=True)
+                fallback_quat = self._identity_quat if ee_quat is None else ee_quat
+                self.ee_quat_des = torch.where(is_valid, normalized_quat, fallback_quat)
 
     def set_joint_pos_limits(self, lower: torch.Tensor, upper: torch.Tensor) -> None:
         """Provide the controlled joints' position limits for null-space joint-limit avoidance.

source/isaaclab/test/cli/test_replay_demos_loop_termination.py

@@ -0,0 +1,111 @@
+# 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
+
+"""Regression test for the replay loop stepping past the end of the recorded data.
+
+``replay_episodes_loop`` used to call ``env.step`` once more after every environment had exhausted
+its episodes, applying the untouched idle action. For a task-space (IK) task that idle action is
+all zeros, i.e. a zero-norm quaternion, which crashed the run after a successful replay.
+
+``scripts/tools/replay_demos.py`` launches the simulator at import time, so the loop function is
+extracted from the source and executed against stub objects instead.
+"""
+
+import ast
+import contextlib
+from pathlib import Path
+
+import pytest
+import torch
+
+from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler
+
+pytestmark = pytest.mark.integration
+
+# This test lives at source/isaaclab/test/cli/test_replay_demos_loop_termination.py.
+_REPLAY_DEMOS_PATH = Path(__file__).resolve().parents[4] / "scripts" / "tools" / "replay_demos.py"
+
+
+def _load_replay_episodes_loop(simulation_app):
+    """Compile ``replay_episodes_loop`` from the script source, bound to the given app stub."""
+    source = _REPLAY_DEMOS_PATH.read_text()
+    tree = ast.parse(source)
+    func = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "replay_episodes_loop")
+    namespace = {
+        "contextlib": contextlib,
+        "torch": torch,
+        "EpisodeData": EpisodeData,
+        "HDF5DatasetFileHandler": HDF5DatasetFileHandler,
+        "simulation_app": simulation_app,
+        "is_paused": False,
+    }
+    exec(compile(ast.Module(body=[func], type_ignores=[]), str(_REPLAY_DEMOS_PATH), "exec"), namespace)
+    return namespace["replay_episodes_loop"]
+
+
+class _SimulationAppStub:
+    def is_running(self):
+        return True
+
+    def is_exiting(self):
+        return False
+
+
+class _EnvStub:
+    """Records the actions passed to :meth:`step`."""
+
+    device = "cpu"
+
+    def __init__(self):
+        self.stepped_actions: list[torch.Tensor] = []
+
+    def reset_to(self, state, env_ids, is_relative=True):
+        pass
+
+    def step(self, actions):
+        self.stepped_actions.append(actions.clone())
+
+
+class _DatasetFileHandlerStub:
+    def __init__(self, actions: torch.Tensor):
+        self._actions = actions
+
+    def load_episode(self, episode_name, device):
+        episode = EpisodeData()
+        episode.data = {"initial_state": {}, "actions": list(self._actions)}
+        return episode
+
+
+def test_replay_loop_does_not_step_after_the_recorded_actions():
+    """The loop steps exactly once per recorded action and never applies the idle action."""
+    # absolute task-space actions: [pos_xyz, quat_xyzw, gripper]
+    recorded_actions = torch.tensor(
+        [
+            [0.30, -0.10, 0.20, 0.0, 0.0, 0.0, 1.0, 0.0],
+            [0.31, -0.10, 0.20, 0.0, 0.0, 0.0, 1.0, 0.0],
+            [0.32, -0.10, 0.20, 0.0, 0.0, 0.0, 1.0, 0.0],
+        ]
+    )
+    idle_action = torch.zeros(1, recorded_actions.shape[-1])
+    env = _EnvStub()
+    replay_episodes_loop = _load_replay_episodes_loop(_SimulationAppStub())
+
+    replayed_episode_count, _, _ = replay_episodes_loop(
+        env,
+        _DatasetFileHandlerStub(recorded_actions),
+        episode_names=["demo_0"],
+        episode_count=1,
+        episode_indices_to_replay=[0],
+        num_envs=1,
+        success_term=None,
+        state_validation_enabled=False,
+        idle_action=idle_action,
+        reset_sim_buffer_each_episode=False,
+    )
+
+    assert replayed_episode_count == 1
+    assert len(env.stepped_actions) == len(recorded_actions)
+    for stepped, recorded in zip(env.stepped_actions, recorded_actions):
+        torch.testing.assert_close(stepped, recorded.unsqueeze(0), atol=1e-6, rtol=0.0)

source/isaaclab/test/controllers/test_differential_ik_features.py

@@ -86,6 +86,39 @@ def test_set_command_renormalizes_quat():
     torch.testing.assert_close(stored, raw / torch.linalg.norm(raw), atol=1e-6, rtol=0.0)
 
 
+@pytest.mark.parametrize("bad_quat", [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1e-38]])
+def test_set_command_unnormalizable_quat_holds_current_orientation(bad_quat):
+    """An unnormalizable commanded quaternion holds that env's current orientation instead of NaN."""
+    c = _make_controller(num_envs=2)
+    good_quat = _quat_xyzw([0.0, 1.0, 0.0], 0.4)
+    held_quat = _quat_xyzw([1.0, 0.0, 0.0], 0.5)
+    ee_pos = torch.tensor([[0.3, -0.1, 0.2], [0.3, -0.1, 0.2]])
+    ee_quat = torch.tensor([_ID_QUAT, held_quat])
+    cmd = torch.tensor([[0.3, -0.1, 0.2] + good_quat, [0.3, -0.1, 0.2] + bad_quat])
+    c.set_command(cmd, ee_pos, ee_quat)
+    torch.testing.assert_close(c.ee_quat_des[0], torch.tensor(good_quat), atol=1e-6, rtol=0.0)
+    torch.testing.assert_close(c.ee_quat_des[1], torch.tensor(held_quat), atol=1e-6, rtol=0.0)
+
+
+def test_set_command_unnormalizable_quat_without_current_orientation_is_identity():
+    """Without a current orientation to hold, an unnormalizable command falls back to identity."""
+    c = _make_controller()
+    cmd = torch.cat([torch.tensor([[0.3, -0.1, 0.2]]), torch.zeros(1, 4)], dim=-1)
+    c.set_command(cmd)
+    torch.testing.assert_close(c.ee_quat_des, torch.tensor([_ID_QUAT]), atol=1e-6, rtol=0.0)
+
+
+@pytest.mark.parametrize("scale", [1e-7, 1e-20])
+def test_set_command_tiny_normalizable_quat_is_still_normalized(scale):
+    """A tiny but normalizable quaternion keeps its meaning instead of taking the fallback."""
+    c = _make_controller()
+    ee_pos = torch.tensor([[0.3, -0.1, 0.2]])
+    ee_quat = torch.tensor([_quat_xyzw([1.0, 0.0, 0.0], 0.5)])  # a fallback that is NOT identity
+    cmd = torch.cat([ee_pos, torch.tensor([[0.0, 0.0, 0.0, scale]])], dim=-1)  # scaled identity
+    c.set_command(cmd, ee_pos, ee_quat)
+    torch.testing.assert_close(c.ee_quat_des, torch.tensor([_ID_QUAT]), atol=1e-5, rtol=0.0)
+
+
 def test_orientation_weight_none_is_unweighted():
     """With no orientation weight, the pose task Jacobian equals the raw Jacobian."""
     c = _make_controller(orientation_weight=None)