[Warp] Restore --frontend warp on Reach and gate the supported task set in CI

.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]

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.

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`.

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.
 """

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.
 """

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.
     """