diff --git a/AGENTS.md b/AGENTS.md index 427e0bc9b0b5..854ac7b8b290 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -322,3 +322,38 @@ To debug Warp kernel behavior: 1. **Write a standalone reproduction script** and run it directly with `./isaaclab.sh -p -c "..."` or `./isaaclab.sh -p script.py`. This keeps stdout visible and avoids the test framework entirely. 2. **Use high-precision format strings** for floating-point debugging (e.g., `wp.printf("val=%.15e\n", x)`) — the default `%f` format hides values smaller than ~1e-6 that can still affect control flow. 3. **Remove all `wp.printf` calls before committing.** + +### Warp mask-first boundaries + +Boolean environment masks are the canonical selection in Warp-first code; host-side +mask-to-ID compaction (`nonzero()` on device data) is a GPU-to-host synchronization +with a data-dependent shape, which both stalls the pipeline and permanently +disqualifies the containing code path from CUDA graph capture. + +1. **Syncs only inside named, sanctioned boundary helpers.** Host syncs + (`nonzero`, `argwhere`, `item`, `cpu`, `tolist`, `numpy`) on the mask-native + path live in small single-purpose functions listed in + `SANCTIONED_BOUNDARIES` of + `source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py` — no + inline marker comments. The gate scans the experimental production tree in + full and `*mask*`-named functions in core/Newton; the list is exact, so + adding or removing a boundary is a conscious, reviewable test change. + Legitimate boundaries are host consumers: recorders, legacy compatibility + fallbacks, checkpoint/restore, and empty-reset predicates. +2. **Runtime backstop: `ISAACLAB_SYNC_DEBUG=1`.** Warp stages then execute under + `torch.cuda.set_sync_debug_mode("error")` (see `WarpGraphCache`), so a hidden + sync in any term — including future ones — raises at the exact call site. + Zero cost when the variable is unset. +3. **Capturability is annotated with `@WarpCapturable(False, reason=...)`** (the + single decorator form; works on term functions and class terms). Every + opt-out is pinned by the inventory test next to the gate. Unannotated terms + default to capturable and are covered by the runtime backstop. +4. **Mask-native Torch means elementwise masked ops.** Use `masked_fill_` or + `torch.where`; boolean advanced indexing (`buf[:, mask] = 0`) gathers indices + with a data-dependent allocation and is not capture-safe. +5. **Dual reset APIs.** Keep the ID-based `reset(env_ids)` signature untouched and + add a sibling `reset_mask(env_mask)`; the base-class fallback compacts (a + sanctioned boundary), mask-native implementations override it. Backends whose + assets accept but ignore `env_mask` (classic PhysX, a develop-era condition) + are a documented hazard on `InteractiveScene.reset`, not guarded machinery — + the warp env only drives Newton scenes. diff --git a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst index 05c5bc848422..5ada16e04bce 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/warp-environments.rst @@ -31,38 +31,41 @@ Available Environments Direct Warp Environments ^^^^^^^^^^^^^^^^^^^^^^^^ -- ``Isaac-Cartpole-Direct-Warp-v0`` — Cartpole balance -- ``Isaac-Ant-Direct-Warp-v0`` — Ant locomotion -- ``Isaac-Humanoid-Direct-Warp-v0`` — Humanoid locomotion -- ``Isaac-Reorient-Cube-Allegro-Direct-Warp-v0`` — Allegro hand cube reorient - - -Manager-Based Warp Environments -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -**Classic** - -- ``Isaac-Cartpole-Warp-v0`` -- ``Isaac-Ant-Warp-v0`` -- ``Isaac-Humanoid-Warp-v0`` - -**Locomotion (Flat)** - -- ``Isaac-Velocity-Flat-AnymalB-Warp-v0`` -- ``Isaac-Velocity-Flat-AnymalC-Warp-v0`` -- ``Isaac-Velocity-Flat-AnymalD-Warp-v0`` -- ``Isaac-Velocity-Flat-Cassie-Warp-v0`` -- ``Isaac-Velocity-Flat-G1-Warp-v0`` -- ``Isaac-Velocity-Flat-G1-Warp-v1`` -- ``Isaac-Velocity-Flat-H1-Warp-v0`` -- ``Isaac-Velocity-Flat-UnitreeA1-Warp-v0`` -- ``Isaac-Velocity-Flat-UnitreeGo1-Warp-v0`` -- ``Isaac-Velocity-Flat-UnitreeGo2-Warp-v0`` - -**Manipulation** - -- ``Isaac-Reach-Franka-Warp-v0`` -- ``Isaac-Reach-UR10-Warp-v0`` +Direct tasks share the stable task configuration: the stable registration +declares its warp implementation through the ``warp_entry_point`` kwarg +(mirroring ``env_cfg_entry_point``), and ``--frontend warp`` swaps only the +environment class. Stable tasks with a declared warp implementation: + +- ``Isaac-Cartpole-Direct`` — Cartpole balance +- ``Isaac-Ant-Direct`` — Ant locomotion +- ``Isaac-Humanoid-Direct`` — Humanoid locomotion +- ``Isaac-Reorient-Cube-Allegro-Direct`` — Allegro hand cube reorient + + +Manager-Based Warp Execution (``--frontend warp``) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Manager-based tasks do not need a parallel warp registration: the shared RL +CLI exposes ``--frontend {torch,warp}``, and ``--frontend warp`` adapts the +*stable* task configuration onto the warp runtime at build time (swapping each +MDP term for its warp twin). Select the Newton solver explicitly with +``presets=newton_mjwarp``. Stable tasks with full twin coverage: + +- ``Isaac-Cartpole`` +- ``Isaac-Ant`` +- ``Isaac-Humanoid`` +- ``Isaac-Reach-Franka``, ``Isaac-Reach-UR10`` +- ``Isaac-Velocity-Flat-AnymalD``, ``-Cassie``, ``-G1``, ``-H1``, ``-UnitreeGo2`` + (and their ``-Play`` variants) +- ``Isaac-Velocity-Rough-AnymalD``, ``-Cassie``, ``-G1``, ``-H1``, ``-UnitreeGo2`` + (and their ``-Play`` variants) — the terrain-levels curriculum runs as a + Warp-native update on the :class:`~isaaclab.terrains.TerrainImporter` Warp + origin views + +A missing twin is a hard error listing the affected terms, so a partially +covered task fails at build time rather than silently changing behavior. +``Isaac-Velocity-Rough-Digit`` is the one rough task not yet covered: its +``desired_contacts`` reward term has no Warp twin. Quick Start @@ -70,13 +73,17 @@ Quick Start .. code-block:: bash - # Direct workflow + # Direct workflow: stable task, warp env implementation ./isaaclab.sh train --rl_library rsl_rl \ - --task Isaac-Cartpole-Direct-Warp-v0 --num_envs 4096 + --task Isaac-Cartpole-Direct --frontend warp presets=newton_mjwarp --num_envs 4096 - # Manager-based workflow + # Manager-based workflow: stable task on the warp runtime ./isaaclab.sh train --rl_library rsl_rl \ - --task Isaac-Velocity-Flat-AnymalC-Warp-v0 --num_envs 4096 + --task Isaac-Cartpole --frontend warp presets=newton_mjwarp --num_envs 4096 + + # Manager-based workflow: velocity task on the warp runtime + ./isaaclab.sh train --rl_library rsl_rl \ + --task Isaac-Velocity-Flat-AnymalD --frontend warp presets=newton_mjwarp --num_envs 4096 All RL libraries with warp-compatible wrappers are supported: RSL-RL, RL Games, SKRL, and Stable-Baselines3. @@ -143,16 +150,6 @@ both running on the Newton physics backend. Measured over 300 iterations with 40 - 11,458 - 7,813 - -31.83% - * - AnymalB - - Manager - - 29,188 - - 21,781 - - -25.38% - * - AnymalC - - Manager - - 30,938 - - 22,228 - - -28.15% * - AnymalD - Manager - 32,294 @@ -173,17 +170,7 @@ both running on the Newton physics backend. Measured over 300 iterations with 40 - 22,202 - 15,864 - -28.55% - * - A1 - - Manager - - 15,257 - - 9,907 - - -35.07% - * - Go1 - - Manager - - 16,515 - - 11,869 - - -28.13% - * - Go2 + * - UnitreeGo2 - Manager - 15,221 - 9,966 @@ -200,7 +187,7 @@ overhead. Reading the table above: - **Manager-based classic RL** (Cartpole, Ant) — biggest gains (-52% to -54%). Many small reward / observation terms with low compute per term, so per-launch CPU overhead dominated the stable baseline. -- **Manager-based locomotion** (Anymal, G1, H1, Cassie, Unitree) — consistent -25% to -38% +- **Manager-based locomotion** (AnymalD, G1, H1, Cassie, UnitreeGo2) — consistent -21% to -38% range. The MDP has more terms but the underlying physics step is heavier, so the relative Python savings shrink. - **Direct workflow** — gains scale with how much the env's step body was Python (Ant -51%, @@ -228,7 +215,7 @@ specific to warp envs; for Newton physics limitations see :doc:`supported-featur ``class_type`` fields resolve to ``isaaclab_physx.*`` classes that depend on ``omni.physics.tensors`` (a Kit module the warp runtime does not initialise), and several warp APIs (env-mask reset, CUDA graph capture) require the Newton articulation. Configure - the cfg with a Newton physics block (or ``physics=newton_mjwarp``). + the cfg with a Newton physics block (or ``presets=newton_mjwarp``). **MDP coverage** @@ -266,25 +253,32 @@ to estimate the gain for your own task before committing to a migration. **Single-task A/B** +Run the same stable task id twice, differing only in ``--frontend``. Pass +``presets=newton_mjwarp`` to *both* runs so the physics backend is identical and the +measured difference isolates the frontend (manager pipeline + CUDA graph capture): + .. code-block:: bash - # Stable variant + # Torch frontend (stable managers) uv run isaaclab benchmark training \ --rl_library rsl_rl \ - --task -v0 \ + --task \ --num_envs 4096 \ --max_iterations 500 \ --benchmark_formatter summary \ - --output_path benchmarks/stable + --output_path benchmarks/torch \ + presets=newton_mjwarp - # Warp variant — same task with -Warp- suffix + # Warp frontend (warp managers, CUDA-graph captured) — same task id uv run isaaclab benchmark training \ --rl_library rsl_rl \ - --task -Warp-v0 \ + --task \ + --frontend warp \ --num_envs 4096 \ --max_iterations 500 \ --benchmark_formatter summary \ - --output_path benchmarks/warp + --output_path benchmarks/warp \ + presets=newton_mjwarp The ``summary`` formatter prints step time (min / mean / max) and total throughput. Compare "step time" between the two runs to estimate the gain per env step. @@ -292,7 +286,7 @@ The ``summary`` formatter prints step time (min / mean / max) and total throughp **Sweep across all available tasks** Run ``isaaclab benchmark training`` for each task in the stable set (cartpole, ant, humanoid, -locomotion, manipulation) and again with the ``-Warp-`` suffixed task ids, then diff the two output +locomotion, manipulation) and again with ``--frontend warp``, then diff the two output directories. **What to look at in the output** diff --git a/docs/source/overview/environments.rst b/docs/source/overview/environments.rst index 8e612678792f..b1ed9639915b 100644 --- a/docs/source/overview/environments.rst +++ b/docs/source/overview/environments.rst @@ -887,6 +887,14 @@ For environments that have a different task name listed under ``Inference Task N provided when running ``play.py`` or any inferencing workflows. These tasks provide more suitable configurations for inferencing, including reading from an already trained checkpoint and disabling runtime perturbations used for training. +.. note:: + + Warp-native environment implementations no longer register separate ``-Warp`` task ids. + The same stable tasks run on the Warp runtime by passing ``--frontend warp`` (together + with ``presets=newton_mjwarp``) to the training and play entrypoints; direct tasks + declare their Warp implementation through the ``warp_entry_point`` registration. See + :doc:`/overview/core-concepts/physical-backends/newton/warp-environments` for details. + .. START-AUTO-GENERATED: comprehensive-environment-list .. list-table:: @@ -907,16 +915,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Ant-Direct-Warp-v0 - - - - Direct - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - - * - Isaac-Ant-Warp-v0 - - - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - * - Isaac-Cartpole - - Manager Based @@ -941,16 +939,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Cartpole-Direct-Warp-v0 - - - - Direct - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - - * - Isaac-Cartpole-Warp-v0 - - - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - * - Isaac-Fourbar-Pole-Swingup - - Manager Based @@ -966,16 +954,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Humanoid-Direct-Warp-v0 - - - - Direct - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - - * - Isaac-Humanoid-Warp-v0 - - - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - * - Isaac-Lift-Cloth-Franka - - Manager Based @@ -1040,11 +1018,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` - * - Isaac-Reach-Franka-Warp-v0 - - Isaac-Reach-Franka-Warp-Play-v0 - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Reach-UR10 - Isaac-Reach-UR10-Play - Manager Based @@ -1060,11 +1033,6 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Reorient-Cube-Allegro-Direct-Warp-v0 - - - - Direct - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Reorient-Cube-Shadow-Camera-Direct - Isaac-Reorient-Cube-Shadow-Camera-Direct-Play - Direct @@ -1111,36 +1079,16 @@ inferencing, including reading from an already trained checkpoint and disabling - Direct - **rl_games** (PPO), **skrl** (PPO, IPPO, MAPPO) - **physics=** ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-AnymalB-Warp-v0 - - Isaac-Velocity-Flat-AnymalB-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - - * - Isaac-Velocity-Flat-AnymalC-Warp-v0 - - Isaac-Velocity-Flat-AnymalC-Warp-Play-v0 - - Manager Based - - **rl_games** (PPO), **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-AnymalD - Isaac-Velocity-Flat-AnymalD-Play - Manager Based - **rsl_rl** (PPO, DISTILLATION, DISTILLATION_RECURRENT, RECURRENT), **skrl** (PPO) - **physics=** ``newton_mjwarp``, ``ovphysx``, ``physx`` - * - Isaac-Velocity-Flat-AnymalD-Warp-v0 - - Isaac-Velocity-Flat-AnymalD-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-Cassie - Isaac-Velocity-Flat-Cassie-Play - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-Cassie-Warp-v0 - - Isaac-Velocity-Flat-Cassie-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-Digit - Isaac-Velocity-Flat-Digit-Play - Manager Based @@ -1151,46 +1099,21 @@ inferencing, including reading from an already trained checkpoint and disabling - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-G1-Warp-v0 - - Isaac-Velocity-Flat-G1-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-H1 - Isaac-Velocity-Flat-H1-Play - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-H1-Warp-v0 - - Isaac-Velocity-Flat-H1-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-Spot - Isaac-Velocity-Flat-Spot-Play - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-UnitreeA1-Warp-v0 - - Isaac-Velocity-Flat-UnitreeA1-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO), **sb3** (PPO) - - - * - Isaac-Velocity-Flat-UnitreeGo1-Warp-v0 - - Isaac-Velocity-Flat-UnitreeGo1-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Flat-UnitreeGo2 - Isaac-Velocity-Flat-UnitreeGo2-Play - Manager Based - **rsl_rl** (PPO), **skrl** (PPO) - **physics=** ``newton_kamino``, ``newton_mjwarp``, ``physx`` - * - Isaac-Velocity-Flat-UnitreeGo2-Warp-v0 - - Isaac-Velocity-Flat-UnitreeGo2-Warp-Play-v0 - - Manager Based - - **rsl_rl** (PPO), **skrl** (PPO) - - * - Isaac-Velocity-Rough-AnymalD - Isaac-Velocity-Rough-AnymalD-Play - Manager Based diff --git a/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst new file mode 100644 index 000000000000..453998f2954b --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -0,0 +1,19 @@ +Added +^^^^^ + +* Added :attr:`~isaaclab.terrains.TerrainImporter.env_origins_pa`, + :attr:`~isaaclab.terrains.TerrainImporter.terrain_levels_pa`, and + :meth:`~isaaclab.terrains.TerrainImporter.update_env_origins_mask` for zero-copy Warp terrain + access and mask-based curriculum updates. +* Added :attr:`~isaaclab.scene.InteractiveScene.env_origins_pa` and a boolean ``env_mask`` + option to :meth:`~isaaclab.scene.InteractiveScene.reset` for mask-based scene resets. + +Changed +^^^^^^^ + +* Changed the uniform pose and velocity command debug visualization to shared private mixins + reused by the Warp-native command terms. +* Changed the :class:`~isaaclab.terrains.TerrainImporter` origin, level, and type buffers to + :class:`~isaaclab.utils.warp.ProxyArray` storage: the Torch accessors and Warp views alias + one owner, so consumer-held views stay pointer-stable across reconfiguration without + per-buffer view caches. diff --git a/source/isaaclab/changelog.d/warp-manager-bridge.rst b/source/isaaclab/changelog.d/warp-manager-bridge.rst new file mode 100644 index 000000000000..8aa615a5294a --- /dev/null +++ b/source/isaaclab/changelog.d/warp-manager-bridge.rst @@ -0,0 +1,8 @@ +Changed +^^^^^^^ + +* Relaxed ``func`` annotation on :class:`~isaaclab.managers.ObservationTermCfg`, + :class:`~isaaclab.managers.RewardTermCfg`, and + :class:`~isaaclab.managers.TerminationTermCfg` to ``Callable[..., torch.Tensor | None]`` + so kernel-style ``func(env, out) -> None`` terms type-check alongside the + existing torch return-tensor form. diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py b/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py new file mode 100644 index 000000000000..6e6c6301f101 --- /dev/null +++ b/source/isaaclab/isaaclab/envs/mdp/commands/_debug_vis.py @@ -0,0 +1,112 @@ +# 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 + +"""Private debug-visualization mixins for command terms.""" + +from __future__ import annotations + +import torch + +import isaaclab.utils.math as math_utils +from isaaclab.markers import VisualizationMarkers + + +class _VelocityCommandDebugVis: + """Debug-visualization mixin shared by uniform velocity command terms.""" + + def _set_debug_vis_impl(self, debug_vis: bool): + # set visibility of markers + # note: parent only deals with callbacks. not their visibility + if debug_vis: + # create markers if necessary for the first time + if not hasattr(self, "goal_vel_visualizer"): + # -- goal + self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) + # -- current + self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) + # set their visibility to true + self.goal_vel_visualizer.set_visibility(True) + self.current_vel_visualizer.set_visibility(True) + else: + if hasattr(self, "goal_vel_visualizer"): + self.goal_vel_visualizer.set_visibility(False) + self.current_vel_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # get marker location + # -- base state + base_pos_w = self.robot.data.root_pos_w.torch.clone() + base_pos_w[:, 2] += 0.5 + # -- resolve the scales and quaternions + vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) + vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( + self.robot.data.root_lin_vel_b.torch[:, :2] + ) + # display markers + self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) + self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) + + def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Converts the XY base velocity command to arrow direction rotation.""" + # obtain default scale of the marker + default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale + # arrow-scale + arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) + arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 + # arrow-direction + heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) + zeros = torch.zeros_like(heading_angle) + arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) + # convert everything back from base to world frame + base_quat_w = self.robot.data.root_quat_w.torch + arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) + + return arrow_scale, arrow_quat + + +class _PoseCommandDebugVis: + """Debug-visualization mixin shared by uniform pose command terms.""" + + def _set_debug_vis_impl(self, debug_vis: bool): + # create markers if necessary for the first time + if debug_vis: + if not hasattr(self, "goal_pose_visualizer"): + # -- goal pose + self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) + # -- current body pose + self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) + # set their visibility to true + self.goal_pose_visualizer.set_visibility(True) + self.current_pose_visualizer.set_visibility(True) + else: + if hasattr(self, "goal_pose_visualizer"): + self.goal_pose_visualizer.set_visibility(False) + self.current_pose_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # update the markers + # -- goal pose + pose_command_w = self._debug_pose_command_w() + self.goal_pose_visualizer.visualize(pose_command_w[:, :3], pose_command_w[:, 3:]) + # -- current body pose + body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] + self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) + + def _debug_pose_command_w(self) -> torch.Tensor: + """Return the world-frame pose command consumed by the goal marker. + + The value is the desired pose in the simulation world frame [m, m, m, quaternion xyzw], + shape ``(num_envs, 7)``. Warp-native terms override this to expose the Torch alias of + their pointer-stable Warp buffer. + """ + return self.pose_command_w diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py index d98a6b5e7f81..efecc3e2a1d0 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py @@ -14,17 +14,18 @@ from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique +from ._debug_vis import _PoseCommandDebugVis + if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv from .commands_cfg import UniformPoseCommandCfg -class UniformPoseCommand(CommandTerm): +class UniformPoseCommand(_PoseCommandDebugVis, CommandTerm): """Command generator for generating pose commands uniformly. The command generator generates poses by sampling positions uniformly within specified @@ -151,31 +152,3 @@ def _resample_command(self, env_ids: Sequence[int]): def _update_command(self): pass - - def _set_debug_vis_impl(self, debug_vis: bool): - # create markers if necessary for the first time - if debug_vis: - if not hasattr(self, "goal_pose_visualizer"): - # -- goal pose - self.goal_pose_visualizer = VisualizationMarkers(self.cfg.goal_pose_visualizer_cfg) - # -- current body pose - self.current_pose_visualizer = VisualizationMarkers(self.cfg.current_pose_visualizer_cfg) - # set their visibility to true - self.goal_pose_visualizer.set_visibility(True) - self.current_pose_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_pose_visualizer"): - self.goal_pose_visualizer.set_visibility(False) - self.current_pose_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # update the markers - # -- goal pose - self.goal_pose_visualizer.visualize(self.pose_command_w[:, :3], self.pose_command_w[:, 3:]) - # -- current body pose - body_link_pose_w = self.robot.data.body_link_pose_w.torch[:, self.body_idx] - self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index 930f663f65d3..31fad8296bc4 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -16,7 +16,8 @@ import isaaclab.utils.math as math_utils from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm -from isaaclab.markers import VisualizationMarkers + +from ._debug_vis import _VelocityCommandDebugVis if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -27,7 +28,7 @@ logger = logging.getLogger(__name__) -class UniformVelocityCommand(CommandTerm): +class UniformVelocityCommand(_VelocityCommandDebugVis, CommandTerm): r"""Command generator that generates a velocity command in SE(2) from uniform distribution. The command comprises of a linear velocity in x and y direction and an angular velocity around @@ -197,63 +198,6 @@ def _update_command(self): standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten() self.vel_command_b[standing_env_ids, :] = 0.0 - def _set_debug_vis_impl(self, debug_vis: bool): - # set visibility of markers - # note: parent only deals with callbacks. not their visibility - if debug_vis: - # create markers if necessary for the first time - if not hasattr(self, "goal_vel_visualizer"): - # -- goal - self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg) - # -- current - self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg) - # set their visibility to true - self.goal_vel_visualizer.set_visibility(True) - self.current_vel_visualizer.set_visibility(True) - else: - if hasattr(self, "goal_vel_visualizer"): - self.goal_vel_visualizer.set_visibility(False) - self.current_vel_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - # check if robot is initialized - # note: this is needed in-case the robot is de-initialized. we can't access the data - if not self.robot.is_initialized: - return - # get marker location - # -- base state - base_pos_w = self.robot.data.root_pos_w.torch.clone() - base_pos_w[:, 2] += 0.5 - # -- resolve the scales and quaternions - vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2]) - vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow( - self.robot.data.root_lin_vel_b.torch[:, :2] - ) - # display markers - self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale) - self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale) - - """ - Internal helpers. - """ - - def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Converts the XY base velocity command to arrow direction rotation.""" - # obtain default scale of the marker - default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale - # arrow-scale - arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1) - arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0 - # arrow-direction - heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0]) - zeros = torch.zeros_like(heading_angle) - arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle) - # convert everything back from base to world frame - base_quat_w = self.robot.data.root_quat_w.torch - arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat) - - return arrow_scale, arrow_quat - class NormalVelocityCommand(UniformVelocityCommand): """Command generator that generates a velocity command in SE(2) from a normal distribution. diff --git a/source/isaaclab/isaaclab/managers/manager_term_cfg.py b/source/isaaclab/isaaclab/managers/manager_term_cfg.py index 29e2281c3c60..3fad3881679d 100644 --- a/source/isaaclab/isaaclab/managers/manager_term_cfg.py +++ b/source/isaaclab/isaaclab/managers/manager_term_cfg.py @@ -152,7 +152,7 @@ class CurriculumTermCfg(ManagerTermBaseCfg): class ObservationTermCfg(ManagerTermBaseCfg): """Configuration for an observation term.""" - func: Callable[..., torch.Tensor] = MISSING + func: Callable[..., torch.Tensor | None] = MISSING """The name of the function to be called. This function should take the environment object and any other parameters @@ -335,7 +335,7 @@ class EventTermCfg(ManagerTermBaseCfg): class RewardTermCfg(ManagerTermBaseCfg): """Configuration for a reward term.""" - func: Callable[..., torch.Tensor] = MISSING + func: Callable[..., torch.Tensor | None] = MISSING """The name of the function to be called. This function should take the environment object and any other parameters @@ -363,7 +363,7 @@ class RewardTermCfg(ManagerTermBaseCfg): class TerminationTermCfg(ManagerTermBaseCfg): """Configuration for a termination term.""" - func: Callable[..., torch.Tensor] = MISSING + func: Callable[..., torch.Tensor | None] = MISSING """The name of the function to be called. This function should take the environment object and any other parameters diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 15b693197a46..000590c4fa0c 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -36,6 +36,7 @@ from isaaclab.sim import SimulationContext from isaaclab.sim.utils.stage import get_current_stage, get_current_stage_id from isaaclab.sim.views import FrameView +from isaaclab.utils.warp import ProxyArray # Note: This is a temporary import for the VisuoTactileSensorCfg class. # It will be removed once the VisuoTactileSensor class is added to the core Isaac Lab framework. @@ -190,6 +191,12 @@ def __init__(self, cfg: InteractiveSceneCfg): self._aggregate_scene_data_requirements(requested_viz_types) + # The terrain importer owns its origins buffer together with its Warp view. + # For clone-plan origins, :attr:`env_origins_pa` caches a zero-copy proxy on + # first access (manager/term initialization, outside capture), so scene + # construction does not require a published clone plan. + self._clone_origins_pa = None + # Collision filtering is PhysX-only (matches both physx and ovphysx). if self.cfg.filter_collisions and "physx" in self.physics_backend and self._is_scene_setup_from_cfg(): self.filter_collisions(self._global_prim_paths) @@ -371,6 +378,20 @@ def env_origins(self) -> torch.Tensor: return self._terrain.env_origins return self.sim.get_clone_plan().positions + @property + def env_origins_pa(self) -> ProxyArray: + """Environment origins proxy [m], shape ``(num_envs,)`` vec3; use ``.warp`` / ``.torch`` views. + + Terrain-backed scenes forward the terrain importer's proxy so both accessors + always share one storage owner. Clone-plan origins are wrapped lazily on + first access and re-wrapped only if a new plan is published. + """ + if self._terrain is not None: + return self._terrain.env_origins_pa + if self._clone_origins_pa is None or self._clone_origins_pa.warp.ptr != self.env_origins.data_ptr(): + self._clone_origins_pa = ProxyArray(wp.from_torch(self.env_origins, dtype=wp.vec3f)) + return self._clone_origins_pa + @property def terrain(self) -> TerrainImporter | None: """The terrain in the scene. If None, then the scene has no terrain. @@ -450,27 +471,49 @@ def state(self) -> dict[str, dict[str, dict[str, torch.Tensor]]]: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None): + def reset( + self, + env_ids: Sequence[int] | None = None, + env_mask: wp.array(dtype=wp.bool) | None = None, + ): """Resets the scene entities. Args: env_ids: The indices of the environments to reset. Defaults to None (all instances). + env_mask: Boolean Warp mask selecting environments. When provided, + it takes precedence over :paramref:`env_ids`. + + .. caution:: + ``env_mask`` is honored by mask-native assets (the Newton backend). + Classic PhysX assets currently accept and ignore it, so passing a + mask in such scenes resets unintended environments — use + :paramref:`env_ids` there. + """ + reset_kwargs = {"env_mask": env_mask} if env_mask is not None else {"env_ids": env_ids} # -- assets for articulation in self._articulations.values(): - articulation.reset(env_ids) + articulation.reset(**reset_kwargs) for deformable_object in self._deformable_objects.values(): - deformable_object.reset(env_ids) + deformable_object.reset(**reset_kwargs) for rigid_object in self._rigid_objects.values(): - rigid_object.reset(env_ids) - for surface_gripper in self._surface_grippers.values(): - surface_gripper.reset(env_ids) + rigid_object.reset(**reset_kwargs) + if env_mask is not None: + # Surface grippers are CPU-only assets driven through a Torch API. + # Hand them the boolean view and let their compatibility layer own + # any index materialization. + gripper_env_mask = wp.to_torch(env_mask) if self._surface_grippers else None + for surface_gripper in self._surface_grippers.values(): + surface_gripper.reset_mask(gripper_env_mask) + else: + for surface_gripper in self._surface_grippers.values(): + surface_gripper.reset(env_ids) for rigid_object_collection in self._rigid_object_collections.values(): - rigid_object_collection.reset(env_ids) + rigid_object_collection.reset(**reset_kwargs) # -- sensors for sensor in self._sensors.values(): - sensor.reset(env_ids) + sensor.reset(**reset_kwargs) def write_data_to_sim(self): """Writes the data of the scene entities to the simulation.""" diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index c78f9ae3344e..b1e0b8cdb6de 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -11,10 +11,12 @@ import numpy as np import torch import trimesh +import warp as wp import isaaclab.sim as sim_utils from isaaclab.markers import VisualizationMarkers from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.utils.warp import ProxyArray from .utils import create_prim_from_mesh @@ -25,6 +27,34 @@ logger = logging.getLogger(__name__) +@wp.kernel +def _update_env_origins_mask( + env_mask: wp.array(dtype=wp.bool), + move_up: wp.array(dtype=wp.bool), + move_down: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + terrain_levels: wp.array(dtype=wp.int64), + terrain_types: wp.array(dtype=wp.int64), + terrain_origins: wp.array(dtype=wp.vec3f, ndim=2), + env_origins: wp.array(dtype=wp.vec3f), + max_terrain_level: int, +): + """Update terrain levels and origins selected by an environment mask.""" + env_id = wp.tid() + if env_mask[env_id]: + state = rng_state[env_id] + random_level = wp.randi(state, 0, max_terrain_level) + # Keep shared terrain indices int64 so their zero-copy Torch views support advanced indexing. + level = terrain_levels[env_id] + wp.int64(move_up[env_id]) - wp.int64(move_down[env_id]) + if level >= wp.int64(max_terrain_level): + level = wp.int64(random_level) + elif level < wp.int64(0): + level = wp.int64(0) + terrain_levels[env_id] = level + env_origins[env_id] = terrain_origins[level, terrain_types[env_id]] + rng_state[env_id] = state + + class TerrainImporter: r"""A class to handle terrain meshes and import them into the simulator. @@ -45,16 +75,6 @@ class TerrainImporter: terrain_prim_paths: list[str] """A list containing the USD prim paths to the imported terrains.""" - terrain_origins: torch.Tensor | None - """The origins of the sub-terrains in the added terrain mesh. Shape is (num_rows, num_cols, 3). - - If terrain origins is not None, the environment origins are computed based on the terrain origins. - Otherwise, the environment origins are computed based on the grid spacing. - """ - - env_origins: torch.Tensor - """The origins of the environments. Shape is (num_envs, 3).""" - def __init__(self, cfg: TerrainImporterCfg): """Initialize the terrain importer. @@ -73,10 +93,14 @@ def __init__(self, cfg: TerrainImporterCfg): self.cfg = cfg self.device = sim_utils.SimulationContext.instance().device # type: ignore - # create buffers for the terrains + # create buffers for the terrains. Each buffer is one ProxyArray owning the + # storage that both the public Torch accessors and the Warp views alias, so + # consumer-held views stay pointer-stable across reconfiguration. self.terrain_prim_paths = list() - self.terrain_origins = None - self.env_origins = None # assigned later when `configure_env_origins` is called + self._terrain_origins_pa: ProxyArray | None = None + self._env_origins_pa: ProxyArray | None = None # assigned when `configure_env_origins` is called + self._terrain_levels_pa: ProxyArray | None = None + self._terrain_types_pa: ProxyArray | None = None # private variables self._terrain_flat_patches = dict() @@ -144,6 +168,38 @@ def terrain_names(self) -> list[str]: """A list of names of the imported terrains.""" return [f"'{path.split('/')[-1]}'" for path in self.terrain_prim_paths] + @property + def terrain_origins(self) -> torch.Tensor | None: + """Sub-terrain origins [m], shape ``(num_rows, num_cols, 3)``; ``None`` for grid-like origins.""" + return self._terrain_origins_pa.torch if self._terrain_origins_pa is not None else None + + @property + def env_origins(self) -> torch.Tensor | None: + """Per-env origins [m], shape ``(num_envs, 3)``; ``None`` until origins are configured.""" + return self._env_origins_pa.torch if self._env_origins_pa is not None else None + + @property + def terrain_levels(self) -> torch.Tensor | None: + """Per-env terrain levels, shape ``(num_envs,)``, int64; ``None`` for grid-like origins.""" + return self._terrain_levels_pa.torch if self._terrain_levels_pa is not None else None + + @property + def terrain_types(self) -> torch.Tensor | None: + """Per-env terrain types, shape ``(num_envs,)``, int64; ``None`` for grid-like origins.""" + return self._terrain_types_pa.torch if self._terrain_types_pa is not None else None + + @property + def terrain_levels_pa(self) -> ProxyArray: + """Terrain levels proxy, shape ``(num_envs,)``, int64; use ``.warp`` / ``.torch`` views.""" + if self._terrain_levels_pa is None: + raise RuntimeError("Terrain levels are unavailable for grid-like environment origins.") + return self._terrain_levels_pa + + @property + def env_origins_pa(self) -> ProxyArray | None: + """Environment origins proxy [m], shape ``(num_envs,)`` vec3; use ``.warp`` / ``.torch`` views.""" + return self._env_origins_pa + """ Operations - Visibility. """ @@ -288,31 +344,50 @@ def import_usd(self, name: str, usd_path: str): Operations - Origins. """ - def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None): + def configure_env_origins(self, origins: np.ndarray | torch.Tensor | None = None) -> None: """Configure the origins of the environments based on the added terrain. + Reconfiguration copies into the existing origin buffers when layouts match, + so cached Torch tensors and Warp views held by consumers stay valid. + Args: - origins: The origins of the sub-terrains. Shape is (num_rows, num_cols, 3). + origins: The origins [m] of the sub-terrains. Shape is (num_rows, num_cols, 3). """ # decide whether to compute origins in a grid or based on curriculum if origins is not None: - # convert to numpy + # convert to torch if isinstance(origins, np.ndarray): origins = torch.from_numpy(origins) # store the origins - self.terrain_origins = origins.to(self.device, dtype=torch.float) + self._terrain_origins_pa = self._assign_proxy_stable( + self._terrain_origins_pa, origins.to(self.device, dtype=torch.float).contiguous(), wp.vec3f + ) # compute environment origins - self.env_origins = self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins) + self._env_origins_pa = self._assign_proxy_stable( + self._env_origins_pa, + self._compute_env_origins_curriculum(self.cfg.num_envs, self.terrain_origins), + wp.vec3f, + ) else: - self.terrain_origins = None + self._terrain_origins_pa = None + self._terrain_levels_pa = None + self._terrain_types_pa = None # check if env spacing is valid if self.cfg.env_spacing is None: raise ValueError("Environment spacing must be specified for configuring grid-like origins.") # compute environment origins - self.env_origins = self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing) + self._env_origins_pa = self._assign_proxy_stable( + self._env_origins_pa, self._compute_env_origins_grid(self.cfg.num_envs, self.cfg.env_spacing), wp.vec3f + ) - def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor): - """Update the environment origins based on the terrain levels.""" + def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_down: torch.Tensor) -> None: + """Update environment origins through the Torch ID-based interface. + + Args: + env_ids: Environment indices to update. + move_up: Flags that increment the selected terrain levels. + move_down: Flags that decrement the selected terrain levels. + """ # check if grid-like spawning if self.terrain_origins is None: return @@ -328,10 +403,88 @@ def update_env_origins(self, env_ids: torch.Tensor, move_up: torch.Tensor, move_ # update the env origins self.env_origins[env_ids] = self.terrain_origins[self.terrain_levels[env_ids], self.terrain_types[env_ids]] + def update_env_origins_mask( + self, + env_mask: wp.array(dtype=wp.bool), + move_up: wp.array(dtype=wp.bool), + move_down: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + ) -> None: + """Update terrain levels and origins through the Warp mask-based interface. + + Moving past the maximum level wraps to a random level, moving below zero clamps to zero, and + unselected environments remain unchanged. + + Args: + env_mask: Boolean mask selecting environments, shape ``(num_envs,)``. + move_up: Flags that increment terrain levels, shape ``(num_envs,)``. + move_down: Flags that decrement terrain levels, shape ``(num_envs,)``. + rng_state: Per-environment random-number-generator state, shape ``(num_envs,)``. + + Raises: + TypeError: If an input is not a Warp array with the required data type. + ValueError: If an input has the wrong shape or device. + """ + if self.terrain_origins is None: + return + num_envs = self.terrain_levels.shape[0] + arrays = { + "env_mask": (env_mask, wp.bool), + "move_up": (move_up, wp.bool), + "move_down": (move_down, wp.bool), + "rng_state": (rng_state, wp.uint32), + } + for name, (array, dtype) in arrays.items(): + if not isinstance(array, wp.array) or array.dtype != dtype: + raise TypeError(f"{name} must be a Warp array with dtype {dtype}; received {array}.") + if array.ndim != 1 or array.shape[0] != num_envs: + raise ValueError(f"{name} must have shape ({num_envs},); received {array.shape}.") + if array.device != wp.get_device(self.device): + raise ValueError(f"{name} must be on device {self.device}; received {array.device}.") + + wp.launch( + kernel=_update_env_origins_mask, + dim=num_envs, + inputs=[ + env_mask, + move_up, + move_down, + rng_state, + self._terrain_levels_pa.warp, + self._terrain_types_pa.warp, + self._terrain_origins_pa.warp, + self._env_origins_pa.warp, + self.max_terrain_level, + ], + device=self.device, + ) + """ Internal helpers. """ + @staticmethod + def _assign_proxy_stable(current: ProxyArray | None, new: torch.Tensor, dtype) -> ProxyArray: + """Copy into the existing proxy storage when layouts match so consumer-held views stay valid. + + Args: + current: The proxy currently owning the buffer, or ``None``. + new: Freshly computed values to store. + dtype: Warp dtype of the proxy's array (e.g. ``wp.vec3f``). + + Returns: + The proxy holding the new values; :paramref:`current` when its layout matches. + """ + if ( + current is not None + and current.torch.shape == new.shape + and current.torch.dtype == new.dtype + and current.torch.device == new.device + ): + current.torch.copy_(new) + return current + return ProxyArray(wp.from_torch(new.contiguous(), dtype=dtype)) + def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) -> torch.Tensor: """Compute the origins of the environments defined by the sub-terrains origins.""" # extract number of rows and cols @@ -344,10 +497,18 @@ def _compute_env_origins_curriculum(self, num_envs: int, origins: torch.Tensor) # store maximum terrain level possible self.max_terrain_level = num_rows # define all terrain levels and types available - self.terrain_levels = torch.randint(0, max_init_level + 1, (num_envs,), device=self.device) - self.terrain_types = torch.div( - torch.arange(num_envs, device=self.device), (num_envs / num_cols), rounding_mode="floor" - ).to(torch.long) + self._terrain_levels_pa = self._assign_proxy_stable( + self._terrain_levels_pa, + torch.randint(0, max_init_level + 1, (num_envs,), device=self.device), + wp.int64, + ) + self._terrain_types_pa = self._assign_proxy_stable( + self._terrain_types_pa, + torch.div(torch.arange(num_envs, device=self.device), (num_envs / num_cols), rounding_mode="floor").to( + torch.long + ), + wp.int64, + ) # create tensor based on number of environments env_origins = torch.zeros(num_envs, 3, device=self.device) env_origins[:] = origins[self.terrain_levels, self.terrain_types] diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index 41aaaf7f9c3d..749f58591e20 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -13,9 +13,11 @@ """Rest everything follows.""" from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch +import warp as wp import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg @@ -132,6 +134,61 @@ def test_reset_to_env_ids_input_types(device, setup_scene): assert_state_equal(prev_state, scene.get_state()) +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_env_origins_pa_is_cached_zero_copy_view(device, setup_scene): + """The scene should expose one persistent Warp view of its Torch origins.""" + make_scene, sim = setup_scene + scene = InteractiveScene(make_scene(num_envs=4)) + sim.reset() + + assert scene.env_origins_pa.warp is scene.env_origins_pa.warp + assert scene.env_origins_pa.warp.dtype == wp.vec3f + assert scene.env_origins_pa.warp.shape == (scene.num_envs,) + assert scene.env_origins_pa.warp.ptr == scene.env_origins.data_ptr() + + +def test_reset_dispatches_warp_mask_to_supported_entities() -> None: + """Mask reset should remain mask-based across assets and sensors.""" + scene = object.__new__(InteractiveScene) + articulation = Mock() + deformable = Mock() + rigid_object = Mock() + rigid_collection = Mock() + sensor = Mock() + scene._articulations = {"articulation": articulation} + scene._deformable_objects = {"deformable": deformable} + scene._rigid_objects = {"rigid_object": rigid_object} + scene._rigid_object_collections = {"collection": rigid_collection} + scene._sensors = {"sensor": sensor} + scene._surface_grippers = {} + env_mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") + + scene.reset(env_mask=env_mask) + + for entity in (articulation, deformable, rigid_object, rigid_collection, sensor): + entity.reset.assert_called_once_with(env_mask=env_mask) + + +def test_mask_reset_forwards_boolean_mask_to_surface_grippers() -> None: + """Mask reset should hand surface grippers a boolean Torch view, not compact IDs.""" + scene = object.__new__(InteractiveScene) + gripper = Mock() + scene._articulations = {} + scene._deformable_objects = {} + scene._rigid_objects = {} + scene._rigid_object_collections = {} + scene._sensors = {} + scene._surface_grippers = {"gripper": gripper} + env_mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") + + scene.reset(env_mask=env_mask) + + gripper.reset.assert_not_called() + gripper_env_mask = gripper.reset_mask.call_args.args[0] + assert gripper_env_mask.dtype == torch.bool + assert torch.equal(gripper_env_mask, torch.tensor([True, False, True])) + + def test_scene_publishes_plan_via_replicate(monkeypatch: pytest.MonkeyPatch): """A cfg-driven scene forwards the right plan and stage to cloner.replicate. diff --git a/source/isaaclab/test/terrains/test_terrain_importer.py b/source/isaaclab/test/terrains/test_terrain_importer.py index 553f01f79233..9b53f80d9167 100644 --- a/source/isaaclab/test/terrains/test_terrain_importer.py +++ b/source/isaaclab/test/terrains/test_terrain_importer.py @@ -52,6 +52,14 @@ def test_terrain_importer_env_origins(device, env_spacing, num_envs): # obtain env origins using terrain importer terrain_importer_origins = terrain_importer.env_origins + # check that the cached Warp view shares the canonical Torch storage + assert isinstance(terrain_importer.env_origins, torch.Tensor) + assert terrain_importer.env_origins_pa.warp.dtype == wp.vec3f + assert terrain_importer.env_origins_pa.warp.shape == (num_envs,) + assert terrain_importer_origins.shape == (num_envs, 3) + assert terrain_importer.env_origins_pa.warp is terrain_importer.env_origins_pa.warp + assert terrain_importer_origins.data_ptr() == terrain_importer.env_origins_pa.warp.ptr + # obtain env origins using Lab's grid_transforms lab_grid_origins, _ = lab_cloner.grid_transforms(num_envs, spacing=env_spacing, device=sim.device) diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst new file mode 100644 index 000000000000..03bcd096bfc8 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -0,0 +1,37 @@ +Added +^^^^^ + +* Added Warp-native uniform pose and velocity command terms with pointer-stable + command arrays, and a Warp-native :class:`NullCommand` term for environments + without commands. +* Added a Warp-native ``height_scan`` observation term (eager until sensor + capture-readiness lands) with ray-count output-dimension inference. +* Added opportunistic curriculum adaptation: curriculum terms swap to warp + twins when one exists (e.g. terrain levels), keeping the legacy ID-based + fallback for stable terms without twins. +* Added :class:`~isaaclab_experimental.managers.CurriculumManager` and a + boolean-mask reset path for Warp manager-based environments, retaining compact + environment IDs only for legacy host consumers. +* Added Warp-native terrain curricula backed by cached zero-copy Warp views of + terrain and scene origins. +* Added a boolean ``env_mask`` option to :meth:`ManagerBasedEnvWarp.reset_to`, the + ``ISAACLAB_WARP_CAPTURE`` environment variable to force eager execution, and + :attr:`WarpGraphCache.captured_stages` introspection for capture-coverage checks. + +Changed +^^^^^^^ + +* **Breaking:** Removed the experimental ``ManagerCallSwitch`` and + ``MANAGER_CALL_CONFIG`` interface. Warp environments now instantiate Warp + managers directly and automatically keep managers with non-capturable terms + on eager execution. + +Fixed +^^^^^ + +* Fixed Warp event state leaking across environments and configurations, and + corrected center-of-mass sampling and termination metrics for partial resets. +* Fixed captured velocity-command terms to read current root state on every + replay instead of stale lazy-derived buffers. +* Fixed identity quaternion initialization for the Warp-native uniform pose + command term. diff --git a/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst new file mode 100644 index 000000000000..5f47328da3b0 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/warp-manager-bridge.rst @@ -0,0 +1,50 @@ +Added +^^^^^ + +* Added ``--frontend {torch,warp}`` to the shared reinforcement-learning + training and play CLIs (all supported RL libraries) for selecting the + environment runtime; default ``torch`` is unchanged. The ``rlinf`` + integration constructs environments inside the external framework and is + not frontend-routable. +* Added :mod:`isaaclab_experimental.envs.frontend` runtime selector and + :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable` used by + ``--frontend=warp`` to adapt stable cfgs onto the warp runtime. +* Added ``warp_entry_point`` registration support: a stable direct task may + declare its warp environment class, which ``--frontend=warp`` constructs + with the stable configuration. +* Added observation-noise conversion to the warp frontend: stable noise cfgs + (constant/uniform/gaussian) swap to their warp-native twins during cfg + adaptation; cfgs without a twin (class-based noise models, customized or + tensor-valued parameters) are a hard error instead of being silently + ignored. The Warp observation manager rejects non-warp noise cfgs and + plumbs the shared per-env RNG state to function-style noise kernels. +* Added warp adapters for the stable + :class:`~isaaclab.envs.mdp.events.randomize_rigid_body_material` and + :class:`~isaaclab.envs.mdp.events.randomize_rigid_body_mass` startup event + terms, so tasks using them run under ``--frontend=warp`` with randomization + active. +* Added deterministic warp twin resolution: a stable symbol's twin is looked + up on the mirrored experimental package tree (``isaaclab.* ↔ + isaaclab_experimental.*``, ``isaaclab_tasks.* ↔ + isaaclab_tasks_experimental.*``); missing twins are collected and reported + in one failure. No registration is needed. + +Changed +^^^^^^^ + +* Changed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` to adapt + its configuration in ``__init__`` via + :meth:`~isaaclab_experimental.envs.frontend.WarpFrontend.adapt_cfg`, so + registered warp task variants can derive from stable configurations instead + of duplicating them. +* Changed the experimental warp ``RewardManager`` to merge dictionaries + returned by class-based reward term ``reset()`` into its episode-log extras + (persistent Warp buffer views, CUDA-graph safe); used to report + ``Metrics/success_rate`` under ``--frontend warp``. + +Fixed +^^^^^ + +* Fixed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` to probe + :meth:`~isaaclab.sim.SimulationContext.has_active_visualizers` after the + ``get_setting`` API was reshaped. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi index ef3d01d6b4ec..109fa5f8d06c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/__init__.pyi @@ -6,13 +6,11 @@ __all__ = [ "mdp", "DirectRLEnvWarp", - "InteractiveSceneWarp", "ManagerBasedEnvWarp", "ManagerBasedRLEnvWarp", ] from . import mdp from .direct_rl_env_warp import DirectRLEnvWarp -from .interactive_scene_warp import InteractiveSceneWarp from .manager_based_env_warp import ManagerBasedEnvWarp from .manager_based_rl_env_warp import ManagerBasedRLEnvWarp diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index d69b13a8feec..2492b964c072 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py @@ -27,17 +27,18 @@ from isaaclab.envs.common import VecEnvObs, VecEnvStepReturn from isaaclab.envs.direct_rl_env import DirectRLEnv from isaaclab.envs.direct_rl_env_cfg import DirectRLEnvCfg +from isaaclab.envs.ui import ViewportCameraController from isaaclab.envs.utils.spaces import sample_space, spec_to_gym_space - -# from isaaclab.envs.ui import ViewportCameraController from isaaclab.managers import EventManager +from isaaclab.scene import InteractiveScene from isaaclab.sim import SimulationContext from isaaclab.sim.utils import use_stage from isaaclab.utils.noise import NoiseModel from isaaclab.utils.seed import configure_seed from isaaclab.utils.timer import Timer +from isaaclab.utils.version import has_kit -from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp +from isaaclab_experimental.utils.warp import increment_all_int32, zero_masked_int32 from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache # from isaacsim.core.simulation_manager import SimulationManager @@ -54,25 +55,6 @@ """Enable all fine-grained inner timers (adds wp.synchronize per sub-phase). Set DEBUG_TIMERS=1 env var to enable.""" -@wp.kernel -def zero_mask_int32( - mask: wp.array(dtype=wp.bool), - data: wp.array(dtype=wp.int32), -): - env_index = wp.tid() - if mask[env_index]: - data[env_index] = 0 - - -@wp.kernel -def add_to_env( - data: wp.array(dtype=wp.int32), - value: wp.int32, -): - env_index = wp.tid() - data[env_index] += value - - class DirectRLEnvWarp(DirectRLEnv): """The superclass for the direct workflow to design environments. @@ -162,21 +144,16 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs with Timer("[INFO]: Time taken for scene creation", "scene_creation"): # set the stage context for scene creation steps which use the stage with use_stage(self.sim.stage): - self.scene = InteractiveSceneWarp(self.cfg.scene) + self.scene = InteractiveScene(self.cfg.scene) self._setup_scene() self.scene.initialize_renderers() # attach_stage_to_usd_context() print("[INFO]: Scene manager: ", self.scene) - # set up camera viewport controller - # viewport is not available in other rendering modes so the function will throw a warning - # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for - # non-rendering modes. - has_gui = bool(self.sim.get_setting("/isaaclab/has_gui")) - offscreen_render = bool(self.sim.get_setting("/isaaclab/render/offscreen")) - if has_gui or offscreen_render: - # self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) - self.viewport_camera_controller = None + # Initialize when a Kit viewport exists. ViewportCameraController uses omni.kit + # (renderer camera); skip in kitless Newton-only runs where no Kit app is running. + if (self.sim.has_gui or self.sim.has_active_visualizers()) and has_kit(): + self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None @@ -243,8 +220,13 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.torch_reset_time_outs: torch.Tensor = None self.torch_episode_length_buf: torch.Tensor = None - # Warp CUDA graph cache for capture-or-replay - self._graph_cache = WarpGraphCache() + # Pure-kernel task stages capture as on develop. Only the reset stage is + # registered eager: it dispatches scene.reset(), whose legacy Torch + # actuator boundary materializes compact IDs on the host, and CUDA graph + # replay would silently skip that Python-side reset work. The follow-up + # launch-replay execution path restores its speed correctly. + self._warp_graph_cache = WarpGraphCache(device=self.device) + self._warp_graph_cache.register_capturability("DirectReset", False) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() @@ -427,10 +409,10 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # set actions into buffers # simulate with Timer(name="apply_action", msg="Action processing step took:", enable=DEBUG_TIMERS): - self._graph_cache.capture_or_replay("action", self.step_warp_action) + self._warp_graph_cache.call("DirectAction_step", self.step_warp_action) - # write_data_to_sim runs outside the CUDA graph because _apply_actuator_model - # uses torch ops (wp.to_torch + torch arithmetic) that cross CUDA streams. + # Keep scene writes outside the task graph until scene, sensor, and + # actuator capturability have been validated as one backend boundary. with Timer(name="write_data_to_sim_loop", msg="Write data to sim (loop) took:", enable=DEBUG_TIMERS): self.scene.write_data_to_sim() @@ -447,12 +429,16 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.common_step_counter += 1 # total step (common for all envs) with Timer(name="end_pre_graph", msg="End pre-graph took:", enable=DEBUG_TIMERS): - self._graph_cache.capture_or_replay("end_pre", self._step_warp_end_pre) - # write_data_to_sim runs uncaptured — it uses torch ops that cross CUDA streams. + self._warp_graph_cache.call("DirectEndPre_compute", self._step_warp_end_pre) + # Reset stage: registered eager (see __init__) because its scene dispatch + # crosses the legacy actuator host boundary. + with Timer(name="reset_stage", msg="Reset stage took:", enable=DEBUG_TIMERS): + self._warp_graph_cache.call("DirectReset_apply", self._step_warp_reset) + # Keep the post-reset scene write at the explicit backend boundary. with Timer(name="write_data_to_sim_post", msg="Write data to sim (post-reset) took:", enable=DEBUG_TIMERS): self.scene.write_data_to_sim() with Timer(name="end_post_graph", msg="End post-graph took:", enable=DEBUG_TIMERS): - self._graph_cache.capture_or_replay("end_post", self._step_warp_end_post) + self._warp_graph_cache.call("DirectEndPost_step", self._step_warp_end_post) # Visualization hook — runs after CUDA graph scope. Override in subclass # to update markers or other non-graphable visual elements. @@ -479,15 +465,13 @@ def _post_step_visualize(self) -> None: def step_warp_action(self) -> None: self._apply_action() - # Note: scene.write_data_to_sim() is called separately outside the CUDA graph - # capture scope because it invokes _apply_actuator_model() which uses torch - # arithmetic (wp.to_torch + torch ops). This would cause a CUDA stream crossing - # error during graph capture. Moving it outside is safe since it runs every step. + # Scene writes remain a separate backend stage. This keeps the Warp-first + # task path correct without making capture support a prerequisite. def _step_warp_end_pre(self) -> None: - """Capturable portion before write_data_to_sim (pure warp kernels).""" + """Capturable portion before the reset stage (pure warp kernels).""" wp.launch( - add_to_env, + increment_all_int32, dim=self.num_envs, inputs=[ self._episode_length_buf_wp, @@ -497,7 +481,8 @@ def _step_warp_end_pre(self) -> None: self._get_dones() self._get_rewards() - # -- reset envs that terminated/timed-out and log the episode information + def _step_warp_reset(self) -> None: + """Reset envs that terminated/timed-out (eager: crosses host boundaries).""" self._reset_idx(mask=self.reset_buf) def _step_warp_end_post(self) -> None: @@ -570,13 +555,8 @@ def render(self, recompute: bool = False) -> np.ndarray | None: if self.render_mode == "human" or self.render_mode is None: return None elif self.render_mode == "rgb_array": - # check that if any render could have happened - has_gui = bool(self.sim.get_setting("/isaaclab/has_gui")) - offscreen_render = bool(self.sim.get_setting("/isaaclab/render/offscreen")) - # Rendering is possible if we have GUI or offscreen rendering enabled - can_render = has_gui or offscreen_render - - if not can_render: + # rendering requires a GUI or offscreen rendering (mirrors the stable env) + if not (self.sim.has_gui or self.sim.has_offscreen_render): render_mode_name = "NO_GUI_OR_RENDERING" raise RuntimeError( f"Cannot render '{self.render_mode}' when the simulation render mode is" @@ -741,7 +721,7 @@ def _reset_idx(self, mask: wp.array | None = None): # reset the episode length buffer wp.launch( - zero_mask_int32, + zero_masked_int32, dim=self.num_envs, inputs=[ mask, diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py new file mode 100644 index 000000000000..91f65daef0d0 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -0,0 +1,589 @@ +# 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 + +"""Warp runtime frontend for IsaacLab tasks (``--frontend warp``). + +:class:`WarpFrontend` encapsulates everything the warp runtime needs from a +stable task definition: + +* Manager-based tasks: adapt the stable cfg in place (Newton physics check, + :class:`SceneEntityCfg` promotion, MDP twin swap) and construct + :class:`ManagerBasedRLEnvWarp`. All warp MDP twins (functions *and* classes) + must exist for the chosen task — a missing twin is a hard failure, not a + silent drop. +* Direct tasks: construct the warp env class the stable registration declares + via its ``warp_entry_point`` kwarg, sharing the stable cfg. + +Twin resolution is purely mechanical: the experimental packages mirror the +stable tree (``isaaclab.* ↔ isaaclab_experimental.*``, +``isaaclab_tasks.* ↔ isaaclab_tasks_experimental.*``), and a stable symbol's +twin is looked up by name on the nearest ``.mdp`` package of the mirrored +path. There is no registration; anything unresolved is collected and reported +in one hard failure. The torch runtime never touches this module — the shared +RL CLI dispatches to plain ``gym.make`` before importing it, so this +(optional) package stays warp-only. +""" + +from __future__ import annotations + +import importlib +import importlib.util +import logging +from collections.abc import Iterable, Iterator +from enum import StrEnum +from types import ModuleType +from typing import Any, ClassVar + +import gymnasium as gym +import torch + +from isaaclab.envs import DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg +from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as _StableSceneEntityCfg + +logger = logging.getLogger(__name__) + + +__all__ = [ + "FrontendIncompatibleError", + "WarpFrontend", + "Workflow", +] + + +class Workflow(StrEnum): + """Manager-based vs direct env workflow.""" + + MANAGER_BASED = "manager_based" + DIRECT = "direct" + + +class FrontendIncompatibleError(RuntimeError): + """Raised when the warp frontend can't run the requested task.""" + + +class WarpFrontend: + """Builds environments on the warp runtime from stable task definitions. + + The class is a stateless namespace: twin resolution is a pure function of + the stable module paths and the mirrored package tree (:attr:`MIRROR_ROOTS`); + there is no registry and no import-order dependence. + """ + + # Stable package roots mirrored by their warp implementation packages. Twin + # resolution is purely mechanical: mirror the root, then look up the nearest + # enclosing ``.mdp`` package on the mirrored path. + MIRROR_ROOTS: ClassVar[dict[str, str]] = { + "isaaclab_tasks": "isaaclab_tasks_experimental", + "isaaclab": "isaaclab_experimental", + } + + # Prefixes that identify warp-origin symbols and registrations. + WARP_ROOT_PREFIXES: ClassVar[tuple[str, ...]] = tuple(MIRROR_ROOTS.values()) + + # Top-level cfg groups whose managers run warp-first. Only terms under these + # are adapted (SceneEntityCfg promotion + MDP twin swap). The event manager + # is warp-first too — it invokes term funcs with a Warp env-mask, so a stable + # event func (which expects torch ``env_ids``) breaks at runtime; its funcs + # must be swapped to warp twins. The command manager is the warp-native + # ``CommandManager`` (captured), so command term ``class_type``s must swap + # to their warp twins as well. A stable term left on a warp manager would + # break, so a missing twin in these groups is a hard error. The recorder + # manager runs on the stable implementation and is skipped entirely; + # curriculum is adapted opportunistically (see ``OPTIONAL_WARP_GROUPS``). + WARP_MANAGED_GROUPS: ClassVar[frozenset[str]] = frozenset( + {"observations", "rewards", "terminations", "actions", "events", "commands"} + ) + + # Groups adapted opportunistically: a warp twin is swapped in when one exists, + # otherwise the stable term stays on the manager's legacy (host-ID) fallback. + # Curriculum is the only such group: its manager supports stable terms by + # design, at the cost of eager execution and one ID materialization per reset. + OPTIONAL_WARP_GROUPS: ClassVar[frozenset[str]] = frozenset({"curriculum"}) + + # ------------------------------------------------------------------ + # Env construction + # ------------------------------------------------------------------ + + @classmethod + def build_env(cls, env_cfg: Any, task_id: str, **construct_kwargs: Any) -> gym.Env: + """Construct the warp env for ``task_id`` from a stable env cfg. + + Args: + env_cfg: Stable env cfg. Manager-based cfgs are mutated in place. + task_id: Gym registration id, e.g. ``"Isaac-Cartpole"``. + **construct_kwargs: Forwarded to the env constructor (``render_mode``, …). + + Raises: + FrontendIncompatibleError: If the warp runtime can't run the task + (wrong physics, missing MDP twins, direct task without a warp + implementation). + """ + if cls._detect_workflow(env_cfg) is Workflow.DIRECT: + return cls._build_direct_env(env_cfg, task_id, **construct_kwargs) + return cls._build_manager_env(env_cfg, **construct_kwargs) + + @classmethod + def _build_manager_env(cls, env_cfg: Any, **construct_kwargs: Any) -> gym.Env: + """Construct :class:`ManagerBasedRLEnvWarp`; the env adapts the cfg itself. + + Imported lazily so that ``--frontend=torch`` callers don't pay the + ``isaaclab_experimental.envs`` import cost. The env runs + :meth:`adapt_cfg` in ``__init__``, so stable cfgs and warp-native cfgs + take the same path. + """ + from isaaclab_experimental.envs import ManagerBasedRLEnvWarp + + return ManagerBasedRLEnvWarp(cfg=env_cfg, **construct_kwargs) + + @classmethod + def _build_direct_env(cls, env_cfg: Any, task_id: str, **construct_kwargs: Any) -> gym.Env: + """Construct a direct warp env. + + Direct workflows aren't cfg-adapted: a hand-written warp env class + implements the task. Preferred path: the stable registration declares + it via ``warp_entry_point`` and shares the stable cfg. Fallback: the + task itself is registered under the warp packages (warp-native cfg). + """ + env_class = cls._resolve_direct_warp_class(task_id) + if env_class is None: + # No warp_entry_point: the task itself must be a warp-native registration. + cls._assert_direct_warp_registration(task_id) + return gym.make(task_id, cfg=env_cfg, **construct_kwargs) + # Declared warp class + stable cfg: swap only the env class. + cls._require_newton_physics(env_cfg, type(env_cfg).__name__) + return env_class(cfg=env_cfg, **construct_kwargs) + + # ------------------------------------------------------------------ + # Cfg adaptation (manager-based) + # ------------------------------------------------------------------ + + @classmethod + def adapt_cfg(cls, cfg: Any) -> None: + """Mutate a stable manager-based ``cfg`` in place for the warp managers. + + Idempotent: re-running on an already-adapted cfg is a no-op (the steps + below skip warp-origin symbols and already-promoted entities). + + Three steps, each independently testable: + + 1. :meth:`_require_newton_physics` — hard check that ``cfg.sim.physics`` + is :class:`~isaaclab_newton.physics.NewtonCfg`. The user is + responsible for selecting the Newton variant of the task's + :class:`PresetCfg` via ``presets=newton_mjwarp``; we don't auto-inject. + 2. :meth:`_promote_scene_entity_cfgs` — replace stable + :class:`~isaaclab.managers.SceneEntityCfg` instances under each + term's ``params`` with the warp variant (which adds warp-cached + ``joint_mask``, ``joint_ids_wp``, ``body_ids_wp`` fields). + 3. :meth:`_swap_mdp` — for every MDP term found anywhere in the cfg tree + (discovered by :meth:`_walk_terms` via :class:`ManagerTermBaseCfg` + subclassing, not by hard-coded attribute names), replace any stable + ``func`` *or* ``class_type`` with its same-named warp twin. A missing + twin raises :class:`FrontendIncompatibleError` — partial coverage is + unsafe under the warp managers' kernel-only signature. + """ + label = type(cfg).__name__ + cls._require_newton_physics(cfg, label) + cls._promote_scene_entity_cfgs(cfg) + cls._swap_mdp(cfg, label) + cls._swap_noise_cfgs(cfg, label) + + @staticmethod + def _require_newton_physics(cfg: Any, label: str) -> None: + """Block unless ``cfg.sim.physics`` is :class:`NewtonCfg`. + + The warp managers' assets read state through :class:`NewtonManager`; + a :class:`PhysxCfg` (or unresolved :class:`PresetCfg`) is a hard + incompatibility. The fix is to pass ``presets=newton_mjwarp`` on the CLI + so Hydra resolves the task's :class:`PresetCfg` wrapper to the Newton + field before construction. + """ + from isaaclab_newton.physics import NewtonCfg + + physics = getattr(getattr(cfg, "sim", None), "physics", None) + if isinstance(physics, NewtonCfg): + return + raise FrontendIncompatibleError( + f"warp env {label!r}: expected cfg.sim.physics to be NewtonCfg," + f" got {type(physics).__name__!r}. Pass `presets=newton_mjwarp` on the CLI so" + f" Hydra resolves the task's PresetCfg wrapper to the Newton variant." + ) + + @classmethod + def _promote_scene_entity_cfgs(cls, cfg: Any) -> None: + """Replace stable :class:`SceneEntityCfg` instances with the warp variant. + + Iterates every term cfg in the tree (via :meth:`_walk_terms`) and + rebuilds any stable :class:`SceneEntityCfg` value under ``term.params`` + through :meth:`isaaclab_experimental.managers.SceneEntityCfg.from_stable`. + The warp variant subclasses the stable one, so type checks elsewhere + stay valid; the new fields (``joint_mask`` / ``joint_ids_wp`` / + ``body_ids_wp``) are filled at :meth:`resolve` time by the warp scene. + """ + from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as _WarpSceneEntityCfg + + promoted: list[str] = [] + for path, term in cls._walk_terms(cfg): + if not path or path[0] not in (cls.WARP_MANAGED_GROUPS | cls.OPTIONAL_WARP_GROUPS): + continue # term runs on a stable manager; keep the stable entity + params = getattr(term, "params", None) + if not isinstance(params, dict): + continue + for key, value in list(params.items()): + # Skip non-entities and already-promoted values (idempotence). + if isinstance(value, _WarpSceneEntityCfg) or not isinstance(value, _StableSceneEntityCfg): + continue + params[key] = _WarpSceneEntityCfg.from_stable(value) + promoted.append(f"{'.'.join(path)}.params[{key!r}] ({value.name!r})") + if promoted: + logger.info( + "frontend.warp: promoted %d SceneEntityCfg instance(s) to warp variant:\n %s", + len(promoted), + "\n ".join(promoted), + ) + + @classmethod + def _swap_mdp(cls, cfg: Any, label: str) -> None: + """Replace ``term.func`` and ``term.class_type`` with their warp twins. + + Iterates every term cfg in the tree (via :meth:`_walk_terms`) and on + each term swaps whichever of ``func`` / ``class_type`` is a + stable-origin callable. Twin lookup is name-based, in the order given by + :meth:`_twin_modules`: the warp mirror of the cfg's own task MDP + namespace first (task twins win, even for symbols defined in core + packages), then the mirror of the symbol's defining package (covers + terms borrowed from another task family, e.g. the Ant task reusing + Humanoid MDP terms), then the shared fallback. Any missing twin raises + :class:`FrontendIncompatibleError` listing every affected term — + partial swaps would leave torch-style callables in the cfg and the warp + managers would call them with the wrong signature. + + The warp-side declarations (``out_dim``, ``axes``, ``observation_type``) + that the warp managers need at init are *not* supplied by this swap; + they travel with the warp twin function itself via its own + ``@generic_io_descriptor_warp(out_dim=…)`` decorator. This method only + substitutes the callable; the manager reads the new func's annotations + when it parses the term cfg. + """ + # Highest-priority twin source: warp mirrors of the cfg's own task namespace. + cfg_route_modules = cls._cfg_route_modules(cfg) + module_cache: dict[str, list[ModuleType]] = {} # defining module -> modules to search + searched: set[str] = set() # module names, for the error message + + swapped = 0 + missing: list[tuple[str, str, str]] = [] # (location, attr, symbol) + for path, term in cls._walk_terms(cfg): + if not path or path[0] not in (cls.WARP_MANAGED_GROUPS | cls.OPTIONAL_WARP_GROUPS): + continue # term runs on a stable manager; leave it stable + optional = path[0] in cls.OPTIONAL_WARP_GROUPS + location = ".".join(path) + for attr in ("func", "class_type"): # a term implements via either attr + stable = getattr(term, attr, None) + if stable is None or not cls._is_swap_candidate(stable): + continue + origin = getattr(stable, "__module__", "") or "" + if origin not in module_cache: + module_cache[origin] = cls._twin_modules(origin, cfg_route_modules) + searched.update(m.__name__ for m in module_cache[origin]) + twin = cls._resolve_warp_twin(stable.__name__, module_cache[origin]) + if twin is None: + if not optional: + missing.append((location, attr, stable.__name__)) # collect every miss; report once below + continue # optional group: the stable term stays on the legacy fallback + setattr(term, attr, twin) + swapped += 1 + + if missing: + lines = "\n ".join(f"{loc}.{attr}: no warp twin for {sym!r}" for loc, attr, sym in missing) + raise FrontendIncompatibleError( + f"warp env {label!r}: missing warp MDP twins (searched {sorted(searched)}):\n {lines}" + ) + + logger.info("frontend.warp: swapped %d MDP symbol(s) to warp twins", swapped) + + # ------------------------------------------------------------------ + # Workflow detection and direct-task resolution + # ------------------------------------------------------------------ + + @staticmethod + def _detect_workflow(cfg: Any) -> Workflow: + """Classify the env cfg into manager-based or direct (used to pick the build path). + + Note: + The env cfg roots (ManagerBasedRLEnvCfg, DirectRLEnvCfg, + DirectMARLEnvCfg) do not share a common base class. When a new cfg + root is added, extend the isinstance tuples below. + """ + if isinstance(cfg, ManagerBasedRLEnvCfg): + return Workflow.MANAGER_BASED + if isinstance(cfg, (DirectRLEnvCfg, DirectMARLEnvCfg)): + return Workflow.DIRECT + raise FrontendIncompatibleError( + f"Unrecognised env cfg type {type(cfg).__name__!r};" + f" expected ManagerBasedRLEnvCfg / DirectRLEnvCfg / DirectMARLEnvCfg subclass." + ) + + @staticmethod + def _resolve_direct_warp_class(task_id: str) -> type | None: + """Return the warp env class a stable direct task declares, if any. + + Stable direct registrations may carry a ``warp_entry_point`` kwarg of + the form ``"module.path:ClassName"`` (mirroring ``env_cfg_entry_point``) + naming the warp implementation of the same task. The class is + constructed with the *stable* cfg, so the task needs no parallel warp + registration or duplicated configuration. + """ + try: + spec = gym.spec(task_id) + except gym.error.NameNotFound as exc: + raise FrontendIncompatibleError( + f"--frontend=warp: task {task_id!r} is not registered with gymnasium." + ) from exc + target = (spec.kwargs or {}).get("warp_entry_point") + if not isinstance(target, str): + return None + module_name, _, class_name = target.partition(":") + try: + module = importlib.import_module(module_name) + except ImportError as exc: + raise FrontendIncompatibleError( + f"--frontend=warp: task {task_id!r} declares warp_entry_point {target!r}," + f" but it failed to import: {exc}" + ) from exc + env_class = getattr(module, class_name, None) + if env_class is None: + raise FrontendIncompatibleError( + f"--frontend=warp: task {task_id!r} declares warp_entry_point {target!r}," + f" but {class_name!r} is not defined in {module_name!r}." + ) + return env_class + + @classmethod + def _assert_direct_warp_registration(cls, task_id: str) -> None: + """For direct workflows without a ``warp_entry_point``, the task must be warp-registered.""" + spec = gym.spec(task_id) + ep = spec.entry_point + module = ep if isinstance(ep, str) else (getattr(ep, "__module__", "") or "") + if not module.startswith(cls.WARP_ROOT_PREFIXES): + raise FrontendIncompatibleError( + f"--frontend=warp on direct task {task_id!r}: entry_point {ep!r}" + f" is not under {list(cls.WARP_ROOT_PREFIXES)}. Direct tasks must either" + f" declare a `warp_entry_point` in their stable registration or be" + f" registered as a warp env class." + ) + + # ------------------------------------------------------------------ + # Twin resolution + # ------------------------------------------------------------------ + + @classmethod + def _mirror_module(cls, module: str) -> str | None: + """Mechanical mirror of a stable module path, or ``None`` if not mirrored.""" + if not isinstance(module, str) or not module: + return None + root, _, rest = module.partition(".") # only the top-level package root is swapped + mirror_root = cls.MIRROR_ROOTS.get(root) + if mirror_root is None: + return None + return f"{mirror_root}.{rest}" if rest else mirror_root + + @staticmethod + def _nearest_mdp_module(mirrored: str) -> str | None: + """Deepest existing ``.mdp`` package on a mirrored module path. + + If the path already contains an ``.mdp`` segment, the search starts at + that boundary (``...velocity.mdp.rewards`` resolves to + ``...velocity.mdp``); otherwise the path is walked upward until a + mirrored ``.mdp`` package exists. Existence is probed with + :func:`importlib.util.find_spec`; a prefix that does not exist simply + ends that probe, so resolution is a pure function of the installed + package tree. + """ + parts = mirrored.split(".") + if "mdp" in parts: + parts = parts[: parts.index("mdp")] # start the walk at the .mdp boundary + for depth in range(len(parts), 0, -1): # deepest prefix first + candidate = ".".join([*parts[:depth], "mdp"]) + try: + if importlib.util.find_spec(candidate) is not None: + return candidate + except ModuleNotFoundError: + continue + return None + + @classmethod + def _swap_noise_cfgs(cls, cfg: Any, label: str) -> None: + """Replace stable observation-noise cfgs with their warp-native twins. + + The warp observation manager applies noise through the experimental noise + cfg family (in-place Warp kernels); a stable noise cfg would otherwise be + silently ignored, training against a different MDP. Twins resolve by class + name; data fields are copied while ``func`` keeps the twin's warp default. + Noise cfgs without a warp twin (e.g. class-based ``NoiseModelCfg``) are a + hard error — silently changing the MDP is never acceptable. + """ + import dataclasses + + from isaaclab.utils import noise as stable_noise + + from isaaclab_experimental.utils import noise as warp_noise + + unsupported: list[str] = [] + swapped = 0 + for path, term in cls._walk_terms(cfg): + if not path or path[0] != "observations": + continue + noise_cfg = getattr(term, "noise", None) + if noise_cfg is None: + continue + if type(noise_cfg).__module__.startswith(cls.WARP_ROOT_PREFIXES): + continue # already warp-native + twin_cls = getattr(warp_noise, type(noise_cfg).__name__, None) + if ( + twin_cls is None + or not twin_cls.__module__.startswith(cls.WARP_ROOT_PREFIXES) + or not isinstance(noise_cfg, stable_noise.NoiseCfg) + ): + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__}") + continue + if noise_cfg.func != type(noise_cfg)().func: + # A user-customized noise func has no warp twin; replacing it with the + # twin's default would silently change the MDP. + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__} with custom func {noise_cfg.func!r}") + continue + fields = {f.name: getattr(noise_cfg, f.name) for f in dataclasses.fields(noise_cfg) if f.name != "func"} + if any(isinstance(value, torch.Tensor) for value in fields.values()): + # The warp noise kernels take scalar parameters; tensor-valued ranges + # (per-element noise) have no warp twin yet. + unsupported.append(f"{'.'.join(path)}: {type(noise_cfg).__name__} with tensor-valued parameters") + continue + term.noise = twin_cls(**fields) + swapped += 1 + if unsupported: + raise FrontendIncompatibleError( + f"warp env {label!r}: observation noise cfgs without warp twins:\n " + "\n ".join(unsupported) + ) + if swapped: + logger.info("frontend.warp: swapped %d observation noise cfg(s) to warp twins", swapped) + + @staticmethod + def _import_twin_module(target: str) -> ModuleType: + """Import a resolved twin module; a module that exists but breaks is a hard error.""" + try: + return importlib.import_module(target) + except ImportError as exc: + raise FrontendIncompatibleError(f"warp twin module {target!r} exists but failed to import: {exc}") from exc + + @classmethod + def _cfg_route_modules(cls, cfg: Any) -> list[ModuleType]: + """Warp MDP mirrors of the cfg's class hierarchy, subclass first. + + A stable cfg consumes terms through its task's ``mdp`` namespace, which + re-exports symbols that may be *defined* in core or shared packages. The + warp mirror of that namespace is therefore the primary place to resolve + twins; it is found by mirroring the modules of ``type(cfg).__mro__`` — + the subclass module first (robot-specific cfgs), then base-cfg modules + (task-family bases). + """ + modules: list[ModuleType] = [] + for klass in type(cfg).__mro__: # subclass first: robot cfg wins over task-family base + mirrored = cls._mirror_module(getattr(klass, "__module__", "") or "") + if mirrored is None: + continue + target = cls._nearest_mdp_module(mirrored) + if target is None: + continue + module = cls._import_twin_module(target) + if module not in modules: + modules.append(module) + return modules + + @classmethod + def _twin_modules(cls, symbol_module: str, cfg_route_modules: list[ModuleType]) -> list[ModuleType]: + """Warp modules to consult for a stable symbol's twin, in preference order. + + 1. The warp mirrors of the cfg's own task MDP namespace + (:meth:`_cfg_route_modules`) — task-specific twins win, including + twins for symbols defined in core/shared packages. + 2. The mirror of the symbol's defining package — covers terms borrowed + from another task family and, for core-defined symbols, the shared + :mod:`isaaclab_experimental.envs.mdp` twins (the mirror of + :mod:`isaaclab.envs.mdp`). + """ + modules = list(cfg_route_modules) + # Fallback: the mirror of the package that defines the symbol. + mirrored = cls._mirror_module(symbol_module) + if mirrored is not None: + target = cls._nearest_mdp_module(mirrored) + if target is not None: + module = cls._import_twin_module(target) + if module not in modules: + modules.append(module) + return modules + + @classmethod + def _resolve_warp_twin(cls, name: str, modules: Iterable[ModuleType]) -> Any | None: + """Return the same-named symbol from ``modules`` that originates under the warp packages.""" + for module in modules: + candidate = getattr(module, name, None) + if candidate is None: + continue + origin = getattr(candidate, "__module__", "") or "" + if origin.startswith(cls.WARP_ROOT_PREFIXES): # re-exported stable symbols don't count + return candidate + return None + + @classmethod + def _is_swap_candidate(cls, value: Any) -> bool: + """Heuristic: callable or class whose origin is *not already* warp.""" + if not callable(value): + return False + origin = getattr(value, "__module__", "") or "" + if origin.startswith(cls.WARP_ROOT_PREFIXES): + return False # already a warp twin (idempotent) + return True + + @staticmethod + def _walk_terms(node: Any, path: tuple[str, ...] = ()) -> Iterator[tuple[tuple[str, ...], Any]]: + """Yield ``(path, term)`` for every MDP term cfg in the cfg tree. + + A "term" is a :class:`ManagerTermBaseCfg` (observation/reward/ + termination/event/curriculum) *or* an :class:`ActionTermCfg` / + :class:`CommandTermCfg` — the latter two are separate bases that are + **not** ``ManagerTermBaseCfg`` subclasses, yet carry a swappable + ``class_type``, so they must be matched explicitly. + + Behavior at each node: + + * Match (a term cfg instance): yield ``(path, node)`` and stop — do not + descend into ``term.params`` / ``term.func`` / ``term.class_type``. + * Configclass: don't yield; recurse into every non-underscore *instance* + attribute (``vars(node)``), extending the path. ``observations``, + ``rewards``, ``events``, ``actions``, sub-groups like + ``observations.policy``, and anything nested deeper are reached + transparently. + * Anything else (plain Python data, callables, non-configclass objects): + stop. No yield, no recursion. + + Iterating the instance ``__dict__`` mirrors how the warp managers + consume group cfgs (they iterate ``cfg.__dict__.items()``), so the + walker sees exactly the terms the managers will see — including terms + assigned in ``__post_init__`` — while never descending into methods or + nested class objects, which live on the class rather than the instance. + + Driven entirely by type — no attribute names are hardcoded — so future + cfg layouts (extra observation groups, new nesting, etc.) are picked up + automatically as long as their terms subclass one of the term base cfgs. + """ + from isaaclab.managers.manager_term_cfg import ActionTermCfg, CommandTermCfg, ManagerTermBaseCfg + + if isinstance(node, (ManagerTermBaseCfg, ActionTermCfg, CommandTermCfg)): + yield path, node # a term: yield and stop; never descend into params/func + return + if not hasattr(node, "__dataclass_fields__"): + return # plain data / callables: not a cfg subtree + for name, value in vars(node).items(): # instance attrs — the same view the managers iterate + if name.startswith("_") or value is None: + continue + yield from WarpFrontend._walk_terms(value, path + (name,)) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py deleted file mode 100644 index f792b994f67b..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py +++ /dev/null @@ -1,44 +0,0 @@ -# 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 - -"""Warp-native interactive scene with env_mask support for reset.""" - -from __future__ import annotations - -from collections.abc import Sequence - -import warp as wp - -from isaaclab.scene import InteractiveScene - - -class InteractiveSceneWarp(InteractiveScene): - """Interactive scene with warp-native env_mask support for reset. - - Extends :class:`InteractiveScene` to accept a boolean warp mask for selective resets, - avoiding the need to convert between env_ids and masks. - """ - - def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): - """Reset scene entities using either env_ids or a boolean env_mask. - - Args: - env_ids: The indices of the environments to reset. Defaults to None (all instances). - env_mask: Boolean warp mask of shape (num_envs,). Defaults to None. - """ - # -- assets (support env_mask) - for articulation in self._articulations.values(): - articulation.reset(env_ids, env_mask=env_mask) - for deformable_object in self._deformable_objects.values(): - deformable_object.reset(env_ids) - for rigid_object in self._rigid_objects.values(): - rigid_object.reset(env_ids, env_mask=env_mask) - for surface_gripper in self._surface_grippers.values(): - surface_gripper.reset(env_ids) - for rigid_object_collection in self._rigid_object_collections.values(): - rigid_object_collection.reset(env_ids, env_mask=env_mask) - # -- sensors (no env_mask support) - for sensor in self._sensors.values(): - sensor.reset(env_ids) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py index 2b4791d22a16..e2eb17835d4a 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_env_warp.py @@ -16,11 +16,9 @@ # import builtins import contextlib -import importlib import logging import warnings from collections.abc import Sequence -from copy import deepcopy from typing import Any import torch @@ -30,6 +28,8 @@ from isaaclab.envs.manager_based_env_cfg import ManagerBasedEnvCfg from isaaclab.envs.ui import ViewportCameraController from isaaclab.envs.utils.io_descriptors import export_articulations_data, export_scene_data +from isaaclab.managers import RecorderManager +from isaaclab.scene import InteractiveScene from isaaclab.sim import SimulationContext from isaaclab.sim.utils import use_stage from isaaclab.ui.widgets import ManagerLiveVisualizer @@ -37,9 +37,10 @@ from isaaclab.utils.timer import Timer from isaaclab.utils.version import has_kit -from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene -from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch +from isaaclab_experimental.managers import ActionManager, EventManager, ObservationManager +from isaaclab_experimental.utils.torch_utils import clone_obs_buffer from isaaclab_experimental.utils.warp import resolve_1d_mask +from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache # import logger logger = logging.getLogger(__name__) @@ -52,6 +53,10 @@ def initialize_rng_state( # output state: wp.array(dtype=wp.uint32), ): + """Initialize each env's persistent RNG state as an independent stream of the + + global ``seed`` (one ``wp.rand_init`` stream per env id). + """ env_id = wp.tid() state[env_id] = wp.rand_init(seed, wp.int32(env_id)) @@ -79,11 +84,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg): self.cfg = cfg # initialize internal variables self._is_closed = False - # temporary debug runtime config for manager source/call switching. - cfg_source: dict | str | None = getattr(self.cfg, "manager_call_config", None) - max_modes: dict[str, int] | None = getattr(self.cfg, "manager_call_max_mode", None) - self._manager_call_switch = ManagerCallSwitch(cfg_source, max_modes=max_modes) - self._apply_manager_term_cfg_profile() + # Manager stages register their capture safety with this environment-owned cache. + self._warp_graph_cache = WarpGraphCache(device=self.cfg.sim.device) # set the seed for the environment if self.cfg.seed is not None: @@ -154,6 +156,9 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # Pre-allocated env masks (shared across managers/terms via `env`). self.ALL_ENV_MASK = wp.ones((self.num_envs,), dtype=wp.bool, device=self.device) self.ENV_MASK = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) + # Reset graphs always read this owner-held pointer, while callers may + # provide a different mask object on each reset. + self.reset_mask_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) # Persistent scalar buffer for global env step count (stable pointer for capture). self._global_env_step_count_wp = wp.zeros((1,), dtype=wp.int32, device=self.device) @@ -172,7 +177,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # create event manager # note: this is needed here (rather than after simulation play) to allow USD-related randomization events # that must happen before the simulation starts. Example: randomizing mesh scale - self.event_manager = self._manager_call_switch.resolve_manager_class("EventManager")(self.cfg.events, self) + self.event_manager = EventManager(self.cfg.events, self) # apply USD-related randomization events if "prestartup" in self.event_manager.available_modes: @@ -273,15 +278,13 @@ def device(self): return self.sim.device @property - def env_origins_wp(self) -> wp.array: - """Scene env origins as a warp ``vec3f`` array. Cached on first access.""" - if not hasattr(self, "_env_origins_wp"): - origins = self.scene.env_origins - if isinstance(origins, wp.array): - self._env_origins_wp = origins - else: - self._env_origins_wp = wp.from_torch(origins, dtype=wp.vec3f) - return self._env_origins_wp + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + """Warp view of the scene-owned environment origins [m], shape ``(num_envs,)``. + + The scene hands out a :class:`~isaaclab.utils.warp.ProxyArray`; this boundary + accessor resolves it once so Warp-native consumers see a plain Warp array. + """ + return self.scene.env_origins_pa.warp def resolve_env_mask( self, @@ -360,21 +363,28 @@ def load_managers(self): :meth:`SimulationContext.reset_async` and it isn't possible to call async functions in the constructor. """ + # Reloading managers replaces pointer-stable buffers captured by prior stages. + self._warp_graph_cache.invalidate() # prepare the managers # -- event manager (we print it here to make the logging consistent) print("[INFO] Event Manager: ", self.event_manager) # -- recorder manager - self.recorder_manager = self._manager_call_switch.resolve_manager_class("RecorderManager")( - self.cfg.recorders, self - ) + self.recorder_manager = RecorderManager(self.cfg.recorders, self) + self._has_recorders = bool(self.recorder_manager.active_terms) print("[INFO] Recorder Manager: ", self.recorder_manager) # -- action manager - self.action_manager = self._manager_call_switch.resolve_manager_class("ActionManager")(self.cfg.actions, self) + self.action_manager = ActionManager(self.cfg.actions, self) + # Manager-dependent persistent buffer: reuse the storage across manager + # reloads when the action dimension is unchanged so captured graphs keep + # reading the same pointer. + action_shape = (self.num_envs, self.action_manager.total_action_dim) + if getattr(self, "_action_in_wp", None) is None or tuple(self._action_in_wp.shape) != action_shape: + self._action_in_wp = wp.zeros(action_shape, dtype=wp.float32, device=self.device) + else: + self._action_in_wp.zero_() print("[INFO] Action Manager: ", self.action_manager) # -- observation manager - self.observation_manager = self._manager_call_switch.resolve_manager_class("ObservationManager")( - self.cfg.observations, self - ) + self.observation_manager = ObservationManager(self.cfg.observations, self) print("[INFO] Observation Manager:", self.observation_manager) # perform events at the start of the simulation @@ -396,11 +406,15 @@ def setup_manager_visualizers(self): """ def reset( - self, seed: int | None = None, env_ids: Sequence[int] | None = None, options: dict[str, Any] | None = None + self, + seed: int | None = None, + env_ids: Sequence[int] | None = None, + options: dict[str, Any] | None = None, + env_mask: wp.array | torch.Tensor | None = None, ) -> tuple[VecEnvObs, dict]: """Resets the specified environments and returns observations. - This function calls the :meth:`_reset_idx` function to reset the specified environments. + This function calls the :meth:`_reset_mask` function to reset the specified environments. However, certain operations, such as procedural terrain generation, that happened during initialization are not repeated. @@ -412,14 +426,19 @@ def reset( Note: This argument is used for compatibility with Gymnasium environment definition. + env_mask: Boolean environment mask. This is the canonical Warp + frontend selection; :paramref:`env_ids` remains supported for + API compatibility. + Returns: A tuple containing the observations and extras. """ - if env_ids is None: - env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) - - # trigger recorder terms for pre-reset calls - self.recorder_manager.record_pre_reset(env_ids) + reset_mask = self.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + if self._has_recorders: + recorder_env_ids = env_ids + if recorder_env_ids is None or isinstance(recorder_env_ids, wp.array): + recorder_env_ids = wp.to_torch(reset_mask).nonzero(as_tuple=False).squeeze(-1) + self.recorder_manager.record_pre_reset(recorder_env_ids) # set the seed if seed is not None: @@ -434,8 +453,9 @@ def reset( device=self.device, ) - # reset state of scene - self._reset_idx(env_ids) + self._reset_mask(env_mask=reset_mask) + if self._has_recorders: + self.extras["log"].update(self.recorder_manager.reset(recorder_env_ids)) # update articulation kinematics self.scene.write_data_to_sim() @@ -445,11 +465,17 @@ def reset( for _ in range(self.cfg.num_rerenders_on_reset): self.sim.render() - # trigger recorder terms for post-reset calls - self.recorder_manager.record_post_reset(env_ids) + if self._has_recorders: + self.recorder_manager.record_post_reset(recorder_env_ids) # compute observations - self.obs_buf = self.observation_manager.compute(update_history=True) + self.obs_buf = self._warp_graph_cache.call( + "ObservationManager_compute_reset", + self.observation_manager.compute, + update_history=True, + return_cloned_output=False, + output=clone_obs_buffer, + ) # return observations return self.obs_buf, self.extras @@ -460,6 +486,7 @@ def reset_to( env_ids: Sequence[int] | None, seed: int | None = None, is_relative: bool = False, + env_mask: wp.array | None = None, ): """Resets specified environments to provided states. @@ -477,19 +504,27 @@ def reset_to( seed: The seed to use for randomization. Defaults to None, in which case the seed is not set. is_relative: If set to True, the state is considered relative to the environment origins. Defaults to False. + env_mask: Boolean Warp mask selecting environments. Takes precedence over + :paramref:`env_ids`. State setting and recorders are host consumers, so + compact IDs are still materialized once inside this method. """ # reset all envs in the scene if env_ids is None - if env_ids is None: + if env_mask is not None: + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + elif env_ids is None: env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device) - # trigger recorder terms for pre-reset calls - self.recorder_manager.record_pre_reset(env_ids) + if self._has_recorders: + self.recorder_manager.record_pre_reset(env_ids) # set the seed if seed is not None: self.seed(seed) - self._reset_idx(env_ids) + env_mask = self.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + self._reset_mask(env_mask=env_mask) + if self._has_recorders: + self.extras["log"].update(self.recorder_manager.reset(env_ids)) # set the state self.scene.reset_to(state, env_ids, is_relative=is_relative) @@ -502,11 +537,17 @@ def reset_to( for _ in range(self.cfg.num_rerenders_on_reset): self.sim.render() - # trigger recorder terms for post-reset calls - self.recorder_manager.record_post_reset(env_ids) + if self._has_recorders: + self.recorder_manager.record_post_reset(env_ids) # compute observations - self.obs_buf = self.observation_manager.compute(update_history=True) + self.obs_buf = self._warp_graph_cache.call( + "ObservationManager_compute_reset", + self.observation_manager.compute, + update_history=True, + return_cloned_output=False, + output=clone_obs_buffer, + ) # return observations return self.obs_buf, self.extras @@ -527,15 +568,16 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: A tuple containing the observations and extras. """ # process actions - action_device = action.to(self.device) - if action_device.dtype != torch.float32: - action_device = action_device.float() - if not action_device.is_contiguous(): - action_device = action_device.contiguous() - action_wp = wp.from_torch(action_device, dtype=wp.float32) - self.action_manager.process_action(action_wp) + action_device = action.to(device=self.device, dtype=torch.float32).contiguous() + wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32)) + self._warp_graph_cache.call( + "ActionManager_process_action", + self.action_manager.process_action, + action=self._action_in_wp, + ) - self.recorder_manager.record_pre_step() + if self._has_recorders: + self.recorder_manager.record_pre_step() # check if we need to do rendering within the physics loop # note: hoisted out of the decimation loop; is_rendering does live settings lookups @@ -545,7 +587,7 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: for _ in range(self.cfg.decimation): self._sim_step_counter += 1 # set actions into buffers - self.action_manager.apply_action() + self._warp_graph_cache.call("ActionManager_apply_action", self.action_manager.apply_action) # set actions into simulator self.scene.write_data_to_sim() # simulate @@ -560,11 +602,23 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: # post-step: step interval event if "interval" in self.event_manager.available_modes: - self.event_manager.apply(mode="interval", dt=self.step_dt) + self._warp_graph_cache.call( + "EventManager_apply_interval", + self.event_manager.apply, + mode="interval", + dt=self.step_dt, + ) # -- compute observations - self.obs_buf = self.observation_manager.compute(update_history=True) - self.recorder_manager.record_post_step() + self.obs_buf = self._warp_graph_cache.call( + "ObservationManager_compute_update_history", + self.observation_manager.compute, + update_history=True, + return_cloned_output=False, + output=clone_obs_buffer, + ) + if self._has_recorders: + self.recorder_manager.record_post_step() # return observations and extras return self.obs_buf, self.extras @@ -613,127 +667,47 @@ def close(self): Helper functions. """ - def _resolve_stable_cfg_counterpart(self) -> ManagerBasedEnvCfg | None: - """Resolve a stable task config counterpart for the current experimental task config. - - The lookup follows a module-name mirror convention: - ``isaaclab_tasks_experimental...`` -> ``isaaclab_tasks...`` with the same config class name. - """ - cfg_cls = self.cfg.__class__ - cfg_module_name = cfg_cls.__module__ - if "isaaclab_tasks_experimental" not in cfg_module_name: - return None - - stable_module_name = cfg_module_name.replace("isaaclab_tasks_experimental", "isaaclab_tasks", 1) - try: - stable_module = importlib.import_module(stable_module_name) - except Exception as exc: - logger.warning( - "Failed to import stable task cfg module '%s' for manager_call_config stable mode: %s", - stable_module_name, - exc, - ) - return None - - stable_cfg_cls = getattr(stable_module, cfg_cls.__name__, None) - if stable_cfg_cls is None: - logger.warning( - "Stable task cfg class '%s' not found in module '%s'.", - cfg_cls.__name__, - stable_module_name, - ) - return None - - try: - return stable_cfg_cls() - except Exception as exc: - logger.warning( - "Failed to instantiate stable task cfg '%s.%s': %s", - stable_module_name, - cfg_cls.__name__, - exc, - ) - return None - - def _apply_manager_term_cfg_profile(self) -> None: - """Align term configs with manager modes for stable manager selections. - - When a manager is configured as STABLE (0), swap its corresponding config subtree - from the stable task counterpart to keep manager-term type/signature compatibility. - """ - manager_to_cfg_attr = { - "ActionManager": "actions", - "ObservationManager": "observations", - "EventManager": "events", - "RecorderManager": "recorders", - "CommandManager": "commands", - "TerminationManager": "terminations", - "RewardManager": "rewards", - "CurriculumManager": "curriculum", - } - - stable_manager_names = [ - manager_name - for manager_name in manager_to_cfg_attr - if self._manager_call_switch.get_mode_for_manager(manager_name) == ManagerCallMode.STABLE - ] - if not stable_manager_names: - return - - stable_cfg = self._resolve_stable_cfg_counterpart() - if stable_cfg is None: - logger.warning( - "Stable managers requested (%s), but no stable cfg counterpart could be resolved." - " Keeping experimental term configs.", - ", ".join(stable_manager_names), - ) - return - - replaced_items: list[str] = [] - for manager_name, cfg_attr in manager_to_cfg_attr.items(): - if self._manager_call_switch.get_mode_for_manager(manager_name) != ManagerCallMode.STABLE: - continue - if not hasattr(self.cfg, cfg_attr) or not hasattr(stable_cfg, cfg_attr): - continue - setattr(self.cfg, cfg_attr, deepcopy(getattr(stable_cfg, cfg_attr))) - replaced_items.append(f"{manager_name} -> cfg.{cfg_attr}") - - if replaced_items: - print("[INFO] Applied stable term config profile for managers:") - for item in replaced_items: - print(f" - {item}") - - def _reset_idx(self, env_ids: Sequence[int]): - """Reset environments based on specified indices. + def _reset_mask( + self, + *, + env_mask: wp.array(dtype=wp.bool), + ) -> None: + """Reset Warp-owned state for selected environments. Args: - env_ids: List of environment ids which must be reset + env_mask: Boolean Warp mask selecting environments to reset. """ + if env_mask is not self.reset_mask_wp: + wp.copy(self.reset_mask_wp, env_mask) + env_mask = self.reset_mask_wp + # reset the internal buffers of the scene elements - self.scene.reset(env_ids) + self.scene.reset(env_mask=env_mask) # apply events such as randomization for environments that need a reset if "reset" in self.event_manager.available_modes: env_step_count = self._sim_step_counter // self.cfg.decimation self._global_env_step_count_wp.fill_(env_step_count) - self.event_manager.apply( - mode="reset", env_ids=env_ids, global_env_step_count=self._global_env_step_count_wp + self._warp_graph_cache.call( + "EventManager_apply_reset", + self.event_manager.apply, + mode="reset", + env_mask_wp=env_mask, + global_env_step_count=self._global_env_step_count_wp, ) # iterate over all managers and reset them # this returns a dictionary of information which is stored in the extras # note: This is order-sensitive! Certain things need be reset before others. self.extras["log"] = dict() - env_mask = self.resolve_env_mask(env_ids=env_ids) # -- observation manager - info = self.observation_manager.reset(env_mask=env_mask) + info = self._warp_graph_cache.call( + "ObservationManager_reset", self.observation_manager.reset, env_mask=env_mask + ) self.extras["log"].update(info) # -- action manager - info = self.action_manager.reset(env_mask=env_mask) + info = self._warp_graph_cache.call("ActionManager_reset", self.action_manager.reset, env_mask=env_mask) self.extras["log"].update(info) # -- event manager - info = self.event_manager.reset(env_mask=env_mask) - self.extras["log"].update(info) - # -- recorder manager - info = self.recorder_manager.reset(env_ids) + info = self._warp_graph_cache.call("EventManager_reset", self.event_manager.reset, env_mask=env_mask) self.extras["log"].update(info) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 2e1894e54003..85390a074577 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py @@ -15,7 +15,6 @@ import math import os -from collections.abc import Sequence from typing import Any, ClassVar import gymnasium as gym @@ -25,12 +24,12 @@ from isaaclab.envs.common import VecEnvStepReturn from isaaclab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg -from isaaclab.managers import CommandManager from isaaclab.ui.widgets import ManagerLiveVisualizer from isaaclab.utils.timer import Timer -from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode +from isaaclab_experimental.managers import CommandManager, CurriculumManager, RewardManager, TerminationManager from isaaclab_experimental.utils.torch_utils import clone_obs_buffer +from isaaclab_experimental.utils.warp import any_env_set, increment_all_int64, zero_masked_int64 from .manager_based_env_warp import ManagerBasedEnvWarp @@ -93,6 +92,14 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, ** render_mode: The render mode for the environment. Defaults to None, which is similar to ``"human"``. """ + # Adapt the cfg for the warp managers (Newton physics check, SceneEntityCfg + # promotion, MDP twin swap). Idempotent: a warp-native cfg passes through + # unchanged, and a stable-derived cfg (``--frontend=warp`` or a registered + # warp task variant subclassing a stable cfg) is adapted in place. + from isaaclab_experimental.envs.frontend import WarpFrontend + + WarpFrontend.adapt_cfg(cfg) + # -- counter for curriculum self.common_step_counter = 0 @@ -106,16 +113,6 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, ** # store the render mode self.render_mode = render_mode - # The persistent reset mask needed for warp capture - # The intended use is to copy into this mask whenever capture is needed - # TODO: termination manager provides the same mask, investigate whether this can be replaced. - self.reset_mask_wp = wp.zeros(cfg.scene.num_envs, dtype=wp.bool, device=cfg.sim.device) - - # Persistent action input buffer to keep pointer stable for captured graphs. - self._action_in_wp: wp.array = wp.zeros( - (self.num_envs, self.action_manager.total_action_dim), dtype=wp.float32, device=self.device - ) - # initialize data and constants # -- set the framerate of the gym video recorder wrapper so that the playback speed # of the produced video matches the simulation @@ -155,7 +152,7 @@ def max_episode_length(self) -> int: def load_managers(self): # note: this order is important since observation manager needs to know the command and action managers # and the reward manager needs to know the termination manager - # -- command manager (stable impl — not routed through ManagerCallSwitch) + # -- Warp-first command manager self.command_manager = CommandManager(self.cfg.commands, self) print("[INFO] Command Manager: ", self.command_manager) @@ -164,17 +161,13 @@ def load_managers(self): # prepare the managers # -- termination manager - self.termination_manager = self._manager_call_switch.resolve_manager_class("TerminationManager")( - self.cfg.terminations, self - ) + self.termination_manager = TerminationManager(self.cfg.terminations, self) print("[INFO] Termination Manager: ", self.termination_manager) - # -- reward manager (experimental fork; Warp-compatible rewards) - self.reward_manager = self._manager_call_switch.resolve_manager_class("RewardManager")(self.cfg.rewards, self) + # -- reward manager + self.reward_manager = RewardManager(self.cfg.rewards, self) print("[INFO] Reward Manager: ", self.reward_manager) - # -- curriculum manager - self.curriculum_manager = self._manager_call_switch.resolve_manager_class("CurriculumManager")( - self.cfg.curriculum, self - ) + # -- Warp-first curriculum manager + self.curriculum_manager = CurriculumManager(self.cfg.curriculum, self) print("[INFO] Curriculum Manager: ", self.curriculum_manager) # setup the action and observation spaces for Gym @@ -205,7 +198,7 @@ def invalidate_wp_graphs(self) -> None: Call this if the captured launch topology changes (e.g. different term list, shapes, etc.). """ - self._manager_call_switch.invalidate_graphs() + self._warp_graph_cache.invalidate() def step_warp_termination_compute(self) -> None: """Captured stage: compute terminations (env-step frequency).""" @@ -238,18 +231,18 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # IMPORTANT: Do NOT re-wrap/replace the `wp.array` used by captured graphs each step. # Instead, copy the latest actions into the persistent buffer. with Timer(name="action_preprocess", msg="Action preprocessing took:", enable=DEBUG_TIMER_STEP, time_unit="us"): - if self._action_in_wp is None: - raise RuntimeError("Action buffer not initialized. Call reset() before step().") - action_device = action.to(self.device) + action_device = action.to(device=self.device, dtype=torch.float32).contiguous() wp.copy(self._action_in_wp, wp.from_torch(action_device, dtype=wp.float32)) - self._manager_call_switch.call_stage( - stage="ActionManager_process_action", - warp_call={"fn": self.action_manager.process_action, "kwargs": {"action": self._action_in_wp}}, + self._warp_graph_cache.call( + "ActionManager_process_action", + self.action_manager.process_action, + action=self._action_in_wp, timer=DEBUG_TIMER_STEP, ) - self.recorder_manager.record_pre_step() + if self._has_recorders: + self.recorder_manager.record_pre_step() # check if we need to do rendering within the physics loop # note: checked here once to avoid multiple checks within the loop @@ -259,21 +252,24 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: for _ in range(self.cfg.decimation): self._sim_step_counter += 1 # set actions into buffers - self._manager_call_switch.call_stage( - stage="ActionManager_apply_action", - warp_call={"fn": self.action_manager.apply_action}, - timer=DEBUG_TIMER_STEP, - ) - self._manager_call_switch.call_stage( - stage="Scene_write_data_to_sim", - warp_call={"fn": self.scene.write_data_to_sim}, + self._warp_graph_cache.call( + "ActionManager_apply_action", + self.action_manager.apply_action, timer=DEBUG_TIMER_STEP, ) + with Timer( + name="Scene_write_data_to_sim", + msg="Scene write took:", + enable=DEBUG_TIMER_STEP, + time_unit="us", + ): + self.scene.write_data_to_sim() # simulate with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMER_STEP, time_unit="us"): self.sim.step(render=False) - self.recorder_manager.record_post_physics_decimation_step() + if self._has_recorders: + self.recorder_manager.record_post_physics_decimation_step() # render between steps only if the GUI or an RTX sensor needs it # note: we assume the render interval to be the shortest accepted rendering interval. # If a camera needs rendering at a faster frequency, this will lead to unexpected behavior. @@ -290,89 +286,65 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # post-step: # -- update env counters (used for curriculum generation) - self.episode_length_buf += 1 # step in current episode (per env) + wp.launch( + increment_all_int64, + dim=self.num_envs, + inputs=[self._episode_length_buf_wp, 1], + device=self.device, + ) self.common_step_counter += 1 # total step (common for all envs) # -- post-processing (termination + reward) as independently configurable stages - self._manager_call_switch.call_stage( - stage="TerminationManager_compute", - warp_call={"fn": self.step_warp_termination_compute}, + self._warp_graph_cache.call( + "TerminationManager_compute", + self.step_warp_termination_compute, timer=DEBUG_TIMER_STEP, ) - self.reward_buf = self._manager_call_switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": self.reward_manager.compute, "kwargs": {"dt": float(self.step_dt)}}, + self.reward_buf = self._warp_graph_cache.call( + "RewardManager_compute", + self.reward_manager.compute, + dt=self.step_dt, timer=DEBUG_TIMER_STEP, ) - if len(self.recorder_manager.active_terms) > 0: + if self._has_recorders: # update observations for recording if needed - self._manager_call_switch.call_stage( - stage="ObservationManager_compute_no_history", - warp_call={"fn": self.observation_manager.compute, "kwargs": {"return_cloned_output": False}}, + self._warp_graph_cache.call( + "ObservationManager_compute_no_history", + self.observation_manager.compute, + return_cloned_output=False, timer=DEBUG_TIMER_STEP, ) self.recorder_manager.record_post_step() - # -- reset envs that terminated/timed-out and log the episode information - # NOTE: Interim path (intentional). - # We still compact `reset_buf` into `env_ids` here because several reset-time managers/recorders - # are still `env_ids`-based. Do NOT remove/replace this until mask-based reset is end-to-end. - with Timer( - name="reset_selection", - msg="Reset selection took:", - enable=DEBUG_TIMER_STEP, - time_unit="us", - ): - # Keep the reset-mask handoff fully in Warp when experimental termination buffers exist. - # Stable termination manager path exposes torch-only dones/reset buffers. - termination_manager_mode = self._manager_call_switch.get_mode_for_manager("TerminationManager") - if termination_manager_mode == ManagerCallMode.STABLE: - # copy still needed as mask will be used if manager is set to mode > 0 - wp.copy(self.reset_mask_wp, wp.from_torch(self.reset_buf, dtype=wp.bool)) - else: - wp.copy(self.reset_mask_wp, self.termination_manager.dones_wp) - reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) - if len(reset_env_ids) > 0: - # trigger recorder terms for pre-reset calls - self.recorder_manager.record_pre_reset(reset_env_ids) - - with Timer( - name="reset_idx", - msg="Reset idx took:", - enable=DEBUG_TIMER_STEP, - time_unit="us", - ): - self._reset_idx(env_ids=reset_env_ids, env_mask=self.reset_mask_wp) - - # if sensors are added to the scene, make sure we render to reflect changes in reset - if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0: - for _ in range(self.cfg.num_rerenders_on_reset): - self.sim.render() - - # trigger recorder terms for post-reset calls - self.recorder_manager.record_post_reset(reset_env_ids) + self._reset_terminated_envs() # -- update command - self.command_manager.compute(dt=float(self.step_dt)) + self._warp_graph_cache.call( + "CommandManager_compute", + self.command_manager.compute, + dt=self.step_dt, + timer=DEBUG_TIMER_STEP, + ) # -- step interval events if "interval" in self.event_manager.available_modes: - self._manager_call_switch.call_stage( - stage="EventManager_apply_interval", - warp_call={"fn": self.event_manager.apply, "kwargs": {"mode": "interval", "dt": float(self.step_dt)}}, + self._warp_graph_cache.call( + "EventManager_apply_interval", + self.event_manager.apply, + mode="interval", + dt=self.step_dt, timer=DEBUG_TIMER_STEP, ) # -- compute observations # note: done after reset to get the correct observations for reset envs - self.obs_buf = self._manager_call_switch.call_stage( - stage="ObservationManager_compute_update_history", - warp_call={ - "fn": self.observation_manager.compute, - "kwargs": {"update_history": True, "return_cloned_output": False}, - "output": lambda r: clone_obs_buffer(r), - }, + self.obs_buf = self._warp_graph_cache.call( + "ObservationManager_compute_update_history", + self.observation_manager.compute, + update_history=True, + return_cloned_output=False, + output=clone_obs_buffer, timer=DEBUG_TIMER_STEP, ) # return observations, rewards, resets and extras @@ -408,10 +380,8 @@ def render(self, recompute: bool = False) -> np.ndarray | None: if self.render_mode == "human" or self.render_mode is None: return None elif self.render_mode == "rgb_array": - # check that if any render could have happened - has_gui = bool(self.sim.get_setting("/isaaclab/has_gui")) - offscreen_render = bool(self.sim.get_setting("/isaaclab/render/offscreen")) - if not (has_gui or offscreen_render): + # rendering requires a GUI or offscreen rendering (mirrors the stable env) + if not (self.sim.has_gui or self.sim.has_offscreen_render): raise RuntimeError( f"Cannot render '{self.render_mode}' when the simulation render mode does not support" " rendering. Please set the simulation render mode to 'PARTIAL_RENDERING' or" @@ -485,68 +455,49 @@ def _configure_gym_env_spaces(self): self.observation_space = gym.vector.utils.batch_space(self.single_observation_space, self.num_envs) self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs) - def _reset_idx( + def _reset_mask( self, - env_ids: Sequence[int] | slice | torch.Tensor, *, - env_mask: wp.array | None = None, - ): - """Reset environments based on specified indices. - - IMPORTANT: - This function always uses the **TerminationManager-produced Warp env mask** (`self.reset_buf`) to select - which envs to reset. The ids/mask conversion is performed in `step()` before calling this function. - - In other words: - - If `env_mask` is provided, it **must** be `self.reset_buf` (Warp bool mask) - - If `env_mask` is not provided, this function will populate `self.reset_buf` from `env_ids` - - When `env_mask` is provided, `env_ids` **must** correspond to the same mask + env_mask: wp.array(dtype=wp.bool), + env_ids: torch.Tensor | None = None, + ) -> None: + """Reset Warp-owned RL state for selected environments. Args: - env_ids: Environment indices to reset. - env_mask: Warp boolean env mask selecting envs to reset. Must be `self.reset_buf`. - If None, uses and populates `self.reset_buf` from `env_ids`. + env_mask: Boolean Warp mask selecting environments to reset. + env_ids: Compact environment IDs matching :paramref:`env_mask`, when a + host consumer (e.g. the recorder boundary) already materialized them. """ - if env_mask is None: - # Base `reset()` / `reset_to()` call-path provides only `env_ids`. - # Populate the stable TerminationManager-owned mask (`self.reset_buf`) from ids. - env_mask = self.reset_mask_wp - # Use the centralized env-id/mask resolution from the base Warp env, then copy into the - # stable TerminationManager-owned buffer (`self.reset_buf`) used by captured graphs. - resolved_mask = self.resolve_env_mask(env_ids=env_ids) - wp.copy(env_mask, resolved_mask) - - if not isinstance(env_mask, wp.array): - raise TypeError(f"env_mask must be a wp.array (got {type(env_mask)}).") - - # update the curriculum for environments that need a reset - with Timer( - name="curriculum_manager.compute_reset", - msg="CurriculumManager.compute (reset) took:", - enable=DEBUG_TIMER_RESET, - time_unit="us", - ): - self.curriculum_manager.compute(env_ids=env_ids) - - # reset the internal buffers of the scene elements - self._manager_call_switch.call_stage( - stage="Scene_reset", - warp_call={"fn": self.scene.reset, "kwargs": {"env_ids": env_ids, "env_mask": env_mask}}, + if env_mask is not self.reset_mask_wp: + wp.copy(self.reset_mask_wp, env_mask) + env_mask = self.reset_mask_wp + + # Legacy curriculum terms consume compact IDs. Materialize them at most + # once per reset and share the recorder boundary's IDs when available. + curriculum_kwargs: dict[str, Any] = {"env_mask": env_mask} + if self.curriculum_manager.requires_host_ids: + if env_ids is None: + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + curriculum_kwargs["env_ids"] = env_ids + + self._warp_graph_cache.call( + "CurriculumManager_compute", + self.curriculum_manager.compute, timer=DEBUG_TIMER_RESET, + **curriculum_kwargs, ) + with Timer(name="Scene_reset", msg="Scene reset took:", enable=DEBUG_TIMER_RESET, time_unit="us"): + self.scene.reset(env_mask=env_mask) + if "reset" in self.event_manager.available_modes: self._global_env_step_count_wp.fill_(self._sim_step_counter // self.cfg.decimation) - self._manager_call_switch.call_stage( - stage="EventManager_apply_reset", - warp_call={ - "fn": self.event_manager.apply, - "kwargs": { - "mode": "reset", - "env_mask_wp": env_mask, - "global_env_step_count": self._global_env_step_count_wp, - }, - }, + self._warp_graph_cache.call( + "EventManager_apply_reset", + self.event_manager.apply, + mode="reset", + env_mask_wp=env_mask, + global_env_step_count=self._global_env_step_count_wp, timer=DEBUG_TIMER_RESET, ) @@ -554,51 +505,58 @@ def _reset_idx( # this returns a dictionary of information which is stored in the extras # note: This is order-sensitive! Certain things need be reset before others. # -- observation manager + action + reward managers - obs_info = self._manager_call_switch.call_stage( - stage="ObservationManager_reset", - warp_call={"fn": self.observation_manager.reset, "kwargs": {"env_mask": env_mask}}, + obs_info = self._warp_graph_cache.call( + "ObservationManager_reset", + self.observation_manager.reset, + env_mask=env_mask, timer=DEBUG_TIMER_RESET, ) - action_info = self._manager_call_switch.call_stage( - stage="ActionManager_reset", - warp_call={"fn": self.action_manager.reset, "kwargs": {"env_mask": env_mask}}, + action_info = self._warp_graph_cache.call( + "ActionManager_reset", + self.action_manager.reset, + env_mask=env_mask, timer=DEBUG_TIMER_RESET, ) - reward_info = self._manager_call_switch.call_stage( - stage="RewardManager_reset", - warp_call={"fn": self.reward_manager.reset, "kwargs": {"env_mask": env_mask}}, + reward_info = self._warp_graph_cache.call( + "RewardManager_reset", + self.reward_manager.reset, + env_mask=env_mask, timer=DEBUG_TIMER_RESET, ) - - # -- curriculum manager - with Timer( - name="curriculum_manager.reset", - msg="CurriculumManager.reset took:", - enable=DEBUG_TIMER_RESET, - time_unit="us", - ): - curriculum_info = self.curriculum_manager.reset(env_ids=env_ids) + curriculum_info = self._warp_graph_cache.call( + "CurriculumManager_reset", + self.curriculum_manager.reset, + timer=DEBUG_TIMER_RESET, + **curriculum_kwargs, + ) # -- command + event + termination managers - command_info = self.command_manager.reset(env_ids=env_ids) - event_info = self._manager_call_switch.call_stage( - stage="EventManager_reset", - warp_call={"fn": self.event_manager.reset, "kwargs": {"env_mask": env_mask}}, - stable_call={"fn": self.event_manager.reset, "kwargs": {"env_ids": env_ids}}, + command_info = self._warp_graph_cache.call( + "CommandManager_reset", + self.command_manager.reset, + env_mask=env_mask, timer=DEBUG_TIMER_RESET, ) - termination_info = self._manager_call_switch.call_stage( - stage="TerminationManager_reset", - warp_call={"fn": self.termination_manager.reset, "kwargs": {"env_mask": env_mask}}, - stable_call={"fn": self.termination_manager.reset, "kwargs": {"env_ids": env_ids}}, + event_info = self._warp_graph_cache.call( + "EventManager_reset", + self.event_manager.reset, + env_mask=env_mask, + timer=DEBUG_TIMER_RESET, + ) + termination_info = self._warp_graph_cache.call( + "TerminationManager_reset", + self.termination_manager.reset, + env_mask=env_mask, timer=DEBUG_TIMER_RESET, ) - - # -- recorder manager - recorder_info = self.recorder_manager.reset(env_ids=env_ids) # reset the episode length buffer - self.episode_length_buf[env_ids] = 0 + wp.launch( + zero_masked_int64, + dim=self.num_envs, + inputs=[env_mask, self._episode_length_buf_wp], + device=self.device, + ) # aggregate logging info log: dict[str, Any] = {} @@ -610,7 +568,45 @@ def _reset_idx( command_info, event_info, termination_info, - recorder_info, ): log.update(info) self.extras["log"] = log + + def _reset_terminated_envs(self) -> None: + """Reset terminated environments using the canonical Warp mask.""" + reset_mask = self.termination_manager.dones_wp + # Keep the mask as the canonical selection, but use one host predicate + # to avoid dispatching the complete reset pipeline when it is empty. + if not any_env_set(self.reset_buf): + return + + # Same-step autoreset exposes terminal observations before any selected + # environment is reset, matching the stable manager-based environment. + if self.cfg.compute_final_obs: + self.extras["final_obs"] = self.observation_manager.compute() + + recorder_env_ids = None + if self._has_recorders: + with Timer( + name="reset_selection_host", + msg="Recorder reset selection took:", + enable=DEBUG_TIMER_STEP, + time_unit="us", + ): + recorder_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) + self.recorder_manager.record_pre_reset(recorder_env_ids) + + with Timer( + name="reset_mask", + msg="Reset mask took:", + enable=DEBUG_TIMER_STEP, + time_unit="us", + ): + self._reset_mask(env_mask=reset_mask, env_ids=recorder_env_ids) + + if self._has_recorders: + self.extras["log"].update(self.recorder_manager.reset(recorder_env_ids)) + self.recorder_manager.record_post_reset(recorder_env_ids) + if self.has_rtx_sensors and self.cfg.num_rerenders_on_reset > 0: + for _ in range(self.cfg.num_rerenders_on_reset): + self.sim.render() diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi index a6dadcfce5bc..9459feae7a18 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/__init__.pyi @@ -19,6 +19,14 @@ from .actions import ( # noqa: F401 JointPositionAction, JointPositionActionCfg, ) +from .commands import ( # noqa: F401 + NullCommand, + NullCommandCfg, + UniformPoseCommand, + UniformPoseCommandCfg, + UniformVelocityCommand, + UniformVelocityCommandCfg, +) # Override stable terms with experimental Warp-first implementations. These leaf # modules are import-clean (no eager backend imports), so re-exporting them here diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py similarity index 67% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py rename to source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py index 3e2b7945ebde..57c6a813798d 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/__init__.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.py @@ -3,8 +3,8 @@ # # SPDX-License-Identifier: BSD-3-Clause -""" -Direct workflow environments. -""" +"""Warp-native command terms.""" -import gymnasium as gym +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi new file mode 100644 index 000000000000..dcc7b0b99b82 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/__init__.pyi @@ -0,0 +1,18 @@ +# 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 + +__all__ = [ + "NullCommandCfg", + "UniformPoseCommandCfg", + "UniformVelocityCommandCfg", + "NullCommand", + "UniformPoseCommand", + "UniformVelocityCommand", +] + +from .commands_cfg import NullCommandCfg, UniformPoseCommandCfg, UniformVelocityCommandCfg +from .null_command import NullCommand +from .pose_command import UniformPoseCommand +from .velocity_command import UniformVelocityCommand diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py new file mode 100644 index 000000000000..483f62cbd65c --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/commands_cfg.py @@ -0,0 +1,41 @@ +# 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 + +"""Configuration classes for Warp-native command terms.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.envs.mdp.commands.commands_cfg import NullCommandCfg as _NullCommandCfg +from isaaclab.envs.mdp.commands.commands_cfg import UniformPoseCommandCfg as _UniformPoseCommandCfg +from isaaclab.envs.mdp.commands.commands_cfg import UniformVelocityCommandCfg as _UniformVelocityCommandCfg +from isaaclab.utils.configclass import configclass + +if TYPE_CHECKING: + from .null_command import NullCommand + from .pose_command import UniformPoseCommand + from .velocity_command import UniformVelocityCommand + + +@configclass +class NullCommandCfg(_NullCommandCfg): + """Configuration for the Warp-native null command generator.""" + + class_type: type[NullCommand] | str = "{DIR}.null_command:NullCommand" + + +@configclass +class UniformVelocityCommandCfg(_UniformVelocityCommandCfg): + """Configuration for the Warp-native uniform velocity command generator.""" + + class_type: type[UniformVelocityCommand] | str = "{DIR}.velocity_command:UniformVelocityCommand" + + +@configclass +class UniformPoseCommandCfg(_UniformPoseCommandCfg): + """Configuration for the Warp-native uniform pose command generator.""" + + class_type: type[UniformPoseCommand] | str = "{DIR}.pose_command:UniformPoseCommand" diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py new file mode 100644 index 000000000000..b210daeefb55 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/null_command.py @@ -0,0 +1,62 @@ +# 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 + +"""Warp-native command term that does nothing.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab_experimental.managers import CommandTerm + +if TYPE_CHECKING: + from .commands_cfg import NullCommandCfg + + +class NullCommand(CommandTerm): + """Command generator that does nothing. + + Warp-native twin of :class:`isaaclab.envs.mdp.commands.NullCommand`. It does not + generate any commands and is used for environments that do not require one. All + inherited bookkeeping runs as device-side kernels on empty metrics with an + infinite resampling time, so the term is trivially graph-capturable. + """ + + cfg: NullCommandCfg + """Configuration for the command generator.""" + + def __str__(self) -> str: + msg = "NullCommand:\n" + msg += "\tCommand dimension: N/A\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}" + return msg + + """ + Properties + """ + + @property + def command(self): + """Null command. + + Raises: + RuntimeError: No command is generated. Always raises this error. + """ + raise RuntimeError("NullCommandTerm does not generate any commands.") + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + pass + + def _resample_command(self, env_mask: wp.array): + pass + + def _update_command(self): + pass diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py new file mode 100644 index 000000000000..bb3944ad07ea --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/pose_command.py @@ -0,0 +1,228 @@ +# 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 + +"""Warp-native pose command generator.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import warp as wp + +from isaaclab.assets import Articulation +from isaaclab.envs.mdp.commands._debug_vis import _PoseCommandDebugVis +from isaaclab.utils.leapp import POSE7_ELEMENT_NAMES + +from isaaclab_experimental.managers import CommandTerm + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .commands_cfg import UniformPoseCommandCfg + + +@wp.kernel +def _resample_pose_command( + env_mask: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + command: wp.array(dtype=wp.float32, ndim=2), + pos_x_min: float, + pos_x_max: float, + pos_y_min: float, + pos_y_max: float, + pos_z_min: float, + pos_z_max: float, + roll_min: float, + roll_max: float, + pitch_min: float, + pitch_max: float, + yaw_min: float, + yaw_max: float, + make_quat_unique: bool, +): + """Draw a new body-frame pose command for selected envs: uniform position [m] + + and roll/pitch/yaw [rad] composed into a quaternion (optionally + sign-canonicalized to the w >= 0 hemisphere). RNG state is written back so + the sequence advances across replays. + """ + env_id = wp.tid() + if env_mask[env_id]: + state = rng_state[env_id] + command[env_id, 0] = wp.randf(state, pos_x_min, pos_x_max) + command[env_id, 1] = wp.randf(state, pos_y_min, pos_y_max) + command[env_id, 2] = wp.randf(state, pos_z_min, pos_z_max) + + roll = wp.randf(state, roll_min, roll_max) + pitch = wp.randf(state, pitch_min, pitch_max) + yaw = wp.randf(state, yaw_min, yaw_max) + qx = wp.quat_from_axis_angle(wp.vec3f(1.0, 0.0, 0.0), roll) + qy = wp.quat_from_axis_angle(wp.vec3f(0.0, 1.0, 0.0), pitch) + qz = wp.quat_from_axis_angle(wp.vec3f(0.0, 0.0, 1.0), yaw) + quat = wp.mul(wp.mul(qz, qy), qx) + if make_quat_unique and quat[3] < 0.0: + quat = wp.quatf(-quat[0], -quat[1], -quat[2], -quat[3]) + command[env_id, 3] = quat[0] + command[env_id, 4] = quat[1] + command[env_id, 5] = quat[2] + command[env_id, 6] = quat[3] + rng_state[env_id] = state + + +@wp.kernel +def _update_pose_metrics( + command_b: wp.array(dtype=wp.float32, ndim=2), + 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), + body_idx: int, + command_w: wp.array(dtype=wp.float32, ndim=2), + position_error: wp.array(dtype=wp.float32), + orientation_error: wp.array(dtype=wp.float32), + success_rate: wp.array(dtype=wp.float32), + track_success: bool, + position_success_threshold: float, +): + """Refresh the world-frame goal pose (root pose composed with the body-frame + + command) and the tracking metrics: end-effector position [m] / orientation + [rad] errors vs the goal, and a position-threshold success flag when success + tracking is enabled. + """ + env_id = wp.tid() + position_b = wp.vec3f(command_b[env_id, 0], command_b[env_id, 1], command_b[env_id, 2]) + orientation_b = wp.quatf(command_b[env_id, 3], command_b[env_id, 4], command_b[env_id, 5], command_b[env_id, 6]) + desired_position_w = root_pos_w[env_id] + wp.quat_rotate(root_quat_w[env_id], position_b) + desired_orientation_w = root_quat_w[env_id] * orientation_b + + command_w[env_id, 0] = desired_position_w[0] + command_w[env_id, 1] = desired_position_w[1] + command_w[env_id, 2] = desired_position_w[2] + command_w[env_id, 3] = desired_orientation_w[0] + command_w[env_id, 4] = desired_orientation_w[1] + command_w[env_id, 5] = desired_orientation_w[2] + command_w[env_id, 6] = desired_orientation_w[3] + + position_delta = body_pos_w[env_id, body_idx] - desired_position_w + position_error_value = wp.length(position_delta) + orientation_delta = body_quat_w[env_id, body_idx] * wp.quat_inverse(desired_orientation_w) + orientation_error_value = 2.0 * wp.acos(wp.clamp(wp.abs(orientation_delta[3]), 0.0, 1.0)) + position_error[env_id] = position_error_value + orientation_error[env_id] = orientation_error_value + if track_success and position_error_value < position_success_threshold: + success_rate[env_id] = 1.0 + + +class UniformPoseCommand(_PoseCommandDebugVis, CommandTerm): + """Generate a pose command from uniform position and Euler-angle distributions.""" + + cfg: UniformPoseCommandCfg + """Configuration for the command generator.""" + + def __init__(self, cfg: UniformPoseCommandCfg, env: ManagerBasedEnv): + """Initialize the command generator. + + Args: + cfg: Configuration for the command generator. + env: Environment containing the commanded articulation. + """ + super().__init__(cfg, env) + + # -- robot / body resolution + self.robot: Articulation = env.scene[cfg.asset_name] + self.body_idx = self.robot.find_bodies(cfg.body_name)[0][0] + + # -- command buffers (Warp-native, pointer-stable) + self.pose_command_b = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) + self.pose_command_w = wp.zeros((self.num_envs, 7), dtype=wp.float32, device=self.device) + + # -- metrics buffers + self.metrics["position_error"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["orientation_error"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._success_rate = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._track_success = cfg.position_success_threshold is not None + if self._track_success: + self.metrics["success_rate"] = self._success_rate + + # adds (optional) cmd kind and element names for leapp export + self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/pose" + self.cfg.element_names = self.cfg.element_names or POSE7_ELEMENT_NAMES + + # -- Torch views (zero-copy aliases for stable consumers) + self.pose_command_b_torch = wp.to_torch(self.pose_command_b) + self.pose_command_w_torch = wp.to_torch(self.pose_command_w) + self.pose_command_b_torch[:, 6] = 1.0 + + def __str__(self) -> str: + """Return a string representation of the command generator.""" + msg = "UniformPoseCommand:\n" + msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" + return msg + + @property + def command(self) -> torch.Tensor: + """Desired pose command [m, m, m, quaternion xyzw], shape ``(num_envs, 7)``.""" + return self.pose_command_b_torch + + @property + def command_wp(self) -> wp.array(dtype=wp.float32, ndim=2): + """Pointer-stable desired pose command [m, m, m, quaternion xyzw].""" + return self.pose_command_b + + def _update_metrics(self): + wp.launch( + kernel=_update_pose_metrics, + dim=self.num_envs, + inputs=[ + self.pose_command_b, + self.robot.data.root_pos_w.warp, + self.robot.data.root_quat_w.warp, + self.robot.data.body_pos_w.warp, + self.robot.data.body_quat_w.warp, + self.body_idx, + self.pose_command_w, + self.metrics["position_error"], + self.metrics["orientation_error"], + self._success_rate, + self._track_success, + self.cfg.position_success_threshold or 0.0, + ], + device=self.device, + ) + + def _resample_command(self, env_mask: wp.array): + wp.launch( + kernel=_resample_pose_command, + dim=self.num_envs, + inputs=[ + env_mask, + self._env.rng_state_wp, + self.pose_command_b, + self.cfg.ranges.pos_x[0], + self.cfg.ranges.pos_x[1], + self.cfg.ranges.pos_y[0], + self.cfg.ranges.pos_y[1], + self.cfg.ranges.pos_z[0], + self.cfg.ranges.pos_z[1], + self.cfg.ranges.roll[0], + self.cfg.ranges.roll[1], + self.cfg.ranges.pitch[0], + self.cfg.ranges.pitch[1], + self.cfg.ranges.yaw[0], + self.cfg.ranges.yaw[1], + self.cfg.make_quat_unique, + ], + device=self.device, + ) + + def _update_command(self): + pass + + def _debug_pose_command_w(self) -> torch.Tensor: + """Return the Torch alias of the pointer-stable world-frame pose command buffer.""" + return self.pose_command_w_torch diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py new file mode 100644 index 000000000000..87707d64f103 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/commands/velocity_command.py @@ -0,0 +1,334 @@ +# 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 + +"""Warp-native velocity command generator.""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch +import warp as wp +from isaaclab_newton.kernels.state_kernels import body_ang_vel_from_root, body_lin_vel_from_root + +from isaaclab.assets import Articulation +from isaaclab.envs.mdp.commands._debug_vis import _VelocityCommandDebugVis + +from isaaclab_experimental.managers import CommandTerm +from isaaclab_experimental.utils.warp import wrap_to_pi + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .commands_cfg import UniformVelocityCommandCfg + +logger = logging.getLogger(__name__) + + +@wp.kernel +def _accumulate_velocity_metrics( + command: wp.array(dtype=wp.float32, ndim=2), + root_pose_w: wp.array(dtype=wp.transformf), + root_vel_w: wp.array(dtype=wp.spatial_vectorf), + error_xy_sum: wp.array(dtype=wp.float32), + error_yaw_sum: wp.array(dtype=wp.float32), + step_count: wp.array(dtype=wp.float32), +): + """Accumulate per-step command-tracking error sums for logging: planar (xy) + + linear-velocity error [m/s] and yaw-rate error [rad/s] vs the command, plus + the step count used later for averaging. Body-frame velocities are derived + inline from the root pose/velocity (capture-safe, no lazy buffers). + """ + env_id = wp.tid() + root_lin_vel_b = body_lin_vel_from_root(root_pose_w[env_id], root_vel_w[env_id]) + root_ang_vel_b = body_ang_vel_from_root(root_pose_w[env_id], root_vel_w[env_id]) + dx = command[env_id, 0] - root_lin_vel_b[0] + dy = command[env_id, 1] - root_lin_vel_b[1] + error_xy_sum[env_id] += wp.sqrt(dx * dx + dy * dy) + error_yaw_sum[env_id] += wp.abs(command[env_id, 2] - root_ang_vel_b[2]) + step_count[env_id] += 1.0 + + +@wp.kernel +def _finalize_velocity_metrics( + env_mask: wp.array(dtype=wp.bool), + error_xy_sum: wp.array(dtype=wp.float32), + error_yaw_sum: wp.array(dtype=wp.float32), + step_count: wp.array(dtype=wp.float32), + error_xy: wp.array(dtype=wp.float32), + error_yaw: wp.array(dtype=wp.float32), + success_rate: wp.array(dtype=wp.float32), + error_xy_threshold: float, + error_yaw_threshold: float, +): + """Turn selected envs' accumulated error sums into episode means and a binary + + success flag (both mean errors under their thresholds), then zero the + accumulators for the next episode. + """ + env_id = wp.tid() + if env_mask[env_id]: + denominator = wp.max(step_count[env_id], 1.0) + mean_error_xy = error_xy_sum[env_id] / denominator + mean_error_yaw = error_yaw_sum[env_id] / denominator + error_xy[env_id] = mean_error_xy + error_yaw[env_id] = mean_error_yaw + success_rate[env_id] = wp.where( + (mean_error_xy < error_xy_threshold) and (mean_error_yaw < error_yaw_threshold), 1.0, 0.0 + ) + error_xy_sum[env_id] = 0.0 + error_yaw_sum[env_id] = 0.0 + step_count[env_id] = 0.0 + + +@wp.kernel +def _resample_velocity_command( + env_mask: wp.array(dtype=wp.bool), + rng_state: wp.array(dtype=wp.uint32), + command: wp.array(dtype=wp.float32, ndim=2), + heading_target: wp.array(dtype=wp.float32), + is_heading_env: wp.array(dtype=wp.bool), + is_standing_env: wp.array(dtype=wp.bool), + lin_vel_x_min: float, + lin_vel_x_max: float, + lin_vel_y_min: float, + lin_vel_y_max: float, + ang_vel_z_min: float, + ang_vel_z_max: float, + heading_min: float, + heading_max: float, + heading_command: bool, + rel_heading_envs: float, + rel_standing_envs: float, +): + """Draw a new SE(2) velocity command (vx, vy [m/s], wz [rad/s]) for selected + + envs, plus their heading target [rad] and heading/standing role flags. + The per-env device RNG state is written back so the random sequence advances + across CUDA-graph replays. + """ + env_id = wp.tid() + if env_mask[env_id]: + state = rng_state[env_id] + command[env_id, 0] = wp.randf(state, lin_vel_x_min, lin_vel_x_max) + command[env_id, 1] = wp.randf(state, lin_vel_y_min, lin_vel_y_max) + command[env_id, 2] = wp.randf(state, ang_vel_z_min, ang_vel_z_max) + if heading_command: + heading_target[env_id] = wp.randf(state, heading_min, heading_max) + is_heading_env[env_id] = wp.randf(state, 0.0, 1.0) <= rel_heading_envs + else: + is_heading_env[env_id] = False + is_standing_env[env_id] = wp.randf(state, 0.0, 1.0) <= rel_standing_envs + rng_state[env_id] = state + + +@wp.kernel +def _update_velocity_command( + command: wp.array(dtype=wp.float32, ndim=2), + heading_target: wp.array(dtype=wp.float32), + is_heading_env: wp.array(dtype=wp.bool), + is_standing_env: wp.array(dtype=wp.bool), + root_pose_w: wp.array(dtype=wp.transformf), + heading_command: bool, + heading_control_stiffness: float, + ang_vel_z_min: float, + ang_vel_z_max: float, +): + """Per-step command shaping: heading-controlled envs convert heading error into + + a clamped yaw-rate command (proportional control); standing envs zero their + command entirely. + """ + env_id = wp.tid() + if heading_command and is_heading_env[env_id]: + root_quat_w = wp.transform_get_rotation(root_pose_w[env_id]) + forward_w = wp.quat_rotate(root_quat_w, wp.vec3f(1.0, 0.0, 0.0)) + heading_w = wp.atan2(forward_w[1], forward_w[0]) + heading_error = wrap_to_pi(heading_target[env_id] - heading_w) + command[env_id, 2] = wp.clamp( + heading_control_stiffness * heading_error, + ang_vel_z_min, + ang_vel_z_max, + ) + if is_standing_env[env_id]: + command[env_id, 0] = 0.0 + command[env_id, 1] = 0.0 + command[env_id, 2] = 0.0 + + +class UniformVelocityCommand(_VelocityCommandDebugVis, CommandTerm): + """Generate an SE(2) velocity command from uniform distributions.""" + + cfg: UniformVelocityCommandCfg + """Configuration for the command generator.""" + + def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedEnv): + """Initialize the command generator. + + Args: + cfg: Configuration for the command generator. + env: Environment containing the commanded articulation. + + Raises: + ValueError: If heading commands are enabled without a heading range. + """ + super().__init__(cfg, env) + + # -- config validation + if cfg.heading_command and cfg.ranges.heading is None: + raise ValueError( + "The velocity command has heading commands active (heading_command=True) but the `ranges.heading`" + " parameter is set to None." + ) + if cfg.ranges.heading and not cfg.heading_command: + logger.warning( + f"The velocity command has the 'ranges.heading' attribute set to '{cfg.ranges.heading}'" + " but the heading command is not active. Consider setting the flag for the heading command to True." + ) + + # -- robot resolution + self.robot: Articulation = env.scene[cfg.asset_name] + + # -- command buffers (Warp-native, pointer-stable) + self.vel_command_b = wp.zeros((self.num_envs, 3), dtype=wp.float32, device=self.device) + self.heading_target = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.is_heading_env = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self.is_standing_env = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + + # -- metrics buffers + self.metrics["error_vel_xy"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["error_vel_yaw"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.metrics["success_rate"] = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._error_xy_sum = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._error_yaw_sum = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self._step_count = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + + # adds (optional) cmd kind and element names for leapp export + self.cfg.cmd_kind = self.cfg.cmd_kind or "command/body/velocity" + self.cfg.element_names = self.cfg.element_names or ["lin_vel_x", "lin_vel_y", "ang_vel_z"] + + # -- Torch views (zero-copy aliases for stable consumers) + self.vel_command_b_torch = wp.to_torch(self.vel_command_b) + self.heading_target_torch = wp.to_torch(self.heading_target) + self.is_heading_env_torch = wp.to_torch(self.is_heading_env) + self.is_standing_env_torch = wp.to_torch(self.is_standing_env) + + def __str__(self) -> str: + """Return a string representation of the command generator.""" + msg = "UniformVelocityCommand:\n" + msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n" + msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n" + msg += f"\tHeading command: {self.cfg.heading_command}\n" + if self.cfg.heading_command: + msg += f"\tHeading probability: {self.cfg.rel_heading_envs}\n" + msg += f"\tStanding probability: {self.cfg.rel_standing_envs}" + return msg + + @property + def command(self) -> torch.Tensor: + """Desired base velocity command [m/s, m/s, rad/s], shape ``(num_envs, 3)``.""" + return self.vel_command_b_torch + + @property + def command_wp(self) -> wp.array(dtype=wp.float32, ndim=2): + """Pointer-stable desired base velocity command [m/s, m/s, rad/s].""" + return self.vel_command_b + + def reset( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + *, + env_mask: wp.array | None = None, + ) -> dict[str, torch.Tensor]: + """Finalize episode metrics and reset selected command state. + + Args: + env_ids: Environment indices used by compatibility call sites. + env_mask: Boolean Warp mask selecting environments to reset. Takes precedence over ``env_ids``. + + Returns: + Persistent scalar metric views for logging. + """ + env_mask = self._resolve_reset_mask(env_ids, env_mask) + wp.launch( + kernel=_finalize_velocity_metrics, + dim=self.num_envs, + inputs=[ + env_mask, + self._error_xy_sum, + self._error_yaw_sum, + self._step_count, + self.metrics["error_vel_xy"], + self.metrics["error_vel_yaw"], + self.metrics["success_rate"], + self.cfg.vel_xy_success_threshold, + self.cfg.vel_yaw_success_threshold, + ], + device=self.device, + ) + return super().reset(env_mask=env_mask) + + def _update_metrics(self): + wp.launch( + kernel=_accumulate_velocity_metrics, + dim=self.num_envs, + inputs=[ + self.vel_command_b, + self.robot.data.root_pose_w.warp, + self.robot.data.root_vel_w.warp, + self._error_xy_sum, + self._error_yaw_sum, + self._step_count, + ], + device=self.device, + ) + + def _resample_command(self, env_mask: wp.array): + heading_range = self.cfg.ranges.heading if self.cfg.ranges.heading is not None else (0.0, 0.0) + wp.launch( + kernel=_resample_velocity_command, + dim=self.num_envs, + inputs=[ + env_mask, + self._env.rng_state_wp, + self.vel_command_b, + self.heading_target, + self.is_heading_env, + self.is_standing_env, + self.cfg.ranges.lin_vel_x[0], + self.cfg.ranges.lin_vel_x[1], + self.cfg.ranges.lin_vel_y[0], + self.cfg.ranges.lin_vel_y[1], + self.cfg.ranges.ang_vel_z[0], + self.cfg.ranges.ang_vel_z[1], + heading_range[0], + heading_range[1], + self.cfg.heading_command, + self.cfg.rel_heading_envs, + self.cfg.rel_standing_envs, + ], + device=self.device, + ) + + def _update_command(self): + wp.launch( + kernel=_update_velocity_command, + dim=self.num_envs, + inputs=[ + self.vel_command_b, + self.heading_target, + self.is_heading_env, + self.is_standing_env, + self.robot.data.root_pose_w.warp, + self.cfg.heading_command, + self.cfg.heading_control_stiffness, + self.cfg.ranges.ang_vel_z[0], + self.cfg.ranges.ang_vel_z[1], + ], + device=self.device, + ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index d9e61041a7ad..faadee5f1770 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -5,7 +5,7 @@ """Warp-first overrides for common event terms. -These functions are intended to be used with the experimental Warp-first +These terms are intended to be used with the experimental Warp-first :class:`isaaclab_experimental.managers.EventManager` (mask-based interval/reset). Why this exists: @@ -17,6 +17,9 @@ These Warp-first implementations avoid that by writing directly into the sim-bound Warp state buffers (`asset.data.joint_pos` / `asset.data.joint_vel`) for the selected envs/joints. +Stateful terms use :class:`~isaaclab_experimental.managers.ManagerTermBase` so +persistent buffers and parsed constants are created once during manager setup. + Notes: - These terms assume the Newton/Warp backend (Warp arrays are available for joint state and defaults). - For best performance, pass :class:`isaaclab_experimental.managers.SceneEntityCfg` so `joint_ids_wp` is cached. @@ -26,13 +29,40 @@ from typing import TYPE_CHECKING +import torch import warp as wp -from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab.envs.mdp.events import randomize_rigid_body_mass as _StableRandomizeRigidBodyMass +from isaaclab.envs.mdp.events import randomize_rigid_body_material as _StableRandomizeRigidBodyMaterial + +from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg +from isaaclab_experimental.managers import ManagerTermBase as _WarpManagerTermBase from isaaclab_experimental.utils.warp import WarpCapturable +__all__ = [ + "apply_external_force_torque", + "push_by_setting_velocity", + "randomize_rigid_body_com", + "randomize_rigid_body_mass", + "randomize_rigid_body_material", + "reset_joints_by_offset", + "reset_joints_by_scale", + "reset_root_state_uniform", +] + if TYPE_CHECKING: - from isaaclab.assets import Articulation + from isaaclab.assets import Articulation, RigidObject + from isaaclab.envs import ManagerBasedEnv + + +def _resolve_body_ids(asset_cfg: SceneEntityCfg, num_bodies: int) -> list[int]: + """Resolve a configured body selection without allocating device storage.""" + if asset_cfg.body_ids is None or asset_cfg.body_ids == slice(None): + return list(range(num_bodies)) + if isinstance(asset_cfg.body_ids, int): + return [asset_cfg.body_ids] + return list(asset_cfg.body_ids) + # --------------------------------------------------------------------------- # Randomize rigid body center of mass @@ -43,66 +73,86 @@ def _randomize_com_kernel( env_mask: wp.array(dtype=wp.bool), rng_state: wp.array(dtype=wp.uint32), + default_body_com_pos_b: wp.array(dtype=wp.vec3f, ndim=2), body_com_pos_b: wp.array(dtype=wp.vec3f, ndim=2), body_ids: wp.array(dtype=wp.int32), com_lo: wp.vec3f, com_hi: wp.vec3f, ): - """Add random offset to center of mass positions for selected bodies.""" + """Add random offsets to the default center-of-mass positions for selected bodies.""" env_id = wp.tid() if not env_mask[env_id]: return state = rng_state[env_id] + dx = wp.randf(state, com_lo[0], com_hi[0]) + dy = wp.randf(state, com_lo[1], com_hi[1]) + dz = wp.randf(state, com_lo[2], com_hi[2]) for k in range(body_ids.shape[0]): b = body_ids[k] - v = body_com_pos_b[env_id, b] - dx = wp.randf(state, com_lo[0], com_hi[0]) - dy = wp.randf(state, com_lo[1], com_hi[1]) - dz = wp.randf(state, com_lo[2], com_hi[2]) + v = default_body_com_pos_b[env_id, b] body_com_pos_b[env_id, b] = wp.vec3f(v[0] + dx, v[1] + dy, v[2] + dz) rng_state[env_id] = state @WarpCapturable(False, reason="set_coms_mask calls SimulationManager.add_model_change") -def randomize_rigid_body_com( - env, - env_mask: wp.array, - com_range: dict[str, tuple[float, float]], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): - """Randomize the center of mass (CoM) of rigid bodies by adding random offsets. +class randomize_rigid_body_com(ManagerTermBase): + """Randomize rigid-body centers of mass from a persistent default baseline. - Warp-first override of :func:`isaaclab.envs.mdp.events.randomize_rigid_body_com`. - Writes directly into the sim-bound ``body_com_pos_b`` buffer, then notifies the solver - via :meth:`set_coms_mask` so it recomputes inertial properties. + This term is not CUDA-graph capturable because notifying the solver of changed + inertial properties calls :meth:`SimulationManager.add_model_change`. """ - asset: Articulation = env.scene[asset_cfg.name] - fn = randomize_rigid_body_com - if not getattr(fn, "_is_warmed_up", False) or fn._asset_name != asset_cfg.name: - fn._asset_name = asset_cfg.name - r = [com_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - fn._com_lo = wp.vec3f(r[0][0], r[1][0], r[2][0]) - fn._com_hi = wp.vec3f(r[0][1], r[1][1], r[2][1]) - fn._is_warmed_up = True - - wp.launch( - kernel=_randomize_com_kernel, - dim=env.num_envs, - inputs=[ - env_mask, - env.rng_state_wp, - asset.data.body_com_pos_b.warp, - asset_cfg.body_ids_wp, - fn._com_lo, - fn._com_hi, - ], - device=env.device, - ) + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: + """Initialize persistent center-of-mass randomization state. + + Args: + cfg: Event term configuration. + env: Environment containing the randomized asset. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + self._default_com = wp.clone(self._asset.data.body_com_pos_b.warp) + body_ids = _resolve_body_ids(asset_cfg, self._asset.num_bodies) + self._body_ids = wp.array(body_ids, dtype=wp.int32, device=env.device) + + com_range = cfg.params["com_range"] + ranges = [com_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z")] + self._com_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) + self._com_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), + com_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Randomize selected center-of-mass offsets [m]. + + Args: + env: Environment containing the randomized asset. + env_mask: Boolean Warp mask selecting environments. + com_range: Per-axis offset ranges [m]. Parsed during initialization. + asset_cfg: Scene entity selection. Resolved during initialization. + """ + wp.launch( + kernel=_randomize_com_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + env.rng_state_wp, + self._default_com, + self._asset.data.body_com_pos_b.warp, + self._body_ids, + self._com_lo, + self._com_hi, + ], + device=env.device, + ) - # Notify the solver that inertial properties changed (COM position affects inertia). - asset.set_coms_mask(coms=asset.data.body_com_pos_b.warp, env_mask=env_mask) + self._asset.set_coms_mask(coms=self._asset.data.body_com_pos_b.warp, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -114,6 +164,7 @@ def randomize_rigid_body_com( def _apply_external_force_torque_kernel( env_mask: wp.array(dtype=wp.bool), rng_state: wp.array(dtype=wp.uint32), + body_ids: wp.array(dtype=wp.int32), force_out: wp.array(dtype=wp.vec3f, ndim=2), torque_out: wp.array(dtype=wp.vec3f, ndim=2), force_lo: float, @@ -121,16 +172,17 @@ def _apply_external_force_torque_kernel( torque_lo: float, torque_hi: float, ): + """Sample uniform random external force [N] and torque [N·m] vectors for + + selected envs' target bodies into the asset's wrench composer buffers. + """ env_id = wp.tid() if not env_mask[env_id]: - # zero out unmasked envs so they don't accumulate stale forces - for b in range(force_out.shape[1]): - force_out[env_id, b] = wp.vec3f(0.0, 0.0, 0.0) - torque_out[env_id, b] = wp.vec3f(0.0, 0.0, 0.0) return state = rng_state[env_id] - for b in range(force_out.shape[1]): + for body_index in range(body_ids.shape[0]): + b = body_ids[body_index] force_out[env_id, b] = wp.vec3f( wp.randf(state, force_lo, force_hi), wp.randf(state, force_lo, force_hi), @@ -144,50 +196,76 @@ def _apply_external_force_torque_kernel( rng_state[env_id] = state -def apply_external_force_torque( - env, - env_mask: wp.array, - force_range: tuple[float, float], - torque_range: tuple[float, float], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): - """Randomize external forces and torques applied to the asset's bodies. - - Warp-first override of :func:`isaaclab.envs.mdp.events.apply_external_force_torque`. - """ - asset: Articulation = env.scene[asset_cfg.name] - - # First-call: allocate scratch and pre-convert constant arguments. - if not getattr(apply_external_force_torque, "_is_warmed_up", False): - apply_external_force_torque._scratch_forces = wp.zeros( - (env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device - ) - apply_external_force_torque._scratch_torques = wp.zeros( - (env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device +class apply_external_force_torque(ManagerTermBase): + """Apply random external forces and torques using persistent wrench buffers.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: + """Initialize persistent external-wrench state. + + Args: + cfg: Event term configuration. + env: Environment containing the randomized asset. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + self._forces = wp.zeros((env.num_envs, self._asset.num_bodies), dtype=wp.vec3f, device=env.device) + self._torques = wp.zeros((env.num_envs, self._asset.num_bodies), dtype=wp.vec3f, device=env.device) + + body_ids = _resolve_body_ids(asset_cfg, self._asset.num_bodies) + body_mask = [False] * self._asset.num_bodies + for body_id in body_ids: + body_mask[body_id] = True + self._body_ids = wp.array(body_ids, dtype=wp.int32, device=env.device) + self._body_mask = wp.array(body_mask, dtype=wp.bool, device=env.device) + + force_range = cfg.params["force_range"] + torque_range = cfg.params["torque_range"] + self._force_lo = float(force_range[0]) + self._force_hi = float(force_range[1]) + self._torque_lo = float(torque_range[0]) + self._torque_hi = float(torque_range[1]) + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), + force_range: tuple[float, float], + torque_range: tuple[float, float], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Apply sampled forces [N] and torques [N·m] to selected bodies. + + Args: + env: Environment containing the randomized asset. + env_mask: Boolean Warp mask selecting environments. + force_range: Component-wise force range [N]. Parsed during initialization. + torque_range: Component-wise torque range [N·m]. Parsed during initialization. + asset_cfg: Scene entity selection. Resolved during initialization. + """ + wp.launch( + kernel=_apply_external_force_torque_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + env.rng_state_wp, + self._body_ids, + self._forces, + self._torques, + self._force_lo, + self._force_hi, + self._torque_lo, + self._torque_hi, + ], + device=env.device, ) - apply_external_force_torque._is_warmed_up = True - - wp.launch( - kernel=_apply_external_force_torque_kernel, - dim=env.num_envs, - inputs=[ - env_mask, - env.rng_state_wp, - apply_external_force_torque._scratch_forces, - apply_external_force_torque._scratch_torques, - force_range[0], - force_range[1], - torque_range[0], - torque_range[1], - ], - device=env.device, - ) - asset.permanent_wrench_composer.set_forces_and_torques_mask( - forces=apply_external_force_torque._scratch_forces, - torques=apply_external_force_torque._scratch_torques, - env_mask=env_mask, - ) + self._asset.permanent_wrench_composer.set_forces_and_torques_mask( + forces=self._forces, + torques=self._torques, + body_mask=self._body_mask, + env_mask=env_mask, + ) # --------------------------------------------------------------------------- @@ -206,6 +284,10 @@ def _push_by_setting_velocity_kernel( ang_lo: wp.vec3f, ang_hi: wp.vec3f, ): + """Add a uniform random velocity kick [m/s, rad/s] to selected envs' current + + root velocity, writing the result for the masked sim write. + """ env_id = wp.tid() if not env_mask[env_id]: return @@ -225,45 +307,60 @@ def _push_by_setting_velocity_kernel( rng_state[env_id] = state -def push_by_setting_velocity( - env, - env_mask: wp.array, - velocity_range: dict[str, tuple[float, float]], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): - """Push the asset by setting the root velocity to a random value within the given ranges. - - Warp-first override of :func:`isaaclab.envs.mdp.events.push_by_setting_velocity`. - """ - asset: Articulation = env.scene[asset_cfg.name] - - # First-call: allocate scratch and pre-parse constant range arguments. - if not getattr(push_by_setting_velocity, "_is_warmed_up", False): - push_by_setting_velocity._scratch_vel = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) - r = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - push_by_setting_velocity._lin_lo = wp.vec3f(r[0][0], r[1][0], r[2][0]) - push_by_setting_velocity._lin_hi = wp.vec3f(r[0][1], r[1][1], r[2][1]) - push_by_setting_velocity._ang_lo = wp.vec3f(r[3][0], r[4][0], r[5][0]) - push_by_setting_velocity._ang_hi = wp.vec3f(r[3][1], r[4][1], r[5][1]) - push_by_setting_velocity._is_warmed_up = True - - wp.launch( - kernel=_push_by_setting_velocity_kernel, - dim=env.num_envs, - inputs=[ - env_mask, - env.rng_state_wp, - asset.data.root_vel_w.warp, - push_by_setting_velocity._scratch_vel, - push_by_setting_velocity._lin_lo, - push_by_setting_velocity._lin_hi, - push_by_setting_velocity._ang_lo, - push_by_setting_velocity._ang_hi, - ], - device=env.device, - ) +class push_by_setting_velocity(ManagerTermBase): + """Push an asset by sampling into a persistent root-velocity buffer.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: + """Initialize persistent root-velocity push state. + + Args: + cfg: Event term configuration. + env: Environment containing the pushed asset. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + self._velocity = wp.zeros(env.num_envs, dtype=wp.spatial_vectorf, device=env.device) + + velocity_range = cfg.params["velocity_range"] + ranges = [velocity_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self._lin_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) + self._lin_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) + self._ang_lo = wp.vec3f(ranges[3][0], ranges[4][0], ranges[5][0]) + self._ang_hi = wp.vec3f(ranges[3][1], ranges[4][1], ranges[5][1]) + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), + velocity_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Add sampled linear [m/s] and angular [rad/s] root velocities. + + Args: + env: Environment containing the pushed asset. + env_mask: Boolean Warp mask selecting environments. + velocity_range: Per-axis velocity ranges [m/s or rad/s]. Parsed during initialization. + asset_cfg: Scene entity selection. Resolved during initialization. + """ + wp.launch( + kernel=_push_by_setting_velocity_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + env.rng_state_wp, + self._asset.data.root_vel_w.warp, + self._velocity, + self._lin_lo, + self._lin_hi, + self._ang_lo, + self._ang_hi, + ], + device=env.device, + ) - asset.write_root_velocity_to_sim_mask(root_velocity=push_by_setting_velocity._scratch_vel, env_mask=env_mask) + self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -289,6 +386,11 @@ def _reset_root_state_uniform_kernel( vel_ang_lo: wp.vec3f, vel_ang_hi: wp.vec3f, ): + """Compose selected envs' reset root state: default pose offset by the env + + origin plus uniform position [m] / roll-pitch-yaw [rad] noise, and default + velocity plus uniform noise [m/s, rad/s]. + """ env_id = wp.tid() if not env_mask[env_id]: return @@ -335,62 +437,81 @@ def _reset_root_state_uniform_kernel( rng_state[env_id] = state -def reset_root_state_uniform( - env, - env_mask: wp.array, - pose_range: dict[str, tuple[float, float]], - velocity_range: dict[str, tuple[float, float]], - asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): - """Reset the asset root state to a random position and velocity uniformly within the given ranges. - - Warp-first override of :func:`isaaclab.envs.mdp.events.reset_root_state_uniform`. - """ - asset: Articulation = env.scene[asset_cfg.name] - - # First-call: allocate scratch and pre-parse range dicts. - if not getattr(reset_root_state_uniform, "_is_warmed_up", False): - reset_root_state_uniform._scratch_pose = wp.zeros((env.num_envs,), dtype=wp.transformf, device=env.device) - reset_root_state_uniform._scratch_vel = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) - # Pre-parse pose_range dict - p = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - reset_root_state_uniform._pos_lo = wp.vec3f(p[0][0], p[1][0], p[2][0]) - reset_root_state_uniform._pos_hi = wp.vec3f(p[0][1], p[1][1], p[2][1]) - reset_root_state_uniform._rot_lo = wp.vec3f(p[3][0], p[4][0], p[5][0]) - reset_root_state_uniform._rot_hi = wp.vec3f(p[3][1], p[4][1], p[5][1]) - # Pre-parse velocity_range dict - v = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - reset_root_state_uniform._vel_lin_lo = wp.vec3f(v[0][0], v[1][0], v[2][0]) - reset_root_state_uniform._vel_lin_hi = wp.vec3f(v[0][1], v[1][1], v[2][1]) - reset_root_state_uniform._vel_ang_lo = wp.vec3f(v[3][0], v[4][0], v[5][0]) - reset_root_state_uniform._vel_ang_hi = wp.vec3f(v[3][1], v[4][1], v[5][1]) - reset_root_state_uniform._is_warmed_up = True - - wp.launch( - kernel=_reset_root_state_uniform_kernel, - dim=env.num_envs, - inputs=[ - env_mask, - env.rng_state_wp, - asset.data.default_root_pose.warp, - asset.data.default_root_vel.warp, - env.env_origins_wp, - reset_root_state_uniform._scratch_pose, - reset_root_state_uniform._scratch_vel, - reset_root_state_uniform._pos_lo, - reset_root_state_uniform._pos_hi, - reset_root_state_uniform._rot_lo, - reset_root_state_uniform._rot_hi, - reset_root_state_uniform._vel_lin_lo, - reset_root_state_uniform._vel_lin_hi, - reset_root_state_uniform._vel_ang_lo, - reset_root_state_uniform._vel_ang_hi, - ], - device=env.device, - ) +class reset_root_state_uniform(ManagerTermBase): + """Reset root pose and velocity using persistent Warp output buffers.""" + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv) -> None: + """Initialize persistent root-state reset state. + + Args: + cfg: Event term configuration. + env: Environment containing the reset asset. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + self._default_root_pose = self._asset.data.default_root_pose.warp + self._default_root_velocity = self._asset.data.default_root_vel.warp + self._env_origins = env.env_origins_wp + self._pose = wp.zeros(env.num_envs, dtype=wp.transformf, device=env.device) + self._velocity = wp.zeros(env.num_envs, dtype=wp.spatial_vectorf, device=env.device) + + pose_range = cfg.params["pose_range"] + pose_ranges = [pose_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self._pos_lo = wp.vec3f(pose_ranges[0][0], pose_ranges[1][0], pose_ranges[2][0]) + self._pos_hi = wp.vec3f(pose_ranges[0][1], pose_ranges[1][1], pose_ranges[2][1]) + self._rot_lo = wp.vec3f(pose_ranges[3][0], pose_ranges[4][0], pose_ranges[5][0]) + self._rot_hi = wp.vec3f(pose_ranges[3][1], pose_ranges[4][1], pose_ranges[5][1]) + + velocity_range = cfg.params["velocity_range"] + velocity_ranges = [velocity_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + self._vel_lin_lo = wp.vec3f(velocity_ranges[0][0], velocity_ranges[1][0], velocity_ranges[2][0]) + self._vel_lin_hi = wp.vec3f(velocity_ranges[0][1], velocity_ranges[1][1], velocity_ranges[2][1]) + self._vel_ang_lo = wp.vec3f(velocity_ranges[3][0], velocity_ranges[4][0], velocity_ranges[5][0]) + self._vel_ang_hi = wp.vec3f(velocity_ranges[3][1], velocity_ranges[4][1], velocity_ranges[5][1]) + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), + pose_range: dict[str, tuple[float, float]], + velocity_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Reset root pose [m, rad] and velocity [m/s, rad/s]. + + Args: + env: Environment containing the reset asset. + env_mask: Boolean Warp mask selecting environments. + pose_range: Position and Euler-angle ranges [m or rad]. Parsed during initialization. + velocity_range: Linear and angular velocity ranges [m/s or rad/s]. Parsed during initialization. + asset_cfg: Scene entity selection. Resolved during initialization. + """ + wp.launch( + kernel=_reset_root_state_uniform_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + env.rng_state_wp, + self._default_root_pose, + self._default_root_velocity, + self._env_origins, + self._pose, + self._velocity, + self._pos_lo, + self._pos_hi, + self._rot_lo, + self._rot_hi, + self._vel_lin_lo, + self._vel_lin_hi, + self._vel_ang_lo, + self._vel_ang_hi, + ], + device=env.device, + ) - asset.write_root_pose_to_sim_mask(root_pose=reset_root_state_uniform._scratch_pose, env_mask=env_mask) - asset.write_root_velocity_to_sim_mask(root_velocity=reset_root_state_uniform._scratch_vel, env_mask=env_mask) + self._asset.write_root_pose_to_sim_mask(root_pose=self._pose, env_mask=env_mask) + self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -414,6 +535,11 @@ def _reset_joints_by_offset_kernel( vel_lo: float, vel_hi: float, ): + """Reset selected envs' selected joints to defaults plus uniform offsets + + [rad or m, depending on joint type], clamped to the soft position and + velocity limits. + """ env_id = wp.tid() if not env_mask[env_id]: return @@ -444,12 +570,12 @@ def _reset_joints_by_offset_kernel( def reset_joints_by_offset( - env, - env_mask: wp.array, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), position_range: tuple[float, float], velocity_range: tuple[float, float], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): +) -> None: """Warp-first reset of joint state by random offsets around defaults. This overrides the stable `isaaclab.envs.mdp.events.reset_joints_by_offset` when importing @@ -512,6 +638,10 @@ def _reset_joints_by_scale_kernel( vel_lo: float, vel_hi: float, ): + """Reset selected envs' selected joints to defaults scaled by uniform random + + factors, clamped to the soft position and velocity limits. + """ env_id = wp.tid() if not env_mask[env_id]: return @@ -540,12 +670,12 @@ def _reset_joints_by_scale_kernel( def reset_joints_by_scale( - env, - env_mask: wp.array, + env: ManagerBasedEnv, + env_mask: wp.array(dtype=wp.bool), position_range: tuple[float, float], velocity_range: tuple[float, float], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), -): +) -> None: """Warp-first reset of joint state by scaling defaults with random factors.""" asset: Articulation = env.scene[asset_cfg.name] @@ -585,3 +715,90 @@ def reset_joints_by_scale( # Sync derived buffers (_previous_joint_vel, joint_acc) for reset envs. asset.write_joint_position_to_sim_mask(position=asset.data.joint_pos.warp, env_mask=env_mask) asset.write_joint_velocity_to_sim_mask(velocity=asset.data.joint_vel.warp, env_mask=env_mask) + + +# --------------------------------------------------------------------------- +# Startup-mode randomization (mask-signature adapters) +# --------------------------------------------------------------------------- +# +# The stable material/mass randomization terms already dispatch to the active +# physics backend (Newton included); only their calling convention differs — +# the warp EventManager invokes terms with a Warp env-mask while the stable +# class terms expect torch env indices. Both terms run in ``startup`` mode, +# once at construction and outside any captured path, so a host-side +# mask-to-ids conversion is acceptable. + + +def _mask_to_env_ids(env_mask: wp.array) -> torch.Tensor: + """Convert a Warp boolean env-mask to the torch index tensor the stable terms expect.""" + return torch.nonzero(wp.to_torch(env_mask), as_tuple=False).squeeze(-1) + + +class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, _WarpManagerTermBase): + """Warp adapter for the stable, backend-dispatched material randomization term. + + Converts the warp event manager's env-mask calling convention to the stable + term's env-ids convention and delegates. Startup mode only. Inherits the + warp :class:`ManagerTermBase` so the warp managers accept it as a class term. + """ + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array, + static_friction_range: tuple[float, float], + dynamic_friction_range: tuple[float, float], + restitution_range: tuple[float, float], + num_buckets: int, + asset_cfg: SceneEntityCfg, + make_consistent: bool = False, + ) -> None: + super().__call__( + env, + _mask_to_env_ids(env_mask), + static_friction_range, + dynamic_friction_range, + restitution_range, + num_buckets, + asset_cfg, + make_consistent, + ) + + def reset(self, env_mask: wp.array | None = None) -> None: + """Nothing to reset; materials are randomized once at startup.""" + return + + +class randomize_rigid_body_mass(_StableRandomizeRigidBodyMass, _WarpManagerTermBase): + """Warp adapter for the stable, backend-dispatched mass randomization term. + + Converts the warp event manager's env-mask calling convention to the stable + term's env-ids convention and delegates. Startup mode only. Inherits the + warp :class:`ManagerTermBase` so the warp managers accept it as a class term. + """ + + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array, + asset_cfg: SceneEntityCfg, + mass_distribution_params: tuple[float, float], + operation: str, + distribution: str = "uniform", + recompute_inertia: bool = True, + min_mass: float = 1e-6, + ) -> None: + super().__call__( + env, + _mask_to_env_ids(env_mask), + asset_cfg, + mass_distribution_params, + operation, + distribution, + recompute_inertia, + min_mass, + ) + + def reset(self, env_mask: wp.array | None = None) -> None: + """Nothing to reset; masses are randomized once at startup.""" + return diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index 808231c8faf6..d534d70dbcec 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -36,10 +36,12 @@ record_shape, ) from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.utils.warp import WarpCapturable if TYPE_CHECKING: from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv + from isaaclab.sensors import RayCaster # --------------------------------------------------------------------------- @@ -52,6 +54,7 @@ def _vec3_to_out3_kernel( src: wp.array(dtype=wp.vec3f), out: wp.array(dtype=wp.float32, ndim=2), ): + """Scatter a vec3 array into three float32 columns of a term's ``out`` block.""" env_id = wp.tid() v = src[env_id] out[env_id, 0] = v[0] @@ -65,6 +68,10 @@ def _joint_gather_kernel( joint_ids: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.float32, ndim=2), ): + """Gather selected joints' values into contiguous term columns. + + Launch with dim=(num_envs, num_selected_joints). + """ env_id, k = wp.tid() j = joint_ids[k] out[env_id, k] = src[env_id, j] @@ -80,6 +87,7 @@ def _base_pos_z_kernel( root_pos_w: wp.array(dtype=wp.vec3f), out: wp.array(dtype=wp.float32, ndim=2), ): + """Write the root height above the world origin [m] into the term's single column.""" env_id = wp.tid() out[env_id, 0] = root_pos_w[env_id][2] @@ -112,6 +120,10 @@ def _base_lin_vel_kernel( root_vel_w: wp.array(dtype=wp.spatial_vectorf), out: wp.array(dtype=wp.float32, ndim=2), ): + """Root linear velocity [m/s] in the body frame, derived inline from the root + + pose and CoM velocity (Tier-1 access; avoids non-capturable lazy buffers). + """ i = wp.tid() v = body_lin_vel_from_root(root_pose_w[i], root_vel_w[i]) out[i, 0] = v[0] @@ -139,6 +151,10 @@ def _base_ang_vel_kernel( root_vel_w: wp.array(dtype=wp.spatial_vectorf), out: wp.array(dtype=wp.float32, ndim=2), ): + """Root angular velocity [rad/s] in the body frame, derived inline from the root + + pose and CoM velocity (Tier-1 access; avoids non-capturable lazy buffers). + """ i = wp.tid() v = body_ang_vel_from_root(root_pose_w[i], root_vel_w[i]) out[i, 0] = v[0] @@ -166,6 +182,7 @@ def _projected_gravity_kernel( gravity_w: wp.array(dtype=wp.vec3f), out: wp.array(dtype=wp.float32, ndim=2), ): + """Unit gravity direction rotated into the body frame (flat pose -> (0, 0, -1)).""" i = wp.tid() g = rotate_vec_to_body_frame(wp.normalize(gravity_w[i]), root_pose_w[i]) out[i, 0] = g[0] @@ -219,6 +236,10 @@ def _joint_rel_gather_kernel( joint_ids: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.float32, ndim=2), ): + """Gather selected joints' values relative to their per-env defaults into + + contiguous term columns. Launch with dim=(num_envs, num_selected_joints). + """ env_id, k = wp.tid() j = joint_ids[k] out[env_id, k] = values[env_id, j] - defaults[env_id, j] @@ -255,6 +276,10 @@ def _joint_pos_limit_normalized_kernel( joint_ids: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.float32, ndim=2), ): + """Selected joints' positions normalized to [-1, 1] within their soft limits. + + Launch with dim=(num_envs, num_selected_joints). + """ env_id, k = wp.tid() j = joint_ids[k] pos = joint_pos[env_id, j] @@ -353,17 +378,91 @@ def generated_commands(env: ManagerBasedEnv, out, command_name: str) -> None: """The generated command from the command manager. Writes into ``out``. Warp-first override of :func:`isaaclab.envs.mdp.observations.generated_commands`. - Uses ``wp.from_torch`` to create a zero-copy warp view of the command tensor on first call. + Reads the command manager's pointer-stable Warp storage directly. + """ + wp.copy(out, env.command_manager.get_command_wp(command_name)) + + +""" +Sensors. +""" + + +@wp.kernel +def _body_incoming_wrench_kernel( + force: wp.array(dtype=wp.vec3f, ndim=2), + torque: wp.array(dtype=wp.vec3f, ndim=2), + body_ids: wp.array(dtype=wp.int32), + out: wp.array(dtype=wp.float32, ndim=2), +): + """Gather selected bodies' incoming force [N] and torque [N·m] into per-body + + ``[fx, fy, fz, tx, ty, tz]`` column blocks. + """ + env_id = wp.tid() + for k in range(body_ids.shape[0]): + b = body_ids[k] + f = force[env_id, b] + t = torque[env_id, b] + out[env_id, k * 6 + 0] = f[0] + out[env_id, k * 6 + 1] = f[1] + out[env_id, k * 6 + 2] = f[2] + out[env_id, k * 6 + 3] = t[0] + out[env_id, k * 6 + 4] = t[1] + out[env_id, k * 6 + 5] = t[2] + + +@generic_io_descriptor_warp(out_dim="body:6", observation_type="BodyState", on_inspect=[record_shape, record_dtype]) +def body_incoming_wrench(env: ManagerBasedEnv, out, sensor_cfg: SceneEntityCfg) -> None: + """Incoming spatial wrench [N, N·m] on selected bodies, laid out per body as ``[fx, fy, fz, tx, ty, tz]``. + + Warp-first override of :func:`isaaclab.envs.mdp.observations.body_incoming_wrench`. Reads the + Newton joint-wrench sensor and gathers the force/torque of the bodies selected by ``sensor_cfg`` + into ``out`` (shape ``(num_envs, 6 * num_selected_bodies)``). + """ + sensor = env.scene.sensors[sensor_cfg.name] + if sensor.data.force is None or sensor.data.torque is None: + raise RuntimeError("Joint wrench sensor data is not initialized. Call sim.reset() before reading observations.") + wp.launch( + kernel=_body_incoming_wrench_kernel, + dim=env.num_envs, + inputs=[sensor.data.force.warp, sensor.data.torque.warp, sensor_cfg.body_ids_wp, out], + device=env.device, + ) + + +@wp.kernel +def _height_scan_kernel( + pos_w: wp.array(dtype=wp.vec3f), + ray_hits_w: wp.array2d(dtype=wp.vec3f), + offset: wp.float32, + out: wp.array2d(dtype=wp.float32), +): + """Height of the scan frame above each ray hit [m]: sensor_z - hit_z - offset. + + Launch with dim=(num_envs, num_rays). Missed hits carry ``inf`` and produce + ``-inf`` heights, matching the stable term (handled by the term's clip). + """ + env_id, ray_id = wp.tid() + out[env_id, ray_id] = pos_w[env_id][2] - ray_hits_w[env_id, ray_id][2] - offset + + +# Sensor reads go through the lazy-update path, whose host-side timestamp bookkeeping +# decides when rays are re-cast and has not been audited for graph capture. +@WarpCapturable(False, reason="sensor lazy-update host bookkeeping is not capture-audited (Part 3 scope)") +@generic_io_descriptor_warp(units="m", out_dim="sensor:rays", observation_type="SensorState", on_inspect=[record_shape]) +def height_scan(env: ManagerBasedEnv, out, sensor_cfg: SceneEntityCfg, offset: float = 0.5) -> None: + """Height scan from the given sensor w.r.t. the sensor's frame [m]. Writes into ``out``. + + Warp-first override of :func:`isaaclab.envs.mdp.observations.height_scan`. + The provided offset (Defaults to 0.5) is subtracted from the returned values. + Missed rays carry ``inf`` hit positions, matching the stable term's semantics + (the resulting ``-inf`` heights are handled by the term's ``clip`` setting). """ - # TODO(warp-migration): Cross-manager access (observation → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - fn = generated_commands - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - fn._cmd_wp = cmd - else: - fn._cmd_wp = wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True - wp.copy(out, fn._cmd_wp) + sensor: RayCaster = env.scene.sensors[sensor_cfg.name] + wp.launch( + kernel=_height_scan_kernel, + dim=(env.num_envs, sensor.num_rays), + inputs=[sensor.data.pos_w.warp, sensor.data.ray_hits_w.warp, float(offset), out], + device=env.device, + ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index 338d27c99c5e..833a656a9a45 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -38,6 +38,7 @@ @wp.kernel def _is_alive_kernel(terminated: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32)): + """1.0 for envs still running, 0.0 for terminated ones.""" i = wp.tid() out[i] = wp.where(terminated[i], 0.0, 1.0) @@ -54,6 +55,7 @@ def is_alive(env: ManagerBasedRLEnv, out: wp.array(dtype=wp.float32)) -> None: @wp.kernel def _is_terminated_kernel(terminated: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32)): + """1.0 for envs terminated by a non-timeout cause, else 0.0.""" i = wp.tid() out[i] = wp.where(terminated[i], 1.0, 0.0) @@ -87,6 +89,10 @@ def _lin_vel_z_l2_kernel( root_vel_w: wp.array(dtype=wp.spatial_vectorf), out: wp.array(dtype=wp.float32), ): + """Squared body-frame vertical velocity [(m/s)^2] per env, derived inline from + + the root state. + """ i = wp.tid() vz = body_lin_vel_from_root(root_pose_w[i], root_vel_w[i])[2] out[i] = vz * vz @@ -109,6 +115,7 @@ def _ang_vel_xy_l2_kernel( root_vel_w: wp.array(dtype=wp.spatial_vectorf), out: wp.array(dtype=wp.float32), ): + """Sum of squared body-frame roll/pitch rates [(rad/s)^2] per env.""" i = wp.tid() v = body_ang_vel_from_root(root_pose_w[i], root_vel_w[i]) out[i] = v[0] * v[0] + v[1] * v[1] @@ -131,6 +138,10 @@ def _flat_orientation_l2_kernel( gravity_w: wp.array(dtype=wp.vec3f), out: wp.array(dtype=wp.float32), ): + """Sum of squared xy components of body-frame projected gravity (zero when the + + base is perfectly flat). + """ # ``gravity_w`` is per-env and may carry magnitude (Newton, m/s^2), so index # per env and normalize before projecting. i = wp.tid() @@ -160,6 +171,7 @@ def flat_orientation_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = def _sum_sq_masked_kernel( x: wp.array(dtype=wp.float32, ndim=2), joint_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32) ): + """Per-env sum of squares over mask-selected joint columns.""" i = wp.tid() s = float(0.0) for j in range(x.shape[1]): @@ -184,6 +196,7 @@ def joint_torques_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Sc def _sum_abs_masked_kernel( x: wp.array(dtype=wp.float32, ndim=2), joint_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32) ): + """Per-env L1 sum over mask-selected joint columns.""" i = wp.tid() s = float(0.0) for j in range(x.shape[1]): @@ -233,6 +246,7 @@ def _sum_abs_diff_masked_kernel( joint_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32), ): + """Per-env L1 sum of ``a - b`` over mask-selected joint columns.""" i = wp.tid() s = float(0.0) for j in range(a.shape[1]): @@ -260,6 +274,11 @@ def _joint_pos_limits_kernel( joint_mask: wp.array(dtype=wp.bool), out: wp.array(dtype=wp.float32), ): + """Per-env sum of soft-limit violations [rad or m, depending on joint type]: + + distance each selected joint sits below its lower or above its upper soft + limit (zero inside the limits). + """ i = wp.tid() s = float(0.0) for j in range(joint_pos.shape[1]): @@ -302,6 +321,7 @@ def _sum_sq_diff_2d_kernel( b: wp.array(dtype=wp.float32, ndim=2), out: wp.array(dtype=wp.float32), ): + """Per-env sum of squared differences across all columns (action-rate penalty).""" i = wp.tid() s = float(0.0) for j in range(a.shape[1]): @@ -323,6 +343,7 @@ def action_rate_l2(env: ManagerBasedRLEnv, out) -> None: # TODO(warp-migration): Revisit 2D kernel + wp.atomic_add vs 1D inner loop. @wp.kernel def _sum_sq_2d_kernel(x: wp.array(dtype=wp.float32, ndim=2), out: wp.array(dtype=wp.float32)): + """Per-env sum of squares across all columns (action-magnitude penalty).""" i = wp.tid() s = float(0.0) for j in range(x.shape[1]): @@ -395,6 +416,10 @@ def _track_lin_vel_xy_exp_kernel( std_sq_inv: float, out: wp.array(dtype=wp.float32), ): + """Exponential tracking shaping ``exp(-err^2 / std^2)`` of planar body-frame + + linear velocity vs the command's (vx, vy). + """ i = wp.tid() v = body_lin_vel_from_root(root_pose_w[i], root_vel_w[i]) dx = command[i, 0] - v[0] @@ -415,24 +440,13 @@ def track_lin_vel_xy_exp( Warp-first override of :func:`isaaclab.envs.mdp.rewards.track_lin_vel_xy_exp`. """ asset: Articulation = env.scene[asset_cfg.name] - # cache the warp view of the command tensor on first call (zero-copy) - # TODO(warp-migration): Cross-manager access (reward → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - if not getattr(track_lin_vel_xy_exp, "_is_warmed_up", False) or track_lin_vel_xy_exp._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - track_lin_vel_xy_exp._cmd_wp = cmd - else: - track_lin_vel_xy_exp._cmd_wp = wp.from_torch(cmd) - track_lin_vel_xy_exp._cmd_name = command_name - track_lin_vel_xy_exp._is_warmed_up = True wp.launch( kernel=_track_lin_vel_xy_exp_kernel, dim=env.num_envs, inputs=[ asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, - track_lin_vel_xy_exp._cmd_wp, + env.command_manager.get_command_wp(command_name), 1.0 / (std * std), out, ], @@ -449,6 +463,10 @@ def _track_ang_vel_z_exp_kernel( std_sq_inv: float, out: wp.array(dtype=wp.float32), ): + """Exponential tracking shaping ``exp(-err^2 / std^2)`` of body-frame yaw rate + + vs the command's wz. + """ i = wp.tid() dz = command[i, cmd_col] - body_ang_vel_from_root(root_pose_w[i], root_vel_w[i])[2] out[i] = wp.exp(-dz * dz * std_sq_inv) @@ -466,23 +484,13 @@ def track_ang_vel_z_exp( Warp-first override of :func:`isaaclab.envs.mdp.rewards.track_ang_vel_z_exp`. """ asset: Articulation = env.scene[asset_cfg.name] - # TODO(warp-migration): Cross-manager access (reward → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - if not getattr(track_ang_vel_z_exp, "_is_warmed_up", False) or track_ang_vel_z_exp._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - track_ang_vel_z_exp._cmd_wp = cmd - else: - track_ang_vel_z_exp._cmd_wp = wp.from_torch(cmd) - track_ang_vel_z_exp._cmd_name = command_name - track_ang_vel_z_exp._is_warmed_up = True wp.launch( kernel=_track_ang_vel_z_exp_kernel, dim=env.num_envs, inputs=[ asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, - track_ang_vel_z_exp._cmd_wp, + env.command_manager.get_command_wp(command_name), 2, 1.0 / (std * std), out, diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.pyi index 18431a745b31..29c0dfa1b494 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/__init__.pyi @@ -8,6 +8,7 @@ __all__ = [ "ActionTerm", "CommandManager", "CommandTerm", + "CurriculumManager", "EventManager", "ManagerBase", "ManagerTermBase", @@ -31,6 +32,7 @@ from isaaclab.managers import * # noqa: F401, F403 from .action_manager import ActionManager, ActionTerm from .command_manager import CommandManager, CommandTerm +from .curriculum_manager import CurriculumManager from .event_manager import EventManager from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import ( diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py index 672af82bc8ab..751cc6740533 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py @@ -20,9 +20,9 @@ from isaaclab.assets import AssetBase from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor +from isaaclab.managers.manager_term_cfg import ActionTermCfg from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import ActionTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -373,32 +373,21 @@ def reset( Returns: An empty dictionary. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "ActionManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) # reset the action history - if env_mask is None: - self._prev_action.fill_(0.0) - self._action.fill_(0.0) - else: - wp.launch( - kernel=_zero_masked_2d, - dim=(self.num_envs, self.total_action_dim), - inputs=[env_mask, self._prev_action], - device=self.device, - ) - wp.launch( - kernel=_zero_masked_2d, - dim=(self.num_envs, self.total_action_dim), - inputs=[env_mask, self._action], - device=self.device, - ) + wp.launch( + kernel=_zero_masked_2d, + dim=(self.num_envs, self.total_action_dim), + inputs=[env_mask, self._prev_action], + device=self.device, + ) + wp.launch( + kernel=_zero_masked_2d, + dim=(self.num_envs, self.total_action_dim), + inputs=[env_mask, self._action], + device=self.device, + ) # reset all action terms for term in self._terms.values(): @@ -484,6 +473,7 @@ def _prepare_terms(self): ) # create the action term term = term_cfg.class_type(term_cfg, self._env) + self._register_term_capturability(term_cfg.class_type) # sanity check if term is valid type if not isinstance(term, ActionTerm): raise TypeError(f"Returned object for the term '{term_name}' is not of type ActionType.") diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index 8b4e8f83dbe7..6a459dcf9b66 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -8,7 +8,6 @@ from __future__ import annotations import inspect -import weakref from abc import abstractmethod from collections.abc import Sequence from typing import TYPE_CHECKING @@ -17,10 +16,11 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import CommandTermCfg + from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import CommandTermCfg # import omni.kit.app @@ -36,6 +36,11 @@ def _sum_and_zero_masked( metric: wp.array(dtype=wp.float32), out_mean: wp.array(dtype=wp.float32), ): + """Fold selected envs' episode metric into a running mean and zero it for the next episode. + + ``scale[0]`` carries the 1/count normalization; the atomic add reduces across all + selected envs into ``out_mean[0]``. + """ env_id = wp.tid() if mask[env_id]: wp.atomic_add(out_mean, 0, metric[env_id] * scale[0]) @@ -44,6 +49,7 @@ def _sum_and_zero_masked( @wp.kernel def _zero_counter_masked(mask: wp.array(dtype=wp.bool), counter: wp.array(dtype=wp.int32)): + """Zero the per-env resample counter for envs selected by ``mask``.""" env_id = wp.tid() if mask[env_id]: counter[env_id] = 0 @@ -55,6 +61,7 @@ def _step_time_left_and_build_resample_mask( dt: wp.float32, out_mask: wp.array(dtype=wp.bool), ): + """Advance each env's resample countdown by ``dt`` [s] and flag expired envs in ``out_mask``.""" env_id = wp.tid() t = time_left[env_id] - dt time_left[env_id] = t @@ -70,6 +77,11 @@ def _resample_time_left_and_increment_counter( lower: wp.float32, upper: wp.float32, ): + """Draw a new resampling interval [s] for flagged envs and count the resample. + + Uses the per-env device RNG state with write-back so the random sequence + advances across CUDA-graph replays. + """ env_id = wp.tid() if mask[env_id]: s = rng_state[env_id] @@ -124,9 +136,11 @@ def __init__(self, cfg: CommandTermCfg, env: ManagerBasedRLEnv): def __del__(self): """Unsubscribe from the callbacks.""" - if self._debug_vis_handle: - self._debug_vis_handle.unsubscribe() - self._debug_vis_handle = None + env = getattr(self, "_env", None) + sim = getattr(env, "sim", None) + registry = getattr(sim, "vis_marker_registry", None) + if registry is not None: + registry.clear_debug_vis_callback(self) """ Properties @@ -138,6 +152,21 @@ def command(self) -> torch.Tensor | wp.array: """The command tensor. Shape is (num_envs, command_dim).""" raise NotImplementedError + @property + def command_wp(self) -> wp.array(dtype=wp.float32, ndim=2): + """Pointer-stable Warp command storage with shape ``(num_envs, command_dim)``. + + Existing custom terms that only implement :attr:`command` receive a + lazily cached zero-copy Warp view. Warp-native terms may override this + property to return their owner-held array directly. + """ + command = self.command + if isinstance(command, wp.array): + return command + if not hasattr(self, "_command_wp_compat"): + self._command_wp_compat = wp.from_torch(command, dtype=wp.float32) + return self._command_wp_compat + @property def has_debug_vis_implementation(self) -> bool: """Whether the command generator has a debug visualization implemented.""" @@ -171,25 +200,10 @@ def set_debug_vis(self, debug_vis: bool) -> bool: self._set_debug_vis_impl(debug_vis) # toggle debug visualization handles if debug_vis: - # only enable debug_vis if omniverse is available - from isaaclab.sim.simulation_context import SimulationContext - - sim_context = SimulationContext.instance() - if not sim_context.has_omniverse_visualizer(): - return False - # create a subscriber for the post update event if it doesn't exist if self._debug_vis_handle is None: - import omni.kit.app - - app_interface = omni.kit.app.get_app_interface() - self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( - lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) - ) + self._debug_vis_handle = self._env.sim.vis_marker_registry.add_debug_vis_callback(self) else: - # remove the subscriber if it exists - if self._debug_vis_handle is not None: - self._debug_vis_handle.unsubscribe() - self._debug_vis_handle = None + self._env.sim.vis_marker_registry.clear_debug_vis_callback(self) # return success return True @@ -213,14 +227,7 @@ def reset( Returns: A dictionary containing the information to log under the "{name}" key. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "CommandTerm.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) # compute selected count and reset scale self._reset_count_wp.zero_() @@ -410,9 +417,14 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): # reset logging extras (persistent holder for orchestrator aggregation) self._reset_extras: dict[str, torch.Tensor] = {} + success_term_count = sum("success_rate" in term.reset_extras for term in self._terms.values()) for term_name, term in self._terms.items(): for metric_name, metric_value in term.reset_extras.items(): - self._reset_extras[f"Metrics/{term_name}/{metric_name}"] = metric_value + if metric_name == "success_rate" and success_term_count == 1: + metric_key = "Metrics/success_rate" + else: + metric_key = f"Metrics/{term_name}/{metric_name}" + self._reset_extras[metric_key] = metric_value def __str__(self) -> str: """Returns: A string representation for the command manager.""" @@ -513,14 +525,7 @@ def reset( Returns: A dictionary containing the information to log under the "Metrics/{term_name}/{metric_name}" key. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "CommandManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) for term in self._terms.values(): # reset the command term @@ -556,6 +561,17 @@ def get_command(self, name: str) -> torch.Tensor: return wp.to_torch(command) return command + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32, ndim=2): + """Return pointer-stable Warp storage for a command term. + + Args: + name: The name of the command term. + + Returns: + The command array with shape ``(num_envs, command_dim)``. + """ + return self._terms[name].command_wp + def get_term(self, name: str) -> CommandTerm: """Returns the command term with the specified name. @@ -590,6 +606,7 @@ def _prepare_terms(self): ) # create the action term term = term_cfg.class_type(term_cfg, self._env) + self._register_term_capturability(term_cfg.class_type) # sanity check if term is valid type if not isinstance(term, CommandTerm): raise TypeError(f"Returned object for the term '{term_name}' is not of type CommandType.") diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py new file mode 100644 index 000000000000..6c358e1c061e --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/curriculum_manager.py @@ -0,0 +1,247 @@ +# 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 + +"""Warp-first curriculum manager with mask-native term dispatch.""" + +from __future__ import annotations + +import inspect +from collections.abc import Sequence +from numbers import Real +from typing import TYPE_CHECKING + +import torch +import warp as wp +from prettytable import PrettyTable + +from isaaclab.managers import CurriculumTermCfg as StableCurriculumTermCfg +from isaaclab.managers import ManagerTermBase as StableManagerTermBase +from isaaclab.utils import string_to_callable + +from .manager_base import ManagerBase, ManagerTermBase + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +class CurriculumManager(ManagerBase): + """Manage curriculum terms that write state into persistent Warp scalar outputs.""" + + _env: ManagerBasedRLEnv + """The environment instance.""" + + def __init__(self, cfg: object, env: ManagerBasedRLEnv): + """Initialize the curriculum manager. + + Args: + cfg: Configuration object or dictionary of curriculum terms. + env: Environment instance. + """ + self._term_names: list[str] = [] + self._term_cfgs: list[StableCurriculumTermCfg] = [] + self._term_modes: list[str] = [] + + super().__init__(cfg, env) + + num_terms = len(self._term_names) + self._term_states_wp = wp.zeros(num_terms, dtype=wp.float32, device=self.device) + self._term_state_views_wp: list[wp.array] = [] + if num_terms > 0: + stride = self._term_states_wp.strides[0] + for term_idx, term_cfg in enumerate(self._term_cfgs): + out_view = wp.array( + ptr=self._term_states_wp.ptr + term_idx * stride, + dtype=wp.float32, + shape=(1,), + strides=(stride,), + device=self.device, + ) + self._term_state_views_wp.append(out_view) + term_cfg.out = out_view + + self._term_states = wp.to_torch(self._term_states_wp) + self._reset_extras = { + f"Curriculum/{term_name}": self._term_states[term_idx] + for term_idx, term_name in enumerate(self._term_names) + } + + def __str__(self) -> str: + """Return a string representation of the curriculum manager.""" + msg = f" contains {len(self._term_names)} active terms.\n" + table = PrettyTable() + table.title = "Active Curriculum Terms" + table.field_names = ["Index", "Name"] + table.align["Name"] = "l" + for index, name in enumerate(self._term_names): + table.add_row([index, name]) + return msg + table.get_string() + "\n" + + @property + def active_terms(self) -> list[str]: + """Names of active curriculum terms.""" + return self._term_names + + @property + def reset_extras(self) -> dict[str, torch.Tensor]: + """Persistent scalar views containing the latest curriculum states.""" + return self._reset_extras + + def reset( + self, + env_mask: wp.array(dtype=wp.bool), + env_ids: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + """Reset selected class terms and return persistent curriculum logging outputs. + + Args: + env_mask: Boolean Warp mask selecting environments to reset. + env_ids: Precomputed compact environment IDs matching :paramref:`env_mask`. + When omitted, legacy terms materialize IDs once at this boundary. + + Returns: + Persistent scalar curriculum states keyed by their logging paths. + """ + env_mask = self._resolve_reset_mask(None, env_mask) + # Dispatch by term kind: Warp class terms consume the mask directly; stable + # class terms receive compact IDs, materialized at most once per call. + compact_env_ids = env_ids + for term_cfg, mode in zip(self._term_cfgs, self._term_modes): + if isinstance(term_cfg.func, ManagerTermBase): + term_cfg.func.reset(env_mask=env_mask) + elif isinstance(term_cfg.func, StableManagerTermBase): + if compact_env_ids is None: + compact_env_ids = self._compact_legacy_env_ids(env_mask) + term_cfg.func.reset(env_ids=compact_env_ids) + return self._reset_extras + + @property + def requires_host_ids(self) -> bool: + """Whether any active legacy term genuinely consumes compact environment IDs.""" + return "legacy" in self._term_modes + + @property + def requires_host_boundary(self) -> bool: + """Whether any active legacy term must run only when a reset occurs.""" + return any(mode != "mask" for mode in self._term_modes) + + def compute( + self, + env_mask: wp.array(dtype=wp.bool), + env_ids: torch.Tensor | None = None, + ) -> None: + """Update curriculum terms for selected environments. + + Args: + env_mask: Boolean Warp mask selecting environments to update. + env_ids: Precomputed compact environment IDs matching :paramref:`env_mask`. + When omitted, legacy terms materialize IDs once at this boundary. + """ + env_mask = self._resolve_reset_mask(None, env_mask) + # Logging states are recomputed every call; mask-native terms write their + # scalar slot on-device through the pointer-stable ``term_cfg.out`` view. + self._term_states_wp.zero_() + compact_env_ids = env_ids + # Term modes: "mask" = Warp-native ``(env, env_mask, out)``; + # "legacy" = stable ``(env, env_ids)`` with compact IDs materialized at most once. + for term_idx, (term_cfg, mode) in enumerate(zip(self._term_cfgs, self._term_modes)): + if mode == "mask": + term_cfg.func(self._env, env_mask, term_cfg.out, **term_cfg.params) + continue + if compact_env_ids is None: + compact_env_ids = self._compact_legacy_env_ids(env_mask) + state = term_cfg.func(self._env, compact_env_ids, **term_cfg.params) + if state is not None: + if isinstance(state, torch.Tensor) and state.numel() != 1: + raise TypeError( + f"Curriculum term '{self._term_names[term_idx]}' must return a scalar state;" + f" received tensor shape {tuple(state.shape)}." + ) + if not isinstance(state, (Real, torch.Tensor)): + raise TypeError( + f"Curriculum term '{self._term_names[term_idx]}' returned {type(state).__name__}." + " Warp CurriculumManager supports scalar logging states only." + ) + self._term_states[term_idx] = state + + def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: + """Return curriculum states for debug inspection. + + Args: + env_idx: Unused environment index retained for the manager interface. + + Returns: + Curriculum term names and their scalar states. + """ + del env_idx + states = self._term_states.detach().cpu().tolist() + return [(name, [states[index]]) for index, name in enumerate(self._term_names)] + + def _prepare_terms(self): + cfg_items = self.cfg.items() if isinstance(self.cfg, dict) else self.cfg.__dict__.items() + for term_name, term_cfg in cfg_items: + if term_cfg is None: + continue + if not isinstance(term_cfg, StableCurriculumTermCfg): + raise TypeError( + f"Configuration for the term '{term_name}' is not of type CurriculumTermCfg." + f" Received: '{type(term_cfg)}'." + ) + if isinstance(term_cfg.func, str): + term_cfg.func = string_to_callable(term_cfg.func) + if self._is_mask_term(term_cfg.func): + mode = "mask" + self._resolve_common_term_cfg(term_name, term_cfg, min_argc=3) + else: + mode = "legacy" + self._resolve_legacy_term_cfg(term_name, term_cfg) + self._term_names.append(term_name) + self._term_cfgs.append(term_cfg) + self._term_modes.append(mode) + + @staticmethod + def _is_mask_term(func) -> bool: + """Return whether a term follows ``(env, env_mask, out, ...)``.""" + func_static = func.__call__ if inspect.isclass(func) else func + parameters = list(inspect.signature(func_static).parameters) + if inspect.isclass(func): + parameters = parameters[1:] + return parameters[:3] == ["env", "env_mask", "out"] + + def _resolve_legacy_term_cfg(self, term_name: str, term_cfg: StableCurriculumTermCfg) -> None: + """Validate and initialize a stable ``(env, env_ids, ...)`` curriculum term.""" + if not callable(term_cfg.func): + raise AttributeError(f"The term '{term_name}' is not callable. Received: {term_cfg.func}") + is_class = inspect.isclass(term_cfg.func) + func_static = term_cfg.func.__call__ if is_class else term_cfg.func + min_argc = 3 if is_class else 2 + signature = inspect.signature(func_static) + parameters = list(signature.parameters.values()) + for parameter in parameters[min_argc:]: + if ( + parameter.default is not inspect.Parameter.empty + and parameter.name not in term_cfg.params + and hasattr(parameter.default, "__dataclass_fields__") + ): + term_cfg.params[parameter.name] = parameter.default.copy() + required = { + parameter.name for parameter in parameters[min_argc:] if parameter.default is inspect.Parameter.empty + } + accepted = {parameter.name for parameter in parameters[min_argc:]} + provided = set(term_cfg.params) + if not required.issubset(provided) or not provided.issubset(accepted): + raise ValueError( + f"The legacy curriculum term '{term_name}' expects parameters {sorted(accepted)}," + f" but received {sorted(provided)}." + ) + graph_cache = getattr(self._env, "_warp_graph_cache", None) + if graph_cache is not None: + graph_cache.register_capturability(type(self).__name__, False) + if self._env.sim.is_playing(): + self._process_term_cfg_at_play(term_name, term_cfg) + + @staticmethod + def _compact_legacy_env_ids(env_mask: wp.array(dtype=wp.bool)) -> torch.Tensor: + """Materialize compact Torch IDs at the legacy curriculum boundary.""" + return wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py index ab8a0264ea9b..b00ffc6daf92 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py @@ -29,8 +29,9 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import EventTermCfg + from .manager_base import ManagerBase -from .manager_term_cfg import EventTermCfg logger = logging.getLogger(__name__) @@ -42,6 +43,7 @@ def _interval_init_per_env( lower: wp.float32, upper: wp.float32, ): + """Seed each env's interval-event countdown with a uniform draw from [``lower``, ``upper``] [s].""" env_id = wp.tid() s = rng_state[env_id] time_left[env_id] = wp.randf(s, lower, upper) @@ -55,6 +57,7 @@ def _interval_init_global( lower: wp.float32, upper: wp.float32, ): + """Seed the single shared countdown of a ``is_global_time`` interval event.""" # single element s = rng_state[0] time_left[0] = wp.randf(s, lower, upper) @@ -70,6 +73,10 @@ def _interval_step_per_env( lower: wp.float32, upper: wp.float32, ): + """Tick each env's interval countdown by ``dt`` [s]; on expiry, flag the env in + + ``trigger_mask`` and immediately draw its next interval (RNG state written back). + """ env_id = wp.tid() t = time_left[env_id] - dt if t < wp.float32(1.0e-6): @@ -91,6 +98,10 @@ def _interval_step_global( lower: wp.float32, upper: wp.float32, ): + """Tick the shared countdown of a ``is_global_time`` interval event; on expiry set + + ``trigger_flag[0]`` and draw the next interval. + """ t = time_left[0] - dt if t < wp.float32(1.0e-6): trigger_flag[0] = True @@ -110,6 +121,7 @@ def _interval_reset_selected( lower: wp.float32, upper: wp.float32, ): + """Re-draw interval countdowns for envs selected by ``env_mask`` (episode-reset path).""" env_id = wp.tid() if env_mask[env_id]: s = rng_state[env_id] @@ -122,6 +134,10 @@ def _seed_global_rng_from_env_rng( env_rng_state: wp.array(dtype=wp.uint32), global_rng_state: wp.array(dtype=wp.uint32), ): + """Derive the global-event RNG stream from env 0's per-env state so global draws + + stay seed-reproducible without a host round-trip. + """ global_rng_state[0] = wp.rand_init(wp.int32(env_rng_state[0]), wp.int32(0)) @@ -134,6 +150,13 @@ def _reset_compute_valid_mask( global_step_count_buf: wp.array(dtype=wp.int32), min_step_count: wp.int32, ): + """Gate a reset-event mask by trigger spacing entirely on device. + + An env stays selected only if at least ``min_step_count`` steps passed since it + last triggered, or it has never triggered (stable-parity: first reset always + fires). Passing envs update ``last_triggered_step``/``triggered_once`` in the + same launch, keeping the spacing bookkeeping replay-safe. + """ env_id = wp.tid() if not in_mask[env_id]: out_mask[env_id] = False @@ -243,15 +266,7 @@ def reset( *, env_mask: wp.array | torch.Tensor | None = None, ) -> dict[str, float]: - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - # Keep all id->mask resolution strictly outside capture. - if wp.get_device().is_capturing: - raise RuntimeError( - "EventManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) # reset class terms (mask-based) for mode_cfg in self._mode_class_term_cfgs.values(): diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 6caea3eebcbb..bb2decbd77e8 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -20,29 +20,54 @@ import inspect import logging from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any import warp as wp import isaaclab.utils.string as string_utils +from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg from isaaclab.utils import class_to_dict, string_to_callable -from isaaclab_experimental.utils.warp import is_warp_capturable +from isaaclab_experimental.utils.warp import WarpCapturable -from .manager_term_cfg import ManagerTermBaseCfg from .scene_entity_cfg import SceneEntityCfg # import omni.timeline if TYPE_CHECKING: + import torch + from isaaclab.envs import ManagerBasedEnv # import logger logger = logging.getLogger(__name__) +def _resolve_reset_mask( + owner: ManagerBase | ManagerTermBase, + env_ids: Sequence[int] | None, + env_mask: wp.array(dtype=wp.bool) | None, +) -> wp.array(dtype=wp.bool): + """Resolve reset input once while keeping ID conversion outside capture.""" + if isinstance(env_mask, wp.array): + if env_mask.dtype != wp.bool or env_mask.ndim != 1 or env_mask.shape[0] != owner.num_envs: + raise ValueError( + f"env_mask must be a Warp boolean array with shape ({owner.num_envs},), received {env_mask}." + ) + expected_device = wp.get_device(owner.device) + if env_mask.device != expected_device: + raise ValueError(f"env_mask must be on {expected_device}, received {env_mask.device}.") + return env_mask + if wp.get_device(owner.device).is_capturing: + raise RuntimeError( + f"{type(owner).__name__}.reset requires env_mask(wp.array[bool]) during capture. " + "Resolve environment IDs before entering the captured stage." + ) + return owner._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + + class ManagerTermBase(ABC): """Base class for manager terms. @@ -107,14 +132,28 @@ def __name__(self) -> str: Operations. """ - def reset(self, env_mask: wp.array | None = None) -> None: + def reset(self, env_mask: wp.array | None = None) -> dict[str, torch.Tensor] | None: """Resets the manager term (mask-based). Args: env_mask: Boolean mask of shape (num_envs,) indicating which envs to reset. If None, all envs are considered. + + Returns: + An optional dictionary of logging values to merge into the owning manager's + reset extras. To stay CUDA-graph capturable, values must be persistent + tensor views over device buffers written by kernels (no host readback); + the same objects must be returned on every call. """ - pass + return None + + def _resolve_reset_mask( + self, + env_ids: Sequence[int] | None, + env_mask: wp.array(dtype=wp.bool) | None, + ) -> wp.array(dtype=wp.bool): + """Resolve a legacy reset selection to the canonical Warp mask.""" + return _resolve_reset_mask(self, env_ids, env_mask) def serialize(self) -> dict: """General serialization call. Includes the configuration dict.""" @@ -226,7 +265,11 @@ def active_terms(self) -> list[str] | dict[str, list[str]]: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> dict[str, float]: + def reset( + self, + env_ids: Sequence[int] | None = None, + env_mask: wp.array(dtype=wp.bool) | None = None, + ) -> dict[str, float]: """Resets the manager and returns logging information for the current time-step. Args: @@ -238,6 +281,14 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None """ return {} + def _resolve_reset_mask( + self, + env_ids: Sequence[int] | None, + env_mask: wp.array(dtype=wp.bool) | None, + ) -> wp.array(dtype=wp.bool): + """Resolve a legacy reset selection to the canonical Warp mask.""" + return _resolve_reset_mask(self, env_ids, env_mask) + def find_terms(self, name_keys: str | Sequence[str]) -> list[str]: """Find terms in the manager based on the names. @@ -403,17 +454,24 @@ def _resolve_common_term_cfg(self, term_name: str, term_cfg: ManagerTermBaseCfg, f" and optional parameters: {args_with_defaults}, but received: {term_params}." ) - # register non-capturable terms with the call switch for mode=2 fallback - if not is_warp_capturable(term_cfg.func): - switch = getattr(self._env, "_manager_call_switch", None) - if switch is not None: - switch.register_manager_capturability(type(self).__name__, False) + self._register_term_capturability(term_cfg.func) # process attributes at runtime # these properties are only resolvable once the simulation starts playing if self._env.sim.is_playing(): self._process_term_cfg_at_play(term_name, term_cfg) + def _register_term_capturability(self, term: Callable) -> None: + """Keep the complete manager eager when a configured term is unsafe.""" + # TODO(#6611): granularity is whole-manager — one non-capturable term forces the + # entire manager eager. Per-term record/replay could keep the capturable terms + # on graphs while only the unsafe term runs eagerly; that is execution-layer + # scope and deliberately not part of the mask-first PR. + if not WarpCapturable.is_capturable(term): + graph_cache = getattr(self._env, "_warp_graph_cache", None) + if graph_cache is not None: + graph_cache.register_capturability(type(self).__name__, False) + def _process_term_cfg_at_play(self, term_name: str, term_cfg: ManagerTermBaseCfg): """Process the term configuration at runtime. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py index ce354be66a2a..9dad78dc25c9 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_term_cfg.py @@ -3,92 +3,50 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Configuration terms for different managers (experimental, Warp-first). +"""Manager term configuration classes used by the Warp-first managers. -This module is a passthrough to :mod:`isaaclab.managers.manager_term_cfg` except for -the following term configs which are overridden for Warp-first execution: +Passthrough to :mod:`isaaclab.managers.manager_term_cfg`: the Warp-first managers +accept the same configuration shapes as the stable managers, so re-exporting the +stable classes preserves type identity and lets a stable task configuration be +adapted in place (by the warp frontend) without rebuilding every term. At +runtime, the adapted term callables still use the Warp-first +``func(env, out, **params) -> None`` signature. -- :class:`ObservationTermCfg` -- :class:`RewardTermCfg` -- :class:`TerminationTermCfg` +:class:`CurriculumTermCfg` is the one override: it extends the stable class with +the Warp-first curriculum term configuration. """ from __future__ import annotations -from collections.abc import Callable -from dataclasses import MISSING - -from isaaclab.managers.manager_term_cfg import * # noqa: F401,F403 -from isaaclab.managers.manager_term_cfg import ManagerTermBaseCfg as _ManagerTermBaseCfg +from isaaclab.managers.manager_term_cfg import ( + ActionTermCfg, + CommandTermCfg, + EventTermCfg, + ManagerTermBaseCfg, + ObservationGroupCfg, + ObservationTermCfg, + RecorderTermCfg, + RewardTermCfg, + TerminationTermCfg, +) +from isaaclab.managers.manager_term_cfg import CurriculumTermCfg as _CurriculumTermCfg from isaaclab.utils.configclass import configclass @configclass -class RewardTermCfg(_ManagerTermBaseCfg): - """Configuration for a reward term. - - The function is expected to write the (unweighted) reward values into a - pre-allocated Warp buffer provided by the manager. - - Expected signature: - - - ``func(env, out, **params) -> None`` - - where ``out`` is a Warp array of shape ``(num_envs,)`` with float32 dtype. - """ - - func: Callable[..., None] = MISSING - """The function to be called to fill the pre-allocated reward buffer.""" - - weight: float = MISSING - """The weight of the reward term.""" - - -@configclass -class TerminationTermCfg(_ManagerTermBaseCfg): - """Configuration for a termination term (experimental, Warp-first). - - The function is expected to write termination flags into a pre-allocated Warp buffer. - - Expected signature: - - - ``func(env, out, **params) -> None`` - - where ``out`` is a Warp array of shape ``(num_envs,)`` with boolean dtype. - """ - - func: Callable[..., None] = MISSING - """The function to be called to fill the pre-allocated termination buffer.""" - - time_out: bool = False - """Whether the termination term contributes towards episodic timeouts. Defaults to False.""" - - -@configclass -class ObservationTermCfg(_ManagerTermBaseCfg): - """Configuration for an observation term (experimental, Warp-first). - - The function is expected to write observation values into a pre-allocated Warp buffer provided - by the observation manager. - - Expected signature: - - - ``func(env, out, **params) -> None`` - - where ``out`` is a Warp array of shape ``(num_envs, obs_term_dim)`` with float32 dtype. - - Notes: - - The stable fields (noise/modifiers/history) are kept for config compatibility, but the - experimental Warp-first observation manager may not support all of them initially. - """ - - func: Callable[..., None] = MISSING - """The function to be called to fill the pre-allocated observation buffer.""" - - # Keep stable configuration fields for compatibility with existing task configs. - modifiers: list[ModifierCfg] | None = None # noqa: F405 - noise: NoiseCfg | NoiseModelCfg | None = None # noqa: F405 - clip: tuple[float, float] | None = None - scale: tuple[float, ...] | float | None = None - history_length: int = 0 - flatten_history_dim: bool = True +class CurriculumTermCfg(_CurriculumTermCfg): + """Configuration for a Warp-mask or legacy curriculum term.""" + + +__all__ = [ + "ActionTermCfg", + "CommandTermCfg", + "CurriculumTermCfg", + "EventTermCfg", + "ManagerTermBaseCfg", + "ObservationGroupCfg", + "ObservationTermCfg", + "RecorderTermCfg", + "RewardTermCfg", + "TerminationTermCfg", +] diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 74d0a9a955be..0c5e9a2ea1de 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -55,6 +55,7 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import ObservationGroupCfg, ObservationTermCfg from isaaclab.utils import class_to_dict from isaaclab_experimental.utils import modifiers, noise @@ -62,7 +63,6 @@ from isaaclab_experimental.utils.torch_utils import clone_obs_buffer from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import ObservationGroupCfg, ObservationTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -70,6 +70,7 @@ @wp.kernel def _apply_clip(out: wp.array(dtype=wp.float32, ndim=2), clip_lo: wp.float32, clip_hi: wp.float32): + """Clamp every column of a term's ``out`` block to [``clip_lo``, ``clip_hi``] in place.""" env_id = wp.tid() for j in range(out.shape[1]): out[env_id, j] = wp.clamp(out[env_id, j], clip_lo, clip_hi) @@ -77,6 +78,7 @@ def _apply_clip(out: wp.array(dtype=wp.float32, ndim=2), clip_lo: wp.float32, cl @wp.kernel def _apply_scale(out: wp.array(dtype=wp.float32, ndim=2), scale: wp.array(dtype=wp.float32)): + """Multiply each column of a term's ``out`` block by its per-column ``scale`` in place.""" env_id = wp.tid() for j in range(out.shape[1]): out[env_id, j] = out[env_id, j] * scale[j] @@ -382,15 +384,7 @@ def reset( *, env_mask: wp.array | None = None, ) -> dict[str, float]: - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - # Keep all id->mask resolution strictly outside capture. - if wp.get_device().is_capturing: - raise RuntimeError( - "ObservationManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) # call all terms that are classes for group_name, group_cfg in self._group_obs_class_term_cfgs.items(): @@ -679,6 +673,16 @@ def _prepare_terms(self): # noqa: C901 # check noise settings if not group_cfg.enable_corruption: term_cfg.noise = None + elif term_cfg.noise is not None and not type(term_cfg.noise).__module__.startswith( + "isaaclab_experimental." + ): + raise TypeError( + f"Observation term '{term_name}' uses noise cfg" + f" '{type(term_cfg.noise).__module__}.{type(term_cfg.noise).__name__}', which the Warp" + " observation manager would silently ignore. Use the warp-native noise cfgs from" + " isaaclab_experimental.utils.noise (the warp frontend converts stable cfgs" + " automatically)." + ) # check group history params and override terms if group_cfg.history_length is not None: term_cfg.history_length = group_cfg.history_length @@ -760,6 +764,10 @@ def _prepare_terms(self): # noqa: C901 f" and optional parameters: {args_with_defaults}, but received: {term_params}." ) + # plumb the shared per-env RNG state so Warp noise kernels can consume it + # (function-style NoiseCfg kernels read it off the cfg at call time) + if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseCfg): + term_cfg.noise.rng_state_wp = self._env.rng_state_wp # prepare noise model classes if term_cfg.noise is not None and isinstance(term_cfg.noise, noise.NoiseModelCfg): # plumb the shared per-env RNG state so Warp noise kernels can consume it @@ -905,10 +913,18 @@ def _resolve_out_dim(self, out_dim: int | str, term_cfg: ObservationTermCfg) -> - ``"body:N"``: ``N`` components per selected body from ``asset_cfg``. - ``"command"``: query ``command_manager.get_command(name).shape[-1]``. - ``"action"``: query ``action_manager.action.shape[-1]``. + - ``"sensor:rays"``: number of rays of the ray-caster sensor from ``sensor_cfg``. """ if isinstance(out_dim, int): return out_dim + if out_dim == "sensor:rays": + sensor_cfg = term_cfg.params.get("sensor_cfg") + if sensor_cfg is None: + raise ValueError("out_dim='sensor:rays' requires a 'sensor_cfg' entry in the term's params.") + sensor = self._env.scene.sensors[sensor_cfg.name] + return int(sensor.num_rays) + if out_dim == "joint": asset_cfg = term_cfg.params.get("asset_cfg") asset = self._env.scene[asset_cfg.name] @@ -922,11 +938,12 @@ def _resolve_out_dim(self, out_dim: int | str, term_cfg: ObservationTermCfg) -> if isinstance(out_dim, str) and out_dim.startswith("body:"): per_body = int(out_dim.split(":")[1]) - asset_cfg = term_cfg.params.get("asset_cfg") - body_ids = getattr(asset_cfg, "body_ids", None) + # Body selection may live on an asset (articulation) or a sensor entity. + entity_cfg = term_cfg.params.get("asset_cfg") or term_cfg.params.get("sensor_cfg") + body_ids = getattr(entity_cfg, "body_ids", None) if body_ids is None or body_ids == slice(None): - asset = self._env.scene[asset_cfg.name] - return per_body * len(asset.body_names) + entity = self._env.scene[entity_cfg.name] + return per_body * len(entity.body_names) return per_body * len(body_ids) if out_dim == "command": diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index 67dcbc055c2f..9f8490ed1c31 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -18,10 +18,11 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import RewardTermCfg + from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import RewardTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv @@ -37,6 +38,13 @@ def _sum_and_zero_masked( # output out_avg: wp.array(dtype=wp.float32), ): + """Fold selected envs' per-term episode reward sums into per-term means and zero + + them for the next episode. + Launch with dim=(num_terms, num_envs); ``scale[0]`` carries the 1/count + normalization and the atomic add reduces across selected envs into + ``out_avg[term]``. + """ term_idx, env_id = wp.tid() if mask[env_id]: wp.atomic_add(out_avg, term_idx, episode_sums[term_idx, env_id] * scale[0]) @@ -54,6 +62,13 @@ def _reward_finalize( episode_sums: wp.array(dtype=wp.float32, ndim=2), step_reward: wp.array(dtype=wp.float32, ndim=2), ): + """Combine per-term outputs into the env step reward and episode bookkeeping. + + Per env: ``step_reward`` records each term's weighted reward rate, + ``reward_buf`` accumulates ``weight * out * dt`` across terms, and + ``episode_sums`` grows by the same amount for episode metrics. Zero-weight + terms are skipped so disabled terms cost nothing. + """ env_id = wp.tid() total = wp.float32(0.0) @@ -127,6 +142,7 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): # call the base class constructor (this will parse the terms config) super().__init__(cfg, env) self._term_name_to_term_idx = {name: i for i, name in enumerate(self._term_names)} + self._device_weight_terms: set[str] = set() num_terms = len(self._term_names) self._num_terms = num_terms @@ -176,6 +192,17 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): self._term_weights_wp = wp.array( [float(term_cfg.weight) for term_cfg in self._term_cfgs], dtype=wp.float32, device=self.device ) + self._term_weight_views_wp: dict[str, wp.array] = {} + if num_terms > 0: + stride = self._term_weights_wp.strides[0] + for term_idx, term_name in enumerate(self._term_names): + self._term_weight_views_wp[term_name] = wp.array( + ptr=self._term_weights_wp.ptr + term_idx * stride, + dtype=wp.float32, + shape=(1,), + strides=(stride,), + device=self.device, + ) # persistent reset-time logging buffers (warp buffers) self._episode_sum_avg_wp = wp.zeros((num_terms,), dtype=wp.float32, device=self.device) @@ -242,14 +269,7 @@ def reset( Returns: A dictionary containing the information to log under the "Reward/{term_name}" key. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "RewardManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) self._episode_sum_avg_wp.zero_() self._reset_count_wp.zero_() @@ -271,9 +291,14 @@ def reset( device=self.device, ) - # reset all the reward terms + # reset all the reward terms; class terms may expose logging values as + # persistent tensor views (see :meth:`ManagerTermBase.reset`), which are + # merged into the manager's reset extras. Merging the same views again on + # subsequent resets is a no-op, so this stays capture-safe. for term_cfg in self._class_term_cfgs: - term_cfg.func.reset(env_mask=env_mask) + term_extras = term_cfg.func.reset(env_mask=env_mask) + if term_extras: + self._reset_extras.update(term_extras) return self._reset_extras @@ -298,9 +323,10 @@ def compute(self, dt: float) -> torch.Tensor: device=self.device, ) # iterate over all the reward terms (Python loop; per-term math is warp) - for term_cfg in self._term_cfgs: - # skip if weight is zero (kind of a micro-optimization) - if term_cfg.weight == 0.0: + for term_name, term_cfg in zip(self._term_names, self._term_cfgs): + # Static zero-weight terms remain free. Terms whose weight is owned + # on-device must still run so a curriculum can enable them later. + if term_cfg.weight == 0.0 and term_name not in self._device_weight_terms: continue # compute term into the persistent warp buffer (raw, unweighted) # NOTE: `out` is pre-zeroed every step by `_reward_pre_compute_reset`. @@ -362,8 +388,30 @@ def get_term_cfg(self, term_name: str) -> RewardTermCfg: """ if term_name not in self._term_names: raise ValueError(f"Reward term '{term_name}' not found.") - # return the configuration - return self._term_cfgs[self._term_names.index(term_name)] + term_idx = self._term_names.index(term_name) + term_cfg = self._term_cfgs[term_idx] + if term_name in self._device_weight_terms: + # Explicit config inspection is a control-plane operation, so it is + # the appropriate place to synchronize a device-owned scalar. + term_cfg.weight = float(self._term_weights_tensor_view[term_idx].item()) + return term_cfg + + def get_term_weight_wp(self, term_name: str) -> wp.array(dtype=wp.float32): + """Return the pointer-stable Warp view of a reward term's scalar weight. + + Args: + term_name: Name of the reward term. + + Returns: + One-element Warp array containing the reward weight. + + Raises: + ValueError: If the term name is not found. + """ + if term_name not in self._term_weight_views_wp: + raise ValueError(f"Reward term '{term_name}' not found.") + self._device_weight_terms.add(term_name) + return self._term_weight_views_wp[term_name] def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequence[float]]]: """Returns the active terms as iterable sequence of tuples. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py index ae97f7f526fb..f6d30489eeca 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py @@ -38,6 +38,15 @@ class SceneEntityCfg(_SceneEntityCfg): """Integer indices of selected bodies — used for subset-sized body gathers.""" body_ids_wp: wp.array | None = None + @classmethod + def from_stable(cls, stable: _SceneEntityCfg) -> SceneEntityCfg: + """Build a warp scene-entity cfg from a stable one. + + Copies every field declared on the stable cfg; the warp-specific fields + stay ``None`` and are filled by :meth:`resolve` at scene build time. + """ + return cls(**{name: getattr(stable, name) for name in _SceneEntityCfg.__dataclass_fields__}) + def resolve(self, scene: InteractiveScene): # run the stable resolution first (fills joint_ids/body_ids from names/regex) super().resolve(scene) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index 8a768651b292..e1657aafb2ed 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -22,8 +22,11 @@ import warp as wp from prettytable import PrettyTable +from isaaclab.managers.manager_term_cfg import TerminationTermCfg + +from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked + from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import TerminationTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv @@ -92,15 +95,16 @@ def _termination_finalize( # TODO(jichuanh): Look into wp.tile for better performance @wp.kernel -def _termination_reset_mean_all_2d( +def _termination_reset_mean_masked_2d( + env_mask: wp.array(dtype=wp.bool), last_episode_dones: wp.array(dtype=wp.bool, ndim=2), + scale: wp.array(dtype=wp.float32), term_done_avg: wp.array(dtype=wp.float32), ): - """Compute mean(done) per term with 2D parallel accumulation.""" + """Compute mean(done) per term over selected environments.""" env_id, term_idx = wp.tid() - num_envs = last_episode_dones.shape[0] - if num_envs > 0 and last_episode_dones[env_id, term_idx]: - wp.atomic_add(term_done_avg, term_idx, 1.0 / float(num_envs)) + if env_mask[env_id] and last_episode_dones[env_id, term_idx]: + wp.atomic_add(term_done_avg, term_idx, scale[0]) class TerminationManager(ManagerBase): @@ -129,6 +133,8 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): num_terms = len(self._term_names) self._term_dones_wp = wp.zeros((self.num_envs, num_terms), dtype=wp.bool, device=self.device) self._term_done_avg_wp = wp.zeros((num_terms,), dtype=wp.float32, device=self.device) + self._reset_count_wp = wp.zeros((1,), dtype=wp.int32, device=self.device) + self._reset_scale_wp = wp.zeros((1,), dtype=wp.float32, device=self.device) self._last_episode_dones_wp = wp.zeros((self.num_envs, num_terms), dtype=wp.bool, device=self.device) self._truncated_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) self._terminated_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) @@ -249,20 +255,27 @@ def reset( Returns: A dictionary containing the information to log under the "Termination/{term_name}" key. """ - # Mask-first path: captured callers must provide env_mask. - if env_mask is None or not isinstance(env_mask, wp.array): - if wp.get_device().is_capturing: - raise RuntimeError( - "TerminationManager.reset requires env_mask(wp.array[bool]) during capture. " - "Do not pass env_ids on captured paths." - ) - env_mask = self._env.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) + env_mask = self._resolve_reset_mask(env_ids, env_mask) if len(self._term_names) > 0: self._term_done_avg_wp.zero_() + self._reset_count_wp.zero_() + self._reset_scale_wp.zero_() + wp.launch( + kernel=count_masked, + dim=self.num_envs, + inputs=[env_mask, self._reset_count_wp], + device=self.device, + ) + wp.launch( + kernel=compute_reset_scale, + dim=1, + inputs=[self._reset_count_wp, 1.0, self._reset_scale_wp], + device=self.device, + ) wp.launch( - kernel=_termination_reset_mean_all_2d, + kernel=_termination_reset_mean_masked_2d, dim=(self.num_envs, len(self._term_names)), - inputs=[self._last_episode_dones_wp, self._term_done_avg_wp], + inputs=[env_mask, self._last_episode_dones_wp, self._reset_scale_wp, self._term_done_avg_wp], device=self.device, ) for term_cfg in self._class_term_cfgs: diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi index de2860732514..00660ba9775e 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi @@ -4,8 +4,6 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ - "ManagerCallMode", - "ManagerCallSwitch", "WarpGraphCache", "clone_obs_buffer", "buffers", @@ -14,7 +12,6 @@ __all__ = [ "warp", ] -from .manager_call_switch import ManagerCallMode, ManagerCallSwitch from .torch_utils import clone_obs_buffer from .warp_graph_cache import WarpGraphCache from . import buffers, modifiers, noise, warp diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py b/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py index f0dcf57b2a30..96bcf21aaf10 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py @@ -42,6 +42,9 @@ def __init__(self, max_len: int, batch_size: int, device: str): self._device = device self._ALL_INDICES = torch.arange(batch_size, device=device) + # max length as a host int: the property is read on the per-step append path, + # so it must not synchronize with the device tensor below. + self._max_len_int = int(max_len) # max length tensor for comparisons self._max_len = torch.full((batch_size,), max_len, dtype=torch.int, device=device) # number of data pushes passed since the last call to :meth:`reset` @@ -69,7 +72,7 @@ def device(self) -> str: @property def max_length(self) -> int: """The maximum length of the ring buffer.""" - return int(self._max_len[0].item()) + return self._max_len_int @property def current_length(self) -> torch.Tensor: diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py deleted file mode 100644 index 89924c789bc3..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py +++ /dev/null @@ -1,262 +0,0 @@ -# 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 - -"""Manager call switch for routing manager stage calls through stable/warp/captured paths.""" - -from __future__ import annotations - -import importlib -import json -import os -from enum import IntEnum -from typing import Any - -from isaaclab.utils.timer import Timer - -from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache - - -class ManagerCallMode(IntEnum): - """Execution mode for manager stage calls. - - * ``STABLE`` (0): Call stable Python manager implementations from :mod:`isaaclab.managers`. - * ``WARP_NOT_CAPTURED`` (1): Call Warp-compatible implementations without CUDA graph capture. - * ``WARP_CAPTURED`` (2): Call Warp implementations with CUDA graph capture/replay. - """ - - STABLE = 0 - WARP_NOT_CAPTURED = 1 - WARP_CAPTURED = 2 - - -class ManagerCallSwitch: - """Per-manager call switch for stable/warp/captured execution. - - Routes each manager stage call through the configured execution path: - stable Python, Warp (eager), or Warp (captured CUDA graph). Optionally - wraps each call in a :class:`Timer` context for profiling. - """ - - DEFAULT_CONFIG: dict[str, int] = {"default": 2} - DEFAULT_KEY = "default" - MANAGER_NAMES: tuple[str, ...] = ( - "ActionManager", - "ObservationManager", - "EventManager", - "RecorderManager", - "CommandManager", - "TerminationManager", - "RewardManager", - "CurriculumManager", - "Scene", - ) - # FIXME: Scene_write_data_to_sim calls articulation._apply_actuator_model which - # uses wp.to_torch + torch indexing -- not capture-safe on this branch. - # Cap Scene stages to WARP_NOT_CAPTURED until the articulation layer is capture-ready. - MAX_MODE_OVERRIDES: dict[str, int] = {"Scene": ManagerCallMode.WARP_NOT_CAPTURED} - - ENV_VAR = "MANAGER_CALL_CONFIG" - """Environment variable name for the JSON config string. - - Example usage:: - - MANAGER_CALL_CONFIG='{"RewardManager": 0, "default": 2}' python train.py ... - """ - - def __init__( - self, - cfg_source: dict | str | None = None, - *, - max_modes: dict[str, int] | None = None, - ): - self._graph_cache = WarpGraphCache() - # Merge caller-supplied max_modes with the class-level MAX_MODE_OVERRIDES. - self._max_modes = dict(self.MAX_MODE_OVERRIDES) - if max_modes is not None: - self._max_modes.update(max_modes) - # Resolve config: prefer explicit cfg_source, fall back to env var. - if cfg_source is None: - cfg_source = os.environ.get(self.ENV_VAR) - self._cfg = self._load_cfg(cfg_source) - print("[INFO] ManagerCallSwitch configuration:") - print(f" - {self.DEFAULT_KEY}: {self._cfg[self.DEFAULT_KEY]}") - for manager_name in self.MANAGER_NAMES: - mode = int(self.get_mode_for_manager(manager_name)) - cap = self._max_modes.get(manager_name) - cap_str = f" (cap={cap})" if cap is not None else "" - print(f" - {manager_name}: {mode}{cap_str}") - - # ------------------------------------------------------------------ - # Graph management - # ------------------------------------------------------------------ - - def invalidate_graphs(self) -> None: - """Invalidate cached capture graphs and their cached return values.""" - self._graph_cache.invalidate() - - # ------------------------------------------------------------------ - # Stage dispatch - # ------------------------------------------------------------------ - - def call_stage( - self, - *, - stage: str, - warp_call: dict[str, Any], - stable_call: dict[str, Any] | None = None, - timer: bool = False, - ) -> Any: - """Run the stage according to configured mode, optionally wrapped in a :class:`Timer`. - - A call spec dict supports the following keys: - - * ``fn`` (required): The callable to invoke. - * ``args`` (optional): Positional arguments tuple. - * ``kwargs`` (optional): Keyword arguments dict. - * ``output`` (optional): A ``Callable[[Any], Any]`` that transforms the raw - return value into the final output. For captured stages the raw value is - ``None``. When omitted, the raw return value is used as-is. - - Args: - stage: Stage identifier in the form ``"ManagerName_function_name"``. - warp_call: Call spec for the warp path (eager or captured). - stable_call: Call spec for the stable (torch) path. Defaults to ``None``. - timer: Whether to wrap execution in a :class:`Timer`. Defaults to ``True`` - (controlled by the global :attr:`Timer.enable` class-level toggle). - Pass a module-level flag like ``TIMER_ENABLED_STEP`` to make timing - conditional on that flag. - - Returns: - The (possibly transformed) return value of the stage. - """ - with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"): - return self._dispatch(stage, stable_call, warp_call) - - def _dispatch( - self, - stage: str, - stable_call: dict[str, Any] | None, - warp_call: dict[str, Any], - ) -> Any: - """Select call path based on mode, execute, and apply output.""" - mode = self.get_mode_for_manager(self._manager_name_from_stage(stage)) - if mode == ManagerCallMode.STABLE: - if stable_call is None: - raise ValueError(f"Stage '{stage}' is configured as STABLE (mode=0) but no stable_call was provided.") - call, result = stable_call, self._run_call(stable_call) - elif mode == ManagerCallMode.WARP_CAPTURED: - call, result = warp_call, self._wp_capture_or_launch(stage, warp_call) - else: - call, result = warp_call, self._run_call(warp_call) - - output_fn = call.get("output") - return output_fn(result) if output_fn is not None else result - - # ------------------------------------------------------------------ - # Manager resolution - # ------------------------------------------------------------------ - - def _manager_name_from_stage(self, stage: str) -> str: - if "_" not in stage: - raise ValueError(f"Invalid stage '{stage}'. Expected '{{manager_name}}_{{function_name}}'.") - return stage.split("_", 1)[0] - - def get_mode_for_manager(self, manager_name: str) -> ManagerCallMode: - """Return the resolved execution mode for the given manager. - - Looks up the manager in the config dict, falls back to the default, - then caps by :attr:`_max_modes` (static overrides + dynamic registrations). - """ - default_key = next(iter(self.DEFAULT_CONFIG)) - mode_value = self._cfg.get(manager_name, self._cfg[default_key]) - cap = self._max_modes.get(manager_name) - if cap is not None: - mode_value = min(mode_value, cap) - return ManagerCallMode(mode_value) - - def resolve_manager_class(self, manager_name: str, mode_override: ManagerCallMode | int | None = None) -> type: - """Import and return the manager class for the configured mode.""" - mode = self.get_mode_for_manager(manager_name) if mode_override is None else ManagerCallMode(mode_override) - module_name = "isaaclab.managers" if mode == ManagerCallMode.STABLE else "isaaclab_experimental.managers" - module = importlib.import_module(module_name) - if not hasattr(module, manager_name): - raise AttributeError(f"Manager '{manager_name}' not found in module '{module_name}'.") - return getattr(module, manager_name) - - def register_manager_capturability(self, manager_name: str, capturable: bool) -> None: - """Register that a manager has non-capturable terms, capping its mode. - - Called by :class:`ManagerBase` during term preparation when a term - is decorated with ``@warp_capturable(False)``. - """ - if not capturable: - self._max_modes[manager_name] = min( - self._max_modes.get(manager_name, ManagerCallMode.WARP_CAPTURED), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - def _run_call(self, call: dict[str, Any]) -> Any: - """Execute a single call spec eagerly.""" - return call["fn"](*call.get("args", ()), **call.get("kwargs", {})) - - def _wp_capture_or_launch(self, stage: str, call: dict[str, Any]) -> Any: - """Capture Warp CUDA graph on first call, then replay. - - Delegates to :class:`WarpGraphCache` which handles warm-up, capture, - caching the return value, and replay. - """ - return self._graph_cache.capture_or_replay( - stage, - call["fn"], - args=call.get("args", ()), - kwargs=call.get("kwargs", {}), - ) - - def _load_cfg(self, cfg_source: dict | str | None) -> dict[str, int]: - if cfg_source is None: - cfg = dict(self.DEFAULT_CONFIG) - elif isinstance(cfg_source, dict): - cfg = dict(cfg_source) - if self.DEFAULT_KEY not in cfg: - cfg[self.DEFAULT_KEY] = self.DEFAULT_CONFIG[self.DEFAULT_KEY] - elif isinstance(cfg_source, str): - if cfg_source.strip() == "": - cfg = dict(self.DEFAULT_CONFIG) - else: - parsed = json.loads(cfg_source) - if not isinstance(parsed, dict): - raise TypeError("manager_call_config must decode to a dict.") - cfg = dict(parsed) - if self.DEFAULT_KEY not in cfg: - cfg[self.DEFAULT_KEY] = self.DEFAULT_CONFIG[self.DEFAULT_KEY] - else: - raise TypeError(f"cfg_source must be a dict, string, or None, got: {type(cfg_source)}") - - # validation - for manager_name, mode_value in cfg.items(): - if not isinstance(mode_value, int): - raise TypeError( - f"manager_call_config value for '{manager_name}' must be int (0/1/2), got: {type(mode_value)}" - ) - try: - ManagerCallMode(mode_value) - except ValueError as exc: - raise ValueError( - f"Invalid manager_call_config value for '{manager_name}': {mode_value}. Expected 0/1/2." - ) from exc - - # Apply max mode caps: bake caps into the resolved config so - # get_mode_for_manager never needs per-call branching. - default_mode = cfg[self.DEFAULT_KEY] - for name, max_mode in self._max_modes.items(): - resolved = cfg.get(name, default_mode) - if resolved > max_mode: - cfg[name] = max_mode - - return cfg diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py index 2d071b823a46..c7e33acf47bf 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py @@ -5,12 +5,21 @@ """Warp utility functions and shared kernels for isaaclab_experimental.""" -from .kernels import compute_reset_scale, count_masked +from isaaclab.utils.warp.utils import resolve_1d_mask + +from .kernels import ( + compute_reset_scale, + count_masked, + increment_all_int32, + increment_all_int64, + zero_masked_int32, + zero_masked_int64, +) from .utils import ( + SYNC_DEBUG_ENV_VAR, WarpCapturable, - is_warp_capturable, - resolve_1d_mask, - warp_capturable, + any_env_set, + sync_debug_enabled, wrap_to_pi, zero_masked_2d, ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py index 8d3f9e65d49a..dbfd3740a61b 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/kernels.py @@ -10,6 +10,36 @@ import warp as wp +@wp.kernel +def increment_all_int32(values: wp.array(dtype=wp.int32), increment: wp.int32): + """Increment every value in a one-dimensional int32 array.""" + index = wp.tid() + values[index] = values[index] + increment + + +@wp.kernel +def increment_all_int64(values: wp.array(dtype=wp.int64), increment: wp.int64): + """Increment every value in a one-dimensional int64 array.""" + index = wp.tid() + values[index] = values[index] + increment + + +@wp.kernel +def zero_masked_int32(mask: wp.array(dtype=wp.bool), values: wp.array(dtype=wp.int32)): + """Zero selected values in a one-dimensional int32 array.""" + index = wp.tid() + if mask[index]: + values[index] = 0 + + +@wp.kernel +def zero_masked_int64(mask: wp.array(dtype=wp.bool), values: wp.array(dtype=wp.int64)): + """Zero selected values in a one-dimensional int64 array.""" + index = wp.tid() + if mask[index]: + values[index] = wp.int64(0) + + @wp.kernel def count_masked( mask: wp.array(dtype=wp.bool), diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py index a90f6b2d6fc6..87c4833f89e9 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py @@ -5,126 +5,36 @@ from __future__ import annotations -from collections.abc import Sequence +import os import torch import warp as wp +SYNC_DEBUG_ENV_VAR = "ISAACLAB_SYNC_DEBUG" +"""Set ``ISAACLAB_SYNC_DEBUG=1`` to trap hidden GPU->host syncs in warp stages (CI/debug).""" -@wp.kernel -def _set_mask_from_ids( - mask: wp.array(dtype=wp.bool), - ids: wp.array(dtype=wp.int32), -): - """Set ``mask[ids[i]] = True`` for each thread *i*.""" - i = wp.tid() - mask[ids[i]] = True - - -def resolve_1d_mask( - *, - ids: Sequence[int] | slice | wp.array | torch.Tensor | None, - mask: wp.array | torch.Tensor | None, - all_mask: wp.array, - scratch_mask: wp.array, - device: str, -) -> wp.array: - """Resolve ids/mask into a warp boolean mask. - - Matches the contract of ``ArticulationData._resolve_1d_mask`` on dev/newton. - Callers must provide pre-allocated ``all_mask`` (all-True) and ``scratch_mask`` - (reusable working buffer). No allocations happen inside this function. - - Args: - ids: Indices to set to ``True``. ``None`` or ``slice(None)`` means all. - mask: Explicit boolean mask. If provided, returned directly (after - torch->warp normalization if needed). Takes precedence over *ids*. - all_mask: Pre-allocated all-True mask of shape ``(size,)``, returned - when both *ids* and *mask* are ``None``. - scratch_mask: Pre-allocated scratch mask of shape ``(size,)``, filled - in-place when *ids* are provided. - device: Warp device string. - - Returns: - A ``wp.array(dtype=wp.bool)`` -- ``mask``, ``all_mask``, or ``scratch_mask``. - """ - # Fast path: explicit mask provided. - if mask is not None: - if isinstance(mask, torch.Tensor): - if mask.dtype != torch.bool: - mask = mask.to(dtype=torch.bool) - if str(mask.device) != device: - mask = mask.to(device) - return wp.from_torch(mask, dtype=wp.bool) - return mask - - # Fast path: all ids. - if ids is None or (isinstance(ids, slice) and ids == slice(None)): - return all_mask - - # Normalize slice into explicit indices. - if isinstance(ids, slice): - start, stop, step = ids.indices(scratch_mask.shape[0]) - ids = list(range(start, stop, step)) - elif not isinstance(ids, (torch.Tensor, wp.array)): - ids = list(ids) - - # Prepare output mask. - scratch_mask.fill_(False) - - # Normalize ids to wp.int32 array and launch kernel. - if isinstance(ids, torch.Tensor): - if ids.numel() == 0: - return scratch_mask - if str(ids.device) != device: - ids = ids.to(device) - if ids.dtype != torch.int32: - ids = ids.to(dtype=torch.int32) - if not ids.is_contiguous(): - ids = ids.contiguous() - ids_wp = wp.from_torch(ids, dtype=wp.int32) - elif isinstance(ids, wp.array): - if ids.shape[0] == 0: - return scratch_mask - ids_wp = ids - else: - if len(ids) == 0: - return scratch_mask - ids_wp = wp.array(ids, dtype=wp.int32, device=device) - - wp.launch(kernel=_set_mask_from_ids, dim=ids_wp.shape[0], inputs=[scratch_mask, ids_wp], device=device) - return scratch_mask - - -def warp_capturable(capturable: bool): - """Annotate an MDP term's CUDA-graph capturability. - - No-wrapper decorator: sets ``_warp_capturable`` directly on the function - and returns it unchanged. Safe to stack with any other decorator in any order. - - By default all MDP terms are assumed capturable (True). Use - ``@warp_capturable(False)`` on terms that call non-capturable external APIs. - """ - def decorator(func): - func._warp_capturable = capturable - return func +def sync_debug_enabled() -> bool: + """Whether the sync-debug trap is requested for this process.""" + return os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1" - return decorator +def any_env_set(mask: torch.Tensor) -> bool: + """Host predicate: does the boolean mask select any environment? -def is_warp_capturable(func) -> bool: - """Check if a term function is CUDA-graph-capturable. - - Checks ``_warp_capturable`` on the function and its ``__wrapped__`` target. - Returns True (capturable) by default if no annotation is found. + This is the single sanctioned per-step host sync of the mask-native reset + pipeline (one predicate instead of dispatching an empty reset). It suspends + the sync-debug trap around itself so the exemption is code, not annotation. """ - for f in (func, getattr(func, "__wrapped__", None)): - if f is not None: - val = getattr(f, "_warp_capturable", None) - if val is not None: - return val - return True + if mask.is_cuda: + prev = torch.cuda.get_sync_debug_mode() + if prev != 0: + torch.cuda.set_sync_debug_mode(0) + try: + return bool(mask.any().item()) + finally: + torch.cuda.set_sync_debug_mode(prev) + return bool(mask.any().item()) @wp.func @@ -140,32 +50,40 @@ def wrap_to_pi(angle: float) -> float: class WarpCapturable: """CUDA graph capture safety: decorator, annotation checker, and runtime guard. - Decorator usage:: - - @WarpCapturable(False) - def reset_root_state_uniform(env, env_mask, ...): - ... + The single capturability annotation (see ``is_capturable``). Decorator usage:: @WarpCapturable(False, reason="calls write_root_pose_to_sim") def push_by_setting_velocity(env, env_mask, ...): ... + @WarpCapturable(False, reason="notifies the solver of model changes") + class randomize_rigid_body_com(ManagerTermBase): + ... + - ``@WarpCapturable(True)`` or no decorator: capturable, returned unwrapped. - - ``@WarpCapturable(False)``: sets ``func._warp_capturable = False``, wraps with - runtime guard that raises if ``wp.get_device().is_capturing`` is ``True``. + - ``@WarpCapturable(False)`` on a function: sets ``_warp_capturable = False`` and + wraps it with a runtime guard that raises if called during CUDA graph capture. + - ``@WarpCapturable(False)`` on a class term: annotates the class (so manager + registration sees it) and guards its ``__call__``. """ def __init__(self, capturable: bool, *, reason: str | None = None): self._capturable = capturable self._reason = reason - def __call__(self, func): - """Decorate *func* with capture safety annotation and optional runtime guard.""" - import functools - - func._warp_capturable = self._capturable + def __call__(self, target): + """Decorate a term function or class with the annotation and optional guard.""" + target._warp_capturable = self._capturable if self._capturable: - return func + return target + if isinstance(target, type): + target.__call__ = self._guarded(target.__call__) + return target + return self._guarded(target) + + def _guarded(self, func): + """Wrap *func* so calling it during CUDA graph capture raises.""" + import functools reason = self._reason @@ -182,12 +100,14 @@ def wrapper(*args, **kwargs): return wrapper @staticmethod - def is_capturable(func) -> bool: - """Check capturability annotation. Default: ``True``. + def is_capturable(term) -> bool: + """Check the capturability annotation on a term function, class, or instance. Checks ``__wrapped__`` for decorated functions to handle stacked decorators. + Unannotated terms default to capturable; hidden syncs in them are trapped by + the ``ISAACLAB_SYNC_DEBUG`` stage sweep and by CUDA graph capture itself. """ - for f in (func, getattr(func, "__wrapped__", None)): + for f in (term, getattr(term, "__wrapped__", None)): if f is not None: val = getattr(f, "_warp_capturable", None) if val is not None: diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py index d446be28d099..0bdb00d02c45 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py @@ -5,20 +5,51 @@ """Warp CUDA graph capture-or-replay utility.""" +import os from collections.abc import Callable from typing import Any +import torch import warp as wp +from isaaclab.utils.timer import Timer + +SYNC_DEBUG_ENV_VAR = "ISAACLAB_SYNC_DEBUG" +"""Set ``ISAACLAB_SYNC_DEBUG=1`` to run eager and warm-up stage invocations under +``torch.cuda.set_sync_debug_mode("error")`` so hidden GPU->host syncs raise (CI/debug).""" + +_SYNC_DEBUG = os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1" + + +def _invoke_stage(fn: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + """Run one stage callable, optionally under the sync-debug trap. + + The trap is the choke point that guards every current and future term inside a + warp stage without per-function annotations; opt-outs suspend it explicitly + (see :func:`isaaclab_experimental.utils.warp.any_env_set`). + """ + if not _SYNC_DEBUG: + return fn(*args, **kwargs) + prev = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode(2) + try: + return fn(*args, **kwargs) + finally: + torch.cuda.set_sync_debug_mode(prev) + + +CAPTURE_ENV_VAR = "ISAACLAB_WARP_CAPTURE" +"""Set ``ISAACLAB_WARP_CAPTURE=0`` to force every stage eager (debug / A-B validation).""" + class WarpGraphCache: - """Caches Warp CUDA graphs by stage name: captures on first call, replays after. + """Execute Warp stages eagerly or through cached CUDA graphs. - On the very first call for a given stage, an **eager warm-up** run - executes *before* graph capture. This lets one-time initialisation - code (memory allocations, torch dtype casts, ``hasattr`` guards, etc.) - run outside the capture context. Only the steady-state kernel - launches are then recorded into the graph. + :meth:`call` uses the first invocation as an eager warm-up, captures the + second invocation, and replays later invocations. This lets one-time + initialization run outside capture without executing stateful manager work + twice in one environment step. :meth:`capture_or_replay` retains the direct + environment behavior of warming up and capturing on its first invocation. The return value from the capture run is cached and returned on every subsequent replay, ensuring captured stages return the same references @@ -27,20 +58,88 @@ class WarpGraphCache: Usage:: cache = WarpGraphCache() - result = cache.capture_or_replay("my_stage", my_warp_function) + result = cache.call("MyManager_step", my_warp_function) # uncaptured work here ... - result2 = cache.capture_or_replay("my_stage_post", my_other_function) + result2 = cache.call("OtherManager_step", my_other_function) """ - def __init__(self): + def __init__(self, *, enabled: bool = True, device: wp.DeviceLike = None): + """Initialize the cache. + + Args: + enabled: Whether eligible stages may use CUDA graph capture. The + :data:`CAPTURE_ENV_VAR` environment variable can force this off. + device: Warp device used by cached stages. CPU devices always run eagerly. + """ + self._enabled = enabled and os.environ.get(CAPTURE_ENV_VAR, "1") != "0" + self._device = wp.get_device(device) if device is not None else None self._graphs: dict[str, Any] = {} self._results: dict[str, Any] = {} + self._capturable: dict[str, bool] = {} + self._warmed: set[str] = set() + + @property + def captured_stages(self) -> tuple[str, ...]: + """Stages currently backed by a captured CUDA graph, sorted by name.""" + return tuple(sorted(self._graphs)) + + @property + def eager_groups(self) -> tuple[str, ...]: + """Stage groups registered as non-capturable, sorted by name.""" + return tuple(sorted(group for group, capturable in self._capturable.items() if not capturable)) + + def call( + self, + stage: str, + fn: Callable[..., Any], + /, + *args: Any, + output: Callable[[Any], Any] | None = None, + timer: bool = False, + **kwargs: Any, + ) -> Any: + """Run a Warp frontend stage eagerly or through its cached CUDA graph. + + The stage prefix before the first underscore identifies its capturability + group. A group stays eager once any of its terms is registered as unsafe. + Eligible stages run eagerly once, capture on their second invocation, and + replay thereafter. + + Args: + stage: Stage identifier in the form ``"GroupName_function_name"``. + fn: Callable implementing the stage. + *args: Positional arguments forwarded to :paramref:`fn`. + output: Optional transform applied to the stage result after execution. + timer: Whether to time the stage execution. + **kwargs: Keyword arguments forwarded to :paramref:`fn`. + + Returns: + The stage result, optionally transformed by :paramref:`output`. + """ + group = stage.partition("_")[0] + with Timer(name=stage, msg=f"{stage} took:", enable=timer, time_unit="us"): + if not self._capture_enabled or not self.is_capturable(group): + result = _invoke_stage(fn, args, kwargs) + elif stage in self._graphs: + wp.capture_launch(self._graphs[stage]) + result = self._results[stage] + elif stage not in self._warmed: + # The first real call doubles as warm-up. Stateful manager stages + # must execute exactly once per environment step. + result = _invoke_stage(fn, args, kwargs) + self._warmed.add(stage) + else: + result = self._capture(stage, fn, args, kwargs) + # CUDA graph capture records launches without executing them. Launch + # the new graph so this logical manager call still advances once. + wp.capture_launch(self._graphs[stage]) + return output(result) if output is not None else result def capture_or_replay( self, stage: str, fn: Callable[..., Any], - args: tuple = (), + args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, ) -> Any: """Capture *fn* into a CUDA graph on the first call, then replay. @@ -58,24 +157,60 @@ def capture_or_replay( """ if kwargs is None: kwargs = {} + if not self._capture_enabled: + return _invoke_stage(fn, args, kwargs) graph = self._graphs.get(stage) if graph is not None: wp.capture_launch(graph) return self._results[stage] # Warm-up: run eagerly to flush first-call allocations / hasattr guards. - fn(*args, **kwargs) - # Capture: allocations already done, only wp.launch calls are recorded. - with wp.ScopedCapture() as capture: + _invoke_stage(fn, args, kwargs) + return self._capture(stage, fn, args, kwargs) + + def _capture( + self, + stage: str, + fn: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> Any: + """Capture one already-warmed stage invocation.""" + with wp.ScopedCapture(device=self._device) as capture: result = fn(*args, **kwargs) self._graphs[stage] = capture.graph self._results[stage] = result + self._warmed.add(stage) + # One grep-able line per stage so runs can assert capture coverage. + print(f"[INFO] WarpGraphCache: captured stage '{stage}'") return result + def register_capturability(self, key: str, capturable: bool) -> None: + """Register conservative capture eligibility for a stage group. + + Args: + key: Stage group identifier. + capturable: Whether every operation registered under the group is safe + during CUDA graph capture. + """ + self._capturable[key] = self._capturable.get(key, True) and capturable + + def is_capturable(self, key: str) -> bool: + """Return whether a stage group is eligible for capture.""" + return self._capturable.get(key, True) + + @property + def _capture_enabled(self) -> bool: + """Return whether this cache can capture on its configured device.""" + device = self._device if self._device is not None else wp.get_device() + return self._enabled and device.is_cuda + def invalidate(self, stage: str | None = None) -> None: """Drop cached graph(s). If *stage* is ``None``, drop all.""" if stage is None: self._graphs.clear() self._results.clear() + self._warmed.clear() else: self._graphs.pop(stage, None) self._results.pop(stage, None) + self._warmed.discard(stage) diff --git a/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py b/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py new file mode 100644 index 000000000000..9d5bef7edfa4 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/mdp/commands/test_commands.py @@ -0,0 +1,334 @@ +# 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 + +"""Tests for Warp-native registered command terms.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import torch +import warp as wp +from isaaclab_experimental.envs import mdp as experimental_mdp +from isaaclab_experimental.envs.mdp.commands import ( + UniformPoseCommand, + UniformPoseCommandCfg, + UniformVelocityCommand, + UniformVelocityCommandCfg, +) +from isaaclab_experimental.managers import CommandManager, CommandTerm + +from isaaclab.envs.mdp.commands import UniformPoseCommandCfg as StableUniformPoseCommandCfg +from isaaclab.envs.mdp.commands import UniformVelocityCommandCfg as StableUniformVelocityCommandCfg +from isaaclab.utils.math import combine_frame_transforms, compute_pose_error, quat_from_euler_xyz, quat_unique + + +class _ArrayProxy: + """Minimal backend-neutral data proxy used by command terms.""" + + def __init__(self, values: np.ndarray, dtype): + self.warp = wp.array(values, dtype=dtype, device="cpu") + self.torch = wp.to_torch(self.warp) + + +class _Robot: + """Minimal articulation data needed by velocity and pose commands.""" + + def __init__(self, num_envs: int): + identity_quat = np.zeros((num_envs, 4), dtype=np.float32) + identity_quat[:, 3] = 1.0 + identity_pose = np.zeros((num_envs, 7), dtype=np.float32) + identity_pose[:, 6] = 1.0 + self.data = SimpleNamespace( + root_pose_w=_ArrayProxy(identity_pose, wp.transformf), + root_vel_w=_ArrayProxy(np.zeros((num_envs, 6), dtype=np.float32), wp.spatial_vectorf), + root_pos_w=_ArrayProxy(np.zeros((num_envs, 3), dtype=np.float32), wp.vec3f), + root_quat_w=_ArrayProxy(identity_quat, wp.quatf), + root_lin_vel_b=_ArrayProxy(np.zeros((num_envs, 3), dtype=np.float32), wp.vec3f), + root_ang_vel_b=_ArrayProxy(np.zeros((num_envs, 3), dtype=np.float32), wp.vec3f), + heading_w=_ArrayProxy(np.zeros(num_envs, dtype=np.float32), wp.float32), + body_pos_w=_ArrayProxy(np.zeros((num_envs, 1, 3), dtype=np.float32), wp.vec3f), + body_quat_w=_ArrayProxy(identity_quat[:, None, :], wp.quatf), + ) + self.is_initialized = True + + def find_bodies(self, body_name: str) -> tuple[list[int], list[str]]: + return [0], [body_name] + + +class _Env: + """Minimal manager-based environment used by command terms.""" + + def __init__(self, num_envs: int = 4): + self.num_envs = num_envs + self.device = "cpu" + self.scene = {"robot": _Robot(num_envs)} + self.rng_state_wp = wp.array(np.arange(num_envs, dtype=np.uint32) + 41, device=self.device) + self.extras = {} + self.sim = SimpleNamespace( + vis_marker_registry=SimpleNamespace( + add_debug_vis_callback=lambda term: None, + clear_debug_vis_callback=lambda term: None, + ) + ) + + +def _velocity_cfg() -> UniformVelocityCommandCfg: + return UniformVelocityCommandCfg( + asset_name="robot", + resampling_time_range=(2.0, 2.0), + heading_command=True, + heading_control_stiffness=2.0, + rel_heading_envs=1.0, + rel_standing_envs=0.0, + vel_xy_success_threshold=0.5, + vel_yaw_success_threshold=0.4, + ranges=UniformVelocityCommandCfg.Ranges( + lin_vel_x=(1.0, 1.0), + lin_vel_y=(2.0, 2.0), + ang_vel_z=(0.3, 0.3), + heading=(0.7, 0.7), + ), + debug_vis=False, + ) + + +def _pose_cfg() -> UniformPoseCommandCfg: + return UniformPoseCommandCfg( + asset_name="robot", + body_name="tool", + resampling_time_range=(3.0, 3.0), + make_quat_unique=True, + position_success_threshold=0.25, + ranges=UniformPoseCommandCfg.Ranges( + pos_x=(1.0, 1.0), + pos_y=(2.0, 2.0), + pos_z=(3.0, 3.0), + roll=(0.2, 0.2), + pitch=(-0.3, -0.3), + yaw=(4.0, 4.0), + ), + debug_vis=False, + ) + + +def test_configs_exports_and_command_accessors_preserve_shared_contracts(): + """Configs should inherit stable fields and terms should expose debug UI and persistent storage.""" + assert issubclass(UniformVelocityCommandCfg, StableUniformVelocityCommandCfg) + assert issubclass(UniformPoseCommandCfg, StableUniformPoseCommandCfg) + assert experimental_mdp.UniformVelocityCommandCfg is UniformVelocityCommandCfg + assert experimental_mdp.UniformPoseCommandCfg is UniformPoseCommandCfg + assert str(_velocity_cfg().class_type).endswith(".commands.velocity_command:UniformVelocityCommand") + assert str(_pose_cfg().class_type).endswith(".commands.pose_command:UniformPoseCommand") + assert UniformVelocityCommand._set_debug_vis_impl.__module__ == "isaaclab.envs.mdp.commands._debug_vis" + assert UniformPoseCommand._set_debug_vis_impl.__module__ == "isaaclab.envs.mdp.commands._debug_vis" + + env = _Env() + velocity_term = UniformVelocityCommand(_velocity_cfg(), env) + pose_term = UniformPoseCommand(_pose_cfg(), env) + manager = CommandManager.__new__(CommandManager) + manager._terms = {"velocity": velocity_term, "pose": pose_term} + + assert velocity_term.command_wp is velocity_term.command_wp + assert pose_term.command_wp is pose_term.command_wp + assert manager.get_command_wp("velocity") is velocity_term.command_wp + assert manager.get_command_wp("pose") is pose_term.command_wp + assert manager.get_command("velocity") is velocity_term.command + assert manager.get_command("pose") is pose_term.command + torch.testing.assert_close(pose_term.command[:, :6], torch.zeros(4, 6)) + torch.testing.assert_close(pose_term.command[:, 6], torch.ones(4)) + + +def test_debug_visualization_uses_current_simulation_registry(): + """Command terms should register visualization through the current simulation API.""" + registry = SimpleNamespace(add_debug_vis_callback=Mock(return_value="handle"), clear_debug_vis_callback=Mock()) + term = SimpleNamespace( + has_debug_vis_implementation=True, + _set_debug_vis_impl=Mock(), + _debug_vis_handle=None, + _env=SimpleNamespace(sim=SimpleNamespace(vis_marker_registry=registry)), + ) + + assert CommandTerm.set_debug_vis(term, True) + assert term._debug_vis_handle == "handle" + registry.add_debug_vis_callback.assert_called_once_with(term) + + assert CommandTerm.set_debug_vis(term, False) + registry.clear_debug_vis_callback.assert_called_once_with(term) + + CommandTerm.__del__(term) + assert registry.clear_debug_vis_callback.call_count == 2 + + +def test_uniform_velocity_command_updates_and_resets_selected_rows(): + """Velocity commands should update in Warp and preserve unselected rows during sparse reset.""" + env = _Env() + term = UniformVelocityCommand(_velocity_cfg(), env) + term._prepare_reset_extras() + command_pointer = term.command.data_ptr() + command_wp = term.command_wp + + term.command[:] = torch.tensor([[1.0, 2.0, 0.3], [2.0, -1.0, -0.2], [0.0, 0.5, 0.1], [-1.0, -2.0, 0.0]]) + env.scene["robot"].data.root_vel_w.torch[:, :3] = torch.tensor( + [[0.0, 0.0, 0.0], [1.0, -1.0, 0.0], [0.0, 0.0, 0.0], [-2.0, -2.0, 0.0]] + ) + env.scene["robot"].data.root_vel_w.torch[:, 5] = torch.tensor([0.1, -0.1, 0.4, 0.0]) + term._update_metrics() + wp.synchronize() + torch.testing.assert_close( + wp.to_torch(term._error_xy_sum), + torch.tensor([np.sqrt(5.0), 1.0, 0.5, 1.0], dtype=torch.float32), + ) + torch.testing.assert_close(wp.to_torch(term._error_yaw_sum), torch.tensor([0.2, 0.1, 0.3, 0.0])) + + term.command.fill_(-9.0) + wp.to_torch(term._error_xy_sum)[:] = torch.tensor([0.2, 9.0, 2.0, 8.0]) + wp.to_torch(term._error_yaw_sum)[:] = torch.tensor([0.1, 9.0, 2.0, 8.0]) + wp.to_torch(term._step_count)[:] = torch.tensor([1.0, 1.0, 2.0, 1.0]) + wp.to_torch(term.time_left_wp)[:] = torch.tensor([-1.0, 6.0, -1.0, 8.0]) + wp.to_torch(term.command_counter_wp)[:] = torch.tensor([5, 6, 7, 8], dtype=torch.int32) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + extras = term.reset(env_mask=env_mask) + wp.synchronize() + + assert term.command_wp is command_wp + assert term.command.data_ptr() == command_pointer + torch.testing.assert_close(term.command[[0, 2]], torch.tensor([[1.0, 2.0, 0.3], [1.0, 2.0, 0.3]])) + torch.testing.assert_close(term.command[[1, 3]], torch.full((2, 3), -9.0)) + torch.testing.assert_close(wp.to_torch(term.time_left_wp), torch.tensor([2.0, 6.0, 2.0, 8.0])) + torch.testing.assert_close(wp.to_torch(term.command_counter_wp), torch.tensor([1, 6, 1, 8], dtype=torch.int32)) + torch.testing.assert_close(extras["error_vel_xy"], torch.tensor(0.6)) + torch.testing.assert_close(extras["error_vel_yaw"], torch.tensor(0.55)) + torch.testing.assert_close(extras["success_rate"], torch.tensor(0.5)) + torch.testing.assert_close(wp.to_torch(term._error_xy_sum), torch.tensor([0.0, 9.0, 0.0, 8.0])) + torch.testing.assert_close(wp.to_torch(term._step_count), torch.tensor([0.0, 1.0, 0.0, 1.0])) + + term.command[:] = torch.tensor([[1.0, 2.0, 0.4]]).repeat(4, 1) + term.heading_target_torch[:] = torch.tensor([np.pi / 2, 0.0, np.pi / 2, 0.0]) + term.is_heading_env_torch[:] = torch.tensor([True, False, True, False]) + term.is_standing_env_torch[:] = torch.tensor([False, False, True, True]) + term._update_command() + wp.synchronize() + torch.testing.assert_close( + term.command, + torch.tensor([[1.0, 2.0, 0.3], [1.0, 2.0, 0.4], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), + ) + + +def test_uniform_velocity_command_replay_reads_current_root_state(): + """Captured command kernels should read current Tier-1 root state on every replay.""" + env = _Env() + cfg = _velocity_cfg() + cfg.ranges.ang_vel_z = (-2.0, 2.0) + term = UniformVelocityCommand(cfg, env) + robot_data = env.scene["robot"].data + + term.command.zero_() + term.command[:, 0] = 1.0 + term.heading_target_torch.fill_(np.pi / 2) + term.is_heading_env_torch.fill_(True) + term.is_standing_env_torch.fill_(False) + + with wp.ScopedCapture(device=env.device) as metrics_capture: + term._update_metrics() + with wp.ScopedCapture(device=env.device) as command_capture: + term._update_command() + + # Identity pose with matching world-frame velocity has zero body-frame XY error. + robot_data.root_vel_w.torch[:, 0] = 1.0 + wp.capture_launch(metrics_capture.graph) + wp.synchronize() + torch.testing.assert_close(wp.to_torch(term._error_xy_sum), torch.zeros(env.num_envs)) + + # A 90-degree root yaw exactly matches the heading target, so replay should command zero yaw rate. + half_sqrt = np.sqrt(0.5) + robot_data.root_pose_w.torch[:, 5] = half_sqrt + robot_data.root_pose_w.torch[:, 6] = half_sqrt + wp.capture_launch(command_capture.graph) + wp.synchronize() + torch.testing.assert_close(term.command[:, 2], torch.zeros(env.num_envs), atol=1.0e-6, rtol=1.0e-6) + + +def test_uniform_pose_command_matches_stable_math_and_tracks_sticky_success(): + """Pose commands should use xyzw math, update metrics, and reset sticky success by mask.""" + env = _Env() + term = UniformPoseCommand(_pose_cfg(), env) + term._prepare_reset_extras() + command_pointer = term.command.data_ptr() + command_wp = term.command_wp + term.command.fill_(-9.0) + wp.to_torch(term.time_left_wp)[:] = torch.tensor([-1.0, 6.0, -1.0, 8.0]) + wp.to_torch(term.command_counter_wp)[:] = torch.tensor([5, 6, 7, 8], dtype=torch.int32) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + term.reset(env_mask=env_mask) + wp.synchronize() + + expected_quat = quat_unique(quat_from_euler_xyz(torch.tensor([0.2]), torch.tensor([-0.3]), torch.tensor([4.0])))[0] + assert term.command_wp is command_wp + assert term.command.data_ptr() == command_pointer + torch.testing.assert_close(term.command[[0, 2], :3], torch.tensor([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]])) + torch.testing.assert_close(term.command[[0, 2], 3:], expected_quat.repeat(2, 1), atol=1.0e-6, rtol=1.0e-6) + torch.testing.assert_close(term.command[[1, 3]], torch.full((2, 7), -9.0)) + torch.testing.assert_close(wp.to_torch(term.time_left_wp), torch.tensor([3.0, 6.0, 3.0, 8.0])) + torch.testing.assert_close(wp.to_torch(term.command_counter_wp), torch.tensor([1, 6, 1, 8], dtype=torch.int32)) + + term.command[[1, 3]] = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]) + + roll = torch.tensor([0.0, 0.1, -0.2, 0.3]) + pitch = torch.tensor([0.0, -0.1, 0.2, -0.3]) + yaw = torch.tensor([0.5, -0.4, 0.3, -0.2]) + root_quat = quat_from_euler_xyz(roll, pitch, yaw) + env.scene["robot"].data.root_pos_w.torch[:] = torch.tensor( + [[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [-5.0, 1.0, 2.0], [1.0, 2.0, 3.0]] + ) + env.scene["robot"].data.root_quat_w.torch[:] = root_quat + expected_position_w, expected_orientation_w = combine_frame_transforms( + env.scene["robot"].data.root_pos_w.torch, + root_quat, + term.command[:, :3], + term.command[:, 3:], + ) + offsets = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.5, 0.0, 0.0]]) + env.scene["robot"].data.body_pos_w.torch[:, 0] = expected_position_w + offsets + identity_quat = torch.zeros(4, 4) + identity_quat[:, 3] = 1.0 + env.scene["robot"].data.body_quat_w.torch[:, 0] = identity_quat + + term._update_metrics() + wp.synchronize() + + expected_position_delta, expected_orientation_delta = compute_pose_error( + expected_position_w, + expected_orientation_w, + env.scene["robot"].data.body_pos_w.torch[:, 0], + env.scene["robot"].data.body_quat_w.torch[:, 0], + ) + torch.testing.assert_close(term.pose_command_w_torch[:, :3], expected_position_w, atol=1.0e-5, rtol=1.0e-5) + torch.testing.assert_close(term.pose_command_w_torch[:, 3:], expected_orientation_w, atol=1.0e-5, rtol=1.0e-5) + torch.testing.assert_close( + wp.to_torch(term.metrics["position_error"]), torch.linalg.norm(expected_position_delta, dim=-1) + ) + torch.testing.assert_close( + wp.to_torch(term.metrics["orientation_error"]), + torch.linalg.norm(expected_orientation_delta, dim=-1), + atol=1.0e-5, + rtol=1.0e-5, + ) + torch.testing.assert_close(wp.to_torch(term._success_rate), torch.tensor([1.0, 0.0, 1.0, 0.0])) + + env.scene["robot"].data.body_pos_w.torch[0, 0] += 2.0 + term._update_metrics() + wp.synchronize() + torch.testing.assert_close(wp.to_torch(term._success_rate), torch.tensor([1.0, 0.0, 1.0, 0.0])) + + reset_mask = wp.array([True, False, False, True], dtype=wp.bool, device="cpu") + extras = term.reset(env_mask=reset_mask) + wp.synchronize() + torch.testing.assert_close(extras["success_rate"], torch.tensor(0.5)) + torch.testing.assert_close(wp.to_torch(term._success_rate), torch.tensor([0.0, 0.0, 1.0, 0.0])) diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index 4ceec2ecb60b..4e452de96a28 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -383,7 +383,16 @@ class MockScene: def __init__(self, assets: dict, env_origins, sensors=None): self._assets = assets - self.env_origins = env_origins + if env_origins is None: + # Action-term tests don't need origins; keep the attributes present but empty. + self.env_origins = None + self.env_origins_pa = None + elif isinstance(env_origins, wp.array): + self.env_origins_pa = ProxyArray(env_origins) + self.env_origins = wp.to_torch(env_origins) + else: + self.env_origins = env_origins + self.env_origins_pa = ProxyArray(wp.from_torch(env_origins, dtype=wp.vec3f)) self.sensors = sensors or {} self.articulations = dict(assets) self.rigid_objects = {} @@ -533,11 +542,15 @@ class MockCommandManager: def __init__(self, command_tensor: torch.Tensor, cmd_term: MockCommandTerm): self._cmd = command_tensor + self._cmd_wp = wp.from_torch(command_tensor, dtype=wp.float32) self._term = cmd_term def get_command(self, name: str) -> torch.Tensor: return self._cmd + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32, ndim=2): + return self._cmd_wp + def get_term(self, name: str): return self._term diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py index 24b9eb022665..aba271420a69 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py @@ -17,6 +17,7 @@ pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") import isaaclab_experimental.envs.mdp.events as warp_evt +from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg from parity_helpers import ( DEVICE, NUM_ACTIONS, @@ -37,25 +38,6 @@ # ============================================================================ -@pytest.fixture(autouse=True) -def _clear_function_caches(): - """Clear first-call caches on warp MDP functions so each test starts fresh. - - Functions that cache warp views via the ``hasattr`` pattern need clearing - between tests to avoid stale references from prior fixtures. - """ - yield - for fn in ( - warp_evt.push_by_setting_velocity, - warp_evt.apply_external_force_torque, - warp_evt.reset_root_state_uniform, - warp_evt.randomize_rigid_body_com, - ): - for attr in list(vars(fn)): - if attr.startswith("_"): - delattr(fn, attr) - - @pytest.fixture() def art_data(): return MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE) @@ -99,6 +81,7 @@ class _Env: env.action_manager = MockActionManagerWarp(action_wp[0], action_wp[1]) env.num_envs = NUM_ENVS env.device = DEVICE + env.env_origins_wp = scene.env_origins_pa.warp env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 @@ -130,6 +113,37 @@ def all_joints_cfg(): return MockSceneEntityCfg("robot", list(range(NUM_JOINTS)), NUM_JOINTS, DEVICE) +def _make_event_env(seed: int, *, num_bodies: int = 1): + """Create an independent Warp event environment for ownership tests.""" + + class _Env: + pass + + data = MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE, seed=seed, num_bodies=num_bodies) + asset = MockArticulation(data, num_bodies=num_bodies) + origins = wp.zeros(NUM_ENVS, dtype=wp.vec3f, device=DEVICE) + env = _Env() + env.scene = MockScene({"robot": asset}, origins) + env.num_envs = NUM_ENVS + env.device = DEVICE + env.env_origins_wp = env.scene.env_origins_pa.warp + env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + seed, device=DEVICE) + return env, data, asset + + +def _make_event_term( + term_type: type[ManagerTermBase], + env: object, + *, + mode: str = "reset", + **params: object, +) -> ManagerTermBase: + """Create a persistent Warp event term for a mock environment.""" + params.setdefault("asset_cfg", SceneEntityCfg("robot")) + cfg = EventTermCfg(func=term_type, mode=mode, params=params) + return term_type(cfg, env) + + # ============================================================================ # Event parity tests: deterministic (zero-width range) warp vs stable # ============================================================================ @@ -268,6 +282,8 @@ def test_reset_joints_by_scale(self, warp_env, art_data, all_joints_cfg): def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): """With zero-width velocity range, scratch == root_vel_w. Mutate root_vel_w -> scratch tracks.""" mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + captured = {} + warp_env.scene["robot"].write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) zero_range = { "x": (0.0, 0.0), "y": (0.0, 0.0), @@ -276,10 +292,16 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): "pitch": (0.0, 0.0), "yaw": (0.0, 0.0), } + term = _make_event_term( + warp_evt.push_by_setting_velocity, + warp_env, + mode="interval", + velocity_range=zero_range, + ) - warp_evt.push_by_setting_velocity(warp_env, mask, velocity_range=zero_range) + term(warp_env, mask, **term.cfg.params) with wp.ScopedCapture() as cap: - warp_evt.push_by_setting_velocity(warp_env, mask, velocity_range=zero_range) + term(warp_env, mask, **term.cfg.params) # Mutate root_vel_w new_vel = np.tile([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], (NUM_ENVS, 1)).astype(np.float32) @@ -288,7 +310,7 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): wp.capture_launch(cap.graph) wp.synchronize() - scratch = wp.to_torch(warp_evt.push_by_setting_velocity._scratch_vel) + scratch = wp.to_torch(captured["root_velocity"]) expected = torch.tensor([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], device=DEVICE).expand(NUM_ENVS, -1) assert_close(scratch, expected) @@ -297,21 +319,73 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): def test_apply_external_force_torque(self, warp_env, art_data, all_joints_cfg): """With zero-width ranges, forces/torques are zero. Non-zero ranges produce non-zero output.""" mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + captured = {} + warp_env.scene["robot"].permanent_wrench_composer.set_forces_and_torques_mask = ( + lambda **kwargs: captured.update(kwargs) + ) + zero_range = (0.0, 0.0) + term = _make_event_term( + warp_evt.apply_external_force_torque, + warp_env, + force_range=zero_range, + torque_range=zero_range, + ) # Zero-range: forces and torques should be zero - warp_evt.apply_external_force_torque(warp_env, mask, force_range=(0.0, 0.0), torque_range=(0.0, 0.0)) + term(warp_env, mask, **term.cfg.params) with wp.ScopedCapture() as cap: - warp_evt.apply_external_force_torque(warp_env, mask, force_range=(0.0, 0.0), torque_range=(0.0, 0.0)) + term(warp_env, mask, **term.cfg.params) wp.capture_launch(cap.graph) wp.synchronize() - forces = wp.to_torch(warp_evt.apply_external_force_torque._scratch_forces) - torques = wp.to_torch(warp_evt.apply_external_force_torque._scratch_torques) + forces = wp.to_torch(captured["forces"]) + torques = wp.to_torch(captured["torques"]) assert_close(forces, torch.zeros_like(forces)) assert_close(torques, torch.zeros_like(torques)) # -- reset_root_state_uniform ----------------------------------------------- + def test_reset_root_state_uniform_composes_orientation_like_stable(self): + """Pin the quaternion path: ``final = default ∘ euler_xyz(roll, pitch, yaw)``. + + Zero-width ranges at non-zero angles make the uniform draw deterministic, + so the composition order, Euler convention, and xyzw storage layout are + all pinned against the stable math reference. + """ + import isaaclab.utils.math as math_utils + + env, data, asset = _make_event_env(1234) + # Non-identity default orientation (90 deg about x), xyzw layout. + default_pose = np.zeros((NUM_ENVS, 7), dtype=np.float32) + default_pose[:, 3] = np.sin(np.pi / 4.0) + default_pose[:, 6] = np.cos(np.pi / 4.0) + copy_np_to_wp(data.default_root_pose, default_pose) + copy_np_to_wp(data.default_root_vel, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = {} + asset.write_root_pose_to_sim_mask = lambda **kwargs: captured.update(kwargs) + asset.write_root_velocity_to_sim_mask = lambda **kwargs: None + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + roll, pitch, yaw = 0.3, -0.4, 1.1 + term = _make_event_term( + warp_evt.reset_root_state_uniform, + env, + pose_range={"roll": (roll, roll), "pitch": (pitch, pitch), "yaw": (yaw, yaw)}, + velocity_range={}, + ) + + term(env, env_mask, **term.cfg.params) + wp.synchronize() + + result_xyzw = wp.to_torch(captured["root_pose"])[:, 3:7] + angles = torch.tensor([[roll], [pitch], [yaw]], device=DEVICE) + # Core math utils use the (x, y, z, w) layout, matching the warp buffers. + delta_xyzw = math_utils.quat_from_euler_xyz(angles[0], angles[1], angles[2]) + default_xyzw = torch.tensor([[np.sin(np.pi / 4.0), 0.0, 0.0, np.cos(np.pi / 4.0)]], device=DEVICE) + expected_xyzw = math_utils.quat_mul(default_xyzw, delta_xyzw).expand(NUM_ENVS, 4) + # Quaternions double-cover rotations: align signs before comparing. + sign = torch.sign((result_xyzw * expected_xyzw).sum(dim=1, keepdim=True)) + assert_close(result_xyzw * sign, expected_xyzw) + # -- env_mask selectivity --------------------------------------------------- def test_reset_joints_mask_selectivity(self, warp_env, art_data, all_joints_cfg): @@ -338,3 +412,277 @@ def test_reset_joints_mask_selectivity(self, warp_env, art_data, all_joints_cfg) assert_close(result[: NUM_ENVS // 2], torch.zeros(NUM_ENVS // 2, NUM_JOINTS, device=DEVICE)) # Unmasked envs: still 999.0 assert_close(result[NUM_ENVS // 2 :], torch.full((NUM_ENVS // 2, NUM_JOINTS), 999.0, device=DEVICE)) + + +class TestEventStateOwnership: + """Verify event scratch and parsed ranges are owned by one environment configuration.""" + + def test_com_term_is_marked_non_capturable(self): + assert warp_evt.randomize_rigid_body_com._warp_capturable is False + + def test_push_state_is_not_shared_between_environments(self): + env_a, data_a, asset_a = _make_event_env(101) + env_b, data_b, asset_b = _make_event_env(202) + copy_np_to_wp(data_a.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + copy_np_to_wp(data_b.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured_a = {} + captured_b = {} + asset_a.write_root_velocity_to_sim_mask = lambda **kwargs: captured_a.update(kwargs) + asset_b.write_root_velocity_to_sim_mask = lambda **kwargs: captured_b.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + range_a = {"x": (1.0, 1.0)} + range_b = {"x": (2.0, 2.0)} + term_a = _make_event_term(warp_evt.push_by_setting_velocity, env_a, mode="interval", velocity_range=range_a) + term_b = _make_event_term(warp_evt.push_by_setting_velocity, env_b, mode="interval", velocity_range=range_b) + + term_a(env_a, env_mask, **term_a.cfg.params) + term_b(env_b, env_mask, **term_b.cfg.params) + wp.synchronize() + + velocity_a = captured_a["root_velocity"] + velocity_b = captured_b["root_velocity"] + assert velocity_a.ptr != velocity_b.ptr + assert_close(wp.to_torch(velocity_a)[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + assert_close(wp.to_torch(velocity_b)[:, 0], torch.full((NUM_ENVS,), 2.0, device=DEVICE)) + + def test_push_state_is_not_shared_between_configurations(self): + env, data, asset = _make_event_env(252) + copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = [] + asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.append(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + range_a = {"x": (1.0, 1.0)} + range_b = {"x": (2.0, 2.0)} + term_a = _make_event_term(warp_evt.push_by_setting_velocity, env, mode="interval", velocity_range=range_a) + term_b = _make_event_term(warp_evt.push_by_setting_velocity, env, mode="interval", velocity_range=range_b) + + term_a(env, env_mask, **term_a.cfg.params) + term_b(env, env_mask, **term_b.cfg.params) + wp.synchronize() + + velocity_a = captured[0]["root_velocity"] + velocity_b = captured[1]["root_velocity"] + assert velocity_a.ptr != velocity_b.ptr + assert_close(wp.to_torch(velocity_a)[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + assert_close(wp.to_torch(velocity_b)[:, 0], torch.full((NUM_ENVS,), 2.0, device=DEVICE)) + + def test_push_term_snapshots_range_during_initialization(self): + env, data, asset = _make_event_env(277) + copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = {} + asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + velocity_range = {"x": (1.0, 1.0)} + term = _make_event_term( + warp_evt.push_by_setting_velocity, + env, + mode="interval", + velocity_range=velocity_range, + ) + + term(env, env_mask, **term.cfg.params) + wp.synchronize() + assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + + velocity_range["x"] = (4.0, 4.0) + term(env, env_mask, **term.cfg.params) + wp.synchronize() + + assert_close(wp.to_torch(captured["root_velocity"])[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + + def test_push_state_respects_sparse_mask(self): + env, data, asset = _make_event_env(303) + copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured = {} + asset.write_root_velocity_to_sim_mask = lambda **kwargs: captured.update(kwargs) + mask_np = np.arange(NUM_ENVS) % 3 == 0 + env_mask = wp.array(mask_np, dtype=wp.bool, device=DEVICE) + term = _make_event_term( + warp_evt.push_by_setting_velocity, + env, + mode="interval", + velocity_range={"x": (3.0, 3.0)}, + ) + + term(env, env_mask, **term.cfg.params) + wp.synchronize() + + velocity = wp.to_torch(captured["root_velocity"]) + mask = torch.from_numpy(mask_np).to(device=DEVICE) + assert_close(velocity[mask, 0], torch.full((int(mask_np.sum()),), 3.0, device=DEVICE)) + assert_close(velocity[~mask], torch.zeros((int((~mask_np).sum()), 6), device=DEVICE)) + + def test_external_wrench_state_is_not_shared_between_environments(self): + env_a, _, asset_a = _make_event_env(404, num_bodies=2) + env_b, _, asset_b = _make_event_env(505, num_bodies=2) + captured_a = {} + captured_b = {} + asset_a.permanent_wrench_composer.set_forces_and_torques_mask = lambda **kwargs: captured_a.update(kwargs) + asset_b.permanent_wrench_composer.set_forces_and_torques_mask = lambda **kwargs: captured_b.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + force_range_a = (1.0, 1.0) + force_range_b = (2.0, 2.0) + torque_range = (0.0, 0.0) + term_a = _make_event_term( + warp_evt.apply_external_force_torque, + env_a, + force_range=force_range_a, + torque_range=torque_range, + ) + term_b = _make_event_term( + warp_evt.apply_external_force_torque, + env_b, + force_range=force_range_b, + torque_range=torque_range, + ) + + term_a(env_a, env_mask, **term_a.cfg.params) + term_b(env_b, env_mask, **term_b.cfg.params) + wp.synchronize() + + forces_a = captured_a["forces"] + forces_b = captured_b["forces"] + assert forces_a.ptr != forces_b.ptr + assert_close(wp.to_torch(forces_a), torch.ones((NUM_ENVS, 2, 3), device=DEVICE)) + assert_close(wp.to_torch(forces_b), torch.full((NUM_ENVS, 2, 3), 2.0, device=DEVICE)) + + def test_external_wrench_respects_body_selection(self): + env, _, asset = _make_event_env(550, num_bodies=3) + captured = {} + asset.permanent_wrench_composer.set_forces_and_torques_mask = lambda **kwargs: captured.update(kwargs) + asset_cfg = SceneEntityCfg("robot", body_ids=[1]) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + term = _make_event_term( + warp_evt.apply_external_force_torque, + env, + force_range=(2.0, 2.0), + torque_range=(0.0, 0.0), + asset_cfg=asset_cfg, + ) + + term(env, env_mask, **term.cfg.params) + wp.synchronize() + + expected_body_mask = torch.tensor([False, True, False], dtype=torch.bool, device=DEVICE) + assert torch.equal(wp.to_torch(captured["body_mask"]), expected_body_mask) + forces = wp.to_torch(captured["forces"]) + assert_close(forces[:, 0], torch.zeros((NUM_ENVS, 3), device=DEVICE)) + assert_close(forces[:, 1], torch.full((NUM_ENVS, 3), 2.0, device=DEVICE)) + assert_close(forces[:, 2], torch.zeros((NUM_ENVS, 3), device=DEVICE)) + + def test_root_reset_state_is_not_shared_between_environments(self): + env_a, data_a, asset_a = _make_event_env(606) + env_b, data_b, asset_b = _make_event_env(707) + default_pose = np.zeros((NUM_ENVS, 7), dtype=np.float32) + default_pose[:, 6] = 1.0 + copy_np_to_wp(data_a.default_root_pose, default_pose) + copy_np_to_wp(data_b.default_root_pose, default_pose) + copy_np_to_wp(data_a.default_root_vel, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + copy_np_to_wp(data_b.default_root_vel, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + captured_a = {} + captured_b = {} + asset_a.write_root_pose_to_sim_mask = lambda **kwargs: captured_a.update(kwargs) + asset_b.write_root_pose_to_sim_mask = lambda **kwargs: captured_b.update(kwargs) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + pose_range_a = {"x": (1.0, 1.0)} + pose_range_b = {"x": (2.0, 2.0)} + velocity_range = {} + term_a = _make_event_term( + warp_evt.reset_root_state_uniform, + env_a, + pose_range=pose_range_a, + velocity_range=velocity_range, + ) + term_b = _make_event_term( + warp_evt.reset_root_state_uniform, + env_b, + pose_range=pose_range_b, + velocity_range=velocity_range, + ) + + term_a(env_a, env_mask, **term_a.cfg.params) + term_b(env_b, env_mask, **term_b.cfg.params) + wp.synchronize() + + pose_a = captured_a["root_pose"] + pose_b = captured_b["root_pose"] + assert pose_a.ptr != pose_b.ptr + assert_close(wp.to_torch(pose_a)[:, 0], torch.ones(NUM_ENVS, device=DEVICE)) + assert_close(wp.to_torch(pose_b)[:, 0], torch.full((NUM_ENVS,), 2.0, device=DEVICE)) + + def test_com_ranges_are_not_shared_between_environments(self): + env_a, data_a, asset_a = _make_event_env(808, num_bodies=2) + env_b, data_b, asset_b = _make_event_env(909, num_bodies=2) + copy_np_to_wp(data_a.body_com_pos_b, np.zeros((NUM_ENVS, 2, 3), dtype=np.float32)) + copy_np_to_wp(data_b.body_com_pos_b, np.zeros((NUM_ENVS, 2, 3), dtype=np.float32)) + asset_a.set_coms_mask = lambda **kwargs: None + asset_b.set_coms_mask = lambda **kwargs: None + asset_cfg_a = SceneEntityCfg("robot") + asset_cfg_b = SceneEntityCfg("robot") + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + com_range_a = {"x": (1.0, 1.0)} + com_range_b = {"x": (2.0, 2.0)} + term_a = _make_event_term( + warp_evt.randomize_rigid_body_com, + env_a, + com_range=com_range_a, + asset_cfg=asset_cfg_a, + ) + term_b = _make_event_term( + warp_evt.randomize_rigid_body_com, + env_b, + com_range=com_range_b, + asset_cfg=asset_cfg_b, + ) + + term_a(env_a, env_mask, **term_a.cfg.params) + term_b(env_b, env_mask, **term_b.cfg.params) + wp.synchronize() + + assert term_a._default_com.ptr != term_b._default_com.ptr + assert_close(data_a.body_com_pos_b.torch[..., 0], torch.ones((NUM_ENVS, 2), device=DEVICE)) + assert_close(data_b.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 2.0, device=DEVICE)) + + def test_com_randomization_uses_persistent_baseline(self): + env, data, asset = _make_event_env(1001, num_bodies=2) + baseline = np.zeros((NUM_ENVS, 2, 3), dtype=np.float32) + baseline[..., 0] = 0.25 + copy_np_to_wp(data.body_com_pos_b, baseline) + asset.set_coms_mask = lambda **kwargs: None + asset_cfg = SceneEntityCfg("robot", body_ids=[0, 1]) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + com_range = {"x": (1.0, 1.0)} + term = _make_event_term( + warp_evt.randomize_rigid_body_com, + env, + com_range=com_range, + asset_cfg=asset_cfg, + ) + + term(env, env_mask, **term.cfg.params) + wp.synchronize() + assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.25, device=DEVICE)) + + term(env, env_mask, **term.cfg.params) + wp.synchronize() + assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.25, device=DEVICE)) + + def test_com_randomization_broadcasts_one_offset_across_selected_bodies(self): + env, data, asset = _make_event_env(1111, num_bodies=3) + baseline = np.zeros((NUM_ENVS, 3, 3), dtype=np.float32) + copy_np_to_wp(data.body_com_pos_b, baseline) + asset.set_coms_mask = lambda **kwargs: None + asset_cfg = SceneEntityCfg("robot", body_ids=[0, 2]) + env_mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + term = _make_event_term( + warp_evt.randomize_rigid_body_com, + env, + com_range={"x": (-1.0, 1.0), "y": (-2.0, 2.0), "z": (-3.0, 3.0)}, + asset_cfg=asset_cfg, + ) + + term(env, env_mask, **term.cfg.params) + wp.synchronize() + + assert_close(data.body_com_pos_b.torch[:, 0], data.body_com_pos_b.torch[:, 2]) + assert_close(data.body_com_pos_b.torch[:, 1], torch.zeros((NUM_ENVS, 3), device=DEVICE)) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py index 9f1eea23bf8d..8654949bbed5 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py @@ -47,15 +47,6 @@ # ============================================================================ -@pytest.fixture(autouse=True) -def _clear_caches(): - yield - for fn in [warp_obs.generated_commands]: - for attr in list(vars(fn)): - if attr.startswith("_"): - delattr(fn, attr) - - @pytest.fixture() def art_data(): return MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE) 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..10c9be10ec6b 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 @@ -55,7 +55,7 @@ @pytest.fixture(autouse=True) def _clear_caches(): yield - for fn in [warp_rew.track_lin_vel_xy_exp, warp_rew.track_ang_vel_z_exp, warp_rew.undesired_contacts]: + for fn in [warp_rew.undesired_contacts]: for attr in list(vars(fn)): if attr.startswith("_"): delattr(fn, attr) diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py new file mode 100644 index 000000000000..20c043180434 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -0,0 +1,579 @@ +# 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 + +"""Tests for :mod:`isaaclab_experimental.envs.frontend`. + +Pure-Python unit tests; no app launch. Covers: + +* :class:`Workflow` enum surface. +* :func:`SceneEntityCfg.from_stable` field copy. +* :func:`_require_newton_physics` hard-check. +* :func:`_walk_terms` recursive ManagerTermBaseCfg discovery over instance + attributes (nested class objects are not descended into). +* :meth:`_mirror_module` / :meth:`_nearest_mdp_module` deterministic mirror + resolution and :meth:`_cfg_route_modules` cfg-hierarchy resolution. +* :func:`_promote_scene_entity_cfgs` walks ``term.params`` dicts. +* :func:`_swap_mdp` swaps ``func`` *and* ``class_type``; raises with a path + list when twins are missing. +* :func:`_resolve_warp_twin` rejects stable-origin re-exports. +* :func:`_assert_direct_warp_registration` accepts warp-rooted entry + points and rejects stable ones. +""" + +from __future__ import annotations + +import contextlib +import types +from typing import Any +from unittest.mock import patch + +import gymnasium as gym +import isaaclab_experimental.envs.frontend as fe +import pytest +from isaaclab_experimental.envs.frontend import ( + FrontendIncompatibleError, + WarpFrontend, + Workflow, +) +from isaaclab_experimental.managers.scene_entity_cfg import SceneEntityCfg as WarpSceneEntityCfg +from isaaclab_newton.physics import NewtonCfg +from isaaclab_physx.physics import PhysxCfg + +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers.manager_term_cfg import EventTermCfg, ObservationTermCfg, RewardTermCfg +from isaaclab.managers.scene_entity_cfg import SceneEntityCfg as StableSceneEntityCfg +from isaaclab.utils.configclass import configclass + +# ====================================================================== +# Fixtures: fake stable/warp symbols and configclass trees. +# +# These mirror the real cfg shape so :func:`_walk_terms` descends into them +# (it only descends into objects with ``__dataclass_fields__``). Term cfgs +# use the real :class:`EventTermCfg`/:class:`RewardTermCfg`/:class:`ObservationTermCfg` +# so the walker's ``isinstance(ManagerTermBaseCfg)`` discriminator yields them. +# ====================================================================== + + +def _stable_func(env, **params): + return None + + +class _StableActionCls: + pass + + +_stable_func.__module__ = "isaaclab_tasks.fake_task.mdp" +_StableActionCls.__module__ = "isaaclab_tasks.fake_task.mdp" + + +def _warp_twin_func(env, out, **params): + return None + + +class _WarpActionCls: + pass + + +_warp_twin_func.__module__ = "isaaclab_experimental.envs.mdp" +_WarpActionCls.__module__ = "isaaclab_experimental.envs.mdp" + + +@configclass +class _PolicyObsGroup: + """Stand-in for a per-task ObservationsCfg sub-group (e.g. PolicyCfg).""" + + o1: ObservationTermCfg | None = None + o2: ObservationTermCfg | None = None + + +@configclass +class _ExtraObsGroup: + """A second obs group (named arbitrarily) to exercise multi-group walks.""" + + o3: ObservationTermCfg | None = None + + +@configclass +class _ObservationsCfg: + policy: _PolicyObsGroup | None = None + perception: _ExtraObsGroup | None = None + + +@configclass +class _RewardsCfg: + r1: RewardTermCfg | None = None + r2: RewardTermCfg | None = None + + +@configclass +class _EventsCfg: + e1: EventTermCfg | None = None + + +@configclass +class _CurriculumCfg: + c1: EventTermCfg | None = None + + +@configclass +class _CfgFixture: + observations: _ObservationsCfg | None = None + rewards: _RewardsCfg | None = None + events: _EventsCfg | None = None + curriculum: _CurriculumCfg | None = None + + +def _term(func=None, params: dict | None = None) -> RewardTermCfg: + """Cheap RewardTermCfg builder for tests; the cfg class is irrelevant for swap/walk logic.""" + return RewardTermCfg(func=func or _stable_func, weight=1.0, params=params or {}) + + +# ====================================================================== +# Enums +# ====================================================================== + + +def test_workflow_values(): + assert Workflow.MANAGER_BASED == "manager_based" + assert Workflow.DIRECT == "direct" + + +# ====================================================================== +# build() +# ====================================================================== + + +def test_warp_manager_build_constructs_warp_env_with_cfg(): + # Cfg adaptation happens inside ManagerBasedRLEnvWarp.__init__, so build() + # only selects the env class and forwards the cfg and construct kwargs. + import isaaclab_experimental.envs as envs + + cfg = object.__new__(ManagerBasedRLEnvCfg) + expected_env = object() + calls: list[tuple[Any, Any]] = [] + + def fake_env(*, cfg: Any, **kwargs: Any) -> Any: + calls.append((cfg, kwargs)) + return expected_env + + with patch.object(envs, "ManagerBasedRLEnvWarp", fake_env): + env = WarpFrontend.build_env(cfg, "Isaac-Test", render_mode="rgb_array") + + assert env is expected_env + assert calls == [(cfg, {"render_mode": "rgb_array"})] + + +# ====================================================================== +# SceneEntityCfg.from_stable +# ====================================================================== + + +def test_from_stable_copies_minimum_fields(): + stable = StableSceneEntityCfg(name="robot") + warp = WarpSceneEntityCfg.from_stable(stable) + assert isinstance(warp, WarpSceneEntityCfg) + assert warp.name == "robot" + # Warp-only fields stay None until :meth:`resolve` runs. + assert warp.joint_mask is None + assert warp.joint_ids_wp is None + assert warp.body_ids_wp is None + + +def test_from_stable_copies_all_selection_fields(): + stable = StableSceneEntityCfg( + name="robot", + joint_names=["lf_hip"], + joint_ids=[0, 1, 2], + fixed_tendon_names=["tendon_a"], + fixed_tendon_ids=[5], + body_names=["base", "lf_foot"], + body_ids=[0, 4], + object_collection_names=["objs"], + object_collection_ids=[7], + preserve_order=True, + ) + warp = WarpSceneEntityCfg.from_stable(stable) + for field in ( + "name", + "joint_names", + "joint_ids", + "fixed_tendon_names", + "fixed_tendon_ids", + "body_names", + "body_ids", + "object_collection_names", + "object_collection_ids", + "preserve_order", + ): + assert getattr(warp, field) == getattr(stable, field), f"field {field!r} mismatch" + + +# ====================================================================== +# _require_newton_physics +# ====================================================================== + + +def _cfg_with_physics(physics: Any) -> Any: + cfg = types.SimpleNamespace() + cfg.sim = types.SimpleNamespace(physics=physics) + return cfg + + +def test_require_newton_passes_for_newton(): + fe.WarpFrontend._require_newton_physics(_cfg_with_physics(NewtonCfg()), "Isaac-Test-v0") # no raise + + +def test_require_newton_rejects_physx(): + with pytest.raises(FrontendIncompatibleError) as exc: + fe.WarpFrontend._require_newton_physics(_cfg_with_physics(PhysxCfg()), "Isaac-Test-v0") + assert "presets=newton_mjwarp" in str(exc.value) + assert "PhysxCfg" in str(exc.value) + + +def test_require_newton_rejects_none(): + with pytest.raises(FrontendIncompatibleError): + fe.WarpFrontend._require_newton_physics(_cfg_with_physics(None), "Isaac-Test-v0") + + +# ====================================================================== +# _walk_terms +# ====================================================================== + + +def test_walk_terms_yields_each_term_with_its_path(): + cfg = _CfgFixture( + rewards=_RewardsCfg(r1=_term(), r2=_term()), + events=_EventsCfg(e1=EventTermCfg(func=_stable_func, mode="reset")), + ) + # Configclass instances aren't hashable, so collect paths only. + paths = {".".join(p) for p, _ in fe.WarpFrontend._walk_terms(cfg)} + assert paths == {"rewards.r1", "rewards.r2", "events.e1"} + + +def test_walk_terms_descends_into_obs_subgroups(): + cfg = _CfgFixture( + observations=_ObservationsCfg( + policy=_PolicyObsGroup(o1=ObservationTermCfg(func=_stable_func)), + perception=_ExtraObsGroup(o3=ObservationTermCfg(func=_stable_func)), + ), + ) + paths = {".".join(p) for p, _ in fe.WarpFrontend._walk_terms(cfg)} + # Discovery is purely type-driven; no obs group name is hardcoded. + assert paths == {"observations.policy.o1", "observations.perception.o3"} + + +def test_walk_terms_stops_at_terms(): + # The walker must not descend into term.params / term.func — yields the term itself. + nested_se_cfg = StableSceneEntityCfg(name="robot") + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": nested_se_cfg}))) + terms = list(fe.WarpFrontend._walk_terms(cfg)) + assert len(terms) == 1 + _, term = terms[0] + assert isinstance(term, RewardTermCfg) + + +def test_walk_terms_skips_non_configclass_attrs(): + # A namespace without __dataclass_fields__ is not descended into. + cfg = types.SimpleNamespace(some_plain_attr="hello") + assert list(fe.WarpFrontend._walk_terms(cfg)) == [] + + +def test_walk_terms_skips_none_subtrees(): + cfg = _CfgFixture(rewards=None, events=None) + assert list(fe.WarpFrontend._walk_terms(cfg)) == [] + + +@configclass +class _GroupWithNestedClass: + """Mirrors the real ``ObservationsCfg.PolicyCfg`` nested-class pattern.""" + + @configclass + class TemplateCfg: + t1: ObservationTermCfg | None = None + + o1: ObservationTermCfg | None = None + + +def test_walk_terms_ignores_nested_class_objects(): + # The walker iterates instance attributes only (mirroring the warp managers' + # ``cfg.__dict__`` consumption), so a nested class object — reachable as an + # attribute on the instance — must never contribute paths. + cfg = _GroupWithNestedClass(o1=ObservationTermCfg(func=_stable_func)) + paths = {".".join(p) for p, _ in fe.WarpFrontend._walk_terms(cfg)} + assert paths == {"o1"} + + +# ====================================================================== +# Mirror resolution +# ====================================================================== + + +def test_mirror_module_maps_stable_roots(): + assert ( + fe.WarpFrontend._mirror_module("isaaclab_tasks.core.velocity.mdp.rewards") + == "isaaclab_tasks_experimental.core.velocity.mdp.rewards" + ) + assert fe.WarpFrontend._mirror_module("isaaclab.envs.mdp.events") == "isaaclab_experimental.envs.mdp.events" + assert fe.WarpFrontend._mirror_module("isaaclab_tasks") == "isaaclab_tasks_experimental" + # Unmirrored roots are not resolvable. + assert fe.WarpFrontend._mirror_module("torch.nn.functional") is None + assert fe.WarpFrontend._mirror_module("") is None + + +def test_nearest_mdp_module_truncates_at_mdp_boundary(): + # A symbol module resolves to its own mdp package on the mirrored path. + assert ( + fe.WarpFrontend._nearest_mdp_module("isaaclab_tasks_experimental.core.locomotion.mdp.rewards") + == "isaaclab_tasks_experimental.core.locomotion.mdp" + ) + assert fe.WarpFrontend._nearest_mdp_module("isaaclab_experimental.envs.mdp.events") == ( + "isaaclab_experimental.envs.mdp" + ) + + +def test_nearest_mdp_module_walks_up_cfg_paths(): + # A cfg module path walks up to the task family's mdp package. + assert ( + fe.WarpFrontend._nearest_mdp_module("isaaclab_tasks_experimental.core.velocity.config.anymal_d.flat_env_cfg") + == "isaaclab_tasks_experimental.core.velocity.mdp" + ) + # Nothing on the path provides an mdp package. + assert fe.WarpFrontend._nearest_mdp_module("isaaclab_tasks_experimental.nonexistent.task_pkg") is None + + +def test_cfg_route_modules_resolves_mirrored_hierarchy(): + # A class whose module mirrors into the experimental tree resolves to the + # task family's mdp package; unmirrored classes contribute nothing. + fixture_cls = type("FakeVelocityCfg", (), {}) + fixture_cls.__module__ = "isaaclab_tasks.core.velocity.config.anymal_d.flat_env_cfg" + modules = fe.WarpFrontend._cfg_route_modules(fixture_cls()) + assert [m.__name__ for m in modules] == ["isaaclab_tasks_experimental.core.velocity.mdp"] + + assert fe.WarpFrontend._cfg_route_modules(_CfgFixture()) == [] + + +# ====================================================================== +# _swap_mdp +# ====================================================================== + + +class _FakeMdpModule: + """Stand-in for a warp mdp module containing twins.""" + + __name__ = "test_fake_warp_mdp" + + +def _fake_mdp_module(name_to_symbol: dict[str, Any]) -> _FakeMdpModule: + module = _FakeMdpModule() + for name, symbol in name_to_symbol.items(): + setattr(module, name, symbol) + return module + + +def _patch_twin_modules(monkeypatch: pytest.MonkeyPatch, modules: list[Any]) -> None: + """Pin twin resolution to ``modules`` and keep unit tests hermetic.""" + monkeypatch.setattr(fe.WarpFrontend, "_cfg_route_modules", classmethod(lambda cls, cfg: [])) + monkeypatch.setattr( + fe.WarpFrontend, "_twin_modules", classmethod(lambda cls, symbol_module, cfg_route_modules: modules) + ) + + +def test_swap_mdp_swaps_func_and_class_type(monkeypatch: pytest.MonkeyPatch): + fake = _fake_mdp_module({"_stable_func": _warp_twin_func, "_StableActionCls": _WarpActionCls}) + _patch_twin_modules(monkeypatch, [fake]) + term_reward = _term(func=_stable_func) + term_action = _term() + term_action.class_type = _StableActionCls # set attr to exercise class_type swap + cfg = _CfgFixture(rewards=_RewardsCfg(r1=term_reward, r2=term_action)) + fe.WarpFrontend._swap_mdp(cfg, "Isaac-Test-v0") + assert cfg.rewards.r1.func is _warp_twin_func + assert cfg.rewards.r2.class_type is _WarpActionCls + + +def test_swap_mdp_missing_twin_raises_with_path_list(monkeypatch: pytest.MonkeyPatch): + _patch_twin_modules(monkeypatch, [_fake_mdp_module({})]) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_stable_func))) + with pytest.raises(FrontendIncompatibleError) as exc: + fe.WarpFrontend._swap_mdp(cfg, "Isaac-Test-v0") + msg = str(exc.value) + assert "rewards.r1.func" in msg + assert "_stable_func" in msg + # The cfg term wasn't mutated for the missing twin. + assert cfg.rewards.r1.func is _stable_func + + +def test_swap_mdp_skips_already_warp(monkeypatch: pytest.MonkeyPatch): + _patch_twin_modules(monkeypatch, [_fake_mdp_module({})]) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(func=_warp_twin_func))) + fe.WarpFrontend._swap_mdp(cfg, "Isaac-Test-v0") # no raise + assert cfg.rewards.r1.func is _warp_twin_func + + +# ====================================================================== +# _promote_scene_entity_cfgs +# ====================================================================== + + +def test_promote_scene_entity_cfgs_promotes_in_params(): + cfg = _CfgFixture( + rewards=_RewardsCfg( + r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="robot", joint_names=["lf_hip"]), "scale": 1.0}) + ) + ) + fe.WarpFrontend._promote_scene_entity_cfgs(cfg) + promoted = cfg.rewards.r1.params["asset_cfg"] + assert isinstance(promoted, WarpSceneEntityCfg) + assert promoted.name == "robot" + assert promoted.joint_names == ["lf_hip"] + # Non-SceneEntityCfg params are untouched. + assert cfg.rewards.r1.params["scale"] == 1.0 + + +def test_promote_scene_entity_cfgs_skips_already_warp(): + # configclass init deep-copies params, so identity won't hold across + # construction; what we actually want to assert is "no re-promotion": + # the asset_cfg remains a WarpSceneEntityCfg (i.e., wasn't passed + # back through `from_stable`). + warp = WarpSceneEntityCfg(name="robot", joint_names=["lf_hip"]) + cfg = _CfgFixture(rewards=_RewardsCfg(r1=_term(params={"asset_cfg": warp}))) + before = cfg.rewards.r1.params["asset_cfg"] + fe.WarpFrontend._promote_scene_entity_cfgs(cfg) + after = cfg.rewards.r1.params["asset_cfg"] + assert isinstance(after, WarpSceneEntityCfg) + # The asset_cfg object was not replaced by another from_stable call. + assert after is before + + +def test_promote_scene_entity_cfgs_walks_all_term_groups(): + cfg = _CfgFixture( + rewards=_RewardsCfg(r1=_term(params={"asset_cfg": StableSceneEntityCfg(name="r")})), + events=_EventsCfg( + e1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="e")}) + ), + observations=_ObservationsCfg( + policy=_PolicyObsGroup( + o1=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-pol")}) + ), + perception=_ExtraObsGroup( + o3=ObservationTermCfg(func=_stable_func, params={"asset_cfg": StableSceneEntityCfg(name="o-per")}) + ), + ), + curriculum=_CurriculumCfg( + c1=EventTermCfg(func=_stable_func, mode="reset", params={"asset_cfg": StableSceneEntityCfg(name="c")}) + ), + ) + fe.WarpFrontend._promote_scene_entity_cfgs(cfg) + # Warp-managed groups (rewards, observations incl. sub-groups, events) are promoted. + assert isinstance(cfg.rewards.r1.params["asset_cfg"], WarpSceneEntityCfg) + assert isinstance(cfg.observations.policy.o1.params["asset_cfg"], WarpSceneEntityCfg) + # The perception sub-group is reached even though its attribute name is + # not hardcoded in the framework. + assert isinstance(cfg.observations.perception.o3.params["asset_cfg"], WarpSceneEntityCfg) + # The event manager is warp-first, so its terms are promoted too. + assert isinstance(cfg.events.e1.params["asset_cfg"], WarpSceneEntityCfg) + # Curriculum is adapted opportunistically (terms swap to warp twins when one + # exists), so its entities are promoted as well; the warp SceneEntityCfg + # subclasses the stable one, keeping legacy fallback terms valid. + assert isinstance(cfg.curriculum.c1.params["asset_cfg"], WarpSceneEntityCfg) + + +# ====================================================================== +# Twin resolution and swap-candidate heuristic +# ====================================================================== + + +def test_resolve_warp_twin_accepts_warp_origin(): + module = types.SimpleNamespace(foo=_warp_twin_func) + assert fe.WarpFrontend._resolve_warp_twin("foo", [module]) is _warp_twin_func + + +def test_resolve_warp_twin_rejects_stable_origin(): + module = types.SimpleNamespace(foo=_stable_func) # same name, stable origin + assert fe.WarpFrontend._resolve_warp_twin("foo", [module]) is None + + +def test_resolve_warp_twin_returns_none_when_absent(): + assert fe.WarpFrontend._resolve_warp_twin("missing", [types.SimpleNamespace()]) is None + + +def test_is_swap_candidate_stable_callable(): + assert fe.WarpFrontend._is_swap_candidate(_stable_func) + + +def test_is_swap_candidate_rejects_warp_callable(): + assert not fe.WarpFrontend._is_swap_candidate(_warp_twin_func) + + +def test_is_swap_candidate_rejects_non_callables(): + assert not fe.WarpFrontend._is_swap_candidate(42) + assert not fe.WarpFrontend._is_swap_candidate("string") + + +# ====================================================================== +# Direct workflow guard +# ====================================================================== + + +_DIRECT_TEST_TASKS = { + "_Frontend-Test-Warp-Direct-v0": ("isaaclab_tasks_experimental.fake:DirectEnv", None), + "_Frontend-Test-Stable-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", None), + "_Frontend-Test-Declared-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", "types:SimpleNamespace"), + "_Frontend-Test-Broken-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", "definitely_not_a_module:Env"), + "_Frontend-Test-Missing-Class-Direct-v0": ("isaaclab_tasks.fake:DirectEnv", "types:NoSuchClass"), +} + + +@pytest.fixture(scope="module", autouse=True) +def _register_direct_stub_tasks(): + # Register stub tasks so ``gym.spec`` resolves them. Entry-point strings + # are never invoked (the guard only inspects them), so a fake import + # path is fine. ``warp_entry_point`` targets are imported, so the valid + # stub points at a real stdlib symbol. + registered = [] + for task_id, (entry_point, warp_entry_point) in _DIRECT_TEST_TASKS.items(): + kwargs = {"warp_entry_point": warp_entry_point} if warp_entry_point else {} + with contextlib.suppress(gym.error.Error): + gym.register(id=task_id, entry_point=entry_point, disable_env_checker=True, kwargs=kwargs) + registered.append(task_id) + yield + # Unregister so later test modules see the real, stub-free registry. + for task_id in registered: + gym.registry.pop(task_id, None) + + +def test_direct_guard_accepts_warp_rooted(): + fe.WarpFrontend._assert_direct_warp_registration("_Frontend-Test-Warp-Direct-v0") # no raise + + +def test_direct_guard_rejects_stable_rooted(): + with pytest.raises(FrontendIncompatibleError) as exc: + fe.WarpFrontend._assert_direct_warp_registration("_Frontend-Test-Stable-Direct-v0") + assert "isaaclab_experimental" in str(exc.value) + + +def test_resolve_direct_warp_class_returns_declared_class(): + import types as types_module + + env_class = fe.WarpFrontend._resolve_direct_warp_class("_Frontend-Test-Declared-Direct-v0") + assert env_class is types_module.SimpleNamespace + + +def test_resolve_direct_warp_class_returns_none_without_declaration(): + assert fe.WarpFrontend._resolve_direct_warp_class("_Frontend-Test-Stable-Direct-v0") is None + + +def test_resolve_direct_warp_class_rejects_broken_module(): + with pytest.raises(FrontendIncompatibleError): + fe.WarpFrontend._resolve_direct_warp_class("_Frontend-Test-Broken-Direct-v0") + + +def test_resolve_direct_warp_class_rejects_missing_class(): + with pytest.raises(FrontendIncompatibleError): + fe.WarpFrontend._resolve_direct_warp_class("_Frontend-Test-Missing-Class-Direct-v0") + + +def test_resolve_direct_warp_class_rejects_unknown_task(): + with pytest.raises(FrontendIncompatibleError): + fe.WarpFrontend._resolve_direct_warp_class("Frontend-Test-NotRegistered-v0") diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py new file mode 100644 index 000000000000..e854739f78c1 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -0,0 +1,264 @@ +# 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 + +"""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``. +""" + +from __future__ import annotations + +import importlib + +import gymnasium as gym +import isaaclab_tasks_experimental # noqa: F401 +import pytest +from isaaclab_experimental.envs.frontend import WarpFrontend +from isaaclab_newton.physics import NewtonCfg + +# 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-Franka-Play", + "Isaac-Reach-UR10", + "Isaac-Reach-UR10-Play", + "Isaac-Velocity-Flat-AnymalD", + "Isaac-Velocity-Flat-AnymalD-Play", + "Isaac-Velocity-Flat-Cassie", + "Isaac-Velocity-Flat-Cassie-Play", + "Isaac-Velocity-Flat-G1", + "Isaac-Velocity-Flat-G1-Play", + "Isaac-Velocity-Flat-H1", + "Isaac-Velocity-Flat-H1-Play", + "Isaac-Velocity-Flat-UnitreeGo2", + "Isaac-Velocity-Flat-UnitreeGo2-Play", + "Isaac-Velocity-Rough-AnymalD", + "Isaac-Velocity-Rough-AnymalD-Play", + "Isaac-Velocity-Rough-Cassie", + "Isaac-Velocity-Rough-Cassie-Play", + "Isaac-Velocity-Rough-G1", + "Isaac-Velocity-Rough-G1-Play", + "Isaac-Velocity-Rough-H1", + "Isaac-Velocity-Rough-H1-Play", + "Isaac-Velocity-Rough-UnitreeGo2", + "Isaac-Velocity-Rough-UnitreeGo2-Play", +] + +# Manager-based warp tasks are exactly those whose entry point is the shared warp +# env class; direct tasks provide their own env class (via ``warp_entry_point`` +# or a warp-native registration) and are not cfg-adapted, so they are excluded here. +_MANAGER_WARP_ENTRY_POINT = "isaaclab_experimental.envs:ManagerBasedRLEnvWarp" + +_WARP_ROOTS = ("isaaclab_experimental", "isaaclab_tasks_experimental") + + +def _manager_warp_tasks() -> list[tuple[str, str]]: + """Return ``(task_id, env_cfg_entry_point)`` for every manager-based warp task.""" + tasks: list[tuple[str, str]] = [] + for task_id, spec in gym.registry.items(): + if spec.entry_point != _MANAGER_WARP_ENTRY_POINT: + continue + cfg_entry = (spec.kwargs or {}).get("env_cfg_entry_point") + if isinstance(cfg_entry, str): + tasks.append((task_id, cfg_entry)) + return sorted(tasks) + + +def _load_adapted_cfg(cfg_entry_point: str): + """Instantiate an env cfg, resolve the Newton preset, and adapt it for warp.""" + module_path, class_name = cfg_entry_point.split(":") + cfg = getattr(importlib.import_module(module_path), class_name)() + cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) + assert isinstance(cfg.sim.physics, NewtonCfg), "task does not provide a newton_mjwarp physics preset" + # Raises FrontendIncompatibleError if any warp-managed term lacks a warp twin. + WarpFrontend.adapt_cfg(cfg) + return cfg + + +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" + return cfg_entry + + +def test_no_manager_warp_variants_remain(): + """End-state pin: manager-based warp execution needs no parallel registrations.""" + assert _manager_warp_tasks() == [], "unexpected ManagerBasedRLEnvWarp registrations reappeared" + + +def test_no_warp_task_registrations_remain(): + """End-state pin: warp execution needs no parallel ``-Warp`` task ids at all.""" + warp_ids = sorted(task_id for task_id in gym.registry if "-Warp" in task_id) + 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) +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)) + + # Action terms carry a ``class_type`` (not a ``func``) and live on a base that + # is not a ManagerTermBaseCfg; guard that the adapter still swaps them to the + # warp ActionTerm, otherwise the warp ActionManager rejects them at runtime. + actions = getattr(cfg, "actions", None) + if actions is not None: + for name, term in vars(actions).items(): + class_type = getattr(term, "class_type", None) + if class_type is not None: + assert class_type.__module__.startswith(_WARP_ROOTS), ( + f"{task_id}: action term '{name}' class_type was not swapped to a warp twin" + f" (got {class_type.__module__}.{class_type.__name__})" + ) + + +def _direct_tasks_with_warp_entry_point() -> list[tuple[str, str]]: + """Return ``(task_id, warp_entry_point)`` for stable direct tasks declaring a warp env.""" + tasks: list[tuple[str, str]] = [] + for task_id, spec in gym.registry.items(): + target = (spec.kwargs or {}).get("warp_entry_point") + if isinstance(target, str): + tasks.append((task_id, target)) + return sorted(tasks) + + +_DIRECT_WARP_TASKS = _direct_tasks_with_warp_entry_point() + + +def test_direct_tasks_declare_expected_warp_entry_points(): + """Sanity: the stable direct tasks with warp implementations declare them.""" + declared = {task_id for task_id, _ in _DIRECT_WARP_TASKS} + assert { + "Isaac-Cartpole-Direct", + "Isaac-Ant-Direct", + "Isaac-Humanoid-Direct", + "Isaac-Reorient-Cube-Allegro-Direct", + } <= declared + + +@pytest.mark.parametrize( + "task_id, warp_entry_point", + [pytest.param(task_id, target, id=task_id) for task_id, target in _DIRECT_WARP_TASKS], +) +def test_declared_direct_warp_entry_point_resolves(task_id: str, warp_entry_point: str): + """Each declared ``warp_entry_point`` imports and names a warp-rooted class.""" + module_path, class_name = warp_entry_point.split(":") + env_class = getattr(importlib.import_module(module_path), class_name) + assert isinstance(env_class, type), f"{task_id}: {warp_entry_point} is not a class" + assert env_class.__module__.startswith(_WARP_ROOTS) + + +def test_stable_cartpole_cfg_adapts_to_current_warp_module_layout(): + """The stable Cartpole cfg resolves task-specific twins in the current package layout.""" + from isaaclab_experimental.managers.action_manager import ActionTerm + + cfg = _load_adapted_cfg(_cfg_entry_point("Isaac-Cartpole")) + + assert cfg.rewards.pole_pos.func.__module__.startswith("isaaclab_tasks_experimental.core.cartpole.mdp") + assert cfg.rewards.success_rate.func.__module__.startswith("isaaclab_tasks_experimental.core.cartpole.mdp") + assert issubclass(cfg.actions.joint_effort.class_type, ActionTerm) + + +def _iter_obs_terms(cfg): + """Yield (name, term cfg) for every observation term that carries a noise slot.""" + from isaaclab.managers.manager_term_cfg import ObservationGroupCfg + + for group in vars(cfg.observations).values(): + if not isinstance(group, ObservationGroupCfg): + continue + for name, term in vars(group).items(): + if hasattr(term, "noise"): + yield name, term + + +def test_stable_observation_noise_converts_to_warp_twins(): + """Stable noise cfgs (e.g. Unoise) become warp-native twins with the same parameters.""" + from isaaclab.utils import noise as stable_noise + + entry = _cfg_entry_point("Isaac-Velocity-Flat-UnitreeGo2") + module_path, class_name = entry.split(":") + raw = getattr(importlib.import_module(module_path), class_name)() + stable_params = { + name: (term.noise.n_min, term.noise.n_max) + for name, term in _iter_obs_terms(raw) + if isinstance(term.noise, stable_noise.UniformNoiseCfg) + } + assert stable_params, "expected uniform noise on the stable velocity observations" + + cfg = _load_adapted_cfg(entry) + converted = dict(_iter_obs_terms(cfg)) + for name, (n_min, n_max) in stable_params.items(): + twin = converted[name].noise + assert type(twin).__module__.startswith("isaaclab_experimental."), name + assert (twin.n_min, twin.n_max) == (n_min, n_max), name + + +def test_noise_cfg_without_warp_twin_is_a_hard_error(): + """Class-based noise models have no warp twin yet: adapting must fail loudly.""" + from isaaclab_experimental.envs.frontend import FrontendIncompatibleError + + from isaaclab.utils.noise import NoiseModelCfg, UniformNoiseCfg + + entry = _cfg_entry_point("Isaac-Velocity-Flat-UnitreeGo2") + module_path, class_name = entry.split(":") + cfg = getattr(importlib.import_module(module_path), class_name)() + cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) + name, term = next(iter(_iter_obs_terms(cfg))) + term.noise = NoiseModelCfg(noise_cfg=UniformNoiseCfg()) + with pytest.raises(FrontendIncompatibleError, match="noise"): + WarpFrontend.adapt_cfg(cfg) + + +@pytest.mark.parametrize( + "task_id", + ["Isaac-Velocity-Rough-AnymalD", "Isaac-Reach-Franka"], + ids=["rough-terrain-levels", "reach-weight-modifiers"], +) +def test_curriculum_terms_swap_to_mask_native_twins(task_id: str): + """Curriculum terms resolve to the mask-native warp classes, not legacy ID paths.""" + cfg = _load_adapted_cfg(_cfg_entry_point(task_id)) + terms = {n: t for n, t in vars(cfg.curriculum).items() if t is not None and hasattr(t, "func")} + assert terms, f"expected curriculum terms on {task_id}" + for name, term in terms.items(): + assert not isinstance(term.func, str) and term.func.__module__.startswith(_WARP_ROOTS), ( + f"curriculum term '{name}' was not swapped to a warp twin (got {term.func!r})" + ) + assert isinstance(term.func, type), ( + f"curriculum term '{name}' resolved to a legacy function instead of a mask-native class" + ) + + +def test_curriculum_term_without_twin_stays_stable(): + """Optional group: a curriculum term lacking a twin stays on the legacy fallback, no error.""" + import isaaclab.utils.math as math_utils + + entry = _cfg_entry_point("Isaac-Velocity-Rough-AnymalD") + module_path, class_name = entry.split(":") + cfg = getattr(importlib.import_module(module_path), class_name)() + cfg = resolve_presets(cfg, selected=("newton_mjwarp",)) + term = next(t for t in vars(cfg.curriculum).values() if t is not None and hasattr(t, "func")) + term.func = math_utils.quat_mul # stable-rooted symbol with no warp mirror + WarpFrontend.adapt_cfg(cfg) + assert term.func is math_utils.quat_mul diff --git a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py new file mode 100644 index 000000000000..93ca80750ec4 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py @@ -0,0 +1,219 @@ +# 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 + +"""Tests for mask-first manager-based RL reset orchestration.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import torch +import warp as wp +from isaaclab_experimental.envs.manager_based_rl_env_warp import ManagerBasedRLEnvWarp + + +def test_reset_without_host_consumers_does_not_compact_ids() -> None: + """The normal Warp reset path should not call :meth:`torch.Tensor.nonzero`.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = Mock() + env.reset_buf.any.return_value.item.return_value = True + env.reset_buf.nonzero.side_effect = AssertionError("reset IDs should not be materialized") + env._reset_mask = Mock() + env._has_recorders = False + env.has_rtx_sensors = False + env.cfg = SimpleNamespace(compute_final_obs=False, num_rerenders_on_reset=0) + env.extras = {"log": {}} + + env._reset_terminated_envs() + + env.reset_buf.nonzero.assert_not_called() + env._reset_mask.assert_called_once_with(env_mask=reset_mask, env_ids=None) + + +def test_curriculum_uses_mask_before_scene_reset() -> None: + """Curriculum updates should stay mask-native and precede scene reset.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + stages: list[tuple[str, dict]] = [] + + def call(stage, fn, /, *args, **kwargs): + del fn, args + stages.append((stage, kwargs)) + return None if stage.endswith("_compute") else {} + + env._warp_graph_cache = SimpleNamespace(call=call) + env.sim = SimpleNamespace(device="cpu") + env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") + env.scene = SimpleNamespace( + num_envs=3, + reset=Mock(side_effect=lambda **kwargs: stages.append(("Scene_reset", kwargs))), + surface_grippers={}, + ) + env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock(), requires_host_ids=False) + env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) + env.observation_manager = SimpleNamespace(reset=Mock()) + env.action_manager = SimpleNamespace(reset=Mock()) + env.reward_manager = SimpleNamespace(reset=Mock()) + env.command_manager = SimpleNamespace(reset=Mock()) + env.termination_manager = SimpleNamespace(reset=Mock()) + env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") + env.extras = {} + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + env._reset_mask(env_mask=reset_mask) + + assert [stage for stage, _ in stages[:2]] == ["CurriculumManager_compute", "Scene_reset"] + assert stages[0][1]["env_mask"] is env.reset_mask_wp + np.testing.assert_array_equal(stages[0][1]["env_mask"].numpy(), reset_mask.numpy()) + assert "env_ids" not in stages[0][1] + assert "CurriculumManager_reset" in [stage for stage, _ in stages] + + +def test_reset_stages_reuse_owner_mask_for_sparse_and_empty_selections() -> None: + """Every reset dispatch should receive the stable owner-held mask pointer.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + seen: list[tuple[wp.array, np.ndarray]] = [] + + def call(stage, fn, /, *args, **kwargs): + del fn, args + if stage == "CurriculumManager_compute": + mask = kwargs["env_mask"] + seen.append((mask, mask.numpy().copy())) + return None if stage.endswith("_compute") or stage == "Scene_reset" else {} + + env._warp_graph_cache = SimpleNamespace(call=call) + env.sim = SimpleNamespace(device="cpu") + env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") + env.scene = SimpleNamespace(num_envs=3, reset=Mock(), surface_grippers={}) + env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock(), requires_host_ids=False) + env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) + env.observation_manager = SimpleNamespace(reset=Mock()) + env.action_manager = SimpleNamespace(reset=Mock()) + env.reward_manager = SimpleNamespace(reset=Mock()) + env.command_manager = SimpleNamespace(reset=Mock()) + env.termination_manager = SimpleNamespace(reset=Mock()) + env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") + env.extras = {} + + env._reset_mask(env_mask=wp.array([False, True, False], dtype=wp.bool, device="cpu")) + env._reset_mask(env_mask=wp.zeros(3, dtype=wp.bool, device="cpu")) + + assert seen[0][0] is env.reset_mask_wp + assert seen[1][0] is env.reset_mask_wp + np.testing.assert_array_equal(seen[0][1], [False, True, False]) + np.testing.assert_array_equal(seen[1][1], [False, False, False]) + + +def test_reset_threads_one_id_materialization_to_legacy_curriculum() -> None: + """Legacy curriculum consumers should share one compact-ID materialization per reset.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + stages: dict[str, dict] = {} + + def call(stage, fn, /, *args, **kwargs): + del fn, args + stages[stage] = kwargs + return None if stage.endswith("_compute") else {} + + env._warp_graph_cache = SimpleNamespace(call=call) + env.sim = SimpleNamespace(device="cpu") + env.reset_mask_wp = wp.zeros(3, dtype=wp.bool, device="cpu") + env.scene = SimpleNamespace(num_envs=3, reset=Mock(), surface_grippers={}) + env.curriculum_manager = SimpleNamespace(compute=Mock(), reset=Mock(), requires_host_ids=True) + env.event_manager = SimpleNamespace(available_modes=[], reset=Mock()) + env.observation_manager = SimpleNamespace(reset=Mock()) + env.action_manager = SimpleNamespace(reset=Mock()) + env.reward_manager = SimpleNamespace(reset=Mock()) + env.command_manager = SimpleNamespace(reset=Mock()) + env.termination_manager = SimpleNamespace(reset=Mock()) + env._episode_length_buf_wp = wp.ones(3, dtype=wp.int64, device="cpu") + env.extras = {} + + env._reset_mask(env_mask=wp.array([False, True, True], dtype=wp.bool, device="cpu")) + + compute_ids = stages["CurriculumManager_compute"]["env_ids"] + torch.testing.assert_close(compute_ids, torch.tensor([1, 2])) + assert stages["CurriculumManager_reset"]["env_ids"] is compute_ids + + +def test_reset_compacts_ids_only_for_active_recorders() -> None: + """Recorder callbacks should receive IDs while the Warp reset keeps its mask.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = torch.tensor([False, True, False]) + env._reset_mask = Mock() + env.recorder_manager = Mock() + env.recorder_manager.reset.return_value = {"recorder": 1.0} + env._has_recorders = True + env.has_rtx_sensors = False + env.cfg = SimpleNamespace(compute_final_obs=False, num_rerenders_on_reset=0) + env.extras = {"log": {}} + + env._reset_terminated_envs() + + reset_ids = env.recorder_manager.record_pre_reset.call_args.args[0] + torch.testing.assert_close(reset_ids, torch.tensor([1])) + env._reset_mask.assert_called_once_with(env_mask=reset_mask, env_ids=reset_ids) + env.recorder_manager.reset.assert_called_once_with(reset_ids) + env.recorder_manager.record_post_reset.assert_called_once_with(reset_ids) + assert env.extras["log"]["recorder"] == 1.0 + + +def test_rtx_rerender_does_not_require_compact_reset_ids() -> None: + """RTX reset rendering should use the nonempty predicate without materializing IDs.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = torch.tensor([False, True, False]) + env._reset_mask = Mock() + env._has_recorders = False + env.has_rtx_sensors = True + env.cfg = SimpleNamespace(compute_final_obs=False, num_rerenders_on_reset=2) + env.sim = SimpleNamespace(render=Mock()) + env.extras = {"log": {}} + + env._reset_terminated_envs() + + env._reset_mask.assert_called_once_with(env_mask=reset_mask, env_ids=None) + assert env.sim.render.call_count == 2 + + +def test_empty_reset_skips_warp_pipeline_without_compacting_ids() -> None: + """An empty mask should skip eager reset stages without materializing IDs.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.zeros(3, dtype=wp.bool, device="cpu") + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = torch.zeros(3, dtype=torch.bool) + env._reset_mask = Mock() + env._has_recorders = False + env.extras = {"log": {"previous": 1.0}} + + env._reset_terminated_envs() + + env._reset_mask.assert_not_called() + assert env.extras["log"] == {"previous": 1.0} + + +def test_reset_captures_final_observation_before_reset() -> None: + """Same-step autoreset should expose the terminal observation before reset.""" + env = ManagerBasedRLEnvWarp.__new__(ManagerBasedRLEnvWarp) + reset_mask = wp.array([False, True, False], dtype=wp.bool, device="cpu") + final_obs = {"policy": torch.tensor([[1.0], [2.0], [3.0]])} + env.termination_manager = SimpleNamespace(dones_wp=reset_mask) + env.reset_buf = torch.tensor([False, True, False]) + env.observation_manager = SimpleNamespace(compute=Mock(return_value=final_obs)) + env._reset_mask = Mock() + env._has_recorders = False + env.has_rtx_sensors = False + env.cfg = SimpleNamespace(compute_final_obs=True, num_rerenders_on_reset=0) + env.extras = {"log": {}} + + env._reset_terminated_envs() + + assert env.extras["final_obs"] is final_obs + env.observation_manager.compute.assert_called_once_with() + env._reset_mask.assert_called_once_with(env_mask=reset_mask, env_ids=None) diff --git a/source/isaaclab_experimental/test/managers/test_curriculum_manager.py b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py new file mode 100644 index 000000000000..9f95f32d412d --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_curriculum_manager.py @@ -0,0 +1,462 @@ +# 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 + +"""Tests for the Warp-first curriculum manager and terrain-level update path.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import warp as wp +from isaaclab_experimental.managers import CurriculumManager, CurriculumTermCfg +from isaaclab_tasks_experimental.core.reach.mdp import modify_reward_weight +from isaaclab_tasks_experimental.core.velocity.mdp import terrain_levels_vel + +from isaaclab.managers import ManagerTermBase as StableManagerTermBase +from isaaclab.terrains import TerrainImporter +from isaaclab.utils.warp import ProxyArray + + +class _GlobalCurriculumTerm(StableManagerTermBase): + """Legacy global term that intentionally ignores compact environment IDs.""" + + def __call__(self, env, env_ids, value: float) -> float: + del env_ids # a stable term is free to ignore its selection + return value + + +class _LegacyIdCurriculumTerm(StableManagerTermBase): + """Legacy term that records the compact IDs received by compute and reset.""" + + def __init__(self, cfg, env): + super().__init__(cfg, env) + self.compute_env_ids = None + self.reset_env_ids = None + + def __call__(self, env, env_ids, scale: float) -> float: + del env + self.compute_env_ids = env_ids.clone() + return len(env_ids) * scale + + def reset(self, env_ids=None) -> None: + self.reset_env_ids = env_ids.clone() + + +def _structured_state(env, env_ids) -> dict[str, float]: + """Legacy curriculum state unsupported by the scalar Warp logging buffer.""" + del env, env_ids + return {"value": 1.0} + + +class _ArrayProxy: + """Minimal Torch/Warp array proxy.""" + + def __init__(self, values: np.ndarray, dtype): + self.warp = wp.array(values, dtype=dtype, device="cpu") + self.torch = wp.to_torch(self.warp) + + +class _Robot: + """Minimal robot exposing root positions.""" + + def __init__(self, root_positions: np.ndarray): + self.cfg = SimpleNamespace() + self.data = SimpleNamespace(root_pos_w=_ArrayProxy(root_positions, wp.vec3f)) + + +class _CommandManager: + """Minimal command manager exposing persistent Warp storage.""" + + def __init__(self, commands: np.ndarray): + self.command_wp = wp.array(commands, dtype=wp.float32, device="cpu") + + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32, ndim=2): + assert name == "base_velocity" + return self.command_wp + + +class _RewardManager: + """Minimal reward manager used by the registered Reach curricula.""" + + def __init__(self): + self._term_names = ["action_rate", "joint_vel"] + self._cfgs = { + "action_rate": SimpleNamespace(weight=-0.01), + "joint_vel": SimpleNamespace(weight=-0.0001), + } + self._weights_wp = wp.array([-0.01, -0.0001], dtype=wp.float32, device="cpu") + stride = self._weights_wp.strides[0] + self._weight_views_wp = { + name: wp.array( + ptr=self._weights_wp.ptr + term_idx * stride, + dtype=wp.float32, + shape=(1,), + strides=(stride,), + device="cpu", + ) + for term_idx, name in enumerate(self._term_names) + } + + def get_term_cfg(self, name: str): + return self._cfgs[name] + + def set_term_cfg(self, name: str, cfg) -> None: + self._cfgs[name] = cfg + + def get_term_weight_wp(self, name: str) -> wp.array(dtype=wp.float32): + return self._weight_views_wp[name] + + +class _Scene: + """Minimal scene mapping with terrain-backed environment origins.""" + + def __init__(self, robot: _Robot, terrain: TerrainImporter): + self._entities = {"robot": robot} + self.terrain = terrain + self.device = "cpu" + + @property + def env_origins(self) -> torch.Tensor: + return self.terrain.env_origins + + @property + def env_origins_pa(self) -> ProxyArray: + return self.terrain.env_origins_pa + + def __getitem__(self, name: str): + return self._entities[name] + + def keys(self): + return self._entities.keys() + + +class _Env: + """Minimal environment for curriculum manager execution.""" + + def __init__(self, terrain: TerrainImporter, root_positions: np.ndarray, commands: np.ndarray): + self.num_envs = root_positions.shape[0] + self.device = "cpu" + self.scene = _Scene(_Robot(root_positions), terrain) + self.command_manager = _CommandManager(commands) + self.rng_state_wp = wp.array(np.arange(self.num_envs, dtype=np.uint32) + 101, device=self.device) + self.max_episode_length_s = 10.0 + self.sim = SimpleNamespace(is_playing=lambda: True) + + @property + def env_origins_wp(self) -> wp.array(dtype=wp.vec3f): + """Boundary accessor mirroring :attr:`ManagerBasedEnvWarp.env_origins_wp`.""" + return self.scene.env_origins_pa.warp + + +def _init_terrain_buffers(terrain: TerrainImporter) -> None: + """Mirror the buffer initializers of :meth:`TerrainImporter.__init__`.""" + terrain._terrain_origins_pa = None + terrain._env_origins_pa = None + terrain._terrain_levels_pa = None + terrain._terrain_types_pa = None + + +def _make_terrain(levels: list[int]) -> TerrainImporter: + num_envs = len(levels) + terrain = TerrainImporter.__new__(TerrainImporter) + terrain.device = "cpu" + terrain.cfg = SimpleNamespace( + num_envs=num_envs, + max_init_terrain_level=2, + terrain_generator=SimpleNamespace(size=(8.0, 8.0)), + ) + _init_terrain_buffers(terrain) + level_values = np.asarray(levels, dtype=np.int64) + type_values = np.asarray([index % 2 for index in range(num_envs)], dtype=np.int64) + origin_values = np.zeros((3, 2, 3), dtype=np.float32) + for level in range(3): + for terrain_type in range(2): + origin_values[level, terrain_type] = [ + 100.0 * level + 10.0 * terrain_type, + float(level), + float(terrain_type), + ] + terrain.configure_env_origins(origin_values) + terrain.terrain_levels.copy_(torch.from_numpy(level_values)) + terrain.terrain_types.copy_(torch.from_numpy(type_values)) + terrain.env_origins.copy_(terrain.terrain_origins[terrain.terrain_levels, terrain.terrain_types]) + return terrain + + +def test_terrain_importer_exposes_persistent_warp_views(): + """Terrain importer should retain one zero-copy Warp view for each Torch buffer.""" + terrain = _make_terrain([0, 1, 2, 1]) + + assert isinstance(terrain.terrain_origins, torch.Tensor) + assert isinstance(terrain.env_origins, torch.Tensor) + assert isinstance(terrain.terrain_levels, torch.Tensor) + assert isinstance(terrain.terrain_types, torch.Tensor) + assert terrain.terrain_origins.data_ptr() == terrain._terrain_origins_pa.warp.ptr + assert terrain.env_origins.data_ptr() == terrain.env_origins_pa.warp.ptr + assert terrain.terrain_levels.data_ptr() == terrain.terrain_levels_pa.warp.ptr + assert terrain.terrain_types.data_ptr() == terrain._terrain_types_pa.warp.ptr + + terrain.terrain_levels[0] = 2 + assert terrain.terrain_levels_pa.warp.numpy()[0] == 2 + replacement_levels = wp.array([1, 0, 2, 1], dtype=wp.int64, device="cpu") + wp.copy(terrain.terrain_levels_pa.warp, replacement_levels) + wp.synchronize() + torch.testing.assert_close(terrain.terrain_levels, torch.tensor([1, 0, 2, 1])) + + +def test_terrain_importer_grid_origins_cache_only_environment_view(): + """Grid origins should cache a Warp view without allocating curriculum buffers.""" + terrain = TerrainImporter.__new__(TerrainImporter) + terrain.device = "cpu" + terrain.cfg = SimpleNamespace(num_envs=4, env_spacing=2.0) + _init_terrain_buffers(terrain) + terrain.configure_env_origins() + + assert terrain.terrain_origins is None + assert isinstance(terrain.env_origins, torch.Tensor) + assert terrain.env_origins.shape == (4, 3) + assert terrain.env_origins_pa.warp.shape == (4,) + assert terrain.env_origins.data_ptr() == terrain.env_origins_pa.warp.ptr + assert terrain.terrain_levels is None + assert terrain.terrain_types is None + assert terrain.terrain_origins is None + + +def test_terrain_importer_reconfigure_preserves_cached_views(): + """Reconfiguring origins must reuse buffer storage so cached views stay valid.""" + terrain = _make_terrain([0, 1, 2, 1]) + env_origins_wp = terrain.env_origins_pa.warp + levels_wp = terrain.terrain_levels_pa.warp + env_origins_ptr = terrain.env_origins.data_ptr() + shifted_origins = terrain.terrain_origins.clone() + 1.0 + + terrain.configure_env_origins(shifted_origins) + + assert terrain.env_origins_pa.warp is env_origins_wp + assert terrain.terrain_levels_pa.warp is levels_wp + assert terrain.env_origins.data_ptr() == env_origins_ptr + torch.testing.assert_close(terrain.terrain_origins, shifted_origins) + + +def test_terrain_importer_mask_update_preserves_sparse_and_stable_semantics(): + """Mask updates should handle up/down/clamp/wrap while stable IDs remain supported.""" + terrain = _make_terrain([0, 1, 2, 2, 1, 0]) + env_mask = wp.array([True, True, True, True, False, False], dtype=wp.bool, device="cpu") + move_up = wp.array([False, True, True, False, True, False], dtype=wp.bool, device="cpu") + move_down = wp.array([True, False, False, True, False, True], dtype=wp.bool, device="cpu") + rng_state = wp.array(np.arange(6, dtype=np.uint32) + 71, device="cpu") + rng_before = rng_state.numpy().copy() + origins_before = terrain.env_origins.clone() + levels_wp = terrain.terrain_levels_pa.warp + origins_wp = terrain.env_origins_pa.warp + levels_ptr = levels_wp.ptr + origins_ptr = origins_wp.ptr + assert terrain.terrain_levels.data_ptr() == levels_ptr + assert terrain.env_origins.data_ptr() == origins_ptr + + terrain.update_env_origins_mask(env_mask, move_up, move_down, rng_state) + wp.synchronize() + + levels = terrain.terrain_levels.tolist() + assert levels[0] == 0 + assert levels[1] == 2 + assert 0 <= levels[2] < terrain.max_terrain_level + assert levels[3] == 1 + assert levels[4:] == [1, 0] + torch.testing.assert_close( + terrain.env_origins[:4], + terrain.terrain_origins[terrain.terrain_levels[:4], terrain.terrain_types[:4]], + ) + torch.testing.assert_close(terrain.env_origins[4:], origins_before[4:]) + assert np.all(rng_state.numpy()[:4] != rng_before[:4]) + np.testing.assert_array_equal(rng_state.numpy()[4:], rng_before[4:]) + assert terrain.terrain_levels_pa.warp is levels_wp + assert terrain.env_origins_pa.warp is origins_wp + assert terrain.terrain_levels_pa.warp.ptr == levels_ptr + assert terrain.env_origins_pa.warp.ptr == origins_ptr + + levels_before_stable_update = terrain.terrain_levels.clone() + env_ids = torch.tensor([1, 3], dtype=torch.int64) + terrain.update_env_origins( + env_ids, + move_up=torch.tensor([False, False]), + move_down=torch.tensor([True, True]), + ) + expected_levels = levels_before_stable_update.clone() + expected_levels[env_ids] -= 1 + torch.testing.assert_close(terrain.terrain_levels, expected_levels) + assert terrain.terrain_levels_pa.warp is levels_wp + assert terrain.env_origins_pa.warp is origins_wp + assert terrain.terrain_levels_pa.warp.ptr == levels_ptr + assert terrain.env_origins_pa.warp.ptr == origins_ptr + + +def test_curriculum_manager_updates_masked_levels_and_logs_without_host_compaction(monkeypatch: pytest.MonkeyPatch): + """Manager compute/reset should remain mask-native and expose persistent scalar logging.""" + terrain = _make_terrain([0, 1, 2, 2, 1, 0]) + root_positions = terrain.env_origins.numpy().copy() + root_positions[:, 0] += np.array([5.0, 1.0, 3.0, 5.0, 9.0, 0.0], dtype=np.float32) + commands = np.array( + [[0.1, 0.0, 0.0], [1.0, 0.0, 0.0], [0.2, 0.0, 0.0], [2.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.1, 0.0, 0.0]], + dtype=np.float32, + ) + env = _Env(terrain, root_positions, commands) + manager = CurriculumManager( + { + "terrain_levels": CurriculumTermCfg(func=terrain_levels_vel), + }, + env, + ) + term = manager._term_cfgs[0].func + env_mask = wp.array([True, True, True, True, False, False], dtype=wp.bool, device="cpu") + move_up_ptr = term._move_up_wp.ptr + move_down_ptr = term._move_down_wp.ptr + state_ptr = manager._term_states_wp.ptr + extras_ref = manager.reset_extras + assert not manager.requires_host_ids + # a purely mask-native manager needs no host boundary at all + assert not manager.requires_host_boundary + + def _fail_host_compaction(*args, **kwargs): + raise AssertionError("Torch host compaction/scalar extraction reached the Warp curriculum path") + + with monkeypatch.context() as context: + context.setattr(torch.Tensor, "nonzero", _fail_host_compaction) + context.setattr(torch.Tensor, "item", _fail_host_compaction) + manager.compute(env_mask) + extras = manager.reset(env_mask) + wp.synchronize() + + levels = terrain.terrain_levels.tolist() + assert levels[0] == 1 + assert levels[1] == 0 + assert levels[2] == 2 + assert 0 <= levels[3] < terrain.max_terrain_level + assert levels[4:] == [1, 0] + np.testing.assert_array_equal(term._move_up_wp.numpy(), [True, False, False, True, False, False]) + np.testing.assert_array_equal(term._move_down_wp.numpy(), [False, True, False, False, False, False]) + assert extras is extras_ref + assert extras["Curriculum/terrain_levels"] is manager.reset_extras["Curriculum/terrain_levels"] + torch.testing.assert_close(extras["Curriculum/terrain_levels"], terrain.terrain_levels.float().mean()) + assert term._move_up_wp.ptr == move_up_ptr + assert term._move_down_wp.ptr == move_down_ptr + assert manager._term_states_wp.ptr == state_ptr + + +def test_curriculum_manager_compacts_ids_inside_legacy_term_boundary(): + """Legacy terms should receive compact IDs without polluting the environment reset API.""" + env = SimpleNamespace( + num_envs=4, + device="cpu", + sim=SimpleNamespace(is_playing=lambda: True), + ) + manager = CurriculumManager( + { + "id_count": CurriculumTermCfg(func=_LegacyIdCurriculumTerm, params={"scale": 2.0}), + "global_state": CurriculumTermCfg(func=_GlobalCurriculumTerm, params={"value": 7.0}), + }, + env, + ) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + term = manager._term_cfgs[0].func + + assert manager.requires_host_ids + assert manager.requires_host_boundary + manager.compute(env_mask) + extras = manager.reset(env_mask) + + torch.testing.assert_close(term.compute_env_ids, torch.tensor([0, 2])) + torch.testing.assert_close(term.reset_env_ids, torch.tensor([0, 2])) + torch.testing.assert_close(extras["Curriculum/id_count"], torch.tensor(4.0)) + torch.testing.assert_close(extras["Curriculum/global_state"], torch.tensor(7.0)) + + +def test_curriculum_manager_consumes_threaded_ids_without_compaction(monkeypatch: pytest.MonkeyPatch): + """Precomputed compact IDs should reach legacy terms without another host materialization.""" + env = SimpleNamespace( + num_envs=4, + device="cpu", + sim=SimpleNamespace(is_playing=lambda: True), + ) + manager = CurriculumManager( + {"id_count": CurriculumTermCfg(func=_LegacyIdCurriculumTerm, params={"scale": 2.0})}, + env, + ) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + precomputed_ids = torch.tensor([0, 2]) + term = manager._term_cfgs[0].func + + def _fail_host_compaction(*args, **kwargs): + raise AssertionError("threaded IDs must not be re-materialized") + + with monkeypatch.context() as context: + context.setattr(torch.Tensor, "nonzero", _fail_host_compaction) + manager.compute(env_mask, env_ids=precomputed_ids) + extras = manager.reset(env_mask, env_ids=precomputed_ids) + + torch.testing.assert_close(term.compute_env_ids, precomputed_ids) + torch.testing.assert_close(term.reset_env_ids, precomputed_ids) + torch.testing.assert_close(extras["Curriculum/id_count"], torch.tensor(4.0)) + + +def test_registered_reach_curricula_update_weights_without_a_host_boundary(): + """Registered Reach reward schedules should update only on a device-selected reset.""" + terrain = _make_terrain([0, 1, 2, 2]) + env = _Env( + terrain, + terrain.env_origins.numpy().copy(), + np.zeros((4, 3), dtype=np.float32), + ) + env.reward_manager = _RewardManager() + env.common_step_counter = 0 + manager = CurriculumManager( + { + "action_rate": CurriculumTermCfg( + func=modify_reward_weight, + params={"term_name": "action_rate", "weight": -0.005, "num_steps": 4500}, + ), + "joint_vel": CurriculumTermCfg( + func=modify_reward_weight, + params={"term_name": "joint_vel", "weight": -0.001, "num_steps": 4500}, + ), + }, + env, + ) + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + assert not manager.requires_host_ids + assert not manager.requires_host_boundary + manager.compute(env_mask) + extras = manager.reset(env_mask) + torch.testing.assert_close(extras["Curriculum/action_rate"], torch.tensor(-0.01)) + torch.testing.assert_close(extras["Curriculum/joint_vel"], torch.tensor(-0.0001)) + + env.common_step_counter = 4501 + empty_mask = wp.zeros(4, dtype=wp.bool, device="cpu") + manager.compute(empty_mask) + np.testing.assert_allclose(env.reward_manager._weights_wp.numpy(), [-0.01, -0.0001]) + + manager.compute(env_mask) + extras = manager.reset(env_mask) + torch.testing.assert_close(extras["Curriculum/action_rate"], torch.tensor(-0.005)) + torch.testing.assert_close(extras["Curriculum/joint_vel"], torch.tensor(-0.001)) + np.testing.assert_allclose(env.reward_manager._weights_wp.numpy(), [-0.005, -0.001]) + + +def test_curriculum_manager_rejects_structured_legacy_state() -> None: + """Unsupported fallback logging should fail explicitly instead of narrowing silently.""" + terrain = _make_terrain([0, 1]) + env = _Env(terrain, terrain.env_origins.numpy().copy(), np.zeros((2, 3), dtype=np.float32)) + manager = CurriculumManager( + {"structured": CurriculumTermCfg(func=_structured_state)}, + env, + ) + env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") + + with pytest.raises(TypeError, match="scalar logging states"): + manager.compute(env_mask) diff --git a/source/isaaclab_experimental/test/managers/test_manager_base.py b/source/isaaclab_experimental/test/managers/test_manager_base.py new file mode 100644 index 000000000000..bef4ad902216 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_manager_base.py @@ -0,0 +1,42 @@ +# 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 + +"""Tests for Warp manager reset-mask contracts.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import warp as wp +from isaaclab_experimental.managers.action_manager import ActionManager +from isaaclab_experimental.managers.manager_base import _resolve_reset_mask +from isaaclab_experimental.utils.warp import WarpCapturable + + +def test_reset_mask_rejects_wrong_device(monkeypatch: pytest.MonkeyPatch) -> None: + """Manager kernels should fail before receiving a mask from another device.""" + owner = SimpleNamespace(num_envs=2, device="expected") + env_mask = wp.zeros(2, dtype=wp.bool, device="cpu") + expected_device = SimpleNamespace(alias="expected") + monkeypatch.setattr(wp, "get_device", lambda device: expected_device) + + with pytest.raises(ValueError, match="env_mask must be on"): + _resolve_reset_mask(owner, None, env_mask) + + +def test_class_term_capturability_is_registered() -> None: + """Class-based manager terms should honor explicit capture metadata.""" + + @WarpCapturable(False) + class HostTerm: + pass + + graph_cache = Mock() + manager = object.__new__(ActionManager) + manager._env = SimpleNamespace(_warp_graph_cache=graph_cache) + + manager._register_term_capturability(HostTerm) + + graph_cache.register_capturability.assert_called_once_with("ActionManager", False) diff --git a/source/isaaclab_experimental/test/managers/test_reward_manager.py b/source/isaaclab_experimental/test/managers/test_reward_manager.py new file mode 100644 index 000000000000..3d056470540f --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py @@ -0,0 +1,42 @@ +# 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 + +"""Tests for device-owned Warp reward weights.""" + +from types import SimpleNamespace + +import torch +import warp as wp +from isaaclab_experimental.managers import RewardManager, RewardTermCfg + + +@wp.kernel +def _fill_unit_reward(out: wp.array(dtype=wp.float32)): + """Test helper: write 1.0 into every env's term output.""" + out[wp.tid()] = 1.0 + + +def _unit_reward(env, out: wp.array(dtype=wp.float32)) -> None: + env.term_call_count += 1 + wp.launch(_fill_unit_reward, dim=env.num_envs, inputs=[out], device=env.device) + + +def test_device_owned_weight_can_enable_zero_weight_term() -> None: + """A curriculum-owned weight should be authoritative over the initial config scalar.""" + env = SimpleNamespace( + num_envs=4, + device="cpu", + sim=SimpleNamespace(is_playing=lambda: False), + term_call_count=0, + ) + manager = RewardManager({"scheduled": RewardTermCfg(func=_unit_reward, weight=0.0)}, env) + + torch.testing.assert_close(manager.compute(dt=1.0), torch.zeros(4)) + assert env.term_call_count == 0 + + manager.get_term_weight_wp("scheduled").fill_(2.0) + torch.testing.assert_close(manager.compute(dt=1.0), torch.full((4,), 2.0)) + assert env.term_call_count == 1 + assert manager.get_term_cfg("scheduled").weight == 2.0 diff --git a/source/isaaclab_experimental/test/managers/test_termination_manager.py b/source/isaaclab_experimental/test/managers/test_termination_manager.py new file mode 100644 index 000000000000..07605af1d853 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_termination_manager.py @@ -0,0 +1,37 @@ +# 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 + +"""Tests for the Warp-first termination manager.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import warp as wp +from isaaclab_experimental.managers.termination_manager import TerminationManager + + +class TestTerminationManager: + """Tests for :class:`TerminationManager`.""" + + def test_reset_statistics_only_include_selected_environments(self): + """Partial reset statistics should be averaged over selected environments only.""" + manager = TerminationManager.__new__(TerminationManager) + manager._env = SimpleNamespace(num_envs=4, device="cpu") + manager._term_names = ["fall", "timeout"] + manager._last_episode_dones_wp = wp.array( + [[True, False], [False, True], [True, True], [False, False]], dtype=wp.bool, device="cpu" + ) + manager._term_done_avg_wp = wp.zeros(2, dtype=wp.float32, device="cpu") + manager._reset_count_wp = wp.zeros(1, dtype=wp.int32, device="cpu") + manager._reset_scale_wp = wp.zeros(1, dtype=wp.float32, device="cpu") + manager._class_term_cfgs = [] + manager._reset_extras = {} + env_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + manager.reset(env_mask=env_mask) + + np.testing.assert_allclose(manager._term_done_avg_wp.numpy(), [1.0, 0.5]) diff --git a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py b/source/isaaclab_experimental/test/utils/test_manager_call_switch.py deleted file mode 100644 index b4745939eaa0..000000000000 --- a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py +++ /dev/null @@ -1,475 +0,0 @@ -# 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 - -"""Tests for ManagerCallSwitch config loading, mode resolution, and dispatch.""" - -from __future__ import annotations - -import json -import os -import unittest - -import warp as wp -from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch - - -@wp.kernel -def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)): - i = wp.tid() - b[i] = a[i] + 1.0 - - -class TestManagerCallMode(unittest.TestCase): - """Tests for the ManagerCallMode enum.""" - - def test_enum_values(self): - self.assertEqual(ManagerCallMode.STABLE, 0) - self.assertEqual(ManagerCallMode.WARP_NOT_CAPTURED, 1) - self.assertEqual(ManagerCallMode.WARP_CAPTURED, 2) - - def test_ordering(self): - self.assertLess(ManagerCallMode.STABLE, ManagerCallMode.WARP_NOT_CAPTURED) - self.assertLess(ManagerCallMode.WARP_NOT_CAPTURED, ManagerCallMode.WARP_CAPTURED) - - -# ====================================================================== -# Config loading -# ====================================================================== - - -class TestConfigLoading(unittest.TestCase): - """Tests for ManagerCallSwitch config parsing from dict, JSON, env var, and None.""" - - def test_none_uses_default(self): - switch = ManagerCallSwitch(cfg_source=None) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED) - - def test_dict_config(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) - - def test_dict_per_manager_override(self): - switch = ManagerCallSwitch(cfg_source={"default": 2, "RewardManager": 0}) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) - self.assertEqual(switch.get_mode_for_manager("ActionManager"), ManagerCallMode.WARP_CAPTURED) - - def test_dict_without_default_key(self): - """A dict missing 'default' should inherit from DEFAULT_CONFIG.""" - switch = ManagerCallSwitch(cfg_source={"RewardManager": 0}) - # default should be 2 (from DEFAULT_CONFIG) - self.assertEqual(switch.get_mode_for_manager("ActionManager"), ManagerCallMode.WARP_CAPTURED) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) - - def test_json_string_config(self): - cfg_str = json.dumps({"default": 1, "ObservationManager": 0}) - switch = ManagerCallSwitch(cfg_source=cfg_str) - self.assertEqual(switch.get_mode_for_manager("ObservationManager"), ManagerCallMode.STABLE) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) - - def test_empty_string_uses_default(self): - switch = ManagerCallSwitch(cfg_source="") - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED) - - def test_env_var_fallback(self): - """When cfg_source is None, should read from MANAGER_CALL_CONFIG env var.""" - old = os.environ.get(ManagerCallSwitch.ENV_VAR) - try: - os.environ[ManagerCallSwitch.ENV_VAR] = json.dumps({"default": 0}) - switch = ManagerCallSwitch(cfg_source=None) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) - finally: - if old is None: - os.environ.pop(ManagerCallSwitch.ENV_VAR, None) - else: - os.environ[ManagerCallSwitch.ENV_VAR] = old - - def test_invalid_json_raises(self): - with self.assertRaises(json.JSONDecodeError): - ManagerCallSwitch(cfg_source="not valid json") - - def test_invalid_mode_value_raises(self): - with self.assertRaises(ValueError): - ManagerCallSwitch(cfg_source={"default": 99}) - - def test_non_int_mode_raises(self): - with self.assertRaises(TypeError): - ManagerCallSwitch(cfg_source={"default": "fast"}) - - def test_invalid_cfg_type_raises(self): - with self.assertRaises(TypeError): - ManagerCallSwitch(cfg_source=42) - - -# ====================================================================== -# Mode resolution and capping -# ====================================================================== - - -class TestModeResolution(unittest.TestCase): - """Tests for mode resolution, MAX_MODE_OVERRIDES, and dynamic capturability.""" - - def test_max_mode_override_caps_default(self): - """Scene is capped to WARP_NOT_CAPTURED by MAX_MODE_OVERRIDES.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - # Scene has a static cap of WARP_NOT_CAPTURED - self.assertEqual( - switch.get_mode_for_manager("Scene"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - # Other managers should not be capped - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_CAPTURED, - ) - - def test_max_modes_kwarg(self): - """max_modes kwarg should cap managers beyond MAX_MODE_OVERRIDES.""" - switch = ManagerCallSwitch( - cfg_source={"default": 2}, - max_modes={"RewardManager": ManagerCallMode.WARP_NOT_CAPTURED}, - ) - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - def test_register_manager_capturability_downgrades(self): - """register_manager_capturability(False) should cap a manager to WARP_NOT_CAPTURED.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED) - - switch.register_manager_capturability("RewardManager", capturable=False) - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - def test_register_capturability_true_is_noop(self): - """register_manager_capturability(True) should not change anything.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - switch.register_manager_capturability("RewardManager", capturable=True) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_CAPTURED) - - def test_register_capturability_does_not_upgrade(self): - """If a manager is already capped to NOT_CAPTURED, registering capturable=False again shouldn't change it.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - switch.register_manager_capturability("RewardManager", capturable=False) - switch.register_manager_capturability("RewardManager", capturable=False) - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - def test_capturability_interacts_with_static_cap(self): - """Dynamic capturability should respect existing static caps.""" - switch = ManagerCallSwitch( - cfg_source={"default": 2}, - max_modes={"RewardManager": ManagerCallMode.WARP_NOT_CAPTURED}, - ) - # Already capped, register_capturability(False) should be harmless - switch.register_manager_capturability("RewardManager", capturable=False) - self.assertEqual( - switch.get_mode_for_manager("RewardManager"), - ManagerCallMode.WARP_NOT_CAPTURED, - ) - - def test_mode_0_not_affected_by_cap(self): - """A manager explicitly set to STABLE (0) should stay at 0 even with cap at 1.""" - switch = ManagerCallSwitch(cfg_source={"default": 2, "RewardManager": 0}) - switch.register_manager_capturability("RewardManager", capturable=False) - self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.STABLE) - - -# ====================================================================== -# Stage dispatch -# ====================================================================== - - -class TestStageDispatch(unittest.TestCase): - """Tests for call_stage routing through stable / warp-eager / warp-captured paths.""" - - def test_stable_mode_calls_stable_fn(self): - switch = ManagerCallSwitch(cfg_source={"default": 0}) - called = {"stable": False, "warp": False} - - def stable_fn(): - called["stable"] = True - return "stable_result" - - def warp_fn(): - called["warp"] = True - return "warp_result" - - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn}, - stable_call={"fn": stable_fn}, - ) - self.assertTrue(called["stable"]) - self.assertFalse(called["warp"]) - self.assertEqual(result, "stable_result") - - def test_stable_mode_without_stable_call_raises(self): - switch = ManagerCallSwitch(cfg_source={"default": 0}) - with self.assertRaises(ValueError): - switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": lambda: None}, - stable_call=None, - ) - - def test_warp_eager_mode_calls_warp_fn(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - called = {"stable": False, "warp": False} - - def stable_fn(): - called["stable"] = True - return "stable_result" - - def warp_fn(): - called["warp"] = True - return "warp_result" - - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn}, - stable_call={"fn": stable_fn}, - ) - self.assertFalse(called["stable"]) - self.assertTrue(called["warp"]) - self.assertEqual(result, "warp_result") - - def test_output_transform(self): - """The 'output' key in call spec should transform the return value.""" - switch = ManagerCallSwitch(cfg_source={"default": 1}) - - def warp_fn(): - return [1, 2, 3] - - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn, "output": len}, - ) - self.assertEqual(result, 3) - - def test_args_and_kwargs_forwarded(self): - """call_stage should forward args and kwargs from the call spec.""" - switch = ManagerCallSwitch(cfg_source={"default": 1}) - received = {} - - def warp_fn(a, b, key=None): - received["a"] = a - received["b"] = b - received["key"] = key - - switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn, "args": (1, 2), "kwargs": {"key": "val"}}, - ) - self.assertEqual(received, {"a": 1, "b": 2, "key": "val"}) - - -class TestStageDispatchCaptured(unittest.TestCase): - """Tests for WARP_CAPTURED mode (requires GPU).""" - - def setUp(self): - self.device = "cuda:0" - - def test_captured_mode_produces_correct_output(self): - """WARP_CAPTURED should capture and replay a warp kernel correctly.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.full(4, value=5.0, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - def warp_fn(): - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - # First call: warm-up + capture - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn}, - ) - self.assertAlmostEqual(result.numpy()[0], 6.0, places=5) - - # Replay - wp.copy(src, wp.full(4, value=10.0, dtype=wp.float32, device=self.device)) - result2 = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": warp_fn}, - ) - self.assertIs(result, result2, "Replay must return same reference") - self.assertAlmostEqual(result2.numpy()[0], 11.0, places=5) - - def test_captured_warmup_call_count(self): - """WARP_CAPTURED first call should invoke fn exactly 2 times (warm-up + capture).""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.zeros(4, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - call_count = [0] - - def warp_fn(): - call_count[0] += 1 - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - # First call_stage: warm-up (1) + capture (2) - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2, "First captured call should invoke fn twice (warm-up + capture)") - - # Second call_stage: replay only — no new fn invocation - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2, "Replay should not invoke fn again") - - # Third call_stage: still replay - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2) - - def test_captured_warmup_handles_hasattr_guard(self): - """Warm-up should flush first-call allocations so capture doesn't record them. - - This simulates the real-world pattern where MDP terms allocate scratch - buffers on first call using hasattr guards. - """ - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.ones(8, dtype=wp.float32, device=self.device) - holder = {} - - def fn_with_guard(): - if "buf" not in holder: - # First-call allocation — must happen during warm-up, not capture - holder["buf"] = wp.zeros(8, dtype=wp.float32, device=self.device) - wp.launch(_add_one, dim=8, inputs=[src, holder["buf"]], device=self.device) - return holder["buf"] - - # Should not raise — warm-up handles the allocation outside capture context - result = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": fn_with_guard}, - ) - result_np = result.numpy() - for val in result_np: - self.assertAlmostEqual(val, 2.0, places=5) - - # Replay should also work (allocation already done, only kernel replays) - wp.copy(src, wp.full(8, value=5.0, dtype=wp.float32, device=self.device)) - result2 = switch.call_stage( - stage="RewardManager_compute", - warp_call={"fn": fn_with_guard}, - ) - for val in result2.numpy(): - self.assertAlmostEqual(val, 6.0, places=5) - - def test_warp_eager_no_warmup(self): - """WARP_NOT_CAPTURED mode should call fn exactly once per call (no warm-up).""" - switch = ManagerCallSwitch(cfg_source={"default": 1}) - call_count = [0] - - def warp_fn(): - call_count[0] += 1 - - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 1, "Eager mode should call fn exactly once") - - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2, "Eager mode should call fn again each time") - - -# ====================================================================== -# Graph invalidation -# ====================================================================== - - -class TestGraphInvalidation(unittest.TestCase): - """Tests for invalidate_graphs on ManagerCallSwitch.""" - - def setUp(self): - self.device = "cuda:0" - - def test_invalidate_forces_recapture(self): - switch = ManagerCallSwitch(cfg_source={"default": 2}) - src = wp.full(4, value=1.0, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - call_count = [0] - - def warp_fn(): - call_count[0] += 1 - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - # First capture - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2) # warm-up + capture - - # Replay — no new calls - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2) - - # Invalidate - switch.invalidate_graphs() - - # Should re-warm-up + re-capture - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 4) - - -# ====================================================================== -# resolve_manager_class -# ====================================================================== - - -class TestResolveManagerClass(unittest.TestCase): - """Tests for resolve_manager_class.""" - - def test_stable_resolves_to_isaaclab_managers(self): - switch = ManagerCallSwitch(cfg_source={"default": 0}) - cls = switch.resolve_manager_class("RewardManager") - self.assertTrue(cls.__module__.startswith("isaaclab.managers")) - - def test_warp_resolves_to_experimental_managers(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - cls = switch.resolve_manager_class("RewardManager") - self.assertTrue(cls.__module__.startswith("isaaclab_experimental")) - - def test_mode_override(self): - """mode_override should bypass the config for class resolution.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - cls = switch.resolve_manager_class("RewardManager", mode_override=ManagerCallMode.STABLE) - self.assertTrue(cls.__module__.startswith("isaaclab.managers")) - - def test_invalid_manager_raises(self): - switch = ManagerCallSwitch(cfg_source={"default": 0}) - with self.assertRaises(AttributeError): - switch.resolve_manager_class("NonexistentManager") - - -# ====================================================================== -# Manager name parsing -# ====================================================================== - - -class TestManagerNameParsing(unittest.TestCase): - """Tests for stage name → manager name extraction.""" - - def test_valid_stage_name(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - # Dispatch a stage with valid name format - called = [False] - - def fn(): - called[0] = True - - switch.call_stage(stage="RewardManager_compute", warp_call={"fn": fn}) - self.assertTrue(called[0]) - - def test_invalid_stage_name_raises(self): - switch = ManagerCallSwitch(cfg_source={"default": 1}) - with self.assertRaises(ValueError): - switch.call_stage(stage="nounderscore", warp_call={"fn": lambda: None}) - - -if __name__ == "__main__": - unittest.main() diff --git a/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py new file mode 100644 index 000000000000..2bdc680d23a8 --- /dev/null +++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py @@ -0,0 +1,217 @@ +# 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 + +"""Static gate for the mask-native execution path. + +The warp frontend's contract is that mask-native code performs no data-dependent +GPU->host synchronization outside sanctioned boundaries. Instead of inline marker +comments, the gate is structural: + +* **Scope** — every function in the experimental production tree that sits on the + step/reset pipeline, plus every ``*mask*``-named function in the core and Newton + trees (the functions that *are* the warp path there). Torch-first core code is + deliberately out of scope. +* **Sanctioned boundaries** — small, named helper functions listed in + ``SANCTIONED_BOUNDARIES``; the list is exact (a stale entry fails the gate). +* **Runtime backstop** — ``ISAACLAB_SYNC_DEBUG=1`` runs every warp stage under + ``torch.cuda.set_sync_debug_mode("error")`` (see ``WarpGraphCache``), trapping + syncs this static scan cannot see. + +The companion inventory test pins every ``@WarpCapturable(False)`` opt-out so +non-capturable terms are documented in exactly one reviewable place. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[4] + +# Attribute calls that synchronize (or copy) device memory to the host. +SYNC_ATTRS = {"nonzero", "argwhere", "item", "cpu", "tolist", "numpy"} + +# Trees scanned in full (production only; every function is in scope unless excluded). +FULL_SCAN_ROOTS = [ + "source/isaaclab_experimental/isaaclab_experimental/envs", + "source/isaaclab_experimental/isaaclab_experimental/managers", + "source/isaaclab_experimental/isaaclab_experimental/utils", +] + +# Trees where only ``*mask*``-named functions are in scope (the warp path inside +# otherwise torch-first packages). +MASK_SCAN_ROOTS = [ + "source/isaaclab/isaaclab/actuators", + "source/isaaclab/isaaclab/scene", + "source/isaaclab/isaaclab/sensors", + "source/isaaclab/isaaclab/terrains", + "source/isaaclab_newton/isaaclab_newton", +] + +# Functions that never run on the per-step path: construction, reporting, and +# debug surfaces. Host syncs there are init-time or human-facing. +NON_STEP_FUNCTIONS = { + "__init__", + "__post_init__", + "__str__", + "__repr__", + "__del__", + "_prepare_terms", + "serialize", + "get_active_iterable_terms", + "_debug_vis_callback", + "_set_debug_vis_impl", + # reporting / config-access surfaces (export- or user-facing, not per-step) + "IO_descriptor", + "get_IO_descriptors", + "get_term_cfg", +} + +# The sanctioned mask->ID / host-sync boundaries, by (path suffix, function name). +# Every entry must exist and contain a sync call; anything else that syncs fails. +SANCTIONED_BOUNDARIES = { + ("managers/curriculum_manager.py", "_compact_legacy_env_ids"), + ("envs/mdp/events.py", "_mask_to_env_ids"), + ("envs/manager_based_env_warp.py", "reset_to"), + # Coarse-grained by review preference: the recorder / legacy-term compactions + # stay inline in their reset functions rather than one-line helpers; the + # ISAACLAB_SYNC_DEBUG runtime trap is the fine-grained net inside them. + ("envs/manager_based_env_warp.py", "reset"), + ("envs/manager_based_rl_env_warp.py", "_reset_terminated_envs"), + ("envs/manager_based_rl_env_warp.py", "_reset_mask"), + ("utils/warp/utils.py", "any_env_set"), + # camera's empty-reset predicate (the sensor-side analogue of any_env_set) + ("isaaclab/sensors/camera/camera.py", "_env_mask_has_any"), + # joint-limit writes mutate the solver model (host-side by nature, event-driven) + ("isaaclab_newton/assets/articulation/articulation.py", "write_joint_position_limit_to_sim_mask"), + # develop-era compaction helpers; deleted by Part 3 (#6646) — its rebase must + # drop these entries (the stale-entry assertion enforces that). + ("isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py", "_resolve_env_mask"), + ("isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py", "_resolve_body_mask"), +} + + +def _sync_calls(func_node: ast.AST) -> list[str]: + """Names of synchronizing attribute calls inside a function body.""" + hits = [] + for node in ast.walk(func_node): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr in SYNC_ATTRS: + hits.append(node.func.attr) + return hits + + +def _iter_functions(tree: ast.Module): + """Yield (function node, name) for every def in the module.""" + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + yield node, node.name + + +def _scan(root: Path, mask_only: bool) -> dict[tuple[str, str], list[str]]: + """Map (relative path, function) -> sync hits for in-scope functions under root.""" + findings: dict[tuple[str, str], list[str]] = {} + for path in sorted(root.rglob("*.py")): + rel = path.relative_to(_REPO_ROOT).as_posix() + if "/test/" in rel: + continue + tree = ast.parse(path.read_text(), filename=rel) + for func, name in _iter_functions(tree): + if name in NON_STEP_FUNCTIONS: + continue + if mask_only and "mask" not in name: + continue + hits = _sync_calls(func) + if hits: + findings[(rel, name)] = hits + return findings + + +def _is_sanctioned(rel: str, name: str) -> bool: + return any(rel.endswith(suffix) and name == func for suffix, func in SANCTIONED_BOUNDARIES) + + +def test_mask_native_code_has_no_unsanctioned_host_syncs(): + """Every sync in mask-native scope is one of the named, sanctioned boundaries.""" + findings: dict[tuple[str, str], list[str]] = {} + for root in FULL_SCAN_ROOTS: + findings.update(_scan(_REPO_ROOT / root, mask_only=False)) + for root in MASK_SCAN_ROOTS: + findings.update(_scan(_REPO_ROOT / root, mask_only=True)) + + violations = {k: v for k, v in findings.items() if not _is_sanctioned(*k)} + assert not violations, "Unsanctioned host syncs on the mask-native path:\n" + "\n".join( + f" {rel}::{name} -> {hits}" for (rel, name), hits in sorted(violations.items()) + ) + + # The sanctioned list is exact: every entry must still exist and still sync. + matched = {k for k in findings if _is_sanctioned(*k)} + stale = { + (suffix, func) + for suffix, func in SANCTIONED_BOUNDARIES + if not any(rel.endswith(suffix) and name == func for rel, name in matched) + } + assert not stale, f"Stale sanctioned-boundary entries (function gone or no longer syncs): {sorted(stale)}" + + +# Expected capturability opt-outs — ``@WarpCapturable(False)`` decorations or +# ``_warp_capturable = False`` class attributes: (path suffix, name). +EXPECTED_NON_CAPTURABLE = { + ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_com"), + ("isaaclab_experimental/envs/mdp/observations.py", "height_scan"), + # Reads the Python-owned common step counter; captured replay would bake it. + ("isaaclab_tasks_experimental/core/reach/mdp/curriculums.py", "modify_reward_weight"), +} + +NON_CAPTURABLE_SCAN_ROOTS = [ + "source/isaaclab_experimental/isaaclab_experimental", + "source/isaaclab_tasks_experimental/isaaclab_tasks_experimental", +] + + +def _warp_capturable_false_targets(tree: ast.Module) -> list[str]: + """Names annotated non-capturable: ``@WarpCapturable(False, ...)`` decorations + or ``_warp_capturable = False`` class attributes.""" + out = [] + for node in ast.walk(tree): + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue + for deco in node.decorator_list: + if ( + isinstance(deco, ast.Call) + and isinstance(deco.func, ast.Name) + and deco.func.id == "WarpCapturable" + and deco.args + and isinstance(deco.args[0], ast.Constant) + and deco.args[0].value is False + ): + out.append(node.name) + if isinstance(node, ast.ClassDef): + for stmt in node.body: + if ( + isinstance(stmt, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "_warp_capturable" for t in stmt.targets) + and isinstance(stmt.value, ast.Constant) + and stmt.value.value is False + ): + out.append(node.name) + return out + + +def test_non_capturable_terms_are_inventoried(): + """Every ``@WarpCapturable(False)`` opt-out is documented here, and only these.""" + found = set() + for scan_root in NON_CAPTURABLE_SCAN_ROOTS: + for path in sorted((_REPO_ROOT / scan_root).rglob("*.py")): + rel = path.relative_to(_REPO_ROOT).as_posix() + if "/test/" in rel: + continue + for name in _warp_capturable_false_targets(ast.parse(path.read_text(), filename=rel)): + found.add((rel, name)) + expected = {(f"source/{suffix.split('/')[0]}/{suffix}", name) for suffix, name in EXPECTED_NON_CAPTURABLE} + assert found == expected, ( + "Non-capturable inventory drift.\n" + f" unexpected: {sorted(found - expected)}\n" + f" missing: {sorted(expected - found)}" + ) diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py index 1b9c10a76895..d49aa5b498cf 100644 --- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py +++ b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py @@ -7,10 +7,12 @@ from __future__ import annotations +import os import unittest +import unittest.mock import warp as wp -from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache +from isaaclab_experimental.utils.warp_graph_cache import CAPTURE_ENV_VAR, WarpGraphCache @wp.kernel @@ -20,12 +22,149 @@ def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)): b[i] = a[i] + 1.0 +@wp.kernel +def _increment(value: wp.array(dtype=wp.int32)): + """Increment a scalar state buffer once.""" + value[0] += 1 + + class TestWarpGraphCache(unittest.TestCase): """Tests for :class:`WarpGraphCache`.""" def setUp(self): self.device = "cuda:0" - self.cache = WarpGraphCache() + self.cache = WarpGraphCache(device=self.device) + + # ------------------------------------------------------------------ + # Capture policy + # ------------------------------------------------------------------ + + def test_env_var_forces_eager(self): + """Setting the capture environment variable to 0 should force every stage eager.""" + with unittest.mock.patch.dict(os.environ, {CAPTURE_ENV_VAR: "0"}): + cache = WarpGraphCache(device=self.device) + call_count = 0 + + def counted_call(): + nonlocal call_count + call_count += 1 + + for _ in range(3): + cache.call("RewardManager_compute", counted_call) + + self.assertEqual(call_count, 3) + self.assertEqual(cache.captured_stages, ()) + + def test_captured_stages_and_eager_groups_reflect_state(self): + """Introspection should report captured stages and non-capturable groups.""" + self.cache.register_capturability("Scene", False) + src = wp.full(4, value=1.0, dtype=wp.float32, device=self.device) + dst = wp.zeros(4, dtype=wp.float32, device=self.device) + + def launch(): + wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) + + self.cache.call("RewardManager_compute", launch) + self.cache.call("RewardManager_compute", launch) + + self.assertEqual(self.cache.captured_stages, ("RewardManager_compute",)) + self.assertEqual(self.cache.eager_groups, ("Scene",)) + + def test_non_capturable_group_stays_eager(self): + """Every stage in a non-capturable group should execute eagerly.""" + call_count = 0 + + def counted_call(): + nonlocal call_count + call_count += 1 + + self.cache.register_capturability("RewardManager", False) + self.cache.call("RewardManager_compute", counted_call) + self.cache.call("RewardManager_reset", counted_call) + + self.assertEqual(call_count, 2) + + def test_cpu_device_stays_eager(self): + """CPU stages should never attempt CUDA graph capture.""" + cache = WarpGraphCache(device="cpu") + call_count = 0 + + def counted_call(): + nonlocal call_count + call_count += 1 + + for _ in range(3): + cache.call("RewardManager_compute", counted_call) + + self.assertEqual(call_count, 3) + + def test_disabled_cache_stays_eager(self): + """An owner should be able to keep an unvalidated stage boundary eager.""" + cache = WarpGraphCache(enabled=False, device=self.device) + call_count = 0 + + def counted_call(): + nonlocal call_count + call_count += 1 + + for _ in range(3): + cache.call("direct_stage", counted_call) + + self.assertEqual(call_count, 3) + + def test_capturability_registration_is_conservative(self): + """A positive registration must not override a known unsafe operation.""" + self.cache.register_capturability("EventManager", False) + self.cache.register_capturability("EventManager", True) + + self.assertFalse(self.cache.is_capturable("EventManager")) + + def test_call_forwards_arguments_and_applies_output(self): + """The frontend call should forward normal arguments and transform output.""" + self.cache.register_capturability("RewardManager", False) + + result = self.cache.call( + "RewardManager_compute", + lambda value, *, scale: value * scale, + 3, + scale=4, + output=lambda value: value + 1, + ) + + self.assertEqual(result, 13) + + def test_call_captures_by_default(self): + """A capturable stage should warm up, capture, and then replay across calls.""" + call_count = 0 + src = wp.full(4, value=2.0, dtype=wp.float32, device=self.device) + dst = wp.zeros(4, dtype=wp.float32, device=self.device) + + def counted_launch(): + nonlocal call_count + call_count += 1 + wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) + return dst + + result = self.cache.call("RewardManager_compute", counted_launch) + captured = self.cache.call("RewardManager_compute", counted_launch) + replay = self.cache.call("RewardManager_compute", counted_launch) + + self.assertEqual(call_count, 2) + self.assertIs(result, captured) + self.assertIs(result, replay) + self.assertAlmostEqual(replay.numpy()[0], 3.0, places=5) + + def test_call_executes_stateful_stage_once_per_logical_call(self): + """Warm-up and capture should not advance manager state twice in one call.""" + state = wp.zeros(1, dtype=wp.int32, device=self.device) + + def increment_state(): + wp.launch(_increment, dim=1, inputs=[state], device=self.device) + return state + + for expected in (1, 2, 3): + result = self.cache.call("RewardManager_stateful", increment_state) + self.assertEqual(result.numpy()[0], expected) # ------------------------------------------------------------------ # Warm-up diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst new file mode 100644 index 000000000000..7ed70c412633 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added boolean environment-mask resets for Newton actuator and external-wrench + state, avoiding host index conversion in Warp environments. diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py index 689232b81e60..b8cc862de278 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py @@ -26,11 +26,12 @@ import warp as wp from newton.actuators import Actuator, Clamping, Delay +from isaaclab.utils.warp.utils import resolve_1d_mask + from .kernels import ( build_implicit_dof_mask, build_per_dof_env_mask_kernel, scatter_gain_kernel, - set_mask_kernel, zero_at_indices_kernel, ) @@ -112,6 +113,8 @@ def __init__( self._states_a = [act.state() for act in actuators] self._states_b = [act.state() for act in actuators] + self._reset_env_mask = wp.zeros(num_envs, dtype=wp.bool, device=device) + self._reset_dof_masks = [wp.zeros(act.indices.shape[0], dtype=wp.bool, device=device) for act in actuators] # Pre-clamp computed effort buffer. Each Newton actuator scatter-adds # its raw controller output to ``sim_control.joint_computed_f`` when @@ -166,14 +169,20 @@ def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: act.step(sim_state, sim_control, sa, sb, dt=dt) self._states_a, self._states_b = self._states_b, self._states_a - def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: + def reset( + self, + env_ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None, + env_mask: wp.array(dtype=wp.bool) | None = None, + ) -> None: """Reset actuator states for the given environments. Args: - env_ids: Environment indices to reset. ``None`` (or - ``slice(None)``, which IsaacLab callers sometimes pass) - resets all environments. Otherwise expects a torch tensor - or sequence of int indices. + env_ids: Environment indices to reset. Accepts a sequence, slice, + Torch tensor, or ``wp.int32`` array. ``None`` and + ``slice(None)`` reset all environments. + env_mask: Boolean Warp mask selecting environments to reset. When + provided, it takes precedence over :paramref:`env_ids`. Shape + is ``(num_envs,)`` on the adapter device. Newton's :meth:`Actuator.State.reset` expects a per-DOF boolean mask of length ``num_actuators`` (= ``num_envs * dofs_per_actuator``), @@ -182,7 +191,7 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: etc.). We therefore build a per-actuator per-DOF mask from the env mask before delegating to each state. """ - if env_ids is None or env_ids == slice(None): + if env_mask is None and (env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None))): for sa, sb in zip(self._states_a, self._states_b): if sa is not None: sa.reset(None) @@ -190,19 +199,22 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: sb.reset(None) return - if isinstance(env_ids, torch.Tensor): - if env_ids.numel() == 0: - return - idx = wp.from_torch(env_ids.to(device=self._device).contiguous().to(torch.int32), dtype=wp.int32) - else: - if len(env_ids) == 0: - return - idx = wp.array(list(env_ids), dtype=wp.int32, device=self._device) - env_mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) - wp.launch(set_mask_kernel, dim=idx.shape[0], inputs=[env_mask, idx], device=self._device) - - for act, sa, sb in zip(self.actuators, self._states_a, self._states_b): - per_dof_mask = wp.zeros(act.indices.shape[0], dtype=wp.bool, device=self._device) + if env_mask is None: + env_mask = resolve_1d_mask( + ids=env_ids, + all_mask=self._reset_env_mask, + scratch_mask=self._reset_env_mask, + device=self._device, + ) + if ( + env_mask.dtype != wp.bool + or env_mask.ndim != 1 + or env_mask.shape[0] != self._num_envs + or env_mask.device != wp.get_device(self._device) + ): + raise ValueError(f"env_mask must have shape ({self._num_envs},) and dtype wp.bool, received {env_mask}.") + + for act, sa, sb, per_dof_mask in zip(self.actuators, self._states_a, self._states_b, self._reset_dof_masks): wp.launch( build_per_dof_env_mask_kernel, dim=act.indices.shape[0], diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index e0c26072c1da..0bf09ce4a210 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -270,9 +270,18 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) - # reset Lab actuators registered on this articulation - for actuator in self.actuators.values(): - actuator.reset(env_ids) + # Reset Lab actuators only outside the Newton-native mode. In native mode + # the per-env actuator state (delay queues, network hidden states) lives + # in the global adapter, which consumes the mask below, and the Lab + # actuator objects only carry configuration; resetting them here would + # touch unselected environments. + if not getattr(self, "_has_newton_actuators", False): + actuator_env_ids = env_ids + if self.actuators and env_mask is not None: + torch_mask = wp.to_torch(env_mask) + actuator_env_ids = torch_mask.nonzero(as_tuple=False).squeeze(-1) + for actuator in self.actuators.values(): + actuator.reset(actuator_env_ids) # reset the global Newton actuator adapter (its ``_states_a/_b`` buffers # carry per-env state — delay queues, neural hidden states — that must # be cleared for the resetting envs). The adapter spans the whole model, @@ -281,7 +290,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None # ``getattr`` guards subclasses (e.g. ``Multirotor``) that override # ``_process_actuators_cfg`` and never initialize ``_has_newton_actuators``. if getattr(self, "_has_newton_actuators", False) and SimulationManager._adapter is not None: - SimulationManager._adapter.reset(env_ids) + SimulationManager._adapter.reset(env_ids, env_mask=env_mask) # reset external wrenches. self._instantaneous_wrench_composer.reset(env_ids, env_mask) self._permanent_wrench_composer.reset(env_ids, env_mask) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py index 3265ad2896dd..3f6eac323887 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -136,8 +136,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None if (env_ids is None) or (env_ids == slice(None)): env_ids = slice(None) # reset external wrench - self._instantaneous_wrench_composer.reset(env_ids) - self._permanent_wrench_composer.reset(env_ids) + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) def write_data_to_sim(self) -> None: """Write external wrench to the simulation. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index 57512ce1b33e..e0eb6276d5ec 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -182,8 +182,8 @@ def reset( if object_ids is None: object_ids = self._ALL_BODY_INDICES # reset external wrench - self._instantaneous_wrench_composer.reset(env_ids) - self._permanent_wrench_composer.reset(env_ids) + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) def write_data_to_sim(self) -> None: """Write external wrench to the simulation. @@ -266,8 +266,8 @@ def write_body_pose_to_sim_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7) @@ -297,8 +297,8 @@ def write_body_pose_to_sim_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body poses in simulation frame. Shape is (num_instances, num_bodies, 7) @@ -339,8 +339,8 @@ def write_body_velocity_to_sim_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body velocities in simulation frame. @@ -374,8 +374,8 @@ def write_body_velocity_to_sim_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body velocities in simulation frame. @@ -422,8 +422,8 @@ def write_body_link_pose_to_sim_index( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body link poses in simulation frame. @@ -479,8 +479,8 @@ def write_body_link_pose_to_sim_mask( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body link poses in simulation frame. Shape is (num_instances, num_bodies, 7) @@ -520,8 +520,8 @@ def write_body_com_pose_to_sim_index( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body center of mass poses in simulation frame. @@ -580,8 +580,8 @@ def write_body_com_pose_to_sim_mask( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_poses: Body center of mass poses in simulation frame. Shape is (num_instances, num_bodies, 7) @@ -624,8 +624,8 @@ def write_body_com_velocity_to_sim_index( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body center of mass velocities in simulation frame. @@ -689,8 +689,8 @@ def write_body_com_velocity_to_sim_mask( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body center of mass velocities in simulation frame. @@ -738,8 +738,8 @@ def write_body_link_velocity_to_sim_index( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body link velocities in simulation frame. @@ -806,8 +806,8 @@ def write_body_link_velocity_to_sim_mask( May trigger per-environment FK recomputation and solver reset (Kamino) for the affected environments. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: body_velocities: Body link velocities in simulation frame. Shape is (num_instances, num_bodies, 6) @@ -848,8 +848,8 @@ def set_masses_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: masses: Masses of all bodies. Shape is (len(env_ids), len(body_ids)). @@ -891,8 +891,8 @@ def set_masses_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: masses: Masses of all bodies. Shape is (num_instances, num_bodies). @@ -935,8 +935,8 @@ def set_coms_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. .. caution:: Unlike the PhysX version of this method, this method does not set the center of mass orientation. @@ -985,8 +985,8 @@ def set_coms_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. .. caution:: Unlike the PhysX version of this method, this method does not set the center of mass orientation. @@ -1036,8 +1036,8 @@ def set_inertias_index( This method expects partial data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: inertias: Inertias of all bodies. Shape is (len(env_ids), len(body_ids), 9). @@ -1079,8 +1079,8 @@ def set_inertias_mask( This method expects full data. .. tip:: - Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. - However, to allow graphed pipelines, the mask method must be used. + Both the index and mask methods route through the same index implementation: mask inputs are + currently compacted to host indices internally, so neither variant is CUDA-graph capturable yet. Args: inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). @@ -1277,6 +1277,8 @@ def _resolve_body_ids(self, body_ids) -> wp.array: return wp.from_torch(body_ids, dtype=wp.int32) return body_ids + # TODO: Replace the host-side compaction below with mask-native write kernels so the + # write_*_to_sim_mask methods become CUDA-graph capturable. def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array | torch.Tensor: """Resolve environment mask to indices via torch.nonzero.""" if env_mask is not None: diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py index 3eb06e007201..2be89399d03a 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py @@ -219,7 +219,8 @@ def _debug_vis_callback(self, event): up_axis = UsdGeom.GetStageUpAxis(self.stage) pos_w_torch = self._data.pos_w.torch accel_w = math_utils.quat_apply(self._data.quat_w.torch, self._data.lin_acc_b.torch) - valid_indices = (torch.linalg.norm(accel_w, dim=-1) > 1e-5).nonzero(as_tuple=True)[0] + accel_valid = torch.linalg.norm(accel_w, dim=-1) > 1e-5 + valid_indices = accel_valid.nonzero(as_tuple=True)[0] if valid_indices.numel() == 0: return pos_filtered = pos_w_torch.index_select(0, valid_indices) diff --git a/source/isaaclab_newton/test/actuators/test_adapter_mask_reset.py b/source/isaaclab_newton/test/actuators/test_adapter_mask_reset.py new file mode 100644 index 000000000000..e33235ba28c7 --- /dev/null +++ b/source/isaaclab_newton/test/actuators/test_adapter_mask_reset.py @@ -0,0 +1,65 @@ +# 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 + +"""Tests for mask-based Newton actuator adapter resets.""" + +from __future__ import annotations + +from unittest.mock import Mock + +import numpy as np +import warp as wp +from isaaclab_newton.actuators.adapter import NewtonActuatorAdapter + + +def test_adapter_expands_environment_mask_without_allocating() -> None: + """Adapter state masks should select every actuator DOF in masked environments.""" + adapter = NewtonActuatorAdapter.__new__(NewtonActuatorAdapter) + actuator = Mock() + actuator.indices = wp.array([0, 1, 2, 3], dtype=wp.uint32, device="cpu") + state_a = Mock() + state_b = Mock() + adapter.actuators = [actuator] + adapter._states_a = [state_a] + adapter._states_b = [state_b] + adapter._num_envs = 2 + adapter._dof_offset = 0 + adapter.num_joints = 2 + adapter._device = "cpu" + adapter._reset_env_mask = wp.zeros(2, dtype=wp.bool, device="cpu") + adapter._reset_dof_masks = [wp.zeros(4, dtype=wp.bool, device="cpu")] + env_mask = wp.array([False, True], dtype=wp.bool, device="cpu") + + adapter.reset(env_mask=env_mask) + + state_a.reset.assert_called_once() + state_b.reset.assert_called_once() + np.testing.assert_array_equal(state_a.reset.call_args.args[0].numpy(), [False, False, True, True]) + assert state_a.reset.call_args.args[0] is adapter._reset_dof_masks[0] + + +def test_adapter_reuses_id_mask_without_stale_selection() -> None: + """Sequential ID resets should clear the reusable environment mask.""" + adapter = NewtonActuatorAdapter.__new__(NewtonActuatorAdapter) + actuator = Mock() + actuator.indices = wp.array([0, 1, 2, 3], dtype=wp.uint32, device="cpu") + state_a = Mock() + adapter.actuators = [actuator] + adapter._states_a = [state_a] + adapter._states_b = [None] + adapter._num_envs = 2 + adapter._dof_offset = 0 + adapter.num_joints = 2 + adapter._device = "cpu" + adapter._reset_env_mask = wp.zeros(2, dtype=wp.bool, device="cpu") + adapter._reset_dof_masks = [wp.zeros(4, dtype=wp.bool, device="cpu")] + + adapter.reset(env_ids=[0]) + first_mask = state_a.reset.call_args.args[0].numpy().copy() + adapter.reset(env_ids=[1]) + second_mask = state_a.reset.call_args.args[0].numpy().copy() + + np.testing.assert_array_equal(first_mask, [True, True, False, False]) + np.testing.assert_array_equal(second_mask, [False, False, True, True]) diff --git a/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py b/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py new file mode 100644 index 000000000000..c01f89c64162 --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_newton_mask_reset_forwarding.py @@ -0,0 +1,86 @@ +# 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 + +"""Tests for Warp reset-mask forwarding to Newton asset wrench storage.""" + +from unittest.mock import Mock, patch + +import warp as wp +from isaaclab_newton.assets.articulation.articulation import Articulation +from isaaclab_newton.assets.rigid_object.rigid_object import RigidObject +from isaaclab_newton.assets.rigid_object_collection.rigid_object_collection import RigidObjectCollection + + +def test_rigid_object_forwards_reset_mask_to_wrench_composers() -> None: + """Rigid-object resets should preserve mask selection at wrench storage.""" + asset = RigidObject.__new__(RigidObject) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_mask = wp.array([False, True], dtype=wp.bool, device="cpu") + + asset.reset(env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) + asset._permanent_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) + + +def test_rigid_object_collection_forwards_reset_mask_to_wrench_composers() -> None: + """Collection resets should preserve mask selection at wrench storage.""" + asset = RigidObjectCollection.__new__(RigidObjectCollection) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_ids = [0] + env_mask = wp.array([True, False], dtype=wp.bool, device="cpu") + + asset.reset(env_ids=env_ids, object_ids=slice(None), env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(env_ids, env_mask) + asset._permanent_wrench_composer.reset.assert_called_once_with(env_ids, env_mask) + + +def test_articulation_without_lab_actuators_keeps_mask_on_device() -> None: + """Mask resets should not materialize IDs when no Torch actuator consumes them.""" + asset = Articulation.__new__(Articulation) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + asset.actuators = {} + asset._has_newton_actuators = False + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_mask = object() + + asset.reset(env_mask=env_mask) + + asset._instantaneous_wrench_composer.reset.assert_called_once_with(slice(None), env_mask) + + +def test_articulation_native_mode_masked_reset_skips_lab_actuators() -> None: + """Native-mode partial resets must not reset Lab actuator state for all environments.""" + asset = Articulation.__new__(Articulation) + asset._initialize_handle = None + asset._invalidate_initialize_handle = None + asset._prim_deletion_handle = None + lab_actuator = Mock() + asset.actuators = {"legs": lab_actuator} + asset._has_newton_actuators = True + asset._instantaneous_wrench_composer = Mock() + asset._permanent_wrench_composer = Mock() + env_mask = wp.array([False, True], dtype=wp.bool, device="cpu") + adapter = Mock() + + from isaaclab_newton.physics import NewtonManager + + with patch.object(NewtonManager, "_adapter", adapter): + asset.reset(env_mask=env_mask) + + lab_actuator.reset.assert_not_called() + adapter.reset.assert_called_once_with(slice(None), env_mask=env_mask) diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rl_games.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rl_games.py index 7b0e4baff460..eb035bcd0f64 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rl_games.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rl_games.py @@ -26,7 +26,12 @@ from isaaclab.utils.dict import print_dict from isaaclab.utils.seed import configure_seed -from isaaclab_rl.entrypoints.common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector +from isaaclab_rl.entrypoints.common import ( + CHECKPOINT_SELECTORS, + add_frontend_args, + create_isaaclab_env, + resolve_checkpoint_selector, +) from isaaclab_rl.rl_games import RlGamesGpuEnv, RlGamesVecEnvWrapper from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint @@ -50,6 +55,7 @@ ) parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") parser.add_argument("--task", type=str, default=None, help="Name of the task.") +add_frontend_args(parser) parser.add_argument( "--agent", type=str, default="rl_games_cfg_entry_point", help="Name of the RL agent configuration entry point." ) @@ -126,12 +132,12 @@ def main(): obs_groups = agent_cfg["params"]["env"].get("obs_groups") concate_obs_groups = agent_cfg["params"]["env"].get("concate_obs_groups", True) - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) - - if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): - from isaaclab.envs import multi_agent_to_single_agent - - env = multi_agent_to_single_agent(env) + env = create_isaaclab_env( + args_cli.task, + env_cfg, + args_cli, + convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg), + ) if args_cli.video: video_kwargs = { diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rsl_rl.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rsl_rl.py index 966010a3bc9f..7b3a1143b969 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rsl_rl.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_rsl_rl.py @@ -25,7 +25,12 @@ from isaaclab.utils.string import list_intersection, string_to_callable from isaaclab_rl.entrypoints.backends import cli_args_rsl_rl as cli_args -from isaaclab_rl.entrypoints.common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector +from isaaclab_rl.entrypoints.common import ( + CHECKPOINT_SELECTORS, + add_frontend_args, + create_isaaclab_env, + resolve_checkpoint_selector, +) from isaaclab_rl.rsl_rl import ( RslRlBaseRunnerCfg, RslRlVecEnvWrapper, @@ -55,6 +60,7 @@ ) parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") parser.add_argument("--task", type=str, default=None, help="Name of the task.") +add_frontend_args(parser) parser.add_argument( "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point." ) @@ -130,12 +136,12 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen env_cfg.log_dir = log_dir - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) - - if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): - from isaaclab.envs import multi_agent_to_single_agent - - env = multi_agent_to_single_agent(env) + env = create_isaaclab_env( + args_cli.task, + env_cfg, + args_cli, + convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg), + ) if args_cli.video: video_kwargs = { diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_sb3.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_sb3.py index 3bef5d69cb1a..00f949bd4a66 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_sb3.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_sb3.py @@ -23,7 +23,12 @@ from isaaclab.utils.dict import print_dict from isaaclab.utils.seed import configure_seed -from isaaclab_rl.entrypoints.common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector +from isaaclab_rl.entrypoints.common import ( + CHECKPOINT_SELECTORS, + add_frontend_args, + create_isaaclab_env, + resolve_checkpoint_selector, +) from isaaclab_rl.sb3 import Sb3VecEnvWrapper, process_sb3_cfg from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint @@ -47,6 +52,7 @@ ) parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") parser.add_argument("--task", type=str, default=None, help="Name of the task.") +add_frontend_args(parser) parser.add_argument( "--agent", type=str, default="sb3_cfg_entry_point", help="Name of the RL agent configuration entry point." ) @@ -116,15 +122,15 @@ def main(): env_cfg.log_dir = log_dir - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + env = create_isaaclab_env( + args_cli.task, + env_cfg, + args_cli, + convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg), + ) agent_cfg = process_sb3_cfg(agent_cfg, env.unwrapped.num_envs) - if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): - from isaaclab.envs import multi_agent_to_single_agent - - env = multi_agent_to_single_agent(env) - if args_cli.video: video_kwargs = { "video_folder": os.path.join(log_dir, "videos", "play"), diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py index cd566c13f23e..6654f28d2364 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/backends/play_skrl.py @@ -27,7 +27,12 @@ from isaaclab.utils.dict import print_dict from isaaclab.utils.seed import configure_seed -from isaaclab_rl.entrypoints.common import CHECKPOINT_SELECTORS, resolve_checkpoint_selector +from isaaclab_rl.entrypoints.common import ( + CHECKPOINT_SELECTORS, + add_frontend_args, + create_isaaclab_env, + resolve_checkpoint_selector, +) from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint import isaaclab_tasks # noqa: F401 @@ -52,6 +57,7 @@ ) parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") parser.add_argument("--task", type=str, default=None, help="Name of the task.") +add_frontend_args(parser) parser.add_argument( "--agent", type=str, @@ -166,12 +172,12 @@ def main(): env_cfg.log_dir = log_dir - env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) - - if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg) and algorithm in ["ppo"]: - from isaaclab.envs import multi_agent_to_single_agent - - env = multi_agent_to_single_agent(env) + env = create_isaaclab_env( + args_cli.task, + env_cfg, + args_cli, + convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg) and algorithm in ["ppo"], + ) try: dt = env.step_dt diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py index a46d1f540400..eed6286410ec 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py @@ -233,6 +233,30 @@ def resolve_play_checkpoint(checkpoint: str | None, framework: str, task: str) - return path +def add_frontend_args(parser: argparse.ArgumentParser) -> None: + """Add the environment-runtime selector argument. + + The flag is always registered so the CLI surface does not depend on + optional packages; the warp runtime itself is imported only when selected + (see :func:`create_isaaclab_env`). + + Args: + parser: The parser to add the argument to. + """ + parser.add_argument( + "--frontend", + type=str, + choices=["torch", "warp"], + default="torch", + help=( + "Runtime that constructs the environment. 'torch' uses the registered stable environment via" + " gym.make. 'warp' (experimental) adapts a manager-based task config onto the Warp runtime, or" + " dispatches a direct task to its registered Warp environment; requires isaaclab_experimental" + " and `presets=newton_mjwarp`." + ), + ) + + def add_common_train_args( parser: argparse.ArgumentParser, *, @@ -257,6 +281,7 @@ def add_common_train_args( ) parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") parser.add_argument("--task", type=str, default=None, help="Name of the task.") + add_frontend_args(parser) if include_agent: parser.add_argument("--agent", type=str, default=agent_default, help=agent_help) parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") @@ -405,7 +430,15 @@ def create_isaaclab_env( Returns: The created Gymnasium environment. """ - env = gym.make(task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + render_mode = "rgb_array" if args_cli.video else None + if args_cli.frontend == "torch": + env = gym.make(task, cfg=env_cfg, render_mode=render_mode) + else: + # Imported lazily: the warp frontend lives in the optional + # isaaclab_experimental package, and the torch path must work without it. + from isaaclab_experimental.envs.frontend import WarpFrontend + + env = WarpFrontend.build_env(env_cfg, task, render_mode=render_mode) if convert_marl_to_single_agent and isinstance(env.unwrapped.cfg, DirectMARLEnvCfg): from isaaclab.envs import multi_agent_to_single_agent diff --git a/source/isaaclab_rl/test/test_entrypoints_common.py b/source/isaaclab_rl/test/test_entrypoints_common.py index 415bb56e27c6..6b7ec1a1ef60 100644 --- a/source/isaaclab_rl/test/test_entrypoints_common.py +++ b/source/isaaclab_rl/test/test_entrypoints_common.py @@ -20,6 +20,7 @@ from isaaclab_rl.entrypoints.common import ( CaptureEnvSensors, add_common_train_args, + create_isaaclab_env, dispatch_library_entrypoint, enable_cameras_for_video, wrap_sensor_capture, @@ -204,6 +205,80 @@ def test_enable_cameras_for_video_enables_cameras_for_sensor_capture() -> None: assert args_cli.enable_cameras +def test_common_train_args_register_frontend_with_torch_default() -> None: + """Every RL library CLI exposes ``--frontend`` and defaults to the torch runtime.""" + parser = argparse.ArgumentParser() + add_common_train_args(parser, agent_default=None, agent_help="", include_agent=False) + + assert parser.parse_args([]).frontend == "torch" + assert parser.parse_args(["--frontend", "warp"]).frontend == "warp" + with pytest.raises(SystemExit): + parser.parse_args(["--frontend", "tensorflow"]) + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[3] + + +def test_play_entrypoints_route_through_frontend_factory() -> None: + """Every dispatched play backend constructs its env via the frontend-aware factory.""" + play_scripts = sorted( + (_repo_root() / "source" / "isaaclab_rl" / "isaaclab_rl" / "entrypoints" / "backends").glob("play_*.py") + ) + # rlinf constructs environments inside the external framework; the frontend cannot + # reach it (documented limitation). + play_scripts = [path for path in play_scripts if path.name != "play_rlinf.py"] + assert len(play_scripts) == 4, sorted(path.name for path in play_scripts) + for script in play_scripts: + source = script.read_text() + assert "create_isaaclab_env(" in source, f"{script.name} bypasses the frontend factory" + assert "gym.make(args_cli.task" not in source, f"{script.name} constructs directly via gym.make" + assert "add_frontend_args(parser)" in source, f"{script.name} does not expose --frontend" + + +def test_create_isaaclab_env_uses_registered_torch_env_by_default(monkeypatch: pytest.MonkeyPatch) -> None: + """The shared factory preserves the existing Gym path when no frontend is selected.""" + expected_env = object() + env_cfg = object() + calls: list[tuple[Any, ...]] = [] + + def fake_make(task: str, **kwargs: Any) -> Any: + calls.append((task, kwargs)) + return expected_env + + monkeypatch.setattr(_rl_common.gym, "make", fake_make) + args_cli = argparse.Namespace(video=False, frontend="torch") + + env = create_isaaclab_env("Isaac-Test", env_cfg, args_cli, convert_marl_to_single_agent=False) + + assert env is expected_env + assert len(calls) == 1 + assert calls[0][0] == "Isaac-Test" + assert calls[0][1]["cfg"] is env_cfg + assert calls[0][1]["render_mode"] is None + + +def test_create_isaaclab_env_uses_selected_warp_frontend(monkeypatch: pytest.MonkeyPatch) -> None: + """The shared factory delegates Warp selection to the experimental frontend.""" + import isaaclab_experimental.envs.frontend as frontend_module + + expected_env = object() + env_cfg = object() + calls: list[tuple[Any, ...]] = [] + + def fake_build_env(cfg: Any, task: str, **kwargs: Any) -> Any: + calls.append((cfg, task, kwargs)) + return expected_env + + monkeypatch.setattr(frontend_module.WarpFrontend, "build_env", fake_build_env) + args_cli = argparse.Namespace(video=True, frontend="warp") + + env = create_isaaclab_env("Isaac-Test", env_cfg, args_cli, convert_marl_to_single_agent=False) + + assert env is expected_env + assert calls == [(env_cfg, "Isaac-Test", {"render_mode": "rgb_array"})] + + def test_dispatch_library_entrypoint_shows_help_without_library( tmp_path: Path, capsys: pytest.CaptureFixture[str] ) -> None: diff --git a/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst b/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst new file mode 100644 index 000000000000..7a9bdd89d0b5 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/warp-manager-bridge.rst @@ -0,0 +1,8 @@ +Added +^^^^^ + +* Added ``warp_entry_point`` declarations to the ``Isaac-Cartpole-Direct``, + ``Isaac-Ant-Direct``, ``Isaac-Humanoid-Direct``, and + ``Isaac-Reorient-Cube-Allegro-Direct`` registrations so ``--frontend warp`` + constructs the Warp implementation of these tasks with the stable task + configuration. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py index bdb55030b3b2..cc8bff65c70a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cartpole/__init__.py @@ -25,6 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.cartpole_direct_env_cfg:CartpoleEnvCfg", + "warp_entry_point": "isaaclab_tasks_experimental.core.cartpole.cartpole_warp_env:CartpoleWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpoleDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py index 4adb37ca8eb6..1e506e5ef4ae 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/ant/__init__.py @@ -25,6 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.ant_direct_env_cfg:AntEnvCfg", + "warp_entry_point": "isaaclab_tasks_experimental.core.locomotion.ant.ant_env_warp:AntWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py index d13b7697595a..e3cee56589b2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/humanoid/__init__.py @@ -25,6 +25,7 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.humanoid_direct_env_cfg:HumanoidEnvCfg", + "warp_entry_point": "isaaclab_tasks_experimental.core.locomotion.humanoid.humanoid_warp_env:HumanoidWarpEnv", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidDirectPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/__init__.py index c5ec71b59cfe..35393451e9e1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/allegro_hand/__init__.py @@ -47,6 +47,9 @@ disable_env_checker=True, kwargs={ "env_cfg_entry_point": f"{__name__}.allegro_hand_direct_env_cfg:AllegroHandEnvCfg", + "warp_entry_point": ( + "isaaclab_tasks_experimental.core.reorient.inhand_manipulation_warp_env:InHandManipulationWarpEnv" + ), "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_direct_ppo_cfg.yaml", "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AllegroHandPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst new file mode 100644 index 000000000000..6cfc728745b0 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.minor.rst @@ -0,0 +1,11 @@ +Added +^^^^^ + +* Added Warp-native terrain and reward-weight curricula and connected locomotion + and reach hot paths to pointer-stable Warp command arrays. + +Fixed +^^^^^ + +* Fixed terrain-boundary termination state leaking between environment + instances with different terrain configurations. diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst new file mode 100644 index 000000000000..f29f10f3e6ab --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-manager-bridge.minor.rst @@ -0,0 +1,39 @@ +Changed +^^^^^^^ + +* **Breaking:** Changed the package layout to mirror :mod:`isaaclab_tasks.core` + (``isaaclab_tasks_experimental.core.``), replacing the previous + ``manager_based``/``direct`` split. Update imports of the old paths to the + ``core`` equivalents (e.g. ``isaaclab_tasks_experimental.core.cartpole.mdp``). + +Removed +^^^^^^^ + +* **Breaking:** Removed all manager-based ``*-Warp-v0`` task registrations — + ``Isaac-Cartpole-Warp-v0``, ``Isaac-Humanoid-Warp-v0``, ``Isaac-Ant-Warp-v0``, + ``Isaac-Reach-Franka-Warp-v0``, ``Isaac-Reach-Franka-Warp-Play-v0``, and the + velocity variants (``Isaac-Velocity-Flat--Warp-v0`` and their + ``-Warp-Play-v0`` forms) — together with their environment configurations. + Run the stable task ids with ``--frontend warp`` and + ``presets=newton_mjwarp`` instead, e.g. + ``--task Isaac-Velocity-Flat-AnymalD --frontend warp presets=newton_mjwarp``. +* **Breaking:** Removed the duplicated direct warp task registrations + ``Isaac-Cartpole-Direct-Warp-v0``, ``Isaac-Ant-Direct-Warp-v0``, + ``Isaac-Humanoid-Direct-Warp-v0``, and + ``Isaac-Reorient-Cube-Allegro-Direct-Warp-v0`` together with their + duplicated environment configurations; the stable registrations declare the + warp environment classes via ``warp_entry_point``. Run the stable task ids + with ``--frontend warp`` and ``presets=newton_mjwarp`` instead, e.g. + ``--task Isaac-Cartpole-Direct --frontend warp presets=newton_mjwarp``. +* Removed the unregistered rough velocity warp configurations; rough-terrain + warp tasks remain unsupported until :class:`~isaaclab.terrains.TerrainImporter` + gains Warp APIs. + +Fixed +^^^^^ + +* Fixed the direct Warp Cartpole task to match the stable task's observations, + reset ranges, termination condition, reward scaling, and scene configuration. +* Fixed the Warp Cartpole ``survival_success_rate`` twin to report the + ``Metrics/success_rate`` value on-device instead of silently dropping the + metric. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py similarity index 68% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py index 0857176a3fc7..bf495309161f 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/__init__.py @@ -3,4 +3,4 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Velocity locomotion experimental task registrations (manager-based).""" +"""Warp-first task implementations mirroring the :mod:`isaaclab_tasks.core` layout.""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py new file mode 100644 index 000000000000..e8bc71c28dc3 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/__init__.py @@ -0,0 +1,12 @@ +# 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 + +"""Warp implementations for the stable Cartpole tasks. + +Holds the MDP twins for the manager-based ``Isaac-Cartpole`` task and the +direct :class:`~isaaclab_tasks_experimental.core.cartpole.cartpole_warp_env.CartpoleWarpEnv` +declared by ``Isaac-Cartpole-Direct``. There is no separate task registration: +run the stable ids with ``--frontend warp`` and ``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py similarity index 71% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py index 14f776c95b0b..8ebae3c986b7 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py @@ -5,11 +5,11 @@ from __future__ import annotations -import math from typing import TYPE_CHECKING import warp as wp from isaaclab_experimental.envs import DirectRLEnvWarp +from isaaclab_experimental.utils.warp.utils import wrap_to_pi import isaaclab.sim as sim_utils from isaaclab import cloner @@ -17,22 +17,24 @@ from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane if TYPE_CHECKING: - from .cartpole_warp_env_cfg import CartpoleWarpEnvCfg + from isaaclab_tasks.core.cartpole.cartpole_direct_env_cfg import CartpoleEnvCfg @wp.kernel def get_observations( joint_pos: wp.array2d(dtype=wp.float32), joint_vel: wp.array2d(dtype=wp.float32), + default_joint_pos: wp.array2d(dtype=wp.float32), + default_joint_vel: wp.array2d(dtype=wp.float32), cart_dof_idx: wp.int32, pole_dof_idx: wp.int32, observations: wp.array(dtype=wp.vec4f), ): env_index = wp.tid() - observations[env_index][0] = joint_pos[env_index, pole_dof_idx] - observations[env_index][1] = joint_vel[env_index, pole_dof_idx] - observations[env_index][2] = joint_pos[env_index, cart_dof_idx] - observations[env_index][3] = joint_vel[env_index, cart_dof_idx] + observations[env_index][0] = joint_pos[env_index, cart_dof_idx] - default_joint_pos[env_index, cart_dof_idx] + observations[env_index][1] = joint_pos[env_index, pole_dof_idx] - default_joint_pos[env_index, pole_dof_idx] + observations[env_index][2] = joint_vel[env_index, cart_dof_idx] - default_joint_vel[env_index, cart_dof_idx] + observations[env_index][3] = joint_vel[env_index, pole_dof_idx] - default_joint_vel[env_index, pole_dof_idx] @wp.kernel @@ -51,7 +53,6 @@ def get_dones( joint_pos: wp.array2d(dtype=wp.float32), episode_length_buf: wp.array(dtype=wp.int32), cart_dof_idx: wp.int32, - pole_dof_idx: wp.int32, max_episode_length: wp.int32, max_cart_pos: wp.float32, out_of_bounds: wp.array(dtype=wp.bool), @@ -59,10 +60,8 @@ def get_dones( reset: wp.array(dtype=wp.bool), ): env_index = wp.tid() - out_of_bounds[env_index] = (wp.abs(joint_pos[env_index, cart_dof_idx]) > max_cart_pos) or ( - wp.abs(joint_pos[env_index, pole_dof_idx]) > math.pi / 2.0 - ) - time_out[env_index] = episode_length_buf[env_index] >= (max_episode_length - 1) + out_of_bounds[env_index] = wp.abs(joint_pos[env_index, cart_dof_idx]) > max_cart_pos + time_out[env_index] = episode_length_buf[env_index] >= max_episode_length reset[env_index] = out_of_bounds[env_index] or time_out[env_index] @@ -85,6 +84,7 @@ def compute_rew_pole_pos( rew_scale_pole_pos: wp.float32, pole_pos: wp.float32, ) -> wp.float32: + pole_pos = wrap_to_pi(pole_pos) return rew_scale_pole_pos * pole_pos * pole_pos @@ -116,10 +116,11 @@ def compute_rewards( cart_dof_idx: wp.int32, pole_dof_idx: wp.int32, reset_terminated: wp.array(dtype=wp.bool), + step_dt: wp.float32, reward: wp.array(dtype=wp.float32), ): env_index = wp.tid() - reward[env_index] = ( + reward[env_index] = step_dt * ( compute_rew_alive(rew_scale_alive, reset_terminated[env_index]) + compute_rew_termination(rew_scale_terminated, reset_terminated[env_index]) + compute_rew_pole_pos(rew_scale_pole_pos, joint_pos[env_index, pole_dof_idx]) @@ -134,21 +135,43 @@ def reset( default_joint_vel: wp.array2d(dtype=wp.float32), joint_pos: wp.array2d(dtype=wp.float32), joint_vel: wp.array2d(dtype=wp.float32), + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), + soft_joint_vel_limits: wp.array2d(dtype=wp.float32), cart_dof_idx: wp.int32, pole_dof_idx: wp.int32, - initial_pose_angle_range: wp.vec2f, + initial_cart_position_range: wp.vec2f, + initial_cart_velocity_range: wp.vec2f, + initial_pole_angle_range: wp.vec2f, + initial_pole_velocity_range: wp.vec2f, env_mask: wp.array(dtype=wp.bool), state: wp.array(dtype=wp.uint32), ): env_index = wp.tid() if env_mask[env_index]: - joint_pos[env_index, cart_dof_idx] = default_joint_pos[env_index, cart_dof_idx] - joint_pos[env_index, pole_dof_idx] = default_joint_pos[env_index, pole_dof_idx] + wp.randf( - state[env_index], initial_pose_angle_range[0] * wp.pi, initial_pose_angle_range[1] * wp.pi + rng_state = state[env_index] + cart_pos = default_joint_pos[env_index, cart_dof_idx] + wp.randf( + rng_state, initial_cart_position_range[0], initial_cart_position_range[1] + ) + pole_pos = default_joint_pos[env_index, pole_dof_idx] + wp.randf( + rng_state, initial_pole_angle_range[0], initial_pole_angle_range[1] + ) + cart_vel = default_joint_vel[env_index, cart_dof_idx] + wp.randf( + rng_state, initial_cart_velocity_range[0], initial_cart_velocity_range[1] ) - joint_vel[env_index, cart_dof_idx] = default_joint_vel[env_index, cart_dof_idx] - joint_vel[env_index, pole_dof_idx] = default_joint_vel[env_index, pole_dof_idx] - state[env_index] += wp.uint32(1) + pole_vel = default_joint_vel[env_index, pole_dof_idx] + wp.randf( + rng_state, initial_pole_velocity_range[0], initial_pole_velocity_range[1] + ) + + cart_pos_limits = soft_joint_pos_limits[env_index, cart_dof_idx] + pole_pos_limits = soft_joint_pos_limits[env_index, pole_dof_idx] + cart_vel_limit = soft_joint_vel_limits[env_index, cart_dof_idx] + pole_vel_limit = soft_joint_vel_limits[env_index, pole_dof_idx] + + joint_pos[env_index, cart_dof_idx] = wp.clamp(cart_pos, cart_pos_limits[0], cart_pos_limits[1]) + joint_pos[env_index, pole_dof_idx] = wp.clamp(pole_pos, pole_pos_limits[0], pole_pos_limits[1]) + joint_vel[env_index, cart_dof_idx] = wp.clamp(cart_vel, -cart_vel_limit, cart_vel_limit) + joint_vel[env_index, pole_dof_idx] = wp.clamp(pole_vel, -pole_vel_limit, pole_vel_limit) + state[env_index] = rng_state @wp.kernel @@ -161,9 +184,9 @@ def initialize_state( class CartpoleWarpEnv(DirectRLEnvWarp): - cfg: CartpoleWarpEnvCfg + cfg: CartpoleEnvCfg - def __init__(self, cfg: CartpoleWarpEnvCfg, render_mode: str | None = None, **kwargs) -> None: + def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs) -> None: super().__init__(cfg, render_mode, **kwargs) # Get the indices (develop API: find_joints returns (indices, names)) @@ -216,8 +239,9 @@ def _setup_scene(self) -> None: # add articulation to scene self.scene.articulations["cartpole"] = self.cartpole # add lights - light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) - light_cfg.func("/World/Light", light_cfg) + light_cfg = sim_utils.DistantLightCfg(intensity=2000.0, color=(1.0, 1.0, 1.0)) + light_orientation = (-0.14644663035869598, -0.3535534143447876, -0.3535534143447876, 0.8535533547401428) + light_cfg.func("/World/Light", light_cfg, orientation=light_orientation) def _pre_physics_step(self, actions: wp.array) -> None: wp.launch( @@ -241,6 +265,8 @@ def _get_observations(self) -> dict: inputs=[ self.joint_pos, self.joint_vel, + self.cartpole.data.default_joint_pos.warp, + self.cartpole.data.default_joint_vel.warp, self._cart_dof_idx[0], self._pole_dof_idx[0], self.observations, @@ -263,6 +289,7 @@ def _get_rewards(self) -> None: self._cart_dof_idx[0], self._pole_dof_idx[0], self.reset_terminated, + self.step_dt, self.rewards, ], ) @@ -275,7 +302,6 @@ def _get_dones(self) -> None: self.joint_pos, self._episode_length_buf_wp, self._cart_dof_idx[0], - self._pole_dof_idx[0], self.max_episode_length, self.cfg.max_cart_pos, self.reset_terminated, @@ -298,9 +324,14 @@ def _reset_idx(self, mask: wp.array | None = None) -> None: self.cartpole.data.default_joint_vel.warp, self.joint_pos, self.joint_vel, + self.cartpole.data.soft_joint_pos_limits.warp, + self.cartpole.data.soft_joint_vel_limits.warp, self._cart_dof_idx[0], self._pole_dof_idx[0], + self.cfg.initial_cart_position_range, + self.cfg.initial_cart_velocity_range, self.cfg.initial_pole_angle_range, + self.cfg.initial_pole_velocity_range, mask, self.states, ], diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi similarity index 82% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi index 377068405c79..657186c2539b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/__init__.pyi @@ -5,10 +5,11 @@ __all__ = [ "joint_pos_target_l2", + "survival_success_rate", ] # Forward stable MDP terms and experimental Warp-first overrides lazily, then # override with cartpole-specific terms below. from isaaclab_experimental.envs.mdp import * # noqa: F401, F403 -from .rewards import joint_pos_target_l2 +from .rewards import joint_pos_target_l2, survival_success_rate diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py new file mode 100644 index 000000000000..6049af38ad78 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/mdp/rewards.py @@ -0,0 +1,113 @@ +# 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 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp +from isaaclab_experimental.managers import ManagerTermBase, SceneEntityCfg +from isaaclab_experimental.utils.warp.utils import wrap_to_pi + +if TYPE_CHECKING: + import torch + + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.managers import RewardTermCfg + + +@wp.kernel +def _joint_pos_target_l2_kernel( + joint_pos: wp.array(dtype=wp.float32, ndim=2), + joint_mask: wp.array(dtype=wp.bool), + out: wp.array(dtype=wp.float32), + target: float, +): + i = wp.tid() + s = float(0.0) + for j in range(joint_pos.shape[1]): + if joint_mask[j]: + a = wrap_to_pi(joint_pos[i, j]) + d = a - target + s += d * d + out[i] = s + + +def joint_pos_target_l2(env: ManagerBasedRLEnv, out, target: float, asset_cfg: SceneEntityCfg) -> None: + """Penalize joint position deviation from a target value. Writes into ``out``.""" + asset: Articulation = env.scene[asset_cfg.name] + assert asset.data.joint_pos.warp.shape[1] == asset_cfg.joint_mask.shape[0] + wp.launch( + kernel=_joint_pos_target_l2_kernel, + dim=env.num_envs, + inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, out, target], + device=env.device, + ) + + +@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), +): + 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), +): + 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.cartpole.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) -> None: + out.zero_() diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py new file mode 100644 index 000000000000..3fba47b723b6 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/__init__.py @@ -0,0 +1,12 @@ +# 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 + +"""Warp implementations for the stable locomotion tasks. + +The ``mdp`` package holds the twins for the shared +:mod:`isaaclab_tasks.core.locomotion.mdp` terms used by the Humanoid and Ant +tasks. There is no separate task registration: run the stable ids with +``--frontend warp`` and ``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py new file mode 100644 index 000000000000..1be981faf7f0 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/__init__.py @@ -0,0 +1,14 @@ +# 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 + +"""Warp implementation of the stable direct Ant task. + +The manager-based MDP twins live in the shared +:mod:`isaaclab_tasks_experimental.core.locomotion.mdp` package (mirroring the +stable layout). The stable ``Isaac-Ant-Direct`` registration declares +:class:`~isaaclab_tasks_experimental.core.locomotion.ant.ant_env_warp.AntWarpEnv` +as its ``warp_entry_point``; run the stable ids with ``--frontend warp`` and +``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/ant_env_warp.py similarity index 54% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/ant_env_warp.py index 2bd7390039b7..7d8957e4d9c7 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/ant/ant_env_warp.py @@ -5,13 +5,13 @@ from __future__ import annotations -from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv +from isaaclab_tasks.core.locomotion.ant.ant_direct_env_cfg import AntEnvCfg -from .ant_env_warp_cfg import AntWarpEnvCfg +from isaaclab_tasks_experimental.core.locomotion.locomotion_env_warp import LocomotionWarpEnv class AntWarpEnv(LocomotionWarpEnv): - cfg: AntWarpEnvCfg + cfg: AntEnvCfg - def __init__(self, cfg: AntWarpEnvCfg, render_mode: str | None = None, **kwargs): + def __init__(self, cfg: AntEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py new file mode 100644 index 000000000000..b87d773f7b6d --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/__init__.py @@ -0,0 +1,14 @@ +# 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 + +"""Warp implementation of the stable direct Humanoid task. + +The manager-based MDP twins live in the shared +:mod:`isaaclab_tasks_experimental.core.locomotion.mdp` package (mirroring the +stable layout). The stable ``Isaac-Humanoid-Direct`` registration declares +:class:`~isaaclab_tasks_experimental.core.locomotion.humanoid.humanoid_warp_env.HumanoidWarpEnv` +as its ``warp_entry_point``; run the stable ids with ``--frontend warp`` and +``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/humanoid_warp_env.py similarity index 52% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/humanoid_warp_env.py index 2fd0ac8433e2..2a843510717a 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/humanoid/humanoid_warp_env.py @@ -5,13 +5,13 @@ from __future__ import annotations -from isaaclab_tasks_experimental.direct.locomotion.locomotion_env_warp import LocomotionWarpEnv +from isaaclab_tasks.core.locomotion.humanoid.humanoid_direct_env_cfg import HumanoidEnvCfg -from .humanoid_warp_env_cfg import HumanoidWarpEnvCfg +from isaaclab_tasks_experimental.core.locomotion.locomotion_env_warp import LocomotionWarpEnv class HumanoidWarpEnv(LocomotionWarpEnv): - cfg: HumanoidWarpEnvCfg + cfg: HumanoidEnvCfg - def __init__(self, cfg: HumanoidWarpEnvCfg, render_mode: str | None = None, **kwargs): + def __init__(self, cfg: HumanoidEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py similarity index 99% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py index 62400c4d48fe..3500e1929d6b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py @@ -385,7 +385,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.rpy = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) self.angle_to_target = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device) self.dof_pos_scaled = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) - self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f) + self.env_origins = self.scene.env_origins_pa.warp self.actions_mapped = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) # Initial states and targets diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.pyi similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/__init__.pyi diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/observations.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/observations.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py new file mode 100644 index 000000000000..1a9be132962a --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/__init__.py @@ -0,0 +1,12 @@ +# 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 + +"""Warp MDP twins for the stable Reach tasks. + +There is no separate task registration: run the stable joint-position reach +tasks (``Isaac-Reach-Franka``, ``Isaac-Reach-UR10``) with ``--frontend warp`` +and ``presets=newton_mjwarp``. The IK and OSC variants require controller +action twins that do not exist yet. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi similarity index 88% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi index c277ec98e076..0fe12a032cc9 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/__init__.pyi @@ -4,6 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "modify_reward_weight", "position_command_error", "position_command_error_tanh", "orientation_command_error", @@ -13,4 +14,5 @@ __all__ = [ # override with reach-specific terms below. from isaaclab_experimental.envs.mdp import * # noqa: F401, F403 +from .curriculums import modify_reward_weight from .rewards import orientation_command_error, position_command_error, position_command_error_tanh diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/curriculums.py new file mode 100644 index 000000000000..a25adf0176bd --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/curriculums.py @@ -0,0 +1,98 @@ +# 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 + +"""Warp-first curriculum terms for reach environments.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp +from isaaclab_experimental.managers import CurriculumTermCfg, ManagerTermBase + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +@wp.kernel +def _mark_reset_requested(env_mask: wp.array(dtype=wp.bool), reset_requested: wp.array(dtype=wp.int32)): + env_id = wp.tid() + if env_mask[env_id]: + wp.atomic_max(reset_requested, 0, 1) + + +@wp.kernel +def _update_reward_weight( + reset_requested: wp.array(dtype=wp.int32), + common_step_counter: int, + num_steps: int, + target_weight: float, + term_weight: wp.array(dtype=wp.float32), + out: wp.array(dtype=wp.float32), +): + if reset_requested[0] != 0: + if common_step_counter > num_steps: + term_weight[0] = target_weight + out[0] = term_weight[0] + reset_requested[0] = 0 + + +class modify_reward_weight(ManagerTermBase): + """Update one reward weight after a reset reaches the configured step threshold.""" + + # The Python-owned common step counter is intentionally read in eager mode. + # A future captured path should first move this counter to persistent device storage. + _warp_capturable = False + + def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv): + """Initialize persistent device state. + + Args: + cfg: Curriculum term configuration. + env: Environment containing the reward manager. + """ + super().__init__(cfg, env) + self._term_weight_wp = env.reward_manager.get_term_weight_wp(cfg.params["term_name"]) + self._reset_requested_wp = wp.zeros(1, dtype=wp.int32, device=self.device) + + def __call__( + self, + env: ManagerBasedRLEnv, + env_mask: wp.array(dtype=wp.bool), + out: wp.array(dtype=wp.float32), + term_name: str, + weight: float, + num_steps: int, + ) -> None: + """Update the reward weight when at least one environment resets. + + Args: + env: Environment containing the common step counter. + env_mask: Boolean mask selecting resetting environments. + out: Persistent one-element output containing the current weight. + term_name: Reward term name, resolved during initialization. + weight: Target reward weight. + num_steps: Step threshold after which the target weight is applied. + """ + del term_name + wp.launch( + _mark_reset_requested, + dim=self.num_envs, + inputs=[env_mask, self._reset_requested_wp], + device=self.device, + ) + wp.launch( + _update_reward_weight, + dim=1, + inputs=[ + self._reset_requested_wp, + env.common_step_counter, + num_steps, + weight, + self._term_weight_wp, + out, + ], + device=self.device, + ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py similarity index 81% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py index 7435752528d6..51a828177a42 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reach/mdp/rewards.py @@ -6,7 +6,6 @@ """Warp-first reward terms for the reach task. All functions follow the ``func(env, out, **params) -> None`` signature. -Command tensors are cached as zero-copy warp views on first call. """ from __future__ import annotations @@ -50,12 +49,7 @@ def _position_command_error_kernel( def position_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: """Penalize tracking of the position error using L2-norm.""" - fn = position_command_error asset: Articulation = env.scene[asset_cfg.name] - 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 wp.launch( kernel=_position_command_error_kernel, dim=env.num_envs, @@ -63,7 +57,7 @@ def position_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, asset.data.body_pos_w.warp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), asset_cfg.body_ids[0], out, ], @@ -101,12 +95,7 @@ def position_command_error_tanh( env: ManagerBasedRLEnv, out, std: float, command_name: str, asset_cfg: SceneEntityCfg ) -> None: """Reward tracking of the position using the tanh kernel.""" - fn = position_command_error_tanh asset: Articulation = env.scene[asset_cfg.name] - 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 wp.launch( kernel=_position_command_error_tanh_kernel, dim=env.num_envs, @@ -114,7 +103,7 @@ def position_command_error_tanh( asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, asset.data.body_pos_w.warp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), asset_cfg.body_ids[0], 1.0 / std, out, @@ -152,15 +141,16 @@ def _orientation_command_error_kernel( def orientation_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: """Penalize tracking orientation error using shortest path.""" - fn = orientation_command_error asset: Articulation = env.scene[asset_cfg.name] - 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 wp.launch( kernel=_orientation_command_error_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w.warp, asset.data.body_quat_w.warp, fn._cmd_wp, asset_cfg.body_ids[0], out], + inputs=[ + asset.data.root_quat_w.warp, + asset.data.body_quat_w.warp, + env.command_manager.get_command_wp(command_name), + asset_cfg.body_ids[0], + out, + ], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/__init__.py similarity index 75% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/__init__.py index 6cd56351b6e5..48bb97b31113 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/__init__.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/__init__.py @@ -3,4 +3,4 @@ # # SPDX-License-Identifier: BSD-3-Clause -from .reach import * +"""Warp-first in-hand reorientation task implementations.""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py similarity index 98% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py index 2082d470ff9b..59a130a8b63e 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/inhand_manipulation_warp_env.py @@ -20,7 +20,7 @@ from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane if TYPE_CHECKING: - from isaaclab_tasks_experimental.direct.allegro_hand.allegro_hand_warp_env_cfg import AllegroHandWarpEnvCfg + from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg @wp.kernel @@ -538,10 +538,10 @@ def rotation_distance(object_rot: wp.quatf, target_rot: wp.quatf) -> wp.float32: class InHandManipulationWarpEnv(DirectRLEnvWarp): - cfg: AllegroHandWarpEnvCfg # | ShadowHandWarpEnvCfg + cfg: AllegroHandEnvCfg # | ShadowHandEnvCfg - # def __init__(self, cfg: AllegroHandWarpEnvCfg | ShadowHandWarpEnvCfg, render_mode: str | None = None, **kwargs): - def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, **kwargs): + # def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | None = None, **kwargs): + def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwargs): super().__init__(cfg, render_mode, **kwargs) # --------------------------------------------------------------------- @@ -582,8 +582,7 @@ def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, * self.y_unit_vec = wp.vec3f(0.0, 1.0, 0.0) self.z_unit_vec = wp.vec3f(0.0, 0.0, 1.0) - # Per-env origins (Warp view for kernels; Torch env uses `self.scene.env_origins` directly). - self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f) + self.env_origins = self.scene.env_origins_pa.warp # --------------------------------------------------------------------- # Warp buffers @@ -667,6 +666,7 @@ def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, * self.torch_reset_terminated = wp.to_torch(self.reset_terminated) self.torch_reset_time_outs = wp.to_torch(self.reset_time_outs) self.torch_episode_length_buf = self.episode_length_buf # already a torch tensor via wp.to_torch + self.torch_consecutive_successes = wp.to_torch(self.consecutive_successes) def _setup_scene(self): # add hand, in-hand object, and goal object @@ -771,7 +771,7 @@ def _get_rewards(self) -> None: if "log" not in self.extras: self.extras["log"] = dict() # .mean() cannot be called here as it causes problems on stream - self.extras["log"]["consecutive_successes"] = wp.to_torch(self.consecutive_successes) + self.extras["log"]["consecutive_successes"] = self.torch_consecutive_successes # Reset goals for envs that reached the target (mask is `reset_goal_buf`). # This avoids Torch-side index extraction and keeps the step graphable. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py new file mode 100644 index 000000000000..f6510ad92c91 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/__init__.py @@ -0,0 +1,13 @@ +# 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 + +"""Warp MDP twins for the stable velocity locomotion tasks. + +There is no separate task registration: run the stable flat velocity ids +(e.g. ``Isaac-Velocity-Flat-AnymalD``) with ``--frontend warp`` and +``presets=newton_mjwarp``. Rough-terrain variants remain unsupported until +:class:`~isaaclab.terrains.TerrainImporter` gains Warp APIs (no +``height_scan`` twin). +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.py similarity index 100% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.py diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi similarity index 97% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi index ae4b6d7af547..e76bfc598bff 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi @@ -5,7 +5,7 @@ __all__ = [ "terrain_levels_vel", - "feet_air_time", + "feet_air_time", "feet_air_time_positive_biped", "feet_slide", "stand_still_joint_deviation_l1", diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py new file mode 100644 index 000000000000..f59c34696154 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/curriculums.py @@ -0,0 +1,127 @@ +# 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 + +"""Warp-first curriculum terms for velocity locomotion environments.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp +from isaaclab_experimental.managers import CurriculumTermCfg, ManagerTermBase, SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.assets import Articulation + from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.terrains import TerrainImporter + + +@wp.kernel +def _compute_terrain_level_updates( + env_mask: wp.array(dtype=wp.bool), + root_pos_w: wp.array(dtype=wp.vec3f), + env_origins: wp.array(dtype=wp.vec3f), + command: wp.array(dtype=wp.float32, ndim=2), + terrain_half_length: float, + max_episode_length_s: float, + move_up: wp.array(dtype=wp.bool), + move_down: wp.array(dtype=wp.bool), +): + env_id = wp.tid() + if env_mask[env_id]: + position_delta = root_pos_w[env_id] - env_origins[env_id] + distance = wp.sqrt(position_delta[0] * position_delta[0] + position_delta[1] * position_delta[1]) + command_distance = ( + wp.sqrt(command[env_id, 0] * command[env_id, 0] + command[env_id, 1] * command[env_id, 1]) + * max_episode_length_s + * 0.5 + ) + should_move_up = distance > terrain_half_length + move_up[env_id] = should_move_up + move_down[env_id] = (distance < command_distance) and (not should_move_up) + else: + move_up[env_id] = False + move_down[env_id] = False + + +@wp.kernel +def _compute_terrain_level_mean( + terrain_levels: wp.array(dtype=wp.int64), + scale: float, + out: wp.array(dtype=wp.float32), +): + env_id = wp.tid() + wp.atomic_add(out, 0, wp.float32(terrain_levels[env_id]) * scale) + + +class terrain_levels_vel(ManagerTermBase): + """Update terrain levels from commanded locomotion progress using a boolean environment mask.""" + + _curriculum_mask_native = True + + def __init__(self, cfg: CurriculumTermCfg, env: ManagerBasedRLEnv): + """Initialize persistent views and decision buffers. + + Args: + cfg: Curriculum term configuration. + env: Environment containing the robot and terrain. + """ + super().__init__(cfg, env) + asset_cfg = cfg.params.get("asset_cfg", SceneEntityCfg("robot")) + self._asset: Articulation = env.scene[asset_cfg.name] + self._terrain: TerrainImporter = env.scene.terrain + self._root_pos_w_wp = self._asset.data.root_pos_w.warp + self._env_origins_wp = env.env_origins_wp + self._command_wp = env.command_manager.get_command_wp("base_velocity") + self._terrain_levels_wp = self._terrain.terrain_levels_pa.warp + self._move_up_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self._move_down_wp = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self._terrain_half_length = float(self._terrain.cfg.terrain_generator.size[0]) * 0.5 + self._max_episode_length_s = float(env.max_episode_length_s) + self._mean_scale = 1.0 / self.num_envs + + def __call__( + self, + env: ManagerBasedRLEnv, + env_mask: wp.array(dtype=wp.bool), + out: wp.array(dtype=wp.float32), + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> None: + """Update selected terrain levels and write the global mean level. + + Args: + env: Environment containing the per-environment random state. + env_mask: Boolean Warp mask selecting environments to update. + out: Persistent scalar output for the mean terrain level. + asset_cfg: Robot scene entity configuration. + """ + del asset_cfg + wp.launch( + kernel=_compute_terrain_level_updates, + dim=self.num_envs, + inputs=[ + env_mask, + self._root_pos_w_wp, + self._env_origins_wp, + self._command_wp, + self._terrain_half_length, + self._max_episode_length_s, + self._move_up_wp, + self._move_down_wp, + ], + device=self.device, + ) + self._terrain.update_env_origins_mask( + env_mask, + self._move_up_wp, + self._move_down_wp, + env.rng_state_wp, + ) + wp.launch( + kernel=_compute_terrain_level_mean, + dim=self.num_envs, + inputs=[self._terrain_levels_wp, self._mean_scale, out], + device=self.device, + ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py similarity index 81% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py index d7e11c517c7d..1b1744c4078b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py @@ -6,8 +6,6 @@ """Warp-first reward functions for the velocity locomotion environment. All functions follow the ``func(env, out, **params) -> None`` signature. -Cross-manager torch tensors (contact sensor, commands) are cached as zero-copy -warp views on first call via ``wp.from_torch``. """ from __future__ import annotations @@ -50,14 +48,7 @@ def _feet_air_time_kernel( def feet_air_time(env: ManagerBasedRLEnv, out, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float) -> None: """Reward long steps taken by the feet using L2-kernel.""" - fn = feet_air_time contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - # Cache command bridge (persistent pointer) - if not getattr(fn, "_is_warmed_up", False) 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 - fn._is_warmed_up = True # Newton contact sensor returns persistent wp.arrays — use directly, no wp.from_torch needed first_contact = contact_sensor.compute_first_contact(env.step_dt) wp.launch( @@ -67,7 +58,7 @@ def feet_air_time(env: ManagerBasedRLEnv, out, command_name: str, sensor_cfg: Sc contact_sensor.data.last_air_time.warp, first_contact, sensor_cfg.body_ids_wp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), threshold, out, ], @@ -114,13 +105,7 @@ def _feet_air_time_positive_biped_kernel( def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg) -> None: """Reward long steps taken by the feet for bipeds.""" - fn = feet_air_time_positive_biped contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - if not getattr(fn, "_is_warmed_up", False) 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 - fn._is_warmed_up = True wp.launch( kernel=_feet_air_time_positive_biped_kernel, dim=env.num_envs, @@ -128,7 +113,7 @@ def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, contact_sensor.data.current_air_time.warp, contact_sensor.data.current_contact_time.warp, sensor_cfg.body_ids_wp, - fn._cmd_wp, + env.command_manager.get_command_wp(command_name), threshold, out, ], @@ -224,17 +209,17 @@ def track_lin_vel_xy_yaw_frame_exp( env, out, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame.""" - fn = track_lin_vel_xy_yaw_frame_exp asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) 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 - fn._is_warmed_up = True wp.launch( kernel=_track_lin_vel_xy_yaw_frame_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w.warp, asset.data.root_lin_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], + inputs=[ + asset.data.root_quat_w.warp, + asset.data.root_lin_vel_w.warp, + env.command_manager.get_command_wp(command_name), + 1.0 / (std * std), + out, + ], device=env.device, ) @@ -260,17 +245,16 @@ def track_ang_vel_z_world_exp( env, out, command_name: str, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Reward tracking of angular velocity commands (yaw) in world frame.""" - fn = track_ang_vel_z_world_exp asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) 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 - fn._is_warmed_up = True wp.launch( kernel=_track_ang_vel_z_world_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_ang_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], + inputs=[ + asset.data.root_ang_vel_w.warp, + env.command_manager.get_command_wp(command_name), + 1.0 / (std * std), + out, + ], device=env.device, ) @@ -304,16 +288,16 @@ def stand_still_joint_deviation_l1( env, out, command_name: str, command_threshold: float = 0.06, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Penalize offsets from the default joint positions when the command is very small.""" - fn = stand_still_joint_deviation_l1 asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) 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 - fn._is_warmed_up = True wp.launch( kernel=_stand_still_joint_deviation_l1_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, fn._cmd_wp, command_threshold, out], + inputs=[ + asset.data.joint_pos.warp, + asset.data.default_joint_pos.warp, + env.command_manager.get_command_wp(command_name), + command_threshold, + out, + ], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py similarity index 57% rename from source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py rename to source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py index 4e0c32c58e64..19583eb4edeb 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/terminations.py @@ -36,31 +36,24 @@ def terrain_out_of_bounds( env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), distance_buffer: float = 3.0 ) -> None: """Terminate when the actor moves too close to the edge of the terrain.""" - fn = terrain_out_of_bounds - if not getattr(fn, "_is_warmed_up", False): - terrain_type = env.scene.cfg.terrain.terrain_type - if terrain_type == "plane": - fn._is_plane = True - elif terrain_type == "generator": - fn._is_plane = False - terrain_gen_cfg = env.scene.terrain.cfg.terrain_generator - grid_width, grid_length = terrain_gen_cfg.size - n_rows, n_cols = terrain_gen_cfg.num_rows, terrain_gen_cfg.num_cols - border_width = terrain_gen_cfg.border_width - fn._half_width = 0.5 * (n_rows * grid_width + 2 * border_width) - fn._half_height = 0.5 * (n_cols * grid_length + 2 * border_width) - else: - raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") - fn._is_warmed_up = True - - if fn._is_plane: + terrain_type = env.scene.cfg.terrain.terrain_type + if terrain_type == "plane": out.zero_() return + if terrain_type != "generator": + raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") + + terrain_gen_cfg = env.scene.terrain.cfg.terrain_generator + grid_width, grid_length = terrain_gen_cfg.size + n_rows, n_cols = terrain_gen_cfg.num_rows, terrain_gen_cfg.num_cols + border_width = terrain_gen_cfg.border_width + half_width = 0.5 * (n_rows * grid_width + 2 * border_width) + half_height = 0.5 * (n_cols * grid_length + 2 * border_width) asset: Articulation = env.scene[asset_cfg.name] wp.launch( kernel=_terrain_out_of_bounds_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w.warp, fn._half_width, fn._half_height, distance_buffer, out], + inputs=[asset.data.root_pos_w.warp, half_width, half_height, distance_buffer, out], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py deleted file mode 100644 index e7fc129c633f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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 - -""" -Allegro Inhand Manipulation environment. -""" - -import gymnasium as gym - -## -# Register Gym environments. -## - -inhand_task_entry = "isaaclab_tasks_experimental.direct.inhand_manipulation" -stable_agents = "isaaclab_tasks.core.reorient.config.allegro_hand.agents" - -gym.register( - id="Isaac-Reorient-Cube-Allegro-Direct-Warp-v0", - entry_point=f"{inhand_task_entry}.inhand_manipulation_warp_env:InHandManipulationWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.allegro_hand_warp_env_cfg:AllegroHandWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_ppo_cfg:AllegroHandPPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/allegro_hand_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/allegro_hand_warp_env_cfg.py deleted file mode 100644 index 573b4f1d302d..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/allegro_hand/allegro_hand_warp_env_cfg.py +++ /dev/null @@ -1,135 +0,0 @@ -# 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 - - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg # , RigidObjectCfg -from isaaclab.envs import DirectRLEnvCfg -from isaaclab.markers import VisualizationMarkersCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.sim.spawners.materials.physics_materials_cfg import RigidBodyMaterialCfg -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.configclass import configclass - -from isaaclab_assets.robots.allegro import ALLEGRO_HAND_CFG - - -@configclass -class AllegroHandWarpEnvCfg(DirectRLEnvCfg): - # env - decimation = 4 - episode_length_s = 10.0 - action_space = 16 - observation_space = 124 # (full) - state_space = 0 - asymmetric_obs = False - obs_type = "full" - - solver_cfg = MJWarpSolverCfg( - solver="newton", - integrator="implicitfast", - njmax=80, - nconmax=70, - impratio=10.0, - cone="elliptic", - update_data_interval=2, - iterations=100, - ls_iterations=15, - # save_to_mjcf="AllegroHand.xml", - ) - - newton_cfg = NewtonCfg( - solver_cfg=solver_cfg, - num_substeps=2, - debug_mode=False, - ) - # simulation - sim: SimulationCfg = SimulationCfg( - dt=1 / 120, - render_interval=decimation, - physics_material=RigidBodyMaterialCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), - physics=newton_cfg, - ) - # robot - robot_cfg: ArticulationCfg = ALLEGRO_HAND_CFG.replace(prim_path="/World/envs/env_.*/Robot") - - actuated_joint_names = [ - "index_joint_0", - "middle_joint_0", - "ring_joint_0", - "thumb_joint_0", - "index_joint_1", - "index_joint_2", - "index_joint_3", - "middle_joint_1", - "middle_joint_2", - "middle_joint_3", - "ring_joint_1", - "ring_joint_2", - "ring_joint_3", - "thumb_joint_1", - "thumb_joint_2", - "thumb_joint_3", - ] - fingertip_body_names = [ - "index_link_3", - "middle_link_3", - "ring_link_3", - "thumb_link_3", - ] - - # in-hand object - object_cfg: ArticulationCfg = ArticulationCfg( - prim_path="/World/envs/env_.*/object", - spawn=sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - mass_props=sim_utils.MassPropertiesCfg(density=400.0), - scale=(1.2, 1.2, 1.2), - ), - # FIXME: it does seem to be a bug for ArticulationCfg for handling empty joint list - init_state=ArticulationCfg.InitialStateCfg( - pos=(0.0, -0.17, 0.565), rot=(0.0, 0.0, 0.0, 1.0), joint_pos={}, joint_vel={} - ), - actuators={}, - articulation_root_prim_path="", - ) - # goal object - goal_object_cfg: VisualizationMarkersCfg = VisualizationMarkersCfg( - prim_path="/Visuals/goal_marker", - markers={ - "goal": sim_utils.UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", - scale=(1.2, 1.2, 1.2), - ) - }, - ) - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=8192, env_spacing=0.75, replicate_physics=True, clone_in_fabric=True - ) - # reset - reset_position_noise = 0.01 # range of position at reset - reset_dof_pos_noise = 0.2 # range of dof pos at reset - reset_dof_vel_noise = 0.0 # range of dof vel at reset - # reward scales - dist_reward_scale = -10.0 - rot_reward_scale = 1.0 - rot_eps = 0.1 - action_penalty_scale = -0.0002 - reach_goal_bonus = 250 - fall_penalty = 0 - fall_dist = 0.24 - vel_obs_scale = 0.2 - success_tolerance = 0.2 - max_consecutive_success = 0 - av_factor = 0.1 - act_moving_average = 1.0 - force_torque_obs_scale = 10.0 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py deleted file mode 100644 index 33d5ac926387..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# 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 - -""" -Ant locomotion environment. -""" - -import gymnasium as gym - -## -# Register Gym environments. -## - -stable_agents = "isaaclab_tasks.core.locomotion.ant.agents" - -gym.register( - id="Isaac-Ant-Direct-Warp-v0", - entry_point=f"{__name__}.ant_env_warp:AntWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.ant_env_warp_cfg:AntWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_ppo_cfg:AntDirectPPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py deleted file mode 100644 index 6e8b7fb81a95..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/ant/ant_env_warp_cfg.py +++ /dev/null @@ -1,82 +0,0 @@ -# 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 - -from __future__ import annotations - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg -from isaaclab.envs import DirectRLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -from isaaclab_assets import ANT_CFG - - -@configclass -class AntWarpEnvCfg(DirectRLEnvCfg): - # env - episode_length_s = 15.0 - decimation = 2 - action_scale = 0.5 - action_space = 8 - observation_space = 36 - state_space = 0 - - solver_cfg = MJWarpSolverCfg( - njmax=45, - nconmax=25, - cone="pyramidal", - integrator="implicitfast", - impratio=1, - ) - newton_cfg = NewtonCfg( - solver_cfg=solver_cfg, - num_substeps=1, - debug_mode=False, - use_cuda_graph=True, - ) - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=newton_cfg) - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="average", - restitution_combine_mode="average", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - debug_vis=False, - ) - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True - ) - - # robot - robot: ArticulationCfg = ANT_CFG.replace(prim_path="/World/envs/env_.*/Robot") - joint_gears: list = [15, 15, 15, 15, 15, 15, 15, 15] - - heading_weight: float = 0.5 - up_weight: float = 0.1 - - energy_cost_scale: float = 0.05 - actions_cost_scale: float = 0.005 - alive_reward_scale: float = 0.5 - dof_vel_scale: float = 0.2 - - death_cost: float = -2.0 - termination_height: float = 0.31 - - angular_velocity_scale: float = 1.0 - contact_force_scale: float = 0.1 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py deleted file mode 100644 index a7fc9492fff6..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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 - -""" -Cartpole balancing environment. -""" - -import gymnasium as gym - -## -# Register Gym environments. -## - -stable_agents = "isaaclab_tasks.core.cartpole.agents" - -gym.register( - id="Isaac-Cartpole-Direct-Warp-v0", - entry_point=f"{__name__}.cartpole_warp_env:CartpoleWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_warp_env_cfg:CartpoleWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_ppo_cfg:CartpoleDirectPPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{stable_agents}:sb3_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py deleted file mode 100644 index 2a99e1cff3e9..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env_cfg.py +++ /dev/null @@ -1,66 +0,0 @@ -# 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 - -from __future__ import annotations - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.assets import ArticulationCfg -from isaaclab.envs import DirectRLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from isaaclab_assets.robots.cartpole import CARTPOLE_CFG - - -@configclass -class CartpoleWarpEnvCfg(DirectRLEnvCfg): - # env - decimation = 2 - episode_length_s = 5.0 - action_scale = 100.0 # [N] - action_space = 1 - observation_space = 4 - state_space = 0 - - solver_cfg = MJWarpSolverCfg( - njmax=5, - nconmax=3, - cone="pyramidal", - integrator="implicitfast", - impratio=1, - ) - - newton_cfg = NewtonCfg( - solver_cfg=solver_cfg, - num_substeps=1, - debug_mode=False, - use_cuda_graph=True, - ) - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=newton_cfg) - - # robot - robot_cfg: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="/World/envs/env_.*/Robot") - cart_dof_name = "slider_to_cart" - pole_dof_name = "cart_to_pole" - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True - ) - - # reset - max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] - initial_pole_angle_range = [-0.25, 0.25] # the range in which the pole angle is sampled from on reset [x pi rad] - - # reward scales - rew_scale_alive = 1.0 - rew_scale_terminated = -2.0 - rew_scale_pole_pos = -1.0 - rew_scale_cart_vel = -0.01 - rew_scale_pole_vel = -0.005 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py deleted file mode 100644 index c38d7dd955e5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -# 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 - -""" -Humanoid locomotion environment. -""" - -import gymnasium as gym - -## -# Register Gym environments. -## - -stable_agents = "isaaclab_tasks.core.locomotion.humanoid.agents" - -gym.register( - id="Isaac-Humanoid-Direct-Warp-v0", - entry_point=f"{__name__}.humanoid_warp_env:HumanoidWarpEnv", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.humanoid_warp_env_cfg:HumanoidWarpEnvCfg", - "rl_games_cfg_entry_point": f"{stable_agents}:rl_games_direct_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{stable_agents}.rsl_rl_ppo_cfg:HumanoidDirectPPORunnerCfg", - "skrl_cfg_entry_point": f"{stable_agents}:skrl_direct_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py deleted file mode 100644 index d7fa473eda38..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/humanoid/humanoid_warp_env_cfg.py +++ /dev/null @@ -1,105 +0,0 @@ -# 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 - -from __future__ import annotations - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg -from isaaclab.envs import DirectRLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -from isaaclab_assets import HUMANOID_CFG - - -@configclass -class HumanoidWarpEnvCfg(DirectRLEnvCfg): - # env - episode_length_s = 15.0 - decimation = 2 - action_scale = 1.0 - action_space = 21 - observation_space = 75 - state_space = 0 - - solver_cfg = MJWarpSolverCfg( - njmax=80, - nconmax=25, - cone="pyramidal", - integrator="implicitfast", - update_data_interval=2, - impratio=1, - ) - newton_cfg = NewtonCfg( - solver_cfg=solver_cfg, - num_substeps=2, - debug_mode=False, - use_cuda_graph=True, - ) - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=newton_cfg) - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="average", - restitution_combine_mode="average", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - debug_vis=False, - ) - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True - ) - - # robot - robot: ArticulationCfg = HUMANOID_CFG.replace(prim_path="/World/envs/env_.*/Robot") - joint_gears: list = [ - 67.5000, # left_upper_arm - 67.5000, # left_upper_arm - 45.0000, # left_lower_arm - 67.5000, # lower_waist - 67.5000, # lower_waist - 67.5000, # pelvis - 45.0000, # left_thigh: x - 135.0000, # left_thigh: y - 45.0000, # left_thigh: z - 90.0000, # left_knee - 22.5, # left_foot - 22.5, # left_foot - 45.0000, # right_thigh: x - 135.0000, # right_thigh: y - 45.0000, # right_thigh: z - 90.0000, # right_knee - 22.5, # right_foot - 22.5, # right_foot - 67.5000, # right_upper_arm - 67.5000, # right_upper_arm - 45.0000, # right_lower_arm - ] - - heading_weight: float = 0.5 - up_weight: float = 0.1 - - energy_cost_scale: float = 0.05 - actions_cost_scale: float = 0.01 - alive_reward_scale: float = 2.0 - dof_vel_scale: float = 0.1 - - death_cost: float = -1.0 - termination_height: float = 0.8 - - angular_velocity_scale: float = 0.25 - contact_force_scale: float = 0.01 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/__init__.py deleted file mode 100644 index 460a30569089..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# 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 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/__init__.py deleted file mode 100644 index 460a30569089..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# 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 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py deleted file mode 100644 index 7f23883e6332..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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 - -"""Experimental registrations for manager-based tasks. - -We intentionally only register new Gym IDs pointing at experimental entry points. -Task definitions (configs/mdp) remain in `isaaclab_tasks` to avoid duplication. -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py deleted file mode 100644 index 79c13e2aa8f4..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# 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 - -"""Classic environments for control. - -These environments are based on the MuJoCo environments provided by OpenAI. - -Reference: - https://github.com/openai/gym/tree/master/gym/envs/mujoco -""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py deleted file mode 100644 index af513e675f03..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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 - -""" -Ant locomotion environment (experimental manager-based entry point). -""" - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.locomotion.ant import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Ant-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.ant_env_cfg:AntEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AntPPORunnerCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py deleted file mode 100644 index 66a9131011f3..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/ant/ant_env_cfg.py +++ /dev/null @@ -1,197 +0,0 @@ -# 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 - -# Ant reuses humanoid's experimental MDP (mirrors stable pattern). -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab_assets.robots.ant import ANT_CFG # isort: skip - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with an ant robot.""" - - # terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="average", - restitution_combine_mode="average", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - debug_vis=False, - ) - - # robot - robot = ANT_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DistantLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg(asset_name="robot", joint_names=[".*"], scale=7.5) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for the policy.""" - - base_height = ObsTerm(func=mdp.base_pos_z) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel) - base_yaw_roll = ObsTerm(func=mdp.base_yaw_roll) - base_angle_to_target = ObsTerm(func=mdp.base_angle_to_target, params={"target_pos": (1000.0, 0.0, 0.0)}) - base_up_proj = ObsTerm(func=mdp.base_up_proj) - base_heading_proj = ObsTerm(func=mdp.base_heading_proj, params={"target_pos": (1000.0, 0.0, 0.0)}) - joint_pos_norm = ObsTerm(func=mdp.joint_pos_limit_normalized) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel, scale=0.2) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "position_range": (-0.2, 0.2), - "velocity_range": (-0.1, 0.1), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Reward for moving forward - progress = RewTerm(func=mdp.progress_reward, weight=1.0, params={"target_pos": (1000.0, 0.0, 0.0)}) - # (2) Stay alive bonus - alive = RewTerm(func=mdp.is_alive, weight=0.5) - # (3) Reward for non-upright posture - upright = RewTerm(func=mdp.upright_posture_bonus, weight=0.1, params={"threshold": 0.93}) - # (4) Reward for moving in the right direction - move_to_target = RewTerm( - func=mdp.move_to_target_bonus, weight=0.5, params={"threshold": 0.8, "target_pos": (1000.0, 0.0, 0.0)} - ) - # (5) Penalty for large action commands - action_l2 = RewTerm(func=mdp.action_l2, weight=-0.005) - # (6) Penalty for energy consumption - energy = RewTerm(func=mdp.power_consumption, weight=-0.05, params={"gear_ratio": {".*": 15.0}}) - # (7) Penalty for reaching close to joint limits - joint_pos_limits = RewTerm( - func=mdp.joint_pos_limits_penalty_ratio, weight=-0.1, params={"threshold": 0.99, "gear_ratio": {".*": 15.0}} - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Terminate if the episode length is exceeded - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Terminate if the robot falls - torso_height = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.31}) - - -## -# Environment configuration -## - - -@configclass -class AntEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the MuJoCo-style Ant walking environment.""" - - # Simulation settings - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=38, - nconmax=15, - ls_iterations=10, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=5.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 16.0 - # simulation settings - self.sim.dt = 1 / 120.0 - self.sim.render_interval = self.decimation - # default friction material - self.sim.physics_material.static_friction = 1.0 - self.sim.physics_material.dynamic_friction = 1.0 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py deleted file mode 100644 index 9f46551de255..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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 - -""" -Cartpole balancing environment (experimental manager-based entry point). -""" - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.cartpole import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Cartpole-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_env_cfg:CartpoleEnvCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_manager_ppo_cfg:CartpolePPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py deleted file mode 100644 index ae187a593a16..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/cartpole_env_cfg.py +++ /dev/null @@ -1,199 +0,0 @@ -# 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 - -import math - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab_assets.robots.cartpole import CARTPOLE_CFG # isort:skip - - -## -# Scene definition -## - - -@configclass -class CartpoleSceneCfg(InteractiveSceneCfg): - """Configuration for a cart-pole scene.""" - - # ground plane - # ground = AssetBaseCfg( - # prim_path="/World/ground", - # spawn=sim_utils.GroundPlaneCfg(size=(100.0, 100.0)), - # ) - - # cartpole - robot: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - dome_light = AssetBaseCfg( - prim_path="/World/DomeLight", - spawn=sim_utils.DomeLightCfg(color=(0.9, 0.9, 0.9), intensity=500.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg(asset_name="robot", joint_names=["slider_to_cart"], scale=100.0) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - joint_pos_rel = ObsTerm(func=mdp.joint_pos_rel) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel) - - def __post_init__(self) -> None: - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - # reset - reset_cart_position = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"]), - "position_range": (-1.0, 1.0), - "velocity_range": (-0.5, 0.5), - }, - ) - - reset_pole_position = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"]), - "position_range": (-0.25 * math.pi, 0.25 * math.pi), - "velocity_range": (-0.25 * math.pi, 0.25 * math.pi), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Constant running reward - alive = RewTerm(func=mdp.is_alive, weight=1.0) - # (2) Failure penalty - terminating = RewTerm(func=mdp.is_terminated, weight=-2.0) - # (3) Primary task: keep pole upright - pole_pos = RewTerm( - func=mdp.joint_pos_target_l2, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"]), "target": 0.0}, - ) - # (4) Shaping tasks: lower cart velocity - cart_vel = RewTerm( - func=mdp.joint_vel_l1, - weight=-0.01, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"])}, - ) - # (5) Shaping tasks: lower pole angular velocity - pole_vel = RewTerm( - func=mdp.joint_vel_l1, - weight=-0.005, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["cart_to_pole"])}, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Time out - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Cart out of bounds - cart_out_of_bounds = DoneTerm( - func=mdp.joint_pos_out_of_manual_limit, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"]), "bounds": (-3.0, 3.0)}, - ) - - -## -# Environment configuration -## - - -@configclass -class CartpoleEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the cartpole environment.""" - - # Scene settings - scene: CartpoleSceneCfg = CartpoleSceneCfg(num_envs=4096, env_spacing=4.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - events: EventCfg = EventCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - # Simulation settings - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=5, - nconmax=3, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - use_cuda_graph=True, - ) - ) - - # Post initialization - def __post_init__(self) -> None: - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 5 - # viewer settings - self.viewer.eye = (8.0, 0.0, 5.0) - # simulation settings - self.sim.dt = 1 / 120 - self.sim.render_interval = self.decimation diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py deleted file mode 100644 index 2678ef03de65..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py +++ /dev/null @@ -1,45 +0,0 @@ -# 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 - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import warp as wp -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.utils.warp.utils import wrap_to_pi - -if TYPE_CHECKING: - from isaaclab.assets import Articulation - from isaaclab.envs import ManagerBasedRLEnv - - -@wp.kernel -def _joint_pos_target_l2_kernel( - joint_pos: wp.array(dtype=wp.float32, ndim=2), - joint_mask: wp.array(dtype=wp.bool), - out: wp.array(dtype=wp.float32), - target: float, -): - i = wp.tid() - s = float(0.0) - for j in range(joint_pos.shape[1]): - if joint_mask[j]: - a = wrap_to_pi(joint_pos[i, j]) - d = a - target - s += d * d - out[i] = s - - -def joint_pos_target_l2(env: ManagerBasedRLEnv, out, target: float, asset_cfg: SceneEntityCfg) -> None: - """Penalize joint position deviation from a target value. Writes into ``out``.""" - asset: Articulation = env.scene[asset_cfg.name] - assert asset.data.joint_pos.warp.shape[1] == asset_cfg.joint_mask.shape[0] - wp.launch( - kernel=_joint_pos_target_l2_kernel, - dim=env.num_envs, - inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, out, target], - device=env.device, - ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py deleted file mode 100644 index 0bcd020bcd4c..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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 - -""" -Humanoid locomotion environment (experimental manager-based entry point). -""" - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.locomotion.humanoid import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Humanoid-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.humanoid_env_cfg:HumanoidEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:HumanoidPPORunnerCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_manager_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py deleted file mode 100644 index 9de689ffb98e..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/humanoid_env_cfg.py +++ /dev/null @@ -1,232 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -import isaaclab.sim as sim_utils -from isaaclab.assets import AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.classic.humanoid.mdp as mdp - -from isaaclab_assets.robots.humanoid import HUMANOID_CFG # isort:skip - - -## -# Scene definition -## - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with a humanoid robot.""" - - # terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg(static_friction=1.0, dynamic_friction=1.0, restitution=0.0), - debug_vis=False, - ) - - # robot - robot = HUMANOID_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DistantLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0), - ) - - -## -# MDP settings -## - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_effort = mdp.JointEffortActionCfg( - asset_name="robot", - joint_names=[".*"], - scale={ - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - }, - ) - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for the policy.""" - - base_height = ObsTerm(func=mdp.base_pos_z) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel, scale=0.25) - base_yaw_roll = ObsTerm(func=mdp.base_yaw_roll) - base_angle_to_target = ObsTerm(func=mdp.base_angle_to_target, params={"target_pos": (1000.0, 0.0, 0.0)}) - base_up_proj = ObsTerm(func=mdp.base_up_proj) - base_heading_proj = ObsTerm(func=mdp.base_heading_proj, params={"target_pos": (1000.0, 0.0, 0.0)}) - joint_pos_norm = ObsTerm(func=mdp.joint_pos_limit_normalized) - joint_vel_rel = ObsTerm(func=mdp.joint_vel_rel, scale=0.1) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = False - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_offset, - mode="reset", - params={ - "position_range": (-0.2, 0.2), - "velocity_range": (-0.1, 0.1), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # (1) Reward for moving forward - progress = RewTerm(func=mdp.progress_reward, weight=1.0, params={"target_pos": (1000.0, 0.0, 0.0)}) - # (2) Stay alive bonus - alive = RewTerm(func=mdp.is_alive, weight=2.0) - # (3) Reward for non-upright posture - upright = RewTerm(func=mdp.upright_posture_bonus, weight=0.1, params={"threshold": 0.93}) - # (4) Reward for moving in the right direction - move_to_target = RewTerm( - func=mdp.move_to_target_bonus, weight=0.5, params={"threshold": 0.8, "target_pos": (1000.0, 0.0, 0.0)} - ) - # (5) Penalty for large action commands - action_l2 = RewTerm(func=mdp.action_l2, weight=-0.01) - # (6) Penalty for energy consumption - energy = RewTerm( - func=mdp.power_consumption, - weight=-0.005, - params={ - "gear_ratio": { - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - } - }, - ) - # (7) Penalty for reaching close to joint limits - joint_pos_limits = RewTerm( - func=mdp.joint_pos_limits_penalty_ratio, - weight=-0.25, - params={ - "threshold": 0.98, - "gear_ratio": { - ".*_waist.*": 67.5, - ".*_upper_arm.*": 67.5, - "pelvis": 67.5, - ".*_lower_arm": 45.0, - ".*_thigh:0": 45.0, - ".*_thigh:1": 135.0, - ".*_thigh:2": 45.0, - ".*_shin": 90.0, - ".*_foot.*": 22.5, - }, - }, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - # (1) Terminate if the episode length is exceeded - time_out = DoneTerm(func=mdp.time_out, time_out=True) - # (2) Terminate if the robot falls - torso_height = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.8}) - - -@configclass -class HumanoidEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the MuJoCo-style Humanoid walking environment.""" - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=5.0, clone_in_fabric=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.episode_length_s = 16.0 - # simulation settings - self.sim: SimulationCfg = SimulationCfg( - dt=1 / 120.0, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=80, - nconmax=25, - ls_iterations=15, - cone="pyramidal", - update_data_interval=2, - impratio=1, - integrator="implicitfast", - ), - num_substeps=2, - debug_mode=False, - ), - ) - # self.sim.dt = 1 / 120.0 - self.sim.render_interval = self.decimation - # default friction material - self.sim.physics_material.static_friction = 1.0 - self.sim.physics_material.dynamic_friction = 1.0 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py deleted file mode 100644 index 0660d38f0658..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# 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 - -"""Locomotion experimental task registrations (manager-based).""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/__init__.py deleted file mode 100644 index 26f3257daef4..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# 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 - -"""Configurations for velocity-based locomotion environments.""" - -# We leave this file empty since we don't want to expose any configs in this package directly. -# We still need this file to import the "config" module in the parent package. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/__init__.py deleted file mode 100644 index 5b79532a456e..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/__init__.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.contrib.velocity.config.a1 import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-UnitreeA1-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeA1FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-UnitreeA1-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeA1FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeA1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py deleted file mode 100644 index f213286d0cfb..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/flat_env_cfg.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import UnitreeA1RoughEnvCfg - - -@configclass -class UnitreeA1FlatEnvCfg(UnitreeA1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=30, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # override rewards - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no height scan - # self.scene.height_scanner = None - # self.observations.policy.height_scan = None - # no terrain curriculum - self.curriculum.terrain_levels = None - - -class UnitreeA1FlatEnvCfg_PLAY(UnitreeA1FlatEnvCfg): - def __post_init__(self) -> None: - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py deleted file mode 100644 index ee0f236043a2..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/a1/rough_env_cfg.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - TerminationsCfg, -) - -from isaaclab_assets.robots.unitree import UNITREE_A1_CFG # isort: skip - - -class TerminationsCfg_A1(TerminationsCfg): - base_too_low = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": 0.2}) - - -@configclass -class UnitreeA1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - terminations: TerminationsCfg_A1 = TerminationsCfg_A1() - - def __post_init__(self): - # post init of parent - super().__post_init__() - - self.scene.robot = UNITREE_A1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - - # reduce action scale - self.actions.joint_pos.scale = 0.25 - - # event - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass.params["mass_distribution_params"] = (-1.0, 3.0) - # self.events.add_base_mass.params["asset_cfg"].body_names = "trunk" - self.events.base_external_force_torque.params["asset_cfg"].body_names = "trunk" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - - # rewards - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts.params["sensor_cfg"].body_names = ".*thigh" - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "trunk" - - -@configclass -class UnitreeA1RoughEnvCfg_PLAY(UnitreeA1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/__init__.py deleted file mode 100644 index a6753cfb95ef..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/__init__.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.contrib.velocity.config.anymal_b import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-AnymalB-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalBFlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerCfg", - "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerWithSymmetryCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-AnymalB-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalBFlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerCfg", - "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalBFlatPPORunnerWithSymmetryCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py deleted file mode 100644 index a5c825bf2e98..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/flat_env_cfg.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import AnymalBRoughEnvCfg - - -@configclass -class AnymalBFlatEnvCfg(AnymalBRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=75, - nconmax=15, - cone="elliptic", - impratio=100, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # override rewards - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no height scan - self.scene.height_scanner = None - self.observations.policy.height_scan = None - # no terrain curriculum - self.curriculum.terrain_levels = None - - -class AnymalBFlatEnvCfg_PLAY(AnymalBFlatEnvCfg): - def __post_init__(self) -> None: - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing event - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py deleted file mode 100644 index 3a7d80af470f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_b/rough_env_cfg.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets import ANYMAL_B_CFG # isort: skip - - -@configclass -class AnymalBRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = ANYMAL_B_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalBRoughEnvCfg_PLAY(AnymalBRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/__init__.py deleted file mode 100644 index 9c9c2530e118..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/__init__.py +++ /dev/null @@ -1,39 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.contrib.velocity.config.anymal_c import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-AnymalC-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", - "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerWithSymmetryCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-AnymalC-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalCFlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerCfg", - "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalCFlatPPORunnerWithSymmetryCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_flat_ppo_cfg.yaml", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py deleted file mode 100644 index 3cc348a97675..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/flat_env_cfg.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import AnymalCRoughEnvCfg - - -@configclass -class AnymalCFlatEnvCfg(AnymalCRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=120, - nconmax=15, - cone="elliptic", - impratio=100, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # override rewards - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - # change terrain to flat - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - # no terrain curriculum - self.curriculum.terrain_levels = None - - -class AnymalCFlatEnvCfg_PLAY(AnymalCFlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py deleted file mode 100644 index 1d2f2676c64a..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_c/rough_env_cfg.py +++ /dev/null @@ -1,40 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets.robots.anymal import ANYMAL_C_CFG # isort: skip - - -@configclass -class AnymalCRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # switch robot to anymal-c - self.scene.robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.robot.actuators["legs"].armature = 0.01 - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalCRoughEnvCfg_PLAY(AnymalCRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py deleted file mode 100644 index 3722de03ed4e..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/__init__.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.anymal_d import agents - -## -# Register Gym environments. -## - -# Rough env disabled: requires isaaclab_physx which is not yet available on dev/newton. -# The package exists on upstream/develop (commit 308400f1d35) but has not been merged. -# Re-enable once dev/newton picks up isaaclab_physx. -# gym.register( -# id="Isaac-Velocity-Rough-AnymalD-Warp-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalDRoughEnvCfg", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDRoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - -# gym.register( -# id="Isaac-Velocity-Rough-AnymalD-Warp-Play-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:AnymalDRoughEnvCfg_PLAY", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDRoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - -gym.register( - id="Isaac-Velocity-Flat-AnymalD-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-AnymalD-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:AnymalDFlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:AnymalDFlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py deleted file mode 100644 index 702443e815e2..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/flat_env_cfg.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import AnymalDRoughEnvCfg - - -@configclass -class AnymalDFlatEnvCfg(AnymalDRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=25, - cone="elliptic", - impratio=100.0, - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.rewards.flat_orientation_l2.weight = -5.0 - self.rewards.dof_torques_l2.weight = -2.5e-5 - self.rewards.feet_air_time.weight = 0.5 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - - -class AnymalDFlatEnvCfg_PLAY(AnymalDFlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py deleted file mode 100644 index 71f49edd2e68..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py +++ /dev/null @@ -1,37 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets.robots.anymal import ANYMAL_D_CFG # isort: skip - - -@configclass -class AnymalDRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - - -@configclass -class AnymalDRoughEnvCfg_PLAY(AnymalDRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/__init__.py deleted file mode 100644 index e3af41334541..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.cassie import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-Cassie-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:CassieFlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieFlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-Cassie-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:CassieFlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CassieFlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py deleted file mode 100644 index 818c7aa30a1e..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/flat_env_cfg.py +++ /dev/null @@ -1,45 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import CassieRoughEnvCfg - - -@configclass -class CassieFlatEnvCfg(CassieRoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=52, - nconmax=15, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 5.0 - self.rewards.joint_deviation_hip.params["asset_cfg"].joint_names = ["hip_rotation_.*"] - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - - -class CassieFlatEnvCfg_PLAY(CassieFlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py deleted file mode 100644 index 2fc44ea36cdc..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/cassie/rough_env_cfg.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -from isaaclab_assets.robots.cassie import CASSIE_CFG # isort: skip - - -@configclass -class CassieRewardsCfg(RewardsCfg): - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=2.5, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*toe"), - "command_name": "base_velocity", - "threshold": 0.3, - }, - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["hip_abduction_.*", "hip_rotation_.*"])}, - ) - joint_deviation_toes = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=["toe_joint_.*"])}, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names="toe_joint_.*")}, - ) - - -@configclass -class CassieRoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: CassieRewardsCfg = CassieRewardsCfg() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = CASSIE_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.actions.joint_pos.scale = 0.5 - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = [".*pelvis"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.terminations.base_contact.params["sensor_cfg"].body_names = [".*pelvis"] - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -5.0e-6 - self.rewards.track_lin_vel_xy_exp.weight = 2.0 - self.rewards.track_ang_vel_z_exp.weight = 1.0 - self.rewards.action_rate_l2.weight *= 1.5 - self.rewards.dof_acc_l2.weight *= 1.5 - - -@configclass -class CassieRoughEnvCfg_PLAY(CassieRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py deleted file mode 100644 index f49e103f431a..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/__init__.py +++ /dev/null @@ -1,58 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.g1 import agents - -## -# Register Gym environments. -## - -# gym.register( -# id="Isaac-Velocity-Rough-G1-Warp-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:G1RoughEnvCfg", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1RoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - - -# gym.register( -# id="Isaac-Velocity-Rough-G1-Warp-Play-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:G1RoughEnvCfg_PLAY", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1RoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - -gym.register( - id="Isaac-Velocity-Flat-G1-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:G1FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-G1-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:G1FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:G1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py deleted file mode 100644 index e99dd95ff82f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/flat_env_cfg.py +++ /dev/null @@ -1,58 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import G1RoughEnvCfg - - -@configclass -class G1FlatEnvCfg(G1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=95, - nconmax=10, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - self.rewards.track_ang_vel_z_exp.weight = 1.0 - self.rewards.lin_vel_z_l2.weight = -0.2 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.0e-7 - self.rewards.feet_air_time.weight = 0.75 - self.rewards.feet_air_time.params["threshold"] = 0.4 - self.rewards.dof_torques_l2.weight = -2.0e-6 - self.rewards.dof_torques_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint"] - ) - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (-0.5, 0.5) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - - -class G1FlatEnvCfg_PLAY(G1FlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py deleted file mode 100644 index 7f9d12d0f201..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/g1/rough_env_cfg.py +++ /dev/null @@ -1,178 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -## -# Pre-defined configs -## -from isaaclab_assets import G1_MINIMAL_CFG # isort: skip - - -@configclass -class G1Rewards(RewardsCfg): - """Reward terms for the MDP.""" - - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_yaw_frame_exp, - weight=1.0, - params={"command_name": "base_velocity", "std": 0.5}, - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_world_exp, weight=2.0, params={"command_name": "base_velocity", "std": 0.5} - ) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=0.25, - params={ - "command_name": "base_velocity", - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_ankle_roll_link"), - "threshold": 0.4, - }, - ) - feet_slide = RewTerm( - func=mdp.feet_slide, - weight=-0.1, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_ankle_roll_link"), - "asset_cfg": SceneEntityCfg("robot", body_names=".*_ankle_roll_link"), - }, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, - weight=-1.0, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_ankle_pitch_joint", ".*_ankle_roll_joint"])}, - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_hip_yaw_joint", ".*_hip_roll_joint"])}, - ) - joint_deviation_arms = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={ - "asset_cfg": SceneEntityCfg( - "robot", - joint_names=[ - ".*_shoulder_pitch_joint", - ".*_shoulder_roll_joint", - ".*_shoulder_yaw_joint", - ".*_elbow_pitch_joint", - ".*_elbow_roll_joint", - ], - ) - }, - ) - joint_deviation_fingers = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.05, - params={ - "asset_cfg": SceneEntityCfg( - "robot", - joint_names=[ - ".*_five_joint", - ".*_three_joint", - ".*_six_joint", - ".*_four_joint", - ".*_zero_joint", - ".*_one_joint", - ".*_two_joint", - ], - ) - }, - ) - joint_deviation_torso = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", joint_names="torso_joint")}, - ) - - -@configclass -class G1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: G1Rewards = G1Rewards() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = G1_MINIMAL_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.events.push_robot = None - # TODO: TEMPORARILY DISABLED - adding this causes NaNs in the simulation - # self.events.add_base_mass = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = ["torso_link"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - - # Rewards - self.rewards.lin_vel_z_l2.weight = 0.0 - self.rewards.undesired_contacts = None - self.rewards.flat_orientation_l2.weight = -1.0 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.25e-7 - self.rewards.dof_acc_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint"] - ) - self.rewards.dof_torques_l2.weight = -1.5e-7 - self.rewards.dof_torques_l2.params["asset_cfg"] = SceneEntityCfg( - "robot", joint_names=[".*_hip_.*", ".*_knee_joint", ".*_ankle_.*"] - ) - - # Commands - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (-0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - - # terminations - self.terminations.base_contact.params["sensor_cfg"].body_names = "torso_link" - - -@configclass -class G1RoughEnvCfg_PLAY(G1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.episode_length_s = 40.0 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - self.commands.base_velocity.ranges.lin_vel_x = (1.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.commands.base_velocity.ranges.heading = (0.0, 0.0) - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/__init__.py deleted file mode 100644 index 417865511f5f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.contrib.velocity.config.go1 import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-UnitreeGo1-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo1FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-UnitreeGo1-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo1FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py deleted file mode 100644 index dd4012882075..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/flat_env_cfg.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import UnitreeGo1RoughEnvCfg - - -@configclass -class UnitreeGo1FlatEnvCfg(UnitreeGo1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=60, - nconmax=25, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - - -class UnitreeGo1FlatEnvCfg_PLAY(UnitreeGo1FlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py deleted file mode 100644 index cb3a602394a6..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go1/rough_env_cfg.py +++ /dev/null @@ -1,61 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets.robots.unitree import UNITREE_GO1_CFG # isort: skip - - -@configclass -class UnitreeGo1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = UNITREE_GO1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.manager_call_max_mode = {"Scene": 1} - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - self.actions.joint_pos.scale = 0.25 - self.events.push_robot = None - self.events.base_external_force_torque.params["asset_cfg"].body_names = "trunk" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "trunk" - - -@configclass -class UnitreeGo1RoughEnvCfg_PLAY(UnitreeGo1RoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/__init__.py deleted file mode 100644 index 2d7f3b3212b8..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/__init__.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.go2 import agents - -## -# Register Gym environments. -## - -gym.register( - id="Isaac-Velocity-Flat-UnitreeGo2-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo2FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-UnitreeGo2-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:UnitreeGo2FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UnitreeGo2FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py deleted file mode 100644 index 349cc450ca9f..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/flat_env_cfg.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import UnitreeGo2RoughEnvCfg - - -@configclass -class UnitreeGo2FlatEnvCfg(UnitreeGo2RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=65, - nconmax=35, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.rewards.flat_orientation_l2.weight = -2.5 - self.rewards.feet_air_time.weight = 0.25 - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - - -class UnitreeGo2FlatEnvCfg_PLAY(UnitreeGo2FlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py deleted file mode 100644 index a25748da9885..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/go2/rough_env_cfg.py +++ /dev/null @@ -1,60 +0,0 @@ -# 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 - -from isaaclab.utils.configclass import configclass - -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg - -from isaaclab_assets.robots.unitree import UNITREE_GO2_CFG # isort: skip - - -@configclass -class UnitreeGo2RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.robot = UNITREE_GO2_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) - self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 - self.actions.joint_pos.scale = 0.25 - self.events.push_robot = None - self.events.base_external_force_torque.params["asset_cfg"].body_names = "base" - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.feet_air_time.params["sensor_cfg"].body_names = ".*_foot" - self.rewards.feet_air_time.weight = 0.01 - self.rewards.undesired_contacts = None - self.rewards.dof_torques_l2.weight = -0.0002 - self.rewards.track_lin_vel_xy_exp.weight = 1.5 - self.rewards.track_ang_vel_z_exp.weight = 0.75 - self.rewards.dof_acc_l2.weight = -2.5e-7 - self.terminations.base_contact.params["sensor_cfg"].body_names = "base" - - -@configclass -class UnitreeGo2RoughEnvCfg_PLAY(UnitreeGo2RoughEnvCfg): - def __post_init__(self): - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.scene.terrain.max_init_terrain_level = None - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py deleted file mode 100644 index c731af7ca986..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/__init__.py +++ /dev/null @@ -1,57 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.velocity.config.h1 import agents - -## -# Register Gym environments. -## - -# gym.register( -# id="Isaac-Velocity-Rough-H1-Warp-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:H1RoughEnvCfg", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1RoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - -# gym.register( -# id="Isaac-Velocity-Rough-H1-Warp-Play-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.rough_env_cfg:H1RoughEnvCfg_PLAY", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1RoughPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", -# }, -# ) - -gym.register( - id="Isaac-Velocity-Flat-H1-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:H1FlatEnvCfg", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Velocity-Flat-H1-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.flat_env_cfg:H1FlatEnvCfg_PLAY", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:H1FlatPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_flat_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py deleted file mode 100644 index a7021343464a..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/flat_env_cfg.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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 - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -from .rough_env_cfg import H1RoughEnvCfg - - -@configclass -class H1FlatEnvCfg(H1RoughEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=65, - nconmax=15, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - super().__post_init__() - self.scene.terrain.terrain_type = "plane" - self.scene.terrain.terrain_generator = None - self.curriculum.terrain_levels = None - self.rewards.feet_air_time.weight = 1.0 - self.rewards.feet_air_time.params["threshold"] = 0.6 - - -class H1FlatEnvCfg_PLAY(H1FlatEnvCfg): - def __post_init__(self) -> None: - super().__post_init__() - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.observations.policy.enable_corruption = False - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py deleted file mode 100644 index 1b65e018c9e4..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/config/h1/rough_env_cfg.py +++ /dev/null @@ -1,131 +0,0 @@ -# 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 - -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg - -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp -from isaaclab_tasks_experimental.manager_based.locomotion.velocity.velocity_env_cfg import ( - LocomotionVelocityRoughEnvCfg, - RewardsCfg, -) - -## -# Pre-defined configs -## -from isaaclab_assets import H1_MINIMAL_CFG # isort: skip - - -@configclass -class H1Rewards(RewardsCfg): - """Reward terms for the MDP.""" - - termination_penalty = RewTerm(func=mdp.is_terminated, weight=-200.0) - lin_vel_z_l2 = None - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_yaw_frame_exp, - weight=1.0, - params={"command_name": "base_velocity", "std": 0.5}, - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_world_exp, weight=1.0, params={"command_name": "base_velocity", "std": 0.5} - ) - feet_air_time = RewTerm( - func=mdp.feet_air_time_positive_biped, - weight=0.25, - params={ - "command_name": "base_velocity", - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle_link"), - "threshold": 0.4, - }, - ) - feet_slide = RewTerm( - func=mdp.feet_slide, - weight=-0.25, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle_link"), - "asset_cfg": SceneEntityCfg("robot", body_names=".*ankle_link"), - }, - ) - dof_pos_limits = RewTerm( - func=mdp.joint_pos_limits, weight=-1.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*_ankle")} - ) - joint_deviation_hip = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_hip_yaw", ".*_hip_roll"])}, - ) - joint_deviation_arms = RewTerm( - func=mdp.joint_deviation_l1, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_shoulder_.*", ".*_elbow"])}, - ) - joint_deviation_torso = RewTerm( - func=mdp.joint_deviation_l1, weight=-0.1, params={"asset_cfg": SceneEntityCfg("robot", joint_names="torso")} - ) - - -@configclass -class H1RoughEnvCfg(LocomotionVelocityRoughEnvCfg): - rewards: H1Rewards = H1Rewards() - - def __post_init__(self): - super().__post_init__() - self.scene.robot = H1_MINIMAL_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - self.events.push_robot = None - self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) - self.events.base_external_force_torque.params["asset_cfg"].body_names = [".*torso_link"] - self.events.reset_base.params = { - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (0.0, 0.0), - "y": (0.0, 0.0), - "z": (0.0, 0.0), - "roll": (0.0, 0.0), - "pitch": (0.0, 0.0), - "yaw": (0.0, 0.0), - }, - } - self.events.base_com = None - self.rewards.undesired_contacts = None - self.rewards.flat_orientation_l2.weight = -1.0 - self.rewards.dof_torques_l2.weight = 0.0 - self.rewards.action_rate_l2.weight = -0.005 - self.rewards.dof_acc_l2.weight = -1.25e-7 - self.commands.base_velocity.ranges.lin_vel_x = (0.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.terminations.base_contact.params["sensor_cfg"].body_names = ".*torso_link" - - -@configclass -class H1RoughEnvCfg_PLAY(H1RoughEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - self.episode_length_s = 40.0 - # spawn the robot randomly in the grid (instead of their terrain levels) - self.scene.terrain.max_init_terrain_level = None - # reduce the number of terrains to save memory - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.num_rows = 5 - self.scene.terrain.terrain_generator.num_cols = 5 - self.scene.terrain.terrain_generator.curriculum = False - - self.commands.base_velocity.ranges.lin_vel_x = (1.0, 1.0) - self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) - self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0) - self.commands.base_velocity.ranges.heading = (0.0, 0.0) - # disable randomization for play - self.observations.policy.enable_corruption = False - # remove random pushing - self.events.base_external_force_torque = None - self.events.push_robot = None diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py deleted file mode 100644 index c1aa55114ab7..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/curriculums.py +++ /dev/null @@ -1,39 +0,0 @@ -# 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 - -"""Curriculum functions for the velocity locomotion environment. - -Curriculum terms are not warp-managed (they run at reset time, not per-step), -so they remain torch-based. -""" - -from __future__ import annotations - -from collections.abc import Sequence -from typing import TYPE_CHECKING - -import torch - -from isaaclab.managers import SceneEntityCfg - -if TYPE_CHECKING: - from isaaclab.assets import Articulation - from isaaclab.envs import ManagerBasedRLEnv - from isaaclab.terrains import TerrainImporter - - -def terrain_levels_vel( - env: ManagerBasedRLEnv, env_ids: Sequence[int], asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") -) -> torch.Tensor: - """Curriculum based on the distance the robot walked when commanded to move at a desired velocity.""" - asset: Articulation = env.scene[asset_cfg.name] - terrain: TerrainImporter = env.scene.terrain - command = env.command_manager.get_command("base_velocity") - distance = torch.norm(asset.data.root_pos_w.torch[env_ids, :2] - env.scene.env_origins[env_ids, :2], dim=1) - move_up = distance > terrain.cfg.terrain_generator.size[0] / 2 - move_down = distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5 - move_down *= ~move_up - terrain.update_env_origins(env_ids, move_up, move_down) - return torch.mean(terrain.terrain_levels.float()) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py deleted file mode 100644 index 52d53c3db268..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/velocity_env_cfg.py +++ /dev/null @@ -1,296 +0,0 @@ -# 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 - -import math -from dataclasses import MISSING - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import CurriculumTermCfg as CurrTerm -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sensors import ContactSensorCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR -from isaaclab.utils.configclass import configclass -from isaaclab.utils.noise import UniformNoiseCfg as Unoise - -import isaaclab_tasks_experimental.manager_based.locomotion.velocity.mdp as mdp - -## -# Pre-defined configs -## -from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG # isort: skip - - -## -# Scene definition -## - - -@configclass -class MySceneCfg(InteractiveSceneCfg): - """Configuration for the terrain scene with a legged robot.""" - - # ground terrain - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="generator", - terrain_generator=ROUGH_TERRAINS_CFG, - max_init_terrain_level=5, - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="multiply", - restitution_combine_mode="multiply", - static_friction=1.0, - dynamic_friction=1.0, - ), - visual_material=sim_utils.MdlFileCfg( - mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl", - project_uvw=True, - texture_scale=(0.25, 0.25), - ), - debug_vis=False, - ) - # robots - robot: ArticulationCfg = MISSING - # sensors - contact_forces = ContactSensorCfg( - prim_path="{ENV_REGEX_NS}/Robot/.*", - filter_prim_paths_expr=[], - history_length=3, - track_air_time=True, - ) - # lights - sky_light = AssetBaseCfg( - prim_path="/World/skyLight", - spawn=sim_utils.DomeLightCfg( - intensity=750.0, - texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr", - ), - ) - - -## -# MDP settings -## - - -@configclass -class CommandsCfg: - """Command specifications for the MDP.""" - - base_velocity = mdp.UniformVelocityCommandCfg( - asset_name="robot", - resampling_time_range=(10.0, 10.0), - rel_standing_envs=0.02, - rel_heading_envs=1.0, - heading_command=True, - heading_control_stiffness=0.5, - debug_vis=True, - ranges=mdp.UniformVelocityCommandCfg.Ranges( - lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi) - ), - ) - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - joint_pos = mdp.JointPositionActionCfg(asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True) - - -@configclass -class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - base_lin_vel = ObsTerm(func=mdp.base_lin_vel, noise=Unoise(n_min=-0.1, n_max=0.1)) - base_ang_vel = ObsTerm(func=mdp.base_ang_vel, noise=Unoise(n_min=-0.2, n_max=0.2)) - projected_gravity = ObsTerm( - func=mdp.projected_gravity, - noise=Unoise(n_min=-0.05, n_max=0.05), - ) - velocity_commands = ObsTerm(func=mdp.generated_commands, params={"command_name": "base_velocity"}) - joint_pos = ObsTerm(func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - joint_vel = ObsTerm(func=mdp.joint_vel_rel, noise=Unoise(n_min=-1.5, n_max=1.5)) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = True - self.concatenate_terms = True - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - # FIXME(warp-migration): COM randomization in exp manager-based locomotion currently causes - # NaNs and is temporarily disabled. - # base_com = EventTerm( - # func=mdp.randomize_rigid_body_com, - # mode="startup", - # params={ - # "asset_cfg": SceneEntityCfg("robot", body_names="base"), - # "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)}, - # }, - # ) - base_com = None - - # reset - base_external_force_torque = EventTerm( - func=mdp.apply_external_force_torque, - mode="reset", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names="base"), - "force_range": (0.0, 0.0), - "torque_range": (-0.0, 0.0), - }, - ) - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={ - "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, - "velocity_range": { - "x": (-0.5, 0.5), - "y": (-0.5, 0.5), - "z": (-0.5, 0.5), - "roll": (-0.5, 0.5), - "pitch": (-0.5, 0.5), - "yaw": (-0.5, 0.5), - }, - }, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_scale, - mode="reset", - params={ - "position_range": (0.5, 1.5), - "velocity_range": (0.0, 0.0), - }, - ) - - # interval - push_robot = EventTerm( - func=mdp.push_by_setting_velocity, - mode="interval", - interval_range_s=(10.0, 15.0), - params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}}, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # -- task - track_lin_vel_xy_exp = RewTerm( - func=mdp.track_lin_vel_xy_exp, weight=1.0, params={"command_name": "base_velocity", "std": math.sqrt(0.25)} - ) - track_ang_vel_z_exp = RewTerm( - func=mdp.track_ang_vel_z_exp, weight=0.5, params={"command_name": "base_velocity", "std": math.sqrt(0.25)} - ) - # -- penalties - lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=-2.0) - ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=-0.05) - dof_torques_l2 = RewTerm(func=mdp.joint_torques_l2, weight=-1.0e-5) - dof_acc_l2 = RewTerm(func=mdp.joint_acc_l2, weight=-2.5e-7) - action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01) - feet_air_time = RewTerm( - func=mdp.feet_air_time, - weight=0.125, - params={ - "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*FOOT"), - "command_name": "base_velocity", - "threshold": 0.5, - }, - ) - undesired_contacts = RewTerm( - func=mdp.undesired_contacts, - weight=-1.0, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*THIGH"), "threshold": 1.0}, - ) - # -- optional penalties - flat_orientation_l2 = RewTerm(func=mdp.flat_orientation_l2, weight=0.0) - dof_pos_limits = RewTerm(func=mdp.joint_pos_limits, weight=0.0) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - base_contact = DoneTerm( - func=mdp.illegal_contact, - params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="base"), "threshold": 1.0}, - ) - - -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - terrain_levels = CurrTerm(func=mdp.terrain_levels_vel) - - -## -# Environment configuration -## - - -@configclass -class LocomotionVelocityRoughEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the locomotion velocity-tracking environment.""" - - # Scene settings - scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=2.5, replicate_physics=True) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - commands: CommandsCfg = CommandsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - curriculum: CurriculumCfg = CurriculumCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 4 - self.episode_length_s = 20.0 - # simulation settings - self.sim.dt = 1.0 / 200.0 - self.sim.render_interval = self.decimation - self.sim.physics_material = self.scene.terrain.physics_material - # update sensor update periods - if self.scene.contact_forces is not None: - self.scene.contact_forces.update_period = self.sim.dt - # check if terrain levels curriculum is enabled - if getattr(self.curriculum, "terrain_levels", None) is not None: - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.curriculum = True - else: - if self.scene.terrain.terrain_generator is not None: - self.scene.terrain.terrain_generator.curriculum = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py deleted file mode 100644 index fe34199f2321..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# 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 - -"""Reach experimental task registrations (manager-based).""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py deleted file mode 100644 index 460a30569089..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# 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 diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py deleted file mode 100644 index 9f01c7b2e729..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/__init__.py +++ /dev/null @@ -1,41 +0,0 @@ -# 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 - -import gymnasium as gym - -# Reuse agent configs from the stable task package. -from isaaclab_tasks.core.reach.config.franka import agents - -## -# Register Gym environments. -## - -## -# Joint Position Control -## - -gym.register( - id="Isaac-Reach-Franka-Warp-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaReachEnvCfg", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", - }, -) - -gym.register( - id="Isaac-Reach-Franka-Warp-Play-v0", - entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", - disable_env_checker=True, - kwargs={ - "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:FrankaReachEnvCfg_PLAY", - "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", - "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg", - "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", - }, -) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py deleted file mode 100644 index d6d36afd15f5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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 - -import math - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp -from isaaclab_tasks_experimental.manager_based.manipulation.reach.reach_env_cfg import ReachEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets import FRANKA_PANDA_CFG # isort: skip - - -## -# Environment configuration -## - - -@configclass -class FrankaReachEnvCfg(ReachEnvCfg): - sim: SimulationCfg = SimulationCfg( - dt=1 / 60, - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=50, - nconmax=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ), - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # switch robot to franka - self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - # override rewards - self.rewards.end_effector_position_tracking.params["asset_cfg"].body_names = ["panda_hand"] - self.rewards.end_effector_position_tracking_fine_grained.params["asset_cfg"].body_names = ["panda_hand"] - self.rewards.end_effector_orientation_tracking.params["asset_cfg"].body_names = ["panda_hand"] - - # override actions - self.actions.arm_action = mdp.JointPositionActionCfg( - asset_name="robot", joint_names=["panda_joint.*"], scale=0.5, use_default_offset=True - ) - # override command generator body - # end-effector is along z-direction - self.commands.ee_pose.body_name = "panda_hand" - self.commands.ee_pose.ranges.pitch = (math.pi, math.pi) - - -@configclass -class FrankaReachEnvCfg_PLAY(FrankaReachEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py deleted file mode 100644 index cdff71bf6523..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/__init__.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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 - -# UR10 env disabled: USD asset has composition errors (broken asset file). -# Fails on both torch baseline and warp with: -# RuntimeError: USD stage has composition errors while loading provided stage -# Re-enable once the UR10 USD asset is fixed. - -# import gymnasium as gym -# from isaaclab_tasks.core.reach.config.ur_10 import agents - -# gym.register( -# id="Isaac-Reach-UR10-Warp-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:UR10ReachEnvCfg", -# "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UR10ReachPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", -# }, -# ) - -# gym.register( -# id="Isaac-Reach-UR10-Warp-Play-v0", -# entry_point="isaaclab_experimental.envs:ManagerBasedRLEnvWarp", -# disable_env_checker=True, -# kwargs={ -# "env_cfg_entry_point": f"{__name__}.joint_pos_env_cfg:UR10ReachEnvCfg_PLAY", -# "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", -# "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:UR10ReachPPORunnerCfg", -# "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", -# }, -# ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py deleted file mode 100644 index 42104e23b1b2..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/config/ur_10/joint_pos_env_cfg.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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 - -import math - -from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg - -from isaaclab.sim import SimulationCfg -from isaaclab.utils.configclass import configclass - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp -from isaaclab_tasks_experimental.manager_based.manipulation.reach.reach_env_cfg import ReachEnvCfg - -## -# Pre-defined configs -## -from isaaclab_assets import UR10_CFG # isort: skip - - -## -# Environment configuration -## - - -@configclass -class UR10ReachEnvCfg(ReachEnvCfg): - sim: SimulationCfg = SimulationCfg( - physics=NewtonCfg( - solver_cfg=MJWarpSolverCfg( - njmax=50, - nconmax=20, - cone="pyramidal", - impratio=1, - integrator="implicitfast", - ), - num_substeps=1, - debug_mode=False, - ) - ) - - def __post_init__(self): - # post init of parent - super().__post_init__() - - # switch robot to ur10 - self.scene.robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") - # override events - self.events.reset_robot_joints.params["position_range"] = (0.75, 1.25) - # override rewards - self.rewards.end_effector_position_tracking.params["asset_cfg"].body_names = ["ee_link"] - self.rewards.end_effector_position_tracking_fine_grained.params["asset_cfg"].body_names = ["ee_link"] - self.rewards.end_effector_orientation_tracking.params["asset_cfg"].body_names = ["ee_link"] - # override actions - self.actions.arm_action = mdp.JointPositionActionCfg( - asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True - ) - # override command generator body - # end-effector is along x-direction - self.commands.ee_pose.body_name = "ee_link" - self.commands.ee_pose.ranges.pitch = (math.pi / 2, math.pi / 2) - - -@configclass -class UR10ReachEnvCfg_PLAY(UR10ReachEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - # make a smaller scene for play - self.scene.num_envs = 50 - self.scene.env_spacing = 2.5 - # disable randomization for play - self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py deleted file mode 100644 index e83a27e44db5..000000000000 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/reach_env_cfg.py +++ /dev/null @@ -1,206 +0,0 @@ -# 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 - -from dataclasses import MISSING - -from isaaclab_experimental.managers import ObservationTermCfg as ObsTerm -from isaaclab_experimental.managers import RewardTermCfg as RewTerm -from isaaclab_experimental.managers import SceneEntityCfg -from isaaclab_experimental.managers import TerminationTermCfg as DoneTerm - -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg, AssetBaseCfg -from isaaclab.envs import ManagerBasedRLEnvCfg -from isaaclab.managers import ActionTermCfg as ActionTerm -from isaaclab.managers import CurriculumTermCfg as CurrTerm -from isaaclab.managers import EventTermCfg as EventTerm -from isaaclab.managers import ObservationGroupCfg as ObsGroup -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.utils.configclass import configclass -from isaaclab.utils.noise import UniformNoiseCfg as Unoise - -import isaaclab_tasks_experimental.manager_based.manipulation.reach.mdp as mdp - -## -# Scene definition -## - - -@configclass -class ReachSceneCfg(InteractiveSceneCfg): - """Configuration for the scene with a robotic arm.""" - - # world - ground = AssetBaseCfg( - prim_path="/World/ground", - spawn=sim_utils.GroundPlaneCfg(), - init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)), - ) - - # robots - robot: ArticulationCfg = MISSING - - # lights - light = AssetBaseCfg( - prim_path="/World/light", - spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=2500.0), - ) - - -## -# MDP settings -## - - -@configclass -class CommandsCfg: - """Command terms for the MDP.""" - - ee_pose = mdp.UniformPoseCommandCfg( - asset_name="robot", - body_name=MISSING, - resampling_time_range=(4.0, 4.0), - debug_vis=True, - ranges=mdp.UniformPoseCommandCfg.Ranges( - pos_x=(0.35, 0.65), - pos_y=(-0.2, 0.2), - pos_z=(0.15, 0.5), - roll=(0.0, 0.0), - pitch=MISSING, # depends on end-effector axis - yaw=(-3.14, 3.14), - ), - ) - - -@configclass -class ActionsCfg: - """Action specifications for the MDP.""" - - arm_action: ActionTerm = MISSING - gripper_action: ActionTerm | None = None - - -@configclass -class ObservationsCfg: - """Observation specifications for the MDP.""" - - @configclass - class PolicyCfg(ObsGroup): - """Observations for policy group.""" - - # observation terms (order preserved) - joint_pos = ObsTerm(func=mdp.joint_pos_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - joint_vel = ObsTerm(func=mdp.joint_vel_rel, noise=Unoise(n_min=-0.01, n_max=0.01)) - pose_command = ObsTerm(func=mdp.generated_commands, params={"command_name": "ee_pose"}) - actions = ObsTerm(func=mdp.last_action) - - def __post_init__(self): - self.enable_corruption = True - self.concatenate_terms = True - - # observation groups - policy: PolicyCfg = PolicyCfg() - - -@configclass -class EventCfg: - """Configuration for events.""" - - reset_base = EventTerm( - func=mdp.reset_root_state_uniform, - mode="reset", - params={"pose_range": {}, "velocity_range": {}}, - ) - - reset_robot_joints = EventTerm( - func=mdp.reset_joints_by_scale, - mode="reset", - params={ - "position_range": (0.5, 1.5), - "velocity_range": (0.0, 0.0), - }, - ) - - -@configclass -class RewardsCfg: - """Reward terms for the MDP.""" - - # task terms - end_effector_position_tracking = RewTerm( - func=mdp.position_command_error, - weight=-0.2, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "command_name": "ee_pose"}, - ) - end_effector_position_tracking_fine_grained = RewTerm( - func=mdp.position_command_error_tanh, - weight=0.1, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "std": 0.1, "command_name": "ee_pose"}, - ) - end_effector_orientation_tracking = RewTerm( - func=mdp.orientation_command_error, - weight=-0.1, - params={"asset_cfg": SceneEntityCfg("robot", body_names=MISSING), "command_name": "ee_pose"}, - ) - - # action penalty - action_rate = RewTerm(func=mdp.action_rate_l2, weight=-0.0001) - joint_vel = RewTerm( - func=mdp.joint_vel_l2, - weight=-0.0001, - params={"asset_cfg": SceneEntityCfg("robot")}, - ) - - -@configclass -class TerminationsCfg: - """Termination terms for the MDP.""" - - time_out = DoneTerm(func=mdp.time_out, time_out=True) - - -@configclass -class CurriculumCfg: - """Curriculum terms for the MDP.""" - - action_rate = CurrTerm( - func=mdp.modify_reward_weight, params={"term_name": "action_rate", "weight": -0.005, "num_steps": 4500} - ) - - joint_vel = CurrTerm( - func=mdp.modify_reward_weight, params={"term_name": "joint_vel", "weight": -0.001, "num_steps": 4500} - ) - - -## -# Environment configuration -## - - -@configclass -class ReachEnvCfg(ManagerBasedRLEnvCfg): - """Configuration for the reach end-effector pose tracking environment.""" - - # Scene settings - scene: ReachSceneCfg = ReachSceneCfg(num_envs=4096, env_spacing=2.5) - # Basic settings - observations: ObservationsCfg = ObservationsCfg() - actions: ActionsCfg = ActionsCfg() - commands: CommandsCfg = CommandsCfg() - # MDP settings - rewards: RewardsCfg = RewardsCfg() - terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() - curriculum: CurriculumCfg = CurriculumCfg() - - def __post_init__(self): - """Post initialization.""" - # general settings - self.decimation = 2 - self.sim.render_interval = self.decimation - self.episode_length_s = 12.0 - self.viewer.eye = (3.5, 3.5, 3.5) - # simulation settings - self.sim.dt = 1.0 / 60.0 diff --git a/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py b/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py new file mode 100644 index 000000000000..f04d225433a8 --- /dev/null +++ b/source/isaaclab_tasks_experimental/test/test_terrain_out_of_bounds.py @@ -0,0 +1,41 @@ +# 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 + +"""Tests for Warp locomotion termination terms.""" + +from types import SimpleNamespace + +import numpy as np +import warp as wp +from isaaclab_tasks_experimental.core.velocity.mdp.terminations import terrain_out_of_bounds + + +class _Scene(dict): + """Minimal scene mapping with terrain configuration.""" + + +def _make_env(terrain_type: str, root_x: float = 0.0): + root_pos_w = wp.array(np.full((2, 3), (root_x, 0.0, 0.0), dtype=np.float32), dtype=wp.vec3f, device="cpu") + scene = _Scene(robot=SimpleNamespace(data=SimpleNamespace(root_pos_w=SimpleNamespace(warp=root_pos_w)))) + scene.cfg = SimpleNamespace(terrain=SimpleNamespace(terrain_type=terrain_type)) + scene.terrain = SimpleNamespace( + cfg=SimpleNamespace( + terrain_generator=SimpleNamespace(size=(2.0, 2.0), num_rows=2, num_cols=2, border_width=0.0) + ) + ) + return SimpleNamespace(scene=scene, num_envs=2, device="cpu") + + +def test_terrain_out_of_bounds_does_not_share_terrain_state_between_envs(): + """A plane environment must not poison bounds used by a later generated-terrain environment.""" + for attribute in ("_is_warmed_up", "_is_plane", "_half_width", "_half_height"): + if hasattr(terrain_out_of_bounds, attribute): + delattr(terrain_out_of_bounds, attribute) + + out = wp.zeros(2, dtype=wp.bool, device="cpu") + terrain_out_of_bounds(_make_env("plane"), out) + terrain_out_of_bounds(_make_env("generator", root_x=2.0), out, distance_buffer=0.5) + + np.testing.assert_array_equal(out.numpy(), np.ones(2, dtype=np.bool_))