diff --git a/scripts/tools/replay_demos.py b/scripts/tools/replay_demos.py index e8ecae95f118..dd3273b9f1c8 100644 --- a/scripts/tools/replay_demos.py +++ b/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: diff --git a/source/isaaclab/changelog.d/rwiltz-diff-ik-degenerate-quat.rst b/source/isaaclab/changelog.d/rwiltz-diff-ik-degenerate-quat.rst new file mode 100644 index 000000000000..ebf049384a8d --- /dev/null +++ b/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. diff --git a/source/isaaclab/isaaclab/controllers/differential_ik.py b/source/isaaclab/isaaclab/controllers/differential_ik.py index 51f399ba4be7..08b469a41ed3 100644 --- a/source/isaaclab/isaaclab/controllers/differential_ik.py +++ b/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. diff --git a/source/isaaclab/test/cli/test_replay_demos_loop_termination.py b/source/isaaclab/test/cli/test_replay_demos_loop_termination.py new file mode 100644 index 000000000000..d390abc14fc1 --- /dev/null +++ b/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) diff --git a/source/isaaclab/test/controllers/test_differential_ik_features.py b/source/isaaclab/test/controllers/test_differential_ik_features.py index e4ac4a041efb..1789fc71da3c 100644 --- a/source/isaaclab/test/controllers/test_differential_ik_features.py +++ b/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)