From babd6890b98c2c450bf321d2edc9145f6ba9a737 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 4 Aug 2026 15:38:45 -0700 Subject: [PATCH 01/10] Add warp MDP twins for the Reach success terms The Reach task configs gained `rewards.success` and `terminations.success`, which resolve to `is_terminated_term` and `pose_command_success`. Neither had a warp counterpart, so `--frontend warp` rejected Isaac-Reach-Franka and Isaac-Reach-UR10 at env construction. `is_terminated_term` is a ManagerTermBase subclass rather than a plain function, so the twin is a class too: the warp RewardManager instantiates it at _prepare_terms and calls the instance. It resolves the term selection to a column mask once at init and reads the per-term done buffer, which is exposed as a new `term_dones_wp` property. `pose_command_success` recomputes the pose error in a kernel instead of copying the stable term's result. `UniformPoseCommand._compute_success` allocates a fresh tensor per call, so its pointer moves between steps; a captured CUDA graph records source pointers at capture time and would replay against a stale buffer. The command's sticky success tracker stays correct because `_update_metrics` already ORs the same value in on every step. Parity tests cover both twins against the stable implementations, eager and under graph capture. The pose-command fixture derives its command from each body's actual pose plus a known error so the errors straddle the thresholds -- a command sampled independently misses every threshold and compares all-False, which passes against almost any kernel. --- .../isaaclab_experimental/envs/mdp/rewards.py | 59 +++++++- .../envs/mdp/terminations.py | 75 ++++++++++ .../managers/termination_manager.py | 10 ++ .../test/envs/mdp/parity_helpers.py | 139 +++++++++++++++++- .../test/envs/mdp/test_rewards_warp_parity.py | 25 ++++ .../envs/mdp/test_terminations_warp_parity.py | 35 +++++ 6 files changed, 340 insertions(+), 3 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index 338d27c99c5e..1eb07bc9c6c0 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -24,11 +24,12 @@ rotate_vec_to_body_frame, ) -from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.managers import ManagerTermBase, SceneEntityCfg if TYPE_CHECKING: from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.managers.manager_term_cfg import RewardTermCfg """ @@ -68,6 +69,62 @@ def is_terminated(env: ManagerBasedRLEnv, out) -> None: ) +@wp.kernel +def _is_terminated_term_kernel( + term_dones: wp.array(dtype=wp.bool, ndim=2), + term_mask: wp.array(dtype=wp.bool), + time_outs: wp.array(dtype=wp.bool), + out: wp.array(dtype=wp.float32), +): + """Sum the selected termination columns, zeroed on episodic timeouts.""" + i = wp.tid() + count = float(0.0) + for t in range(term_mask.shape[0]): + if term_mask[t] and term_dones[i, t]: + count += 1.0 + out[i] = wp.where(time_outs[i], 0.0, count) + + +class is_terminated_term(ManagerTermBase): + """Penalize termination for specific terms that don't correspond to episodic timeouts. + + Warp-first override of :class:`isaaclab.envs.mdp.rewards.is_terminated_term`. + + The parameters are as follows: + + * attr:`term_keys`: The termination terms to penalize. This can be a string, a list of strings + or regular expressions. Default is ".*" which penalizes all terminations. + + The reward is the number of active selected termination terms, or 0 when the episode ended on + an episodic timeout. + """ + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + # resolve the selection once, as a per-column mask over the termination manager's terms + selected = set(env.termination_manager.find_terms(cfg.params.get("term_keys", ".*"))) + self._term_mask_wp = wp.array( + [name in selected for name in env.termination_manager.active_terms], + dtype=wp.bool, + device=env.device, + ) + + def __call__( + self, env: ManagerBasedRLEnv, out: wp.array(dtype=wp.float32), term_keys: str | list[str] = ".*" + ) -> None: + wp.launch( + kernel=_is_terminated_term_kernel, + dim=env.num_envs, + inputs=[ + env.termination_manager.term_dones_wp, + self._term_mask_wp, + env.termination_manager.time_outs_wp, + out, + ], + device=env.device, + ) + + """ Root penalties. """ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py index d309a70538bb..2cd189e8b18a 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -49,6 +49,81 @@ def time_out(env: ManagerBasedRLEnv, out) -> None: ) +@wp.kernel +def _pose_command_success_kernel( + root_pos_w: wp.array(dtype=wp.vec3f), + root_quat_w: wp.array(dtype=wp.quatf), + body_pos_w: wp.array(dtype=wp.vec3f, ndim=2), + body_quat_w: wp.array(dtype=wp.quatf, ndim=2), + cmd: wp.array(dtype=wp.float32, ndim=2), + body_idx: int, + position_threshold: float, + orientation_threshold: float, + out: wp.array(dtype=wp.bool), +): + """Flag envs whose body pose is within every configured threshold. A negative threshold is unset.""" + i = wp.tid() + success = bool(True) + if position_threshold >= 0.0: + des_b = wp.vec3f(cmd[i, 0], cmd[i, 1], cmd[i, 2]) + des_w = root_pos_w[i] + wp.quat_rotate(root_quat_w[i], des_b) + cur_w = body_pos_w[i, body_idx] + dx = cur_w[0] - des_w[0] + dy = cur_w[1] - des_w[1] + dz = cur_w[2] - des_w[2] + if wp.sqrt(dx * dx + dy * dy + dz * dz) >= position_threshold: + success = False + if orientation_threshold >= 0.0: + des_q_b = wp.quatf(cmd[i, 3], cmd[i, 4], cmd[i, 5], cmd[i, 6]) + des_q_w = root_quat_w[i] * des_q_b + q_err = wp.quat_inverse(body_quat_w[i, body_idx]) * des_q_w + if 2.0 * wp.acos(wp.clamp(wp.abs(q_err[3]), 0.0, 1.0)) >= orientation_threshold: + success = False + out[i] = success + + +def pose_command_success(env: ManagerBasedRLEnv, out, command_name: str) -> None: + """Terminate environments whose pose command satisfies all configured success thresholds. + + Warp-first override of :func:`isaaclab.envs.mdp.terminations.pose_command_success`. + + Note: + The stable term also ORs the result into the command's per-episode success tracker. That + tracker is already maintained by ``UniformPoseCommand._update_metrics`` on every step, so + recomputing the thresholds here keeps ``Metrics/success_rate`` intact. + """ + fn = pose_command_success + command = env.command_manager.get_term(command_name) + position_threshold = command.cfg.position_success_threshold + orientation_threshold = command.cfg.orientation_success_threshold + # matches the stable term: with no threshold configured, no env ever succeeds + if position_threshold is None and orientation_threshold is None: + out.zero_() + return + # cache the warp view of the command tensor on first call (zero-copy) + if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: + cmd = env.command_manager.get_command(command_name) + fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) + fn._cmd_name = command_name + asset: Articulation = command.robot + wp.launch( + kernel=_pose_command_success_kernel, + dim=env.num_envs, + inputs=[ + asset.data.root_pos_w.warp, + asset.data.root_quat_w.warp, + asset.data.body_pos_w.warp, + asset.data.body_quat_w.warp, + fn._cmd_wp, + command.body_idx, + -1.0 if position_threshold is None else position_threshold, + -1.0 if orientation_threshold is None else orientation_threshold, + out, + ], + device=env.device, + ) + + """ Root terminations. """ diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index 732edf9f2fe2..368d9ca5d130 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -229,6 +229,16 @@ def terminated_wp(self) -> wp.array: """The terminated signal (reaching a terminal state). Shape is (num_envs,).""" return self._terminated_wp + @property + def term_dones_wp(self) -> wp.array: + """The per-term done signals. Shape is (num_envs, num_terms). + + Columns follow the order of :attr:`active_terms`. Reward terms that aggregate a + subset of terminations (e.g. :class:`~isaaclab_experimental.envs.mdp.is_terminated_term`) + read this buffer directly instead of the per-term views. + """ + return self._term_dones_wp + """ Operations. """ diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index 4ceec2ecb60b..9da36f29bf5a 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -12,6 +12,9 @@ from __future__ import annotations +import re +from types import SimpleNamespace + import numpy as np import torch import warp as wp @@ -241,6 +244,11 @@ def __init__(self, num_envs=NUM_ENVS, num_joints=NUM_JOINTS, device=DEVICE, seed bpose_np[:, :, 3:7] = [0.0, 0.0, 0.0, 1.0] self.body_pose_w = proxy_array(bpose_np, dtype=wp.transformf, device=device) + # body_pos_w / body_quat_w: the position and rotation slices of body_pose_w, + # mirroring the shorthands on the real articulation data. + self.body_pos_w = proxy_array(bpose_np[:, :, :3].copy(), dtype=wp.vec3f, device=device) + self.body_quat_w = proxy_array(bpose_np[:, :, 3:7].copy(), dtype=wp.quatf, device=device) + # body_lin_acc_w: (num_envs, num_bodies) vec3f self.body_lin_acc_w = proxy_array( rng.randn(num_envs, num_bodies, 3).astype(np.float32), dtype=wp.vec3f, device=device @@ -542,6 +550,109 @@ def get_term(self, name: str): return self._term +def quat_mul_np(q1_xyzw: np.ndarray, q2_xyzw: np.ndarray) -> np.ndarray: + """Hamilton product of two ``(N, 4)`` quaternion arrays in ``(x, y, z, w)`` order.""" + x1, y1, z1, w1 = q1_xyzw.T + x2, y2, z2, w2 = q2_xyzw.T + return np.stack( + [ + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + ], + axis=-1, + ).astype(np.float32) + + +# Error scales for the generated pose commands. Errors are spread over twice these values so +# roughly half the envs fall inside a threshold of the same size — a command sampled +# independently of the body poses would miss every threshold and compare all-False. +POSE_CMD_POSITION_SCALE = 0.5 +POSE_CMD_ORIENTATION_SCALE = 1.0 + + +def make_pose_command_term( + articulation, + body_idx: int = 0, + num_envs: int = NUM_ENVS, + device: str = DEVICE, + position_success_threshold: float | None = POSE_CMD_POSITION_SCALE, + orientation_success_threshold: float | None = POSE_CMD_ORIENTATION_SCALE, + seed: int = 71, +): + """Build a real :class:`UniformPoseCommand` around mock articulation data. + + Bypasses ``__init__`` (which needs a live scene) but keeps the genuine class, so + ``compute_success`` under test is the shipped stable implementation rather than a + reimplementation of it. + + The command is derived from the body's current pose plus a known error, so the + resulting position and orientation errors straddle + :data:`POSE_CMD_POSITION_SCALE` and :data:`POSE_CMD_ORIENTATION_SCALE`. + + Returns: + The command term, whose ``pose_command_b`` is the (num_envs, 7) command buffer. + """ + from isaaclab.envs.mdp.commands.commands_cfg import UniformPoseCommandCfg + from isaaclab.envs.mdp.commands.pose_command import UniformPoseCommand + + rng = np.random.RandomState(seed) + data = articulation.data + root_pos = data.root_pos_w.torch.cpu().numpy() + root_quat = data.root_quat_w.torch.cpu().numpy() + body_pos = data.body_pos_w.torch.cpu().numpy()[:, body_idx] + body_quat = data.body_quat_w.torch.cpu().numpy()[:, body_idx] + + # position: place the target a known distance off the body, then express it in the root frame + offset_dir = rng.randn(num_envs, 3).astype(np.float32) + offset_dir /= np.linalg.norm(offset_dir, axis=1, keepdims=True) + distance = rng.uniform(0.0, 2.0 * POSE_CMD_POSITION_SCALE, num_envs).astype(np.float32) + target_pos_w = body_pos + offset_dir * distance[:, None] + + # orientation: rotate the body frame by a known angle, then express it in the root frame + axis = rng.randn(num_envs, 3).astype(np.float32) + axis /= np.linalg.norm(axis, axis=1, keepdims=True) + angle = rng.uniform(0.0, 2.0 * POSE_CMD_ORIENTATION_SCALE, num_envs).astype(np.float32) + perturb = np.concatenate([axis * np.sin(angle / 2.0)[:, None], np.cos(angle / 2.0)[:, None]], axis=1) + target_quat_w = quat_mul_np(body_quat, perturb.astype(np.float32)) + + root_quat_conj = root_quat * np.array([-1.0, -1.0, -1.0, 1.0], dtype=np.float32) + command_np = np.zeros((num_envs, 7), dtype=np.float32) + command_np[:, :3] = quat_rotate_inv_np(root_quat, target_pos_w - root_pos) + command_np[:, 3:] = quat_mul_np(root_quat_conj, target_quat_w) + + term = object.__new__(UniformPoseCommand) + term.cfg = UniformPoseCommandCfg( + asset_name="robot", + body_name="body", + resampling_time_range=(1.0, 1.0), + position_success_threshold=position_success_threshold, + orientation_success_threshold=orientation_success_threshold, + ) + term.robot = articulation + term.body_idx = body_idx + term.pose_command_b = torch.tensor(command_np, device=device) + term.pose_command_w = torch.zeros_like(term.pose_command_b) + term._env = SimpleNamespace(num_envs=num_envs, device=device) + term._track_success = position_success_threshold is not None or orientation_success_threshold is not None + term._succeeded = torch.zeros(num_envs, dtype=torch.bool, device=device) + return term + + +class MockPoseCommandManager: + """Command manager serving a single :func:`make_pose_command_term` pose command.""" + + def __init__(self, term): + self._term = term + + def get_command(self, name: str) -> torch.Tensor: + return self._term.pose_command_b + + def get_term(self, name: str): + return self._term + + class MockBodyCfg: """SceneEntityCfg-like object for body-level reward/termination terms.""" @@ -564,12 +675,34 @@ def __init__(self, name="contact_sensor", body_ids=None, device=DEVICE): class MockTerminationManager: - """Mock termination manager providing both torch and warp terminated buffers.""" + """Mock termination manager providing both torch and warp terminated buffers. - def __init__(self, num_envs=NUM_ENVS, device=DEVICE): + Also carries per-term done columns and timeout flags, which terms that aggregate a + subset of terminations (``is_terminated_term``) read on both the stable and warp paths. + """ + + def __init__(self, num_envs=NUM_ENVS, device=DEVICE, term_names=("success", "fell_over", "time_out"), seed=31): self.terminated = torch.zeros(num_envs, dtype=torch.bool, device=device) self.terminated_wp = wp.from_torch(self.terminated) + self.active_terms = list(term_names) + rng = np.random.RandomState(seed) + dones_np = rng.rand(num_envs, len(self.active_terms)) < 0.3 + self.term_dones = torch.tensor(dones_np, dtype=torch.bool, device=device) + self.term_dones_wp = wp.from_torch(self.term_dones) + # the last column doubles as the timeout term, so the masking path is exercised + self.time_outs = self.term_dones[:, -1].clone() + self.time_outs_wp = wp.from_torch(self.time_outs) + + def find_terms(self, name_keys) -> list[str]: + """Resolve term-name patterns the way :meth:`ManagerBase.find_terms` does.""" + if isinstance(name_keys, str): + name_keys = [name_keys] + return [name for name in self.active_terms if any(re.fullmatch(key, name) for key in name_keys)] + + def get_term(self, name: str) -> torch.Tensor: + return self.term_dones[:, self.active_terms.index(name)] + # --------------------------------------------------------------------------- # Art-data mutation helpers (previously duplicated in obs/rew/term test files) @@ -626,5 +759,7 @@ def mutate_body_data( pose_np[:, :, :3] = rng.randn(num_envs, num_bodies, 3).astype(np.float32) pose_np[:, :, 3:7] = [0.0, 0.0, 0.0, 1.0] copy_np_to_wp(art_data.body_pose_w, pose_np) + copy_np_to_wp(art_data.body_pos_w, pose_np[:, :, :3]) + copy_np_to_wp(art_data.body_quat_w, pose_np[:, :, 3:7]) wp.synchronize() diff --git a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py index 69cabcd14a36..82270d6a4f6b 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py @@ -247,6 +247,31 @@ def test_is_terminated(self, warp_env, stable_env, term_mgr): assert_close(actual, expected) assert_close(actual_cap, expected) + @pytest.mark.parametrize( + "term_keys", + [".*", ["success"], ["success", "fell_over"], ["time_out"]], + ids=["all", "single", "subset", "timeout_only"], + ) + def test_is_terminated_term(self, warp_env, stable_env, term_keys): + """The warp twin resolves the term selection to a column mask at init, so every + selection shape is compared against the stable term's per-term lookup.""" + from isaaclab.managers.manager_term_cfg import RewardTermCfg + + params = {"term_keys": term_keys} + stable_fn = stable_rew.is_terminated_term( + RewardTermCfg(func=stable_rew.is_terminated_term, weight=1.0, params=params), stable_env + ) + warp_fn = warp_rew.is_terminated_term( + RewardTermCfg(func=warp_rew.is_terminated_term, weight=1.0, params=params), warp_env + ) + expected = stable_fn(stable_env, term_keys=term_keys) + # selecting only the timeout term must score zero everywhere; every other + # selection must score somewhere, otherwise the comparison proves nothing + assert expected.any() == (term_keys != ["time_out"]) + + assert_close(run_warp_rew(warp_fn, warp_env, term_keys=term_keys), expected) + assert_close(run_warp_rew_captured(warp_fn, warp_env, term_keys=term_keys), expected) + # -- Root penalties --------------------------------------------------------- def test_lin_vel_z_l2(self, warp_env, stable_env, all_joints_cfg): diff --git a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py index d82339e01228..9dc7be922d4a 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py @@ -7,6 +7,8 @@ from __future__ import annotations +from types import SimpleNamespace + import numpy as np import pytest import torch @@ -33,11 +35,13 @@ MockCommandTerm, MockContactSensor, MockContactSensorData, + MockPoseCommandManager, MockScene, MockSceneEntityCfg, MockSensorCfg, MockTerminationManager, assert_equal, + make_pose_command_term, mutate_art_data, mutate_body_data, run_warp_term, @@ -268,6 +272,37 @@ def test_illegal_contact(self, warp_env_bodies, stable_env_bodies, sensor_cfg): assert_equal(actual_cap, expected) +class TestPoseCommandSuccessParity: + """The warp twin recomputes the pose error in a kernel, so it is checked against the + stable term for every threshold combination — a frame or quaternion-convention slip + would otherwise resolve cleanly and silently train against a different MDP.""" + + @pytest.mark.parametrize( + ("position_threshold", "orientation_threshold"), + [(0.5, 1.0), (0.5, None), (None, 1.0), (None, None)], + ids=["both", "position_only", "orientation_only", "neither"], + ) + def test_pose_command_success(self, scene_bodies, position_threshold, orientation_threshold): + command = make_pose_command_term( + scene_bodies["robot"], + position_success_threshold=position_threshold, + orientation_success_threshold=orientation_threshold, + ) + env = SimpleNamespace( + scene=scene_bodies, + command_manager=MockPoseCommandManager(command), + num_envs=NUM_ENVS, + device=DEVICE, + ) + expected = stable_term.pose_command_success(env, command_name="ee_pose") + # a degenerate all-False expectation would pass against almost any kernel + if position_threshold is not None or orientation_threshold is not None: + assert expected.any(), "thresholds produced no successes; the comparison would be vacuous" + + assert_equal(run_warp_term(warp_term.pose_command_success, env, command_name="ee_pose"), expected) + assert_equal(run_warp_term_captured(warp_term.pose_command_success, env, command_name="ee_pose"), expected) + + # ============================================================================ # Capture-then-mutate-then-replay termination tests (from test_mdp_warp_parity.py) # ============================================================================ From 0a8a27d294be71da153b8ce3022cbea462323c33 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 4 Aug 2026 15:39:00 -0700 Subject: [PATCH 02/10] Derive the warp-supported task set from the registry The list of tasks that `--frontend warp` accepts was maintained by hand, so it only covered a task once someone remembered to add it. Sweeping the gym registry instead shows the record had already drifted: four contrib velocity tasks adapt cleanly and were never listed. The sweep runs `check_compatibility` over every registered stable manager-based task. That is pure config-tree work with no simulator, so all 104 resolve in about ten seconds. The checked-in set is now a regression pin rather than an allowlist. Losing a task always fails. Gaining one fails too, but only when every candidate config was importable, since an environment missing optional dependencies cannot see the full set; those tasks are named in the skip message rather than silently dropped. `check_compatibility` calls `adapt_cfg` and returns the reason instead of raising, so the verdict cannot drift from what the frontend actually does. --- .../jichuanh-warp-reach-mdp-twins.minor.rst | 11 ++ .../isaaclab_experimental/envs/frontend.py | 26 ++++ .../test/envs/test_frontend_cfg_conversion.py | 139 ++++++++++++++---- 3 files changed, 147 insertions(+), 29 deletions(-) create mode 100644 source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst new file mode 100644 index 000000000000..1e962ceb28b0 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst @@ -0,0 +1,11 @@ +Added +^^^^^ + +* Added Warp MDP twins for :class:`~isaaclab.envs.mdp.is_terminated_term` and + :func:`~isaaclab.envs.mdp.pose_command_success`, so ``Isaac-Reach-Franka`` and + ``Isaac-Reach-UR10`` run under ``--frontend warp``. +* Added :meth:`~isaaclab_experimental.envs.frontend.WarpFrontend.check_compatibility`, which + reports why an environment configuration cannot run on the Warp frontend instead of raising, + so several configurations can be surveyed in one pass. +* Added :attr:`~isaaclab_experimental.managers.TerminationManager.term_dones_wp`, exposing the + per-term done buffer to reward terms that aggregate a subset of termination terms. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index 467bd718bc3f..9da98ad24cba 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -194,6 +194,32 @@ def adapt_cfg(cls, cfg: Any) -> None: cls._swap_mdp(cfg, label) cls._swap_noise_cfgs(cfg, label) + @classmethod + def check_compatibility(cls, cfg: Any) -> str | None: + """Report whether ``cfg`` can run on the warp frontend, without raising. + + Answers the question ``--frontend warp`` answers at env construction, but as a value + instead of an exception, so callers can survey many cfgs and report on all of them. + It runs :meth:`adapt_cfg` itself rather than re-deriving the rules, so the verdict + cannot drift from the real code path. + + Args: + cfg: A stable manager-based env cfg with its physics preset already resolved + (``presets=newton_mjwarp``); an unresolved preset is reported as incompatible. + + Returns: + ``None`` when the cfg adapts, otherwise the reason, listing every missing twin. + + Note: + ``cfg`` is adapted in place when compatible, and may be left partially adapted + when not. Pass a freshly constructed cfg when only the verdict is wanted. + """ + try: + cls.adapt_cfg(cfg) + except FrontendIncompatibleError as exc: + return str(exc) + return None + @staticmethod def _require_newton_physics(cfg: Any, label: str) -> None: """Block unless ``cfg.sim.physics`` is :class:`NewtonCfg`. diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 76a50b6211ad..8a22e3ff9ab1 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -5,23 +5,24 @@ """Coverage tests for warp task configuration adaptation. -Two sweeps, both running :meth:`WarpFrontend.adapt_cfg` — the same adaptation -:class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` runs in its -``__init__``: - -* Every *stable* task id with declared warp twin coverage: resolve the - ``newton_mjwarp`` preset and adapt the stable cfg, exactly what - ``--frontend warp`` does. -* Every registered manager-based ``*-Warp-v0`` variant (velocity deltas whose - stable cfgs still contain terms without warp twins): load its env cfg, - resolve presets, and adapt. - -A dedicated stable Cartpole case also guards the stable-to-experimental -module routing used by ``--frontend warp``. +Which tasks ``--frontend warp`` accepts is *derived*, not declared: every registered stable +manager-based task is swept through :meth:`WarpFrontend.check_compatibility` — the same +:meth:`WarpFrontend.adapt_cfg` that +:class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` runs in its ``__init__``. The +sweep is pure config-tree work, so the whole registry resolves in seconds without a +simulator, and a task gains coverage the moment its terms are twinned. + +:data:`_WARP_SUPPORTED_TASKS` is the checked-in record of that derived set. It is a +regression pin, not an allowlist: losing a task always fails, and gaining one fails too +whenever every candidate cfg was importable, so the record cannot quietly rot. + +A dedicated stable Cartpole case also guards the stable-to-experimental module routing used +by ``--frontend warp``. """ from __future__ import annotations +import functools import importlib import gymnasium as gym @@ -34,21 +35,31 @@ import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils.hydra import resolve_presets -# Stable task ids whose MDP terms are fully twinned: ``--frontend warp`` adapts -# their cfgs directly, with no parallel ``*-Warp-v0`` registration. Extend this -# list when a new task family gains full twin coverage. -_STABLE_TASKS_WITH_WARP_COVERAGE = [ - "Isaac-Ant", - "Isaac-Cartpole", - "Isaac-Humanoid", - "Isaac-Reach-Franka", - "Isaac-Reach-UR10", - "Isaac-Velocity-Flat-AnymalD", - "Isaac-Velocity-Flat-Cassie", - "Isaac-Velocity-Flat-G1", - "Isaac-Velocity-Flat-H1", - "Isaac-Velocity-Flat-UnitreeGo2", -] +# Stable manager-based tasks resolve to this env class; direct tasks provide their own and +# take the :meth:`WarpFrontend._resolve_direct_warp_class` path instead of cfg adaptation. +_STABLE_MANAGER_ENTRY_POINT = "isaaclab.envs:ManagerBasedRLEnv" + +# Stable task ids that adapt cleanly under ``--frontend warp``, as produced by +# :func:`_sweep_warp_support`. Do not curate this by hand: when the test fails it prints the +# exact set to paste back. +_WARP_SUPPORTED_TASKS = frozenset( + { + "Isaac-Ant", + "Isaac-Cartpole", + "Isaac-Humanoid", + "Isaac-Reach-Franka", + "Isaac-Reach-UR10", + "Isaac-Velocity-Flat-AnymalD", + "Isaac-Velocity-Flat-Cassie", + "Isaac-Velocity-Flat-G1", + "Isaac-Velocity-Flat-H1", + "Isaac-Velocity-Flat-UnitreeGo2", + "IsaacContrib-Velocity-Flat-AnymalB", + "IsaacContrib-Velocity-Flat-AnymalC", + "IsaacContrib-Velocity-Flat-UnitreeA1", + "IsaacContrib-Velocity-Flat-UnitreeGo1", + } +) # Manager-based warp tasks are exactly those whose entry point is the shared warp # env class; direct tasks provide their own env class (resolved by name-based @@ -82,6 +93,76 @@ def _load_adapted_cfg(cfg_entry_point: str): return cfg +@functools.lru_cache(maxsize=1) +def _sweep_warp_support() -> tuple[frozenset[str], dict[str, str], dict[str, str]]: + """Ask every stable manager-based task whether it adapts for warp. + + Cached: the sweep instantiates every registered cfg, so it runs once per session. + + Returns: + ``(supported, incompatible, unimportable)`` — task ids that adapt, task ids that + do not mapped to the reason, and task ids whose cfg could not be built at all + (an optional dependency missing from this environment). + """ + supported: set[str] = set() + incompatible: dict[str, str] = {} + unimportable: dict[str, str] = {} + for task_id, spec in gym.registry.items(): + if spec.entry_point != _STABLE_MANAGER_ENTRY_POINT: + continue + cfg_entry = (spec.kwargs or {}).get("env_cfg_entry_point") + if not isinstance(cfg_entry, str): + continue + try: + module_path, class_name = cfg_entry.split(":") + cfg = getattr(importlib.import_module(module_path), class_name)() + cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) + except Exception as exc: # noqa: BLE001 - any cfg import failure means "cannot judge" + unimportable[task_id] = f"{type(exc).__name__}: {exc}" + continue + reason = WarpFrontend.check_compatibility(cfg) + if reason is None: + supported.add(task_id) + else: + incompatible[task_id] = reason + return frozenset(supported), incompatible, unimportable + + +def _format_task_set(task_ids) -> str: + """Render a task-id set as a paste-ready literal for :data:`_WARP_SUPPORTED_TASKS`.""" + return "{\n" + "".join(f" {task_id!r},\n" for task_id in sorted(task_ids)) + "}" + + +def test_warp_supported_task_set_matches_the_registry(): + """The recorded warp-supported set still matches what the registry actually adapts. + + Losing a task is always a failure — that is a task that used to train under + ``--frontend warp`` and no longer does. Gaining one is a failure only when every + candidate cfg was importable, since a partial environment cannot see the full set. + """ + supported, _, unimportable = _sweep_warp_support() + + lost = _WARP_SUPPORTED_TASKS - supported + assert not lost, "tasks lost warp frontend support:\n " + "\n ".join( + f"{task_id}: {_sweep_warp_support()[1].get(task_id, 'cfg no longer importable')}" for task_id in sorted(lost) + ) + + gained = supported - _WARP_SUPPORTED_TASKS + if unimportable: + # Never claim coverage that was not measured: name what this environment could not judge. + assert not gained, ( + f"newly warp-supported tasks: {sorted(gained)}. Update _WARP_SUPPORTED_TASKS to:\n" + f"{_format_task_set(supported)}" + ) + pytest.skip( + f"{len(unimportable)} task cfg(s) not importable here, so the set was not checked for" + f" completeness: {sorted(unimportable)}" + ) + assert not gained, ( + f"tasks gained warp frontend support. Update _WARP_SUPPORTED_TASKS to:\n{_format_task_set(supported)}" + ) + + def _cfg_entry_point(task_id: str) -> str: cfg_entry = (gym.spec(task_id).kwargs or {}).get("env_cfg_entry_point") assert isinstance(cfg_entry, str), f"{task_id}: no env_cfg_entry_point" @@ -99,7 +180,7 @@ def test_no_warp_task_registrations_remain(): assert warp_ids == [], f"unexpected warp task registrations: {warp_ids}" -@pytest.mark.parametrize("task_id", _STABLE_TASKS_WITH_WARP_COVERAGE, ids=_STABLE_TASKS_WITH_WARP_COVERAGE) +@pytest.mark.parametrize("task_id", sorted(_WARP_SUPPORTED_TASKS), ids=sorted(_WARP_SUPPORTED_TASKS)) def test_stable_task_cfg_adapts_to_warp(task_id: str): """Each covered stable task adapts without a missing twin (the --frontend warp path).""" cfg = _load_adapted_cfg(_cfg_entry_point(task_id)) From 4ad010ed2ff9746ed8efc182aaed943f46395daa Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 4 Aug 2026 15:39:03 -0700 Subject: [PATCH 03/10] Run the isaaclab_experimental tests in CI Every test job either names a package with a positive filter or, for the core jobs, excludes anything matching `isaaclab_`. No job named this package, so all nine of its test files fell between the two forms and were collected by nothing. Replaying the filter logic over the tree shows test files matched by no job dropping from thirteen to four; the remainder belong to isaaclab_ppisp and are left for a separate change. --- .github/workflows/build.yaml | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 2524575de36f..04bf93602155 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -529,6 +529,27 @@ jobs: filter-pattern: "isaaclab_assets" container-name: isaac-lab-assets-test + test-isaaclab-experimental: + name: isaaclab_experimental + runs-on: [self-hosted, gpu] + timeout-minutes: 180 + needs: [build, config] + if: >- + github.event_name != 'push' && + needs.build.result == 'success' + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 1 + lfs: true + - uses: ./.github/actions/run-package-tests + with: + image-tag: ${{ needs.config.outputs.ci_image_tag }} + isaacsim-base-image: ${{ needs.config.outputs.isaacsim_image_name }} + isaacsim-version: ${{ needs.config.outputs.isaacsim_image_tag }} + filter-pattern: "isaaclab_experimental" + container-name: isaac-lab-experimental-test + test-isaaclab-newton: name: isaaclab_newton runs-on: [self-hosted, gpu] From 55029cfd600b65945e72869b2d8d51f978a40989 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 4 Aug 2026 15:56:03 -0700 Subject: [PATCH 04/10] Make pose_command_success a class term Resolving the command term, its thresholds, the tracked body index and the warp view of the command buffer per call left non-capturable work on the hot path, including a hasattr-guarded warm-up cache on the function object. All of it is fixed for the run, so it belongs in __init__; __call__ is now a single kernel launch. The command manager is constructed before the termination manager, so the lookup is available at init. --- .../envs/mdp/terminations.py | 76 +++++++++++-------- .../envs/mdp/test_terminations_warp_parity.py | 9 ++- 2 files changed, 52 insertions(+), 33 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py index 2cd189e8b18a..6cc820378713 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -19,11 +19,17 @@ import warp as wp -from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.managers import ManagerTermBase, SceneEntityCfg if TYPE_CHECKING: from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.managers.manager_term_cfg import TerminationTermCfg + + +def _threshold_or_unset(threshold: float | None) -> float: + """Map an optional success threshold to the kernel's negative-means-unset encoding.""" + return -1.0 if threshold is None else threshold """ @@ -82,46 +88,54 @@ def _pose_command_success_kernel( out[i] = success -def pose_command_success(env: ManagerBasedRLEnv, out, command_name: str) -> None: +class pose_command_success(ManagerTermBase): """Terminate environments whose pose command satisfies all configured success thresholds. Warp-first override of :func:`isaaclab.envs.mdp.terminations.pose_command_success`. + The command term, its thresholds, the tracked body index and the warp view of the command + buffer are all resolved at init, leaving :meth:`__call__` a single kernel launch. + Note: The stable term also ORs the result into the command's per-episode success tracker. That tracker is already maintained by ``UniformPoseCommand._update_metrics`` on every step, so recomputing the thresholds here keeps ``Metrics/success_rate`` intact. """ - fn = pose_command_success - command = env.command_manager.get_term(command_name) - position_threshold = command.cfg.position_success_threshold - orientation_threshold = command.cfg.orientation_success_threshold - # matches the stable term: with no threshold configured, no env ever succeeds - if position_threshold is None and orientation_threshold is None: - out.zero_() - return - # cache the warp view of the command tensor on first call (zero-copy) - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: + + def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + command_name = cfg.params["command_name"] + command = env.command_manager.get_term(command_name) + # a negative threshold marks "not configured" for the kernel + self._position_threshold = _threshold_or_unset(command.cfg.position_success_threshold) + self._orientation_threshold = _threshold_or_unset(command.cfg.orientation_success_threshold) + # matches the stable term: with no threshold configured, no env ever succeeds + self._any_threshold = self._position_threshold >= 0.0 or self._orientation_threshold >= 0.0 + self._body_idx = command.body_idx + self._asset: Articulation = command.robot cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - asset: Articulation = command.robot - wp.launch( - kernel=_pose_command_success_kernel, - dim=env.num_envs, - inputs=[ - asset.data.root_pos_w.warp, - asset.data.root_quat_w.warp, - asset.data.body_pos_w.warp, - asset.data.body_quat_w.warp, - fn._cmd_wp, - command.body_idx, - -1.0 if position_threshold is None else position_threshold, - -1.0 if orientation_threshold is None else orientation_threshold, - out, - ], - device=env.device, - ) + self._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) + + def __call__(self, env: ManagerBasedRLEnv, out: wp.array(dtype=wp.bool), command_name: str) -> None: + if not self._any_threshold: + out.zero_() + return + wp.launch( + kernel=_pose_command_success_kernel, + dim=env.num_envs, + inputs=[ + self._asset.data.root_pos_w.warp, + self._asset.data.root_quat_w.warp, + self._asset.data.body_pos_w.warp, + self._asset.data.body_quat_w.warp, + self._cmd_wp, + self._body_idx, + self._position_threshold, + self._orientation_threshold, + out, + ], + device=env.device, + ) """ diff --git a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py index 9dc7be922d4a..4aa7d9c269f1 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py @@ -49,6 +49,7 @@ ) import isaaclab.envs.mdp.terminations as stable_term +from isaaclab.managers.manager_term_cfg import TerminationTermCfg # ============================================================================ # Fixtures @@ -299,8 +300,12 @@ def test_pose_command_success(self, scene_bodies, position_threshold, orientatio if position_threshold is not None or orientation_threshold is not None: assert expected.any(), "thresholds produced no successes; the comparison would be vacuous" - assert_equal(run_warp_term(warp_term.pose_command_success, env, command_name="ee_pose"), expected) - assert_equal(run_warp_term_captured(warp_term.pose_command_success, env, command_name="ee_pose"), expected) + params = {"command_name": "ee_pose"} + warp_fn = warp_term.pose_command_success( + TerminationTermCfg(func=warp_term.pose_command_success, params=params), env + ) + assert_equal(run_warp_term(warp_fn, env, **params), expected) + assert_equal(run_warp_term_captured(warp_fn, env, **params), expected) # ============================================================================ From 446f4171267a7f7b618a3c7c227a4494e3778a6a Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 4 Aug 2026 15:57:27 -0700 Subject: [PATCH 05/10] Inline the success-threshold sentinel encoding A two-line ternary used twice does not need a named helper. --- .../isaaclab_experimental/envs/mdp/terminations.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py index 6cc820378713..e7ec0ce81516 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -27,11 +27,6 @@ from isaaclab.managers.manager_term_cfg import TerminationTermCfg -def _threshold_or_unset(threshold: float | None) -> float: - """Map an optional success threshold to the kernel's negative-means-unset encoding.""" - return -1.0 if threshold is None else threshold - - """ MDP terminations. """ @@ -107,8 +102,10 @@ def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedRLEnv): command_name = cfg.params["command_name"] command = env.command_manager.get_term(command_name) # a negative threshold marks "not configured" for the kernel - self._position_threshold = _threshold_or_unset(command.cfg.position_success_threshold) - self._orientation_threshold = _threshold_or_unset(command.cfg.orientation_success_threshold) + position_threshold = command.cfg.position_success_threshold + orientation_threshold = command.cfg.orientation_success_threshold + self._position_threshold = -1.0 if position_threshold is None else position_threshold + self._orientation_threshold = -1.0 if orientation_threshold is None else orientation_threshold # matches the stable term: with no threshold configured, no env ever succeeds self._any_threshold = self._position_threshold >= 0.0 or self._orientation_threshold >= 0.0 self._body_idx = command.body_idx From fe34694a5f23ace69339e4816074a6d9f26bb161 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 5 Aug 2026 00:29:25 -0700 Subject: [PATCH 06/10] Match the stable pose-command success semantics Three divergences from the stable term, each with a regression test that fails when the change is reverted. The sticky per-episode success bit was never set. The stable term ORs it inside compute_success(), which runs before the episode is reset; the twin only wrote the termination output, and _update_metrics() runs after reset() has already read and cleared the tracker, so a terminating step went unrecorded. The write now happens in the kernel rather than after the launch, because the termination manager runs graph-captured by default and host-side work between launches is not replayed. The threshold comparisons were reject predicates ("veto when error >= bound"), which leave a non-finite error reported as success -- terminating and rewarding a diverged environment. They are accept predicates now, matching the stable `error < bound` and failing closed. A negative threshold was encoded as "not configured". UniformPoseCommandCfg accepts any float without validation, so a configured negative bound was silently ignored instead of denying success. Configuredness is carried as its own flag. Thresholds stay resolved at init; the docstring now records that a curriculum mutating them at runtime is not honoured, since a kernel scalar is baked into the graph at capture. --- .../envs/mdp/terminations.py | 51 +++++++++++++------ .../envs/mdp/test_terminations_warp_parity.py | 39 ++++++++++++++ 2 files changed, 74 insertions(+), 16 deletions(-) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py index e7ec0ce81516..1a80d78c4942 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -58,29 +58,40 @@ def _pose_command_success_kernel( body_quat_w: wp.array(dtype=wp.quatf, ndim=2), cmd: wp.array(dtype=wp.float32, ndim=2), body_idx: int, + check_position: bool, position_threshold: float, + check_orientation: bool, orientation_threshold: float, + succeeded: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.bool), ): - """Flag envs whose body pose is within every configured threshold. A negative threshold is unset.""" + """Flag envs whose body pose is within every configured threshold. + + Each threshold is an accept predicate, matching the stable term's ``error < threshold``: a + non-finite error compares false and denies success, where a reject form would let it through. + """ i = wp.tid() success = bool(True) - if position_threshold >= 0.0: + if check_position: des_b = wp.vec3f(cmd[i, 0], cmd[i, 1], cmd[i, 2]) des_w = root_pos_w[i] + wp.quat_rotate(root_quat_w[i], des_b) cur_w = body_pos_w[i, body_idx] dx = cur_w[0] - des_w[0] dy = cur_w[1] - des_w[1] dz = cur_w[2] - des_w[2] - if wp.sqrt(dx * dx + dy * dy + dz * dz) >= position_threshold: - success = False - if orientation_threshold >= 0.0: + success = success and (wp.sqrt(dx * dx + dy * dy + dz * dz) < position_threshold) + if check_orientation: des_q_b = wp.quatf(cmd[i, 3], cmd[i, 4], cmd[i, 5], cmd[i, 6]) des_q_w = root_quat_w[i] * des_q_b q_err = wp.quat_inverse(body_quat_w[i, body_idx]) * des_q_w - if 2.0 * wp.acos(wp.clamp(wp.abs(q_err[3]), 0.0, 1.0)) >= orientation_threshold: - success = False + angle = 2.0 * wp.acos(wp.clamp(wp.abs(q_err[3]), 0.0, 1.0)) + success = success and (angle < orientation_threshold) out[i] = success + # sticky per-episode tracker, mirroring the stable term's ``self._succeeded |= success``. + # It must be written here rather than after the launch: the termination manager runs + # graph-captured by default, and host-side work between launches is not replayed. + if success: + succeeded[i] = True class pose_command_success(ManagerTermBase): @@ -88,30 +99,35 @@ class pose_command_success(ManagerTermBase): Warp-first override of :func:`isaaclab.envs.mdp.terminations.pose_command_success`. - The command term, its thresholds, the tracked body index and the warp view of the command - buffer are all resolved at init, leaving :meth:`__call__` a single kernel launch. + The command term, its thresholds, the tracked body index and the warp views of the command + and success buffers are all resolved at init, leaving :meth:`__call__` a single kernel launch. Note: - The stable term also ORs the result into the command's per-episode success tracker. That - tracker is already maintained by ``UniformPoseCommand._update_metrics`` on every step, so - recomputing the thresholds here keeps ``Metrics/success_rate`` intact. + Thresholds are read once. A curriculum that mutates ``position_success_threshold`` or + ``orientation_success_threshold`` at runtime is not honoured here: the termination manager + runs graph-captured by default, and a kernel scalar is baked into the graph at capture. """ def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedRLEnv): super().__init__(cfg, env) command_name = cfg.params["command_name"] command = env.command_manager.get_term(command_name) - # a negative threshold marks "not configured" for the kernel + # carry "configured" as its own flag: a threshold is any float, so no value can encode absence position_threshold = command.cfg.position_success_threshold orientation_threshold = command.cfg.orientation_success_threshold - self._position_threshold = -1.0 if position_threshold is None else position_threshold - self._orientation_threshold = -1.0 if orientation_threshold is None else orientation_threshold + self._check_position = position_threshold is not None + self._check_orientation = orientation_threshold is not None + self._position_threshold = float(position_threshold or 0.0) + self._orientation_threshold = float(orientation_threshold or 0.0) # matches the stable term: with no threshold configured, no env ever succeeds - self._any_threshold = self._position_threshold >= 0.0 or self._orientation_threshold >= 0.0 + self._any_threshold = self._check_position or self._check_orientation self._body_idx = command.body_idx self._asset: Articulation = command.robot cmd = env.command_manager.get_command(command_name) self._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) + # zero-copy alias of the command's persistent sticky-success buffer (allocated once, never + # reallocated), so the kernel's write lands in the tensor ``reset()`` reads and clears. + self._succeeded_wp = wp.from_torch(command._succeeded) if self._any_threshold else None def __call__(self, env: ManagerBasedRLEnv, out: wp.array(dtype=wp.bool), command_name: str) -> None: if not self._any_threshold: @@ -127,8 +143,11 @@ def __call__(self, env: ManagerBasedRLEnv, out: wp.array(dtype=wp.bool), command self._asset.data.body_quat_w.warp, self._cmd_wp, self._body_idx, + self._check_position, self._position_threshold, + self._check_orientation, self._orientation_threshold, + self._succeeded_wp, out, ], device=env.device, diff --git a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py index 4aa7d9c269f1..15637d2bb0e4 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py @@ -301,12 +301,51 @@ def test_pose_command_success(self, scene_bodies, position_threshold, orientatio assert expected.any(), "thresholds produced no successes; the comparison would be vacuous" params = {"command_name": "ee_pose"} + command._succeeded.zero_() warp_fn = warp_term.pose_command_success( TerminationTermCfg(func=warp_term.pose_command_success, params=params), env ) assert_equal(run_warp_term(warp_fn, env, **params), expected) + # the stable term ORs into the sticky tracker; the twin must too, or the terminating + # step is never recorded before ``reset()`` reads and clears it + assert_equal(command._succeeded, expected) assert_equal(run_warp_term_captured(warp_fn, env, **params), expected) + def _run_both(self, scene_bodies, **thresholds): + """Return ``(stable, warp)`` results for one threshold configuration.""" + command = make_pose_command_term(scene_bodies["robot"], **thresholds) + env = SimpleNamespace( + scene=scene_bodies, + command_manager=MockPoseCommandManager(command), + num_envs=NUM_ENVS, + device=DEVICE, + ) + params = {"command_name": "ee_pose"} + expected = stable_term.pose_command_success(env, **params) + warp_fn = warp_term.pose_command_success( + TerminationTermCfg(func=warp_term.pose_command_success, params=params), env + ) + return expected, run_warp_term(warp_fn, env, **params), command + + def test_non_finite_error_denies_success(self, scene_bodies): + """A diverged env must not be reported as successful (and rewarded for it).""" + art_data = scene_bodies["robot"].data + art_data.body_pos_w.torch[:8, 0, :] = float("nan") + + expected, actual, _ = self._run_both(scene_bodies) + + assert not expected[:8].any(), "stable must deny success on NaN" + assert_equal(actual, expected) + + def test_negative_threshold_is_configured_not_unset(self, scene_bodies): + """A negative threshold is a real (unsatisfiable) bound, not an absent one.""" + expected, actual, _ = self._run_both( + scene_bodies, position_success_threshold=-0.5, orientation_success_threshold=None + ) + + assert not expected.any(), "no distance is below a negative threshold" + assert_equal(actual, expected) + # ============================================================================ # Capture-then-mutate-then-replay termination tests (from test_mdp_warp_parity.py) From cb4437d0f3efec7a382cb4ee502a6d9e091da69b Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 5 Aug 2026 00:29:27 -0700 Subject: [PATCH 07/10] Add a CUDA-graph capture-safety gate for warp MDP terms Managers run graph-captured by default, so a term that is not capture-safe is on the normal execution path, not an edge case. Capture-and-replay alone does not prove safety: replaying against unchanged memory re-reads the same valid bytes, so a term holding a stale pointer still looks correct. Each term here is captured, has its inputs overwritten in place, then replayed and compared against the stable implementation on the new data. Terms are discovered from the warp MDP modules rather than listed, so a new term without a declaration fails the gate. Declaring one takes a CaptureSpec naming its modality and a builder that wires up the environment, parameters and in-place mutation; the runner is shared. Enumerating also separated three groups that convention had blurred: terms already covered by the hand-written mutation classes, three observation terms with no capture coverage anywhere, and eight event terms that cannot be captured at all because they launch over a variable env_ids subset at reset cadence. --- .../test/envs/mdp/test_capture_safety.py | 285 ++++++++++++++++++ 1 file changed, 285 insertions(+) create mode 100644 source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py diff --git a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py new file mode 100644 index 000000000000..a39ef0e3acb1 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -0,0 +1,285 @@ +# 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 + +"""CUDA-graph capture safety for warp MDP terms. + +The manager call switch runs every manager graph-captured by default +(``ManagerCallSwitch.DEFAULT_CONFIG``), so a warp MDP term that is not capture-safe is not an +edge case — it is the normal execution path. A captured graph records launch topology and +pointer values and replays them, so a term must launch a fixed dim, touch only buffers that +were allocated before capture and never move, and do all of its work in kernels: host-side +Python between launches is simply not replayed. + +Capture-and-replay alone cannot prove that. Replaying against unchanged memory re-reads the +same still-valid bytes, so a term that baked a stale pointer or a stale scalar still looks +correct. Every term here is therefore captured, then has its input buffers overwritten +**in place**, then replayed, and the result is compared against the stable implementation on +the *new* data. + +Terms are discovered from the warp MDP modules rather than listed, so +:func:`test_every_warp_mdp_term_is_declared` fails as soon as a new term is added without a +:data:`CAPTURE_SPECS` entry. +""" + +from __future__ import annotations + +import dataclasses +import inspect +from collections.abc import Callable +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import warp as wp + +wp.init() +pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") + +import isaaclab_experimental.envs.mdp.events as warp_events +import isaaclab_experimental.envs.mdp.observations as warp_obs +import isaaclab_experimental.envs.mdp.rewards as warp_rew +import isaaclab_experimental.envs.mdp.terminations as warp_term +from parity_helpers import ( + DEVICE, + NUM_BODIES, + NUM_ENVS, + MockArticulation, + MockArticulationData, + MockPoseCommandManager, + MockScene, + MockTerminationManager, + assert_close, + assert_equal, + make_pose_command_term, + mutate_body_data, +) + +import isaaclab.envs.mdp.rewards as stable_rew +import isaaclab.envs.mdp.terminations as stable_term +from isaaclab.managers.manager_term_cfg import RewardTermCfg, TerminationTermCfg + + +@dataclasses.dataclass(frozen=True) +class CaptureCase: + """One term, wired up and ready to be captured. + + Attributes: + warp_fn: The warp term, already instantiated when it is a class term. + stable_fn: The stable counterpart, called as ``stable_fn(stable_env, **params)``. + warp_env: Environment passed to the warp term. + stable_env: Environment passed to the stable term; may be the same object. + params: Term parameters forwarded to both sides. + mutate: Overwrites every input buffer in place, so replay sees new values. + """ + + warp_fn: Callable + stable_fn: Callable + warp_env: object + stable_env: object + params: dict + mutate: Callable[[], None] + + +@dataclasses.dataclass(frozen=True) +class CaptureSpec: + """Declares how to exercise one warp MDP term under capture. + + Attributes: + name: The term's ``__name__``, matched against the discovered modules. + modality: ``"reward"`` (float32 output) or ``"termination"`` (bool output). + build: Builds a fresh :class:`CaptureCase`; called once per test. + """ + + name: str + modality: str + build: Callable[[], CaptureCase] + + +# --------------------------------------------------------------------------- +# Case builders +# --------------------------------------------------------------------------- + + +def _build_is_terminated_term() -> CaptureCase: + """A reward term reading the termination manager's per-term done columns.""" + manager = MockTerminationManager() + manager.term_dones[::3] = True + manager.time_outs[:] = manager.term_dones[:, -1] + env = SimpleNamespace(termination_manager=manager, num_envs=NUM_ENVS, device=DEVICE) + params = {"term_keys": ["success", "fell_over"]} + + def mutate() -> None: + rng = np.random.RandomState(404) + fresh = torch.tensor(rng.rand(NUM_ENVS, len(manager.active_terms)) < 0.6, device=DEVICE) + manager.term_dones[:] = fresh + manager.time_outs[:] = fresh[:, -1] + wp.synchronize() + + return CaptureCase( + warp_fn=warp_rew.is_terminated_term( + RewardTermCfg(func=warp_rew.is_terminated_term, weight=1.0, params=params), env + ), + stable_fn=stable_rew.is_terminated_term( + RewardTermCfg(func=stable_rew.is_terminated_term, weight=1.0, params=params), env + ), + warp_env=env, + stable_env=env, + params=params, + mutate=mutate, + ) + + +def _build_pose_command_success() -> CaptureCase: + """A termination term reading articulation buffers plus a cached command view.""" + art_data = MockArticulationData(num_bodies=NUM_BODIES) + scene = MockScene( + {"robot": MockArticulation(art_data, num_bodies=NUM_BODIES)}, + wp.array(np.zeros((NUM_ENVS, 3), dtype=np.float32), dtype=wp.vec3f, device=DEVICE), + ) + command = make_pose_command_term(scene["robot"]) + env = SimpleNamespace( + scene=scene, + command_manager=MockPoseCommandManager(command), + num_envs=NUM_ENVS, + device=DEVICE, + ) + params = {"command_name": "ee_pose"} + + def mutate() -> None: + # both the articulation state and the command buffer move, so a stale view of either + # (the term caches a warp alias of the command tensor at init) shows up on replay + mutate_body_data(art_data) + fresh = make_pose_command_term(scene["robot"], seed=505) + command.pose_command_b[:] = fresh.pose_command_b + wp.synchronize() + + return CaptureCase( + warp_fn=warp_term.pose_command_success( + TerminationTermCfg(func=warp_term.pose_command_success, params=params), env + ), + stable_fn=stable_term.pose_command_success, + warp_env=env, + stable_env=env, + params=params, + mutate=mutate, + ) + + +CAPTURE_SPECS: list[CaptureSpec] = [ + CaptureSpec("is_terminated_term", "reward", _build_is_terminated_term), + CaptureSpec("pose_command_success", "termination", _build_pose_command_success), +] + +# Warp MDP terms not yet exercised by this harness. Pre-existing terms only: a *new* term must +# arrive with a CAPTURE_SPECS entry instead of a row here. Several are already covered by the +# hand-written ``TestCapturedDataMutation*`` classes in the parity test files; the rest are an +# explicit backlog rather than a silent gap. +CAPTURE_UNAUDITED: dict[str, str] = { + "action_l2": "covered by TestCapturedDataMutationRewards", + "action_rate_l2": "covered by TestCapturedDataMutationRewards", + "ang_vel_xy_l2": "covered by TestCapturedDataMutationRewards", + "flat_orientation_l2": "covered by TestCapturedDataMutationRewards", + "is_alive": "covered by TestCapturedDataMutationRewards", + "is_terminated": "covered by TestCapturedDataMutationRewards", + "joint_acc_l2": "covered by TestCapturedDataMutationRewards", + "joint_deviation_l1": "covered by TestCapturedDataMutationRewards", + "joint_pos_limits": "covered by TestCapturedDataMutationRewards", + "joint_torques_l2": "covered by TestCapturedDataMutationRewards", + "joint_vel_l1": "covered by TestCapturedDataMutationRewards", + "joint_vel_l2": "covered by TestCapturedDataMutationRewards", + "lin_vel_z_l2": "covered by TestCapturedDataMutationRewards", + "track_ang_vel_z_exp": "covered by TestCapturedDataMutationRewardsNewTerms", + "track_lin_vel_xy_exp": "covered by TestCapturedDataMutationRewardsNewTerms", + "undesired_contacts": "covered by TestCapturedDataMutationRewardsNewTerms", + "illegal_contact": "covered by TestCapturedDataMutationTerminationsNewTerms", + "joint_pos_out_of_manual_limit": "covered by TestCapturedDataMutationTerminations", + "root_height_below_minimum": "covered by TestCapturedDataMutationTerminations", + "time_out": "covered by TestCapturedDataMutationTerminationsNewTerms", + "base_ang_vel": "covered by TestCapturedDataMutationObservations", + "base_lin_vel": "covered by TestCapturedDataMutationObservations", + "base_pos_z": "covered by TestCapturedDataMutationObservations", + "generated_commands": "covered by TestCapturedDataMutationObservations", + "joint_pos": "covered by TestCapturedDataMutationObservations", + "joint_pos_limit_normalized": "covered by TestCapturedDataMutationObservations", + "joint_vel": "covered by TestCapturedDataMutationObservations", + "last_action": "covered by TestCapturedDataMutationObservations", + "projected_gravity": "covered by TestCapturedDataMutationObservations", + # no capture coverage anywhere — these three are the real backlog this gate exposes + "body_incoming_wrench": "no capture coverage yet", + "joint_pos_rel": "no capture coverage yet", + "joint_vel_rel": "no capture coverage yet", + # Event terms run at reset cadence over an env_ids subset, so their launch dim varies per + # call and cannot be recorded into a graph. Capture safety does not apply to them as written. + "apply_external_force_torque": "reset-cadence event term; variable env_ids launch dim", + "push_by_setting_velocity": "reset-cadence event term; variable env_ids launch dim", + "randomize_rigid_body_com": "reset-cadence event term; variable env_ids launch dim", + "randomize_rigid_body_mass": "reset-cadence event term; variable env_ids launch dim", + "randomize_rigid_body_material": "reset-cadence event term; variable env_ids launch dim", + "reset_joints_by_offset": "reset-cadence event term; variable env_ids launch dim", + "reset_joints_by_scale": "reset-cadence event term; variable env_ids launch dim", + "reset_root_state_uniform": "reset-cadence event term; variable env_ids launch dim", +} + + +def _discover_warp_mdp_terms() -> dict[str, str]: + """Return every public warp MDP term, mapped to its defining module name.""" + terms: dict[str, str] = {} + for module in (warp_rew, warp_term, warp_obs, warp_events): + for name, obj in vars(module).items(): + if name.startswith("_") or not (inspect.isfunction(obj) or inspect.isclass(obj)): + continue + if getattr(obj, "__module__", "") == module.__name__: + terms[name] = module.__name__ + return terms + + +def test_every_warp_mdp_term_is_declared(): + """Every warp MDP term is either exercised here or listed as an unaudited pre-existing term.""" + declared = {spec.name for spec in CAPTURE_SPECS} | set(CAPTURE_UNAUDITED) + undeclared = sorted(set(_discover_warp_mdp_terms()) - declared) + + assert not undeclared, ( + "warp MDP terms with no capture declaration: " + + ", ".join(undeclared) + + ". Add a CAPTURE_SPECS entry; CAPTURE_UNAUDITED is for pre-existing terms only." + ) + + +def test_no_term_is_both_specified_and_unaudited(): + """A term moved into the harness must lose its unaudited row, so the backlog stays truthful.""" + overlap = sorted({spec.name for spec in CAPTURE_SPECS} & set(CAPTURE_UNAUDITED)) + + assert not overlap, f"remove from CAPTURE_UNAUDITED, now exercised here: {overlap}" + + +def test_unaudited_terms_still_exist(): + """Drop rows for terms that no longer exist, so the list cannot rot.""" + stale = sorted(set(CAPTURE_UNAUDITED) - set(_discover_warp_mdp_terms())) + + assert not stale, f"CAPTURE_UNAUDITED lists terms that no longer exist: {stale}" + + +@pytest.mark.parametrize("spec", CAPTURE_SPECS, ids=lambda spec: spec.name) +def test_term_is_capture_safe(spec: CaptureSpec): + """Capture the term, overwrite its inputs in place, replay, and match stable on the new data.""" + case = spec.build() + dtype = wp.float32 if spec.modality == "reward" else wp.bool + out = wp.zeros((NUM_ENVS,), dtype=dtype, device=DEVICE) + + case.warp_fn(case.warp_env, out, **case.params) # warm-up outside the capture + with wp.ScopedCapture() as capture: + case.warp_fn(case.warp_env, out, **case.params) + + case.mutate() + wp.capture_launch(capture.graph) + + expected = case.stable_fn(case.stable_env, **case.params) + actual = wp.to_torch(out).clone() + if spec.modality == "reward": + assert_close(actual, expected) + else: + assert_equal(actual, expected) + assert expected.any(), "mutated inputs produced a degenerate expectation; the replay proves little" From 034524ca09c68f32d900e366fedc693035bbed3a Mon Sep 17 00:00:00 2001 From: jichuanh Date: Wed, 5 Aug 2026 08:32:52 -0700 Subject: [PATCH 08/10] Twin the locomotion success and termination reward terms Isaac-Ant and Isaac-Humanoid stopped adapting under --frontend warp when the locomotion reward cfg gained `terminated_penalty` and `survival_success_rate`; neither had a warp counterpart, which the registry sweep reports as lost support. `survival_success_rate` follows the cartpole twin: counts are accumulated on device and the rate is exposed as a persistent tensor view through the reward manager's reset extras, so no host readback happens where the stable term calls .item(). The capture-safety gate now also walks the per-task warp mirrors under isaaclab_tasks_experimental, not just the shared modules, so a twin added to a task mirror is held to the same declaration rule. Enumerating them shows that package has no test directory at all: all twenty pre-existing per-task twins are recorded as having no coverage rather than assumed fine. CaptureSpec gains expect_nonzero so a term whose output is constant by design -- survival_success_rate writes zeros, its real work being in reset() -- can be exercised without the non-degeneracy guard rejecting it. --- .../test/envs/mdp/test_capture_safety.py | 101 +++++++++++++++++- .../core/locomotion/mdp/__init__.pyi | 4 + .../core/locomotion/mdp/rewards.py | 101 ++++++++++++++++++ 3 files changed, 204 insertions(+), 2 deletions(-) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py index a39ef0e3acb1..ad8c5888e008 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -26,7 +26,9 @@ from __future__ import annotations import dataclasses +import importlib import inspect +import pkgutil from collections.abc import Callable from types import SimpleNamespace @@ -42,6 +44,8 @@ import isaaclab_experimental.envs.mdp.observations as warp_obs import isaaclab_experimental.envs.mdp.rewards as warp_rew import isaaclab_experimental.envs.mdp.terminations as warp_term +import isaaclab_tasks_experimental.core as tasks_experimental_core +import isaaclab_tasks_experimental.core.locomotion.mdp.rewards as warp_loco_rew from parity_helpers import ( DEVICE, NUM_BODIES, @@ -61,6 +65,8 @@ import isaaclab.envs.mdp.terminations as stable_term from isaaclab.managers.manager_term_cfg import RewardTermCfg, TerminationTermCfg +import isaaclab_tasks.core.locomotion.mdp.rewards as stable_loco_rew + @dataclasses.dataclass(frozen=True) class CaptureCase: @@ -96,6 +102,12 @@ class CaptureSpec: name: str modality: str build: Callable[[], CaptureCase] + expect_nonzero: bool = True + """Whether the mutated inputs should produce a non-zero expectation. + + Guards against a comparison that holds trivially. Set ``False`` only for terms whose + output is a constant by design, such as a pure-metric term that writes zeros. + """ # --------------------------------------------------------------------------- @@ -168,9 +180,54 @@ def mutate() -> None: ) +def _build_terminated_penalty() -> CaptureCase: + """A reward term reading the termination manager's aggregate terminated flag.""" + manager = MockTerminationManager() + manager.terminated[::4] = True + env = SimpleNamespace(termination_manager=manager, num_envs=NUM_ENVS, device=DEVICE, step_dt=0.02) + + def mutate() -> None: + rng = np.random.RandomState(606) + manager.terminated[:] = torch.tensor(rng.rand(NUM_ENVS) < 0.5, device=DEVICE) + wp.synchronize() + + return CaptureCase( + warp_fn=warp_loco_rew.terminated_penalty, + stable_fn=stable_loco_rew.terminated_penalty, + warp_env=env, + stable_env=env, + params={}, + mutate=mutate, + ) + + +def _build_survival_success_rate() -> CaptureCase: + """A pure-metric reward term: the output is zeros, the work happens on reset.""" + manager = MockTerminationManager() + env = SimpleNamespace(termination_manager=manager, num_envs=NUM_ENVS, device=DEVICE, step_dt=0.02) + cfg = RewardTermCfg(func=warp_loco_rew.survival_success_rate, weight=0.0, params={}) + + def mutate() -> None: + manager.time_outs[:] = True + wp.synchronize() + + return CaptureCase( + warp_fn=warp_loco_rew.survival_success_rate(cfg, env), + stable_fn=stable_loco_rew.survival_success_rate( + RewardTermCfg(func=stable_loco_rew.survival_success_rate, weight=0.0, params={}), env + ), + warp_env=env, + stable_env=env, + params={}, + mutate=mutate, + ) + + CAPTURE_SPECS: list[CaptureSpec] = [ CaptureSpec("is_terminated_term", "reward", _build_is_terminated_term), CaptureSpec("pose_command_success", "termination", _build_pose_command_success), + CaptureSpec("terminated_penalty", "reward", _build_terminated_penalty), + CaptureSpec("survival_success_rate", "reward", _build_survival_success_rate, expect_nonzero=False), ] # Warp MDP terms not yet exercised by this harness. Pre-existing terms only: a *new* term must @@ -221,13 +278,52 @@ def mutate() -> None: "reset_joints_by_offset": "reset-cadence event term; variable env_ids launch dim", "reset_joints_by_scale": "reset-cadence event term; variable env_ids launch dim", "reset_root_state_uniform": "reset-cadence event term; variable env_ids launch dim", + # Per-task warp mirrors under isaaclab_tasks_experimental. That package has no test + # directory, so none of these has parity or capture coverage anywhere in the repo. + **{ + name: "per-task warp mirror; no test coverage in the repo yet" + for name in ( + "base_angle_to_target", + "base_heading_proj", + "base_up_proj", + "base_yaw_roll", + "feet_air_time", + "feet_air_time_positive_biped", + "feet_slide", + "joint_pos_limits_penalty_ratio", + "joint_pos_target_l2", + "move_to_target_bonus", + "orientation_command_error", + "position_command_error", + "position_command_error_tanh", + "power_consumption", + "progress_reward", + "stand_still_joint_deviation_l1", + "terrain_out_of_bounds", + "track_ang_vel_z_world_exp", + "track_lin_vel_xy_yaw_frame_exp", + "upright_posture_bonus", + ) + }, } +def _warp_mdp_modules() -> list: + """The shared warp MDP modules plus every per-task warp MDP mirror.""" + modules = [warp_rew, warp_term, warp_obs, warp_events] + for entry in pkgutil.iter_modules(tasks_experimental_core.__path__): + for leaf in ("rewards", "terminations", "observations", "events"): + try: + modules.append(importlib.import_module(f"{tasks_experimental_core.__name__}.{entry.name}.mdp.{leaf}")) + except ImportError: + continue # not every task mirror defines every modality + return modules + + def _discover_warp_mdp_terms() -> dict[str, str]: """Return every public warp MDP term, mapped to its defining module name.""" terms: dict[str, str] = {} - for module in (warp_rew, warp_term, warp_obs, warp_events): + for module in _warp_mdp_modules(): for name, obj in vars(module).items(): if name.startswith("_") or not (inspect.isfunction(obj) or inspect.isclass(obj)): continue @@ -282,4 +378,5 @@ def test_term_is_capture_safe(spec: CaptureSpec): assert_close(actual, expected) else: assert_equal(actual, expected) - assert expected.any(), "mutated inputs produced a degenerate expectation; the replay proves little" + if spec.expect_nonzero: + assert expected.any(), "mutated inputs produced a degenerate expectation; the replay proves little" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.pyi index f476fe7b000a..a3db8c2d719b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.pyi @@ -13,6 +13,8 @@ __all__ = [ "progress_reward", "joint_pos_limits_penalty_ratio", "power_consumption", + "terminated_penalty", + "survival_success_rate", ] # Forward stable MDP terms and experimental Warp-first overrides lazily, then @@ -25,5 +27,7 @@ from .rewards import ( move_to_target_bonus, power_consumption, progress_reward, + survival_success_rate, + terminated_penalty, upright_posture_bonus, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py index 902b88c6094f..8a78413f4ca5 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py @@ -325,3 +325,104 @@ def __call__( inputs=[env.action_manager.action, asset.data.joint_vel.warp, self._gear_ratio_scaled_wp, out], device=env.device, ) + + +# --------------------------------------------------------------------------- +# terminated_penalty +# --------------------------------------------------------------------------- + + +@wp.kernel +def _terminated_penalty_kernel( + terminated: wp.array(dtype=wp.bool), + inv_dt: float, + out: wp.array(dtype=wp.float32), +): + """Emit ``1 / step_dt`` on the terminating step so the manager's dt scaling cancels.""" + i = wp.tid() + out[i] = wp.where(terminated[i], inv_dt, 0.0) + + +def terminated_penalty(env: ManagerBasedRLEnv, out: wp.array(dtype=wp.float32)) -> None: + """One-off penalty for terminating early, independent of the environment step size. + + Warp-first override of :func:`isaaclab_tasks.core.locomotion.mdp.rewards.terminated_penalty`. + """ + wp.launch( + kernel=_terminated_penalty_kernel, + dim=env.num_envs, + inputs=[env.termination_manager.terminated_wp, 1.0 / env.step_dt, out], + device=env.device, + ) + + +# --------------------------------------------------------------------------- +# survival_success_rate +# --------------------------------------------------------------------------- + + +@wp.kernel +def _survival_counts_kernel( + env_mask: wp.array(dtype=wp.bool), + time_outs: wp.array(dtype=wp.bool), + counts: wp.array(dtype=wp.int32), +): + """Atomically count masked (just-reset) envs and how many of them timed out.""" + env_index = wp.tid() + if env_mask[env_index]: + wp.atomic_add(counts, 0, 1) + if time_outs[env_index]: + wp.atomic_add(counts, 1, 1) + + +@wp.kernel +def _survival_rate_kernel( + counts: wp.array(dtype=wp.int32), + success_rate: wp.array(dtype=wp.float32), +): + """Compute the survival success rate as the fraction of just-reset envs that timed out.""" + if counts[0] > 0: + success_rate[0] = wp.float32(counts[1]) / wp.float32(counts[0]) + else: + success_rate[0] = 0.0 + + +class survival_success_rate(ManagerTermBase): + """Tracks episode survival as the success metric (Warp-first). + + Twin of :class:`isaaclab_tasks.core.locomotion.mdp.rewards.survival_success_rate`. + Returns zero reward (pure metric tracking). On reset, computes the fraction of + just-reset environments that timed out (survived the full episode) entirely on-device + and exposes it as ``Metrics/success_rate`` through the reward manager's reset extras. + Unlike the stable term, there is no host readback, so the computation stays + CUDA-graph capturable. + """ + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + # [0] number of just-reset envs, [1] number of those that timed out + self._counts_wp = wp.zeros((2,), dtype=wp.int32, device=env.device) + self._success_rate_wp = wp.zeros((1,), dtype=wp.float32, device=env.device) + # persistent 0-dim tensor view: kernels refresh the value on every reset/replay + self._reset_extras = {"Metrics/success_rate": wp.to_torch(self._success_rate_wp)[0]} + + def reset(self, env_mask: wp.array | None = None) -> dict[str, torch.Tensor]: + if env_mask is None: + env_mask = self._env.resolve_env_mask() + self._counts_wp.zero_() + wp.launch( + kernel=_survival_counts_kernel, + dim=self.num_envs, + inputs=[env_mask, self._env.termination_manager.time_outs_wp, self._counts_wp], + device=self.device, + ) + wp.launch( + kernel=_survival_rate_kernel, + dim=1, + inputs=[self._counts_wp, self._success_rate_wp], + device=self.device, + ) + return self._reset_extras + + def __call__(self, env: ManagerBasedRLEnv, out: wp.array(dtype=wp.float32)) -> None: + out.zero_() From a7c34234525e5569fe717614c572e2d7f3eb91d5 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 11 Aug 2026 11:48:40 -0700 Subject: [PATCH 09/10] Compute warp orientation error with atan2 rather than acos 2*acos(|w|) reads the scalar part as a cosine, which holds only for a unit quaternion and is ill-conditioned as the error approaches zero: w = cos(theta/2) carries theta only at second order, so in float32 the value rounds to exactly 1 and the reported angle collapses to 0 below roughly 1e-4 rad. Measured against the analytic angle, 2*acos is 2.3% low at 1e-3 rad and returns zero at 1e-4. 2*atan2(|xyz|, |w|) recovers the angle from the ratio, where the numerator carries theta at first order. It is exact across the same range and independent of the quaternion's norm. This matters most for orientation_command_error, a continuous reward whose error shrinks toward zero as a policy converges -- precisely where acos stops resolving. pose_command_success gates on a threshold three orders of magnitude above the affected range, but computes the same quantity, so both use one form. survival_success_rate no longer substitutes 0.0 when a reset selects no environments: the stable term's mean over an empty selection is NaN, and publishing a real 0% sample misreports it as a measurement. --- .../jichuanh-warp-reach-mdp-twins.minor.rst | 9 +++++++++ .../isaaclab_experimental/envs/mdp/terminations.py | 6 +++++- .../jichuanh-warp-reach-mdp-twins.minor.rst | 13 +++++++++++++ .../core/locomotion/mdp/rewards.py | 12 +++++++----- .../core/reach/mdp/rewards.py | 9 +++++---- 5 files changed, 39 insertions(+), 10 deletions(-) create mode 100644 source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst index 1e962ceb28b0..dd89c3ddb09c 100644 --- a/source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst @@ -9,3 +9,12 @@ Added so several configurations can be surveyed in one pass. * Added :attr:`~isaaclab_experimental.managers.TerminationManager.term_dones_wp`, exposing the per-term done buffer to reward terms that aggregate a subset of termination terms. + +Fixed +^^^^^ + +* Fixed the Warp :func:`~isaaclab_experimental.envs.mdp.pose_command_success` twin disagreeing + with the stable term's orientation error for non-unit quaternions, by computing the angle as + ``2*atan2(|xyz|, |w|)`` instead of ``2*acos(|w|)``. Both arguments scale with the quaternion, + so the angle is now norm-invariant and environments near the configured + ``orientation_success_threshold`` are classified the same way as on the stable path. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py index 1a80d78c4942..84928d12656c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -84,7 +84,11 @@ def _pose_command_success_kernel( des_q_b = wp.quatf(cmd[i, 3], cmd[i, 4], cmd[i, 5], cmd[i, 6]) des_q_w = root_quat_w[i] * des_q_b q_err = wp.quat_inverse(body_quat_w[i, body_idx]) * des_q_w - angle = 2.0 * wp.acos(wp.clamp(wp.abs(q_err[3]), 0.0, 1.0)) + # 2*atan2(|xyz|, |w|), matching the stable axis-angle magnitude. Both arguments scale + # with the quaternion, so this is norm-invariant, where 2*acos(|w|) would understate + # the error for a non-unit input. |w| takes the shortest path, as stable's sign flip does. + axis = wp.vec3f(q_err[0], q_err[1], q_err[2]) + angle = 2.0 * wp.atan2(wp.length(axis), wp.abs(q_err[3])) success = success and (angle < orientation_threshold) out[i] = success # sticky per-episode tracker, mirroring the stable term's ``self._succeeded |= success``. diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst new file mode 100644 index 000000000000..918335e23546 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst @@ -0,0 +1,13 @@ +Added +^^^^^ + +* Added Warp MDP twins for :func:`~isaaclab_tasks.core.locomotion.mdp.terminated_penalty` and + :class:`~isaaclab_tasks.core.locomotion.mdp.survival_success_rate`, so ``Isaac-Ant`` and + ``Isaac-Humanoid`` run under ``--frontend warp``. + +Fixed +^^^^^ + +* Fixed the Warp orientation-error twins reporting a smaller rotation error than the stable + terms for non-unit quaternions, by computing the angle as ``2*atan2(|xyz|, |w|)`` instead of + ``2*acos(|w|)``. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py index 8a78413f4ca5..e180774bb9b7 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py @@ -380,11 +380,13 @@ def _survival_rate_kernel( counts: wp.array(dtype=wp.int32), success_rate: wp.array(dtype=wp.float32), ): - """Compute the survival success rate as the fraction of just-reset envs that timed out.""" - if counts[0] > 0: - success_rate[0] = wp.float32(counts[1]) / wp.float32(counts[0]) - else: - success_rate[0] = 0.0 + """Compute the survival success rate as the fraction of just-reset envs that timed out. + + Left undivided when no env reset: 0/0 is NaN, which is what the stable term's mean over an + empty selection yields. Substituting 0.0 would log a real 0% sample for a reset that + selected nothing. + """ + success_rate[0] = wp.float32(counts[1]) / wp.float32(counts[0]) class survival_success_rate(ManagerTermBase): diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py index cbd3759e4603..cdab28c55625 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py @@ -147,10 +147,11 @@ def _orientation_command_error_kernel( cur_w = body_quat_w[i, body_idx] # shortest-path error: angle of q_err = cur^-1 * des q_err = wp.quat_inverse(cur_w) * des_w - # error magnitude = 2 * acos(|w|) (w component of the error quaternion) - qw = wp.abs(q_err[3]) - qw = wp.clamp(qw, 0.0, 1.0) - out[i] = 2.0 * wp.acos(qw) + # 2*atan2(|xyz|, |w|), matching the stable axis-angle magnitude. Both arguments scale with + # the quaternion, so this is norm-invariant, where 2*acos(|w|) would understate the error + # for a non-unit input. |w| takes the shortest path, as stable's sign flip does. + axis = wp.vec3f(q_err[0], q_err[1], q_err[2]) + out[i] = 2.0 * wp.atan2(wp.length(axis), wp.abs(q_err[3])) def orientation_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: From 5e0afc549f1a90a948fbbe32da61f40e8a604df3 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Tue, 11 Aug 2026 11:48:42 -0700 Subject: [PATCH 10/10] Close false-green paths in the warp MDP gates Both gates could report success while the property they guard was violated. Terms were keyed by bare name, but survival_success_rate is defined by both the cartpole and locomotion mirrors, so a spec for one marked the other declared. Identity is now module-qualified throughout discovery, specs and the backlog. Discovery skipped action terms even though ActionManager is graph-captured like every other manager, leaving three of them unable to fail the declaration check. The event-term rows claimed a variable env_ids launch dim made them uncapturable. That is not what they do -- every one launches over env.num_envs with a boolean mask, which is the capture-safe form -- and four already have capture coverage in the events parity file. The rows now say which are covered and which are simply untested. The runner captured only each term's output call, so a stateful write moved to host code would replay as a no-op unnoticed. Terms may now assert their side effects after replay, and that assertion runs before the stable reference is called: the stable terms write the same shared state they are compared against, so calling the reference first repaired exactly what the assertion checks. The registry sweep matched only string cfg entry points and silently omitted any other supported form; it now loads through load_cfg_from_registry. --- .../test/envs/mdp/test_capture_safety.py | 276 +++++++++++++----- .../test/envs/test_frontend_cfg_conversion.py | 11 +- 2 files changed, 208 insertions(+), 79 deletions(-) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py index ad8c5888e008..1bd0d1cfae35 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -40,8 +40,6 @@ wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") -import isaaclab_experimental.envs.mdp.events as warp_events -import isaaclab_experimental.envs.mdp.observations as warp_obs import isaaclab_experimental.envs.mdp.rewards as warp_rew import isaaclab_experimental.envs.mdp.terminations as warp_term import isaaclab_tasks_experimental.core as tasks_experimental_core @@ -79,6 +77,9 @@ class CaptureCase: stable_env: Environment passed to the stable term; may be the same object. params: Term parameters forwarded to both sides. mutate: Overwrites every input buffer in place, so replay sees new values. + after_replay: Asserts the term's side effects landed during replay, given the replayed + output. A term whose stateful write moved to host code would pass the output + comparison and fail here. """ warp_fn: Callable @@ -87,6 +88,7 @@ class CaptureCase: stable_env: object params: dict mutate: Callable[[], None] + after_replay: Callable[[torch.Tensor], None] | None = None @dataclasses.dataclass(frozen=True) @@ -99,6 +101,7 @@ class CaptureSpec: build: Builds a fresh :class:`CaptureCase`; called once per test. """ + module: str name: str modality: str build: Callable[[], CaptureCase] @@ -109,6 +112,11 @@ class CaptureSpec: output is a constant by design, such as a pure-metric term that writes zeros. """ + @property + def qualified(self) -> str: + """The ``":"`` identity this spec declares.""" + return f"{self.module}:{self.name}" + # --------------------------------------------------------------------------- # Case builders @@ -166,9 +174,16 @@ def mutate() -> None: mutate_body_data(art_data) fresh = make_pose_command_term(scene["robot"], seed=505) command.pose_command_b[:] = fresh.pose_command_b + # cleared so the replay has to rewrite it: the sticky write lives in the kernel, and a + # host-side equivalent would simply not run on replay + command._succeeded.zero_() wp.synchronize() + def after_replay(out_torch) -> None: + assert_equal(command._succeeded, out_torch) + return CaptureCase( + after_replay=after_replay, warp_fn=warp_term.pose_command_success( TerminationTermCfg(func=warp_term.pose_command_success, params=params), env ), @@ -223,96 +238,151 @@ def mutate() -> None: ) +_SHARED_REW = "isaaclab_experimental.envs.mdp.rewards" +_SHARED_TERM = "isaaclab_experimental.envs.mdp.terminations" +_LOCO_REW = "isaaclab_tasks_experimental.core.locomotion.mdp.rewards" + CAPTURE_SPECS: list[CaptureSpec] = [ - CaptureSpec("is_terminated_term", "reward", _build_is_terminated_term), - CaptureSpec("pose_command_success", "termination", _build_pose_command_success), - CaptureSpec("terminated_penalty", "reward", _build_terminated_penalty), - CaptureSpec("survival_success_rate", "reward", _build_survival_success_rate, expect_nonzero=False), + CaptureSpec(_SHARED_REW, "is_terminated_term", "reward", _build_is_terminated_term), + CaptureSpec(_SHARED_TERM, "pose_command_success", "termination", _build_pose_command_success), + CaptureSpec(_LOCO_REW, "terminated_penalty", "reward", _build_terminated_penalty), + CaptureSpec(_LOCO_REW, "survival_success_rate", "reward", _build_survival_success_rate, expect_nonzero=False), ] + # Warp MDP terms not yet exercised by this harness. Pre-existing terms only: a *new* term must # arrive with a CAPTURE_SPECS entry instead of a row here. Several are already covered by the # hand-written ``TestCapturedDataMutation*`` classes in the parity test files; the rest are an # explicit backlog rather than a silent gap. +def _ids(module: str, *names: str) -> list[str]: + """Qualified ``":"`` identities for several terms in one module.""" + return [f"{module}:{name}" for name in names] + + +_SH = "isaaclab_experimental.envs.mdp" +_TE = "isaaclab_tasks_experimental.core" +_MUT = "covered by the mutate-replay class in the matching parity test file" +_NONE = "capturable (fixed dim=num_envs launch over a mask) but no capture coverage yet" +_MIRROR = "per-task warp mirror; isaaclab_tasks_experimental has no test directory" + +# Warp MDP terms not exercised by this harness. Pre-existing terms only: a *new* term must +# arrive with a CAPTURE_SPECS entry instead of a row here. CAPTURE_UNAUDITED: dict[str, str] = { - "action_l2": "covered by TestCapturedDataMutationRewards", - "action_rate_l2": "covered by TestCapturedDataMutationRewards", - "ang_vel_xy_l2": "covered by TestCapturedDataMutationRewards", - "flat_orientation_l2": "covered by TestCapturedDataMutationRewards", - "is_alive": "covered by TestCapturedDataMutationRewards", - "is_terminated": "covered by TestCapturedDataMutationRewards", - "joint_acc_l2": "covered by TestCapturedDataMutationRewards", - "joint_deviation_l1": "covered by TestCapturedDataMutationRewards", - "joint_pos_limits": "covered by TestCapturedDataMutationRewards", - "joint_torques_l2": "covered by TestCapturedDataMutationRewards", - "joint_vel_l1": "covered by TestCapturedDataMutationRewards", - "joint_vel_l2": "covered by TestCapturedDataMutationRewards", - "lin_vel_z_l2": "covered by TestCapturedDataMutationRewards", - "track_ang_vel_z_exp": "covered by TestCapturedDataMutationRewardsNewTerms", - "track_lin_vel_xy_exp": "covered by TestCapturedDataMutationRewardsNewTerms", - "undesired_contacts": "covered by TestCapturedDataMutationRewardsNewTerms", - "illegal_contact": "covered by TestCapturedDataMutationTerminationsNewTerms", - "joint_pos_out_of_manual_limit": "covered by TestCapturedDataMutationTerminations", - "root_height_below_minimum": "covered by TestCapturedDataMutationTerminations", - "time_out": "covered by TestCapturedDataMutationTerminationsNewTerms", - "base_ang_vel": "covered by TestCapturedDataMutationObservations", - "base_lin_vel": "covered by TestCapturedDataMutationObservations", - "base_pos_z": "covered by TestCapturedDataMutationObservations", - "generated_commands": "covered by TestCapturedDataMutationObservations", - "joint_pos": "covered by TestCapturedDataMutationObservations", - "joint_pos_limit_normalized": "covered by TestCapturedDataMutationObservations", - "joint_vel": "covered by TestCapturedDataMutationObservations", - "last_action": "covered by TestCapturedDataMutationObservations", - "projected_gravity": "covered by TestCapturedDataMutationObservations", - # no capture coverage anywhere — these three are the real backlog this gate exposes - "body_incoming_wrench": "no capture coverage yet", - "joint_pos_rel": "no capture coverage yet", - "joint_vel_rel": "no capture coverage yet", - # Event terms run at reset cadence over an env_ids subset, so their launch dim varies per - # call and cannot be recorded into a graph. Capture safety does not apply to them as written. - "apply_external_force_torque": "reset-cadence event term; variable env_ids launch dim", - "push_by_setting_velocity": "reset-cadence event term; variable env_ids launch dim", - "randomize_rigid_body_com": "reset-cadence event term; variable env_ids launch dim", - "randomize_rigid_body_mass": "reset-cadence event term; variable env_ids launch dim", - "randomize_rigid_body_material": "reset-cadence event term; variable env_ids launch dim", - "reset_joints_by_offset": "reset-cadence event term; variable env_ids launch dim", - "reset_joints_by_scale": "reset-cadence event term; variable env_ids launch dim", - "reset_root_state_uniform": "reset-cadence event term; variable env_ids launch dim", - # Per-task warp mirrors under isaaclab_tasks_experimental. That package has no test - # directory, so none of these has parity or capture coverage anywhere in the repo. - **{ - name: "per-task warp mirror; no test coverage in the repo yet" - for name in ( + **dict.fromkeys( + _ids( + f"{_SH}.rewards", + "action_l2", + "action_rate_l2", + "ang_vel_xy_l2", + "flat_orientation_l2", + "is_alive", + "is_terminated", + "joint_acc_l2", + "joint_deviation_l1", + "joint_pos_limits", + "joint_torques_l2", + "joint_vel_l1", + "joint_vel_l2", + "lin_vel_z_l2", + "track_ang_vel_z_exp", + "track_lin_vel_xy_exp", + "undesired_contacts", + ) + + _ids( + f"{_SH}.terminations", + "illegal_contact", + "joint_pos_out_of_manual_limit", + "root_height_below_minimum", + "time_out", + ) + + _ids( + f"{_SH}.observations", + "base_ang_vel", + "base_lin_vel", + "base_pos_z", + "generated_commands", + "joint_pos", + "joint_pos_limit_normalized", + "joint_vel", + "last_action", + "projected_gravity", + ) + + _ids( + f"{_SH}.events", + "apply_external_force_torque", + "push_by_setting_velocity", + "reset_joints_by_offset", + "reset_joints_by_scale", + ), + _MUT, + ), + **dict.fromkeys( + _ids(f"{_SH}.observations", "body_incoming_wrench", "joint_pos_rel", "joint_vel_rel") + + _ids( + f"{_SH}.events", + "randomize_rigid_body_com", + "randomize_rigid_body_mass", + "randomize_rigid_body_material", + "reset_root_state_uniform", + ) + + _ids( + f"{_SH}.actions.joint_actions", + "JointAction", + "JointPositionAction", + "JointEffortAction", + ), + _NONE, + ), + **dict.fromkeys( + _ids(f"{_TE}.cartpole.mdp.rewards", "joint_pos_target_l2", "survival_success_rate") + + _ids( + f"{_TE}.locomotion.mdp.observations", "base_angle_to_target", "base_heading_proj", "base_up_proj", "base_yaw_roll", - "feet_air_time", - "feet_air_time_positive_biped", - "feet_slide", + ) + + _ids( + f"{_TE}.locomotion.mdp.rewards", "joint_pos_limits_penalty_ratio", - "joint_pos_target_l2", "move_to_target_bonus", + "power_consumption", + "progress_reward", + "upright_posture_bonus", + ) + + _ids( + f"{_TE}.reach.mdp.rewards", "orientation_command_error", "position_command_error", "position_command_error_tanh", - "power_consumption", - "progress_reward", + ) + + _ids( + f"{_TE}.velocity.mdp.rewards", + "feet_air_time", + "feet_air_time_positive_biped", + "feet_slide", "stand_still_joint_deviation_l1", - "terrain_out_of_bounds", "track_ang_vel_z_world_exp", "track_lin_vel_xy_yaw_frame_exp", - "upright_posture_bonus", ) - }, + + _ids(f"{_TE}.velocity.mdp.terminations", "terrain_out_of_bounds"), + _MIRROR, + ), } +# Every modality a manager can run. ``actions`` is included because ``ActionManager`` is +# graph-captured like the rest, so an unsafe action term would otherwise bypass every check. +_MDP_LEAF_MODULES = ("rewards", "terminations", "observations", "events", "actions.joint_actions") + + def _warp_mdp_modules() -> list: """The shared warp MDP modules plus every per-task warp MDP mirror.""" - modules = [warp_rew, warp_term, warp_obs, warp_events] + modules = [] + for leaf in _MDP_LEAF_MODULES: + modules.append(importlib.import_module(f"isaaclab_experimental.envs.mdp.{leaf}")) for entry in pkgutil.iter_modules(tasks_experimental_core.__path__): - for leaf in ("rewards", "terminations", "observations", "events"): + for leaf in _MDP_LEAF_MODULES: try: modules.append(importlib.import_module(f"{tasks_experimental_core.__name__}.{entry.name}.mdp.{leaf}")) except ImportError: @@ -320,22 +390,27 @@ def _warp_mdp_modules() -> list: return modules -def _discover_warp_mdp_terms() -> dict[str, str]: - """Return every public warp MDP term, mapped to its defining module name.""" - terms: dict[str, str] = {} +def _discover_warp_mdp_terms() -> set[str]: + """Return every public warp MDP term as a ``":"`` identity. + + Qualified rather than bare: the same term name legitimately appears in more than one task + mirror (``survival_success_rate`` is defined by both cartpole and locomotion), and keying + by name alone would let a spec for one of them mark the other as declared. + """ + terms: set[str] = set() for module in _warp_mdp_modules(): for name, obj in vars(module).items(): if name.startswith("_") or not (inspect.isfunction(obj) or inspect.isclass(obj)): continue if getattr(obj, "__module__", "") == module.__name__: - terms[name] = module.__name__ + terms.add(f"{module.__name__}:{name}") return terms def test_every_warp_mdp_term_is_declared(): """Every warp MDP term is either exercised here or listed as an unaudited pre-existing term.""" - declared = {spec.name for spec in CAPTURE_SPECS} | set(CAPTURE_UNAUDITED) - undeclared = sorted(set(_discover_warp_mdp_terms()) - declared) + declared = {spec.qualified for spec in CAPTURE_SPECS} | set(CAPTURE_UNAUDITED) + undeclared = sorted(_discover_warp_mdp_terms() - declared) assert not undeclared, ( "warp MDP terms with no capture declaration: " @@ -346,14 +421,14 @@ def test_every_warp_mdp_term_is_declared(): def test_no_term_is_both_specified_and_unaudited(): """A term moved into the harness must lose its unaudited row, so the backlog stays truthful.""" - overlap = sorted({spec.name for spec in CAPTURE_SPECS} & set(CAPTURE_UNAUDITED)) + overlap = sorted({spec.qualified for spec in CAPTURE_SPECS} & set(CAPTURE_UNAUDITED)) assert not overlap, f"remove from CAPTURE_UNAUDITED, now exercised here: {overlap}" def test_unaudited_terms_still_exist(): """Drop rows for terms that no longer exist, so the list cannot rot.""" - stale = sorted(set(CAPTURE_UNAUDITED) - set(_discover_warp_mdp_terms())) + stale = sorted(set(CAPTURE_UNAUDITED) - _discover_warp_mdp_terms()) assert not stale, f"CAPTURE_UNAUDITED lists terms that no longer exist: {stale}" @@ -371,12 +446,65 @@ def test_term_is_capture_safe(spec: CaptureSpec): case.mutate() wp.capture_launch(capture.graph) + actual = wp.to_torch(out).clone() + + # Side effects are checked before the stable reference runs. The stable terms write the same + # shared state they are being compared against — ``compute_success`` ORs into the command's + # sticky buffer — so calling the reference first would repair exactly what this asserts. + if case.after_replay is not None: + case.after_replay(actual) expected = case.stable_fn(case.stable_env, **case.params) - actual = wp.to_torch(out).clone() if spec.modality == "reward": assert_close(actual, expected) else: assert_equal(actual, expected) if spec.expect_nonzero: assert expected.any(), "mutated inputs produced a degenerate expectation; the replay proves little" + + +def test_survival_success_rate_reset_is_capture_safe(): + """The metric is produced by ``reset()``, not ``__call__``, so capture it separately. + + ``__call__`` writes constant zeros; capturing only that would leave the counting kernels + and the persistent metric view unexercised, and a host-side ``.item()`` readback — which is + what the stable term does and what cannot survive capture — would go unnoticed. + """ + manager = MockTerminationManager() + env = SimpleNamespace(termination_manager=manager, num_envs=NUM_ENVS, device=DEVICE, step_dt=0.02) + term = warp_loco_rew.survival_success_rate( + RewardTermCfg(func=warp_loco_rew.survival_success_rate, weight=0.0, params={}), env + ) + env_mask = wp.array(np.ones(NUM_ENVS, dtype=bool), dtype=wp.bool, device=DEVICE) + metric = term.reset(env_mask)["Metrics/success_rate"] + + manager.time_outs[:] = False + manager.time_outs[: NUM_ENVS // 4] = True + term.reset(env_mask) # warm-up outside the capture + with wp.ScopedCapture() as capture: + term.reset(env_mask) + + # a different survival fraction must show up through the same persistent view on replay + manager.time_outs[:] = False + manager.time_outs[: NUM_ENVS // 2] = True + wp.synchronize() + wp.capture_launch(capture.graph) + wp.synchronize() + + assert metric.item() == pytest.approx(0.5), "replayed reset did not recompute the metric on device" + + +def test_survival_success_rate_reports_nan_for_an_empty_reset(): + """An empty selection must not publish a real 0% sample; stable's mean of nothing is NaN.""" + manager = MockTerminationManager() + env = SimpleNamespace(termination_manager=manager, num_envs=NUM_ENVS, device=DEVICE, step_dt=0.02) + term = warp_loco_rew.survival_success_rate( + RewardTermCfg(func=warp_loco_rew.survival_success_rate, weight=0.0, params={}), env + ) + empty = wp.array(np.zeros(NUM_ENVS, dtype=bool), dtype=wp.bool, device=DEVICE) + + metric = term.reset(empty)["Metrics/success_rate"] + wp.synchronize() + + expected = torch.tensor([], device=DEVICE).mean() + assert torch.isnan(metric) == torch.isnan(expected), "empty reset must match the stable empty-mean result" diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 8a22e3ff9ab1..c6d345c54934 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -34,6 +34,7 @@ # Registering the task packages is the whole point — import for side effects. import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils.hydra import resolve_presets +from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry # Stable manager-based tasks resolve to this env class; direct tasks provide their own and # take the :meth:`WarpFrontend._resolve_direct_warp_class` path instead of cfg adaptation. @@ -110,14 +111,14 @@ def _sweep_warp_support() -> tuple[frozenset[str], dict[str, str], dict[str, str for task_id, spec in gym.registry.items(): if spec.entry_point != _STABLE_MANAGER_ENTRY_POINT: continue - cfg_entry = (spec.kwargs or {}).get("env_cfg_entry_point") - if not isinstance(cfg_entry, str): + if (spec.kwargs or {}).get("env_cfg_entry_point") is None: continue try: - module_path, class_name = cfg_entry.split(":") - cfg = getattr(importlib.import_module(module_path), class_name)() + # the canonical loader, so every registry form the runtime accepts is surveyed; + # matching only ``str`` entry points here would skip callable ones silently + cfg = load_cfg_from_registry(task_id, "env_cfg_entry_point") cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) - except Exception as exc: # noqa: BLE001 - any cfg import failure means "cannot judge" + except Exception as exc: # noqa: BLE001 - any cfg load failure means "cannot judge" unimportable[task_id] = f"{type(exc).__name__}: {exc}" continue reason = WarpFrontend.check_compatibility(cfg)