diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 1b2db35ebfe6..2b5cb7d41226 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -552,6 +552,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] 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..dd89c3ddb09c --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-reach-mdp-twins.minor.rst @@ -0,0 +1,20 @@ +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. + +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/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/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..84928d12656c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -19,11 +19,12 @@ 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 """ @@ -49,6 +50,114 @@ 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, + 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. + + 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 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] + 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 + # 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``. + # 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): + """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 views of the command + and success buffers are all resolved at init, leaving :meth:`__call__` a single kernel launch. + + Note: + 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) + # 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._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._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: + 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._check_position, + self._position_threshold, + self._check_orientation, + self._orientation_threshold, + self._succeeded_wp, + 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_capture_safety.py b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py new file mode 100644 index 000000000000..1bd0d1cfae35 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -0,0 +1,510 @@ +# 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 importlib +import inspect +import pkgutil +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.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, + 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 + +import isaaclab_tasks.core.locomotion.mdp.rewards as stable_loco_rew + + +@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. + 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 + stable_fn: Callable + warp_env: object + stable_env: object + params: dict + mutate: Callable[[], None] + after_replay: Callable[[torch.Tensor], None] | None = 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. + """ + + module: str + 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. + """ + + @property + def qualified(self) -> str: + """The ``":"`` identity this spec declares.""" + return f"{self.module}:{self.name}" + + +# --------------------------------------------------------------------------- +# 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 + # 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 + ), + stable_fn=stable_term.pose_command_success, + warp_env=env, + stable_env=env, + params=params, + mutate=mutate, + ) + + +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, + ) + + +_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(_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] = { + **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", + ) + + _ids( + f"{_TE}.locomotion.mdp.rewards", + "joint_pos_limits_penalty_ratio", + "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", + ) + + _ids( + f"{_TE}.velocity.mdp.rewards", + "feet_air_time", + "feet_air_time_positive_biped", + "feet_slide", + "stand_still_joint_deviation_l1", + "track_ang_vel_z_world_exp", + "track_lin_vel_xy_yaw_frame_exp", + ) + + _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 = [] + 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 _MDP_LEAF_MODULES: + 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() -> 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.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.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: " + + ", ".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.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) - _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) + 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) + 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/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..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 @@ -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, @@ -45,6 +49,7 @@ ) import isaaclab.envs.mdp.terminations as stable_term +from isaaclab.managers.manager_term_cfg import TerminationTermCfg # ============================================================================ # Fixtures @@ -268,6 +273,80 @@ 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" + + 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) # ============================================================================ 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..c6d345c54934 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 @@ -33,22 +34,33 @@ # 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 - -# 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", -] +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. +_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 +94,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 + if (spec.kwargs or {}).get("env_cfg_entry_point") is None: + continue + try: + # 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 load 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 +181,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)) 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/__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..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 @@ -325,3 +325,106 @@ 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. + + 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): + """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_() 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: