diff --git a/docs/source/_static/css/environment-browser.js b/docs/source/_static/css/environment-browser.js index edb297ec0bf1..27e8fabc11d6 100644 --- a/docs/source/_static/css/environment-browser.js +++ b/docs/source/_static/css/environment-browser.js @@ -38,7 +38,7 @@ ["Isaac-Reach-UR10", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "environment-previews/ur10_reach.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Reorient-Cube-Allegro-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "environment-previews/allegro_cube.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Reorient-Cube-Allegro", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "randomized,reset_only", "environment-previews/allegro_cube.jpg", false, {"*": ["reset_only"]}, {}, "", null, ["||", "||randomized", "||reset_only", "isaacsim_physx||", "isaacsim_physx||randomized", "isaacsim_physx||reset_only", "newton_mjwarp||", "newton_mjwarp||randomized", "newton_mjwarp||reset_only", "ovphysx||", "ovphysx||randomized", "ovphysx||reset_only"]], - ["Isaac-Reorient-Cube-Shadow-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "environment-previews/shadow_cube.jpg", false, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_mjwarp||", "ovphysx||"]], + ["Isaac-Reorient-Cube-Shadow-Direct", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "", "environment-previews/shadow_cube.jpg", true, {}, {}, "", null, ["||", "isaacsim_physx||", "newton_mjwarp||", "ovphysx||"]], ["Isaac-Reorient-Cube-Shadow", "rl_games,rsl_rl,skrl", "isaacsim_physx,newton_mjwarp,ovphysx", "", "asymmetric,randomized", "environment-previews/shadow_cube.jpg", false, {}, {}, "", null, ["||", "||asymmetric", "||randomized", "isaacsim_physx||", "isaacsim_physx||asymmetric", "isaacsim_physx||randomized", "newton_mjwarp||", "newton_mjwarp||asymmetric", "newton_mjwarp||randomized", "ovphysx||", "ovphysx||asymmetric", "ovphysx||randomized"]], ["Isaac-Reorient-Cube-Shadow-Camera-Direct", "rl_games,rsl_rl", "isaacsim_physx,newton_mjwarp,ovphysx", "isaacsim_rtx,newton_renderer,ovrtx", "albedo,full,rgb,semantic_segmentation,simple_shading_constant_diffuse,simple_shading_diffuse_mdl,simple_shading_full_mdl", "tasks/manipulation/shadow_cube.jpg", false, {"*": ["full"]}, {}, "", null, ["||", "||albedo", "||full", "||rgb", "||semantic_segmentation", "|isaacsim_rtx|", "|isaacsim_rtx|albedo", "|isaacsim_rtx|full", "|isaacsim_rtx|rgb", "|isaacsim_rtx|semantic_segmentation", "|isaacsim_rtx|simple_shading_constant_diffuse", "|isaacsim_rtx|simple_shading_diffuse_mdl", "|isaacsim_rtx|simple_shading_full_mdl", "|newton_renderer|", "|newton_renderer|albedo", "|newton_renderer|full", "|newton_renderer|rgb", "|newton_renderer|semantic_segmentation", "|ovrtx|", "|ovrtx|albedo", "|ovrtx|full", "|ovrtx|rgb", "|ovrtx|semantic_segmentation", "|ovrtx|simple_shading_constant_diffuse", "|ovrtx|simple_shading_diffuse_mdl", "|ovrtx|simple_shading_full_mdl", "isaacsim_physx||", "isaacsim_physx||albedo", "isaacsim_physx||full", "isaacsim_physx||rgb", "isaacsim_physx||semantic_segmentation", "isaacsim_physx|isaacsim_rtx|", "isaacsim_physx|isaacsim_rtx|albedo", "isaacsim_physx|isaacsim_rtx|full", "isaacsim_physx|isaacsim_rtx|rgb", "isaacsim_physx|isaacsim_rtx|semantic_segmentation", "isaacsim_physx|isaacsim_rtx|simple_shading_constant_diffuse", "isaacsim_physx|isaacsim_rtx|simple_shading_diffuse_mdl", "isaacsim_physx|isaacsim_rtx|simple_shading_full_mdl", "isaacsim_physx|newton_renderer|", "isaacsim_physx|newton_renderer|albedo", "isaacsim_physx|newton_renderer|full", "isaacsim_physx|newton_renderer|rgb", "isaacsim_physx|newton_renderer|semantic_segmentation", "newton_mjwarp||", "newton_mjwarp||albedo", "newton_mjwarp||full", "newton_mjwarp||rgb", "newton_mjwarp||semantic_segmentation", "newton_mjwarp|isaacsim_rtx|", "newton_mjwarp|isaacsim_rtx|albedo", "newton_mjwarp|isaacsim_rtx|full", "newton_mjwarp|isaacsim_rtx|rgb", "newton_mjwarp|isaacsim_rtx|semantic_segmentation", "newton_mjwarp|isaacsim_rtx|simple_shading_constant_diffuse", "newton_mjwarp|isaacsim_rtx|simple_shading_diffuse_mdl", "newton_mjwarp|isaacsim_rtx|simple_shading_full_mdl", "newton_mjwarp|newton_renderer|", "newton_mjwarp|newton_renderer|albedo", "newton_mjwarp|newton_renderer|full", "newton_mjwarp|newton_renderer|rgb", "newton_mjwarp|newton_renderer|semantic_segmentation", "newton_mjwarp|ovrtx|", "newton_mjwarp|ovrtx|albedo", "newton_mjwarp|ovrtx|full", "newton_mjwarp|ovrtx|rgb", "newton_mjwarp|ovrtx|semantic_segmentation", "newton_mjwarp|ovrtx|simple_shading_constant_diffuse", "newton_mjwarp|ovrtx|simple_shading_diffuse_mdl", "newton_mjwarp|ovrtx|simple_shading_full_mdl", "ovphysx||", "ovphysx||albedo", "ovphysx||full", "ovphysx||rgb", "ovphysx||semantic_segmentation", "ovphysx|newton_renderer|", "ovphysx|newton_renderer|albedo", "ovphysx|newton_renderer|full", "ovphysx|newton_renderer|rgb", "ovphysx|newton_renderer|semantic_segmentation", "ovphysx|ovrtx|", "ovphysx|ovrtx|albedo", "ovphysx|ovrtx|full", "ovphysx|ovrtx|rgb", "ovphysx|ovrtx|semantic_segmentation", "ovphysx|ovrtx|simple_shading_constant_diffuse", "ovphysx|ovrtx|simple_shading_diffuse_mdl", "ovphysx|ovrtx|simple_shading_full_mdl"]], ["Isaac-Reorient-Cube-Shadow-Camera", "rl_games,rsl_rl", "isaacsim_physx,newton_mjwarp,ovphysx", "isaacsim_rtx,newton_renderer,ovrtx", "albedo,full,randomized,rgb,semantic_segmentation,simple_shading_constant_diffuse,simple_shading_diffuse_mdl,simple_shading_full_mdl", "tasks/manipulation/shadow_cube.jpg", false, {"*": ["full"]}, {}, "", null, ["||", "||albedo", "||full", "||randomized", "||rgb", "||semantic_segmentation", "|isaacsim_rtx|", "|isaacsim_rtx|albedo", "|isaacsim_rtx|full", "|isaacsim_rtx|randomized", "|isaacsim_rtx|rgb", "|isaacsim_rtx|semantic_segmentation", "|isaacsim_rtx|simple_shading_constant_diffuse", "|isaacsim_rtx|simple_shading_diffuse_mdl", "|isaacsim_rtx|simple_shading_full_mdl", "|newton_renderer|", "|newton_renderer|albedo", "|newton_renderer|full", "|newton_renderer|randomized", "|newton_renderer|rgb", "|newton_renderer|semantic_segmentation", "|ovrtx|", "|ovrtx|albedo", "|ovrtx|full", "|ovrtx|randomized", "|ovrtx|rgb", "|ovrtx|semantic_segmentation", "|ovrtx|simple_shading_constant_diffuse", "|ovrtx|simple_shading_diffuse_mdl", "|ovrtx|simple_shading_full_mdl", "isaacsim_physx||", "isaacsim_physx||albedo", "isaacsim_physx||full", "isaacsim_physx||randomized", "isaacsim_physx||rgb", "isaacsim_physx||semantic_segmentation", "isaacsim_physx|isaacsim_rtx|", "isaacsim_physx|isaacsim_rtx|albedo", "isaacsim_physx|isaacsim_rtx|full", "isaacsim_physx|isaacsim_rtx|randomized", "isaacsim_physx|isaacsim_rtx|rgb", "isaacsim_physx|isaacsim_rtx|semantic_segmentation", "isaacsim_physx|isaacsim_rtx|simple_shading_constant_diffuse", "isaacsim_physx|isaacsim_rtx|simple_shading_diffuse_mdl", "isaacsim_physx|isaacsim_rtx|simple_shading_full_mdl", "isaacsim_physx|newton_renderer|", "isaacsim_physx|newton_renderer|albedo", "isaacsim_physx|newton_renderer|full", "isaacsim_physx|newton_renderer|randomized", "isaacsim_physx|newton_renderer|rgb", "isaacsim_physx|newton_renderer|semantic_segmentation", "newton_mjwarp||", "newton_mjwarp||albedo", "newton_mjwarp||full", "newton_mjwarp||randomized", "newton_mjwarp||rgb", "newton_mjwarp||semantic_segmentation", "newton_mjwarp|isaacsim_rtx|", "newton_mjwarp|isaacsim_rtx|albedo", "newton_mjwarp|isaacsim_rtx|full", "newton_mjwarp|isaacsim_rtx|randomized", "newton_mjwarp|isaacsim_rtx|rgb", "newton_mjwarp|isaacsim_rtx|semantic_segmentation", "newton_mjwarp|isaacsim_rtx|simple_shading_constant_diffuse", "newton_mjwarp|isaacsim_rtx|simple_shading_diffuse_mdl", "newton_mjwarp|isaacsim_rtx|simple_shading_full_mdl", "newton_mjwarp|newton_renderer|", "newton_mjwarp|newton_renderer|albedo", "newton_mjwarp|newton_renderer|full", "newton_mjwarp|newton_renderer|randomized", "newton_mjwarp|newton_renderer|rgb", "newton_mjwarp|newton_renderer|semantic_segmentation", "newton_mjwarp|ovrtx|", "newton_mjwarp|ovrtx|albedo", "newton_mjwarp|ovrtx|full", "newton_mjwarp|ovrtx|randomized", "newton_mjwarp|ovrtx|rgb", "newton_mjwarp|ovrtx|semantic_segmentation", "newton_mjwarp|ovrtx|simple_shading_constant_diffuse", "newton_mjwarp|ovrtx|simple_shading_diffuse_mdl", "newton_mjwarp|ovrtx|simple_shading_full_mdl", "ovphysx||", "ovphysx||albedo", "ovphysx||full", "ovphysx||randomized", "ovphysx||rgb", "ovphysx||semantic_segmentation", "ovphysx|newton_renderer|", "ovphysx|newton_renderer|albedo", "ovphysx|newton_renderer|full", "ovphysx|newton_renderer|randomized", "ovphysx|newton_renderer|rgb", "ovphysx|newton_renderer|semantic_segmentation", "ovphysx|ovrtx|", "ovphysx|ovrtx|albedo", "ovphysx|ovrtx|full", "ovphysx|ovrtx|randomized", "ovphysx|ovrtx|rgb", "ovphysx|ovrtx|semantic_segmentation", "ovphysx|ovrtx|simple_shading_constant_diffuse", "ovphysx|ovrtx|simple_shading_diffuse_mdl", "ovphysx|ovrtx|simple_shading_full_mdl"]], diff --git a/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst b/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst index cd0989d66e31..bb9a863b2b4f 100644 --- a/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst +++ b/docs/source/api/lab_experimental/isaaclab_experimental.utils.rst @@ -13,15 +13,13 @@ The following classes are part of the public :mod:`isaaclab_experimental.utils` .. autosummary:: :nosignatures: - ManagerCallMode - ManagerCallSwitch - WarpGraphCache + CapturedStage + captured + eager -.. autoclass:: ManagerCallMode +.. autoclass:: CapturedStage :show-inheritance: -.. autoclass:: ManagerCallSwitch - :show-inheritance: +.. autofunction:: captured -.. autoclass:: WarpGraphCache - :show-inheritance: +.. autofunction:: eager diff --git a/docs/source/concepts/warp_environments.rst b/docs/source/concepts/warp_environments.rst index 1debce871a09..34a486508fd0 100644 --- a/docs/source/concepts/warp_environments.rst +++ b/docs/source/concepts/warp_environments.rst @@ -73,6 +73,7 @@ implementation: - ``Isaac-Ant-Direct`` — Ant locomotion - ``Isaac-Humanoid-Direct`` — Humanoid locomotion - ``Isaac-Reorient-Cube-Allegro-Direct`` — Allegro hand cube reorient +- ``Isaac-Reorient-Cube-Shadow-Direct`` — Shadow hand cube reorient Manager-Based Warp Execution (``--frontend warp``) @@ -493,7 +494,8 @@ checks the mask to skip non-selected environments: return # ... modify state for selected envs only -- RNG uses per-env ``env.rng_state_wp`` (``wp.uint32``) instead of ``torch.rand`` +- RNG uses the per-env :attr:`WarpRng.state ` (``wp.uint32``) instead of + ``torch.rand`` - **Startup/prestartup** events use the torch convention ``(env, env_ids, **params)`` - **Reset/interval** events use the warp convention ``(env, env_mask, **params)`` diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst new file mode 100644 index 000000000000..935dc70f7fa4 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.added.rst @@ -0,0 +1,2 @@ +* Added :func:`~isaaclab.utils.warp.utils.resolve_ids`, which resolves an index selection given as ids or as a + boolean mask to ids on a device, as :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` does to a mask. diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..2ae1cf9e1daa --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,11 @@ +* Changed the env selection helpers of :class:`~isaaclab.utils.wrench_composer.WrenchComposer`, + :class:`~isaaclab.sensors.SensorBase` and :class:`~isaaclab.sensors.Camera` to resolve ``env_ids`` and ``env_mask`` + through :func:`~isaaclab.utils.warp.utils.resolve_ids` and :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`. + Any integer sequence and slices with a negative step are accepted, ids on another device are moved to the owner's + device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected with a :class:`TypeError`. +* Changed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to resolve ``body_ids`` through + :func:`~isaaclab.utils.warp.utils.resolve_ids` as well, so slices other than ``slice(None)`` are accepted. +* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to make ``scratch_mask`` optional for callers that never + pass ids. Ids without a ``scratch_mask`` raise a :class:`ValueError`. +* Changed :func:`~isaaclab.utils.warp.utils.resolve_1d_mask` to copy a Warp mask on another device to ``device``, and + to accept a boolean torch mask on ``device`` inside a CUDA graph capture. diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst new file mode 100644 index 000000000000..9ed262f73ec2 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-id-mask.fixed.rst @@ -0,0 +1,2 @@ +* Fixed the ``write_nodal_*_to_sim_mask`` methods of :class:`~isaaclab.assets.BaseDeformableObject`, which the + OVPhysX deformable object inherits, raising a :class:`RuntimeError` whenever ``env_mask`` was given. diff --git a/source/isaaclab/changelog.d/jichuanh-id-mask.minor b/source/isaaclab/changelog.d/jichuanh-id-mask.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.added.rst new file mode 100644 index 000000000000..a90ea5eafdf7 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.added.rst @@ -0,0 +1,3 @@ +* Added :class:`~isaaclab.utils.seed.WarpRng`, the per-environment Warp random number generator state that + environments and sensors share, one per process. :func:`~isaaclab.utils.seed.configure_seed` reseeds it in place, + and environments call :meth:`~isaaclab.utils.seed.WarpRng.initialize` at construction. diff --git a/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst new file mode 100644 index 000000000000..2fde38574e92 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.changed.rst @@ -0,0 +1,2 @@ +* Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to draw its drift from the shared + :class:`~isaaclab.utils.seed.WarpRng` state instead of random states of its own. diff --git a/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.minor b/source/isaaclab/changelog.d/jichuanh-warp-rng-shared.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst new file mode 100644 index 000000000000..37303ca182c2 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.added.rst @@ -0,0 +1,7 @@ +* Added ``env_mask`` to :meth:`~isaaclab.actuators.ActuatorCollection.reset` and + :meth:`~isaaclab.actuators.ActuatorControl.reset_native_actuators`. Isaac Lab actuator models and backend-native + actuator state reset only the masked environments from the mask itself, without converting it to indices. +* Added ``batch_mask`` to :meth:`~isaaclab.utils.buffers.DelayBuffer.set_time_lag` and + :meth:`~isaaclab.utils.buffers.DelayBuffer.reset`, which select batches on the device without synchronizing the + host. Tensor lags set through a mask are full-sized and are not range-checked. +* Added ``env_mask`` to :meth:`~isaaclab.assets.BaseCableObject.reset`, matching the other scene assets. diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst new file mode 100644 index 000000000000..ae017575fe71 --- /dev/null +++ b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.changed.rst @@ -0,0 +1,17 @@ +* Changed masked :meth:`~isaaclab.sensors.ray_caster.BaseRayCaster.reset` to resample drift in a Warp kernel from + per-environment random states seeded from torch's global generator, so it neither synchronizes the host nor + allocates and can be captured in a CUDA graph. When both ``env_ids`` and ``env_mask`` are given, the mask selects + the resampled environments, matching :meth:`~isaaclab.sensors.SensorBase.reset`. +* Changed :meth:`~isaaclab.actuators.newton.NewtonActuatorAdapter.reset` to reuse preallocated per-DOF masks, so + masked resets of Newton-native actuator state do not allocate and can be captured in a CUDA graph. +* **Breaking:** Changed :meth:`~isaaclab.actuators.ActuatorBase.reset` to take an ``env_mask`` argument, the + boolean mask of the environments to reset that masked resets pass instead of indices. Custom actuator models must + accept ``env_mask`` and reset from it without synchronizing the host, for example through + :func:`~isaaclab.utils.array.index_fill_` or :func:`torch.where`. +* Changed :class:`~isaaclab.actuators.DelayedPDActuator` to raise :class:`ValueError` at construction unless + ``0 <= min_delay <= max_delay``, instead of at a reset, because masked resets apply the sampled delays without + checking them on the host. +* Changed the ``limit`` annotation of + :meth:`~isaaclab.assets.BaseArticulation.set_fixed_tendon_position_limit_index` and + :meth:`~isaaclab.assets.BaseArticulation.set_fixed_tendon_position_limit_mask` to exclude ``float``, matching the + backends, which raise :class:`ValueError` for a float because it cannot hold both bounds. diff --git a/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.major b/source/isaaclab/changelog.d/jichuanh-warp-scene-mask-native.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/isaaclab/actuators/actuator_base.py b/source/isaaclab/isaaclab/actuators/actuator_base.py index 8b031f6fb534..a75221b05574 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_base.py +++ b/source/isaaclab/isaaclab/actuators/actuator_base.py @@ -279,11 +279,17 @@ def joint_indices(self) -> slice | torch.Tensor: """ @abstractmethod - def reset(self, env_ids: Sequence[int]): + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None): """Reset the internals within the group. + .. caution:: + If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. + Args: - env_ids: List of environment IDs to reset. + env_ids: Environment indices or a slice selecting the environments to reset. None resets all + environments. + env_mask: Boolean mask of the environments to reset. Shape is (num_envs,). Implementations reset from it + without synchronizing the host, for example through :func:`~isaaclab.utils.array.index_fill_`. """ raise NotImplementedError diff --git a/source/isaaclab/isaaclab/actuators/actuator_collection.py b/source/isaaclab/isaaclab/actuators/actuator_collection.py index a3b970a983f0..cd9366b6095a 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_collection.py +++ b/source/isaaclab/isaaclab/actuators/actuator_collection.py @@ -180,18 +180,27 @@ def applied_effort(self) -> ProxyArray: # Lifecycle. - def reset(self, env_ids: Sequence[int] | slice | None = None) -> None: + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array | None = None) -> None: """Reset all actuator group states. + Isaac Lab actuator models receive the mask as a zero-copy boolean tensor through the ``env_mask`` + argument of :meth:`~isaaclab.actuators.ActuatorBase.reset`. Backend-native actuator state resets from + the mask directly. + + .. caution:: + If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. + Args: env_ids: Environment indices to reset. Defaults to all environments. + env_mask: Environment mask. Shape is (num_instances,). Defaults to None. """ group_env_ids = self._control._normalize_index_sequence(env_ids) + group_env_mask = None if env_mask is None else wp.to_torch(env_mask) for actuator in self._groups.values(): # Newton-executed groups are reset through the backend below. if isinstance(actuator, ActuatorBase): - actuator.reset(group_env_ids) - self._control.reset_native_actuators(slice(None) if group_env_ids is None else group_env_ids) + actuator.reset(group_env_ids, env_mask=group_env_mask) + self._control.reset_native_actuators(slice(None) if group_env_ids is None else group_env_ids, env_mask) def compute(self, dt: float = 0.0) -> None: """Compute processed actuator commands and telemetry. diff --git a/source/isaaclab/isaaclab/actuators/actuator_control.py b/source/isaaclab/isaaclab/actuators/actuator_control.py index 35e84e956302..64f7d4b37261 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_control.py +++ b/source/isaaclab/isaaclab/actuators/actuator_control.py @@ -307,11 +307,15 @@ def submit_commands(self, collection: ActuatorCollection) -> None: """ raise NotImplementedError - def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None: + def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None: """Reset backend-native actuator state. + .. caution:: + If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. + Args: env_ids: Environments to reset. + env_mask: Environment mask. Shape is (num_instances,). Defaults to None. """ diff --git a/source/isaaclab/isaaclab/actuators/actuator_net.py b/source/isaaclab/isaaclab/actuators/actuator_net.py index c5dca36a609d..9a026fe07d5a 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_net.py +++ b/source/isaaclab/isaaclab/actuators/actuator_net.py @@ -68,11 +68,12 @@ def __init__(self, cfg: ActuatorNetLSTMCfg, *args, **kwargs): Operations. """ - def reset(self, env_ids: Sequence[int]): + def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None): # reset the hidden and cell states for the specified environments + selection = env_ids if env_mask is None else env_mask with torch.no_grad(): - index_fill_(self.sea_hidden_state_per_env, env_ids, 0.0, dim=1) - index_fill_(self.sea_cell_state_per_env, env_ids, 0.0, dim=1) + index_fill_(self.sea_hidden_state_per_env, selection, 0.0, dim=1) + index_fill_(self.sea_cell_state_per_env, selection, 0.0, dim=1) def compute( self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor @@ -138,10 +139,11 @@ def __init__(self, cfg: ActuatorNetMLPCfg, *args, **kwargs): Operations. """ - def reset(self, env_ids: Sequence[int]): + def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None): # reset the history for the specified environments - index_fill_(self._joint_pos_error_history, env_ids, 0.0) - index_fill_(self._joint_vel_history, env_ids, 0.0) + selection = env_ids if env_mask is None else env_mask + index_fill_(self._joint_pos_error_history, selection, 0.0) + index_fill_(self._joint_vel_history, selection, 0.0) def compute( self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor diff --git a/source/isaaclab/isaaclab/actuators/actuator_pd.py b/source/isaaclab/isaaclab/actuators/actuator_pd.py index bc20ea721fb7..c2d2e5a2c192 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_pd.py +++ b/source/isaaclab/isaaclab/actuators/actuator_pd.py @@ -190,7 +190,7 @@ def _warn_articulation_owned_write(self, name: str, writer_name: str) -> None: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None): + def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None): # This is a no-op. There is no state to reset for implicit actuators. pass @@ -313,7 +313,7 @@ def __init__( Operations. """ - def reset(self, env_ids: Sequence[int]): + def reset(self, env_ids: Sequence[int] | None = None, env_mask: torch.Tensor | None = None): pass def compute( @@ -451,18 +451,25 @@ class DelayedPDActuator(IdealPDActuator): def __init__(self, cfg: DelayedPDActuatorCfg, *args, **kwargs): super().__init__(cfg, *args, **kwargs) + # masked resets apply the sampled delays without range-checking them on the host + if not 0 <= cfg.min_delay <= cfg.max_delay: + raise ValueError(f"Expected 0 <= min_delay <= max_delay, received {cfg.min_delay} and {cfg.max_delay}.") # instantiate the delay buffers self.positions_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) self.velocities_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) self._delay_buffers = (self.positions_delay_buffer, self.velocities_delay_buffer, self.efforts_delay_buffer) - def reset(self, env_ids: Sequence[int]): - super().reset(env_ids) + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: torch.Tensor | None = None): + super().reset(env_ids, env_mask) # number of environments (since env_ids can be a slice) if env_ids is None: env_ids = slice(None) - num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) + if env_mask is not None: + # a mask samples every environment and the delay buffers keep the masked ones on the device + num_envs = self._num_envs + else: + num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) # set a new random delay for environments in env_ids time_lags = torch.randint( low=self.cfg.min_delay, @@ -473,8 +480,8 @@ def reset(self, env_ids: Sequence[int]): ) # set delays and reset buffers for delay_buffer in self._delay_buffers: - delay_buffer.set_time_lag(time_lags, env_ids) - delay_buffer.reset(env_ids) + delay_buffer.set_time_lag(time_lags, env_ids, batch_mask=env_mask) + delay_buffer.reset(env_ids, batch_mask=env_mask) def compute( self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor diff --git a/source/isaaclab/isaaclab/actuators/newton/adapter.py b/source/isaaclab/isaaclab/actuators/newton/adapter.py index cb0c8807c568..7e0be3c44448 100644 --- a/source/isaaclab/isaaclab/actuators/newton/adapter.py +++ b/source/isaaclab/isaaclab/actuators/newton/adapter.py @@ -32,12 +32,8 @@ from newton.selection import ArticulationView from ...utils import index_fill_ -from .kernels import ( - build_implicit_dof_mask, - build_per_dof_env_mask_kernel, - set_mask_kernel, - zero_at_indices_kernel, -) +from ...utils.warp.utils import resolve_1d_mask +from .kernels import build_implicit_dof_mask, build_per_dof_env_mask_kernel, zero_at_indices_kernel if TYPE_CHECKING: from .. import ActuatorCollection @@ -106,6 +102,12 @@ def __init__( self._states_a = [act.state() for act in actuators] self._states_b = [act.state() for act in actuators] + # Newton actuator states reset by a per-DOF mask, which every partial reset expands from the env mask. + # Allocated once and fully rewritten on each reset, so a reset does not allocate. + self._reset_dof_masks = [wp.zeros(act.indices.shape[0], dtype=wp.bool, device=device) for act in actuators] + # Env masks for resolving env ids in :meth:`reset`. + self._all_env_mask = wp.ones(num_envs, dtype=wp.bool, device=device) + self._scratch_env_mask = wp.zeros(num_envs, dtype=wp.bool, device=device) # Pre-clamp computed effort buffer. Each Newton actuator scatter-adds # its raw drive output to ``sim_control.joint_computed_f`` when @@ -172,7 +174,11 @@ def _swap_state_buffers(self) -> None: """Advance the actuator state ping-pong after an eager step or graph replay.""" self._states_a, self._states_b = self._states_b, self._states_a - def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + def reset( + self, + env_ids: Sequence[int] | torch.Tensor | slice | None = None, + env_mask: wp.array | None = None, + ) -> None: """Reset actuator states for the given environments. Args: @@ -180,6 +186,8 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> ``slice(None)``, which IsaacLab callers sometimes pass) resets all environments. Otherwise expects a torch tensor or sequence of int indices. + env_mask: Boolean environment mask with shape ``(num_envs,)``. + Takes precedence over :paramref:`env_ids`. Newton's :meth:`Actuator.State.reset` expects a per-DOF boolean mask of length ``num_actuators`` (= ``num_envs * dofs_per_actuator``), @@ -188,7 +196,7 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | 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 env_ids == slice(None)): for sa, sb in zip(self._states_a, self._states_b): if sa is not None: sa.reset(None) @@ -196,23 +204,18 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> sb.reset(None) return - num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) - if num_envs == 0: - return - env_mask = wp.zeros(self._num_envs, dtype=wp.bool, device=self._device) - if isinstance(env_ids, slice): - wp.to_torch(env_mask)[env_ids] = True - else: - # Torch indexed scalar assignment uploads True and synchronizes on CUDA. - indices = ( - wp.from_torch(env_ids) - if isinstance(env_ids, torch.Tensor) - else wp.array(env_ids, dtype=wp.int32, device=self._device) + if env_mask is None: + num_envs = len(range(self._num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) + if num_envs == 0: + return + env_mask = resolve_1d_mask( + ids=env_ids, + all_mask=self._all_env_mask, + scratch_mask=self._scratch_env_mask, + device=self._device, ) - wp.launch(set_mask_kernel, dim=num_envs, inputs=[env_mask, indices], 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) + 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/isaaclab/actuators/newton/kernels.py b/source/isaaclab/isaaclab/actuators/newton/kernels.py index 5272de15f539..c88da490a00e 100644 --- a/source/isaaclab/isaaclab/actuators/newton/kernels.py +++ b/source/isaaclab/isaaclab/actuators/newton/kernels.py @@ -6,7 +6,6 @@ """Shared Warp kernels for the Newton actuator fast path.""" from collections.abc import Sequence -from typing import Any import torch import warp as wp @@ -14,7 +13,7 @@ from ...utils import index_fill_ # --------------------------------------------------------------------------- -# Adapter / per-actuator helper kernels: per-DOF zeroing, env-mask building, +# Adapter / per-actuator helper kernels: per-DOF zeroing # and per-DOF env-mask projection (used by :meth:`NewtonActuatorAdapter.reset`). # Parameter reads and writes go through Newton's selection API instead. # --------------------------------------------------------------------------- @@ -27,13 +26,6 @@ def zero_at_indices_kernel(data: wp.array(dtype=wp.float32), indices: wp.array(d data[indices[i]] = 0.0 -@wp.kernel(enable_backward=False) -def set_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=Any)): - """Set selected mask entries without converting or uploading the indices.""" - i = wp.tid() - mask[indices[i]] = True - - @wp.kernel(enable_backward=False) def build_per_dof_env_mask_kernel( indices: wp.array(dtype=wp.uint32), diff --git a/source/isaaclab/isaaclab/actuators/newton/physx_runtime.py b/source/isaaclab/isaaclab/actuators/newton/physx_runtime.py index 76bc8a8d1885..a8455ef4c3b1 100644 --- a/source/isaaclab/isaaclab/actuators/newton/physx_runtime.py +++ b/source/isaaclab/isaaclab/actuators/newton/physx_runtime.py @@ -125,10 +125,15 @@ def compute(self, collection: ActuatorCollection, dt: float) -> None: return self._run_native_actuator_kernels(collection, dt) - def reset(self, env_ids: Sequence[int] | slice) -> None: - """Reset selected native actuator state.""" + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array | None = None) -> None: + """Reset selected native actuator state. + + Args: + env_ids: Environments to reset. Defaults to all environments. + env_mask: Environment mask. Takes precedence over ``env_ids``. Shape is (num_instances,). + """ if self.adapter is not None: - self.adapter.reset(env_ids) + self.adapter.reset(env_ids, env_mask) def _run_native_actuator_kernels(self, collection: ActuatorCollection, dt: float) -> None: from . import kernels as actuator_kernels # noqa: PLC0415 diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index 215f216005e0..dddb3c4c2b60 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -2091,7 +2091,7 @@ def set_fixed_tendon_limit_stiffness_mask( def set_fixed_tendon_position_limit_index( self, *, - limit: float | torch.Tensor | wp.array, + limit: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -2109,7 +2109,8 @@ def set_fixed_tendon_position_limit_index( Some backends may provide optimized implementations for masks / indices. Args: - limit: Fixed tendon limit. Shape is (len(env_ids), len(fixed_tendon_ids)). + limit: Fixed tendon position limits ``[lower, upper]`` [m]. Shape is (len(env_ids), len(fixed_tendon_ids)) + with dtype wp.vec2f. A torch tensor has a trailing dimension of 2. fixed_tendon_ids: The tendon indices to set the limit for. Defaults to None (all fixed tendons). env_ids: The environment indices to set the limit for. Defaults to None (all instances). """ @@ -2119,7 +2120,7 @@ def set_fixed_tendon_position_limit_index( def set_fixed_tendon_position_limit_mask( self, *, - limit: float | torch.Tensor | wp.array, + limit: torch.Tensor | wp.array, fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: @@ -2137,7 +2138,8 @@ def set_fixed_tendon_position_limit_mask( Some backends may provide optimized implementations for masks / indices. Args: - limit: Fixed tendon limit. Shape is (num_instances, num_fixed_tendons). + limit: Fixed tendon position limits ``[lower, upper]`` [m]. Shape is (num_instances, num_fixed_tendons) + with dtype wp.vec2f. A torch tensor has a trailing dimension of 2. fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated. Shape is (num_fixed_tendons,). env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). diff --git a/source/isaaclab/isaaclab/assets/cable_object/base_cable_object.py b/source/isaaclab/isaaclab/assets/cable_object/base_cable_object.py index 4bc2b6951f3c..49a50c93fe26 100644 --- a/source/isaaclab/isaaclab/assets/cable_object/base_cable_object.py +++ b/source/isaaclab/isaaclab/assets/cable_object/base_cable_object.py @@ -61,11 +61,12 @@ def root_view(self): raise NotImplementedError() @abstractmethod - def reset(self, env_ids: Sequence[int] | None = None) -> None: + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: """Reset the cable object. Args: env_ids: Environment indices. If None, all instances are used. + env_mask: Environment mask. If None, all instances are used. Shape is (num_instances,). """ raise NotImplementedError() diff --git a/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py b/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py index 0542a460af48..4467403b79cf 100644 --- a/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py +++ b/source/isaaclab/isaaclab/assets/deformable_object/base_deformable_object.py @@ -15,6 +15,7 @@ from ...utils import math as math_utils from ...utils.warp import ProxyArray +from ...utils.warp.utils import resolve_ids from ..asset_base import AssetBase if TYPE_CHECKING: @@ -200,6 +201,9 @@ def write_nodal_kinematic_target_to_sim_index( """ Operations - Write to simulation (mask variants). + + These convert the mask to ids, which waits for the device, and call the index variants. + Backends with a mask write path override them. """ def write_nodal_state_to_sim_mask( @@ -214,10 +218,7 @@ def write_nodal_state_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 6). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True) def write_nodal_pos_to_sim_mask( @@ -232,10 +233,7 @@ def write_nodal_pos_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 3). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True) def write_nodal_velocity_to_sim_mask( @@ -250,10 +248,7 @@ def write_nodal_velocity_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 3). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True) def write_nodal_kinematic_target_to_sim_mask( @@ -268,10 +263,7 @@ def write_nodal_kinematic_target_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 4). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True) """ diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index d052dd59aa73..b1e577520bbb 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -26,6 +26,7 @@ quat_from_matrix, ) from ...utils.warp import ProxyArray +from ...utils.warp.utils import resolve_ids from ..sensor_base import SensorBase from .camera_data import CameraData, RenderBufferKind @@ -1028,27 +1029,7 @@ def _resolve_env_ids_wp(self, env_ids: Sequence[int] | torch.Tensor | wp.array | """Resolve camera indices to a Warp ``int32`` array.""" if env_ids is None: return None - if isinstance(env_ids, wp.array): - if env_ids.dtype != wp.int32: - raise TypeError(f"Unsupported wp.array dtype for env_ids: {env_ids.dtype}. Expected wp.int32.") - if env_ids.ndim != 1: - raise ValueError("Warp camera indices must be a one-dimensional array.") - if str(env_ids.device) == str(self._device): - return env_ids - return env_ids.to(self._device) - elif isinstance(env_ids, torch.Tensor): - env_ids = env_ids.to(device=self._device, dtype=torch.int32).reshape(-1) - if not env_ids.is_contiguous(): - env_ids = env_ids.contiguous() - return wp.from_torch(env_ids, dtype=wp.int32) - elif isinstance(env_ids, slice) and (env_ids.step or 1) > 0: - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - elif isinstance(env_ids, slice): - # Torch cannot slice with a negative step. - env_ids = np.arange(self._view.count, dtype=np.int32)[env_ids] - else: - env_ids = np.asarray(env_ids, dtype=np.int32).reshape(-1) - return wp.array(env_ids, dtype=wp.int32, device=self._device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32) def _env_mask_has_any(self, env_mask: wp.array) -> bool: """Return whether the mask selects any camera.""" diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py index f8c5aba6af6c..5a8882487c0f 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -20,6 +20,7 @@ from ...terrains.trimesh.utils import make_plane from ...utils import math as math_utils from ...utils.mesh import convert_faces_to_triangles +from ...utils.seed import WarpRng from ...utils.warp import ProxyArray, convert_to_warp_mesh from ...utils.warp.kernels import raycast_mesh_masked_kernel from ..sensor_base import SensorBase @@ -119,12 +120,27 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None self._drift_sampled |= sample_drift or sample_ray_cast_drift if not self._drift_sampled: return + if env_mask is not None: + # Resample in a kernel from the shared per-environment random states with the ranges as launch + # arguments, so the masked reset neither synchronizes nor allocates and can be captured. + # A zero range draws zero, which clears drift that was sampled before it was disabled. + wp.launch( + ray_caster_kernels.resample_drift_masked_kernel, + dim=self._view_count, + inputs=[ + env_mask, + wp.vec2f(*self.cfg.drift_range), + wp.vec3f(*(low for low, _ in ray_cast_range_list)), + wp.vec3f(*(high for _, high in ray_cast_range_list)), + WarpRng.state, + ], + outputs=[self.drift.warp, self.ray_cast_drift.warp], + device=self._device, + ) + return # determine the selected batch size if env_ids is not None: num_envs_ids = len(range(self._view_count)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) - elif env_mask is not None: - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - num_envs_ids = len(env_ids) else: env_ids = slice(None) num_envs_ids = self._view_count diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py index 98e3f145d97c..6fa4c65ffcc1 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py @@ -170,6 +170,49 @@ def apply_z_drift_kernel( ray_hits[env, ray] = wp.vec3f(hit[0], hit[1], hit[2] + ray_cast_drift[env][2]) +@wp.kernel(enable_backward=False) +def resample_drift_masked_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + drift_range: wp.vec2f, + ray_cast_drift_min: wp.vec3f, + ray_cast_drift_max: wp.vec3f, + # input/output + rng_state: wp.array(dtype=wp.uint32), + # output + drift: wp.array(dtype=wp.vec3f), + ray_cast_drift: wp.array(dtype=wp.vec3f), +): + """Resample the position drift and ray-cast drift of the masked environments. + + Launch with dim=(num_envs,). + + Args: + env_mask: Boolean mask for which environments to resample. Shape is (num_envs,). + drift_range: Uniform range of every position drift component [m]. + ray_cast_drift_min: Per-axis lower bound of the ray-cast drift [m]. + ray_cast_drift_max: Per-axis upper bound of the ray-cast drift [m]. + rng_state: Random number generator states, advanced for the resampled environments. Shape is (num_envs,). + drift: Position drift [m]. Shape is (num_envs,). + ray_cast_drift: Ray-cast drift [m]. Shape is (num_envs,). + """ + env = wp.tid() + if not env_mask[env]: + return + state = rng_state[env] + drift[env] = wp.vec3f( + wp.randf(state, drift_range[0], drift_range[1]), + wp.randf(state, drift_range[0], drift_range[1]), + wp.randf(state, drift_range[0], drift_range[1]), + ) + ray_cast_drift[env] = wp.vec3f( + wp.randf(state, ray_cast_drift_min[0], ray_cast_drift_max[0]), + wp.randf(state, ray_cast_drift_min[1], ray_cast_drift_max[1]), + wp.randf(state, ray_cast_drift_min[2], ray_cast_drift_max[2]), + ) + rng_state[env] = state + + @wp.kernel(enable_backward=False) def fill_ray_hits_distance_inf_kernel( # input diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py index 5bcec23eb91e..f0e39055540d 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base.py @@ -29,7 +29,7 @@ from ..physics import PhysicsEvent, PhysicsManager from ..sim.utils.queries import get_first_matching_ancestor_prim from ..sim.utils.transforms import resolve_prim_pose -from ..utils import index_fill_ +from ..utils.warp.utils import resolve_1d_mask from .kernels import invalidate_buffers_kernel, reset_envs_kernel, update_outdated_envs_kernel, update_timestamp_kernel if TYPE_CHECKING: @@ -506,14 +506,9 @@ def _resolve_indices_and_mask( self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None ) -> wp.array: """Resolve environment indices to a warp array and mask.""" - if (env_ids is None or env_ids == slice(None)) and env_mask is None: - return self._ALL_ENV_MASK - elif env_mask is not None: - return env_mask - else: - self._reset_mask.zero_() - index_fill_(self._reset_mask_torch, env_ids, True) - return self._reset_mask + return resolve_1d_mask( + ids=env_ids, mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self._device, scratch_mask=self._reset_mask + ) def _resolve_rigid_body_ancestor_expr( self, diff --git a/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py b/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py index 17398b58869e..9ab8c652994f 100644 --- a/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py +++ b/source/isaaclab/isaaclab/utils/buffers/delay_buffer.py @@ -122,18 +122,29 @@ def time_lags(self) -> torch.Tensor: Operations. """ - def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] | None = None): - """Sets the time lag for the delay buffer across the provided batch indices. + def set_time_lag( + self, + time_lag: int | torch.Tensor, + batch_ids: Sequence[int] | None = None, + batch_mask: torch.Tensor | None = None, + ): + """Sets the time lag for the delay buffer across the provided batch indices or mask. + + .. caution:: + If both ``batch_ids`` and ``batch_mask`` are provided, ``batch_mask`` takes precedence. Args: time_lag: The desired delay for the buffer. - * If an integer is provided, the same delay is set for the provided batch indices. - * If a tensor is provided, the delay is set for each batch index separately. The shape of the tensor - should be (len(batch_ids),). + * If an integer is provided, the same delay is set for the selected batches. + * If a tensor is provided, the delay is set for each batch separately. The shape of the tensor + should be (len(batch_ids),), or (batch_size,) with ``batch_mask``. batch_ids: The batch indices for which the time lag is set. Default is None, which sets the time lag for all batch indices. + batch_mask: Boolean mask of the batches for which the time lag is set. Shape is (batch_size,). + It selects on the device without synchronizing the host, so tensor lags are not range-checked and + must lie in ``[0, history_length]``. Raises: TypeError: If the type of the :attr:`time_lag` is not int or integer tensor. @@ -150,6 +161,10 @@ def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] | # check valid dtype for time_lag: must be int or long if time_lag.dtype not in [torch.int, torch.long]: raise TypeError(f"Invalid dtype for time_lag: {time_lag.dtype}. Expected torch.int or torch.long.") + if batch_mask is not None: + time_lag = time_lag.to(device=self.device, dtype=self._time_lags.dtype) + torch.where(batch_mask, time_lag, self._time_lags, out=self._time_lags) + return min_time_lag = int(time_lag.min().item()) if time_lag.numel() else 0 max_time_lag = int(time_lag.max().item()) if time_lag.numel() else 0 else: @@ -163,20 +178,37 @@ def set_time_lag(self, time_lag: int | torch.Tensor, batch_ids: Sequence[int] | if isinstance(time_lag, torch.Tensor): self._time_lags[batch_ids] = time_lag.to(device=self.device, dtype=self._time_lags.dtype) else: - index_fill_(self._time_lags, batch_ids, time_lag) + index_fill_(self._time_lags, batch_ids if batch_mask is None else batch_mask, time_lag) - def reset(self, batch_ids: Sequence[int] | None = None): - """Reset the data in the delay buffer at the specified batch indices. + def reset(self, batch_ids: Sequence[int] | None = None, batch_mask: torch.Tensor | None = None): + """Reset the data in the delay buffer at the specified batch indices or mask. Automatically sampled lags are redrawn for those batches regardless of ``hold_prob``. Externally selected lags are preserved. + .. caution:: + If both ``batch_ids`` and ``batch_mask`` are provided, ``batch_mask`` takes precedence. + Args: batch_ids: Elements to reset in the batch dimension. Default is None, which resets all the batch indices. + batch_mask: Boolean mask of the batches to reset. Shape is (batch_size,). It selects on the device + without synchronizing the host. """ indices = slice(None) if batch_ids is None else batch_ids - index_fill_(self._num_pushes, indices, 0) - if self._hold_prob is not None: + index_fill_(self._num_pushes, indices if batch_mask is None else batch_mask, 0) + if self._hold_prob is None: + return + if batch_mask is not None: + # Draw for every batch and keep the masked ones, so the selection stays on the device. + lags = torch.randint( + self._min_lag, + self.history_length + 1, + (self.batch_size,), + dtype=self._time_lags.dtype, + device=self.device, + ) + torch.where(batch_mask, lags, self._time_lags, out=self._time_lags) + else: self._time_lags[indices] = torch.randint( self._min_lag, self.history_length + 1, diff --git a/source/isaaclab/isaaclab/utils/seed.py b/source/isaaclab/isaaclab/utils/seed.py index 535789305164..6d8f13de85dd 100644 --- a/source/isaaclab/isaaclab/utils/seed.py +++ b/source/isaaclab/isaaclab/utils/seed.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + import os import random @@ -11,6 +13,62 @@ import warp as wp +class WarpRng: + """Process-wide per-environment random number generator state for Warp kernels. + + Warp has no global generator: :func:`warp.rand_init` hashes a seed and an offset into a state that lives only + for one kernel launch. Kernels that draw random numbers per environment therefore read ``state[env]``, draw, and + store the advanced state back, so the next launch or CUDA-graph replay draws new numbers. This class holds that + array once per process, so environments, managers and sensors share it instead of each allocating a copy. + + :meth:`initialize` must be called once the number of environments and the device are known; environments do this + at construction, so users read :attr:`state` directly. :meth:`seed` is optional; without it, the state is seeded + from :func:`torch.initial_seed`. Only kernels launched over environments, one thread per environment, may advance + the state. + """ + + state: wp.array | None = None + """The per-environment random number generator state, or None before :meth:`initialize`. Shape is (num_envs,).""" + + _seed: int | None = None + + @classmethod + def seed(cls, seed: int) -> None: + """Set the seed, and reseed the state in place if it is initialized. + + Args: + seed: The random seed value. + """ + cls._seed = seed + if cls.state is not None: + cls._reset_state() + + @classmethod + def initialize(cls, num_envs: int, device: str) -> None: + """Allocate and seed a new state. Environments call it once at construction. + + Args: + num_envs: The number of environments. + device: The device of the state. + """ + cls.state = wp.empty(num_envs, dtype=wp.uint32, device=device) + cls._reset_state() + + @classmethod + def _reset_state(cls) -> None: + """Seed every environment's state from the set seed, or from torch's initial seed without one.""" + from .warp.kernels import initialize_rng_state + + seed = cls._seed if cls._seed is not None else torch.initial_seed() + wp.launch( + initialize_rng_state, + dim=cls.state.shape[0], + inputs=[seed % 2**31], + outputs=[cls.state], + device=cls.state.device, + ) + + def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int: """Set seed across all random number generators (torch, numpy, random, warp). @@ -31,7 +89,7 @@ def configure_seed(seed: int | None, torch_deterministic: bool = False) -> int: os.environ["PYTHONHASHSEED"] = str(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) - wp.rand_init(seed) + WarpRng.seed(seed) if torch_deterministic: # refer to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py index 3d046abe2bba..a7818691165e 100644 --- a/source/isaaclab/isaaclab/utils/warp/kernels.py +++ b/source/isaaclab/isaaclab/utils/warp/kernels.py @@ -780,3 +780,28 @@ def gather_subtract_2d( def gather_subtract_2d_kernel(env_ids: "wp.array | torch.Tensor", item_ids: "wp.array | torch.Tensor") -> wp.Kernel: """Select the gather-subtract worker for the selector dtypes.""" return _GATHER_SUBTRACT_2D_DISPATCHER.select(env_ids, item_ids) + + +## +# Random number generation. +## + + +@wp.kernel(enable_backward=False) +def initialize_rng_state( + # input + seed: wp.int32, + # output + state: wp.array(dtype=wp.uint32), +): + """Initialize one random number generator state per environment. + + Launch with ``dim=(num_envs,)``. Kernels that draw from ``state[i]`` should store the advanced state back, + so the next launch or graph replay draws new numbers. + + Args: + seed: Seed shared by all environments; each environment uses its index as the offset. + state: Random number generator states. Shape is (num_envs,). + """ + env_id = wp.tid() + state[env_id] = wp.rand_init(seed, env_id) diff --git a/source/isaaclab/isaaclab/utils/warp/utils.py b/source/isaaclab/isaaclab/utils/warp/utils.py index afd31e9c1b38..24e23707615c 100644 --- a/source/isaaclab/isaaclab/utils/warp/utils.py +++ b/source/isaaclab/isaaclab/utils/warp/utils.py @@ -11,6 +11,8 @@ import torch import warp as wp +from .proxy_array import ProxyArray + ## # Mask resolution - ids/mask to warp boolean mask. ## @@ -30,7 +32,7 @@ def resolve_1d_mask( ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None, mask: wp.array | torch.Tensor | None = None, all_mask: wp.array, - scratch_mask: wp.array, + scratch_mask: wp.array | None = None, device: str, ) -> wp.array: """Resolve ids/mask into a warp boolean mask. @@ -41,30 +43,36 @@ def resolve_1d_mask( Args: ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``, ``wp.array(dtype=wp.int32)``, or ``None`` (all elements). - mask: Direct boolean mask. ``wp.array`` is returned as-is; + mask: Direct boolean mask. ``wp.array`` is returned as-is, or copied when on another device; ``torch.Tensor`` is converted. all_mask: Pre-allocated all-True mask returned when both *ids* and *mask* are ``None``. scratch_mask: Pre-allocated scratch buffer populated in-place when *ids* - are provided. Not re-entrant (shared buffer). + are provided. Not re-entrant (shared buffer). Defaults to None, for callers + that never pass *ids*. device: Warp device string (e.g. ``"cuda:0"``). Returns: A ``wp.array(dtype=wp.bool)`` mask. + + Raises: + ValueError: If *ids* select a subset of elements and no ``scratch_mask`` is given. """ # Normalize slice(None) to None so the capture guard treats it identically to ids=None. if isinstance(ids, slice) and ids == slice(None): ids = None if wp.get_device().is_capturing: - if ids is not None or (mask is not None and not isinstance(mask, wp.array)): + if ids is not None or (mask is not None and _mask_needs_copy(mask, device)): raise RuntimeError( - "resolve_1d_mask is only capturable when mask is a wp.array or both ids and mask are None." + "resolve_1d_mask is only capturable when ids are None and mask is None or a boolean mask on device." ) # --- Direct mask input --- if mask is not None: if isinstance(mask, wp.array): + if mask.device != wp.get_device(device): + mask = wp.clone(mask, device=device) return mask if isinstance(mask, torch.Tensor): if mask.dtype != torch.bool: @@ -78,6 +86,9 @@ def resolve_1d_mask( if ids is None: return all_mask + if scratch_mask is None: + raise ValueError("resolve_1d_mask needs a scratch_mask to write ids into.") + if isinstance(ids, slice): scratch_mask.zero_() wp.to_torch(scratch_mask)[ids] = True @@ -115,6 +126,104 @@ def resolve_1d_mask( return scratch_mask +def _mask_needs_copy(mask: wp.array | torch.Tensor, device: str) -> bool: + """Whether :func:`resolve_1d_mask` casts or copies ``mask``, which a graph capture cannot record.""" + if isinstance(mask, wp.array): + return mask.device != wp.get_device(device) + if isinstance(mask, torch.Tensor): + return mask.dtype != torch.bool or str(mask.device) != device + return True + + +## +# Id resolution - ids/mask to ids. +## + + +def resolve_ids( + *, + ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None, + mask: wp.array | torch.Tensor | None = None, + all_ids: wp.array, + device: str, + dtype: type | None = None, +) -> wp.array | torch.Tensor: + """Resolve ids/mask into ids on ``device``. The mask takes precedence over the ids. + + Args: + ids: Index ids. Accepts ``Sequence[int]``, ``slice``, ``torch.Tensor``, ``wp.array``, + or ``None`` (all elements). Tensors and warp arrays keep their integer dtype + unless ``dtype`` is given. Pass ids from ``find_*(as_proxy=True)`` as ``.warp`` or ``.torch``. + mask: Boolean mask. Its ids are found with :func:`torch.nonzero`, which waits for + the device. + all_ids: Pre-allocated ids of every element, returned as-is when *ids* and *mask* + select every element. + device: Warp device string (e.g. ``"cuda:0"``). + dtype: Warp integer dtype the ids must have, for callers whose kernels take one + index dtype. The ids are then a contiguous 1-D ``wp.array`` of this dtype. + Defaults to None: a mask resolves to an int64 ``torch.Tensor``, and ids keep + their type and dtype. + + Returns: + ``all_ids``, a ``torch.Tensor`` or a ``wp.array`` of ids on ``device``. + + Raises: + TypeError: If *ids* is a :class:`ProxyArray`, which is output-only, *mask* is + neither a ``wp.array`` nor a ``torch.Tensor``, or ``dtype`` is given and the + ids are not 32-bit or 64-bit integers. + ValueError: If ``dtype`` is given and the ids are not one-dimensional. + """ + if mask is not None: + if isinstance(mask, wp.array): + mask = wp.to_torch(mask) + elif not isinstance(mask, torch.Tensor): + raise TypeError(f"Unsupported mask type: {type(mask)}") + ids = torch.nonzero(mask.to(wp.device_to_torch(device))).squeeze(-1) + else: + ids = _ids_on_device(ids, all_ids, device) + return ids if dtype is None else _ids_as_dtype(ids, dtype) + + +def _ids_on_device( + ids: Sequence[int] | slice | torch.Tensor | wp.array | None, all_ids: wp.array, device: str +) -> wp.array | torch.Tensor: + """Ids on ``device``, keeping the dtype of tensors and warp arrays; lists and slices become ``wp.int32``.""" + if isinstance(ids, torch.Tensor): + return ids if str(ids.device) == str(wp.get_device(device)) else ids.to(wp.device_to_torch(device)) + if ids is None: + return all_ids + if isinstance(ids, wp.array): + return ids if ids.device == wp.get_device(device) else ids.to(device) + if isinstance(ids, ProxyArray): + raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") + if isinstance(ids, slice): + if ids == slice(None): + return all_ids + start, stop, step = ids.indices(all_ids.shape[0]) + if step > 0: + # a strided view of the caller's all-ids, no copy + return wp.from_torch(wp.to_torch(all_ids)[start:stop:step]) + return wp.array(list(range(start, stop, step)), dtype=wp.int32, device=device) + return wp.array(list(ids), dtype=wp.int32, device=device) + + +def _ids_as_dtype(ids: wp.array | torch.Tensor, dtype: type) -> wp.array: + """``ids`` as a contiguous 1-D ``wp.array`` of ``dtype`` on their device.""" + if dtype not in (wp.int32, wp.int64): + raise TypeError(f"Unsupported id dtype: {dtype}. Expected wp.int32 or wp.int64.") + if ids.ndim != 1: + raise ValueError(f"Ids must be a one-dimensional array, got {ids.ndim} dimensions.") + if isinstance(ids, wp.array): + if ids.dtype == dtype and ids.is_contiguous: + return ids + if ids.dtype not in (wp.int32, wp.int64): + raise TypeError(f"Warp ids must be 32-bit or 64-bit signed integers, got {ids.dtype}.") + ids = wp.to_torch(ids) + elif ids.dtype.is_floating_point or ids.dtype.is_complex or ids.dtype == torch.bool: + raise TypeError(f"Ids must be integers, got {ids.dtype}.") + return wp.from_torch(ids.to(wp.dtype_to_torch(dtype)).contiguous(), dtype=dtype) + + ## # Capture safety — property guard. ## diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py index 6331716cad5d..ace1570dc4dc 100644 --- a/source/isaaclab/isaaclab/utils/wrench_composer.py +++ b/source/isaaclab/isaaclab/utils/wrench_composer.py @@ -24,6 +24,7 @@ set_forces_to_dual_buffers_index_kernel, set_forces_to_dual_buffers_mask, ) +from .warp.utils import resolve_ids if TYPE_CHECKING: from ..assets import BaseArticulation, BaseRigidObject, BaseRigidObjectCollection @@ -721,24 +722,9 @@ def _resolve_env_ids(self, env_ids: wp.array | torch.Tensor | list | slice | Non Environment indices. Raises: - TypeError: If ``env_ids`` is an unsupported type. + TypeError: If ``env_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`. """ - if env_ids is None: - return self._ALL_ENV_INDICES - # Check tensor types before slice comparison (tensor == slice crashes) - if isinstance(env_ids, torch.Tensor): - return env_ids - if isinstance(env_ids, wp.array): - return env_ids - if env_ids == slice(None): - return self._ALL_ENV_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - raise TypeError( - f"env_ids must be None, slice(None), list, torch.Tensor, or wp.array, got {type(env_ids).__name__}" - ) + return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device) def _resolve_body_ids( self, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None @@ -752,23 +738,9 @@ def _resolve_body_ids( Body indices. Raises: - TypeError: If ``body_ids`` is an unsupported type. + TypeError: If ``body_ids`` is a :class:`~isaaclab.utils.warp.ProxyArray`. """ - if body_ids is None: - return self._ALL_BODY_INDICES - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, torch.Tensor): - return body_ids - if isinstance(body_ids, wp.array): - return body_ids - if body_ids == slice(None): - return self._ALL_BODY_INDICES - if isinstance(body_ids, Sequence): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - raise TypeError( - f"body_ids must be None, slice(None), a sequence, torch.Tensor, or wp.array, got {type(body_ids).__name__}" - ) + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) def _ensure_composed(self): """Compose input buffers into output buffers if dirty.""" diff --git a/source/isaaclab/test/actuators/test_delayed_pd_actuator.py b/source/isaaclab/test/actuators/test_delayed_pd_actuator.py new file mode 100644 index 000000000000..5c2d5a6cd38c --- /dev/null +++ b/source/isaaclab/test/actuators/test_delayed_pd_actuator.py @@ -0,0 +1,59 @@ +# 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 pytest +import torch + +from isaaclab.actuators import DelayedPDActuatorCfg +from isaaclab.test.utils import DeviceScope, test_devices +from isaaclab.utils import instantiate +from isaaclab.utils.types import ArticulationActions + +pytestmark = pytest.mark.unit + +NUM_ENVS, NUM_JOINTS = 4, 2 + + +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_masked_reset_matches_index_reset(device): + """A boolean-mask reset delays the same environments as the equivalent index reset, without synchronizing CUDA.""" + joint_names = [f"joint_{j}" for j in range(NUM_JOINTS)] + cfg = DelayedPDActuatorCfg(joint_names_expr=joint_names, stiffness=10.0, damping=1.0, min_delay=2, max_delay=2) + by_ids, by_mask = ( + instantiate(cfg, joint_names=joint_names, joint_ids=slice(None), num_envs=NUM_ENVS, device=device) + for _ in range(2) + ) + env_ids = torch.tensor([1, 2], device=device) + env_mask = torch.zeros(NUM_ENVS, dtype=torch.bool, device=device) + env_mask[env_ids] = True + zeros = torch.zeros(NUM_ENVS, NUM_JOINTS, device=device) + + def step(actuator, target: float) -> torch.Tensor: + action = ArticulationActions(joint_positions=zeros + target, joint_velocities=zeros, joint_efforts=zeros) + actuator.compute(action, joint_pos=zeros, joint_vel=zeros) + return actuator.computed_effort.clone() + + # A lag of one in every environment differs from the configured lag the reset samples. + for actuator in (by_ids, by_mask): + for delay_buffer in ( + actuator.positions_delay_buffer, + actuator.velocities_delay_buffer, + actuator.efforts_delay_buffer, + ): + delay_buffer.set_time_lag(1) + for target in range(3): + step(actuator, float(target)) + + by_ids.reset(env_ids) + torch.cuda.synchronize(device) + previous = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + by_mask.reset(env_mask=env_mask) + finally: + torch.cuda.set_sync_debug_mode(previous) + + for target in range(3, 7): + torch.testing.assert_close(step(by_mask, float(target)), step(by_ids, float(target))) diff --git a/source/isaaclab/test/assets/contract/articulation_factory.py b/source/isaaclab/test/assets/contract/articulation_factory.py index 8ff52e4fdafb..6b0488bb006b 100644 --- a/source/isaaclab/test/assets/contract/articulation_factory.py +++ b/source/isaaclab/test/assets/contract/articulation_factory.py @@ -237,6 +237,7 @@ def _newton_articulation(N, J, B, FT, ST, device, is_fixed_base, joint_ordering, articulation._ALL_BODY_INDICES = indices(B, device) articulation._ALL_JOINT_INDICES = indices(J, device) articulation._ALL_ENV_MASK = wp.ones((N,), dtype=wp.bool, device=device) + articulation._scratch_env_mask = wp.zeros((N,), dtype=wp.bool, device=device) articulation._ALL_BODY_MASK = wp.ones((B,), dtype=wp.bool, device=device) articulation._ALL_JOINT_MASK = wp.ones((J,), dtype=wp.bool, device=device) articulation._resolve_and_install_ordering_maps() diff --git a/source/isaaclab/test/assets/contract/rigid_object_factory.py b/source/isaaclab/test/assets/contract/rigid_object_factory.py index 0bf4741ed9f5..cb108da59a1c 100644 --- a/source/isaaclab/test/assets/contract/rigid_object_factory.py +++ b/source/isaaclab/test/assets/contract/rigid_object_factory.py @@ -195,6 +195,7 @@ def _newton_collection(num_instances: int, num_bodies: int, device: str, monkeyp collection._ALL_BODY_INDICES = indices(num_bodies, device) collection._ALL_ENV_MASK = wp.ones((num_instances,), dtype=wp.bool, device=device) collection._ALL_BODY_MASK = wp.ones((num_bodies,), dtype=wp.bool, device=device) + collection._scratch_body_mask = wp.zeros((num_bodies,), dtype=wp.bool, device=device) return collection, view diff --git a/source/isaaclab/test/assets/contract/test_articulation_contract.py b/source/isaaclab/test/assets/contract/test_articulation_contract.py index 99abc45ae632..42878092b782 100644 --- a/source/isaaclab/test/assets/contract/test_articulation_contract.py +++ b/source/isaaclab/test/assets/contract/test_articulation_contract.py @@ -536,11 +536,6 @@ def read() -> ProxyArray: @requires("newton") -@pytest.mark.xfail( - raises=pytest.fail.Exception, - strict=True, - reason="Newton fixed-tendon mask setters slice full data without validating its shape", -) def test_newton_fixed_tendon_mask_setters_reject_extra_environments(monkeypatch): art, _ = _articulation("newton", "cpu", monkeypatch) with pytest.raises((AssertionError, RuntimeError)): diff --git a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py index 1d294d9a8efc..514f4811ad57 100644 --- a/source/isaaclab/test/assets/contract/test_rigid_object_contract.py +++ b/source/isaaclab/test/assets/contract/test_rigid_object_contract.py @@ -314,6 +314,27 @@ def test_pose_and_velocity_writer_aliases_forward_to_canonical_writers(monkeypat torch.testing.assert_close(getattr(obj.data, getter).torch, values, atol=1e-5, rtol=1e-5) +@pytest.mark.parametrize("quantity", ["link_pose", "com_pose", "com_velocity", "link_velocity"]) +def test_collection_mask_writers_select_bodies_by_ids(monkeypatch, backend, device, quantity): + """Collection mask writers given ``body_ids`` write only those bodies of the masked environments.""" + if backend == "ovphysx": + pytest.skip("OVPhysX collection mask writers take no body_ids.") + obj, _, _, num_bodies = _create("collection", backend, monkeypatch, device) + obj.data.update(dt=0.01) + trailing = (7,) if quantity.endswith("pose") else (6,) + kwarg = "body_poses" if quantity.endswith("pose") else "body_velocities" + getter = f"body_{quantity.replace('velocity', 'vel')}_w" + method = getattr(obj, f"write_body_{quantity}_to_sim_mask") + full = _payload((_NUM_INSTANCES, num_bodies), trailing, device=device) + method(**{kwarg: full}) + update = _payload((_NUM_INSTANCES, num_bodies), trailing, offset=100.0, device=device) + env_mask = wp.array(np.arange(_NUM_INSTANCES) == _NUM_INSTANCES - 1, dtype=wp.bool, device=str(device)) + method(**{kwarg: update}, env_mask=env_mask, body_ids=[num_bodies - 1]) + expected = full.clone() + expected[-1, -1] = update[-1, -1] + torch.testing.assert_close(getattr(obj.data, getter).torch, expected, atol=1e-5, rtol=1e-5) + + # backend-stored quantity -> (written quantity, per-element shape) _PARTIAL_WRITES = {"pose": ("link_pose", (7,)), "velocity": ("com_velocity", (6,)), "mass": ("masses", ())} diff --git a/source/isaaclab/test/renderers/test_simulation_render_context.py b/source/isaaclab/test/renderers/test_simulation_render_context.py index fdcc8d4f3779..cf16bc4a2d76 100644 --- a/source/isaaclab/test/renderers/test_simulation_render_context.py +++ b/source/isaaclab/test/renderers/test_simulation_render_context.py @@ -317,6 +317,7 @@ def __init__(self, renderer, name, update_period=0.0): self._frame = ProxyArray(wp.zeros(2, dtype=wp.int64, device="cpu")) self._ALL_INDICES = wp.array([0, 1], dtype=wp.int32, device="cpu") self._ALL_ENV_MASK = wp.ones(2, dtype=wp.bool, device="cpu") + self._reset_mask = wp.zeros(2, dtype=wp.bool, device="cpu") self._create_timing_buffers() self._invalidate_buffers(self._ALL_ENV_MASK) self.pose = 0.0 diff --git a/source/isaaclab/test/sensors/test_sensor_update_schedule.py b/source/isaaclab/test/sensors/test_sensor_update_schedule.py index 66224a3fea17..39353ecefdb3 100644 --- a/source/isaaclab/test/sensors/test_sensor_update_schedule.py +++ b/source/isaaclab/test/sensors/test_sensor_update_schedule.py @@ -70,6 +70,7 @@ def __init__(self, update_period: float, device: str, num_envs: int = 2): self._device = device self._num_envs = num_envs self._ALL_ENV_MASK = wp.ones(num_envs, dtype=wp.bool, device=device) + self._reset_mask = wp.zeros(num_envs, dtype=wp.bool, device=device) self._is_initialized = True self._is_visualizing = False self._create_timing_buffers() diff --git a/source/isaaclab/test/utils/test_delay_buffer.py b/source/isaaclab/test/utils/test_delay_buffer.py index d2e46d19d13d..6055fc0ffde0 100644 --- a/source/isaaclab/test/utils/test_delay_buffer.py +++ b/source/isaaclab/test/utils/test_delay_buffer.py @@ -101,6 +101,41 @@ def test_random_time_lags(delay_buffer, hold_prob, monkeypatch): torch.testing.assert_close(delay_buffer.compute(data + 100, update_history=False), expected) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_mask_selects_like_indices(device, monkeypatch): + """A boolean mask sets lags and resets histories of the same batches as indices, without synchronizing CUDA.""" + batch_ids = [1, 3] + by_ids, by_mask = (DelayBuffer(4, 10, device, hold_prob=1.0) for _ in range(2)) + mask = torch.zeros(by_mask.batch_size, dtype=torch.bool, device=device) + mask[batch_ids] = True + # Masked lags are full-sized; unmasked entries differ from the current lags so applying them shows. + masked_lags = torch.ones(by_mask.batch_size, dtype=torch.long, device=device) + masked_lags[batch_ids] = torch.tensor([4, 0], device=device) + monkeypatch.setattr(torch, "randint", lambda low, high, size, **kwargs: torch.full(size, 2, **kwargs)) + for delay_buffer in (by_ids, by_mask): + delay_buffer.set_time_lag(torch.arange(delay_buffer.batch_size, device=device) % 5) + for data in _generate_data(delay_buffer.batch_size, 3, device): + delay_buffer.compute(data) + + def without_sync(call): + torch.cuda.synchronize(device) + previous = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + call() + finally: + torch.cuda.set_sync_debug_mode(previous) + + by_ids.set_time_lag(masked_lags[batch_ids], batch_ids) + without_sync(lambda: by_mask.set_time_lag(masked_lags, batch_mask=mask)) + torch.testing.assert_close(by_mask.time_lags, by_ids.time_lags) + by_ids.reset(batch_ids) + without_sync(lambda: by_mask.reset(batch_mask=mask)) + torch.testing.assert_close(by_mask.time_lags, by_ids.time_lags) + for data in _generate_data(by_mask.batch_size, 4, device): + torch.testing.assert_close(by_mask.compute(data), by_ids.compute(data)) + + @pytest.mark.parametrize( ("time_lag", "batch_ids"), [ diff --git a/source/isaaclab/test/utils/test_warp_rng.py b/source/isaaclab/test/utils/test_warp_rng.py new file mode 100644 index 000000000000..2bc3815a2b8c --- /dev/null +++ b/source/isaaclab/test/utils/test_warp_rng.py @@ -0,0 +1,63 @@ +# 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 process-wide per-environment Warp random number generator state.""" + +import numpy as np +import pytest +import torch + +from isaaclab.test.utils import test_devices +from isaaclab.utils.seed import WarpRng, configure_seed + + +@pytest.fixture(autouse=True) +def fresh_warp_rng(): + """Start and end every test without a state or a set seed.""" + WarpRng.state = None + WarpRng._seed = None + yield + WarpRng.state = None + WarpRng._seed = None + + +@pytest.mark.parametrize("device", test_devices()) +def test_initialize_allocates_a_new_state(device): + """Every :meth:`WarpRng.initialize` call allocates and seeds a new state of the given size.""" + WarpRng.initialize(4, device) + state = WarpRng.state + WarpRng.initialize(5, device) + assert WarpRng.state is not state + assert WarpRng.state.shape == (5,) + + +@pytest.mark.parametrize("device", test_devices()) +def test_configure_seed_reseeds_in_place(device): + """configure_seed reseeds the existing state in place and reproducibly.""" + configure_seed(7) + WarpRng.initialize(4, device) + state = WarpRng.state + seeded = state.numpy().copy() + + state.fill_(0) + configure_seed(7) + assert WarpRng.state.ptr == state.ptr + np.testing.assert_array_equal(state.numpy(), seeded) + + configure_seed(8) + assert not np.array_equal(state.numpy(), seeded) + + +@pytest.mark.parametrize("device", test_devices()) +def test_default_seed_follows_torch(device): + """Without a set seed, the state is seeded from torch's initial seed.""" + torch.manual_seed(5) + WarpRng.initialize(4, device) + default = WarpRng.state.numpy().copy() + + WarpRng.state = None + WarpRng.seed(5) + WarpRng.initialize(4, device) + np.testing.assert_array_equal(WarpRng.state.numpy(), default) diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index 10daf0d06612..b236b052707a 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -477,8 +477,9 @@ def test_deprecated_writers_warn_and_keep_add_and_set_behavior() -> None: def test_invalid_selection_and_empty_wrench_input_are_reported() -> None: composer = _make_composer() - with pytest.raises(TypeError, match="env_ids must be"): - composer.add_forces_and_torques_index(forces=_vectors(torch.ones((1, 2, 3))), env_ids=(0,)) + env_ids = ProxyArray(wp.array([0], dtype=wp.int32, device=composer.device)) + with pytest.raises(TypeError, match="ProxyArray is output-only"): + composer.add_forces_and_torques_index(forces=_vectors(torch.ones((1, 2, 3))), env_ids=env_ids) for method in ("add_forces_and_torques_index", "set_forces_and_torques_mask"): with pytest.warns(UserWarning, match="No forces or torques"): getattr(composer, method)() diff --git a/source/isaaclab/test/utils/warp/test_resolve_selection.py b/source/isaaclab/test/utils/warp/test_resolve_selection.py new file mode 100644 index 000000000000..fc951df8f614 --- /dev/null +++ b/source/isaaclab/test/utils/warp/test_resolve_selection.py @@ -0,0 +1,155 @@ +# 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 resolving index selections with :func:`resolve_ids` and :func:`resolve_1d_mask`.""" + +import pytest +import torch +import warp as wp + +from isaaclab.test.utils import test_devices +from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids + + +def _all_ids(n: int, device: str) -> wp.array: + return wp.array(list(range(n)), dtype=wp.int32, device=device) + + +@pytest.mark.parametrize("device", test_devices()) +def test_resolve_ids(device): + """Every selection resolves to device ids; a selection of every element is the shared constant itself.""" + all_ids = _all_ids(10, device) + assert resolve_ids(all_ids=all_ids, device=device) is all_ids + assert resolve_ids(ids=slice(None), all_ids=all_ids, device=device) is all_ids + assert resolve_ids(ids=[3, 1], all_ids=all_ids, device=device).numpy().tolist() == [3, 1] + assert resolve_ids(ids=slice(2, 8, 3), all_ids=all_ids, device=device).numpy().tolist() == [2, 5] + assert resolve_ids(ids=slice(None, None, -4), all_ids=all_ids, device=device).numpy().tolist() == [9, 5, 1] + ids_torch = torch.tensor([4, 2], dtype=torch.int64, device=wp.device_to_torch(device)) + assert resolve_ids(ids=ids_torch, all_ids=all_ids, device=device) is ids_torch + ids_wp = wp.array([6], dtype=wp.int32, device=device) + assert resolve_ids(ids=ids_wp, all_ids=all_ids, device=device) is ids_wp + with pytest.raises(TypeError, match="ProxyArray"): + resolve_ids(ids=ProxyArray(ids_wp), all_ids=all_ids, device=device) + + +@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("mask_type", ["warp", "torch_bool", "torch_int"]) +def test_resolve_ids_from_a_mask(device, mask_type): + """A mask resolves to its ascending ids, takes precedence over ids, and honours ``dtype``.""" + n = 4096 + expected = (torch.rand(n, generator=torch.Generator().manual_seed(0)) < 0.5).to(wp.device_to_torch(device)) + mask = { + "warp": wp.from_torch(expected, dtype=wp.bool), + "torch_bool": expected, + "torch_int": expected.to(torch.int64), + }[mask_type] + all_ids = _all_ids(n, device) + ids = resolve_ids(ids=[0, 1], mask=mask, all_ids=all_ids, device=device) + assert ids.dtype == torch.int64 and torch.equal(ids, torch.nonzero(expected).flatten()) + ids_int32 = resolve_ids(mask=mask, all_ids=all_ids, device=device, dtype=wp.int32) + assert isinstance(ids_int32, wp.array) and ids_int32.dtype == wp.int32 and ids_int32.is_contiguous + assert torch.equal(wp.to_torch(ids_int32), torch.nonzero(expected).flatten().to(torch.int32)) + + +@pytest.mark.parametrize("device", test_devices()) +def test_resolve_ids_with_dtype(device): + """``dtype=wp.int32`` gives a contiguous one-dimensional int32 Warp array for every kind of ids.""" + all_ids = _all_ids(10, device) + torch_device = wp.device_to_torch(device) + assert resolve_ids(all_ids=all_ids, device=device, dtype=wp.int32) is all_ids + for ids, expected in ( + (torch.tensor([4, 2], dtype=torch.int64, device=torch_device), [4, 2]), + (torch.arange(10, device=torch_device)[::3], [0, 3, 6, 9]), + (wp.array([7, 1], dtype=wp.int64, device=device), [7, 1]), + (slice(1, 8, 3), [1, 4, 7]), + (slice(None, None, -4), [9, 5, 1]), + ([3, 8], [3, 8]), + ): + resolved = resolve_ids(ids=ids, all_ids=all_ids, device=device, dtype=wp.int32) + assert isinstance(resolved, wp.array) and resolved.dtype == wp.int32 and resolved.is_contiguous + assert resolved.numpy().tolist() == expected + ids_int32 = wp.array([5], dtype=wp.int32, device=device) + assert resolve_ids(ids=ids_int32, all_ids=all_ids, device=device, dtype=wp.int32) is ids_int32 + with pytest.raises(TypeError): + resolve_ids(ids=torch.tensor([1.0], device=torch_device), all_ids=all_ids, device=device, dtype=wp.int32) + with pytest.raises(TypeError): + resolve_ids( + ids=wp.array([1.0], dtype=wp.float32, device=device), all_ids=all_ids, device=device, dtype=wp.int32 + ) + with pytest.raises(ValueError): + resolve_ids( + ids=torch.zeros((2, 2), dtype=torch.int64, device=torch_device), + all_ids=all_ids, + device=device, + dtype=wp.int32, + ) + + +@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("ids_type", ["list", "slice", "torch_int64", "torch_int32", "warp_int32"]) +def test_resolve_1d_mask_from_ids(device, ids_type): + """Ids of every type are written into the scratch mask, which is cleared first.""" + n = 12 + all_mask = wp.ones(n, dtype=wp.bool, device=device) + scratch = wp.ones(n, dtype=wp.bool, device=device) + selected = [1, 4, 5, 11] + torch_device = wp.device_to_torch(device) + ids = { + "list": selected, + "slice": slice(1, 12, 3), + "torch_int64": torch.tensor(selected, dtype=torch.int64, device=torch_device), + "torch_int32": torch.tensor(selected, dtype=torch.int32, device=torch_device), + "warp_int32": wp.array(selected, dtype=wp.int32, device=device), + }[ids_type] + expected = [False] * n + for i in range(n)[ids] if isinstance(ids, slice) else selected: + expected[i] = True + mask = resolve_1d_mask(ids=ids, all_mask=all_mask, scratch_mask=scratch, device=device) + assert mask is scratch + assert mask.numpy().tolist() == expected + + +@pytest.mark.parametrize("device", test_devices()) +def test_resolve_1d_mask_needs_a_scratch_mask_only_for_ids(device): + """A mask or no selection resolves without a scratch mask; ids without one raise.""" + all_mask = wp.ones(5, dtype=wp.bool, device=device) + mask = wp.array([True, False, True, False, False], dtype=wp.bool, device=device) + assert resolve_1d_mask(ids=[0], mask=mask, all_mask=all_mask, device=device) is mask + assert resolve_1d_mask(all_mask=all_mask, device=device) is all_mask + assert resolve_1d_mask(ids=slice(None), all_mask=all_mask, device=device) is all_mask + mask_torch = torch.tensor([0, 1, 1, 0, 0], device=wp.device_to_torch(device)) + resolved = resolve_1d_mask(mask=mask_torch, all_mask=all_mask, device=device) + assert resolved.dtype == wp.bool and resolved.numpy().tolist() == [False, True, True, False, False] + with pytest.raises(ValueError, match="scratch_mask"): + resolve_1d_mask(ids=[1], all_mask=all_mask, device=device) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_resolve_1d_mask_copies_a_mask_on_another_device(): + """A Warp mask on another device is copied to ``device``.""" + all_mask = wp.ones(3, dtype=wp.bool, device="cuda:0") + mask = wp.array([True, False, True], dtype=wp.bool, device="cpu") + resolved = resolve_1d_mask(mask=mask, all_mask=all_mask, device="cuda:0") + assert str(resolved.device) == "cuda:0" and resolved.numpy().tolist() == [True, False, True] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_resolve_1d_mask_in_a_graph_capture(): + """Inside a capture, a boolean mask on device resolves; ids or a mask that needs a cast or copy raise.""" + device = "cuda:0" + all_mask = wp.ones(3, dtype=wp.bool, device=device) + scratch = wp.zeros(3, dtype=wp.bool, device=device) + mask_wp = wp.array([True, False, True], dtype=wp.bool, device=device) + mask_torch = torch.tensor([False, True, True], device=device) + mask_int = mask_torch.int() + mask_cpu = wp.array([True, False, True], dtype=wp.bool, device="cpu") + with wp.ScopedDevice(device), wp.ScopedCapture(): + assert resolve_1d_mask(mask=mask_wp, all_mask=all_mask, device=device) is mask_wp + resolved = resolve_1d_mask(mask=mask_torch, all_mask=all_mask, device=device) + for selection in ({"mask": mask_int}, {"mask": mask_cpu}, {"ids": [0]}): + with pytest.raises(RuntimeError, match="capturable"): + resolve_1d_mask(**selection, all_mask=all_mask, scratch_mask=scratch, device=device) + assert resolved.numpy().tolist() == [False, True, True] diff --git a/source/isaaclab_contrib/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_contrib/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst new file mode 100644 index 000000000000..177eb50253d4 --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst @@ -0,0 +1,2 @@ +* Fixed :class:`~isaaclab_contrib.assets.Multirotor` resets with only ``env_mask`` resetting the thrusters of every + environment. diff --git a/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py b/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py index 064fb1a17858..f65ded646726 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py +++ b/source/isaaclab_contrib/isaaclab_contrib/assets/multirotor/multirotor.py @@ -42,12 +42,19 @@ class _ThrusterCollection(dict): exposing the :meth:`reset` entry point that :meth:`isaaclab.assets.Articulation.reset` invokes. """ - def reset(self, env_ids: Sequence[int] | slice | None = None) -> None: + def reset(self, env_ids: Sequence[int] | slice | None = None, env_mask: wp.array | None = None) -> None: """Reset every thruster actuator for the given environments. + .. caution:: + If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. + Args: env_ids: Environment indices to reset. Defaults to None (all environments). + env_mask: Environment mask. Shape is (num_instances,). Defaults to None. """ + if env_mask is not None: + # thrusters index their buffers by environment + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) for actuator in self.values(): actuator.reset(env_ids) diff --git a/source/isaaclab_contrib/test/assets/test_multirotor.py b/source/isaaclab_contrib/test/assets/test_multirotor.py index 5d195eeffe7a..d0bf04b2baf4 100644 --- a/source/isaaclab_contrib/test/assets/test_multirotor.py +++ b/source/isaaclab_contrib/test/assets/test_multirotor.py @@ -18,12 +18,14 @@ import pytest import torch +import warp as wp from isaaclab_physx.physics import PhysxCfg import isaaclab.sim.utils.prims as prim_utils from isaaclab.sim import SimulationCfg, build_simulation_context from isaaclab_contrib.assets import Multirotor, MultirotorCfg +from isaaclab_contrib.assets.multirotor.multirotor import _ThrusterCollection # Best-effort: suppress unraisable destructor warnings emitted during # teardown of partially-constructed assets in CI/dev environments. We still @@ -156,6 +158,16 @@ def test_set_thrust_target_env_slice_unit(): assert torch.allclose(m._data.thrust_target[3], original[3]) +def test_thruster_collection_resets_masked_envs_unit(): + """A mask-only reset, as the articulation forwards it, resets the thrusters of the masked envs only.""" + reset_env_ids = [] + thrusters = _ThrusterCollection(thrusters=types.SimpleNamespace(reset=reset_env_ids.append)) + + thrusters.reset(env_mask=wp.array([True, False, True, False], dtype=wp.bool, device="cpu")) + + assert [env_ids.tolist() for env_ids in reset_env_ids] == [[0, 2]] + + def generate_multirotor( multirotor_cfg: MultirotorCfg, num_multirotors: int, device: str ) -> tuple[Multirotor, torch.Tensor]: diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..ce57c9e7851f --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,2 @@ +* Changed ``isaaclab_experimental.utils.warp.resolve_1d_mask`` to be + :func:`isaaclab.utils.warp.utils.resolve_1d_mask` instead of a copy of it. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.skip b/source/isaaclab_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst new file mode 100644 index 000000000000..03346fb5b32e --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.added.rst @@ -0,0 +1,15 @@ +* Added :class:`~isaaclab_experimental.utils.CapturedStage` and the :func:`~isaaclab_experimental.utils.captured` + decorator, which run a Warp manager or environment stage eagerly or through CUDA graphs recorded from its own + calls. A stage records on its first call without an eager warm-up, records again when an argument array or + scalar changes or the physics backend rebinds its buffers, and runs eagerly on CPU or while + :attr:`CapturedStage.enabled ` is False. The Warp + environments enable capture. +* Added :func:`~isaaclab_experimental.utils.eager` for calls that cannot be recorded inside a captured stage. + While the stage records, the call ends the current CUDA graph and runs eagerly before the next graph starts. + Every replay repeats the call at the same place with the same arguments. +* Added the ``ISAACLAB_SYNC_DEBUG`` environment variable: ``ISAACLAB_SYNC_DEBUG=1`` raises on a hidden + GPU-to-host synchronization inside an eagerly executed Warp environment stage. +* Added a boolean ``env_mask`` argument to :meth:`~isaaclab_experimental.envs.ManagerBasedEnvWarp.reset`. +* Added parameter predicates to :class:`~isaaclab_experimental.utils.warp.WarpCapturable`: + ``@WarpCapturable(lambda params: ...)`` marks a term capturable only for the parameters the predicate + accepts. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst new file mode 100644 index 000000000000..7eb3c5d84f4a --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst @@ -0,0 +1,29 @@ +* **Breaking:** Removed ``ManagerCallSwitch``, ``ManagerCallMode`` and the ``MANAGER_CALL_CONFIG`` + environment variable. Warp environments construct their Warp managers directly. +* **Breaking:** Removed ``WarpGraphCache``. The Warp managers and environments record their stages through + :class:`~isaaclab_experimental.utils.CapturedStage`, and the environments call the managers as the stable + environments do. +* **Breaking:** Removed ``ManagerBasedRLEnvWarp.step_warp_termination_compute``. The environment calls + :meth:`TerminationManager.compute ` directly. +* **Breaking:** Removed ``DirectRLEnvWarp.step_warp_action``. :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` + calls the task methods in the order of the stable :class:`~isaaclab.envs.DirectRLEnv` and no longer records them. + To keep recording them as CUDA graphs, decorate a task's ``_apply_action``, ``_get_dones``, ``_get_rewards``, + ``_get_observations`` and ``_reset_idx`` with :func:`~isaaclab_experimental.utils.captured`. +* **Breaking:** Changed the Warp + :meth:`ActionManager.process_action ` to take a + torch tensor, as the stable :meth:`~isaaclab.managers.ActionManager.process_action` does. +* Changed the Warp managers to decide CUDA graph capture per term: an operation records its capturable terms + and runs each term decorated with ``@WarpCapturable(False)``, and each observation term with history, eagerly + between the recorded parts, instead of running the whole operation eagerly. A class term's ``reset`` is + recorded unless the ``reset`` method itself is decorated with ``@WarpCapturable(False)``. +* Changed the Warp manager-based environments to reset environments by boolean mask. The mask is converted + to environment indices once per reset, and only when a command, curriculum or recorder term is active. +* Changed :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` to skip its reset stage on steps where no + environment terminated, deciding on the GPU (a conditional CUDA graph node, CUDA 12.4+), so a step no longer + synchronizes with the host. With ``compute_final_obs``, ``extras["final_obs"]`` is now set on every step. +* Changed the Warp ``push_by_setting_velocity``, ``apply_external_force_torque``, ``reset_root_state_uniform`` + and ``randomize_rigid_body_com`` event terms to classes that allocate their buffers when the term is + created. They read their ranges on every call, so a changed range applies to the next event. +* **Breaking:** Renamed :attr:`SceneEntityCfg.joint_mask ` + to ``joint_mask_wp``, matching the ``_wp`` suffix of its other Warp fields, and added ``body_mask_wp``, the + boolean mask of the selected bodies. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst new file mode 100644 index 000000000000..f295460afdc2 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.fixed.rst @@ -0,0 +1,23 @@ +* Fixed Warp environments replaying CUDA graphs that read freed simulation buffers after the simulation was + reset or stopped. The stages now record again on their next call. +* Fixed the Warp environments applying actions, stepping the simulation and updating the scene once per physics + step when the physics backend runs the whole decimation loop. Like the stable environments, they now do it + once per environment step, so contact histories, air times and the rewards and terminations reading them + match the stable environments. +* Fixed :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` and + :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` ignoring ``compute_final_obs``. A step that resets + environments now stores the observations from before the reset in ``extras["final_obs"]``, and the + environments declare Same-Step autoreset in their metadata, like the stable environments. +* Fixed :class:`~isaaclab_experimental.envs.ManagerBasedEnvWarp`, + :class:`~isaaclab_experimental.envs.ManagerBasedRLEnvWarp` and + :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` returning ``extras["log"]`` values that later steps + overwrite. Loggers that keep the logs of an iteration and average them, such as the rsl_rl logger, reported + the latest value for every step. Each step now returns copies, like the stable environments. +* Fixed :attr:`~isaaclab_experimental.utils.buffers.CircularBuffer.max_length` synchronizing with the GPU. +* Fixed the Warp ``randomize_rigid_body_com`` event accumulating offsets across calls. It now offsets the + center of mass the term finds on its first call, like the stable term. +* Fixed the Warp ``apply_external_force_torque`` event clearing the permanent wrenches of reset environments + when its force and torque ranges are zero. +* Fixed :meth:`EventManager.set_term_cfg ` and + :meth:`RewardManager.set_term_cfg ` replaying + recorded stages with the replaced term. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.major b/source/isaaclab_experimental/changelog.d/jichuanh-warp-frontend-cleanup.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst new file mode 100644 index 000000000000..aa642663b908 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst @@ -0,0 +1,4 @@ +* Changed :class:`~isaaclab_experimental.envs.ManagerBasedEnvWarp` and + :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` to initialize the process-wide + :class:`~isaaclab.utils.seed.WarpRng` state, which the Warp MDP terms, managers and noise models now draw from. + An environment created without a seed now seeds it from :func:`torch.initial_seed` instead of ``-1``. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.major b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.major new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.removed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.removed.rst new file mode 100644 index 000000000000..8e486bedf7bd --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-rng-shared.removed.rst @@ -0,0 +1,2 @@ +* **Breaking:** Removed ``ManagerBasedEnvWarp.rng_state_wp``. Read :attr:`~isaaclab.utils.seed.WarpRng.state` + instead. diff --git a/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst new file mode 100644 index 000000000000..b712f5b4a8d4 --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst @@ -0,0 +1,2 @@ +* Fixed :meth:`~isaaclab_experimental.envs.InteractiveSceneWarp.reset` with ``env_mask`` resetting deformable + objects and sensors in every environment, and skipping cable objects. 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 2c032cea3c22..b906771136e7 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 @@ -9,7 +9,6 @@ import inspect import logging import math -import os from abc import abstractmethod from dataclasses import MISSING from typing import Any, ClassVar @@ -32,10 +31,11 @@ from isaaclab.sim.utils import use_stage from isaaclab.utils import instantiate, validate from isaaclab.utils.noise import NoiseModel -from isaaclab.utils.seed import configure_seed +from isaaclab.utils.seed import WarpRng, configure_seed from isaaclab.utils.timer import Timer -from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache +from isaaclab_experimental.utils.warp import count_masked +from isaaclab_experimental.utils.warp_capture import CapturedStage, captured, reset_captured_stages # from isaacsim.core.simulation_manager import SimulationManager # from isaacsim.core.version import get_version @@ -44,12 +44,6 @@ # import logger logger = logging.getLogger(__name__) -DEBUG_TIMER_STEP = os.environ.get("DEBUG_TIMER_STEP", "0") == "1" -"""Enable outer step() timer only. Set DEBUG_TIMER_STEP=1 env var to enable.""" - -DEBUG_TIMERS = os.environ.get("DEBUG_TIMERS", "0") == "1" -"""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( @@ -86,6 +80,10 @@ class DirectRLEnvWarp(DirectRLEnv): here and leave it up to library-defined wrappers to take care of wrapping this environment for their agents. + Tasks record their GPU work as CUDA graphs by decorating :meth:`_apply_action`, :meth:`_get_dones`, + :meth:`_get_rewards`, :meth:`_get_observations` and :meth:`_reset_idx` with + :func:`~isaaclab_experimental.utils.captured`. + Note: For vectorized environments, it is recommended to **only** call the :meth:`reset` method once before the first call to :meth:`step`, i.e. after the environment is created. @@ -98,6 +96,7 @@ class DirectRLEnvWarp(DirectRLEnv): """Whether the environment is a vectorized environment.""" metadata: ClassVar[dict[str, Any]] = { "render_modes": [None, "human", "rgb_array"], + "autoreset_mode": gym.vector.AutoresetMode.SAME_STEP, # "isaac_sim_version": get_version(), } """Metadata for the environment.""" @@ -125,7 +124,7 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # set the seed for the environment if self.cfg.seed is not None: - self.cfg.seed = self.seed(self.cfg.seed) + self.cfg.seed = configure_seed(self.cfg.seed) else: logger.warning("Seed not set for the environment. The environment creation may not be deterministic.") @@ -166,6 +165,9 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # attach_stage_to_usd_context() logger.info(f"Scene manager: {self.scene}") + # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors. + WarpRng.initialize(self.num_envs, self.device) + # 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 @@ -191,6 +193,10 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # this shouldn't cause an issue since later on, users do a reset over all the # environments so the lazy buffers would be reset. self.scene.update(dt=self.physics_dt) + # let the physics backend know about the env decimation so it can + # fold the full loop into a single step() when possible + self.sim.backend_manager.set_decimation(self.cfg.decimation) + self._physics_handles_decimation = self.sim.backend_manager.handles_decimation() # check if debug visualization is has been implemented by the environment source_code = inspect.getsource(self._set_debug_vis_impl) @@ -208,6 +214,8 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # allocate dictionary to store metrics self.extras = {} + # the log the task writes; the environment returns copies of it (see :meth:`step`) + self._task_log: dict | None = None # initialize data and constants # -- counter for simulation steps @@ -220,6 +228,8 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.reset_terminated = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) self.reset_time_outs = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) self.reset_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + # the number of environments to reset, counted on the device so the step does not synchronize with the host + self._reset_count_wp = wp.zeros(1, dtype=wp.int32, device=self.device) self._ALL_ENV_MASK = wp.ones(self.num_envs, dtype=wp.bool, device=self.device) # Expected bindings: @@ -229,8 +239,9 @@ 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() + # the task stages record CUDA graphs, and record again after the physics buffers rebind + CapturedStage.enabled = True + CapturedStage.invalidate_on(self.sim.backend_manager) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() @@ -345,7 +356,7 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) """ # set the seed if seed is not None: - self.seed(seed) + configure_seed(seed) # reset state of scene self._reset_idx(self._ALL_ENV_MASK) @@ -367,7 +378,6 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) self.obs_buf = {"policy": self.torch_obs_buf.clone()} return self.obs_buf, self.extras - @Timer(name="env_step", msg="Step took:", enable=DEBUG_TIMER_STEP or DEBUG_TIMERS) def step(self, action: torch.Tensor) -> VecEnvStepReturn: """Execute one time-step of the environment's dynamics. @@ -393,59 +403,82 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: A tuple containing the observations, rewards, resets (terminated and truncated) and extras. """ + # task code that runs during this step writes into its own log, not into the copy returned last step + if self._task_log is not None: + self.extras["log"] = self._task_log action = action.to(self.device) # add action noise if self.cfg.action_noise_model: action = self._action_noise_model(action) - # process actions, #TODO pass the torch tensor directly. - with Timer(name="pre_physics", msg="Pre-physics step took:", enable=DEBUG_TIMERS): - self._pre_physics_step( - wp.from_torch(action) - ) # Creates a tensor and discards it. Not graphable unless training loop reuses the same pointer. + # process actions + self._pre_physics_step(wp.from_torch(action)) # check if we need to do rendering within the physics loop # note: hoisted out of the decimation loop; is_rendering does live settings lookups is_rendering = self.sim.is_rendering - # perform physics stepping - with Timer(name="physics_loop", msg="Physics loop took:", enable=DEBUG_TIMERS): - for _ in range(self.cfg.decimation): - self._sim_step_counter += 1 - # 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) - - # 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. - with Timer(name="write_data_to_sim_loop", msg="Write data to sim (loop) took:", enable=DEBUG_TIMERS): - self.scene.write_data_to_sim() - - with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMERS): - self.sim.step(render=False) - # 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. - if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: - self.sim.render() - # update buffers at sim dt - with Timer(name="scene_update", msg="Scene update took:", enable=DEBUG_TIMERS): - self.scene.update(dt=self.physics_dt) + # physics-owned decimation covers all substeps in one call + steps_per_call = self.cfg.decimation if self._physics_handles_decimation else 1 + for _ in range(self.cfg.decimation // steps_per_call): + self._sim_step_counter += steps_per_call + self._apply_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. + self.scene.write_data_to_sim() + self.sim.step(render=False) + # 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. + if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: + self.sim.render() + self.scene.update(dt=self.physics_dt * steps_per_call) + + # post-step: + # -- update env counters (used for curriculum generation) + wp.launch(add_to_env, dim=self.num_envs, inputs=[self._episode_length_buf_wp, 1]) # step in current episode 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) + + self._get_dones() + self._get_rewards() + + # capture the terminal observation before reset and expose it for Same-Step autoreset. + # whether any environment resets is only known on the device, so it is kept on every step + if self.cfg.compute_final_obs: + self._get_observations() + self.extras["final_obs"] = {"policy": self.torch_obs_buf.clone()} + # -- reset envs that terminated/timed-out + self._reset_terminated() + # if sensors are added to the scene, make sure we render to reflect changes in reset + # if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: + # self.sim.render() # write_data_to_sim runs uncaptured — it uses torch ops that cross CUDA streams. - 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.scene.write_data_to_sim() + + # post-step: step interval event + # if self.cfg.events: + # if "interval" in self.event_manager.available_modes: + # self.event_manager.apply(mode="interval", dt=self.step_dt) + + self._get_observations() + + # add observation noise + # note: we apply no noise to the state space (since it is used for critic networks) + # if self.cfg.observation_noise_model: + # self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"]) # Visualization hook — runs after CUDA graph scope. Override in subclass # to update markers or other non-graphable visual elements. - with Timer(name="visualize", msg="Visualize took:", enable=DEBUG_TIMERS): - self._post_step_visualize() + self._post_step_visualize() + + # tasks log persistent buffers that later steps overwrite, and RL libraries keep the log of every step + # until they average it, so hand out copies like the fresh tensors of stable tasks + self._task_log = self.extras.get("log") + if self._task_log: + self.extras["log"] = { + key: value.clone() if torch.is_tensor(value) else value for key, value in self._task_log.items() + } # return observations, rewards, resets and extras # store the returned buffer so RslRlVecEnvWrapper.get_observations() can read env.obs_buf @@ -467,49 +500,6 @@ def _post_step_visualize(self) -> None: """ pass - 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. - - def _step_warp_end_pre(self) -> None: - """Capturable portion before write_data_to_sim (pure warp kernels).""" - wp.launch( - add_to_env, - dim=self.num_envs, - inputs=[ - self._episode_length_buf_wp, - 1, - ], - ) - self._get_dones() - self._get_rewards() - - # -- reset envs that terminated/timed-out and log the episode information - self._reset_idx(mask=self.reset_buf) - - def _step_warp_end_post(self) -> None: - """Capturable portion after write_data_to_sim (pure warp kernels).""" - # if sensors are added to the scene, make sure we render to reflect changes in reset - # if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset: - # self.sim.render() - - # TODO We could split it out. - # post-step: step interval event - # if self.cfg.events: - # if "interval" in self.event_manager.available_modes: - # self.event_manager.apply(mode="interval", dt=self.step_dt) - - # update observations - self._get_observations() - - # add observation noise - # note: we apply no noise to the state space (since it is used for critic networks) - # if self.cfg.observation_noise_model: - # self.obs_buf["policy"] = self._observation_noise_model(self.obs_buf["policy"]) - @staticmethod def seed(seed: int = -1) -> int: """Set the seed for the environment. @@ -577,6 +567,8 @@ def close(self): if self.cfg.events: del self.event_manager del self.scene + CapturedStage.disable() + reset_captured_stages(self) # # clear callbacks and instance # if float(".".join(get_version()[2])) >= 5: @@ -665,6 +657,17 @@ def _configure_gym_env_spaces(self): # instantiate actions (needed for tasks for which the observations computation is dependent on the actions) self.actions = sample_space(self.single_action_space, self.sim.device, batch_size=self.num_envs, fill_value=0) + @captured + def _reset_terminated(self) -> None: + """Reset the environments in :attr:`reset_buf`, deciding on the device whether any must reset. + + Recorded, the decision is a conditional graph node (CUDA 12.4+), so the step does not synchronize with the host; + run eagerly, :func:`warp.capture_if` reads the count on the host. + """ + self._reset_count_wp.zero_() + wp.launch(count_masked, dim=self.num_envs, inputs=[self.reset_buf, self._reset_count_wp], device=self.device) + wp.capture_if(self._reset_count_wp, on_true=lambda: self._reset_idx(mask=self.reset_buf)) + def _reset_idx(self, mask: wp.array | None = None): """Reset environments based on a boolean mask. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py index cf2183c9310c..e97a796eadda 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/frontend.py @@ -182,7 +182,7 @@ def adapt_cfg(cls, cfg: Any) -> None: 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). + ``joint_mask_wp``, ``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 @@ -247,7 +247,7 @@ def _promote_scene_entity_cfgs(cls, cfg: Any) -> None: 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`` / + stay valid; the new fields (``joint_mask_wp`` / ``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 diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py index f792b994f67b..5fa1a2eda227 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/interactive_scene_warp.py @@ -24,21 +24,26 @@ class InteractiveSceneWarp(InteractiveScene): 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. + .. caution:: + If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. + 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) + # -- assets for articulation in self._articulations.values(): articulation.reset(env_ids, env_mask=env_mask) + for cable_object in self._cable_objects.values(): + cable_object.reset(env_ids, env_mask=env_mask) for deformable_object in self._deformable_objects.values(): - deformable_object.reset(env_ids) + deformable_object.reset(env_ids, env_mask=env_mask) 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) + # -- sensors for sensor in self._sensors.values(): - sensor.reset(env_ids) + sensor.reset(env_ids, env_mask=env_mask) 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 8b6710633ae4..932020581dc0 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 @@ -33,31 +31,22 @@ export_articulations_data, export_scene_data, ) +from isaaclab.managers import RecorderManager from isaaclab.sim import SimulationContext from isaaclab.sim.utils import use_stage from isaaclab.utils import validate -from isaaclab.utils.seed import configure_seed +from isaaclab.utils.seed import WarpRng, configure_seed from isaaclab.utils.timer import Timer 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.warp import resolve_1d_mask +from isaaclab_experimental.utils.warp_capture import CapturedStage # import logger logger = logging.getLogger(__name__) -@wp.kernel -def initialize_rng_state( - # input - seed: wp.int32, - # output - state: wp.array(dtype=wp.uint32), -): - env_id = wp.tid() - state[env_id] = wp.rand_init(seed, wp.int32(env_id)) - - class ManagerBasedEnvWarp: """The base environment for the manager-based workflow (experimental fork). @@ -89,15 +78,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg): ) # 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() # set the seed for the environment if self.cfg.seed is not None: - self.cfg.seed = self.seed(self.cfg.seed) + self.cfg.seed = configure_seed(self.cfg.seed) else: logger.warning("Seed not set for the environment. The environment creation may not be deterministic.") @@ -117,6 +101,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg): if "cuda" in self.device: torch.cuda.set_device(self.device) + # the managers' stages record CUDA graphs, and record again after the physics buffers rebind + CapturedStage.enabled = True + CapturedStage.invalidate_on(self.sim.backend_manager) + # print useful information logger.info( "Base environment:\n" @@ -149,22 +137,14 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # attach_stage_to_usd_context() logger.info(f"Scene manager: {self.scene}") - # Shared per-env Warp RNG state (accessible to all managers/terms via `env`). - # This is a single stream per env (no lookup) and is initialized once when `num_envs` is known. - self.rng_state_wp = wp.zeros((self.num_envs,), dtype=wp.uint32, device=self.device) - seed_val = int(self.cfg.seed) if self.cfg.seed is not None else -1 - wp.launch( - kernel=initialize_rng_state, - dim=self.num_envs, - inputs=[seed_val, self.rng_state_wp], - device=self.device, - ) + # Process-wide per-env Warp RNG state (isaaclab.utils.seed.WarpRng), shared by terms, tasks and sensors. + WarpRng.initialize(self.num_envs, self.device) - # TODO(jichuanh): this is problematic as warp capture requires stable pointers, - # using different masks for different managers/terms will cause problems. # 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 stages read this mask; every reset selection is copied into it. + 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 +152,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: @@ -193,6 +173,10 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # this shouldn't cause an issue since later on, users do a reset over all the environments so the lazy # buffers would be reset. self.scene.update(dt=self.physics_dt) + # let the physics backend know about the env decimation so it can + # fold the full loop into a single step() when possible + self.sim.backend_manager.set_decimation(self.cfg.decimation) + self._physics_handles_decimation = self.sim.backend_manager.handles_decimation() # TODO(jichuanh): This is a temporary solution for event_manager only, but it should be general for all managers # Resolve SceneEntityCfg-dependent term params once before any captured event paths. @@ -379,17 +363,13 @@ def load_managers(self): # -- event manager (we print it here to make the logging consistent) logger.info(f"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) logger.info(f"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) logger.info(f"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) logger.info(f"Observation Manager: {self.observation_manager}") # perform events at the start of the simulation @@ -415,11 +395,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. @@ -431,30 +415,30 @@ def reset( Note: This argument is used for compatibility with Gymnasium environment definition. + env_mask: Boolean mask of the environments to reset, of shape ``(num_envs,)``. Takes precedence + over :paramref:`env_ids`. Defaults to None. + 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) + if env_mask is not None and self.recorder_manager.active_terms: + # the stable recorder manager hooks take indices; a mask takes precedence over the caller's ids + env_mask = self.resolve_env_mask(env_mask=env_mask) + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + env_mask = self.resolve_env_mask(env_ids=env_ids, env_mask=env_mask) # trigger recorder terms for pre-reset calls self.recorder_manager.record_pre_reset(env_ids) # set the seed if seed is not None: - used_seed = self.seed(seed) + # also reseeds the per-env Warp RNG state in place (stable pointer for capture) + used_seed = configure_seed(seed) # keep cfg seed in sync for downstream users self.cfg.seed = used_seed - # re-initialize per-env Warp RNG state without reallocating (stable pointer for capture) - wp.launch( - kernel=initialize_rng_state, - dim=self.num_envs, - inputs=[int(used_seed), self.rng_state_wp], - device=self.device, - ) # reset state of scene - self._reset_idx(env_ids) + self._reset_mask(env_mask) # update articulation kinematics self.scene.write_data_to_sim() @@ -506,9 +490,9 @@ def reset_to( # set the seed if seed is not None: - self.seed(seed) + configure_seed(seed) - self._reset_idx(env_ids) + self._reset_mask(self.resolve_env_mask(env_ids=env_ids)) # set the state self.scene.reset_to(state, env_ids, is_relative=is_relative) @@ -546,13 +530,7 @@ 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) + self.action_manager.process_action(action.to(self.device)) self.recorder_manager.record_pre_step() @@ -560,9 +538,10 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: # note: hoisted out of the decimation loop; is_rendering does live settings lookups is_rendering = self.sim.is_rendering - # perform physics stepping - for _ in range(self.cfg.decimation): - self._sim_step_counter += 1 + # physics-owned decimation covers all substeps in one call + steps_per_call = self.cfg.decimation if self._physics_handles_decimation else 1 + for _ in range(self.cfg.decimation // steps_per_call): + self._sim_step_counter += steps_per_call # set actions into buffers self.action_manager.apply_action() # set actions into simulator @@ -574,8 +553,7 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: # If a camera needs rendering at a faster frequency, this will lead to unexpected behavior. if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: self.sim.render() - # update buffers at sim dt - self.scene.update(dt=self.physics_dt) + self.scene.update(dt=self.physics_dt * steps_per_call) # post-step: step interval event if "interval" in self.event_manager.available_modes: @@ -609,6 +587,7 @@ def close(self): del self.event_manager del self.recorder_manager del self.scene + CapturedStage.disable() # self.sim.clear_all_callbacks() self.sim.clear_instance() @@ -623,117 +602,36 @@ 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: - items = "\n".join(f" - {item}" for item in replaced_items) - logger.info(f"Applied stable term config profile for managers:\n{items}") - - def _reset_idx(self, env_ids: Sequence[int]): - """Reset environments based on specified indices. + def _reset_mask(self, env_mask: wp.array): + """Reset the environments selected by a boolean mask. Args: - env_ids: List of environment ids which must be reset + env_mask: Boolean mask of the environments to reset, of shape ``(num_envs,)``. """ + # recorded reset stages read the environment-owned mask + if env_mask is not self.reset_mask_wp: + wp.copy(self.reset_mask_wp, env_mask) + env_mask = self.reset_mask_wp + # the stable recorder manager resets by index + env_ids = ( + wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) if self.recorder_manager.active_terms else None + ) + # 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 + 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) self.extras["log"].update(info) @@ -746,3 +644,8 @@ def _reset_idx(self, env_ids: Sequence[int]): # -- recorder manager info = self.recorder_manager.reset(env_ids) self.extras["log"].update(info) + # the managers log persistent buffers that the next reset overwrites, and RL libraries keep the log of + # every step until they average it, so hand out copies like the stable managers' fresh tensors + self.extras["log"] = { + key: value.clone() if torch.is_tensor(value) else value for key, value in self.extras["log"].items() + } 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 be522e073a2b..a4dfde185259 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,8 +15,6 @@ import logging import math -import os -from collections.abc import Sequence from typing import Any, ClassVar import gymnasium as gym @@ -26,26 +24,14 @@ from isaaclab.envs.common import VecEnvStepReturn from isaaclab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg -from isaaclab.managers import CommandManager -from isaaclab.utils import index_fill_ -from isaaclab.utils.timer import Timer +from isaaclab.managers import CommandManager, CurriculumManager -from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode -from isaaclab_experimental.utils.torch_utils import clone_obs_buffer +from isaaclab_experimental.managers import RewardManager, TerminationManager from .manager_based_env_warp import ManagerBasedEnvWarp logger = logging.getLogger(__name__) -DEBUG_TIMERS = os.environ.get("DEBUG_TIMERS", "0") == "1" -"""Enable outer step() timer. Set DEBUG_TIMERS=1 env var to enable.""" - -DEBUG_TIMER_STEP = os.environ.get("DEBUG_TIMER_STEP", "0") == "1" -"""Enable step sub-phase timers. Set DEBUG_TIMER_STEP=1 env var to enable.""" - -DEBUG_TIMER_RESET = os.environ.get("DEBUG_TIMER_RESET", "0") == "1" -"""Enable reset sub-phase timers. Set DEBUG_TIMER_RESET=1 env var to enable.""" - class ManagerBasedRLEnvWarp(ManagerBasedEnvWarp, gym.Env): """The superclass for the manager-based workflow reinforcement learning-based environments. @@ -81,6 +67,7 @@ class ManagerBasedRLEnvWarp(ManagerBasedEnvWarp, gym.Env): """Whether the environment is a vectorized environment.""" metadata: ClassVar[dict[str, Any]] = { "render_modes": [None, "human", "rgb_array"], + "autoreset_mode": gym.vector.AutoresetMode.SAME_STEP, # "isaac_sim_version": get_version(), } """Metadata for the environment.""" @@ -117,16 +104,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 @@ -166,7 +143,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) + # -- command manager (stable implementation) self.command_manager = CommandManager(self.cfg.commands, self) logger.info(f"Command Manager: {self.command_manager}") @@ -175,17 +152,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) logger.info(f"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) + self.reward_manager = RewardManager(self.cfg.rewards, self) logger.info(f"Reward Manager: {self.reward_manager}") - # -- curriculum manager - self.curriculum_manager = self._manager_call_switch.resolve_manager_class("CurriculumManager")( - self.cfg.curriculum, self - ) + # -- curriculum manager (stable implementation) + self.curriculum_manager = CurriculumManager(self.cfg.curriculum, self) logger.info(f"Curriculum Manager: {self.curriculum_manager}") # setup the action and observation spaces for Gym @@ -215,20 +188,6 @@ def setup_manager_visualizers(self): Operations - MDP """ - def invalidate_wp_graphs(self) -> None: - """Invalidate all cached Warp graphs. - - Call this if the captured launch topology changes (e.g. different term list, shapes, etc.). - """ - self._manager_call_switch.invalidate_graphs() - - def step_warp_termination_compute(self) -> None: - """Captured stage: compute terminations (env-step frequency).""" - self.reset_buf = self.termination_manager.compute() - self.reset_terminated = self.termination_manager.terminated - self.reset_time_outs = self.termination_manager.time_outs - - @Timer(name="env_step", msg="Step took:", enable=DEBUG_TIMER_STEP, time_unit="us") def step(self, action: torch.Tensor) -> VecEnvStepReturn: """Execute one time-step of the environment's dynamics and reset terminated environments. @@ -249,20 +208,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: A tuple containing the observations, rewards, resets (terminated and truncated) and extras. """ # process actions - # NOTE: keep a persistent action input buffer for graph pointer stability. - # 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) - 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}}, - timer=DEBUG_TIMER_STEP, - ) + self.action_manager.process_action(action.to(self.device)) self.recorder_manager.record_pre_step() @@ -270,95 +216,60 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: # note: checked here once to avoid multiple checks within the loop is_rendering = self.sim.is_rendering - # perform physics stepping - for _ in range(self.cfg.decimation): - self._sim_step_counter += 1 + # physics-owned decimation covers all substeps in one call + steps_per_call = self.cfg.decimation if self._physics_handles_decimation else 1 + for _ in range(self.cfg.decimation // steps_per_call): + self._sim_step_counter += steps_per_call # 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}, - timer=DEBUG_TIMER_STEP, - ) + self.action_manager.apply_action() + # scene writes cross the actuator and asset host boundaries, so they stay eager + 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.sim.step(render=False) 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. if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering: self.sim.render() - # update buffers at sim dt - with Timer( - name="scene.update", - msg="Scene.update took:", - enable=DEBUG_TIMER_STEP, - time_unit="us", - ): - self.scene.update(dt=self.physics_dt) + self.scene.update(dt=self.physics_dt * steps_per_call) # post-step: # -- update env counters (used for curriculum generation) self.episode_length_buf += 1 # step in current episode (per env) 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}, - 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)}}, - timer=DEBUG_TIMER_STEP, - ) + # -- post-processing: check terminations + self.reset_buf = self.termination_manager.compute() + self.reset_terminated = self.termination_manager.terminated + self.reset_time_outs = self.termination_manager.time_outs + # -- reward computation + self.reward_buf = self.reward_manager.compute(dt=self.step_dt) if len(self.recorder_manager.active_terms) > 0: # 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}}, - timer=DEBUG_TIMER_STEP, - ) + self.observation_manager.compute(return_cloned_output=False) 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: + # the reset pipeline advances curriculum and logging state even for an empty mask, so one + # host synchronization per step decides whether it runs + if self.reset_buf.any().item(): + reset_mask = self.termination_manager.dones_wp + # the stable recorder manager hooks take indices + reset_env_ids = ( + wp.to_torch(reset_mask).nonzero(as_tuple=False).squeeze(-1) + if self.recorder_manager.active_terms + else None + ) + # capture the terminal observation before reset and expose it for Same-Step autoreset. + # the reset overwrites the persistent observation buffers, so keep a copy + if self.cfg.compute_final_obs: + self.extras["final_obs"] = self.observation_manager.compute() # 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) + self._reset_mask(reset_mask) # 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: @@ -369,27 +280,15 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.recorder_manager.record_post_reset(reset_env_ids) # -- update command - self.command_manager.compute(dt=float(self.step_dt)) + self.command_manager.compute(dt=self.step_dt) # -- 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)}}, - timer=DEBUG_TIMER_STEP, - ) + self.event_manager.apply(mode="interval", dt=self.step_dt) # -- 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), - }, - timer=DEBUG_TIMER_STEP, - ) + self.obs_buf = self.observation_manager.compute(update_history=True) # return observations, rewards, resets and extras return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras @@ -482,120 +381,58 @@ 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( - 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 + def _reset_mask(self, env_mask: wp.array): + """Reset the environments selected by a boolean mask. 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 mask of the environments to reset, of shape ``(num_envs,)``. """ - 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)}).") + # recorded reset stages read the environment-owned mask + if env_mask is not self.reset_mask_wp: + wp.copy(self.reset_mask_wp, env_mask) + env_mask = self.reset_mask_wp + # the stable command, curriculum and recorder managers reset by index + env_ids = ( + wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + if self.command_manager.active_terms + or self.curriculum_manager.active_terms + or self.recorder_manager.active_terms + else None + ) # 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) + 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}}, - timer=DEBUG_TIMER_RESET, - ) + 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, - }, - }, - timer=DEBUG_TIMER_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. # -- 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}}, - 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}}, - 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}}, - timer=DEBUG_TIMER_RESET, - ) + obs_info = self.observation_manager.reset(env_mask=env_mask) + action_info = self.action_manager.reset(env_mask=env_mask) + reward_info = self.reward_manager.reset(env_mask=env_mask) # -- 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.curriculum_manager.reset(env_ids=env_ids) # -- 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}}, - 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}}, - timer=DEBUG_TIMER_RESET, - ) + event_info = self.event_manager.reset(env_mask=env_mask) + termination_info = self.termination_manager.reset(env_mask=env_mask) # -- recorder manager recorder_info = self.recorder_manager.reset(env_ids=env_ids) # reset the episode length buffer - index_fill_(self.episode_length_buf, env_ids, 0) + self.episode_length_buf.masked_fill_(wp.to_torch(env_mask), 0) # aggregate logging info log: dict[str, Any] = {} @@ -610,4 +447,6 @@ def _reset_idx( recorder_info, ): log.update(info) - self.extras["log"] = log + # the managers log persistent buffers that the next reset overwrites, and RL libraries keep the log of + # every step until they average it, so hand out copies like the stable managers' fresh tensors + self.extras["log"] = {key: value.clone() if torch.is_tensor(value) else value for key, value in log.items()} diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index 90e5696a6b10..8e7cf9c9ffe0 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -17,6 +17,10 @@ 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. +Terms that need scratch buffers are classes that allocate them in ``__init__``: a recorded stage runs its +terms for the first time inside CUDA graph capture, where an allocation would become graph-owned memory. +Ranges are read from the call arguments, so a changed term configuration takes effect on the next call. + 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. @@ -31,15 +35,22 @@ 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.utils.seed import WarpRng -from isaaclab_experimental.managers import ManagerTermBase as _WarpManagerTermBase -from isaaclab_experimental.managers import SceneEntityCfg +from isaaclab_experimental.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg from isaaclab_experimental.utils.warp import WarpCapturable if TYPE_CHECKING: - from isaaclab.assets import Articulation + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedEnv + +def _vec3_range(ranges: dict[str, tuple[float, float]], keys: tuple[str, str, str]) -> tuple[wp.vec3f, wp.vec3f]: + """Lower and upper bounds of three named ranges; a missing key is the zero range.""" + bounds = [ranges.get(key, (0.0, 0.0)) for key in keys] + return wp.vec3f(*(bound[0] for bound in bounds)), wp.vec3f(*(bound[1] for bound in bounds)) + + # --------------------------------------------------------------------------- # Randomize rigid body center of mass # --------------------------------------------------------------------------- @@ -48,67 +59,86 @@ @wp.kernel def _randomize_com_kernel( env_mask: wp.array(dtype=wp.bool), + needs_default: 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.""" + """Offset the default center of mass of the selected bodies by one random offset per environment.""" env_id = wp.tid() + if needs_default[0]: + for k in range(body_ids.shape[0]): + default_body_com_pos_b[env_id, body_ids[k]] = body_com_pos_b[env_id, body_ids[k]] 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 the center of mass (CoM) of rigid bodies around their default positions. - 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. + Warp-first override of :class:`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. """ - 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 + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize the term. + + Args: + cfg: The configuration of the event term. + env: The environment instance. + """ + super().__init__(cfg, env) + asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self._asset: RigidObject | Articulation = env.scene[asset_cfg.name] + # Other startup terms may move the CoM after this term is built, so the defaults are copied on the + # first call. Randomizing from them keeps repeated calls from accumulating offsets. + self._default_com = wp.zeros_like(self._asset.data.body_com_pos_b.warp) + self._needs_default_com = wp.ones(1, dtype=wp.bool, device=env.device) - 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, - ) - - # 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) + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array, + com_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ): + """Offset the selected environments' CoM positions [m] by a sample from ``com_range``.""" + body_com_pos_b = self._asset.data.body_com_pos_b.warp + com_lo, com_hi = _vec3_range(com_range, ("x", "y", "z")) + wp.launch( + kernel=_randomize_com_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + self._needs_default_com, + WarpRng.state, + self._default_com, + body_com_pos_b, + asset_cfg.body_ids_wp, + com_lo, + com_hi, + ], + device=env.device, + ) + # the kernel only reads the flag, so it is cleared after the launch + self._needs_default_com.zero_() + # Notify the solver that inertial properties changed (COM position affects inertia). + self._asset.set_coms_mask(coms=body_com_pos_b, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -120,6 +150,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, @@ -127,16 +158,14 @@ def _apply_external_force_torque_kernel( torque_lo: float, torque_hi: float, ): + """Sample uniform forces [N] and torques [N·m] for the selected bodies of the selected environments.""" 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 k in range(body_ids.shape[0]): + b = body_ids[k] force_out[env_id, b] = wp.vec3f( wp.randf(state, force_lo, force_hi), wp.randf(state, force_lo, force_hi), @@ -150,50 +179,57 @@ 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. +class apply_external_force_torque(ManagerTermBase): + """Randomize the 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 - ) - apply_external_force_torque._is_warmed_up = True + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize the term. - 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, - ) + Args: + cfg: The configuration of the event term. + env: The environment instance. + """ + 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) - 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, - ) + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array, + force_range: tuple[float, float], + torque_range: tuple[float, float], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ): + """Set the selected bodies' wrenches to samples from ``force_range`` [N] and ``torque_range`` [N·m].""" + # Zero ranges configure a placeholder term. Skipping it keeps wrenches that other sources wrote to the + # permanent composer instead of clearing them on every reset. + if tuple(force_range) == (0.0, 0.0) and tuple(torque_range) == (0.0, 0.0): + return + wp.launch( + kernel=_apply_external_force_torque_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + WarpRng.state, + asset_cfg.body_ids_wp, + self._forces, + self._torques, + float(force_range[0]), + float(force_range[1]), + float(torque_range[0]), + float(torque_range[1]), + ], + device=env.device, + ) + self._asset.permanent_wrench_composer.set_forces_and_torques_mask( + forces=self._forces, torques=self._torques, body_mask=asset_cfg.body_mask_wp, env_mask=env_mask + ) # --------------------------------------------------------------------------- @@ -231,45 +267,49 @@ 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. +class push_by_setting_velocity(ManagerTermBase): + """Push the asset by adding a random velocity to its root velocity. 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 + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize the term. - 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, - ) + Args: + cfg: The configuration of the event term. + env: The environment instance. + """ + super().__init__(cfg, env) + self._asset: RigidObject | Articulation = env.scene[cfg.params["asset_cfg"].name] + self._velocity = wp.zeros(env.num_envs, dtype=wp.spatial_vectorf, device=env.device) - asset.write_root_velocity_to_sim_mask(root_velocity=push_by_setting_velocity._scratch_vel, env_mask=env_mask) + def __call__( + self, + env: ManagerBasedEnv, + env_mask: wp.array, + velocity_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ): + """Add samples from ``velocity_range`` [m/s, rad/s] to the selected environments' root velocity.""" + lin_lo, lin_hi = _vec3_range(velocity_range, ("x", "y", "z")) + ang_lo, ang_hi = _vec3_range(velocity_range, ("roll", "pitch", "yaw")) + wp.launch( + kernel=_push_by_setting_velocity_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + WarpRng.state, + self._asset.data.root_vel_w.warp, + self._velocity, + lin_lo, + lin_hi, + ang_lo, + ang_hi, + ], + device=env.device, + ) + self._asset.write_root_velocity_to_sim_mask(root_velocity=self._velocity, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -341,62 +381,61 @@ 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"), -): +class reset_root_state_uniform(ManagerTermBase): """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 + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + """Initialize the term. - 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, - ) + Args: + cfg: The configuration of the event term. + env: The environment instance. + """ + super().__init__(cfg, env) + self._asset: RigidObject | Articulation = env.scene[cfg.params["asset_cfg"].name] + 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) - 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) + def __call__( + self, + env: ManagerBasedEnv, + 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 selected environments' root pose [m, rad] and velocity [m/s, rad/s].""" + pos_lo, pos_hi = _vec3_range(pose_range, ("x", "y", "z")) + rot_lo, rot_hi = _vec3_range(pose_range, ("roll", "pitch", "yaw")) + vel_lin_lo, vel_lin_hi = _vec3_range(velocity_range, ("x", "y", "z")) + vel_ang_lo, vel_ang_hi = _vec3_range(velocity_range, ("roll", "pitch", "yaw")) + wp.launch( + kernel=_reset_root_state_uniform_kernel, + dim=env.num_envs, + inputs=[ + env_mask, + WarpRng.state, + self._asset.data.default_root_pose.warp, + self._asset.data.default_root_vel.warp, + env.env_origins_wp, + self._pose, + self._velocity, + pos_lo, + pos_hi, + rot_lo, + rot_hi, + vel_lin_lo, + vel_lin_hi, + vel_ang_lo, + vel_ang_hi, + ], + device=env.device, + ) + 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) # --------------------------------------------------------------------------- @@ -470,11 +509,6 @@ def reset_joints_by_offset( f"but got None for asset '{asset_cfg.name}'. " "Use isaaclab_experimental.managers.SceneEntityCfg and ensure joint_names are set." ) - if not hasattr(env, "rng_state_wp") or env.rng_state_wp is None: - raise AttributeError( - "reset_joints_by_offset requires env.rng_state_wp to be initialized. " - "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment." - ) wp.launch( kernel=_reset_joints_by_offset_kernel, @@ -482,7 +516,7 @@ def reset_joints_by_offset( inputs=[ env_mask, asset_cfg.joint_ids_wp, - env.rng_state_wp, + WarpRng.state, asset.data.default_joint_pos.warp, asset.data.default_joint_vel.warp, asset.data.joint_pos.warp, @@ -561,11 +595,6 @@ def reset_joints_by_scale( f"but got None for asset '{asset_cfg.name}'. " "Use isaaclab_experimental.managers.SceneEntityCfg and ensure joint_names are set." ) - if not hasattr(env, "rng_state_wp") or env.rng_state_wp is None: - raise AttributeError( - "reset_joints_by_scale requires env.rng_state_wp to be initialized. " - "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment." - ) wp.launch( kernel=_reset_joints_by_scale_kernel, @@ -573,7 +602,7 @@ def reset_joints_by_scale( inputs=[ env_mask, asset_cfg.joint_ids_wp, - env.rng_state_wp, + WarpRng.state, asset.data.default_joint_pos.warp, asset.data.default_joint_vel.warp, asset.data.joint_pos.warp, @@ -600,9 +629,9 @@ def reset_joints_by_scale( # 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. +# class terms expect torch env indices. Both terms are marked non-capturable, so +# in reset or interval mode the event stages run them eagerly and the host-side +# mask-to-ids conversion stays outside any recorded graph. def _mask_to_env_ids(env_mask: wp.array) -> torch.Tensor: @@ -610,12 +639,13 @@ def _mask_to_env_ids(env_mask: wp.array) -> torch.Tensor: return torch.nonzero(wp.to_torch(env_mask), as_tuple=False).squeeze(-1) -class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, _WarpManagerTermBase): +@WarpCapturable(False, reason="converts the environment mask to indices on the host") +class randomize_rigid_body_material(_StableRandomizeRigidBodyMaterial, ManagerTermBase): """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. + term's env-ids convention and delegates. Inherits the warp :class:`ManagerTermBase` + so the warp managers accept it as a class term. """ def __call__( @@ -641,16 +671,17 @@ def __call__( ) def reset(self, env_mask: wp.array | None = None) -> None: - """Nothing to reset; materials are randomized once at startup.""" + """Nothing to reset; the term keeps no per-environment state.""" return -class randomize_rigid_body_mass(_StableRandomizeRigidBodyMass, _WarpManagerTermBase): +@WarpCapturable(False, reason="converts the environment mask to indices on the host") +class randomize_rigid_body_mass(_StableRandomizeRigidBodyMass, ManagerTermBase): """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. + term's env-ids convention and delegates. Inherits the warp :class:`ManagerTermBase` + so the warp managers accept it as a class term. """ def __call__( @@ -676,5 +707,5 @@ def __call__( ) def reset(self, env_mask: wp.array | None = None) -> None: - """Nothing to reset; masses are randomized once at startup.""" + """Nothing to reset; the term keeps no per-environment state.""" return diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index 4cc597989021..db9488e09d92 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -268,7 +268,7 @@ def joint_torques_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Sc wp.launch( kernel=_sum_sq_masked_kernel, dim=env.num_envs, - inputs=[asset.actuators.applied_effort.warp, asset_cfg.joint_mask, out], + inputs=[asset.actuators.applied_effort.warp, asset_cfg.joint_mask_wp, out], device=env.device, ) @@ -292,7 +292,7 @@ def joint_vel_l1(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg) -> None wp.launch( kernel=_sum_abs_masked_kernel, dim=env.num_envs, - inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask_wp, out], device=env.device, ) @@ -303,7 +303,7 @@ def joint_vel_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneE wp.launch( kernel=_sum_sq_masked_kernel, dim=env.num_envs, - inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask_wp, out], device=env.device, ) @@ -314,7 +314,7 @@ def joint_acc_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneE wp.launch( kernel=_sum_sq_masked_kernel, dim=env.num_envs, - inputs=[asset.data.joint_acc.warp, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_acc.warp, asset_cfg.joint_mask_wp, out], device=env.device, ) @@ -341,7 +341,7 @@ def joint_deviation_l1(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = wp.launch( kernel=_sum_abs_diff_masked_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, asset_cfg.joint_mask_wp, out], device=env.device, ) @@ -379,7 +379,7 @@ def joint_pos_limits(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = Sc wp.launch( kernel=_joint_pos_limits_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos.warp, asset.data.soft_joint_pos_limits.warp, asset_cfg.joint_mask, out], + inputs=[asset.data.joint_pos.warp, asset.data.soft_joint_pos_limits.warp, asset_cfg.joint_mask_wp, out], device=env.device, ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py index 6d538e86bb1a..5b25f6382f75 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -199,12 +199,15 @@ def _joint_pos_out_of_manual_limit_kernel( upper: float, out: wp.array(dtype=wp.bool), ): - """2D kernel (num_envs, num_joints). ``out`` is pre-zeroed; only writes True.""" - i, j = wp.tid() - if joint_mask[j]: - v = joint_pos[i, j] - if v < lower or v > upper: - out[i] = True + """1D kernel (num_envs,): whether any selected joint is outside the bounds.""" + i = wp.tid() + violated = bool(False) + for j in range(joint_pos.shape[1]): + if joint_mask[j]: + v = joint_pos[i, j] + if v < lower or v > upper: + violated = True + out[i] = violated def joint_pos_out_of_manual_limit( @@ -212,20 +215,20 @@ def joint_pos_out_of_manual_limit( ) -> None: """Terminate when joint positions are outside configured bounds. Writes into ``out``.""" asset: Articulation = env.scene[asset_cfg.name] - if asset_cfg.joint_mask is None: + if asset_cfg.joint_mask_wp is None: raise ValueError( - f"joint_pos_out_of_manual_limit requires SceneEntityCfg with resolved joint_mask, " + f"joint_pos_out_of_manual_limit requires SceneEntityCfg with resolved joint_mask_wp, " f"but got None for asset '{asset_cfg.name}'." ) - if asset.data.joint_pos.warp.shape[1] != asset_cfg.joint_mask.shape[0]: + if asset.data.joint_pos.warp.shape[1] != asset_cfg.joint_mask_wp.shape[0]: raise ValueError( - f"joint_mask length ({asset_cfg.joint_mask.shape[0]}) does not match " + f"joint_mask_wp length ({asset_cfg.joint_mask_wp.shape[0]}) does not match " f"joint_pos dim ({asset.data.joint_pos.warp.shape[1]}) for asset '{asset_cfg.name}'." ) wp.launch( kernel=_joint_pos_out_of_manual_limit_kernel, - dim=(env.num_envs, asset.data.joint_pos.warp.shape[1]), - inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, bounds[0], bounds[1], out], + dim=env.num_envs, + inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask_wp, bounds[0], bounds[1], out], device=env.device, ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py index fa0ce31f4219..e201b2f68cbf 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py @@ -22,7 +22,10 @@ from isaaclab.managers.manager_term_cfg import ActionTermCfg from isaaclab.utils import instantiate -from .manager_base import ManagerBase, ManagerTermBase +from isaaclab_experimental.utils.warp import is_warp_capturable +from isaaclab_experimental.utils.warp_capture import captured, eager + +from .manager_base import ManagerBase, ManagerTermBase, split_terms if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -227,6 +230,15 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): self._action_torch = wp.to_torch(self._action) self._prev_action_torch = wp.to_torch(self._prev_action) + terms = list(self._terms.values()) + self._term_split = split_terms(terms, is_warp_capturable) + # an annotation on a term class covers its processing; its reset is judged by its own annotation + self._reset_split = split_terms(terms, lambda term: is_warp_capturable(term.reset)) + # each term's slice of the action vector + self._term_offsets = { + term: sum(other.action_dim for other in terms[:index]) for index, term in enumerate(terms) + } + # check if any term has debug visualization implemented self.cfg.debug_vis = False for term in self._terms.values(): @@ -366,57 +378,39 @@ def set_debug_vis(self, debug_vis: bool): for term in self._terms.values(): term.set_debug_vis(debug_vis) - def reset( - self, - env_ids: Sequence[int] | torch.Tensor | None = None, - *, - env_mask: wp.array | None = None, - ) -> dict[str, Any]: + @captured + def _reset(self, env_mask: wp.array) -> dict[str, Any]: """Resets the action history. Args: - env_ids: The specific environment indices to reset. - If None, all environments are considered. env_mask: Boolean Warp mask of shape (num_envs,) indicating which envs to reset. - If provided, takes precedence over ``env_ids``. 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) - # 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(): + for term in self._reset_split.eager: + eager(term.reset, env_mask=env_mask) + for term in self._reset_split.captured: term.reset(env_mask=env_mask) # nothing to log here return {} - def process_action(self, action: wp.array): + def process_action(self, action: torch.Tensor): """Processes the actions sent to the environment. Note: @@ -431,21 +425,29 @@ def process_action(self, action: wp.array): # store the input actions wp.copy(self._prev_action, self._action) - wp.copy(self._action, action) + wp.copy(self._action, wp.from_torch(action.float())) # split the actions and apply to each term - idx = 0 - for term in self._terms.values(): - term.process_actions(self._action, idx) - idx += term.action_dim + self._process_action() + @captured + def _process_action(self) -> None: + """Processes each term's slice of the stored actions.""" + for term in self._term_split.eager: + eager(term.process_actions, self._action, self._term_offsets[term]) + for term in self._term_split.captured: + term.process_actions(self._action, self._term_offsets[term]) + + @captured def apply_action(self) -> None: """Applies the actions to the environment/simulation. Note: This should be called at every simulation step. """ - for term in self._terms.values(): + for term in self._term_split.eager: + eager(term.apply_actions) + for term in self._term_split.captured: term.apply_actions() def get_term(self, name: str) -> ActionTerm: diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index e7d8b4b746b0..2b4cd8de1828 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -18,6 +18,7 @@ from isaaclab.managers.manager_term_cfg import CommandTermCfg from isaaclab.utils import instantiate +from isaaclab.utils.seed import WarpRng from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked @@ -306,8 +307,6 @@ def _resample(self, env_mask: wp.array): raise TypeError(f"env_mask must be a wp.array (got {type(env_mask)}).") if env_mask.dtype != wp.bool or env_mask.ndim != 1: raise TypeError(f"env_mask must be wp.bool 1D (got dtype={env_mask.dtype}, ndim={env_mask.ndim}).") - if self._env.rng_state_wp is None: - raise RuntimeError("Environment rng_state_wp is not initialized.") # resample time-left and increment command-counter for masked envs wp.launch( @@ -317,7 +316,7 @@ def _resample(self, env_mask: wp.array): env_mask, self.time_left_wp, self.command_counter_wp, - self._env.rng_state_wp, + WarpRng.state, float(self.cfg.resampling_time_range[0]), float(self.cfg.resampling_time_range[1]), ], diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py index 5c3832fb803b..b051f7e3d695 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py @@ -30,8 +30,12 @@ from prettytable import PrettyTable from isaaclab.managers.manager_term_cfg import EventTermCfg +from isaaclab.utils.seed import WarpRng -from .manager_base import ManagerBase +from isaaclab_experimental.utils.warp import is_warp_capturable +from isaaclab_experimental.utils.warp_capture import captured, eager, reset_captured_stages + +from .manager_base import ManagerBase, ManagerTermBase, split_resets logger = logging.getLogger(__name__) @@ -170,7 +174,6 @@ def __init__(self, cfg: object, env): # Warp buffers for interval/reset modes (populated in _prepare_terms) self._interval_term_time_left_wp: list[wp.array] = [] - self._interval_term_ranges: list[tuple[float, float]] = [] self._interval_term_is_global: list[bool] = [] # Scalar RNG state for global interval timers (allocated lazily if needed). self._interval_global_rng_state_wp: wp.array | None = None @@ -179,6 +182,7 @@ def __init__(self, cfg: object, env): self._reset_term_triggered_once_wp: list[wp.array] = [] super().__init__(cfg, env) + self._split_terms() # persistent scratch mask for per-term interval/reset triggering (must be stable pointer for capture) self._scratch_term_mask_wp = wp.zeros((self.num_envs,), dtype=wp.bool, device=self.device) @@ -223,14 +227,30 @@ def available_modes(self) -> list[str]: return list(self._mode_term_names.keys()) def set_term_cfg(self, term_name: str, cfg: EventTermCfg): + """Sets the configuration of the specified term into the manager. + + The method finds the term by name by searching through all the modes. + It then updates the configuration of the term with the first matching name. + + Args: + term_name: The name of the event term. + cfg: The configuration for the event term. + + Raises: + ValueError: If the term name is not found. + """ term_found = False for mode, terms in self._mode_term_names.items(): if term_name in terms: + cfg.capturable = is_warp_capturable(cfg.func, cfg.params) self._mode_term_cfgs[mode][terms.index(term_name)] = cfg term_found = True break if not term_found: raise ValueError(f"Event term '{term_name}' not found.") + # the class terms derive from the term configurations, and the recorded stages hold the replaced term + self._split_terms() + reset_captured_stages(self) def get_term_cfg(self, term_name: str) -> EventTermCfg: for mode, terms in self._mode_term_names.items(): @@ -238,26 +258,14 @@ def get_term_cfg(self, term_name: str) -> EventTermCfg: return self._mode_term_cfgs[mode][terms.index(term_name)] raise ValueError(f"Event term '{term_name}' not found.") - def reset( - self, - env_ids: Sequence[int] | slice | torch.Tensor | wp.array | None = None, - *, - 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) - + @captured + def _reset(self, env_mask: wp.array) -> dict[str, float]: + """Resets the class terms and the interval timers of the environments selected by a boolean mask.""" # reset class terms (mask-based) - for mode_cfg in self._mode_class_term_cfgs.values(): - for term_cfg in mode_cfg: - term_cfg.func.reset(env_mask=env_mask) + for term_cfg in self._class_term_split.eager: + eager(term_cfg.func.reset, env_mask=env_mask) + for term_cfg in self._class_term_split.captured: + term_cfg.func.reset(env_mask=env_mask) # reset interval timers for non-global interval events # note: terms with resample_interval_on_reset=False keep their per-environment counter @@ -266,16 +274,16 @@ def reset( for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]): if term_cfg.is_global_time or not term_cfg.resample_interval_on_reset: continue - lower, upper = self._interval_term_ranges[i] + lower, upper = term_cfg.interval_range_s wp.launch( kernel=_interval_reset_selected, dim=self.num_envs, inputs=[ env_mask, self._interval_term_time_left_wp[i], - self._env.rng_state_wp, - float(lower), - float(upper), + WarpRng.state, + lower, + upper, ], device=self.device, ) @@ -311,7 +319,7 @@ def apply( raise ValueError( f"Event mode '{mode}' does not require environment indices. This is an undefined behavior." ) - self._apply_interval(float(dt)) + self._apply_interval(dt) return # resolve the environment mask @@ -330,13 +338,11 @@ def apply( for term_cfg in self._mode_term_cfgs[mode]: term_cfg.func(self._env, env_mask_wp, **term_cfg.params) + @captured def _apply_interval(self, dt: float) -> None: - if self._env.rng_state_wp is None: - raise RuntimeError("EventManager._apply_interval: env.rng_state_wp is not initialized.") - # iterate over all the interval terms (fixed list; captured graph-friendly) for i, term_cfg in enumerate(self._mode_term_cfgs["interval"]): - lower, upper = self._interval_term_ranges[i] + lower, upper = term_cfg.interval_range_s if self._interval_term_is_global[i]: if self._interval_global_rng_state_wp is None: raise RuntimeError( @@ -350,36 +356,37 @@ def _apply_interval(self, dt: float) -> None: self._interval_term_time_left_wp[i], self._interval_global_rng_state_wp, self._scratch_interval_trigger_flag_wp, - float(dt), - float(lower), - float(upper), + dt, + lower, + upper, ], device=self.device, ) - term_cfg.func(self._env, self._scratch_interval_trigger_mask_view_wp, **term_cfg.params) + term_mask = self._scratch_interval_trigger_mask_view_wp else: wp.launch( kernel=_interval_step_per_env, dim=self.num_envs, inputs=[ self._interval_term_time_left_wp[i], - self._env.rng_state_wp, + WarpRng.state, self._scratch_term_mask_wp, - float(dt), - float(lower), - float(upper), + dt, + lower, + upper, ], device=self.device, ) - term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params) + term_mask = self._scratch_term_mask_wp + if term_cfg.capturable: + term_cfg.func(self._env, term_mask, **term_cfg.params) + else: + eager(term_cfg.func, self._env, term_mask, **term_cfg.params) + @captured def _apply_reset(self, env_mask_wp: wp.array, global_env_step_count_wp: wp.array) -> None: - if self._scratch_term_mask_wp is None: - raise RuntimeError("EventManager._apply_reset: _scratch_term_mask_wp is not initialized.") - # iterate over all the reset terms for index, term_cfg in enumerate(self._mode_term_cfgs["reset"]): - min_step_count = int(term_cfg.min_step_count_between_reset) wp.launch( kernel=_reset_compute_valid_mask, dim=self.num_envs, @@ -389,11 +396,22 @@ def _apply_reset(self, env_mask_wp: wp.array, global_env_step_count_wp: wp.array self._reset_term_triggered_once_wp[index], self._scratch_term_mask_wp, global_env_step_count_wp, - int(min_step_count), + term_cfg.min_step_count_between_reset, ], device=self.device, ) - term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params) + if term_cfg.capturable: + term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params) + else: + eager(term_cfg.func, self._env, self._scratch_term_mask_wp, **term_cfg.params) + + def _split_terms(self) -> None: + """Collect the class terms of every mode and split them by whether their resets can be recorded.""" + self._mode_class_term_cfgs = { + mode: [cfg for cfg in cfgs if inspect.isclass(cfg.func) or isinstance(cfg.func, ManagerTermBase)] + for mode, cfgs in self._mode_term_cfgs.items() + } + self._class_term_split = split_resets([cfg for cfgs in self._mode_class_term_cfgs.values() for cfg in cfgs]) def _prepare_terms(self): # check if config is dict already @@ -430,14 +448,10 @@ def _prepare_terms(self): if term_cfg.mode not in self._mode_term_names: self._mode_term_names[term_cfg.mode] = [] self._mode_term_cfgs[term_cfg.mode] = [] - self._mode_class_term_cfgs[term_cfg.mode] = [] # add term name and parameters self._mode_term_names[term_cfg.mode].append(term_name) self._mode_term_cfgs[term_cfg.mode].append(term_cfg) - if inspect.isclass(term_cfg.func): - self._mode_class_term_cfgs[term_cfg.mode].append(term_cfg) - # per-mode Warp buffers if term_cfg.mode == "interval": if term_cfg.interval_range_s is None: @@ -445,25 +459,22 @@ def _prepare_terms(self): f"Event term '{term_name}' has mode 'interval' but 'interval_range_s' is not specified." ) lower, upper = term_cfg.interval_range_s - self._interval_term_ranges.append((float(lower), float(upper))) if term_cfg.is_global_time: # allocate and seed scalar global RNG state if needed (avoid consuming env0 RNG stream) if self._interval_global_rng_state_wp is None: - if self._env.rng_state_wp is None: - raise RuntimeError("EventManager._prepare_terms: env.rng_state_wp is not initialized.") self._interval_global_rng_state_wp = wp.zeros((1,), dtype=wp.uint32, device=self.device) wp.launch( kernel=_seed_global_rng_from_env_rng, dim=1, - inputs=[self._env.rng_state_wp, self._interval_global_rng_state_wp], + inputs=[WarpRng.state, self._interval_global_rng_state_wp], device=self.device, ) time_left = wp.zeros((1,), dtype=wp.float32, device=self.device) wp.launch( kernel=_interval_init_global, dim=1, - inputs=[time_left, self._interval_global_rng_state_wp, float(lower), float(upper)], + inputs=[time_left, self._interval_global_rng_state_wp, lower, upper], device=self.device, ) self._interval_term_time_left_wp.append(time_left) @@ -473,7 +484,7 @@ def _prepare_terms(self): wp.launch( kernel=_interval_init_per_env, dim=self.num_envs, - inputs=[time_left, self._env.rng_state_wp, float(lower), float(upper)], + inputs=[time_left, WarpRng.state, lower, upper], device=self.device, ) self._interval_term_time_left_wp.append(time_left) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py index 0607fe111f20..36f5cd4a033c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/manager_base.py @@ -20,8 +20,8 @@ import inspect import logging from abc import ABC, abstractmethod -from collections.abc import Sequence -from typing import TYPE_CHECKING, Any +from collections.abc import Callable, Sequence +from typing import TYPE_CHECKING, Any, NamedTuple import warp as wp @@ -31,6 +31,7 @@ from isaaclab.utils import string_to_callable, to_dict from isaaclab_experimental.utils.warp import is_warp_capturable +from isaaclab_experimental.utils.warp_capture import reset_captured_stages from .scene_entity_cfg import SceneEntityCfg @@ -46,6 +47,38 @@ logger = logging.getLogger(__name__) +class TermSplit(NamedTuple): + """Terms of one operation, split by whether they can be recorded into a CUDA graph.""" + + eager: list + """Terms that run eagerly, before the recorded part.""" + + captured: list + """Terms recorded with the rest of the operation.""" + + +def split_terms(terms: Sequence[Any], capturable: Callable[[Any], bool] = lambda term: term.capturable) -> TermSplit: + """Split terms by capturability, keeping their order. + + Args: + terms: The terms (term configurations by default). + capturable: Whether a term can be recorded. Defaults to the term configuration's ``capturable`` flag. + """ + return TermSplit([term for term in terms if not capturable(term)], [term for term in terms if capturable(term)]) + + +def split_resets(term_cfgs: Sequence[Any]) -> TermSplit: + """Split class terms by whether their ``reset`` can be recorded, keeping their order. + + An annotation on a term class covers its ``__call__``; the ``reset`` is judged by the annotation on the + ``reset`` method itself. + + Args: + term_cfgs: The configurations of the class terms. + """ + return split_terms(term_cfgs, lambda term_cfg: is_warp_capturable(term_cfg.func.reset, term_cfg.params)) + + class ManagerTermBase(ABC): """Base class for manager terms. @@ -237,17 +270,34 @@ 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] | torch.Tensor | None = None, + *, + env_mask: wp.array | None = None, + ) -> dict[str, float]: """Resets the manager and returns logging information for the current time-step. + The ids resolve to a mask on the host, which a CUDA graph cannot record, so the recorded part is :meth:`_reset`. + Args: env_ids: The environment ids for which to log data. Defaults None, which logs data for all environments. + env_mask: Boolean Warp mask of shape (num_envs,) selecting the environments. + If provided, takes precedence over ``env_ids``. Returns: Dictionary containing the logging information. """ - return {} + # 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( + f"{type(self).__name__}.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) + return self._reset(env_mask) def find_terms(self, name_keys: str | Sequence[str]) -> list[str]: """Find terms in the manager based on the names. @@ -295,6 +345,10 @@ def _prepare_terms(self): """Prepare terms information from the configuration object.""" raise NotImplementedError + def _reset(self, env_mask: wp.array) -> dict[str, float]: + """Resets the environments selected by a boolean mask of shape (num_envs,) and returns logging information.""" + return {} + """ Internal callbacks. """ @@ -324,6 +378,8 @@ def _resolve_terms_callback(self, event): # set the flag self._is_scene_entities_resolved = True + # class terms were just created + reset_captured_stages(self) """ Internal functions. @@ -407,11 +463,8 @@ 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) + # whether the term can be recorded into a CUDA graph, resolved once per term + term_cfg.capturable = is_warp_capturable(term_cfg.func, term_cfg.params) # process attributes at runtime # these properties are only resolvable once the simulation starts playing diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index b58bfeebe788..1cbec8abbd19 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -62,8 +62,9 @@ from isaaclab_experimental.utils import noise, tensor_processors from isaaclab_experimental.utils.buffers import CircularBuffer from isaaclab_experimental.utils.torch_utils import clone_obs_buffer +from isaaclab_experimental.utils.warp_capture import captured, eager -from .manager_base import ManagerBase, ManagerTermBase +from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -194,6 +195,14 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): # Stores the latest observations. self._obs_buffer: dict[str, torch.Tensor | dict[str, torch.Tensor]] | None = None + # (group, index) of every term, and the class terms, split by capturability + terms = [ + (group_name, index) for group_name, cfgs in self._group_obs_term_cfgs.items() for index in range(len(cfgs)) + ] + self._term_split = split_terms(terms, lambda term: self._group_obs_term_cfgs[term[0]][term[1]].capturable) + self._class_term_split = split_resets( + [cfg for cfgs in self._group_obs_class_term_cfgs.values() for cfg in cfgs] + ) # Note: Persistent Warp output buffers (`_group_out_wp` / `_group_out_torch`) and per-term post-processing # buffers are allocated during `_prepare_terms()` since they are per-term/per-group setup. @@ -385,30 +394,17 @@ def _collect_io_descriptors(self, group_names_to_export: list[str] = ["policy"]) Operations. """ - def reset( - self, - env_ids: Sequence[int] | torch.Tensor | None = None, - *, - 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) - + @captured + def _reset(self, env_mask: wp.array) -> dict[str, float]: + """Resets the observation terms of the environments selected by a boolean mask of shape (num_envs,).""" # call all terms that are classes - for group_name, group_cfg in self._group_obs_class_term_cfgs.items(): - for term_cfg in group_cfg: - term_cfg.func.reset(env_mask=env_mask) - # reset terms with history - for term_name in self._group_obs_term_names[group_name]: - if term_name in self._group_obs_term_history_buffer[group_name]: - self._group_obs_term_history_buffer[group_name][term_name].reset(env_mask=env_mask) + for term_cfg in self._class_term_split.eager: + eager(term_cfg.func.reset, env_mask=env_mask) + for term_cfg in self._class_term_split.captured: + term_cfg.func.reset(env_mask=env_mask) + # reset terms with history + if any(self._group_obs_term_history_buffer.values()): + eager(self._reset_history, env_mask) # call all tensor processors/noise models that are classes for tensor_processor in self._group_obs_class_instances: tensor_processor.reset(env_mask=env_mask) @@ -416,6 +412,16 @@ def reset( # nothing to log here return {} + def _reset_history(self, env_mask: wp.array) -> None: + """Resets the history buffers of the selected environments, which are not capture-safe. + + The buffers are looked up at call time because :meth:`_compute_term` can replace them, while :func:`eager` + replays its recorded arguments. + """ + for history_buffers in self._group_obs_term_history_buffer.values(): + for circular_buffer in history_buffers.values(): + circular_buffer.reset(env_mask=env_mask) + def compute( self, update_history: bool = False, return_cloned_output: bool = True ) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]: @@ -436,9 +442,18 @@ def compute( The observations are either concatenated into a single tensor or returned as a dictionary with keys corresponding to the term's name. """ + obs_buffer = self._compute(update_history) + # clone outside the recorded stage + return clone_obs_buffer(obs_buffer) if return_cloned_output else obs_buffer + + @captured + def _compute(self, update_history: bool) -> dict[str, torch.Tensor | dict[str, torch.Tensor]]: + """Compute every term into the persistent buffers and return the persistent observation buffer.""" # Launch kernels for every group (writes into persistent buffers in-place). - for group_name in self._group_obs_term_names: - self.compute_group(group_name, update_history=update_history) + for group_name, index in self._term_split.eager: + eager(self._compute_term, group_name, index, update_history=update_history) + for group_name, index in self._term_split.captured: + self._compute_term(group_name, index, update_history=update_history) # Build the obs buffer once (persistent refs to in-place-updated tensors/dicts). if self._obs_buffer is None: self._obs_buffer = { @@ -449,8 +464,6 @@ def compute( ) for group_name in self._group_obs_term_names } - if return_cloned_output: - return clone_obs_buffer(self._obs_buffer) return self._obs_buffer def compute_group(self, group_name: str, update_history: bool = False) -> torch.Tensor | dict[str, torch.Tensor]: @@ -497,74 +510,78 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch. # iterate over all the terms in each group group_term_names = self._group_obs_term_names[group_name] - # Persistent per-term obs dict (pre-allocated in _prepare_terms). - group_obs = self._group_obs_dict[group_name] - # evaluate terms: compute, apply tensor processors, add noise, clip, scale - for term_name, term_cfg in zip(group_term_names, self._group_obs_term_cfgs[group_name]): - # compute term's value into pre-allocated Warp output - term_cfg.func(self._env, term_cfg.out_wp, **term_cfg.params) - - # apply custom tensor processors (in-place on out_wp) - if term_cfg.tensor_processors is not None: - for tensor_processor_cfg in term_cfg.tensor_processors: - tensor_processor_cfg.func(term_cfg.out_wp, **tensor_processor_cfg.params) - - # apply noise (Warp in-place on out_wp) - if isinstance(term_cfg.noise, noise.NoiseCfg): - term_cfg.noise.func(term_cfg.out_wp, term_cfg.noise) - elif isinstance(term_cfg.noise, noise.NoiseModelCfg) and term_cfg.noise.func is not None: - term_cfg.noise.func(term_cfg.out_wp) - - # clip then scale (stable semantics); implementation may use Warp kernels - if term_cfg.clip is not None: - wp.launch( - kernel=_apply_clip, - dim=self.num_envs, - inputs=[term_cfg.out_wp, float(term_cfg.clip[0]), float(term_cfg.clip[1])], - device=self.device, - ) - if term_cfg.scale is not None: - wp.launch( - kernel=_apply_scale, - dim=self.num_envs, - inputs=[term_cfg.out_wp, term_cfg.scale_wp], - device=self.device, - ) - - # TODO(jichuanh): This is not migrated yet. Need revisit. - # Update the history buffer if observation term has history enabled - if term_cfg.history_length > 0: - # circular buffer is not capture safe - if wp.get_device().is_capturing: - raise RuntimeError( - "Observation terms with history (circular buffer) are not CUDA-graph-capture-safe yet. " - "Disable history for observation terms used inside a captured graph, or restructure " - "the graph to exclude history-buffered terms." - ) - circular_buffer = self._group_obs_term_history_buffer[group_name][term_name] - if update_history: - circular_buffer.append(wp.to_torch(term_cfg.out_wp)) - elif circular_buffer._buffer is None: - # because circular buffer only exits after the simulation steps, - # this guards history buffer from corruption by external calls before simulation start - circular_buffer = CircularBuffer( - max_len=circular_buffer.max_length, - batch_size=circular_buffer.batch_size, - device=circular_buffer.device, - ) - self._group_obs_term_history_buffer[group_name][term_name] = circular_buffer - circular_buffer.append(wp.to_torch(term_cfg.out_wp)) - - if term_cfg.flatten_history_dim: - group_obs[term_name] = circular_buffer.buffer.reshape(self._env.num_envs, -1) - else: - group_obs[term_name] = circular_buffer.buffer + for index in range(len(group_term_names)): + self._compute_term(group_name, index, update_history=update_history) # return persistent output (updated in-place by kernels above) if self._group_use_warp_concat[group_name]: return self._group_out_torch[group_name] - return group_obs + return self._group_obs_dict[group_name] + + def _compute_term(self, group_name: str, index: int, update_history: bool = False) -> None: + """Compute one observation term into its slice of the group buffer.""" + term_name = self._group_obs_term_names[group_name][index] + term_cfg = self._group_obs_term_cfgs[group_name][index] + group_obs = self._group_obs_dict[group_name] + # compute term's value into pre-allocated Warp output + term_cfg.func(self._env, term_cfg.out_wp, **term_cfg.params) + + # apply custom tensor processors (in-place on out_wp) + if term_cfg.tensor_processors is not None: + for tensor_processor_cfg in term_cfg.tensor_processors: + tensor_processor_cfg.func(term_cfg.out_wp, **tensor_processor_cfg.params) + + # apply noise (Warp in-place on out_wp) + if isinstance(term_cfg.noise, noise.NoiseCfg): + term_cfg.noise.func(term_cfg.out_wp, term_cfg.noise) + elif isinstance(term_cfg.noise, noise.NoiseModelCfg) and term_cfg.noise.func is not None: + term_cfg.noise.func(term_cfg.out_wp) + + # clip then scale (stable semantics); implementation may use Warp kernels + if term_cfg.clip is not None: + wp.launch( + kernel=_apply_clip, + dim=self.num_envs, + inputs=[term_cfg.out_wp, float(term_cfg.clip[0]), float(term_cfg.clip[1])], + device=self.device, + ) + if term_cfg.scale is not None: + wp.launch( + kernel=_apply_scale, + dim=self.num_envs, + inputs=[term_cfg.out_wp, term_cfg.scale_wp], + device=self.device, + ) + + # TODO(jichuanh): This is not migrated yet. Need revisit. + # Update the history buffer if observation term has history enabled + if term_cfg.history_length > 0: + # circular buffer is not capture safe + if wp.get_device().is_capturing: + raise RuntimeError( + "Observation terms with history (circular buffer) are not CUDA-graph-capture-safe yet. " + "Disable history for observation terms used inside a captured graph, or restructure " + "the graph to exclude history-buffered terms." + ) + circular_buffer = self._group_obs_term_history_buffer[group_name][term_name] + if update_history: + circular_buffer.append(wp.to_torch(term_cfg.out_wp)) + elif circular_buffer._buffer is None: + # because circular buffer only exits after the simulation steps, + # this guards history buffer from corruption by external calls before simulation start + circular_buffer = CircularBuffer( + max_len=circular_buffer.max_length, + batch_size=circular_buffer.batch_size, + device=circular_buffer.device, + ) + self._group_obs_term_history_buffer[group_name][term_name] = circular_buffer + circular_buffer.append(wp.to_torch(term_cfg.out_wp)) + + if term_cfg.flatten_history_dim: + group_obs[term_name] = circular_buffer.buffer.reshape(self._env.num_envs, -1) + else: + group_obs[term_name] = circular_buffer.buffer def serialize(self) -> dict: """Serialize the observation term configurations for all active groups. @@ -704,6 +721,8 @@ def _prepare_terms(self): # noqa: C901 if group_cfg.history_length is not None: term_cfg.history_length = group_cfg.history_length term_cfg.flatten_history_dim = group_cfg.flatten_history_dim + # history buffers are not capture-safe yet + term_cfg.capturable = term_cfg.capturable and term_cfg.history_length == 0 # add term config to list self._group_obs_term_names[group_name].append(term_name) self._group_obs_term_cfgs[group_name].append(term_cfg) @@ -786,14 +805,8 @@ 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 - term_cfg.noise.rng_state_wp = self._env.rng_state_wp noise_model_cls = term_cfg.noise.class_type if not issubclass(noise_model_cls, noise.NoiseModel): raise TypeError( diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index 7dd6b85bd6ed..381bad1e0386 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -20,9 +20,11 @@ from isaaclab.managers.manager_term_cfg import RewardTermCfg +from isaaclab_experimental.utils.warp import is_warp_capturable from isaaclab_experimental.utils.warp.kernels import compute_reset_scale, count_masked +from isaaclab_experimental.utils.warp_capture import captured, eager, reset_captured_stages -from .manager_base import ManagerBase, ManagerTermBase +from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv @@ -68,29 +70,12 @@ def _reward_finalize( val = weighted * dt total += val episode_sums[term_idx, env_id] += val + else: + step_reward[env_id, term_idx] = 0.0 reward_buf[env_id] = total -@wp.kernel -def _reward_pre_compute_reset( - # output - reward_buf: wp.array(dtype=wp.float32), - step_reward: wp.array(dtype=wp.float32, ndim=2), - term_outs: wp.array(dtype=wp.float32, ndim=2), -): - """Reset per-step reward buffers. - - Launched with dim = (num_envs,) to reset `reward_buf` and clear the corresponding row in `step_reward`. - This works even when `step_reward.shape[1] == 0` (no terms). - """ - env_id = wp.tid() - reward_buf[env_id] = 0.0 - for term_idx in range(term_outs.shape[0]): - step_reward[env_id, term_idx] = 0.0 - term_outs[term_idx, env_id] = 0.0 - - class RewardManager(ManagerBase): """Manager for computing reward signals for a given world. @@ -128,6 +113,9 @@ 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)} + # skip terms with zero weight (kind of a micro-optimization) + self._term_split = split_terms([term_cfg for term_cfg in self._term_cfgs if term_cfg.weight != 0.0]) + self._class_term_split = split_resets(self._class_term_cfgs) num_terms = len(self._term_names) self._num_terms = num_terms @@ -226,32 +214,16 @@ def active_terms(self) -> list[str]: Operations. """ - def reset( - self, - env_ids: Sequence[int] | torch.Tensor | None = None, - *, - env_mask: wp.array | None = None, - ) -> dict[str, torch.Tensor]: + @captured + def _reset(self, env_mask: wp.array) -> dict[str, torch.Tensor]: """Computes/reset episodic reward sums for masked envs (capturable core). Args: - env_ids: The specific environment indices to reset. - If None, all environments are considered. env_mask: Boolean Warp mask of shape (num_envs,) selecting reset environments. - If provided, takes precedence over ``env_ids``. 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) - self._episode_sum_avg_wp.zero_() self._reset_count_wp.zero_() self._reset_scale_wp.zero_() @@ -276,13 +248,18 @@ def reset( # 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: + for term_cfg in self._class_term_split.eager: + term_extras = eager(term_cfg.func.reset, env_mask=env_mask) + if term_extras: + self._reset_extras.update(term_extras) + for term_cfg in self._class_term_split.captured: term_extras = term_cfg.func.reset(env_mask=env_mask) if term_extras: self._reset_extras.update(term_extras) return self._reset_extras + @captured def compute(self, dt: float) -> torch.Tensor: """Computes the reward signal as a weighted sum of individual terms. @@ -295,21 +272,11 @@ def compute(self, dt: float) -> torch.Tensor: Returns: The net reward signal of shape (num_envs,). """ - # TODO: Investigate performance diff between two .fill_ and kernel launch - # reset computation (Warp buffers) in a single kernel launch - wp.launch( - kernel=_reward_pre_compute_reset, - dim=self.num_envs, - inputs=[self._reward_wp, self._step_reward_wp, self._term_outs_wp], - 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: - continue - # compute term into the persistent warp buffer (raw, unweighted) - # NOTE: `out` is pre-zeroed every step by `_reward_pre_compute_reset`. + # compute each term into the persistent warp buffer (raw, unweighted) + for term_cfg in self._term_split.eager: + eager(term_cfg.func, self._env, term_cfg.out, **term_cfg.params) + for term_cfg in self._term_split.captured: term_cfg.func(self._env, term_cfg.out, **term_cfg.params) # update total reward, episodic sums and step rewards in a single kernel launch @@ -353,6 +320,10 @@ def set_term_cfg(self, term_name: str, cfg: RewardTermCfg): self._term_cfgs[term_idx] = cfg # keep on-device weights in sync (call this to update weights used in compute) self._term_weights_tensor_view[term_idx] = float(cfg.weight) + # recorded compute graphs hold the replaced term, which may also differ in capturability + cfg.capturable = is_warp_capturable(cfg.func, cfg.params) + self._term_split = split_terms([term_cfg for term_cfg in self._term_cfgs if term_cfg.weight != 0.0]) + reset_captured_stages(self) def get_term_cfg(self, term_name: str) -> RewardTermCfg: """Gets the configuration for the specified term. 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 f6d30489eeca..1f47d2c4facd 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/scene_entity_cfg.py @@ -26,10 +26,10 @@ class SceneEntityCfg(_SceneEntityCfg): """Scene entity configuration with an optional Warp joint mask. Notes: - - `joint_mask` is intended for Warp kernels only. + - `joint_mask_wp` is intended for Warp kernels only. """ - joint_mask: wp.array | None = None + joint_mask_wp: wp.array | None = None """Integer indices of selected joints — used for subset-sized gathers where a boolean mask cannot provide the mapping from output index k to joint index.""" @@ -38,6 +38,9 @@ class SceneEntityCfg(_SceneEntityCfg): """Integer indices of selected bodies — used for subset-sized body gathers.""" body_ids_wp: wp.array | None = None + """Boolean mask of selected bodies — used for masked per-body writes.""" + body_mask_wp: wp.array | None = None + @classmethod def from_stable(cls, stable: _SceneEntityCfg) -> SceneEntityCfg: """Build a warp scene-entity cfg from a stable one. @@ -63,11 +66,18 @@ def resolve(self, scene: InteractiveScene): mask_list = [False] * entity.num_joints for idx in joint_ids_list: mask_list[idx] = True - self.joint_mask = wp.array(mask_list, dtype=wp.bool, device=scene.device) + self.joint_mask_wp = wp.array(mask_list, dtype=wp.bool, device=scene.device) self.joint_ids_wp = wp.array(joint_ids_list, dtype=wp.int32, device=scene.device) - # -- Warp body ids + # -- Warp body ids / mask if self.body_ids is not None and self.body_ids != slice(None): - self.body_ids_wp = wp.array(list(self.body_ids), dtype=wp.int32, device=scene.device) + body_ids_list = list(self.body_ids) elif hasattr(entity, "num_bodies"): - self.body_ids_wp = wp.array(list(range(entity.num_bodies)), dtype=wp.int32, device=scene.device) + body_ids_list = list(range(entity.num_bodies)) + else: + return + self.body_ids_wp = wp.array(body_ids_list, dtype=wp.int32, device=scene.device) + if hasattr(entity, "num_bodies"): + selected = set(body_ids_list) + mask_list = [idx in selected for idx in range(entity.num_bodies)] + self.body_mask_wp = wp.array(mask_list, dtype=wp.bool, device=scene.device) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index 368d9ca5d130..9191679aec12 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -24,33 +24,14 @@ from isaaclab.managers.manager_term_cfg import TerminationTermCfg -from .manager_base import ManagerBase, ManagerTermBase +from isaaclab_experimental.utils.warp_capture import captured, eager + +from .manager_base import ManagerBase, ManagerTermBase, split_resets, split_terms if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv -@wp.kernel -def _termination_pre_compute_reset( - # output - term_dones: wp.array(dtype=wp.bool, ndim=2), - truncated: wp.array(dtype=wp.bool), - terminated: wp.array(dtype=wp.bool), - dones: wp.array(dtype=wp.bool), -): - """Reset per-step termination buffers. - - Launched with dim = (num_envs,) to reset per-env flags and clear the corresponding row in `term_dones`. - This works even when `term_dones.shape[1] == 0` (no terms). - """ - env_id = wp.tid() - truncated[env_id] = False - terminated[env_id] = False - dones[env_id] = False - for term_idx in range(term_dones.shape[1]): - term_dones[env_id, term_idx] = False - - @wp.kernel def _termination_finalize( # input @@ -125,6 +106,8 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): super().__init__(cfg, env) self._term_name_to_term_idx = {name: i for i, name in enumerate(self._term_names)} + self._term_split = split_terms(self._term_cfgs) + self._class_term_split = split_resets(self._class_term_cfgs) # persistent buffers (Warp) num_terms = len(self._term_names) @@ -243,31 +226,16 @@ def term_dones_wp(self) -> wp.array: Operations. """ - def reset( - self, - env_ids: Sequence[int] | torch.Tensor | None = None, - *, - env_mask: wp.array | None = None, - ) -> dict[str, torch.Tensor]: + @captured + def _reset(self, env_mask: wp.array) -> dict[str, torch.Tensor]: """Reset termination stats and class terms; return pre-allocated extras. Args: - env_ids: The specific environment indices to reset. - If None, all environments are considered. env_mask: Boolean Warp mask of shape (num_envs,) selecting reset environments. - If provided, takes precedence over ``env_ids``. 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) if len(self._term_names) > 0: self._term_done_avg_wp.zero_() wp.launch( @@ -276,7 +244,9 @@ def reset( inputs=[self._last_episode_dones_wp, self._term_done_avg_wp], device=self.device, ) - for term_cfg in self._class_term_cfgs: + for term_cfg in self._class_term_split.eager: + eager(term_cfg.func.reset, env_mask=env_mask) + for term_cfg in self._class_term_split.captured: term_cfg.func.reset(env_mask=env_mask) return self._reset_extras @@ -285,22 +255,17 @@ def episode_termination_extras(self) -> dict[str, torch.Tensor]: """Pre-allocated reset logging extras for termination terms.""" return self._reset_extras + @captured def compute(self) -> torch.Tensor: """Computes the termination signal as union of individual terms. Returns: The combined termination signal of shape (num_envs,). """ - # reset computation (Warp buffers) in a single kernel launch - wp.launch( - kernel=_termination_pre_compute_reset, - dim=self.num_envs, - inputs=[self._term_dones_wp, self._truncated_wp, self._terminated_wp, self._dones_wp], - device=self.device, - ) - # iterate over all the termination terms (fixed list; per-term math is Warp) - for term_cfg in self._term_cfgs: + for term_cfg in self._term_split.eager: + eager(term_cfg.func, self._env, term_cfg.out, **term_cfg.params) + for term_cfg in self._term_split.captured: term_cfg.func(self._env, term_cfg.out, **term_cfg.params) # finalize dones and update last-episode term flags (single kernel launch) diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi index 8211d4785061..4f4de2a2e281 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/__init__.pyi @@ -4,9 +4,9 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ - "ManagerCallMode", - "ManagerCallSwitch", - "WarpGraphCache", + "CapturedStage", + "captured", + "eager", "clone_obs_buffer", "buffers", "modifiers", @@ -15,7 +15,6 @@ __all__ = [ "warp", ] -from .manager_call_switch import ManagerCallMode, ManagerCallSwitch from .torch_utils import clone_obs_buffer -from .warp_graph_cache import WarpGraphCache +from .warp_capture import CapturedStage, captured, eager from . import buffers, modifiers, tensor_processors, 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 3eba7590624c..ad9ba04950a5 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/buffers/circular_buffer.py @@ -44,6 +44,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 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` @@ -71,7 +74,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 a6dc7c449b89..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py +++ /dev/null @@ -1,265 +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 logging -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 - -logger = logging.getLogger(__name__) - - -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) - lines = [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 "" - lines.append(f" - {manager_name}: {mode}{cap_str}") - logger.info("ManagerCallSwitch configuration:\n" + "\n".join(lines)) - - # ------------------------------------------------------------------ - # 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/noise/noise_model.py b/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py index acba926ea488..cbc933324830 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/noise/noise_model.py @@ -10,10 +10,8 @@ noise_cfg.func(data_wp, noise_cfg) -> None -Random noise kernels (gaussian, uniform) consume the shared per-env Warp RNG state -(``rng_state_wp``) that is set on the config at manager prep time from -``env.rng_state_wp``. See :func:`initialize_rng_state` in -``isaaclab_experimental.envs.manager_based_env_warp`` for the initialization pattern. +Random noise kernels (gaussian, uniform) draw from the process-wide per-env Warp RNG state +:attr:`isaaclab.utils.seed.WarpRng.state`. """ from __future__ import annotations @@ -22,6 +20,8 @@ import warp as wp +from isaaclab.utils.seed import WarpRng + if TYPE_CHECKING: from . import noise_cfg @@ -108,7 +108,7 @@ def uniform_noise(data: wp.array, cfg: noise_cfg.UniformNoiseCfg) -> None: wp.launch( _apply_uniform_noise, dim=data.shape[0], - inputs=[data, cfg.rng_state_wp, float(cfg.n_min), float(cfg.n_max), _OPERATION_MAP[cfg.operation]], + inputs=[data, WarpRng.state, float(cfg.n_min), float(cfg.n_max), _OPERATION_MAP[cfg.operation]], device=data.device, ) @@ -149,7 +149,7 @@ def gaussian_noise(data: wp.array, cfg: noise_cfg.GaussianNoiseCfg) -> None: wp.launch( _apply_gaussian_noise, dim=data.shape[0], - inputs=[data, cfg.rng_state_wp, float(cfg.mean), float(cfg.std), _OPERATION_MAP[cfg.operation]], + inputs=[data, WarpRng.state, float(cfg.mean), float(cfg.std), _OPERATION_MAP[cfg.operation]], device=data.device, ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py index 2d071b823a46..757a45e8b825 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/__init__.py @@ -5,11 +5,12 @@ """Warp utility functions and shared kernels for isaaclab_experimental.""" +from isaaclab.utils.warp.utils import resolve_1d_mask + from .kernels import compute_reset_scale, count_masked from .utils import ( WarpCapturable, is_warp_capturable, - resolve_1d_mask, warp_capturable, wrap_to_pi, zero_masked_2d, diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py index a90f6b2d6fc6..0cac9bd65ac5 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp/utils.py @@ -5,97 +5,13 @@ from __future__ import annotations -from collections.abc import Sequence +import inspect +from collections.abc import Callable +from typing import Any -import torch import warp as wp -@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. @@ -113,17 +29,24 @@ def decorator(func): return decorator -def is_warp_capturable(func) -> bool: - """Check if a term function is CUDA-graph-capturable. +def is_warp_capturable(func, params: dict[str, Any] | None = None) -> bool: + """Check if a term is CUDA-graph-capturable. - Checks ``_warp_capturable`` on the function and its ``__wrapped__`` target. - Returns True (capturable) by default if no annotation is found. + Checks ``_warp_capturable`` on the term function or class (the class of an instance) and on its + ``__wrapped__`` target. The annotation is a bool, or a predicate that decides from the term's + parameters (see :class:`WarpCapturable`). Returns True (capturable) if no annotation is found. + + Args: + func: The term function, class, or class instance. + params: The term's parameters, passed to a predicate annotation. Defaults to None (no parameters). """ + if not (isinstance(func, type) or inspect.isroutine(func)): + func = type(func) 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 val if isinstance(val, bool) else bool(val(params or {})) return True @@ -146,53 +69,67 @@ class WarpCapturable: def reset_root_state_uniform(env, env_mask, ...): ... - @WarpCapturable(False, reason="calls write_root_pose_to_sim") - def push_by_setting_velocity(env, env_mask, ...): + @WarpCapturable(False, reason="notifies the solver of a model change") + class randomize_rigid_body_com(ManagerTermBase): + ... + + @WarpCapturable(lambda params: params.get("sensor_cfg") is None, reason="refreshes the sensor on the host") + def base_height_l2(env, out, target_height, asset_cfg=..., sensor_cfg=None): ... - ``@WarpCapturable(True)`` or no decorator: capturable, returned unwrapped. - - ``@WarpCapturable(False)``: sets ``func._warp_capturable = False``, wraps with + - ``@WarpCapturable(False)`` on a function: sets ``func._warp_capturable = False`` and wraps it with a runtime guard that raises if ``wp.get_device().is_capturing`` is ``True``. + - ``@WarpCapturable(False)`` on a class term: annotates the class and guards its ``__call__``. The term's + ``reset`` is still recorded; decorate the ``reset`` method itself to keep it out of the recording. + - ``@WarpCapturable(predicate)``: the term is capturable for the parameters the predicate accepts. The + predicate receives the term's parameters as a dict, and the guard raises only for rejected ones. """ - def __init__(self, capturable: bool, *, reason: str | None = None): + def __init__(self, capturable: bool | Callable[[dict[str, Any]], 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.""" + def __call__(self, target): + """Decorate a term function or class with the capture safety annotation and runtime guard.""" + if isinstance(self._capturable, bool): + target._warp_capturable = self._capturable + if self._capturable: + return target + else: + # a plain function stored on a class would bind to instances + target._warp_capturable = staticmethod(self._capturable) if isinstance(target, type) else self._capturable + if isinstance(target, type): + target.__call__ = self._guarded(target.__call__) + return target + return self._guarded(target) + + def _guarded(self, func): + """Wrap *func* so that calling it during CUDA graph capture with rejected parameters raises.""" import functools - func._warp_capturable = self._capturable - if self._capturable: - return func - + capturable = self._capturable reason = self._reason @functools.wraps(func) def wrapper(*args, **kwargs): - if wp.get_device().is_capturing: - msg = f"'{func.__qualname__}' is marked @WarpCapturable(False) but called during CUDA graph capture." + if wp.get_device().is_capturing and (isinstance(capturable, bool) or not capturable(kwargs)): + msg = ( + f"'{func.__qualname__}' is not capturable with these parameters but was called during CUDA" + " graph capture." + ) if reason: msg = f"{msg} {reason}" raise RuntimeError(msg) return func(*args, **kwargs) - wrapper._warp_capturable = False + wrapper._warp_capturable = capturable return wrapper @staticmethod - def is_capturable(func) -> bool: - """Check capturability annotation. Default: ``True``. - - Checks ``__wrapped__`` for decorated functions to handle stacked decorators. - """ - 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 + def is_capturable(func, params: dict[str, Any] | None = None) -> bool: + """Check the capturability annotation. See :func:`is_warp_capturable`.""" + return is_warp_capturable(func, params) @wp.kernel diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py new file mode 100644 index 000000000000..7ac964333549 --- /dev/null +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_capture.py @@ -0,0 +1,298 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""CUDA graph capture of Warp frontend stages: a :func:`captured` method records its calls and replays them, and +:func:`eager` runs a call that cannot be recorded at its place in the recording. + +.. code-block:: python + + @captured + def compute(self, dt): + eager(term_cfg.func, self._env, term_cfg.out, **term_cfg.params) +""" + +from __future__ import annotations + +import contextlib +import functools +import gc +import os +from collections.abc import Callable +from typing import Any, ClassVar + +import torch +import warp as wp + +from isaaclab.physics import PhysicsEvent, PhysicsManager +from isaaclab.utils.version import has_kit + +SYNC_DEBUG_ENV_VAR = "ISAACLAB_SYNC_DEBUG" +"""Set ``ISAACLAB_SYNC_DEBUG=1`` to run eager stage calls under ``torch.cuda.set_sync_debug_mode("error")``, +so a hidden GPU-to-host synchronization inside a stage raises instead of stalling silently.""" + + +def captured(method: Callable[..., Any]) -> Callable[..., Any]: + """Run a method through a :class:`CapturedStage` stored on its instance. + + The method body is the plain eager code of the stage, with the calls that cannot be recorded wrapped in + :func:`eager`. Its instance must expose ``device``. The instance keeps its stages in ``_captured_stages``, + keyed by the decorated method, so an override and the method it extends own separate stages. + + Args: + method: The method to decorate. + """ + + @functools.wraps(method) + def wrapper(self, *args: Any, **kwargs: Any) -> Any: + stages = self.__dict__.setdefault("_captured_stages", {}) + stage = stages.get(wrapper) + if stage is None: + stage = stages[wrapper] = CapturedStage(method, self.device) + return stage(self, *args, **kwargs) + + return wrapper + + +def eager(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + """Call a function that cannot be recorded from inside a :class:`CapturedStage`. + + Outside a recording this is a plain call. While a stage records, the call ends the current graph, runs eagerly + and starts the next graph, and every replay of the stage repeats it at the same place with the same arguments, + so the arguments must stay valid across calls like the arrays a graph reads. Replays discard its result. + + Args: + fn: The function to call. + *args: Positional arguments of the call. + **kwargs: Keyword arguments of the call. + + Returns: + The result of the call. + + Raises: + RuntimeError: If the current device is capturing a graph that no :class:`CapturedStage` records. + """ + recording = CapturedStage._recording + if recording is not None and recording.capturing: + return recording.split(fn, args, kwargs) + if wp.get_device().is_capturing: + name = getattr(fn, "__qualname__", type(fn).__name__) + raise RuntimeError( + f"'{name}' cannot be recorded, so it cannot run inside a graph capture that no CapturedStage started." + ) + return fn(*args, **kwargs) + + +def reset_captured_stages(owner: Any) -> None: + """Drop the graphs recorded by the :func:`captured` methods of an instance.""" + for stage in owner.__dict__.get("_captured_stages", {}).values(): + stage.reset() + + +class CapturedStage: + """A stage of the Warp frontend, run eagerly or through CUDA graphs recorded from its own calls. + + * The first call with a given argument layout records the stage without an eager warm-up: the recording is its + first execution. Later calls with the same array pointers and scalar values replay that recording and return + the result of the recording call, so a stage returns buffers its owner keeps. + * Calls that cannot be recorded go through :func:`eager`. A recording ends its graph at each of them, runs the + call and starts the next graph; a replay launches the graphs and repeats the calls with their recorded + arguments, in the recorded order. + * :attr:`enabled` is process-wide and set by the Warp environments. A stage runs eagerly while it is False and + on CPU devices. + * :attr:`generation` advances when the physics backend rebinds its buffers or stops (see :meth:`invalidate_on`), + and every stage records again on its next call. + * A stage called while another stage records runs inline, so its graphs and eager calls become part of that + recording. + """ + + enabled: ClassVar[bool] = False + """Whether stages record CUDA graphs. The Warp environments set it; the torch frontend never captures.""" + + generation: ClassVar[int] = 0 + """Counter of physics rebinds and stops; a stage recorded under an older value records again.""" + + _callback_handles: ClassVar[list] = [] + + _recording: ClassVar[_Recording | None] = None + """The recording in progress, which :func:`eager` calls split.""" + + def __init__(self, fn: Callable[..., Any], device: wp.DeviceLike): + """Initialize the stage. + + Args: + fn: Callable implementing the stage. Its GPU work must be capturable, except for :func:`eager` calls. + device: Warp device that runs the stage. + """ + self._fn = fn + self._device = wp.get_device(device) + self._sync_debug = os.environ.get(SYNC_DEBUG_ENV_VAR, "0") == "1" + self._recordings: dict[tuple, tuple[list, Any]] = {} + self._generation = CapturedStage.generation + + @property + def num_graphs(self) -> int: + """Number of graphs currently recorded for this stage.""" + return sum(isinstance(step, wp.Graph) for steps, _ in self._recordings.values() for step in steps) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Run the stage and return its result.""" + if self._device.is_capturing: + return self._fn(*args, **kwargs) + if not (CapturedStage.enabled and self._device.is_cuda): + return self._run_eager(self._fn, args, kwargs) + if self._generation != CapturedStage.generation: + self.reset() + self._generation = CapturedStage.generation + key = _argument_key(args, kwargs) + recorded = self._recordings.get(key) + if recorded is None: + self._recordings[key] = recorded = self._record(args, kwargs) + return recorded[1] + steps, result = recorded + for step in steps: + if isinstance(step, wp.Graph): + wp.capture_launch(step) + else: + self._run_eager(*step) + return result + + def reset(self) -> None: + """Drop the recordings; the stage records again on its next call.""" + if self._recordings: + # graphs may still be executing; drain the device before releasing them + wp.synchronize_device(self._device) + self._recordings.clear() + + @classmethod + def invalidate_on(cls, backend_manager: PhysicsManager) -> None: + """Record every stage again after the physics backend rebinds its buffers or stops. + + Args: + backend_manager: Physics manager that dispatches the lifecycle events. + """ + cls._callback_handles = [ + backend_manager.register_callback(cls.invalidate, event) + for event in (PhysicsEvent.PHYSICS_READY, PhysicsEvent.STOP) + ] + + @classmethod + def disable(cls) -> None: + """Run every stage eagerly again and stop following the physics backend.""" + for handle in cls._callback_handles: + handle.deregister() + cls._callback_handles = [] + cls.enabled = False + + @classmethod + def invalidate(cls, payload: Any = None) -> None: + """Record every stage again on its next call.""" + cls.generation += 1 + + def _record(self, args: tuple, kwargs: dict[str, Any]) -> tuple[list, Any]: + """Record the stage as its first run and return the recorded steps and the result.""" + recording = _Recording(self._device, self._run_eager) + enclosing, CapturedStage._recording = CapturedStage._recording, recording + try: + with _paused_gc(), wp.ScopedStream(recording.capture_stream): + recording.begin() + try: + result = self._fn(*args, **kwargs) + except BaseException: + recording.abort() + raise + recording.end() + finally: + CapturedStage._recording = enclosing + return recording.steps, result + + def _run_eager(self, fn: Callable[..., Any], args: tuple, kwargs: dict[str, Any]) -> Any: + """Run eagerly, under the synchronization trap when requested.""" + if not self._sync_debug: + return fn(*args, **kwargs) + previous = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + return fn(*args, **kwargs) + finally: + torch.cuda.set_sync_debug_mode(previous) + + +class _Recording: + """The steps of one stage recording: CUDA graphs split around :func:`eager` calls, launched as they end.""" + + def __init__(self, device: wp.Device, run_eager: Callable[[Callable[..., Any], tuple, dict], Any]): + # graphs are captured on the capture stream and launched on the stream the stage runs on + self.stream = wp.get_stream(device) + self.capture_stream = self.stream + self.mode = wp.CaptureMode.THREAD_LOCAL + if has_kit(): + # RTX uses the legacy CUDA stream. A nonblocking stream avoids implicit synchronization. + self.capture_stream = wp.stream_from_torch(torch.cuda.Stream(device=device.alias)) + self.mode = wp.CaptureMode.RELAXED + self.steps: list = [] + self.capturing = False + self._run_eager = run_eager + + def begin(self) -> None: + """Start the next graph.""" + wp.capture_begin(stream=self.capture_stream, capture_mode=self.mode) + self.capturing = True + + def end(self) -> None: + """End the current graph and launch it, since the recording is the first run of the stage.""" + self.capturing = False + graph = wp.capture_end(stream=self.capture_stream) + wp.capture_launch(graph, stream=self.stream) + self.steps.append(graph) + + def abort(self) -> None: + """End the current graph without launching it, after the stage raised.""" + if self.capturing: + self.capturing = False + with contextlib.suppress(Exception): + wp.capture_end(stream=self.capture_stream) + + def split(self, fn: Callable[..., Any], args: tuple, kwargs: dict[str, Any]) -> Any: + """Run a call that cannot be recorded between the current graph and the next one.""" + self.end() + with wp.ScopedStream(self.stream, sync_enter=False): + result = self._run_eager(fn, args, kwargs) + self.steps.append((fn, args, kwargs)) + self.begin() + return result + + +@contextlib.contextmanager +def _paused_gc(): + """Keep collection-driven frees out of a recording, as ``NewtonQueries.capture_graph`` does. + + Objects allocated during the recording stay in the youngest generation, so collecting that generation afterwards + frees the recording's garbage without the full collection of the whole heap. + """ + was_enabled = gc.isenabled() + gc.disable() + try: + yield + finally: + if was_enabled: + gc.enable() + gc.collect(0) + + +def _argument_key(args: tuple, kwargs: dict[str, Any]) -> tuple: + """Key a stage call by the pointers of its arrays and the values of its scalars.""" + return tuple(_value_key(value) for value in args) + tuple( + (name, _value_key(value)) for name, value in kwargs.items() + ) + + +def _value_key(value: Any) -> Any: + if isinstance(value, wp.array): + return value.ptr + if isinstance(value, torch.Tensor): + return value.data_ptr() + if value is None or isinstance(value, (bool, int, float, str)): + return value + return id(value) diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py deleted file mode 100644 index d446be28d099..000000000000 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py +++ /dev/null @@ -1,81 +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 CUDA graph capture-or-replay utility.""" - -from collections.abc import Callable -from typing import Any - -import warp as wp - - -class WarpGraphCache: - """Caches Warp CUDA graphs by stage name: captures on first call, replays after. - - 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. - - The return value from the capture run is cached and returned on every - subsequent replay, ensuring captured stages return the same references - (e.g. tensor views) as eager stages. - - Usage:: - - cache = WarpGraphCache() - result = cache.capture_or_replay("my_stage", my_warp_function) - # uncaptured work here ... - result2 = cache.capture_or_replay("my_stage_post", my_other_function) - """ - - def __init__(self): - self._graphs: dict[str, Any] = {} - self._results: dict[str, Any] = {} - - def capture_or_replay( - self, - stage: str, - fn: Callable[..., Any], - args: tuple = (), - kwargs: dict[str, Any] | None = None, - ) -> Any: - """Capture *fn* into a CUDA graph on the first call, then replay. - - Args: - stage: Unique name identifying this captured scope. - fn: The callable to capture. Must contain only CUDA-graph-safe - operations (pure warp kernels, no Python-level branching on - GPU data). - args: Positional arguments forwarded to *fn*. Defaults to ``()``. - kwargs: Keyword arguments forwarded to *fn*. Defaults to ``None``. - - Returns: - The cached return value from the first (capture) invocation. - """ - if kwargs is None: - 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: - result = fn(*args, **kwargs) - self._graphs[stage] = capture.graph - self._results[stage] = result - return result - - 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() - else: - self._graphs.pop(stage, None) - self._results.pop(stage, None) diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index 758b2f9ecbee..a3ef24ec7ac8 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -279,15 +279,19 @@ def resolve_joint_mask(self, joint_ids=None): class MockWrenchComposer: - """Mock wrench composer that records the last forces and torques it was handed.""" + """Mock wrench composer that records the last wrenches and selection it was handed.""" def __init__(self): self.last_forces = None self.last_torques = None + self.last_body_mask = None + self.last_env_mask = None - def set_forces_and_torques_mask(self, forces=None, torques=None, **kw): + def set_forces_and_torques_mask(self, forces=None, torques=None, body_mask=None, env_mask=None, **kw): self.last_forces = forces self.last_torques = torques + self.last_body_mask = body_mask + self.last_env_mask = env_mask class MockArticulation: @@ -484,7 +488,7 @@ def prev_action(self) -> torch.Tensor: class MockSceneEntityCfg: - """Unified cfg that works for both stable (joint_ids) and experimental (joint_mask / joint_ids_wp).""" + """Unified cfg that works for both stable (joint_ids) and experimental (joint_mask_wp / joint_ids_wp).""" def __init__(self, name: str, joint_ids: list[int], num_joints: int, device: str): self.name = name @@ -494,7 +498,7 @@ def __init__(self, name: str, joint_ids: list[int], num_joints: int, device: str mask = [False] * num_joints for idx in joint_ids: mask[idx] = True - self.joint_mask = wp.array(mask, dtype=wp.bool, device=device) + self.joint_mask_wp = wp.array(mask, dtype=wp.bool, device=device) self.joint_ids_wp = wp.array(joint_ids, dtype=wp.int32, device=device) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py index ea69b3308fc4..02c8e488ad52 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -5,9 +5,9 @@ """CUDA-graph capture safety for warp MDP terms. -The manager call switch runs every manager graph-captured by default -(``ManagerCallSwitch.DEFAULT_CONFIG``), so a warp MDP term that is not capture-safe is not an -edge case — it is the normal execution path. A captured graph records launch topology and +The Warp environments record every capturable manager stage into a CUDA graph by default +(:class:`~isaaclab_experimental.utils.CapturedStage`), so a warp MDP term that is not capture-safe +is not an edge case — it is the normal execution path. A captured graph records launch topology and pointer values and replays them, so a term must launch a fixed dim, touch only buffers that were allocated before capture and never move, and do all of its work in kernels: host-side Python between launches is simply not replayed. 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 e93e73267938..a9b2455039e9 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 @@ -7,16 +7,22 @@ from __future__ import annotations +from types import SimpleNamespace + import numpy as np import pytest import torch import warp as wp +from isaaclab.utils.seed import WarpRng + # Skip entire module if no CUDA device available wp.init() 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 EventManager, EventTermCfg, ManagerTermBase, SceneEntityCfg +from isaaclab_experimental.utils import CapturedStage from parity_helpers import ( DEVICE, NUM_ACTIONS, @@ -37,22 +43,10 @@ @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. - """ +def _reset_warp_rng(): + """Drop the Warp RNG state each test sets.""" 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) + WarpRng.state = None @pytest.fixture() @@ -101,8 +95,9 @@ class _Env: env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 + env.env_origins_wp = scene.env_origins # RNG state for events (seeded deterministically) - env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) + WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) return env @@ -111,6 +106,23 @@ def all_joints_cfg(): return MockSceneEntityCfg("robot", list(range(NUM_JOINTS)), NUM_JOINTS, DEVICE) +def _make_event_term(term_type: type[ManagerTermBase], env, mode: str = "reset", **params) -> ManagerTermBase: + """Build a class event term the way the event manager does, after resolving its asset's Warp body fields.""" + params.setdefault("asset_cfg", SceneEntityCfg("robot")) + cfg = EventTermCfg(func=term_type, mode=mode, params=params) + # the mock scene cannot run SceneEntityCfg.resolve, so fill the fields it would set on the stored cfg + asset_cfg = cfg.params["asset_cfg"] + num_bodies = env.scene[asset_cfg.name].num_bodies + body_ids = list(range(num_bodies)) if asset_cfg.body_ids in (None, slice(None)) else list(asset_cfg.body_ids) + asset_cfg.body_ids_wp = wp.array(body_ids, dtype=wp.int32, device=DEVICE) + asset_cfg.body_mask_wp = wp.array([i in body_ids for i in range(num_bodies)], dtype=wp.bool, device=DEVICE) + return term_type(cfg, env) + + +def _all_envs() -> wp.array: + return wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + + # ============================================================================ # Event capture-mutate-replay tests # ============================================================================ @@ -192,7 +204,7 @@ 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, the written root velocity == root_vel_w. Mutate root_vel_w -> tracks.""" - mask = wp.array([True] * NUM_ENVS, dtype=wp.bool, device=DEVICE) + mask = _all_envs() zero_range = { "x": (0.0, 0.0), "y": (0.0, 0.0), @@ -201,10 +213,10 @@ 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) 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) @@ -220,26 +232,34 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): # -- apply_external_force_torque -------------------------------------------- def test_apply_external_force_torque(self, warp_env, art_data, all_joints_cfg): - """A degenerate non-zero range reaches masked envs through the replayed kernel; unmasked envs get zero.""" + """A degenerate non-zero range reaches masked envs through the replayed kernel; the composer gets the mask.""" mask = wp.array([i < NUM_ENVS // 2 for i in range(NUM_ENVS)], dtype=wp.bool, device=DEVICE) + term = _make_event_term( + warp_evt.apply_external_force_torque, warp_env, force_range=(2.0, 2.0), torque_range=(3.0, 3.0) + ) - warp_evt.apply_external_force_torque(warp_env, mask, force_range=(2.0, 2.0), torque_range=(3.0, 3.0)) with wp.ScopedCapture() as cap: - warp_evt.apply_external_force_torque(warp_env, mask, force_range=(2.0, 2.0), torque_range=(3.0, 3.0)) + term(warp_env, mask, **term.cfg.params) composer = warp_env.scene["robot"].permanent_wrench_composer forces = wp.to_torch(composer.last_forces) torques = wp.to_torch(composer.last_torques) - # clear the buffers so only the replayed kernel can produce the expected values - forces.fill_(-1.0) - torques.fill_(-1.0) wp.capture_launch(cap.graph) wp.synchronize() half = NUM_ENVS // 2 assert_close(forces[:half], torch.full_like(forces[:half], 2.0)) assert_close(torques[:half], torch.full_like(torques[:half], 3.0)) - assert_close(forces[half:], torch.zeros_like(forces[half:])) - assert_close(torques[half:], torch.zeros_like(torques[half:])) + assert composer.last_env_mask is mask + + def test_apply_external_force_torque_zero_ranges_are_no_op(self, warp_env, art_data, all_joints_cfg): + """Zero ranges keep the composer's wrenches: a placeholder term must not clear them on every reset.""" + term = _make_event_term( + warp_evt.apply_external_force_torque, warp_env, force_range=(0.0, 0.0), torque_range=(0.0, 0.0) + ) + + term(warp_env, _all_envs(), **term.cfg.params) + + assert warp_env.scene["robot"].permanent_wrench_composer.last_forces is None # -- env_mask selectivity --------------------------------------------------- @@ -267,3 +287,129 @@ 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 event terms: persistent state and live configuration +# ============================================================================ + + +def _make_event_env(seed: int, num_bodies: int = 1): + """An independent Warp event environment with its own asset.""" + 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) + scene = MockScene({"robot": asset}, origins) + env = SimpleNamespace(scene=scene, num_envs=NUM_ENVS, device=DEVICE, env_origins_wp=origins) + WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + seed, device=DEVICE) + return env, data, asset + + +@pytest.mark.parametrize( + ("term_type", "params"), + [ + (warp_evt.push_by_setting_velocity, {"velocity_range": {"x": (1.0, 1.0)}}), + (warp_evt.apply_external_force_torque, {"force_range": (1.0, 1.0), "torque_range": (2.0, 2.0)}), + (warp_evt.reset_root_state_uniform, {"pose_range": {}, "velocity_range": {}}), + ], + ids=["push", "external_wrench", "root_reset"], +) +def test_class_terms_allocate_only_during_initialization(term_type, params, monkeypatch): + """A recorded stage runs its terms for the first time inside CUDA graph capture, so a call must not allocate.""" + env, _, _ = _make_event_env(81) + term = _make_event_term(term_type, env, **params) + + def fail(*args, **kwargs): + pytest.fail("the term allocated during a call") + + monkeypatch.setattr(wp, "zeros", fail) + monkeypatch.setattr(wp, "empty", fail) + monkeypatch.setattr(wp, "clone", fail) + term(env, _all_envs(), **term.cfg.params) + term(env, _all_envs(), **term.cfg.params) + + +def test_push_range_change_applies_on_the_next_call(): + """A curriculum that edits the term's range in place changes the next push, as in the stable term.""" + env, data, asset = _make_event_env(277) + copy_np_to_wp(data.root_vel_w, np.zeros((NUM_ENVS, 6), dtype=np.float32)) + term = _make_event_term(warp_evt.push_by_setting_velocity, env, mode="interval", velocity_range={"x": (1.0, 1.0)}) + term(env, _all_envs(), **term.cfg.params) + + term.cfg.params["velocity_range"]["x"] = (4.0, 4.0) + term(env, _all_envs(), **term.cfg.params) + wp.synchronize() + + assert_close(wp.to_torch(asset.last_root_velocity)[:, 0], torch.full((NUM_ENVS,), 4.0, device=DEVICE)) + + +def test_com_randomization_uses_the_first_call_baseline_without_accumulating(): + """Startup terms may move the CoM after construction; each call offsets that baseline, not the last result.""" + env, data, asset = _make_event_env(84, num_bodies=2) + baseline = np.zeros((NUM_ENVS, 2, 3), dtype=np.float32) + copy_np_to_wp(data.body_com_pos_b, baseline) + asset.set_coms_mask = lambda **kwargs: None + term = _make_event_term( + warp_evt.randomize_rigid_body_com, + env, + mode="startup", + com_range={"x": (1.0, 1.0)}, + asset_cfg=SceneEntityCfg("robot", body_ids=[0, 1]), + ) + # another startup term moves the CoM between construction and the first call + baseline[..., 0] = 0.75 + copy_np_to_wp(data.body_com_pos_b, baseline) + + term(env, _all_envs(), **term.cfg.params) + term(env, _all_envs(), **term.cfg.params) + wp.synchronize() + + assert_close(data.body_com_pos_b.torch[..., 0], torch.full((NUM_ENVS, 2), 1.75, device=DEVICE)) + + +@pytest.mark.parametrize( + ("term_type", "params"), + [ + ( + warp_evt.randomize_rigid_body_material, + { + "static_friction_range": (0.5, 0.5), + "dynamic_friction_range": (0.5, 0.5), + "restitution_range": (0.0, 0.0), + "num_buckets": 1, + }, + ), + (warp_evt.randomize_rigid_body_mass, {"mass_distribution_params": (1.0, 1.0), "operation": "scale"}), + ], + ids=["material", "mass"], +) +def test_mask_adapters_run_eagerly_in_reset_mode(term_type, params, monkeypatch): + """The adapters convert the mask to indices on the host, so reset events run them outside the recorded graph.""" + stable_term_type = term_type.__mro__[1] + calls = [] + + def stable_init(self, cfg, env): + self.cfg, self._env = cfg, env + + monkeypatch.setattr(stable_term_type, "__init__", stable_init) + monkeypatch.setattr(stable_term_type, "__call__", lambda self, env, env_ids, *args: calls.append(env_ids.tolist())) + monkeypatch.setattr(CapturedStage, "enabled", True) + env, _, _ = _make_event_env(5) + env.sim = SimpleNamespace(is_playing=lambda: True) + manager = EventManager( + { + "randomize": EventTermCfg( + func=term_type, mode="reset", params={**params, "asset_cfg": SimpleNamespace(name="robot")} + ) + }, + env, + ) + env_mask = wp.array([i % 2 == 0 for i in range(NUM_ENVS)], dtype=wp.bool, device=DEVICE) + selected = [i for i in range(NUM_ENVS) if i % 2 == 0] + + for _ in range(2): + manager.apply( + mode="reset", env_mask_wp=env_mask, global_env_step_count=wp.zeros(1, dtype=wp.int32, device=DEVICE) + ) + + assert calls == [selected, selected], "the adapter must run, eagerly, on every reset" 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 0f014cc2b3fd..91a5f8fe138f 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 @@ -12,6 +12,8 @@ import torch import warp as wp +from isaaclab.utils.seed import WarpRng + # Skip entire module if no CUDA device available wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") @@ -50,6 +52,7 @@ @pytest.fixture(autouse=True) def _clear_caches(): yield + WarpRng.state = None for fn in [warp_obs.generated_commands]: for attr in list(vars(fn)): if attr.startswith("_"): @@ -130,7 +133,7 @@ class _Env: env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 - env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) + WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) return env @@ -169,7 +172,7 @@ class _Env: env.step_dt = 0.02 env.max_episode_length = 500 env.max_episode_length_s = 10.0 - env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) + WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) return env 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 c81f6f73ec89..9f2d39253d4e 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 @@ -12,6 +12,8 @@ import torch import warp as wp +from isaaclab.utils.seed import WarpRng + # Skip entire module if no CUDA device available wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") @@ -58,6 +60,7 @@ @pytest.fixture(autouse=True) def _clear_caches(): yield + WarpRng.state = None for fn in [warp_rew.track_lin_vel_xy_exp, warp_rew.track_ang_vel_z_exp, warp_rew.undesired_contacts]: for attr in list(vars(fn)): if attr.startswith("_"): @@ -144,7 +147,7 @@ class _Env: env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 - env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) + WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) return env @@ -184,7 +187,7 @@ class _Env: env.step_dt = 0.02 env.max_episode_length = 500 env.max_episode_length_s = 10.0 - env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) + WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) return env diff --git a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py index 9676f2b06044..e0fbee333530 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py @@ -14,6 +14,8 @@ import torch import warp as wp +from isaaclab.utils.seed import WarpRng + # Skip entire module if no CUDA device available wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") @@ -59,6 +61,7 @@ @pytest.fixture(autouse=True) def _clear_caches(): yield + WarpRng.state = None for fn in [warp_term.illegal_contact]: for attr in list(vars(fn)): if attr.startswith("_"): @@ -139,7 +142,7 @@ class _Env: env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 - env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) + WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) return env @@ -179,7 +182,7 @@ class _Env: env.step_dt = 0.02 env.max_episode_length = 500 env.max_episode_length_s = 10.0 - env.rng_state_wp = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) + WarpRng.state = wp.array(np.arange(NUM_ENVS, dtype=np.uint32) + 42, device=DEVICE) return env diff --git a/source/isaaclab_experimental/test/envs/test_frontend.py b/source/isaaclab_experimental/test/envs/test_frontend.py index ff71cdb49b0d..67cf269db3a8 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend.py +++ b/source/isaaclab_experimental/test/envs/test_frontend.py @@ -206,7 +206,7 @@ def test_from_stable_copies_all_selection_fields(): assert getattr(warp, field) == getattr(stable, field), f"field {field!r} mismatch" assert isinstance(warp, WarpSceneEntityCfg) # Warp-only fields stay None until :meth:`resolve` runs. - assert warp.joint_mask is None + assert warp.joint_mask_wp is None assert warp.joint_ids_wp is None assert warp.body_ids_wp is None diff --git a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py index 7496d4c7824f..3a0026522462 100644 --- a/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py +++ b/source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py @@ -201,6 +201,7 @@ def test_no_warp_task_registrations_remain(): "Isaac-Ant-Direct", "Isaac-Humanoid-Direct", "Isaac-Reorient-Cube-Allegro-Direct", + "Isaac-Reorient-Cube-Shadow-Direct", ] diff --git a/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py new file mode 100644 index 000000000000..440cc82a29a9 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_interactive_scene_warp.py @@ -0,0 +1,63 @@ +# 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 masked resets of the Warp interactive scene on Newton.""" + +import torch +import warp as wp +from isaaclab_experimental.envs import InteractiveSceneWarp +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors.imu import ImuCfg +from isaaclab.sim import SimulationCfg +from isaaclab.terrains import TerrainImporterCfg +from isaaclab.utils import configclass + + +@configclass +class ImuSceneCfg(InteractiveSceneCfg): + """Scene with a cube resting on the ground and an IMU on the cube.""" + + env_spacing = 2.0 + terrain = TerrainImporterCfg(prim_path="/World/ground", terrain_type="plane") + + cube = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube", + spawn=sim_utils.CuboidCfg( + size=(0.2, 0.2, 0.2), + rigid_props=sim_utils.UsdPhysicsRigidBodyCfg(), + mass_props=sim_utils.MassCfg(mass=1.0), + collision_props=sim_utils.UsdPhysicsCollisionCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.1)), + ) + + imu = ImuCfg(prim_path="{ENV_REGEX_NS}/Cube") + + +def test_masked_reset_keeps_sensor_state_of_unselected_envs(): + """A masked scene reset zeroes the selected environment's IMU reading and keeps the other one.""" + sim_cfg = SimulationCfg(dt=1.0 / 200.0, physics=NewtonCfg(solver_cfg=MJWarpSolverCfg(), num_substeps=1)) + with sim_utils.build_simulation_context(sim_cfg=sim_cfg) as sim: + sim._app_control_on_stop_handle = None + scene = InteractiveSceneWarp(ImuSceneCfg(num_envs=2)) + sim.reset() + scene.reset() + imu = scene["imu"] + + for _ in range(50): + scene.write_data_to_sim() + sim.step(render=False) + scene.update(dt=sim.get_physics_dt()) + pre_reset_lin_acc = imu.data.lin_acc_b.torch.clone() + assert (pre_reset_lin_acc[:, 2] > 5.0).all(), f"Expected a resting gravity reading, got {pre_reset_lin_acc}" + + scene.reset(env_mask=wp.array([True, False], dtype=wp.bool, device=imu.device)) + lin_acc = imu.data.lin_acc_b.torch + torch.testing.assert_close(lin_acc[0], torch.zeros_like(lin_acc[0])) + torch.testing.assert_close(lin_acc[1], pre_reset_lin_acc[1]) diff --git a/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py b/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py new file mode 100644 index 000000000000..85c57918bb03 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_reorient_warp_env.py @@ -0,0 +1,159 @@ +# 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 + +"""Same-state parity of the Warp in-hand reorientation twins against the stable Direct task. + +Each twin is built from its stable registration through :class:`WarpFrontend` and stepped on its +default captured path. After every step, the stable torch implementation recomputes observations +and rewards from the twin's own state, so the comparison isolates the twin's kernels from +simulation noise. +""" + +from types import SimpleNamespace + +import pytest +import torch +import warp as wp +from isaaclab_experimental.envs.frontend import WarpFrontend + +from isaaclab.sim import SimulationContext +from isaaclab.utils.math import quat_error_magnitude + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.core.reorient.reorient_direct_env import ReorientDirectEnv, reorient_reward +from isaaclab_tasks.utils import resolve_task_config + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="The Warp twins capture CUDA graphs.") + +_SHADOW_TASK = "Isaac-Reorient-Cube-Shadow-Direct" +_TASKS = ["Isaac-Reorient-Cube-Allegro-Direct", _SHADOW_TASK] + + +@pytest.fixture(scope="module", params=_TASKS) +def env(request): + """A 16-environment twin of the parametrized stable task, shared by this module's tests.""" + cfg, _ = resolve_task_config(request.param, "", overrides=("physics=newton_mjwarp",)) + cfg.scene.num_envs = 16 + cfg.seed = 42 + env = WarpFrontend.build_env(cfg, request.param) + yield env + env.close() + SimulationContext.clear_instance() + + +def _stable_state(env, actions: torch.Tensor) -> SimpleNamespace: + """The state the stable task derives its observations and rewards from, read off the twin.""" + state = SimpleNamespace( + cfg=env.cfg, + num_envs=env.num_envs, + hand=env.hand, + object=env.object, + scene=env.scene, + finger_bodies=env.hand.find_bodies(env.cfg.fingertip_body_names)[0], + actions=actions, + goal_rot=wp.to_torch(env.goal_rot), + ) + state.num_fingertips = len(state.finger_bodies) + ReorientDirectEnv._compute_intermediate_values(state) + joint_limits = env.hand.data.joint_limits.torch + state.hand_dof_lower_limits, state.hand_dof_upper_limits = joint_limits[..., 0], joint_limits[..., 1] + state.in_hand_pos = env.object.data.default_root_pose.torch[:, :3] + torch.tensor( + env.cfg.in_hand_pos_offset, device=env.device + ) + return state + + +def test_observations_and_rewards_match_stable_task(env): + """Observations and rewards equal the stable task's on the same state. + + Rewards are computed before the step's resets, so they are compared on the environments whose + state no episode reset or goal resample changed afterwards. + """ + env.reset() + generator = torch.Generator(device=env.device).manual_seed(0) + num_compared = 0 + for _ in range(40): + goal_rot = wp.to_torch(env.goal_rot).clone() + # unit Gaussian: a policy's samples leave the [-1, 1] action range too + actions = torch.randn((env.num_envs, env.cfg.action_space), device=env.device, generator=generator) + obs, reward, terminated, truncated, _ = env.step(actions) + stable = _stable_state(env, actions) + + torch.testing.assert_close(obs["policy"], ReorientDirectEnv.compute_full_observations(stable)) + + untouched = ~(terminated | truncated) & (stable.goal_rot == goal_rot).all(dim=-1) + error = quat_error_magnitude(stable.object_rot, stable.goal_rot) + no_reset = torch.zeros_like(untouched) + expected, *_ = reorient_reward( + no_reset, + no_reset, + torch.zeros_like(error), + torch.zeros(1, device=env.device), + stable.object_pos, + stable.in_hand_pos, + error <= env.cfg.success_tolerance, + error, + actions, + env.cfg.dist_reward_scale, + env.cfg.rot_reward_scale, + env.cfg.rot_eps, + env.cfg.action_penalty_scale, + env.cfg.reach_goal_bonus, + env.cfg.fall_dist, + env.cfg.fall_penalty, + env.cfg.av_factor, + ) + torch.testing.assert_close(reward[untouched], expected[untouched], rtol=1e-4, atol=1e-4) + num_compared += int(untouched.sum()) + assert num_compared > 0 + + +@pytest.mark.parametrize("env", [_SHADOW_TASK], indirect=True) +def test_tendon_actions_curl_the_coupled_finger_joints(env): + """The tendon action columns drive the middle and distal finger joints, step after step. + + Every finger is first opened, then half of the environments curl theirs; a tendon command + that never reached the solver, or reached it only once, leaves both halves open. + """ + env.reset() + coupled_joints, _ = env.hand.find_joints("rh_(FF|MF|RF|LF)J[12]") + num_joint_actions = len(env.cfg.actuated_joint_names) + curling = torch.arange(env.num_envs, device=env.device) < env.num_envs // 2 + actions = torch.zeros((env.num_envs, env.cfg.action_space), device=env.device) + actions[:, num_joint_actions:] = -1.0 + for _ in range(10): + env.step(actions) + actions[curling, num_joint_actions:] = 1.0 + was_reset = torch.zeros_like(curling) + for _ in range(30): + _, _, terminated, truncated, _ = env.step(actions) + was_reset |= terminated | truncated + + curl = env.hand.data.joint_pos.torch[:, coupled_joints].sum(dim=-1) + assert (curl[curling & ~was_reset].min() - curl[~curling & ~was_reset].max()) > 1.0 + + +@pytest.mark.parametrize("env", [_SHADOW_TASK], indirect=True) +def test_assigned_episode_lengths_time_out_only_their_envs(env): + """An episode-length buffer assigned from outside, as RSL-RL does, times out only the envs it ends. + + The masked reset restarts those episodes and resamples their goals, leaving the others' alone. + """ + env.reset() + episode_lengths = torch.zeros_like(env.episode_length_buf) + episode_lengths[0] = env.max_episode_length - 2 + env.episode_length_buf = episode_lengths + goal_rot = wp.to_torch(env.goal_rot).clone() + + _, _, terminated, truncated, _ = env.step(torch.zeros((env.num_envs, env.cfg.action_space), device=env.device)) + + assert truncated[0] and not truncated[1:].any() + continuing = ~(terminated | truncated) + assert env.episode_length_buf[0] == 0 + assert (env.episode_length_buf[continuing] == 1).all() + new_goal_rot = wp.to_torch(env.goal_rot) + assert not torch.equal(new_goal_rot[0], goal_rot[0]) + unreached = continuing & ~wp.to_torch(env.goal_reached) + torch.testing.assert_close(new_goal_rot[unreached], goal_rot[unreached]) diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_capture.py b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py new file mode 100644 index 000000000000..9beeae991935 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_warp_env_capture.py @@ -0,0 +1,152 @@ +# 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 + +"""Eager and CUDA-graph execution of the Warp environments must produce the same rollout. + +Each task runs the same seeded action sequence twice, once with capture disabled +(:attr:`CapturedStage.enabled` False) and once with stage capture on. For manager-based tasks, a hard +:meth:`SimulationContext.reset` halfway through rebuilds the Newton model and reallocates every simulation +buffer; recorded graphs still point at the freed buffers, so the captured rollout only stays equal when the +stages record again. A variant marks one term of every manager as not capturable: its stage must record the +other terms and run that term eagerly. +""" + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +from isaaclab.sim import SimulationCfg +from isaaclab.test.utils import launch_test_simulation + +launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg()))) + +import isaaclab_tasks_experimental # noqa: F401 +import pytest +import torch +from isaaclab_experimental.envs.frontend import WarpFrontend +from isaaclab_experimental.utils import CapturedStage +from isaaclab_experimental.utils.warp import WarpCapturable + +import isaaclab.sim as sim_utils + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import resolve_task_config + +_NUM_ENVS = 32 +_MANAGERS = ("action_manager", "observation_manager", "event_manager", "reward_manager", "termination_manager") +_STEPS = 80 +_REBIND_STEP = 40 + +# One function term per staged manager of the manager-based Cartpole. +_CARTPOLE_EAGER_TERMS = ( + "observations.policy.joint_pos_rel", + "rewards.pole_pos", + "terminations.cart_out_of_bounds", + "events.reset_cart_position", +) + + +def _recorded_stages(env) -> tuple[str, ...]: + """``Owner.method`` of every stage currently backed by a recorded graph, sorted.""" + owners = [env, *(getattr(env, name) for name in _MANAGERS if hasattr(env, name))] + return tuple( + sorted( + f"{type(owner).__name__}.{method.__name__}" + for owner in owners + for method, stage in owner.__dict__.get("_captured_stages", {}).items() + if stage.num_graphs + ) + ) + + +def _stale_stages(env) -> list[str]: + """Stages whose graphs were recorded before the latest physics rebind.""" + owners = [env, *(getattr(env, name) for name in _MANAGERS if hasattr(env, name))] + return [ + method.__name__ + for owner in owners + for method, stage in owner.__dict__.get("_captured_stages", {}).items() + if stage.num_graphs and stage._generation != CapturedStage.generation + ] + + +def _rollout( + task_id: str, capture: bool, rebind: bool, monkeypatch: pytest.MonkeyPatch, eager_terms: tuple[str, ...] = () +) -> dict: + """Run the seeded action sequence and return the rollout with the recorded stages.""" + env_cfg, _ = resolve_task_config(task_id, "", overrides=("physics=newton_mjwarp",)) + env_cfg.seed = 7 + env_cfg.scene.num_envs = _NUM_ENVS + # contacts are only reproducible between runs with deterministic contact ordering + env_cfg.sim.physics.deterministic_mode = "run_to_run" + env_cfg.sim.physics.solver_cfg.disable_sensors = True + WarpFrontend.adapt_cfg(env_cfg) + for path in eager_terms: + term_cfg = env_cfg + for name in path.split("."): + term_cfg = getattr(term_cfg, name) + # the guard raises if the term is ever recorded into a graph + term_cfg.func = WarpCapturable(False, reason="marked eager by the test")(term_cfg.func) + sim_utils.create_new_stage() + env = WarpFrontend.build_env(env_cfg, task_id).unwrapped + # the environment enables capture when it is built + monkeypatch.setattr(CapturedStage, "enabled", capture) + rollout = {"obs": [], "reward": [], "terminated": [], "truncated": []} + try: + env.reset() + generator = torch.Generator().manual_seed(123) + # a per-env bias drives environments to terminate at different steps + bias = torch.rand((_NUM_ENVS, env.action_space.shape[-1]), generator=generator) * 2.0 - 1.0 + for step in range(_STEPS): + if rebind and step == _REBIND_STEP: + rollout["stages_before_rebind"] = _recorded_stages(env) + generation = CapturedStage.generation + env.sim.reset() + rollout["rebind_advanced_generation"] = CapturedStage.generation > generation + env.reset() + noise = torch.rand((_NUM_ENVS, bias.shape[1]), generator=generator) * 2.0 - 1.0 + obs, reward, terminated, truncated, _ = env.step((bias + 0.5 * noise).to(env.device)) + rollout["obs"].append(obs["policy"].clone()) + rollout["reward"].append(reward.clone()) + rollout["terminated"].append(terminated.clone()) + rollout["truncated"].append(truncated.clone()) + rollout["stages_at_end"] = _recorded_stages(env) + rollout["stale_stages"] = _stale_stages(env) + finally: + env.close() + for key in ("obs", "reward", "terminated", "truncated"): + rollout[key] = torch.stack(rollout[key]) + return rollout + + +# Direct tasks keep the simulation arrays they bind in ``__init__``, so they do not survive a rebind. +@pytest.mark.parametrize( + ("task_id", "rebind", "eager_terms"), + [ + ("Isaac-Cartpole", True, ()), + ("Isaac-Cartpole", False, _CARTPOLE_EAGER_TERMS), + ("Isaac-Cartpole-Direct", False, ()), + ], + ids=["manager-cartpole", "manager-cartpole-eager-terms", "direct-cartpole"], +) +def test_captured_rollout_matches_eager( + task_id: str, rebind: bool, eager_terms: tuple[str, ...], monkeypatch: pytest.MonkeyPatch +): + eager = _rollout(task_id, capture=False, rebind=rebind, monkeypatch=monkeypatch, eager_terms=eager_terms) + captured = _rollout(task_id, capture=True, rebind=rebind, monkeypatch=monkeypatch, eager_terms=eager_terms) + + assert eager["stages_at_end"] == () + assert captured["stages_at_end"], "no stage was recorded; the comparison proves nothing" + if eager_terms: + # every stage holding an eager term still records the terms around it + recorded_groups = {stage.partition(".")[0] for stage in captured["stages_at_end"]} + assert {"ObservationManager", "RewardManager", "TerminationManager", "EventManager"} <= recorded_groups + if rebind: + assert captured["rebind_advanced_generation"], "the rebind must retire graphs that read freed buffers" + assert captured["stale_stages"] == [], "every stage must record again after the rebind" + assert captured["stages_at_end"] == captured["stages_before_rebind"] + for window in (slice(0, _REBIND_STEP), slice(_REBIND_STEP, _STEPS)): + dones = captured["terminated"][window] | captured["truncated"][window] + assert dones.any(), "each half must reset environments so the reset stages replay" + for key in ("obs", "reward", "terminated", "truncated"): + assert torch.equal(captured[key], eager[key]), f"{key} differs between eager and captured execution" diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py b/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py new file mode 100644 index 000000000000..8ce9fc741750 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_warp_env_decimation.py @@ -0,0 +1,153 @@ +# 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 + +"""The Warp manager-based environment steps physics at the stable environment's cadence. + +When every actuator runs inside Newton, the physics backend folds the whole decimation loop into one +step: actions are applied, the simulation stepped and the scene updated once per environment step. +Contact sensors then sample once per environment step, so their history, the air times and every +reward and termination reading them depend on that cadence. The same seeded Go2 rollout on the stable +and the Warp frontend must make the same calls and produce the same contact data. +""" + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +from isaaclab.sim import SimulationCfg +from isaaclab.test.utils import launch_test_simulation + +launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg()))) + +import math + +import isaaclab_tasks_experimental # noqa: F401 +import pytest +import torch +from isaaclab_experimental.envs.frontend import WarpFrontend + +import isaaclab.sim as sim_utils +from isaaclab.envs import ManagerBasedRLEnv + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.core.velocity.velocity_env_cfg import RewardsCfg +from isaaclab_tasks.utils import resolve_task_config + +_TASK_ID = "Isaac-Velocity-Flat-UnitreeGo2" +_NUM_ENVS = 64 +_STEPS = 200 + + +def _go2_cfg(): + """Go2 flat with every random draw removed, so both frontends see identical physics inputs.""" + env_cfg, _ = resolve_task_config(_TASK_ID, "", overrides=("physics=newton_mjwarp",)) + env_cfg.seed = 7 + env_cfg.scene.num_envs = _NUM_ENVS + # contacts are only reproducible between runs with deterministic contact ordering + env_cfg.sim.physics.deterministic_mode = "run_to_run" + env_cfg.sim.physics.solver_cfg.disable_sensors = True + env_cfg.events.add_base_mass = None + env_cfg.events.base_com = None + env_cfg.events.push_robot = None + env_cfg.events.reset_base.params["pose_range"] = {} + env_cfg.events.reset_base.params["velocity_range"] = {} + env_cfg.events.reset_robot_joints.params["position_range"] = (1.0, 1.0) + command = env_cfg.commands.base_velocity + command.rel_standing_envs = 0.0 + command.rel_heading_envs = 0.0 + command.heading_command = False + command.debug_vis = False + command.ranges.lin_vel_x = (0.5, 0.5) + command.ranges.lin_vel_y = (0.0, 0.0) + command.ranges.ang_vel_z = (0.0, 0.0) + # Go2 drops the thigh-contact penalty; restore it to cover a second contact-history reader + env_cfg.rewards.undesired_contacts = RewardsCfg().undesired_contacts + env_cfg.rewards.undesired_contacts.params["sensor_cfg"].body_names = ".*_thigh" + return env_cfg + + +def _rollout(frontend: str) -> dict: + """Step a fixed action sequence and record the physics calls and the contact data of every step.""" + env_cfg = _go2_cfg() + sim_utils.create_new_stage() + if frontend == "warp": + env = WarpFrontend.build_env(env_cfg, _TASK_ID).unwrapped + else: + env = ManagerBasedRLEnv(cfg=env_cfg) + calls = {"sim_step": 0, "apply_action": 0} + sim_step = env.sim.step + + def count_sim_step(*args, **kwargs): + calls["sim_step"] += 1 + return sim_step(*args, **kwargs) + + env.sim.step = count_sim_step + apply_action = env.action_manager.apply_action + + def count_apply_action(): + calls["apply_action"] += 1 + apply_action() + + env.action_manager.apply_action = count_apply_action + if frontend == "warp": + step_reward = env.reward_manager._step_reward_tensor_view + term_dones = env.termination_manager._term_dones_tensor_view + else: + step_reward = env.reward_manager._step_reward + term_dones = env.termination_manager._term_dones + reward_names = env.reward_manager._term_names + base_contact = env.termination_manager._term_names.index("base_contact") + sensor = env.scene["contact_forces"].data + rollout = {"calls": [], "handles_decimation": env._physics_handles_decimation} + records = { + "net_forces_w_history": lambda: sensor.net_forces_w_history.torch, + "current_air_time": lambda: sensor.current_air_time.torch, + "last_air_time": lambda: sensor.last_air_time.torch, + "undesired_contacts": lambda: step_reward[:, reward_names.index("undesired_contacts")], + "feet_air_time": lambda: step_reward[:, reward_names.index("feet_air_time")], + "base_contact": lambda: term_dones[:, base_contact], + } + try: + env.reset() + action_dim = env.action_space.shape[-1] + envs = torch.arange(_NUM_ENVS, dtype=torch.float32)[:, None] + joints = torch.arange(action_dim, dtype=torch.float32)[None, :] + # growing amplitudes make environments fall, and reset, at different steps + amplitude = 1.0 + 3.0 * envs / (_NUM_ENVS - 1) + for step in range(_STEPS): + if step < 20: + action = torch.zeros((_NUM_ENVS, action_dim)) + else: + action = amplitude * torch.sin(2 * math.pi * step / 25.0 + 0.7 * joints + 0.05 * envs) + calls["sim_step"] = calls["apply_action"] = 0 + env.step(action.to(env.device)) + rollout["calls"].append((calls["sim_step"], calls["apply_action"])) + for key, read in records.items(): + rollout.setdefault(key, []).append(read().clone()) + finally: + env.close() + for key in records: + rollout[key] = torch.stack(rollout[key]) + return rollout + + +@pytest.fixture(scope="module") +def rollouts() -> dict[str, dict]: + return {frontend: _rollout(frontend) for frontend in ("stable", "warp")} + + +def test_physics_loop_cadence_matches_stable(rollouts: dict[str, dict]): + stable, warp = rollouts["stable"], rollouts["warp"] + assert stable["handles_decimation"], "Go2 runs its actuators inside Newton, so physics must own decimation" + assert warp["handles_decimation"] + assert warp["calls"] == stable["calls"] == [(1, 1)] * _STEPS + + +def test_contact_data_matches_stable(rollouts: dict[str, dict]): + stable, warp = rollouts["stable"], rollouts["warp"] + assert stable["base_contact"].any(), "no environment terminated; the resets are not exercised" + for key in ("net_forces_w_history", "current_air_time", "last_air_time", "undesired_contacts", "base_contact"): + assert torch.equal(warp[key], stable[key]), key + # stable stores ``raw * weight * dt / dt`` and reduces the feet with torch, the Warp manager stores + # ``raw * weight`` and its term sums the feet in a loop: the two differ in the last bit only + torch.testing.assert_close(warp["feet_air_time"], stable["feet_air_time"], rtol=0.0, atol=1e-6) diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py b/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py new file mode 100644 index 000000000000..07c2e4413244 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_warp_env_final_obs.py @@ -0,0 +1,129 @@ +# 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 + +"""Terminal observations of the Warp environments match the stable environments. + +With ``compute_final_obs`` set, a step that resets environments exposes the observation computed before the +reset in ``extras["final_obs"]`` and returns the observation after the reset, as the stable environments do +for Same-Step autoreset. +""" + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +from isaaclab.sim import SimulationCfg +from isaaclab.test.utils import launch_test_simulation + +launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg()))) + +import gymnasium as gym +import isaaclab_tasks_experimental # noqa: F401 +import pytest +import torch +from isaaclab_experimental.envs.frontend import WarpFrontend + +import isaaclab.sim as sim_utils + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import resolve_task_config + +_NUM_ENVS = 16 +_STEPS = 120 +_CART_BOUND = 3.0 +_RESET_CART_POS = 0.5 +_TASK_IDS = ("Isaac-Cartpole", "Isaac-Cartpole-Direct") + + +def _cartpole_cfg(task_id: str, compute_final_obs: bool): + """Cartpole with fixed reset states, so the stable and the Warp random draws agree.""" + env_cfg, _ = resolve_task_config(task_id, "", overrides=("physics=newton_mjwarp",)) + env_cfg.seed = 7 + env_cfg.scene.num_envs = _NUM_ENVS + env_cfg.compute_final_obs = compute_final_obs + if task_id == "Isaac-Cartpole": + env_cfg.events.reset_cart_position.params.update( + position_range=(_RESET_CART_POS, _RESET_CART_POS), velocity_range=(0.0, 0.0) + ) + env_cfg.events.reset_pole_position.params.update(position_range=(0.1, 0.1), velocity_range=(0.0, 0.0)) + else: + env_cfg.initial_cart_position_range = (_RESET_CART_POS, _RESET_CART_POS) + env_cfg.initial_cart_velocity_range = (0.0, 0.0) + env_cfg.initial_pole_angle_range = (0.1, 0.1) + env_cfg.initial_pole_velocity_range = (0.0, 0.0) + return env_cfg + + +def _rollout(task_id: str, frontend: str, compute_final_obs: bool) -> dict: + """Push every cart with a constant force and record the observations of every step.""" + env_cfg = _cartpole_cfg(task_id, compute_final_obs) + sim_utils.create_new_stage() + if frontend == "warp": + env = WarpFrontend.build_env(env_cfg, task_id).unwrapped + else: + env = gym.make(task_id, cfg=env_cfg).unwrapped + rollout = {"obs": [], "dones": [], "final_obs": {}, "final_obs_key": []} + try: + env.reset() + # different pushes drive the carts out of bounds at different steps + push = torch.linspace(-0.3, 0.3, _NUM_ENVS, device=env.device)[:, None] + for step in range(_STEPS): + obs, _, terminated, truncated, extras = env.step(push) + dones = terminated | truncated + rollout["obs"].append(obs["policy"].clone()) + rollout["dones"].append(dones.clone()) + rollout["final_obs_key"].append("final_obs" in extras) + if dones.any() and "final_obs" in extras: + rollout["final_obs"][step] = extras["final_obs"]["policy"].clone() + rollout["cart_index"] = ( + 0 if task_id != "Isaac-Cartpole" else env.scene["robot"].find_joints("slider_to_cart")[0][0] + ) + finally: + env.close() + rollout["obs"] = torch.stack(rollout["obs"]) + rollout["dones"] = torch.stack(rollout["dones"]) + return rollout + + +@pytest.fixture(scope="module") +def rollouts() -> dict[tuple[str, str, bool], dict]: + arms = [(task_id, frontend, True) for task_id in _TASK_IDS for frontend in ("stable", "warp")] + arms += [(task_id, "warp", False) for task_id in _TASK_IDS] + return {arm: _rollout(*arm) for arm in arms} + + +@pytest.mark.parametrize("task_id", _TASK_IDS) +def test_final_obs_is_the_observation_before_reset(task_id: str, rollouts: dict): + rollout = rollouts[(task_id, "warp", True)] + cart = rollout["cart_index"] + reset_steps = torch.nonzero(rollout["dones"].any(dim=1)).squeeze(-1).tolist() + assert reset_steps, "no environment terminated; the terminal observation is not exercised" + assert sorted(rollout["final_obs"]) == reset_steps + for step in reset_steps: + dones = rollout["dones"][step] + final_obs, obs = rollout["final_obs"][step], rollout["obs"][step] + # the terminal observation still has the cart out of bounds, the returned one is reset + assert (final_obs[dones, cart].abs() > _CART_BOUND).all() + assert torch.equal(obs[dones, cart], torch.full_like(obs[dones, cart], _RESET_CART_POS)) + assert torch.equal(final_obs[~dones], obs[~dones]) + + +@pytest.mark.parametrize("task_id", _TASK_IDS) +def test_final_obs_is_absent_when_disabled(task_id: str, rollouts: dict): + rollout = rollouts[(task_id, "warp", False)] + assert rollout["dones"].any() + assert not any(rollout["final_obs_key"]) + + +@pytest.mark.parametrize( + ("task_id", "atol"), + # the direct twin differs from the stable task in the last bit only (at most 4.8e-7 at |obs| < 8) + [("Isaac-Cartpole", 0.0), ("Isaac-Cartpole-Direct", 1e-6)], +) +def test_final_obs_matches_stable(task_id: str, atol: float, rollouts: dict): + stable, warp = rollouts[(task_id, "stable", True)], rollouts[(task_id, "warp", True)] + assert torch.equal(warp["dones"], stable["dones"]) + torch.testing.assert_close(warp["obs"], stable["obs"], rtol=0.0, atol=atol) + assert sorted(warp["final_obs"]) == sorted(stable["final_obs"]) + for step, final_obs in stable["final_obs"].items(): + torch.testing.assert_close(warp["final_obs"][step], final_obs, rtol=0.0, atol=atol, msg=f"step {step}") diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_log.py b/source/isaaclab_experimental/test/envs/test_warp_env_log.py new file mode 100644 index 000000000000..fd5a88e0549a --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_warp_env_log.py @@ -0,0 +1,102 @@ +# 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 + +"""Logged values of the Warp environments keep the value of the step that returned them. + +Warp managers and tasks log persistent buffers that later steps overwrite. RL libraries keep the +``extras["log"]`` of every step and average them at the end of an iteration, so the environments must return +copies, like the fresh tensors of the stable environments. +""" + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +from isaaclab.sim import SimulationCfg +from isaaclab.test.utils import launch_test_simulation + +launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg()))) + +import isaaclab_tasks_experimental # noqa: F401 +import pytest +import torch +from isaaclab_experimental.envs import ManagerBasedEnvWarp +from isaaclab_experimental.envs.frontend import WarpFrontend + +import isaaclab.sim as sim_utils + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import resolve_task_config + +_NUM_ENVS = 16 +_STEPS = 120 + + +def _make_env(task_id: str): + env_cfg, _ = resolve_task_config(task_id, "", overrides=("physics=newton_mjwarp",)) + env_cfg.seed = 7 + env_cfg.scene.num_envs = _NUM_ENVS + sim_utils.create_new_stage() + return WarpFrontend.build_env(env_cfg, task_id).unwrapped + + +def test_reset_logs_keep_their_values(): + env = _make_env("Isaac-Cartpole") + logged = [] + try: + env.reset() + # different pushes drive the carts out of bounds at different steps, so episode sums differ + push = torch.linspace(-0.3, 0.3, _NUM_ENVS, device=env.device)[:, None] + for _ in range(_STEPS): + _, _, terminated, truncated, extras = env.step(push) + if (terminated | truncated).any(): + log = extras["log"] + logged.append((log, {key: value.clone() for key, value in log.items() if torch.is_tensor(value)})) + finally: + env.close() + + first_values = logged[0][1] + assert any(not torch.equal(values[key], first_values[key]) for _, values in logged[1:] for key in first_values) + for log, values in logged: + for key, value in values.items(): + assert torch.equal(log[key], value), key + + +def test_base_env_reset_logs_keep_their_values(monkeypatch: pytest.MonkeyPatch): + env_cfg, _ = resolve_task_config("Isaac-Cartpole", "", overrides=("physics=newton_mjwarp",)) + env_cfg.seed = 7 + env_cfg.scene.num_envs = _NUM_ENVS + WarpFrontend.adapt_cfg(env_cfg) + sim_utils.create_new_stage() + env = ManagerBasedEnvWarp(cfg=env_cfg) + logged = [] + try: + # a manager logs a persistent buffer that it refreshes on every reset + metric = torch.zeros((), device=env.device) + monkeypatch.setattr(env.recorder_manager, "reset", lambda env_ids=None: {"Metrics/value": metric}) + for value in range(3): + metric.fill_(value) + _, extras = env.reset() + logged.append(extras["log"]) + finally: + env.close() + + assert [log["Metrics/value"].item() for log in logged] == [0.0, 1.0, 2.0] + + +def test_direct_logs_keep_their_values(): + env = _make_env("Isaac-Cartpole-Direct") + logged = [] + try: + env.reset() + # a task logs a persistent buffer that its kernels refresh + metric = torch.zeros((), device=env.device) + env.extras["log"] = {"Metrics/value": metric} + for step in range(5): + metric.fill_(step) + _, _, _, _, extras = env.step(torch.zeros((_NUM_ENVS, 1), device=env.device)) + logged.append(extras["log"]) + finally: + env.close() + + assert [log["Metrics/value"].item() for log in logged] == [0.0, 1.0, 2.0, 3.0, 4.0] diff --git a/source/isaaclab_experimental/test/envs/test_warp_env_reset.py b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py new file mode 100644 index 000000000000..bf0ee211d098 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_warp_env_reset.py @@ -0,0 +1,80 @@ +# 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 + +"""Resets of the manager-based Warp environment.""" + +from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg + +from isaaclab.sim import SimulationCfg +from isaaclab.test.utils import launch_test_simulation + +launch_test_simulation(SimulationCfg(physics=NewtonCfg(solver_cfg=MJWarpSolverCfg()))) + +import torch +from isaaclab_experimental.envs.frontend import WarpFrontend +from isaaclab_experimental.managers import EventManager + +import isaaclab.sim as sim_utils + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.utils import resolve_task_config + +_NUM_ENVS = 8 + + +def _make_cartpole(): + env_cfg, _ = resolve_task_config("Isaac-Cartpole", "", overrides=("physics=newton_mjwarp",)) + env_cfg.seed = 7 + env_cfg.scene.num_envs = _NUM_ENVS + sim_utils.create_new_stage() + return WarpFrontend.build_env(env_cfg, "Isaac-Cartpole").unwrapped + + +def test_reset_with_mask_resets_only_selected_environments(): + env = _make_cartpole() + try: + env.reset() + actions = torch.ones((_NUM_ENVS, env.action_space.shape[-1]), device=env.device) + for _ in range(5): + env.step(actions) + robot = env.scene["robot"] + joint_pos = robot.data.joint_pos.torch.clone() + selected = torch.arange(_NUM_ENVS, device=env.device) % 2 == 0 + + env.reset(env_mask=selected) + + assert torch.equal(env.episode_length_buf[selected], torch.zeros_like(env.episode_length_buf[selected])) + assert torch.equal(env.episode_length_buf[~selected], torch.full_like(env.episode_length_buf[~selected], 5)) + assert torch.equal(robot.data.joint_pos.torch[~selected], joint_pos[~selected]) + # the reset events sample new joint positions for every selected environment + assert not torch.isclose(robot.data.joint_pos.torch[selected], joint_pos[selected]).all(dim=1).any() + finally: + env.close() + + +def test_set_term_cfg_applies_to_recorded_reset_events(): + """A replaced event term configuration must not replay the scalars recorded from the old one.""" + env = _make_cartpole() + try: + env.reset() + env.step(torch.zeros((_NUM_ENVS, env.action_space.shape[-1]), device=env.device)) + selected = torch.arange(_NUM_ENVS, device=env.device) % 2 == 0 + # records the reset events with the configured random cart offsets + env.reset(env_mask=selected) + assert env.event_manager._captured_stages[EventManager._apply_reset].num_graphs + + term_cfg = env.event_manager.get_term_cfg("reset_cart_position") + term_cfg.params["position_range"] = (0.5, 0.5) + term_cfg.params["velocity_range"] = (0.0, 0.0) + env.event_manager.set_term_cfg("reset_cart_position", term_cfg) + env.reset(env_mask=selected) + + robot = env.scene["robot"] + cart = robot.find_joints("slider_to_cart")[0][0] + cart_pos = robot.data.joint_pos.torch[selected, cart] + default_pos = robot.data.default_joint_pos.torch[selected, cart] + assert torch.allclose(cart_pos, default_pos + 0.5) + finally: + env.close() diff --git a/source/isaaclab_experimental/test/managers/test_event_manager.py b/source/isaaclab_experimental/test/managers/test_event_manager.py new file mode 100644 index 000000000000..c1ee9b0b13df --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_event_manager.py @@ -0,0 +1,99 @@ +# 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 event manager.""" + +from collections import Counter +from types import SimpleNamespace + +import pytest +import warp as wp +from isaaclab_experimental.managers import EventManager, EventTermCfg, ManagerTermBase +from isaaclab_experimental.utils import CapturedStage +from isaaclab_experimental.utils.warp import WarpCapturable + +wp.init() +pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") + +DEVICE = "cuda:0" +NUM_ENVS = 4 + + +CALLS: Counter[str] = Counter() + + +@WarpCapturable(False, reason="synchronizes with the host") +def host_synchronizing_term(env, env_mask: wp.array): + """Event term that cannot run inside a recorded stage.""" + + +def device_term(env, env_mask: wp.array): + """Event term whose work is entirely on the device.""" + CALLS["device_term"] += 1 + + +@WarpCapturable(False, reason="synchronizes with the host") +class host_call_term(ManagerTermBase): + """Class event term whose call cannot be recorded but whose reset can.""" + + def __call__(self, env, env_mask: wp.array): + pass + + def reset(self, env_mask: wp.array | None = None): + CALLS["host_call_term.reset"] += 1 + + +class host_reset_term(ManagerTermBase): + """Class event term whose reset cannot be recorded.""" + + def __call__(self, env, env_mask: wp.array): + pass + + @WarpCapturable(False, reason="synchronizes with the host") + def reset(self, env_mask: wp.array | None = None): + CALLS["host_reset_term.reset"] += 1 + + +def test_startup_only_term_does_not_keep_the_reset_stage_eager(monkeypatch): + """Startup terms run once while the environment is built, so they must not stop reset events from recording.""" + monkeypatch.setattr(CapturedStage, "enabled", True) + env = SimpleNamespace(num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True)) + manager = EventManager( + { + "startup_term": EventTermCfg(func=host_synchronizing_term, mode="startup"), + "reset_term": EventTermCfg(func=device_term, mode="reset"), + }, + env, + ) + env_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE) + step_count = wp.zeros(1, dtype=wp.int32, device=DEVICE) + CALLS.clear() + + for _ in range(2): + manager.apply(mode="reset", env_mask_wp=env_mask, global_env_step_count=step_count) + + assert CALLS["device_term"] == 1, "the second call must replay the recorded reset events" + assert manager._captured_stages[EventManager._apply_reset].num_graphs == 1 + + +def test_class_annotation_does_not_keep_the_reset_eager(monkeypatch): + """A class annotation covers the term's call; only an annotation on ``reset`` keeps the reset out of the graph.""" + monkeypatch.setattr(CapturedStage, "enabled", True) + env = SimpleNamespace(num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True)) + manager = EventManager( + { + "host_call_term": EventTermCfg(func=host_call_term, mode="startup"), + "host_reset_term": EventTermCfg(func=host_reset_term, mode="startup"), + }, + env, + ) + env_mask = wp.ones(NUM_ENVS, dtype=wp.bool, device=DEVICE) + CALLS.clear() + + for _ in range(2): + manager.reset(env_mask=env_mask) + + assert CALLS["host_call_term.reset"] == 1, "the second reset must replay the recorded reset" + assert CALLS["host_reset_term.reset"] == 2, "an annotated reset must run on every reset" diff --git a/source/isaaclab_experimental/test/managers/test_observation_manager.py b/source/isaaclab_experimental/test/managers/test_observation_manager.py new file mode 100644 index 000000000000..3094a9863da0 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_observation_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 observation manager configuration parsing.""" + +from types import SimpleNamespace + +import warp as wp +from isaaclab_experimental.managers import ObservationManager + +from isaaclab.managers import ObservationGroupCfg, ObservationTermCfg +from isaaclab.utils import configclass + + +def constant_obs(env, out: wp.array, out_dim: int): + """Warp-first observation term; the manager reads its width from ``out_dim``.""" + out.fill_(1.0) + + +@configclass +class PolicyCfg(ObservationGroupCfg): + """Group with two terms and every group-level setting at its default.""" + + first = ObservationTermCfg(func=constant_obs, params={"out_dim": 2}) + second = ObservationTermCfg(func=constant_obs, params={"out_dim": 3}) + + +def test_group_settings_are_not_parsed_as_terms(): + """Every :class:`ObservationGroupCfg` field, including ``history_order``, is a group setting, not a term.""" + env = SimpleNamespace(num_envs=4, device="cpu", sim=SimpleNamespace(is_playing=lambda: True)) + + manager = ObservationManager({"policy": PolicyCfg()}, env) + + assert manager.active_terms["policy"] == ["first", "second"] + assert manager.group_obs_dim["policy"] == (5,) 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..3901abb2c702 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_reward_manager.py @@ -0,0 +1,71 @@ +# 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 reward manager.""" + +from collections import Counter +from types import SimpleNamespace + +import pytest +import torch +import warp as wp +from isaaclab_experimental.managers import ManagerTermBase, RewardManager +from isaaclab_experimental.utils import CapturedStage +from isaaclab_experimental.utils.warp import WarpCapturable + +from isaaclab.managers import RewardTermCfg + +wp.init() +pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") + +DEVICE = "cuda:0" +NUM_ENVS = 4 +CALLS: Counter[float] = Counter() + + +@WarpCapturable(lambda params: params.get("sensor_cfg") is None, reason="refreshes the sensor on the host") +def constant_reward(env, out: wp.array, value: float, sensor_cfg: str | None = None): + """Function term that is capturable only without a sensor.""" + CALLS[value] += 1 + out.fill_(value) + + +@WarpCapturable(False, reason="synchronizes with the host") +class ConstantReward(ManagerTermBase): + """Class term that is never capturable.""" + + def __call__(self, env, out: wp.array, value: float): + CALLS[value] += 1 + out.fill_(value) + + +def test_capturability_is_decided_per_term_from_its_parameters(monkeypatch): + """The same function term is recorded for one parameter set and runs eagerly for another. + + An eager term runs its Python on every step; a recorded one only while its graph is recorded. + The guard of the predicate term raises if the instance with a sensor is ever recorded. + """ + monkeypatch.setattr(CapturedStage, "enabled", True) + env = SimpleNamespace( + num_envs=NUM_ENVS, device=DEVICE, sim=SimpleNamespace(is_playing=lambda: True), max_episode_length_s=1.0 + ) + cfg = { + "function_recorded": RewardTermCfg(func=constant_reward, weight=1.0, params={"value": 1.0}), + "function_eager": RewardTermCfg( + func=constant_reward, weight=1.0, params={"value": 2.0, "sensor_cfg": "height_scanner"} + ), + "class_eager": RewardTermCfg(func=ConstantReward, weight=1.0, params={"value": 4.0}), + } + manager = RewardManager(cfg, env) + CALLS.clear() + + for _ in range(3): + reward = manager.compute(dt=0.5) + wp.synchronize() + + assert CALLS == {1.0: 1, 2.0: 3, 4.0: 3} + # each eager term splits the recording, so it holds one graph before each of the two eager terms and one after + assert manager._captured_stages[RewardManager.compute].num_graphs == 3 + assert torch.equal(reward, torch.full((NUM_ENVS,), (1.0 + 2.0 + 4.0) * 0.5, device=DEVICE)) 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..8d60dbad3329 --- /dev/null +++ b/source/isaaclab_experimental/test/managers/test_termination_manager.py @@ -0,0 +1,39 @@ +# 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 termination manager.""" + +from types import SimpleNamespace + +import pytest +import warp as wp +from isaaclab_experimental.managers import TerminationManager + +from isaaclab.managers import TerminationTermCfg + + +def copy_flags(env, out: wp.array): + """Warp-first termination term that reports the environment's preset flags.""" + wp.copy(out, env.flags) + + +def test_episode_termination_metric_averages_over_all_environments(): + """Like the stable manager, the metric is the share of all environments whose last episode the term ended. + + Only environment 0 is reset. It terminated, so averaging over the reset selection would log 1.0. + """ + env = SimpleNamespace( + num_envs=4, + device="cpu", + sim=SimpleNamespace(is_playing=lambda: True), + flags=wp.array([True, True, False, False], dtype=wp.bool, device="cpu"), + ) + manager = TerminationManager({"fell": TerminationTermCfg(func=copy_flags)}, env) + manager.compute() + + env_mask = wp.array([True, False, False, False], dtype=wp.bool, device="cpu") + extras = manager.reset(env_mask=env_mask) + + assert float(extras["Episode_Termination/fell"]) == pytest.approx(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 8f4bb904c7c1..000000000000 --- a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py +++ /dev/null @@ -1,363 +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 - - -# ====================================================================== -# 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): - """Once capped to NOT_CAPTURED, a later capturable=True registration must not lift the cap.""" - switch = ManagerCallSwitch(cfg_source={"default": 2}) - switch.register_manager_capturability("RewardManager", capturable=False) - switch.register_manager_capturability("RewardManager", capturable=True) - 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): - """WARP_NOT_CAPTURED calls the warp fn exactly once per call (no warm-up, no stable call).""" - switch = ManagerCallSwitch(cfg_source={"default": 1}) - called = {"stable": 0, "warp": 0} - - def stable_fn(): - called["stable"] += 1 - return "stable_result" - - def warp_fn(): - called["warp"] += 1 - return "warp_result" - - call = {"stage": "RewardManager_compute", "warp_call": {"fn": warp_fn}, "stable_call": {"fn": stable_fn}} - result = switch.call_stage(**call) - self.assertEqual(called, {"stable": 0, "warp": 1}, "Eager mode should call the warp fn exactly once") - self.assertEqual(result, "warp_result") - - switch.call_stage(**call) - self.assertEqual(called["warp"], 2, "Eager mode should call fn again each time") - - 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_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) - - -# ====================================================================== -# 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_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..62e809ed2bf1 --- /dev/null +++ b/source/isaaclab_experimental/test/utils/test_mask_boundary_gate.py @@ -0,0 +1,180 @@ +# 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-first execution path. + +Warp frontend code performs no data-dependent GPU-to-host synchronization outside named boundaries: + +* **Scope**: every function in the experimental production tree, plus every ``*mask*``-named function in + the core and Newton trees (the functions that are the Warp path there). Torch-first core code is out of + scope. +* **Sanctioned boundaries**: the functions in :data:`SANCTIONED_BOUNDARIES`. The list is exact; a stale + entry fails the gate. +* **Runtime backstop**: ``ISAACLAB_SYNC_DEBUG=1`` runs eager Warp stages under + ``torch.cuda.set_sync_debug_mode("error")``, trapping synchronizations this scan cannot see. + +A second test pins every ``@WarpCapturable(False)`` opt-out, so non-capturable terms are listed in one place. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[4] + +# Attribute calls that synchronize with, or copy from, device memory. +SYNC_ATTRS = {"nonzero", "argwhere", "item", "cpu", "tolist", "numpy"} + +# Trees scanned in full (production code only). +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. +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. +NON_STEP_FUNCTIONS = { + "__init__", + "__post_init__", + "__str__", + "__repr__", + "__del__", + "_prepare_terms", + "serialize", + "get_active_iterable_terms", + "_debug_vis_callback", + "_set_debug_vis_impl", + # deprecated IO-descriptor export surfaces + "IO_descriptor", + "_collect_io_descriptors", + "get_term_cfg", +} + +# Sanctioned host-synchronization boundaries, by (path suffix, function name). +SANCTIONED_BOUNDARIES = { + # one host predicate per step decides whether the reset pipeline runs (W7); the direct env decides on the device + ("envs/manager_based_rl_env_warp.py", "step"), + # mask to index compaction for the managers that reset by index (W11) + ("envs/manager_based_env_warp.py", "reset"), + ("envs/manager_based_env_warp.py", "_reset_mask"), + ("envs/manager_based_rl_env_warp.py", "_reset_mask"), + # the stable randomization adapters take environment indices; they run outside recorded graphs + ("envs/mdp/events.py", "_mask_to_env_ids"), + # the camera's empty-reset predicate + ("isaaclab/sensors/camera/camera.py", "_env_mask_has_any"), + # joint-limit writes mutate the solver model (event-driven, not per-step) + ("isaaclab_newton/assets/articulation/articulation.py", "write_joint_position_limit_to_sim_mask"), +} + + +def _sync_calls(func_node: ast.AST) -> list[str]: + """Names of synchronizing attribute calls inside a function body.""" + return [ + node.func.attr + 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 + ] + + +def _scan(root: Path, mask_only: bool) -> dict[tuple[str, str], list[str]]: + """Map (relative path, function) to synchronizing calls for the 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 + for node in ast.walk(ast.parse(path.read_text(), filename=rel)): + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) or node.name in NON_STEP_FUNCTIONS: + continue + if mask_only and "mask" not in node.name: + continue + hits = _sync_calls(node) + if hits: + findings[(rel, node.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_first_code_has_no_unsanctioned_host_syncs(): + """Every synchronization in 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 = {key: hits for key, hits in findings.items() if not _is_sanctioned(*key)} + assert not violations, "Unsanctioned host syncs on the mask-first path:\n" + "\n".join( + f" {rel}::{name} -> {hits}" for (rel, name), hits in sorted(violations.items()) + ) + + matched = {key for key in findings if _is_sanctioned(*key)} + 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 boundaries (function gone or no longer syncs): {sorted(stale)}" + + +# ``@WarpCapturable(False)`` and parameter-predicate opt-outs, by (path suffix, name). +EXPECTED_NON_CAPTURABLE = { + ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_com"), + ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_mass"), + ("isaaclab_experimental/envs/mdp/events.py", "randomize_rigid_body_material"), +} + +NON_CAPTURABLE_SCAN_ROOTS = [ + "source/isaaclab_experimental/isaaclab_experimental", + "source/isaaclab_tasks_experimental/isaaclab_tasks_experimental", +] + + +def _non_capturable_targets(tree: ast.Module) -> list[str]: + """Names of functions and classes decorated with ``@WarpCapturable(False)`` or a parameter predicate.""" + targets = [] + for node in ast.walk(tree): + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue + for decorator in node.decorator_list: + if ( + isinstance(decorator, ast.Call) + and isinstance(decorator.func, ast.Name) + and decorator.func.id == "WarpCapturable" + and decorator.args + and not (isinstance(decorator.args[0], ast.Constant) and decorator.args[0].value is True) + ): + targets.append(node.name) + return targets + + +def test_non_capturable_terms_are_inventoried(): + """Every ``@WarpCapturable(False)`` or predicate opt-out is listed 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 _non_capturable_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, ( + f"Non-capturable inventory drift.\n unexpected: {sorted(found - expected)}\n" + f" missing: {sorted(expected - found)}" + ) diff --git a/source/isaaclab_experimental/test/utils/test_warp_capture.py b/source/isaaclab_experimental/test/utils/test_warp_capture.py new file mode 100644 index 000000000000..0c8c3d49c8c9 --- /dev/null +++ b/source/isaaclab_experimental/test/utils/test_warp_capture.py @@ -0,0 +1,241 @@ +# 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 capture of Warp frontend stages.""" + +from __future__ import annotations + +import pytest +import torch +import warp as wp +from isaaclab_experimental.utils.warp_capture import ( + SYNC_DEBUG_ENV_VAR, + CapturedStage, + captured, + eager, + reset_captured_stages, +) + +wp.init() +pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") + +DEVICE = "cuda:0" + + +@wp.kernel +def _affine(values: wp.array(dtype=wp.int32), scale: wp.int32, shift: wp.int32): + values[wp.tid()] = values[wp.tid()] * scale + shift + + +class _AffineStage: + """Stage that applies ``x * scale + shift`` on the device and counts its Python invocations.""" + + def __init__(self, scale: int = 1, shift: int = 1): + self.scale, self.shift, self.calls = scale, shift, 0 + + def __call__(self, values: wp.array, *args) -> wp.array: + self.calls += 1 + wp.launch(_affine, dim=values.shape[0], inputs=[values, self.scale, self.shift], device=values.device) + return values + + +def _read(values: wp.array) -> list[int]: + wp.synchronize_device(values.device) + return values.numpy().tolist() + + +@pytest.fixture(autouse=True) +def enabled(monkeypatch): + monkeypatch.delenv(SYNC_DEBUG_ENV_VAR, raising=False) + monkeypatch.setattr(CapturedStage, "enabled", True) + + +def test_first_call_records_and_runs_the_stage_once(): + """Recording has no eager warm-up: a stateful stage advances once per call, including the first.""" + fn = _AffineStage() + stage = CapturedStage(fn, DEVICE) + values = wp.zeros(2, dtype=wp.int32, device=DEVICE) + + assert stage(values) is values + assert _read(values) == [1, 1] + assert stage(values) is values + assert _read(values) == [2, 2] + assert (fn.calls, stage.num_graphs) == (1, 1), "the second call must replay the recorded graph" + + +def test_changed_argument_records_the_stage_again(): + """A graph replays its recorded pointers and scalars, so a new array or scalar records a new graph.""" + fn = _AffineStage() + stage = CapturedStage(fn, DEVICE) + old, new = wp.zeros(1, dtype=wp.int32, device=DEVICE), wp.zeros(1, dtype=wp.int32, device=DEVICE) + + stage(old) + stage(new) + stage(new) + stage(new, 5) + assert (_read(old), _read(new), fn.calls, stage.num_graphs) == ([1], [3], 3, 3) + + +def test_physics_rebind_records_the_stage_again(): + fn = _AffineStage() + stage = CapturedStage(fn, DEVICE) + values = wp.zeros(1, dtype=wp.int32, device=DEVICE) + stage(values) + + CapturedStage.invalidate() + stage(values) + stage(values) + assert (fn.calls, _read(values), stage.num_graphs) == (2, [3], 1) + + +def test_stage_runs_eagerly_when_disabled_or_on_cpu(monkeypatch): + on_cpu = _AffineStage() + cpu_stage = CapturedStage(on_cpu, "cpu") + cpu_values = wp.zeros(1, dtype=wp.int32, device="cpu") + cpu_stage(cpu_values) + cpu_stage(cpu_values) + + monkeypatch.setattr(CapturedStage, "enabled", False) + disabled = _AffineStage() + stage = CapturedStage(disabled, DEVICE) + values = wp.zeros(1, dtype=wp.int32, device=DEVICE) + stage(values) + stage(values) + + assert (on_cpu.calls, _read(cpu_values), cpu_stage.num_graphs) == (2, [2], 0) + assert (disabled.calls, _read(values), stage.num_graphs) == (2, [2], 0) + + +def test_eager_call_runs_between_its_graphs_on_every_call(): + """The recording runs the graph before an eager call first, and every replay keeps that order.""" + before, between, after = _AffineStage(2, 0), _AffineStage(1, 10), _AffineStage(3, 0) + + def body(values: wp.array) -> wp.array: + before(values) + eager(between, values) + return after(values) + + stage = CapturedStage(body, DEVICE) + values = wp.ones(1, dtype=wp.int32, device=DEVICE) + + stage(values) + # the non-commuting steps only give these values in order: 1 * 2 + 10 = 12 -> 36, then 36 * 2 + 10 = 82 -> 246 + assert _read(values) == [36] + stage(values) + assert (_read(values), between.calls, before.calls, after.calls, stage.num_graphs) == ([246], 2, 1, 1, 2) + + +def test_eager_call_first_leaves_one_graph_after_it(): + """A stage making its eager call first records all its capturable work into the single graph after it.""" + eager_fn, first, second = _AffineStage(1, 10), _AffineStage(2, 0), _AffineStage(3, 0) + + def body(values: wp.array) -> wp.array: + eager(eager_fn, values) + first(values) + return second(values) + + stage = CapturedStage(body, DEVICE) + values = wp.ones(1, dtype=wp.int32, device=DEVICE) + stage(values) + stage(values) + + assert _read(values) == [((1 + 10) * 6 + 10) * 6] + ((steps, _),) = stage._recordings.values() + assert [isinstance(step, wp.Graph) for step in steps] == [True, False, True] + # the graph before the eager call is empty; the last one multiplies by 6 + wp.capture_launch(steps[0]) + assert _read(values) == [456] + wp.capture_launch(steps[-1]) + assert _read(values) == [456 * 6] + + +class _Owner: + device = DEVICE + + def __init__(self): + self.inner_fn, self.outer_fn = _AffineStage(2, 0), _AffineStage(1, 1) + + @captured + def inner(self, values: wp.array) -> wp.array: + return self.inner_fn(values) + + @captured + def outer(self, values: wp.array) -> wp.array: + self.outer_fn(values) + return self.inner(values) + + +def test_decorated_methods_own_their_stages_and_nest_into_an_enclosing_recording(): + owner = _Owner() + values = wp.zeros(1, dtype=wp.int32, device=DEVICE) + + assert owner.outer(values) is values + owner.outer(values) + # (0 + 1) * 2 = 2 -> (2 + 1) * 2 = 6: the inner call is part of the outer graph + assert _read(values) == [6] + assert (owner.outer_fn.calls, owner.inner_fn.calls) == (1, 1) + stages = owner._captured_stages + assert (stages[_Owner.outer].num_graphs, stages[_Owner.inner].num_graphs) == (1, 0) + + reset_captured_stages(owner) + owner.outer(values) + assert (owner.outer_fn.calls, _read(values)) == (2, [14]) + + +def test_eager_call_of_a_nested_stage_splits_the_enclosing_recording(): + eager_fn, inner_fn, outer_fn = _AffineStage(1, 10), _AffineStage(2, 0), _AffineStage(3, 0) + nested = CapturedStage(lambda values: (eager(eager_fn, values), inner_fn(values)), DEVICE) + outer = CapturedStage(lambda values: (outer_fn(values), nested(values)), DEVICE) + values = wp.ones(1, dtype=wp.int32, device=DEVICE) + + outer(values) + outer(values) + # 1 * 3 + 10 = 13 -> 26, then 26 * 3 + 10 = 88 -> 176: the nested eager call runs between the outer graphs + assert (_read(values), eager_fn.calls, inner_fn.calls, outer_fn.calls) == ([176], 2, 1, 1) + assert (outer.num_graphs, nested.num_graphs) == (2, 0) + + +def test_eager_call_raises_inside_a_capture_no_stage_started(): + """A call that cannot be recorded must not be captured silently into a foreign graph.""" + fn = _AffineStage() + values = wp.zeros(1, dtype=wp.int32, device=DEVICE) + + with pytest.raises(RuntimeError, match="no CapturedStage started"): + with wp.ScopedDevice(DEVICE), wp.ScopedCapture(): + eager(fn, values) + assert fn.calls == 0 + + +def test_stage_raising_while_recording_ends_the_capture(): + """A failed recording must not leave the device capturing; the next call records again.""" + recorded, eager_fn = _AffineStage(2, 0), _AffineStage(1, 10) + + def body(values: wp.array) -> wp.array: + recorded(values) + eager(eager_fn, values) + if recorded.calls == 1: + raise ValueError("first recording fails") + return values + + stage = CapturedStage(body, DEVICE) + values = wp.ones(1, dtype=wp.int32, device=DEVICE) + + with pytest.raises(ValueError, match="first recording fails"): + stage(values) + assert not wp.get_device(DEVICE).is_capturing + stage(values) + # the failed recording ran its first graph and the eager call: 1 * 2 + 10 = 12, then 12 * 2 + 10 = 34 + assert (_read(values), stage.num_graphs) == ([34], 2) + + +def test_sync_debug_trap_raises_on_a_hidden_host_sync(monkeypatch): + monkeypatch.setattr(CapturedStage, "enabled", False) + monkeypatch.setenv(SYNC_DEBUG_ENV_VAR, "1") + flags = torch.ones(4, dtype=torch.bool, device=DEVICE) + stage = CapturedStage(lambda: flags.any().item(), DEVICE) + + with pytest.raises(RuntimeError, match="synchroniz"): + stage() + assert torch.cuda.get_sync_debug_mode() == 0, "the trap must be released after the stage" diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py deleted file mode 100644 index 7a715c639d44..000000000000 --- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.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 - -"""Tests for WarpGraphCache capture/replay and warm-up behavior.""" - -from __future__ import annotations - -import unittest - -import warp as wp -from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache - - -@wp.kernel -def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)): - """Simple warp kernel: b[i] = a[i] + 1.""" - i = wp.tid() - b[i] = a[i] + 1.0 - - -class TestWarpGraphCache(unittest.TestCase): - """Tests for :class:`WarpGraphCache`.""" - - def setUp(self): - self.device = "cuda:0" - self.cache = WarpGraphCache() - - # ------------------------------------------------------------------ - # Warm-up - # ------------------------------------------------------------------ - - def test_warmup_runs_before_capture(self): - """The function should be called eagerly (warm-up) before graph capture. - - We verify this by counting total invocations on the first call. - Warm-up = 1, capture = 1, so fn should be called exactly 2 times. - """ - call_count = [0] - src = wp.zeros(4, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - def counted_launch(): - call_count[0] += 1 - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - # First call: warm-up + capture - self.cache.capture_or_replay("stage_a", counted_launch) - self.assertEqual(call_count[0], 2, "First call should invoke fn twice (warm-up + capture)") - - # Second call: replay only - self.cache.capture_or_replay("stage_a", counted_launch) - self.assertEqual(call_count[0], 2, "Replay should NOT invoke fn again") - - # ------------------------------------------------------------------ - # Capture / replay correctness - # ------------------------------------------------------------------ - - def test_replay_uses_updated_input(self): - """Replay should re-read from the same input buffer (pointer-stable). - - CUDA graph replay re-executes the same kernel on the same memory - addresses. If we update the input buffer in-place, the output - should reflect the new values. - """ - src = wp.full(4, value=1.0, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - def my_fn(): - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - # Capture - first = self.cache.capture_or_replay("replay_test", my_fn) - - # Update input in-place - wp.copy(src, wp.full(4, value=10.0, dtype=wp.float32, device=self.device)) - - # Replay — should see updated input - result = self.cache.capture_or_replay("replay_test", my_fn) - # Replay returns the cached object from the capture call, not a new buffer. - self.assertIs(result, first) - result_np = result.numpy() - for val in result_np: - self.assertAlmostEqual(val, 11.0, places=5) - - def test_multiple_stages_independent(self): - """Different stages should be captured and replayed independently.""" - src_a = wp.full(4, value=1.0, dtype=wp.float32, device=self.device) - dst_a = wp.zeros(4, dtype=wp.float32, device=self.device) - src_b = wp.full(4, value=5.0, dtype=wp.float32, device=self.device) - dst_b = wp.zeros(4, dtype=wp.float32, device=self.device) - - def fn_a(): - wp.launch(_add_one, dim=4, inputs=[src_a, dst_a], device=self.device) - return dst_a - - def fn_b(): - wp.launch(_add_one, dim=4, inputs=[src_b, dst_b], device=self.device) - return dst_b - - result_a = self.cache.capture_or_replay("stage_a", fn_a) - result_b = self.cache.capture_or_replay("stage_b", fn_b) - - self.assertAlmostEqual(result_a.numpy()[0], 2.0, places=5) - self.assertAlmostEqual(result_b.numpy()[0], 6.0, places=5) - - # Replay both - result_a2 = self.cache.capture_or_replay("stage_a", fn_a) - result_b2 = self.cache.capture_or_replay("stage_b", fn_b) - self.assertIs(result_a, result_a2) - self.assertIs(result_b, result_b2) - - # ------------------------------------------------------------------ - # Invalidation - # ------------------------------------------------------------------ - - def test_invalidate_all(self): - """invalidate() with no args should drop all cached graphs.""" - src = wp.zeros(4, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - call_count = [0] - - def my_fn(): - call_count[0] += 1 - wp.launch(_add_one, dim=4, inputs=[src, dst], device=self.device) - return dst - - self.cache.capture_or_replay("s1", my_fn) - self.assertEqual(call_count[0], 2) # warm-up + capture - - self.cache.invalidate() - - # After invalidation, next call should re-warm-up and re-capture - self.cache.capture_or_replay("s1", my_fn) - self.assertEqual(call_count[0], 4) # 2 more (warm-up + capture) - - def test_invalidate_single_stage(self): - """invalidate(stage) should only drop the named stage.""" - src = wp.zeros(4, dtype=wp.float32, device=self.device) - dst_a = wp.zeros(4, dtype=wp.float32, device=self.device) - dst_b = wp.zeros(4, dtype=wp.float32, device=self.device) - - count_a = [0] - count_b = [0] - - def fn_a(): - count_a[0] += 1 - wp.launch(_add_one, dim=4, inputs=[src, dst_a], device=self.device) - return dst_a - - def fn_b(): - count_b[0] += 1 - wp.launch(_add_one, dim=4, inputs=[src, dst_b], device=self.device) - return dst_b - - self.cache.capture_or_replay("a", fn_a) - self.cache.capture_or_replay("b", fn_b) - self.assertEqual(count_a[0], 2) - self.assertEqual(count_b[0], 2) - - # Invalidate only "a"; a stage that was never captured is a no-op - self.cache.invalidate("a") - self.cache.invalidate("nonexistent") - - self.cache.capture_or_replay("a", fn_a) - self.cache.capture_or_replay("b", fn_b) - self.assertEqual(count_a[0], 4, "Stage 'a' should re-capture after invalidation") - self.assertEqual(count_b[0], 2, "Stage 'b' should replay (not re-capture)") - - # ------------------------------------------------------------------ - # Args / kwargs forwarding - # ------------------------------------------------------------------ - - def test_args_and_kwargs_forwarded(self): - """capture_or_replay should forward args and kwargs to fn.""" - src = wp.full(4, value=2.0, dtype=wp.float32, device=self.device) - dst = wp.zeros(4, dtype=wp.float32, device=self.device) - - def my_fn(a, b, device="cuda:0"): - wp.launch(_add_one, dim=4, inputs=[a, b], device=device) - return b - - result = self.cache.capture_or_replay( - "args_test", - my_fn, - args=(src, dst), - kwargs={"device": self.device}, - ) - self.assertAlmostEqual(result.numpy()[0], 3.0, places=5) - - -if __name__ == "__main__": - unittest.main() diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..0c69e3841f70 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,13 @@ +* Changed the env selection helpers of the Newton assets to resolve ``env_ids`` through + :func:`~isaaclab.utils.warp.utils.resolve_ids` and ``env_mask`` through + :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`. Any integer sequence and slices with a negative step are accepted, + ids on another device are moved to the asset's device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected + with a :class:`TypeError`. +* Changed the body, joint and tendon selection helpers of the Newton assets to resolve ids and masks through + :func:`~isaaclab.utils.warp.utils.resolve_ids`, as the env selection does, so slices other than ``slice(None)`` are + accepted. +* Changed the joint-coordinate writes of :class:`~isaaclab_newton.assets.Articulation` to write env ids into a + scratch mask of the articulation instead of allocating a mask on every call. +* Changed the mask writers of :class:`~isaaclab_newton.assets.RigidObjectCollection` and the fixed-tendon mask setters + of :class:`~isaaclab_newton.assets.Articulation` to write through the masks, as the other Newton mask writers do, + instead of converting the masks to indices, which synchronized the host with the GPU on every call. diff --git a/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst new file mode 100644 index 000000000000..638c02730a49 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-id-mask.fixed.rst @@ -0,0 +1,2 @@ +* Fixed the fixed-tendon mask setters of :class:`~isaaclab_newton.assets.Articulation` accepting data whose shape + does not match the articulation. diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-rng-shared.skip b/source/isaaclab_newton/changelog.d/jichuanh-warp-rng-shared.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.added.rst new file mode 100644 index 000000000000..bce8fc255e4e --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.added.rst @@ -0,0 +1 @@ +* Added ``env_mask`` to :meth:`~isaaclab_newton.assets.CableObject.reset`. diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst new file mode 100644 index 000000000000..4f283403787d --- /dev/null +++ b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst @@ -0,0 +1,5 @@ +* Fixed :meth:`~isaaclab_newton.assets.Articulation.reset` with only ``env_mask`` resetting the actuator state + (delay buffers, network history, and Newton-native actuator state) of every environment. +* Fixed :meth:`~isaaclab_newton.assets.RigidObject.reset` and + :meth:`~isaaclab_newton.assets.RigidObjectCollection.reset` ignoring ``env_mask`` and clearing the external + wrenches of every environment. diff --git a/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.minor b/source/isaaclab_newton/changelog.d/jichuanh-warp-scene-mask-native.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py index 4d90ca825ca6..56ef68b261bf 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/actuator_control.py @@ -232,9 +232,9 @@ def submit_commands(self, collection: ActuatorCollection) -> None: articulation._ALL_ENV_MASK, ) - def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None: + def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None: if self._native_actuator_path_active and self._articulation._backend_manager._adapter is not None: - self._articulation._backend_manager._adapter.reset(env_ids) + self._articulation._backend_manager._adapter.reset(env_ids, env_mask) def _bam_actuators(self) -> list[tuple[Actuator, list[str]]]: """Return the Newton BAM actuators that drive this articulation, each with the BAM groups it drives. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index f6b03aa5b673..d1bacfb77805 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -12,7 +12,7 @@ import re import warnings from collections.abc import Sequence -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import numpy as np import torch @@ -33,6 +33,7 @@ from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -382,7 +383,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None env_ids = slice(None) # reset actuators, including backend-native actuator state. None selects all # environments; delayed-actuator buffers do not accept a slice. - self.actuators.reset(None if env_ids == slice(None) else env_ids) + self.actuators.reset(None if env_ids == slice(None) else 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) @@ -2298,15 +2299,11 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( Operations - Newton Actuator Parameter Writers. """ - @staticmethod - @wp.kernel(enable_backward=False) - def _build_env_mask_kernel(mask: wp.array(dtype=wp.bool), indices: wp.array(dtype=Any)): - i = wp.tid() - mask[wp.int32(indices[i])] = True - def _env_ids_to_mask(self, env_ids: wp.array | torch.Tensor) -> wp.array: """Convert env_ids to a boolean Warp mask. + Needed because :func:`scatter_joint_coordinates` only takes a mask. + Args: env_ids: Environment indices as returned by :meth:`_resolve_env_ids`, which may be a warp array or a torch tensor of any integer width. @@ -2316,11 +2313,11 @@ def _env_ids_to_mask(self, env_ids: wp.array | torch.Tensor) -> wp.array: """ if env_ids is self._ALL_INDICES: return self._ALL_ENV_MASK - if isinstance(env_ids, torch.Tensor): - env_ids = wp.from_torch(env_ids) - mask = wp.zeros(self.num_instances, dtype=wp.bool, device=self.device) - wp.launch(self._build_env_mask_kernel, dim=env_ids.shape[0], inputs=[mask, env_ids], device=self.device) - return mask + if isinstance(env_ids, wp.array) and env_ids.dtype != wp.int32: + env_ids = wp.to_torch(env_ids) + return resolve_1d_mask( + ids=env_ids, all_mask=self._ALL_ENV_MASK, scratch_mask=self._scratch_env_mask, device=self.device + ) """ Operations - Setters. @@ -2716,12 +2713,17 @@ def set_fixed_tendon_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) - self.set_fixed_tendon_stiffness_index( - stiffness=self._select_full_data(stiffness, env_ids, fixed_tendon_ids), - fixed_tendon_ids=fixed_tendon_ids, - env_ids=env_ids, + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK) + self.assert_shape_and_dtype_mask(stiffness, (env_mask, fixed_tendon_mask), wp.float32, "stiffness") + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask + if isinstance(stiffness, float) + else shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self.num_instances, self.num_fixed_tendons), + inputs=[stiffness, env_mask, fixed_tendon_mask], + outputs=[self.data._fixed_tendon_stiffness], + device=self.device, ) def set_fixed_tendon_damping_index( @@ -2813,12 +2815,17 @@ def set_fixed_tendon_damping_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) - self.set_fixed_tendon_damping_index( - damping=self._select_full_data(damping, env_ids, fixed_tendon_ids), - fixed_tendon_ids=fixed_tendon_ids, - env_ids=env_ids, + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK) + self.assert_shape_and_dtype_mask(damping, (env_mask, fixed_tendon_mask), wp.float32, "damping") + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask + if isinstance(damping, float) + else shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(self.num_instances, self.num_fixed_tendons), + inputs=[damping, env_mask, fixed_tendon_mask], + outputs=[self.data._fixed_tendon_damping], + device=self.device, ) def set_fixed_tendon_limit_stiffness_index( @@ -2930,12 +2937,17 @@ def set_fixed_tendon_position_limit_mask( fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) - self.set_fixed_tendon_position_limit_index( - limit=self._select_full_data(limit, env_ids, fixed_tendon_ids), - fixed_tendon_ids=fixed_tendon_ids, - env_ids=env_ids, + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK) + self.assert_shape_and_dtype_mask(limit, (env_mask, fixed_tendon_mask), wp.vec2f, "limit") + if isinstance(limit, float): + raise ValueError("Fixed tendon position limits must be a tensor or array, not a float.") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.vec2f), + dim=(self.num_instances, self.num_fixed_tendons), + inputs=[limit, env_mask, fixed_tendon_mask], + outputs=[self.data._fixed_tendon_pos_limits], + device=self.device, ) def set_fixed_tendon_rest_length_index( @@ -3136,10 +3148,27 @@ def write_fixed_tendon_properties_to_sim_mask( fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are updated. env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.write_fixed_tendon_properties_to_sim_index( - fixed_tendon_ids=self._resolve_fixed_tendon_mask(fixed_tendon_mask), - env_ids=self._resolve_env_mask(env_mask), + env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) + fixed_tendon_mask = self._resolve_mask(fixed_tendon_mask, self._ALL_FIXED_TENDON_MASK) + wp.launch( + articulation_kernels.copy_fixed_tendon_properties_with_mask, + dim=(self.num_instances, self.num_fixed_tendons), + inputs=[ + env_mask, + fixed_tendon_mask, + self.data._fixed_tendon_stiffness, + self.data._fixed_tendon_damping, + self.data._fixed_tendon_pos_limits, + ], + outputs=[ + self.data._sim_bind_fixed_tendon_stiffness, + self.data._sim_bind_fixed_tendon_damping, + self.data._sim_bind_fixed_tendon_pos_limits, + ], + device=self.device, ) + # the solver keeps its own copy of the tendon properties and only re-reads them when notified + self._backend_manager.add_model_change(ModelFlags.TENDON_PROPERTIES) def set_spatial_tendon_stiffness_index( self, @@ -3445,6 +3474,7 @@ def _clear_callbacks(self) -> None: def _create_buffers(self): self._ALL_INDICES = wp.array(np.arange(self.num_instances, dtype=np.int32), device=self.device) self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device) + self._scratch_env_mask = wp.zeros((self.num_instances,), dtype=wp.bool, device=self.device) self._ALL_JOINT_INDICES = wp.array(np.arange(self.num_joints, dtype=np.int32), device=self.device) self._ALL_JOINT_MASK = wp.ones((self.num_joints,), dtype=wp.bool, device=self.device) self._ALL_BODY_INDICES = wp.array(np.arange(self.num_bodies, dtype=np.int32), device=self.device) @@ -3720,15 +3750,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No Returns: Environment indices. """ - if isinstance(env_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device) def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve joint indices to a warp array or tensor. @@ -3739,13 +3761,7 @@ def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array Returns: A warp array of joint indices or a tensor of joint indices. """ - if isinstance(joint_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(joint_ids, list): - return wp.array(joint_ids, dtype=wp.int32, device=self.device) - if (joint_ids is None) or (joint_ids == slice(None)): - return self._ALL_JOINT_INDICES - return joint_ids + return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES, device=self.device) def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve body indices to a warp array or tensor. @@ -3756,13 +3772,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | Returns: A warp array of body indices or a tensor of body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) def _resolve_fixed_tendon_ids( self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None @@ -3775,13 +3785,7 @@ def _resolve_fixed_tendon_ids( Returns: A warp array of tendon indices or a tensor of tendon indices. """ - if isinstance(tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(tendon_ids, list): - return wp.array(tendon_ids, dtype=wp.int32, device=self.device) - if (tendon_ids is None) or (tendon_ids == slice(None)): - return self._ALL_FIXED_TENDON_INDICES - return tendon_ids + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self.device) def _resolve_spatial_tendon_ids( self, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None @@ -3794,27 +3798,7 @@ def _resolve_spatial_tendon_ids( Returns: A warp array of spatial tendon indices or a tensor of spatial tendon indices. """ - if isinstance(spatial_tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(spatial_tendon_ids, list): - return wp.array(spatial_tendon_ids, dtype=wp.int32, device=self.device) - if (spatial_tendon_ids is None) or (spatial_tendon_ids == slice(None)): - return self._ALL_SPATIAL_TENDON_INDICES - return spatial_tendon_ids - - def _resolve_env_mask(self, env_mask: wp.array | torch.Tensor | None) -> wp.array | torch.Tensor: - """Resolve an environment mask to environment indices.""" - return self._ALL_INDICES if env_mask is None else self._mask_to_ids(env_mask) - - def _resolve_fixed_tendon_mask(self, fixed_tendon_mask: wp.array | torch.Tensor | None) -> wp.array | torch.Tensor: - """Resolve a fixed tendon mask to fixed tendon indices.""" - return self._ALL_FIXED_TENDON_INDICES if fixed_tendon_mask is None else self._mask_to_ids(fixed_tendon_mask) - - @staticmethod - def _mask_to_ids(mask: wp.array | torch.Tensor) -> torch.Tensor: - """Convert a boolean mask to int32 indices.""" - mask = wp.to_torch(mask) if isinstance(mask, wp.array) else mask - return torch.nonzero(mask)[:, 0].to(torch.int32) + return resolve_ids(ids=spatial_tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES, device=self.device) @staticmethod def _to_torch_ids(ids: wp.array | torch.Tensor) -> torch.Tensor: @@ -3839,12 +3823,7 @@ def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.arra Returns: A warp array of mask. """ - if mask is None: - return full_mask - - if isinstance(mask, torch.Tensor): - return wp.from_torch(mask, dtype=wp.bool) - return mask + return resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device) """ Deprecated methods. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py index 60544e3a4100..7ad2bd59f057 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py @@ -553,6 +553,36 @@ def float_data_to_buffer_with_indices_kernel( return _FLOAT_DATA_TO_BUFFER_WITH_INDICES_DISPATCHER.select(env_ids, joint_ids) +@wp.kernel +def copy_fixed_tendon_properties_with_mask( + env_mask: wp.array(dtype=wp.bool), + tendon_mask: wp.array(dtype=wp.bool), + stiffness: wp.array2d(dtype=wp.float32), + damping: wp.array2d(dtype=wp.float32), + pos_limits: wp.array2d(dtype=wp.vec2f), + sim_stiffness: wp.array2d(dtype=wp.float32), + sim_damping: wp.array2d(dtype=wp.float32), + sim_pos_limits: wp.array2d(dtype=wp.vec2f), +): + """Copy the staged fixed-tendon stiffness, damping and position limits to the simulation where both masks are true. + + Args: + env_mask: Environment mask. Shape is (num_instances,). + tendon_mask: Fixed-tendon mask. Shape is (num_fixed_tendons,). + stiffness: Staged stiffness. Shape is (num_instances, num_fixed_tendons). + damping: Staged damping. Shape is (num_instances, num_fixed_tendons). + pos_limits: Staged position limits. Shape is (num_instances, num_fixed_tendons). + sim_stiffness: Simulation stiffness. Shape is (num_instances, num_fixed_tendons). + sim_damping: Simulation damping. Shape is (num_instances, num_fixed_tendons). + sim_pos_limits: Simulation position limits. Shape is (num_instances, num_fixed_tendons). + """ + i, j = wp.tid() + if env_mask[i] and tendon_mask[j]: + sim_stiffness[i, j] = stiffness[i, j] + sim_damping[i, j] = damping[i, j] + sim_pos_limits[i, j] = pos_limits[i, j] + + @wp.kernel def float_data_to_buffer_with_mask( in_data: wp.float32, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py b/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py index 86bd5d8590cb..1b4e9c84d1b8 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/cable_object/cable_object.py @@ -18,6 +18,7 @@ from isaaclab.assets.cable_object.base_cable_object import BaseCableObject from isaaclab.sim.utils.queries import has_deformable_curve_api, path_expr_to_glob, resolve_matching_prims_from_source from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from .cable_object_data import CableObjectData from .kernels import ( @@ -58,13 +59,14 @@ def root_view(self) -> ArticulationView: """Articulation view for the cable.""" return self._root_view - def reset(self, env_ids: Sequence[int] | None = None) -> None: + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: """Reset the cable object's internal buffers. Args: env_ids: Environment indices. Defaults to all instances. + env_mask: Environment mask. Defaults to all instances. Shape is (num_instances,). """ - del env_ids + del env_ids, env_mask def write_data_to_sim(self) -> None: """Write buffered commands to the simulation.""" @@ -207,15 +209,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array(dtyp dtype=wp.int32 ): """Resolve environment indices to a Warp array.""" - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, torch.Tensor): - return wp.from_torch(env_ids.to(device=self.device, dtype=torch.int32).contiguous(), dtype=wp.int32) - if isinstance(env_ids, Sequence): - return wp.array(list(env_ids), dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) def _write_segment_state_to_sim( self, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py b/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py index 06e96694147e..d2016b4186ed 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/deformable_object/deformable_object.py @@ -17,6 +17,7 @@ from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject from isaaclab.markers import VisualizationMarkers from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from .deformable_object_data import DeformableObjectData from .kernels import ( @@ -402,23 +403,11 @@ def write_nodal_kinematic_target_to_sim_mask( def _resolve_env_ids(self, env_ids): """Resolve environment indices to a warp int32 array.""" - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - elif isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - elif isinstance(env_ids, torch.Tensor): - return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) def _resolve_mask(self, mask: wp.array | torch.Tensor | None, full_mask: wp.array) -> wp.array: """Resolve an environment mask to a warp bool array.""" - if mask is None: - return full_mask - if isinstance(mask, torch.Tensor): - if mask.dtype != torch.bool: - mask = mask.to(torch.bool) - return wp.from_torch(mask, dtype=wp.bool) - return mask + return resolve_1d_mask(mask=mask, all_mask=full_mask, device=self.device) def _iter_particle_states(self): """Yield active Newton states.""" diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py index eec3479f1fdc..c9f43ef3adec 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py @@ -5,6 +5,7 @@ from __future__ import annotations +import functools from typing import TYPE_CHECKING, Any import torch @@ -960,6 +961,89 @@ def set_body_link_velocity_to_sim( body_acc_w[env_id, body_id] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) +@wp.kernel +def set_body_com_pose_to_sim_mask( + data: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + body_com_pose_w: wp.array2d(dtype=wp.transformf), + body_link_pose_w: wp.array2d(dtype=wp.transformf), +): + """Write body COM pose data, and the link poses derived from it, for the selected environments and bodies. + + Args: + data: Input array of body COM poses. Shape is (num_envs, num_bodies). + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + env_mask: Input array of environment mask. Shape is (num_envs,). + body_mask: Input array of body mask. Shape is (num_bodies,). + body_com_pose_w: Output array where body COM poses are written. Shape is (num_envs, num_bodies). + body_link_pose_w: Output array where body link poses (derived from COM) are written. + Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + body_com_pose_w[i, j] = data[i, j] + body_link_pose_w[i, j] = get_com_pose_in_link_frame_func(body_com_pose_w[i, j], body_com_pos_b[i, j]) + + +@wp.kernel +def set_body_com_velocity_to_sim_mask( + data: wp.array2d(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), +): + """Write body COM velocity data, and zero the body accelerations, for the selected environments and bodies. + + Args: + data: Input array of body COM spatial velocities. Shape is (num_envs, num_bodies). + env_mask: Input array of environment mask. Shape is (num_envs,). + body_mask: Input array of body mask. Shape is (num_bodies,). + body_com_velocity_w: Output array where body COM velocities are written. Shape is (num_envs, num_bodies). + body_acc_w: Output array where body accelerations are zeroed. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + body_com_velocity_w[i, j] = data[i, j] + body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def set_body_link_velocity_to_sim_mask( + data: wp.array2d(dtype=wp.spatial_vectorf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + body_link_pose_w: wp.array2d(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + body_link_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), +): + """Write body link velocity data, the COM velocities derived from it, and zero the body accelerations, for the + selected environments and bodies. + + Args: + data: Input array of body link spatial velocities. Shape is (num_envs, num_bodies). + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + body_link_pose_w: Input array of body link poses in world frame. Shape is (num_envs, num_bodies). + env_mask: Input array of environment mask. Shape is (num_envs,). + body_mask: Input array of body mask. Shape is (num_bodies,). + body_link_velocity_w: Output array where body link velocities are written. Shape is (num_envs, num_bodies). + body_com_velocity_w: Output array where body COM velocities (derived from link) are written. + Shape is (num_envs, num_bodies). + body_acc_w: Output array where body accelerations are zeroed. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + body_link_velocity_w[i, j] = data[i, j] + body_com_velocity_w[i, j] = get_link_velocity_in_com_frame_func( + body_link_velocity_w[i, j], body_link_pose_w[i, j], body_com_pos_b[i, j] + ) + body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + _SET_BODY_LINK_POSE_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_link_pose_to_sim, ("env_ids", "body_ids")) _SET_BODY_COM_POSE_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_com_pose_to_sim, ("env_ids", "body_ids")) _SET_BODY_COM_VELOCITY_TO_SIM_DISPATCHER = IndexKernelDispatcher(set_body_com_velocity_to_sim, ("env_ids", "body_ids")) @@ -1034,21 +1118,18 @@ def write_2d_data_to_buffer_with_indices_kernel( @wp.kernel -def write_2d_data_to_buffer_with_mask( - in_data: wp.array2d(dtype=wp.float32), +def _write_2d_data_to_buffer_with_mask( + in_data: wp.array2d(dtype=Any), env_mask: wp.array(dtype=wp.bool), joint_mask: wp.array(dtype=wp.bool), - out_data: wp.array2d(dtype=wp.float32), + out_data: wp.array2d(dtype=Any), ): - """Write 2D float data to a buffer at specified indices. - - This kernel copies float data from an input array to an output buffer at the specified - environment and joint/body indices. + """Write 2D data to a buffer where both masks are true. Args: - in_data: Input array containing float data. Shape is (num_instances, num_joints). + in_data: Input array. Shape is (num_instances, num_joints). env_mask: Input array of environment mask. Shape is (num_instances,). - joint_mask: Input array of joint/body mask. Shape is (num_instances, num_joints). + joint_mask: Input array of joint/body/tendon mask. Shape is (num_joints,). out_data: Output array where data is written. Shape is (num_instances, num_joints). """ i, j = wp.tid() @@ -1056,6 +1137,20 @@ def write_2d_data_to_buffer_with_mask( out_data[i, j] = in_data[i, j] +@functools.cache +def write_2d_data_to_buffer_with_mask_kernel(dtype: type) -> wp.Kernel: + """Select the masked 2-D buffer writer for the buffer's ``dtype``. + + Launching the generic kernel directly resolves its types on every launch (about twice the host cost), and + cannot take torch input. + """ + array_type = wp.array2d(dtype=dtype) + return wp.overload(_write_2d_data_to_buffer_with_mask, {"in_data": array_type, "out_data": array_type}) + + +write_2d_data_to_buffer_with_mask = write_2d_data_to_buffer_with_mask_kernel(wp.float32) + + @wp.func def _update_body_inertial_properties_inverse( env_id: int, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py index 3d1329cca985..8a725ca9f1ea 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py @@ -15,6 +15,7 @@ from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab_newton.sim.spawners.mpm.mpm import _SIMULATION_POINTS_SUFFIX @@ -283,24 +284,10 @@ def _create_buffers(self): self._data.default_particle_state_w = self._data.default_nodal_state_w def _resolve_env_ids(self, env_ids): - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, torch.Tensor): - return wp.from_torch(env_ids.to(device=self.device, dtype=torch.int32), dtype=wp.int32) - if isinstance(env_ids, Sequence): - return wp.array(list(env_ids), dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) def _resolve_mask(self, mask: wp.array | torch.Tensor | None) -> wp.array: - if mask is None: - return self._ALL_ENV_MASK - if isinstance(mask, torch.Tensor): - if mask.dtype != torch.bool: - mask = mask.to(torch.bool) - return wp.from_torch(mask.to(device=self.device).contiguous(), dtype=wp.bool) - return mask + return resolve_1d_mask(mask=mask, all_mask=self._ALL_ENV_MASK, device=self.device) def _iter_particle_states(self): """Yield the Newton states whose particle arrays must receive writes.""" 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 ae04842b9967..ac801b5b8802 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 @@ -21,6 +21,7 @@ from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -129,8 +130,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. @@ -1122,13 +1123,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No Returns: A warp array of environment indices or a tensor of environment indices. """ - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device) def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve body indices to a warp array or tensor. @@ -1139,13 +1134,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | Returns: A warp array of body indices or a tensor of body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) """ Internal simulation callbacks. 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 268e7a3e74d3..a44922079df5 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 @@ -22,6 +22,7 @@ from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection from isaaclab.utils import clone, validate from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -175,8 +176,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. @@ -308,16 +309,8 @@ def write_body_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES - self.write_body_link_pose_to_sim_index( - body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward + self.write_body_link_pose_to_sim_mask( + body_poses=body_poses, body_mask=body_mask, env_mask=env_mask, skip_forward=skip_forward ) def write_body_velocity_to_sim_index( @@ -386,20 +379,8 @@ def write_body_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES - self.write_body_com_velocity_to_sim_index( - body_velocities=body_velocities, - env_ids=env_ids, - body_ids=body_ids, - full_data=True, - skip_forward=skip_forward, + self.write_body_com_velocity_to_sim_mask( + body_velocities=body_velocities, body_mask=body_mask, env_mask=env_mask, skip_forward=skip_forward ) def write_body_link_pose_to_sim_index( @@ -492,16 +473,25 @@ def write_body_link_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_mask = self._ALL_ENV_MASK - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - self.write_body_link_pose_to_sim_index( - body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward + env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device) + # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies + body_mask = resolve_1d_mask( + ids=body_ids, + mask=body_mask, + all_mask=self._ALL_BODY_MASK, + scratch_mask=self._scratch_body_mask, + device=self.device, + ) + self.assert_shape_and_dtype_mask(body_poses, (env_mask, body_mask), wp.transformf, "body_poses") + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask_kernel(wp.transformf), + dim=(self.num_instances, self.num_bodies), + inputs=[body_poses, env_mask, body_mask], + outputs=[self.data.body_link_pose_w], + device=self.device, ) + if not skip_forward: + self.data._reset_pose(env_mask=env_mask) def write_body_com_pose_to_sim_index( self, @@ -597,17 +587,25 @@ def write_body_com_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_mask = self._ALL_ENV_MASK - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - # The index writer owns the invalidation (with from_link=False); forward skip_forward to it. - self.write_body_com_pose_to_sim_index( - body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward + env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device) + # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies + body_mask = resolve_1d_mask( + ids=body_ids, + mask=body_mask, + all_mask=self._ALL_BODY_MASK, + scratch_mask=self._scratch_body_mask, + device=self.device, + ) + self.assert_shape_and_dtype_mask(body_poses, (env_mask, body_mask), wp.transformf, "body_poses") + wp.launch( + shared_kernels.set_body_com_pose_to_sim_mask, + dim=(self.num_instances, self.num_bodies), + inputs=[body_poses, self.data.body_com_pos_b, env_mask, body_mask], + outputs=[self.data.body_com_pose_w, self.data.body_link_pose_w], + device=self.device, ) + if not skip_forward: + self.data._reset_pose(env_mask=env_mask, from_link=False) def write_body_com_velocity_to_sim_index( self, @@ -711,21 +709,25 @@ def write_body_com_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_mask = self._ALL_ENV_MASK - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - # The index writer owns the invalidation; forward skip_forward to it. - self.write_body_com_velocity_to_sim_index( - body_velocities=body_velocities, - env_ids=env_ids, - body_ids=body_ids, - full_data=True, - skip_forward=skip_forward, + env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device) + # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies + body_mask = resolve_1d_mask( + ids=body_ids, + mask=body_mask, + all_mask=self._ALL_BODY_MASK, + scratch_mask=self._scratch_body_mask, + device=self.device, ) + self.assert_shape_and_dtype_mask(body_velocities, (env_mask, body_mask), wp.spatial_vectorf, "body_velocities") + wp.launch( + shared_kernels.set_body_com_velocity_to_sim_mask, + dim=(self.num_instances, self.num_bodies), + inputs=[body_velocities, env_mask, body_mask], + outputs=[self.data.body_com_vel_w, self.data.body_com_acc_w], + device=self.device, + ) + if not skip_forward: + self.data._reset_velocity(env_mask=env_mask) def write_body_link_velocity_to_sim_index( self, @@ -831,21 +833,25 @@ def write_body_link_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_mask = self._ALL_ENV_MASK - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - # The index writer owns the invalidation (with from_com=False); forward skip_forward to it. - self.write_body_link_velocity_to_sim_index( - body_velocities=body_velocities, - env_ids=env_ids, - body_ids=body_ids, - full_data=True, - skip_forward=skip_forward, + env_mask = resolve_1d_mask(mask=env_mask, all_mask=self._ALL_ENV_MASK, device=self.device) + # a body selection given as ids becomes a body mask on the device, so the same mask kernel applies + body_mask = resolve_1d_mask( + ids=body_ids, + mask=body_mask, + all_mask=self._ALL_BODY_MASK, + scratch_mask=self._scratch_body_mask, + device=self.device, + ) + self.assert_shape_and_dtype_mask(body_velocities, (env_mask, body_mask), wp.spatial_vectorf, "body_velocities") + wp.launch( + shared_kernels.set_body_link_velocity_to_sim_mask, + dim=(self.num_instances, self.num_bodies), + inputs=[body_velocities, self.data.body_com_pos_b, self.data.body_link_pose_w, env_mask, body_mask], + outputs=[self.data.body_link_vel_w, self.data.body_com_vel_w, self.data.body_com_acc_w], + device=self.device, ) + if not skip_forward: + self.data._reset_velocity(env_mask=env_mask, from_com=False) """ Operations - Setters. @@ -1244,6 +1250,7 @@ def _create_buffers(self): ) self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device) self._ALL_BODY_MASK = wp.ones((self.num_bodies,), dtype=wp.bool, device=self.device) + self._scratch_body_mask = wp.zeros((self.num_bodies,), dtype=wp.bool, device=self.device) # external wrench composer self._instantaneous_wrench_composer = WrenchComposer(self) @@ -1285,13 +1292,7 @@ def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor: Returns: Environment indices. """ - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_ENV_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self.device) def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: """Resolve body indices. @@ -1302,37 +1303,7 @@ def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: Returns: Body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - if isinstance(body_ids, slice): - return wp.from_torch( - torch.arange(self.num_bodies, dtype=torch.int32, device=self.device)[body_ids], dtype=wp.int32 - ) - return body_ids - - 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: - if isinstance(env_mask, wp.array): - env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0] - else: - env_ids = self._ALL_ENV_INDICES - return env_ids - - def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array | torch.Tensor: - """Resolve body mask to indices via torch.nonzero.""" - if body_mask is not None: - if isinstance(body_mask, wp.array): - body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0] - else: - body_ids = self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self.device) """ Internal simulation callbacks. diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index 8f234e202fff..dbf4f1c23c51 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -16,11 +16,11 @@ world for its own duration. Configurations that fail initialization, nested-root IK, the rebind scenario, which swaps the live Newton -state, and the CUDA-graph check build their own small scenes. Only one simulation context can be alive, so +state, and the CUDA-graph checks build their own small scenes. Only one simulation context can be alive, so these tests come first and fail if selected after a composite-scene test. -The composite scene runs on the CPU without CUDA-graph capture. On CUDA, one own-scene test covers the -ordered-state republish recorded into a captured graph. +The composite scene runs on the CPU without CUDA-graph capture. On CUDA, own-scene tests cover the ordered-state +republish and the fixed-tendon mask writes recorded into a captured graph. """ from isaaclab_newton.physics import NewtonCfg, NewtonMJWarpManager @@ -707,6 +707,56 @@ def test_newton_ordered_state_publishes_inside_captured_cuda_graph(device: str) ) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_fixed_tendon_mask_writes_replay_from_cuda_graph(device: str) -> None: + """Fixed-tendon mask setters and the mask writer recorded into a CUDA graph act on every replay. + + The mask variants are the documented graph-capturable path. The inputs are filled only after capture, so the + replay must read them at launch time and write only the selected cells. + """ + with build_simulation_context(sim_cfg=newton_sim_cfg(device, use_newton_actuators=False)) as sim: + articulation = spawn_assets({"tendon": _island_cfgs()["tendon"]})["tendon"] + sim.reset() + num_envs, num_tendons = articulation.num_instances, articulation.num_fixed_tendons + env_mask = wp.array([False, True], dtype=wp.bool, device=device) + tendon_mask = wp.array([False, True], dtype=wp.bool, device=device) + selected = torch.zeros(num_envs, num_tendons, 1, dtype=torch.bool, device=device) + selected[1, 1] = True + stiffness = torch.zeros(num_envs, num_tendons, device=device) + damping = torch.zeros_like(stiffness) + limits = torch.zeros(num_envs, num_tendons, 2, device=device) + data = articulation.data + before = { + "stiffness": data.fixed_tendon_stiffness.torch.clone(), + "damping": data.fixed_tendon_damping.torch.clone(), + "pos_limits": data.fixed_tendon_pos_limits.torch.clone(), + } + with wp.ScopedCapture(device) as capture: + articulation.set_fixed_tendon_stiffness_mask( + stiffness=stiffness, fixed_tendon_mask=tendon_mask, env_mask=env_mask + ) + articulation.set_fixed_tendon_damping_mask( + damping=damping, fixed_tendon_mask=tendon_mask, env_mask=env_mask + ) + articulation.set_fixed_tendon_position_limit_mask( + limit=limits, fixed_tendon_mask=tendon_mask, env_mask=env_mask + ) + articulation.write_fixed_tendon_properties_to_sim_mask(fixed_tendon_mask=tendon_mask, env_mask=env_mask) + stiffness.fill_(12.0) + damping.fill_(3.0) + limits[..., 0] = -0.1 + limits[..., 1] = 0.2 + wp.capture_launch(capture.graph) + + written = {"stiffness": stiffness, "damping": damping, "pos_limits": limits} + for name, value in written.items(): + after = getattr(data, f"fixed_tendon_{name}").torch + expected = torch.where( + selected, value.view(*selected.shape[:2], -1), before[name].view(*selected.shape[:2], -1) + ) + torch.testing.assert_close(after.view_as(expected), expected, msg=name) + + ## # Composite scene. ## @@ -1254,9 +1304,9 @@ def one_step_response() -> tuple[torch.Tensor, torch.Tensor]: actuator_reset = actuator.reset reset_env_ids = [] - def record_actuator_reset(env_ids: torch.Tensor | None = None) -> None: + def record_actuator_reset(env_ids: torch.Tensor | None = None, env_mask: torch.Tensor | None = None) -> None: reset_env_ids.append(env_ids) - actuator_reset(env_ids) + actuator_reset(env_ids, env_mask=env_mask) monkeypatch.setattr(actuator, "reset", record_actuator_reset) articulation.reset() diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index b96c41ed3aa4..94d592c03083 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py @@ -383,35 +383,96 @@ def _run( return _Run(sim=sim, articulations=articulations, results=results) -def _record_lab_state_reset(run: _Run) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: - """Reset one Lab-path environment of the delayed island and record each group's effort and PD demand. +_RESET_SELECTORS = ("env_ids", "env_mask") +"""Partial-reset selectors: an index tensor and a Warp mask.""" + + +def _reset_env_0(articulation: Articulation, selector: str) -> None: + """Reset environment 0 alone through the given selector.""" + if selector == "env_ids": + articulation.reset(env_ids=torch.tensor([0], device=articulation.device, dtype=torch.long)) + else: + env_mask = torch.arange(articulation.num_instances, device=articulation.device) == 0 + articulation.reset(env_mask=wp.from_torch(env_mask)) + + +def _record_lab_state_reset(run: _Run) -> dict[str, dict[str, tuple[torch.Tensor, torch.Tensor]]]: + """Reset one Lab-path environment of the delayed island per selector and record each group's effort and PD demand. Reset environments must accept fresh commands while the remaining environments retain their delay history. """ articulation = run.articulations["delayed"] commands = articulation.actuators.target_command old_target = commands.position.torch.clone() - # Episode reset samples the configured lag; construction alone leaves actuator lag at zero. - articulation.reset() - commands.set_position_index(value=old_target) - articulation.write_data_to_sim() new_target = old_target + 0.02 - articulation.reset(env_ids=torch.tensor([0], device=articulation.device)) - commands.set_position_index(value=new_target) - articulation.write_data_to_sim() expected_target = old_target.clone() expected_target[0] = new_target[0] recorded = {} - for name, actuator in articulation.actuators.items(): - joints = actuator.joint_indices - demand = ( - actuator.stiffness * (expected_target[:, joints] - articulation.data.joint_pos.torch[:, joints]) - - actuator.damping * articulation.data.joint_vel.torch[:, joints] - ) - recorded[name] = (actuator.computed_effort.clone(), demand) + for selector in _RESET_SELECTORS: + # Episode reset samples the configured lag; construction alone leaves actuator lag at zero. + articulation.reset() + commands.set_position_index(value=old_target) + articulation.write_data_to_sim() + _reset_env_0(articulation, selector) + commands.set_position_index(value=new_target) + articulation.write_data_to_sim() + recorded[selector] = {} + for name, actuator in articulation.actuators.items(): + joints = actuator.joint_indices + demand = ( + actuator.stiffness * (expected_target[:, joints] - articulation.data.joint_pos.torch[:, joints]) + - actuator.damping * articulation.data.joint_vel.torch[:, joints] + ) + recorded[selector][name] = (actuator.computed_effort.clone(), demand) return recorded +def _warm_delayed_island(run: _Run) -> list: + """Step the delayed island until its delay buffers hold history and return its delayed actuators' state pairs. + + Only the reset articulation's own delayed actuators are returned. The adapter is model-wide, so the other + islands' actuators share its state buffers; their state is not part of this articulation's contract. + """ + articulation = run.articulations["delayed"] + for _ in range(3): + articulation.write_data_to_sim() + run.sim.step() + articulation.update(DT) + adapter = run.sim.backend_manager._adapter + assert adapter is not None + own_actuators = [] + for group_name in articulation.actuators._native_group_names: + group_actuators = articulation.actuators[group_name] + own_actuators.extend(group_actuators if isinstance(group_actuators, tuple) else (group_actuators,)) + stateful_pairs = [ + (act, st) + for act, st in zip(adapter.actuators, adapter._states_a) + if any(act is own for own in own_actuators) and st is not None and st.delay_state is not None + ] + assert len(stateful_pairs) > 0, "expected at least one DelayedPD actuator with delay_state" + for act, state in stateful_pairs: + pushes_before = state.delay_state.num_pushes.numpy() + assert (pushes_before > 0).all(), "expected non-zero num_pushes for all DOFs after warmup" + return stateful_pairs + + +def _assert_delay_state_cleared_for_env_0_only(run: _Run, stateful_pairs: list) -> None: + """Assert ``num_pushes`` is zero for environment 0's DOFs and positive for every other environment's DOFs.""" + # Map each entry of ``act.indices`` to its env via the adapter's per-env DOF count. The adapter is + # model-wide (includes free-joint DOFs on floating-base articulations), so ``adapter.num_joints`` is the + # stride. + adapter = run.sim.backend_manager._adapter + for act, state in stateful_pairs: + pushes_after = state.delay_state.num_pushes.numpy() + indices_np = act.indices.numpy() + for i, global_dof in enumerate(indices_np): + env = int(global_dof) // adapter.num_joints + if env == 0: + assert int(pushes_after[i]) == 0, f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}" + else: + assert int(pushes_after[i]) > 0, f"DOF {i} (env {env}) was NOT in the reset selection but is 0" + + @pytest.fixture(scope="module", params=test_devices(DeviceScope.CUDA)) def device(request: pytest.FixtureRequest) -> str: """Simulation device of the actuator scenes.""" @@ -756,6 +817,31 @@ def test_sync_torque_telemetry_keeps_user_order_effort_buffers_unmapped() -> Non np.testing.assert_allclose(applied.numpy(), np.asarray([[100.0, 200.0, 300.0]], dtype=np.float32)) +# --------------------------------------------------------------------------- +# Own scene; runs before the module-scoped scenes exist, since only one simulation context can be alive +# --------------------------------------------------------------------------- + + +def test_newton_masked_reset_replays_from_cuda_graph_without_allocation(device: str) -> None: + """A masked reset records into a CUDA graph without allocating and, on replay, clears only the masked env. + + The scene holds only the delayed island: the adapter resets every actuator of the model, and Newton's neural + actuator state reset cannot be recorded into a graph. + """ + sim_cfg = newton_sim_cfg(device, dt=DT, use_newton_actuators=True, solver_cfg=_solver_cfg()) + with build_simulation_context(sim_cfg=sim_cfg) as sim: + run = _run(sim, {"delayed": _islands()["delayed"]}, use_newton_actuators=True) + articulation = run.articulations["delayed"] + env_mask = wp.from_torch(torch.arange(articulation.num_instances, device=device) == 0) + # Allocating inside the capture raises while the memory pool is disabled. + with wp.ScopedMempool(device, False), wp.ScopedCapture(device) as capture: + articulation.reset(env_mask=env_mask) + # Repopulate the delay buffers, so the replay rather than the capture must clear environment 0. + stateful_pairs = _warm_delayed_island(run) + wp.capture_launch(capture.graph) + _assert_delay_state_cleared_for_env_0_only(run, stateful_pairs) + + # --------------------------------------------------------------------------- # Equivalence tests with different actuator types # --------------------------------------------------------------------------- @@ -943,47 +1029,16 @@ def gains(name: str) -> torch.Tensor: # --------------------------------------------------------------------------- -def test_newton_state_reset_isolated_to_reset_env(newton_run: _Run) -> None: - """Newton: ``num_pushes`` zeroes for env 0's DOFs only after reset of [0]. - - Only the reset articulation's own delayed actuators are checked. The adapter is model-wide, so the other - islands' actuators share its state buffers; their state is not part of this articulation's contract. - """ - articulation = newton_run.articulations["delayed"] - adapter = newton_run.sim.backend_manager._adapter - assert adapter is not None - own_actuators = [] - for group_name in articulation.actuators._native_group_names: - group_actuators = articulation.actuators[group_name] - own_actuators.extend(group_actuators if isinstance(group_actuators, tuple) else (group_actuators,)) - stateful_pairs = [ - (act, st) - for act, st in zip(adapter.actuators, adapter._states_a) - if any(act is own for own in own_actuators) and st is not None and st.delay_state is not None - ] - assert len(stateful_pairs) > 0, "expected at least one DelayedPD actuator with delay_state" - - for act, state in stateful_pairs: - pushes_before = state.delay_state.num_pushes.numpy() - assert (pushes_before > 0).all(), "expected non-zero num_pushes for all DOFs after warmup" - - articulation.reset(env_ids=torch.tensor([0], device=articulation.device, dtype=torch.long)) - - # Map each entry of ``act.indices`` to its env via the adapter's per-env DOF count. The adapter is - # model-wide (includes free-joint DOFs on floating-base articulations), so ``adapter.num_joints`` is the - # stride. - for act, state in stateful_pairs: - pushes_after = state.delay_state.num_pushes.numpy() - indices_np = act.indices.numpy() - for i, global_dof in enumerate(indices_np): - env = int(global_dof) // adapter.num_joints - if env == 0: - assert int(pushes_after[i]) == 0, f"DOF {i} (env {env}) should be reset to 0, got {pushes_after[i]}" - else: - assert int(pushes_after[i]) > 0, f"DOF {i} (env {env}) was NOT in reset env_ids but num_pushes is 0" +@pytest.mark.parametrize("selector", _RESET_SELECTORS) +def test_newton_state_reset_isolated_to_reset_env(newton_run: _Run, selector: str) -> None: + """Newton: ``num_pushes`` zeroes for env 0's DOFs only after a partial reset of env 0.""" + stateful_pairs = _warm_delayed_island(newton_run) + _reset_env_0(newton_run.articulations["delayed"], selector) + _assert_delay_state_cleared_for_env_0_only(newton_run, stateful_pairs) -def test_lab_state_reset_isolated_to_reset_env(lab_run: dict) -> None: +@pytest.mark.parametrize("selector", _RESET_SELECTORS) +def test_lab_state_reset_isolated_to_reset_env(lab_run: dict, selector: str) -> None: """Reset environments accept fresh commands while the remaining environments retain their history.""" - for computed_effort, demand in lab_run["state_reset"].values(): + for computed_effort, demand in lab_run["state_reset"][selector].values(): torch.testing.assert_close(computed_effort, demand) diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index adaa18fe03fd..4a9299774b02 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -11,9 +11,9 @@ Every environment holds a rigid-object cube and a three-cube collection, all locally spawned shapes. The cube also stands beside the collection as an unrelated rigid body, so the collection view must select exactly its configured bodies. Scene gravity is off; a test that needs gravity applies it to every world for its own duration. -Configurations that fail initialization, VBD state publication, and the single-instance cases build their own -scenes. Only one simulation context can be alive, so those tests come first and fail if selected after a -shared-scene test. +Configurations that fail initialization, VBD state publication, the single-instance cases, and the CUDA-graph +check of the collection mask writers build their own scenes. Only one simulation context can be alive, so those +tests come first and fail if selected after a shared-scene test. """ from isaaclab_newton.physics import NewtonCfg @@ -247,6 +247,46 @@ def test_collection_single_instance_initialization(device: str) -> None: torch.testing.assert_close(object_collection.data.body_link_pose_w.torch, pose) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) +def test_collection_mask_writes_replay_from_cuda_graph_on_selected_cells(device: str) -> None: + """Mask writers record into a CUDA graph and, on replay, write only the selected environment-body cells. + + The mask variants are the documented graph-capturable write path. The replay reads the input buffer's + contents at launch time, so the input is filled only after capture. + """ + with build_simulation_context(sim_cfg=newton_sim_cfg(device)) as sim: + collection = spawn_assets({"collection": _collection_cfg(2)})["collection"] + sim.reset() + sim.step() + collection.update(sim.cfg.dt) + + env_mask = wp.array([False, True], dtype=wp.bool, device=device) + body_mask = wp.array([True, False], dtype=wp.bool, device=device) + selected = torch.zeros(collection.num_instances, collection.num_bodies, 1, dtype=torch.bool, device=device) + selected[1, 0] = True + identity_pose = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], device=device) + writers = ( + ("write_body_link_pose_to_sim_mask", "body_poses", "body_link_pose_w"), + ("write_body_com_pose_to_sim_mask", "body_poses", "body_com_pose_w"), + ("write_body_com_velocity_to_sim_mask", "body_velocities", "body_com_vel_w"), + ("write_body_link_velocity_to_sim_mask", "body_velocities", "body_link_vel_w"), + ) + for writer, data_arg, buffer_name in writers: + before = getattr(collection.data, buffer_name).torch.clone() + target = torch.zeros_like(before) + with wp.ScopedCapture(device) as capture: + getattr(collection, writer)(**{data_arg: target}, env_mask=env_mask, body_mask=body_mask) + if data_arg == "body_poses": + target.copy_(identity_pose.expand_as(target)) + target[..., :3] = torch.rand_like(target[..., :3]) + 2.0 + else: + target.copy_(torch.rand_like(target) + 1.0) + wp.capture_launch(capture.graph) + + written = getattr(collection.data, buffer_name).torch + torch.testing.assert_close(written, torch.where(selected, target, before), msg=writer) + + ## # Shared scene. ## @@ -342,7 +382,8 @@ def test_external_force_on_single_body(scene: _Scene) -> None: """A permanent wrench reaches the solver in the local and the global frame, and a reset clears it. Cube 0 receives a force equal to its weight and hovers while cube 1 falls. Cube 0 then receives a force set - and added 1 m off its center along y and must turn about x while cube 1 keeps falling. + and added 1 m off its center along y and must turn about x while cube 1 keeps falling. A masked reset clears + the wrench of the selected environment only. """ cube = scene.cube composer = cube.permanent_wrench_composer @@ -371,6 +412,12 @@ def test_external_force_on_single_body(scene: _Scene) -> None: assert cube.data.root_ang_vel_b.torch[0, 0].abs() > 0.1 assert cube.data.root_pos_w.torch[1, 2] < 1.0 + ones = torch.ones_like(zeros) + composer.set_forces_and_torques_index(forces=ones, torques=ones) + cube.reset(env_mask=wp.from_torch(torch.arange(cube.num_instances, device=scene.device) == 0)) + assert torch.count_nonzero(composer.out_force_b.torch[0]) == 0 + assert torch.count_nonzero(composer.out_force_b.torch[1:]) == composer.out_force_b.torch[1:].numel() + @pytest.mark.isaacsim_ci def test_rigid_body_set_material_properties(scene: _Scene) -> None: @@ -861,7 +908,8 @@ def test_collection_external_force_on_single_body(scene: _Scene) -> None: """A permanent wrench reaches the solver in the local and the global frame, and a reset clears it. Every other cube receives a force equal to its weight and hovers while the rest fall. Those cubes then - receive a force set and added 1 m off their centers along y and must turn about x. + receive a force set and added 1 m off their centers along y and must turn about x. A masked reset clears the + wrenches of the selected environment only. """ collection = scene.collection composer = collection.permanent_wrench_composer @@ -890,6 +938,12 @@ def test_collection_external_force_on_single_body(scene: _Scene) -> None: assert torch.all(collection.data.body_com_ang_vel_b.torch[:, 0::2, 0] > 0.1) assert torch.all(height[:, 1::2] < 1.0) + ones = torch.ones_like(zeros) + composer.set_forces_and_torques_index(forces=ones, torques=ones) + collection.reset(env_mask=wp.from_torch(torch.arange(collection.num_instances, device=scene.device) == 0)) + assert torch.count_nonzero(composer.out_force_b.torch[0]) == 0 + assert torch.count_nonzero(composer.out_force_b.torch[1:]) == composer.out_force_b.torch[1:].numel() + @pytest.mark.isaacsim_ci def test_collection_gravity_vec_w_tracks_model_gravity(scene: _Scene) -> None: diff --git a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py index 4216a21eb42d..8cf45b728922 100644 --- a/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py +++ b/source/isaaclab_newton/test/sensors/test_newton_raycast_sensor.py @@ -40,6 +40,7 @@ from isaaclab.sim import SimulationCfg from isaaclab.terrains import TerrainImporterCfg from isaaclab.utils import configclass +from isaaclab.utils.seed import WarpRng SENSOR_HEIGHT = 2.0 RAY_OFFSET = 0.2 @@ -211,6 +212,52 @@ def test_legacy_multi_mesh_tracks_ad_hoc_regex_target(sim): torch.testing.assert_close(sensor.data.pos_w.torch, pose[:, :3]) +@pytest.mark.parametrize("sim", [pytest.param(False, id="eager")], indirect=True) +def test_partial_reset_resamples_drift_of_selected_envs_only(sim): + """Mask and slice resets resample the drift of the selected environments only, also on graph replay.""" + scene_cfg = RaycastTestSceneCfg(num_envs=3) + scene_cfg.raycast.drift_range = (0.01, 0.05) + scene_cfg.raycast.ray_cast_drift_range = {"z": (0.1, 0.2)} + scene = InteractiveScene(scene_cfg) + sim.reset() + sensor: NewtonRaycastSensor = scene["raycast"] + sensor.reset() + + # A masked reset draws from the process-wide Warp RNG state, which an environment initializes. + WarpRng.initialize(scene.num_envs, sensor.device) + + drift = sensor.drift.torch.clone() + ray_cast_drift = sensor.ray_cast_drift.torch.clone() + sensor.reset(env_mask=wp.array([False, True, False], dtype=wp.bool, device=sensor.device)) + torch.testing.assert_close(sensor.drift.torch[[0, 2]], drift[[0, 2]]) + torch.testing.assert_close(sensor.ray_cast_drift.torch[[0, 2]], ray_cast_drift[[0, 2]]) + assert not torch.allclose(sensor.drift.torch[1], drift[1]) + assert not torch.allclose(sensor.ray_cast_drift.torch[1, 2], ray_cast_drift[1, 2]) + assert ((sensor.drift.torch >= 0.01) & (sensor.drift.torch <= 0.05)).all() + assert ((sensor.ray_cast_drift.torch[:, 2] >= 0.1) & (sensor.ray_cast_drift.torch[:, 2] <= 0.2)).all() + + drift = sensor.drift.torch.clone() + sensor.reset(env_ids=slice(2, None)) + torch.testing.assert_close(sensor.drift.torch[:2], drift[:2]) + assert not torch.allclose(sensor.drift.torch[2], drift[2]) + + # The masked reset records into a CUDA graph without allocating (the disabled memory pool makes an + # allocation raise), and every replay draws new drift for the masked environment only. + env_mask = wp.array([True, False, False], dtype=wp.bool, device=sensor.device) + with wp.ScopedMempool(sensor.device, False), wp.ScopedCapture(sensor.device) as capture: + sensor.reset(env_mask=env_mask) + for _ in range(2): + drift = sensor.drift.torch.clone() + ray_cast_drift = sensor.ray_cast_drift.torch.clone() + wp.capture_launch(capture.graph) + torch.testing.assert_close(sensor.drift.torch[1:], drift[1:]) + torch.testing.assert_close(sensor.ray_cast_drift.torch[1:], ray_cast_drift[1:]) + assert not torch.allclose(sensor.drift.torch[0], drift[0]) + assert not torch.allclose(sensor.ray_cast_drift.torch[0, 2], ray_cast_drift[0, 2]) + assert ((sensor.drift.torch >= 0.01) & (sensor.drift.torch <= 0.05)).all() + assert ((sensor.ray_cast_drift.torch[:, 2] >= 0.1) & (sensor.ray_cast_drift.torch[:, 2] <= 0.2)).all() + + @configclass class MultiMeshTargetSceneCfg(RaycastTestSceneCfg): """Read the same box through a static target and a tracked target.""" diff --git a/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..3f1ff2245d66 --- /dev/null +++ b/source/isaaclab_ov/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,7 @@ +* Changed the env selection helpers of the OVPhysX assets to resolve ``env_ids`` through + :func:`~isaaclab.utils.warp.utils.resolve_ids` and ``env_mask`` through + :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`. Any integer sequence and slices with a negative step are accepted, + and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected with a :class:`TypeError`. +* Changed the body, joint and tendon selection helpers of the OVPhysX assets to resolve ids and masks through + :func:`~isaaclab.utils.warp.utils.resolve_ids` and :func:`~isaaclab.utils.warp.utils.resolve_1d_mask`, as the env + selection does, so slices other than ``slice(None)`` are accepted. diff --git a/source/isaaclab_ov/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_ov/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst new file mode 100644 index 000000000000..776f5fc75723 --- /dev/null +++ b/source/isaaclab_ov/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst @@ -0,0 +1,2 @@ +* Fixed :meth:`~isaaclab_ov.assets.Articulation.reset` with only ``env_mask`` resetting the actuator state of every + environment. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py index 9676c1a562da..15ddafab9200 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/actuator_control.py @@ -98,9 +98,9 @@ def compute_native_actuators(self, collection: ActuatorCollection, dt: float) -> self._actuator_runtime.compute(collection, dt) return True - def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None: + def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None: if self._native_actuator_path_active and self._actuator_runtime is not None: - self._actuator_runtime.reset(env_ids) + self._actuator_runtime.reset(env_ids, env_mask) def stage_user_command( self, diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py index 07bc71d3f611..ef11c93d5d0f 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py @@ -33,6 +33,7 @@ from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray from isaaclab.utils.warp import kernels as warp_kernels +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ov import tensor_types as TT @@ -220,7 +221,7 @@ def reset( env_ids = slice(None) # reset actuators, including backend-native actuator state. None selects all # environments; delayed-actuator buffers do not accept a slice. - self.actuators.reset(None if env_ids == slice(None) else env_ids) + self.actuators.reset(None if env_ids == slice(None) else 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) @@ -4373,79 +4374,27 @@ def _log_articulation_info(self) -> None: def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve environment indices on ``self._device``.""" - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self._device) - if isinstance(env_ids, torch.Tensor): - return env_ids.to(device=self._device) - if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device: - env_ids = wp.clone(env_ids, device=self._device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device) def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve body indices to a Warp signed-integer array on ``self._device``.""" - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if body_ids is None or body_ids == slice(None): - return self._ALL_BODY_INDICES - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self._device) - if isinstance(body_ids, torch.Tensor): - return body_ids.to(device=self._device) - if isinstance(body_ids, wp.array) and str(body_ids.device) != self._device: - body_ids = wp.clone(body_ids, device=self._device) - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: """Resolve joint indices to a Warp signed-integer array on ``self._device``.""" - if isinstance(joint_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if joint_ids is None or joint_ids == slice(None): - return self._ALL_JOINT_INDICES - if isinstance(joint_ids, list): - return wp.array(joint_ids, dtype=wp.int32, device=self._device) - if isinstance(joint_ids, torch.Tensor): - return joint_ids.to(device=self._device) - if isinstance(joint_ids, wp.array) and str(joint_ids.device) != self._device: - joint_ids = wp.clone(joint_ids, device=self._device) - return joint_ids + return resolve_ids(ids=joint_ids, all_ids=self._ALL_JOINT_INDICES, device=self._device) def _resolve_fixed_tendon_ids( self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None ) -> wp.array | torch.Tensor: """Resolve fixed-tendon indices to a Warp signed-integer array on ``self._device``.""" - if isinstance(tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if tendon_ids is None or tendon_ids == slice(None): - return self._ALL_FIXED_TENDON_INDICES - if isinstance(tendon_ids, list): - return wp.array(tendon_ids, dtype=wp.int32, device=self._device) - if isinstance(tendon_ids, torch.Tensor): - return tendon_ids.to(device=self._device) - if isinstance(tendon_ids, wp.array) and str(tendon_ids.device) != self._device: - tendon_ids = wp.clone(tendon_ids, device=self._device) - return tendon_ids + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self._device) def _resolve_spatial_tendon_ids( self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None ) -> wp.array | torch.Tensor: """Resolve spatial-tendon indices to a Warp signed-integer array on ``self._device``.""" - if isinstance(tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if tendon_ids is None or tendon_ids == slice(None): - return self._ALL_SPATIAL_TENDON_INDICES - if isinstance(tendon_ids, list): - return wp.array(tendon_ids, dtype=wp.int32, device=self._device) - if isinstance(tendon_ids, torch.Tensor): - return tendon_ids.to(device=self._device) - if isinstance(tendon_ids, wp.array) and str(tendon_ids.device) != self._device: - tendon_ids = wp.clone(tendon_ids, device=self._device) - return tendon_ids + return resolve_ids(ids=tendon_ids, all_ids=self._ALL_SPATIAL_TENDON_INDICES, device=self._device) def _broadcast_scalar_to_2d( self, value: float | torch.Tensor | wp.array, shape: tuple[int, int] @@ -4483,13 +4432,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: no ``torch.nonzero`` conversion is needed. ``None`` returns the pre-allocated all-true mask. """ - if env_mask is None: - return self._ALL_TRUE_ENV_MASK - if isinstance(env_mask, torch.Tensor): - return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool) - if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device: - env_mask = wp.clone(env_mask, device=self._device) - return env_mask + return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device) def _sim_env_ids_view(self, count: int) -> wp.array: """Return a cached prefix view of the simulator-index scratch buffer.""" @@ -4505,43 +4448,19 @@ def _sim_env_ids_view(self, count: int) -> wp.array: def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: """Resolve a body mask to a ``wp.bool`` array on ``self._device`` (Newton-style).""" - if body_mask is None: - return self._ALL_TRUE_BODY_MASK - if isinstance(body_mask, torch.Tensor): - return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool) - if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device: - body_mask = wp.clone(body_mask, device=self._device) - return body_mask + return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device) def _resolve_joint_mask(self, joint_mask: wp.array | None) -> wp.array: """Resolve a joint mask to a ``wp.bool`` array on ``self._device``.""" - if joint_mask is None: - return self._ALL_TRUE_JOINT_MASK - if isinstance(joint_mask, torch.Tensor): - return wp.from_torch(joint_mask.to(torch.bool), dtype=wp.bool) - if isinstance(joint_mask, wp.array) and str(joint_mask.device) != self._device: - joint_mask = wp.clone(joint_mask, device=self._device) - return joint_mask + return resolve_1d_mask(mask=joint_mask, all_mask=self._ALL_TRUE_JOINT_MASK, device=self._device) def _resolve_fixed_tendon_mask(self, tendon_mask: wp.array | None) -> wp.array: """Resolve a fixed-tendon mask to a ``wp.bool`` array on ``self._device``.""" - if tendon_mask is None: - return self._ALL_TRUE_FIXED_TENDON_MASK - if isinstance(tendon_mask, torch.Tensor): - return wp.from_torch(tendon_mask.to(torch.bool), dtype=wp.bool) - if isinstance(tendon_mask, wp.array) and str(tendon_mask.device) != self._device: - tendon_mask = wp.clone(tendon_mask, device=self._device) - return tendon_mask + return resolve_1d_mask(mask=tendon_mask, all_mask=self._ALL_TRUE_FIXED_TENDON_MASK, device=self._device) def _resolve_spatial_tendon_mask(self, tendon_mask: wp.array | None) -> wp.array: """Resolve a spatial-tendon mask to a ``wp.bool`` array on ``self._device``.""" - if tendon_mask is None: - return self._ALL_TRUE_SPATIAL_TENDON_MASK - if isinstance(tendon_mask, torch.Tensor): - return wp.from_torch(tendon_mask.to(torch.bool), dtype=wp.bool) - if isinstance(tendon_mask, wp.array) and str(tendon_mask.device) != self._device: - tendon_mask = wp.clone(tendon_mask, device=self._device) - return tendon_mask + return resolve_1d_mask(mask=tendon_mask, all_mask=self._ALL_TRUE_SPATIAL_TENDON_MASK, device=self._device) def _get_cpu_env_mask(self, env_mask: wp.array) -> wp.array: """Return a pinned-host CPU copy of :paramref:`env_mask` for a CPU-only binding write. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py b/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py index ccca84b3cd0a..ceef9f0e2049 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/deformable_object/deformable_object.py @@ -22,6 +22,7 @@ from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject from isaaclab.markers import VisualizationMarkers from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab_ov import tensor_types as TT from isaaclab_ov.sim.views.ovphysx_view import OvPhysxView @@ -528,22 +529,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array(dtyp dtype=wp.int32 ): """Resolve environment indices to a device-resident int32 Warp array.""" - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - # OVPhysX consumes these IDs directly and requires contiguous DLPack tensors. - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids].contiguous()) - if isinstance(env_ids, torch.Tensor): - values = env_ids.to(device=self.device, dtype=torch.int32).contiguous() - return wp.from_torch(values, dtype=wp.int32) - if isinstance(env_ids, wp.array): - if env_ids.dtype != wp.int32: - values = wp.to_torch(env_ids).to(device=self.device, dtype=torch.int32).contiguous() - return wp.from_torch(values, dtype=wp.int32) - if str(env_ids.device) != self.device: - return wp.clone(env_ids, device=self.device) - return env_ids - return wp.array(list(env_ids), dtype=wp.int32, device=self.device) + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) @staticmethod def _arrays_overlap(first: wp.array, second: wp.array) -> bool: diff --git a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py index 7e19bc75109d..a44665429701 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object.py @@ -22,6 +22,7 @@ from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ov import tensor_types as TT @@ -1076,17 +1077,7 @@ def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | No Returns: Environment indices on ``self._device``. """ - if env_ids is None or env_ids == slice(None): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self._device) - if isinstance(env_ids, torch.Tensor): - return env_ids.to(device=self._device) - if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device: - env_ids = wp.clone(env_ids, device=self._device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self._device) def _sim_env_ids_view(self, count: int) -> wp.array: """Return a cached prefix of the simulator-index scratch buffer.""" @@ -1109,13 +1100,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | Returns: Body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if body_ids is None or body_ids == slice(None): - return self._ALL_BODY_INDICES - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self._device) - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: """Resolve an environment mask to a ``wp.bool`` array. @@ -1126,13 +1111,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: Returns: A warp ``wp.bool`` array on ``self._device``. """ - if env_mask is None: - return self._ALL_TRUE_ENV_MASK - if isinstance(env_mask, torch.Tensor): - return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool) - if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device: - env_mask = wp.clone(env_mask, device=self._device) - return env_mask + return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device) def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: """Resolve a body mask to a ``wp.bool`` array. @@ -1143,13 +1122,7 @@ def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: Returns: A warp ``wp.bool`` array on ``self._device``. """ - if body_mask is None: - return self._ALL_TRUE_BODY_MASK - if isinstance(body_mask, torch.Tensor): - return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool) - if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device: - body_mask = wp.clone(body_mask, device=self._device) - return body_mask + return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device) def _get_cpu_env_mask(self, env_mask: wp.array) -> wp.array: """Return a pinned-host CPU copy of *env_mask* for a CPU-only binding write. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py index c39c8c6d2f2b..1d4066c2709a 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object_collection/rigid_object_collection.py @@ -21,6 +21,7 @@ from isaaclab.utils import clone, validate from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_1d_mask, resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_ov import tensor_types as TT @@ -439,16 +440,8 @@ def write_body_link_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask - body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0]) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) self.assert_shape_and_dtype(body_poses, (self._num_instances, self._num_bodies), wp.transformf, "body_poses") wp.launch( shared_kernels.set_body_link_pose_to_sim_kernel(env_ids, body_ids), @@ -535,16 +528,8 @@ def write_body_com_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask - body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0]) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) self.assert_shape_and_dtype(body_poses, (self._num_instances, self._num_bodies), wp.transformf, "body_poses") wp.launch( shared_kernels.set_body_com_pose_to_sim_kernel(env_ids, body_ids), @@ -643,16 +628,8 @@ def write_body_com_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask - body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0]) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) self.assert_shape_and_dtype( body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities" ) @@ -758,16 +735,8 @@ def write_body_link_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_mask_t = wp.to_torch(body_mask) if isinstance(body_mask, wp.array) else body_mask - body_ids = self._resolve_body_ids(torch.nonzero(body_mask_t)[:, 0]) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) + body_ids = resolve_ids(mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self._device) self.assert_shape_and_dtype( body_velocities, (self._num_instances, self._num_bodies), wp.spatial_vectorf, "body_velocities" ) @@ -869,11 +838,7 @@ def set_masses_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) self.assert_shape_and_dtype(masses, (self._num_instances, self._num_bodies), wp.float32, "masses") wp.launch( shared_kernels.write_2d_data_to_buffer_with_mask, @@ -955,11 +920,7 @@ def set_coms_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) self.assert_shape_and_dtype(coms, (self._num_instances, self._num_bodies), wp.transformf, "coms") wp.launch( shared_kernels.write_body_com_pose_to_buffer_mask, @@ -1051,11 +1012,7 @@ def set_inertias_mask( env_mask: Environment mask. If None, all instances are updated. Shape is (num_instances,). """ - if env_mask is not None: - env_mask_t = wp.to_torch(env_mask) if isinstance(env_mask, wp.array) else env_mask - env_ids = self._resolve_env_ids(torch.nonzero(env_mask_t)[:, 0]) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self._device) self.assert_shape_and_dtype(inertias, (self._num_instances, self._num_bodies, 9), wp.float32, "inertias") wp.launch( shared_kernels.write_body_inertia_to_buffer_mask, @@ -1388,27 +1345,11 @@ def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor: resolved array onto ``self._device`` so kernel launches don't fail on a device mismatch. """ - if env_ids is None or env_ids == slice(None): - return self._ALL_ENV_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self._device) - if isinstance(env_ids, torch.Tensor): - return env_ids.to(device=self._device) - if isinstance(env_ids, wp.array) and str(env_ids.device) != self._device: - env_ids = wp.clone(env_ids, device=self._device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_ENV_INDICES, device=self._device) def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: """Resolve body indices on ``self._device``.""" - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if body_ids is None or body_ids == slice(None): - return self._ALL_BODY_INDICES - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self._device) - return body_ids + return resolve_ids(ids=body_ids, all_ids=self._ALL_BODY_INDICES, device=self._device) def _env_body_ids_to_view_ids( self, env_ids: torch.Tensor | wp.array, body_ids: torch.Tensor | wp.array, device: str = "cuda:0" @@ -1418,6 +1359,8 @@ def _env_body_ids_to_view_ids( Computes ``view_id = body_id * num_instances + env_id`` for each (env_id, body_id) pair. The output array is laid out column-major over the (env, body) grid, matching the PhysX ``root_view`` ordering. + The mask writers convert their masks to ids for this: a binding write mask must cover all + ``num_bodies * num_instances`` rows, so an env mask does not fit. Args: env_ids: Environment indices. @@ -1466,13 +1409,7 @@ def _resolve_env_mask(self, env_mask: wp.array | None) -> wp.array: Returns: A ``wp.bool`` array of shape (num_instances,) on ``self._device``. """ - if env_mask is None: - return self._ALL_TRUE_ENV_MASK - if isinstance(env_mask, torch.Tensor): - return wp.from_torch(env_mask.to(torch.bool), dtype=wp.bool) - if isinstance(env_mask, wp.array) and str(env_mask.device) != self._device: - env_mask = wp.clone(env_mask, device=self._device) - return env_mask + return resolve_1d_mask(mask=env_mask, all_mask=self._ALL_TRUE_ENV_MASK, device=self._device) def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: """Resolve a body mask to a ``wp.bool`` array on ``self._device``. @@ -1485,13 +1422,7 @@ def _resolve_body_mask(self, body_mask: wp.array | None) -> wp.array: Returns: A ``wp.bool`` array of shape (num_bodies,) on ``self._device``. """ - if body_mask is None: - return self._ALL_TRUE_BODY_MASK - if isinstance(body_mask, torch.Tensor): - return wp.from_torch(body_mask.to(torch.bool), dtype=wp.bool) - if isinstance(body_mask, wp.array) and str(body_mask.device) != self._device: - body_mask = wp.clone(body_mask, device=self._device) - return body_mask + return resolve_1d_mask(mask=body_mask, all_mask=self._ALL_TRUE_BODY_MASK, device=self._device) def _sim_env_ids_view(self, count: int) -> wp.array: """Return a cached prefix of the simulator-index scratch buffer.""" diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index e9b758f99eff..2085881cd1a6 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -1651,6 +1651,13 @@ def __getitem__(self, name): # The delay state is flattened environment-major over the group's joints. assert stateful_pairs[0].delay_state.num_pushes.numpy().tolist() == [0] * num_joints + [1] * num_joints + scene.step(articulation) + articulation.reset(env_mask=wp.array([False, True], dtype=wp.bool, device=device)) + stateful_pairs = [ + state for state in adapter._states_a if state is not None and getattr(state, "delay_state", None) is not None + ] + assert stateful_pairs[0].delay_state.num_pushes.numpy().tolist() == [1] * num_joints + [0] * num_joints + env = Env(articulation) asset_cfg = SceneEntityCfg("robot") event_params = { diff --git a/source/isaaclab_ov/test/assets/test_deformable_object.py b/source/isaaclab_ov/test/assets/test_deformable_object.py index bc9fc6a2e04b..5ea0b8ea84f6 100644 --- a/source/isaaclab_ov/test/assets/test_deformable_object.py +++ b/source/isaaclab_ov/test/assets/test_deformable_object.py @@ -386,6 +386,28 @@ def test_indexed_partial_write_preserves_retained_aliased_slice( torch.testing.assert_close(retained.torch, expected) +@pytest.mark.usefixtures("host_shell") +def test_mask_writes_update_only_masked_rows(): + """The mask writers take full-size data and write only the environments the mask selects.""" + asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) + zeros = torch.zeros((3, 2, 3), device=asset.device) + asset.root_view.positions = wp.from_torch(zeros.clone(), dtype=wp.float32) + asset.root_view.velocities = wp.from_torch(zeros.clone(), dtype=wp.float32) + asset.update(0.1) + env_mask = wp.array([False, True, False], dtype=wp.bool, device=asset.device) + + state = torch.cat((torch.full((3, 2, 3), 5.0), torch.full((3, 2, 3), -5.0)), dim=-1).to(asset.device) + asset.write_nodal_state_to_sim_mask(state, env_mask=env_mask) + expected = torch.zeros((3, 2, 6), device=asset.device) + expected[1] = state[1] + torch.testing.assert_close(asset.data.nodal_state_w.torch, expected) + + positions = torch.full((3, 2, 3), 9.0, device=asset.device) + asset.write_nodal_pos_to_sim_mask(positions, env_mask=env_mask) + expected[1, :, :3] = 9.0 + torch.testing.assert_close(asset.data.nodal_pos_w.torch, expected[..., :3]) + + @pytest.mark.usefixtures("host_shell") def test_full_overwrite_stale_cache_does_not_read_simulator() -> None: asset = _make_asset_shell(deformable_type="volume", num_instances=3, num_vertices=2) diff --git a/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst new file mode 100644 index 000000000000..fb8da731a750 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.changed.rst @@ -0,0 +1,7 @@ +* Changed the env selection helpers of the PhysX assets to resolve ``env_ids`` through + :func:`~isaaclab.utils.warp.utils.resolve_ids`. Any integer sequence and slices with a negative step are accepted, + ids on another device are moved to the asset's device, and a :class:`~isaaclab.utils.warp.ProxyArray` is rejected + with a :class:`TypeError`. +* Changed the body, joint and tendon selection helpers of the PhysX assets to resolve ids and masks through + :func:`~isaaclab.utils.warp.utils.resolve_ids`, as the env selection does, so slices other than ``slice(None)`` are + accepted. diff --git a/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst new file mode 100644 index 000000000000..96c018a902e7 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/jichuanh-id-mask.fixed.rst @@ -0,0 +1,2 @@ +* Fixed the ``write_nodal_*_to_sim_mask`` methods of :class:`~isaaclab_physx.assets.DeformableObject` raising a + :class:`RuntimeError` whenever ``env_mask`` was given. diff --git a/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.added.rst b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.added.rst new file mode 100644 index 000000000000..e51504073ab7 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.added.rst @@ -0,0 +1,3 @@ +* Added ``env_mask`` support to the Newton-native actuator reset of PhysX articulations, so + :meth:`~isaaclab.actuators.ActuatorCollection.reset` with a mask resets the native actuator state of only the masked + environments, without converting the mask to indices. diff --git a/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst new file mode 100644 index 000000000000..fcfe2bc98ec3 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/jichuanh-warp-scene-mask-native.fixed.rst @@ -0,0 +1,5 @@ +* Fixed :meth:`~isaaclab_physx.assets.Articulation.reset` with only ``env_mask`` resetting the actuator state of every + environment. +* Fixed :meth:`~isaaclab_physx.assets.RigidObject.reset` and + :meth:`~isaaclab_physx.assets.RigidObjectCollection.reset` ignoring ``env_mask`` and clearing the external + wrenches of every environment. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/actuator_control.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/actuator_control.py index da1714701482..13f0d39ef883 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/actuator_control.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/actuator_control.py @@ -22,6 +22,8 @@ from isaaclab.sim.utils.stage import get_current_stage if TYPE_CHECKING: + import warp as wp + from .articulation import Articulation logger = logging.getLogger(__name__) @@ -155,6 +157,6 @@ def submit_commands(self, collection: ActuatorCollection) -> None: articulation.root_view.set_dof_position_targets(pos_target, articulation._ALL_INDICES) articulation.root_view.set_dof_velocity_targets(vel_target, articulation._ALL_INDICES) - def reset_native_actuators(self, env_ids: Sequence[int] | slice) -> None: + def reset_native_actuators(self, env_ids: Sequence[int] | slice, env_mask: wp.array | None = None) -> None: if self._native_actuator_path_active and self._actuator_runtime is not None: - self._actuator_runtime.reset(env_ids) + self._actuator_runtime.reset(env_ids, env_mask) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 3d2cb4f1d946..89f627dd452c 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -27,6 +27,7 @@ from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -231,7 +232,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None env_ids = slice(None) # reset actuators, including backend-native actuator state. None selects all # environments; delayed-actuator buffers do not accept a slice. - self.actuators.reset(None if env_ids == slice(None) else env_ids) + self.actuators.reset(None if env_ids == slice(None) else 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) @@ -590,7 +591,7 @@ def write_root_link_pose_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_root_link_pose_to_sim_index( root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -692,7 +693,7 @@ def write_root_com_pose_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_root_com_pose_to_sim_index( root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -866,7 +867,7 @@ def write_root_com_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_root_com_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -977,7 +978,7 @@ def write_root_link_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_root_link_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -1091,8 +1092,8 @@ def write_joint_state_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # set into simulation - env_ids = self._resolve_env_mask(env_mask) - joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask) self.write_joint_state_to_sim_index( position=position, velocity=velocity, @@ -1196,8 +1197,8 @@ def write_joint_position_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_joint_position_to_sim_index( position=position, joint_ids=joint_ids, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -1301,8 +1302,8 @@ def write_joint_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = None if joint_mask is None else self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = None if joint_mask is None else self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_joint_velocity_to_sim_index( velocity=velocity, joint_ids=joint_ids, env_ids=env_ids, full_data=True, skip_forward=skip_forward @@ -1412,8 +1413,8 @@ def write_joint_stiffness_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the the right code path is taken inside the kernel. self.write_joint_stiffness_to_sim_index( stiffness=stiffness, joint_ids=joint_ids, env_ids=env_ids, full_data=True @@ -1560,8 +1561,8 @@ def write_joint_damping_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_damping_to_sim_index(damping=damping, joint_ids=joint_ids, env_ids=env_ids, full_data=True) @@ -1672,8 +1673,8 @@ def write_joint_position_limit_to_sim_mask( exceed the new limits. Defaults to True. """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_position_limit_to_sim_index( limits=limits, @@ -1787,8 +1788,8 @@ def write_joint_velocity_limit_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_velocity_limit_to_sim_index( limits=limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True @@ -1897,8 +1898,8 @@ def write_joint_effort_limit_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_effort_limit_to_sim_index(limits=limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True) @@ -2005,8 +2006,8 @@ def write_joint_armature_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_armature_to_sim_index(armature=armature, joint_ids=joint_ids, env_ids=env_ids, full_data=True) @@ -2173,8 +2174,8 @@ def write_joint_friction_coefficient_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_friction_coefficient_to_sim_index( joint_friction_coeff=joint_friction_coeff, @@ -2283,8 +2284,8 @@ def write_joint_dynamic_friction_coefficient_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_dynamic_friction_coefficient_to_sim_index( joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, @@ -2394,8 +2395,8 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks - env_ids = self._resolve_env_mask(env_mask) - joint_ids = self._resolve_joint_mask(joint_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + joint_ids = self._resolve_joint_ids(joint_mask=joint_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_viscous_friction_coefficient_to_sim_index( joint_viscous_friction_coeff=joint_viscous_friction_coeff, @@ -2496,8 +2497,8 @@ def set_masses_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - body_ids = self._resolve_body_mask(body_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -2603,8 +2604,8 @@ def set_coms_mask( all_envs_selected = env_mask is None all_bodies_selected = body_mask is None # Resolve masks. - env_ids = None if all_envs_selected else self._resolve_env_mask(env_mask) - body_ids = None if all_bodies_selected else self._resolve_body_mask(body_mask) + env_ids = None if all_envs_selected else self._resolve_env_ids(env_mask=env_mask) + body_ids = None if all_bodies_selected else self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -2700,8 +2701,8 @@ def set_inertias_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - body_ids = self._resolve_body_mask(body_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -2806,8 +2807,8 @@ def set_fixed_tendon_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_stiffness_index( stiffness=stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -2906,8 +2907,8 @@ def set_fixed_tendon_damping_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_damping_index( damping=damping, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3010,8 +3011,8 @@ def set_fixed_tendon_limit_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_limit_stiffness_index( limit_stiffness=limit_stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3020,7 +3021,7 @@ def set_fixed_tendon_limit_stiffness_mask( def set_fixed_tendon_position_limit_index( self, *, - limit: float | torch.Tensor | wp.array, + limit: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, @@ -3064,7 +3065,7 @@ def set_fixed_tendon_position_limit_index( def set_fixed_tendon_position_limit_mask( self, *, - limit: float | torch.Tensor | wp.array, + limit: torch.Tensor | wp.array, fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: @@ -3088,8 +3089,8 @@ def set_fixed_tendon_position_limit_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_position_limit_index( limit=limit, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3192,8 +3193,8 @@ def set_fixed_tendon_rest_length_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_rest_length_index( rest_length=rest_length, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3343,8 +3344,8 @@ def set_fixed_tendon_position_target_mask( """ self.set_fixed_tendon_position_target_index( target=target, - fixed_tendon_ids=self._resolve_fixed_tendon_mask(fixed_tendon_mask), - env_ids=self._resolve_env_mask(env_mask), + fixed_tendon_ids=self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask), + env_ids=self._resolve_env_ids(env_mask=env_mask), full_data=True, ) @@ -3374,8 +3375,8 @@ def set_fixed_tendon_offset_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - fixed_tendon_ids = self._resolve_fixed_tendon_mask(fixed_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + fixed_tendon_ids = self._resolve_fixed_tendon_ids(fixed_tendon_mask=fixed_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_offset_index( offset=offset, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True @@ -3430,7 +3431,7 @@ def write_fixed_tendon_properties_to_sim_mask( only the selected cached values before this method pushes the complete rows. .. tip:: - For maximum performance we recommend using the mask method. This is because in PhysX, the tensor API + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API is only supporting indexing, hence masks need to be converted to indices. Args: @@ -3439,7 +3440,7 @@ def write_fixed_tendon_properties_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_fixed_tendon_properties_to_sim_index(env_ids=env_ids) def set_spatial_tendon_stiffness_index( @@ -3539,8 +3540,8 @@ def set_spatial_tendon_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_stiffness_index( stiffness=stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True @@ -3639,8 +3640,8 @@ def set_spatial_tendon_damping_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_damping_index( damping=damping, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True @@ -3744,8 +3745,8 @@ def set_spatial_tendon_limit_stiffness_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_limit_stiffness_index( limit_stiffness=limit_stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True @@ -3844,8 +3845,8 @@ def set_spatial_tendon_offset_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) - spatial_tendon_ids = self._resolve_spatial_tendon_mask(spatial_tendon_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) + spatial_tendon_ids = self._resolve_spatial_tendon_ids(spatial_tendon_mask=spatial_tendon_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_offset_index( offset=offset, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True @@ -3898,7 +3899,7 @@ def write_spatial_tendon_properties_to_sim_mask( this method pushes the complete rows. .. tip:: - For maximum performance we recommend using the mask method. This is because in PhysX, the tensor API + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API is only supporting indexing, hence masks need to be converted to indices. Args: @@ -3907,7 +3908,7 @@ def write_spatial_tendon_properties_to_sim_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - env_ids = self._resolve_env_mask(env_mask) + env_ids = self._resolve_env_ids(env_mask=env_mask) # Write spatial tendon properties to the simulation. self.write_spatial_tendon_properties_to_sim_index(env_ids=env_ids) @@ -4463,62 +4464,23 @@ def _cpu_env_ids_view(self, count: int) -> wp.array: ) return self._cpu_env_ids_views[count] - def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array: - """ - Resolve environment mask to a torch tensor. - - Args: - env_mask: Environment mask. If None, then all indices are used. - - Returns: - A torch tensor of environment indices. - """ - # resolve masks - if env_mask is not None: - if isinstance(env_mask, wp.array): - env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0] - else: - env_ids = self._ALL_INDICES - return env_ids - - def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + def _resolve_env_ids( + self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve environment indices. Args: env_ids: Environment indices. If None, then all indices are used. + env_mask: Boolean mask; takes precedence over ``env_ids``. Returns: Environment indices. """ - if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids - - def _resolve_joint_mask(self, joint_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve joint mask to a torch tensor. - - Args: - joint_mask: Joint mask. If None, then all indices are used. + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) - Returns: - A torch tensor of joint indices. - """ - if joint_mask is not None: - if isinstance(joint_mask, wp.array): - joint_mask = wp.to_torch(joint_mask) - joint_ids = torch.nonzero(joint_mask)[:, 0] - else: - joint_ids = self._ALL_JOINT_INDICES - return joint_ids - - def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + def _resolve_joint_ids( + self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, joint_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve joint indices to a warp array or tensor. .. note:: @@ -4526,123 +4488,65 @@ def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array Args: joint_ids: Joint indices. If None, then all indices are used. + joint_mask: Boolean mask; takes precedence over ``joint_ids``. Returns: A warp array of joint indices or a tensor of joint indices. """ - if isinstance(joint_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(joint_ids, list): - return wp.array(joint_ids, dtype=wp.int32, device=self.device) - if (joint_ids is None) or (joint_ids == slice(None)): - return self._ALL_JOINT_INDICES - return joint_ids - - def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve body mask to a torch tensor. - - Args: - body_mask: Body mask. If None, then all indices are used. - - Returns: - A torch tensor of body indices. - """ - if body_mask is not None: - if isinstance(body_mask, wp.array): - body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0] - else: - body_ids = self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=joint_ids, mask=joint_mask, all_ids=self._ALL_JOINT_INDICES, device=self.device) - def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + def _resolve_body_ids( + self, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve body indices to a warp array or tensor. Args: body_ids: Body indices. If None, then all indices are used. + body_mask: Boolean mask; takes precedence over ``body_ids``. Returns: A warp array of body indices or a tensor of body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - return body_ids - - def _resolve_fixed_tendon_mask(self, fixed_tendon_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve fixed tendon mask to a torch tensor. - - Args: - fixed_tendon_mask: Fixed tendon mask. If None, then all indices are used. - - Returns: - A torch tensor of fixed tendon indices. - """ - if fixed_tendon_mask is not None: - if isinstance(fixed_tendon_mask, wp.array): - fixed_tendon_mask = wp.to_torch(fixed_tendon_mask) - fixed_tendon_ids = torch.nonzero(fixed_tendon_mask)[:, 0] - else: - fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES - return fixed_tendon_ids + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device) def _resolve_fixed_tendon_ids( - self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None + self, + tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + fixed_tendon_mask: wp.array | None = None, ) -> wp.array | torch.Tensor: """Resolve tendon indices to a warp array or tensor. Args: tendon_ids: Tendon indices. If None, then all indices are used. + fixed_tendon_mask: Boolean mask; takes precedence over ``tendon_ids``. Returns: A warp array of tendon indices or a tensor of tendon indices. """ - if isinstance(tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(tendon_ids, list): - return wp.array(tendon_ids, dtype=wp.int32, device=self.device) - if (tendon_ids is None) or (tendon_ids == slice(None)): - return self._ALL_FIXED_TENDON_INDICES - return tendon_ids - - def _resolve_spatial_tendon_mask(self, spatial_tendon_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve spatial tendon mask to a torch tensor. - - Args: - spatial_tendon_mask: Spatial tendon mask. If None, then all indices are used. - - Returns: - A torch tensor of spatial tendon indices. - """ - if spatial_tendon_mask is not None: - if isinstance(spatial_tendon_mask, wp.array): - spatial_tendon_mask = wp.to_torch(spatial_tendon_mask) - spatial_tendon_ids = torch.nonzero(spatial_tendon_mask)[:, 0] - else: - spatial_tendon_ids = self._ALL_SPATIAL_TENDON_INDICES - return spatial_tendon_ids + return resolve_ids( + ids=tendon_ids, mask=fixed_tendon_mask, all_ids=self._ALL_FIXED_TENDON_INDICES, device=self.device + ) def _resolve_spatial_tendon_ids( - self, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None + self, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + spatial_tendon_mask: wp.array | None = None, ) -> wp.array | torch.Tensor: """Resolve spatial tendon indices to a warp array or tensor. Args: spatial_tendon_ids: Spatial tendon indices. If None, then all indices are used. + spatial_tendon_mask: Boolean mask; takes precedence over ``spatial_tendon_ids``. Returns: A warp array of spatial tendon indices or a tensor of spatial tendon indices. """ - if isinstance(spatial_tendon_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(spatial_tendon_ids, list): - return wp.array(spatial_tendon_ids, dtype=wp.int32, device=self.device) - if (spatial_tendon_ids is None) or (spatial_tendon_ids == slice(None)): - return self._ALL_SPATIAL_TENDON_INDICES - return spatial_tendon_ids + return resolve_ids( + ids=spatial_tendon_ids, + mask=spatial_tendon_mask, + all_ids=self._ALL_SPATIAL_TENDON_INDICES, + device=self.device, + ) """ Deprecated methods. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py index 019d50bdb1a8..d8dc3d9e7bc3 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py @@ -21,6 +21,7 @@ from isaaclab.assets.asset_base import AssetBase from isaaclab.markers import VisualizationMarkers from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from .deformable_object_data import DeformableObjectData from .kernels import ( @@ -218,10 +219,7 @@ def write_nodal_state_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 6). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids, full_data=True) def write_nodal_pos_to_sim_index( @@ -288,10 +286,7 @@ def write_nodal_pos_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 3). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_pos_to_sim_index(nodal_pos, env_ids=env_ids, full_data=True) def write_nodal_velocity_to_sim_index( @@ -359,10 +354,7 @@ def write_nodal_velocity_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 3). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_velocity_to_sim_index(nodal_vel, env_ids=env_ids, full_data=True) def write_nodal_kinematic_target_to_sim_index( @@ -438,10 +430,7 @@ def write_nodal_kinematic_target_to_sim_mask( Shape is (num_instances, max_sim_vertices_per_body, 4). env_mask: Environment mask. If None, then all indices are used. """ - if env_mask is not None: - env_ids = wp.nonzero(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = resolve_ids(mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) self.write_nodal_kinematic_target_to_sim_index(targets, env_ids=env_ids, full_data=True) """ @@ -566,17 +555,7 @@ def _resolve_env_ids(self, env_ids): Returns: A warp array of environment indices. """ - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, torch.Tensor): - if env_ids.dtype == torch.int64: - env_ids = env_ids.to(torch.int32) - return wp.from_torch(env_ids, dtype=wp.int32) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids + return resolve_ids(ids=env_ids, all_ids=self._ALL_INDICES, device=self.device, dtype=wp.int32) def _initialize_impl(self): # obtain global simulation view diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py index 47c7240c510e..af8a0096890e 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py @@ -19,6 +19,7 @@ from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -137,8 +138,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. @@ -397,10 +398,7 @@ def write_root_link_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_root_link_pose_to_sim_index( root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward ) @@ -491,10 +489,7 @@ def write_root_com_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_root_com_pose_to_sim_index( root_pose=root_pose, env_ids=env_ids, full_data=True, skip_forward=skip_forward ) @@ -593,10 +588,7 @@ def write_root_com_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_root_com_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward ) @@ -699,10 +691,7 @@ def write_root_link_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.write_root_link_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, full_data=True, skip_forward=skip_forward ) @@ -793,14 +782,8 @@ def set_masses_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -885,14 +868,8 @@ def set_coms_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -980,14 +957,8 @@ def set_inertias_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -1068,40 +1039,6 @@ def _process_cfg(self) -> None: self._data.default_root_pose = wp.array(default_root_pose, dtype=wp.transformf, device=self.device) self._data.default_root_vel = wp.array(default_root_vel, dtype=wp.spatial_vectorf, device=self.device) - def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve environment mask to a torch tensor. - - Args: - env_mask: Environment mask. If None, then all indices are used. - - Returns: - A torch tensor of environment indices. - """ - if env_mask is not None: - if isinstance(env_mask, wp.array): - env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0] - else: - env_ids = self._ALL_INDICES - return env_ids - - def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve body mask to a torch tensor. - - Args: - body_mask: Body mask. If None, then all indices are used. - - Returns: - A torch tensor of body indices. - """ - if body_mask is not None: - if isinstance(body_mask, wp.array): - body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0] - else: - body_ids = self._ALL_BODY_INDICES - return body_ids - def _get_cpu_env_ids(self, env_ids: wp.array | torch.Tensor, sim_env_ids: wp.array | None = None) -> wp.array: """Get CPU environment indices. Uses pre-allocated pinned buffer for full-index case. @@ -1164,39 +1101,33 @@ def _sim_env_ids_view(self, count: int) -> wp.array: ) return self._sim_env_ids_views[count] - def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + def _resolve_env_ids( + self, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve environment indices to a warp array or tensor. Args: env_ids: Environment indices. If None, then all indices are used. + env_mask: Boolean mask; takes precedence over ``env_ids``. Returns: A warp array of environment indices or a tensor of environment indices. """ - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids - - def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self.device) + + def _resolve_body_ids( + self, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_mask: wp.array | None = None + ) -> wp.array | torch.Tensor: """Resolve body indices to a warp array or tensor. Args: body_ids: Body indices. If None, then all indices are used. + body_mask: Boolean mask; takes precedence over ``body_ids``. Returns: A warp array of body indices or a tensor of body indices. """ - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if (body_ids is None) or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - return body_ids + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device) """ Internal simulation callbacks. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index 47d5e1cb460b..88c914e485ad 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -22,6 +22,7 @@ from isaaclab.assets.rigid_object_collection.base_rigid_object_collection import BaseRigidObjectCollection from isaaclab.utils import clone, validate from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.utils import resolve_ids from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -174,8 +175,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. @@ -307,14 +308,8 @@ def write_body_pose_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) self.write_body_link_pose_to_sim_index( body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward ) @@ -391,14 +386,8 @@ def write_body_velocity_to_sim_mask( skip_forward: Whether to skip invalidating cached data after the write. When True, the caller must invalidate stale cached data before reading it back. Defaults to False. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) self.write_body_com_velocity_to_sim_index( body_velocities=body_velocities, env_ids=env_ids, @@ -509,10 +498,7 @@ def write_body_link_pose_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids) self.write_body_link_pose_to_sim_index( body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward @@ -624,10 +610,7 @@ def write_body_com_pose_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids) self.write_body_com_pose_to_sim_index( body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True, skip_forward=skip_forward @@ -744,10 +727,7 @@ def write_body_com_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids) self.write_body_com_velocity_to_sim_index( body_velocities=body_velocities, @@ -871,10 +851,7 @@ def write_body_link_velocity_to_sim_mask( must invalidate stale cached data before reading it back. Defaults to False. body_ids: Deprecated, use :attr:`body_mask` instead. Body indices. Defaults to None. """ - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) body_ids = self._resolve_mask_writer_body_ids(body_mask, body_ids) self.write_body_link_velocity_to_sim_index( body_velocities=body_velocities, @@ -968,14 +945,8 @@ def set_masses_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -1062,14 +1033,8 @@ def set_coms_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -1157,14 +1122,8 @@ def set_inertias_mask( env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. - if env_mask is not None: - env_ids = self._resolve_env_mask(env_mask) - else: - env_ids = self._ALL_ENV_INDICES - if body_mask is not None: - body_ids = self._resolve_body_mask(body_mask) - else: - body_ids = self._ALL_BODY_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) + body_ids = self._resolve_body_ids(body_mask=body_mask) # Set full data to True to ensure the right code path is taken inside the kernel. self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) @@ -1311,46 +1270,20 @@ def reshape_data_to_view_3d( Internal helper. """ - def _resolve_env_ids(self, env_ids) -> wp.array | torch.Tensor: + def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | None = None) -> wp.array | torch.Tensor: """Resolve environment indices.""" - if (env_ids is None) or (env_ids == slice(None)): - return self._ALL_ENV_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_ENV_INDICES)[env_ids]) - if isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self.device) - return env_ids - - def _resolve_body_ids(self, body_ids) -> wp.array | torch.Tensor: - """Resolve body indices.""" - if isinstance(body_ids, ProxyArray): - raise TypeError("ProxyArray is output-only; pass .warp or .torch explicitly.") - if isinstance(body_ids, list): - return wp.array(body_ids, dtype=wp.int32, device=self.device) - if body_ids is None or (body_ids == slice(None)): - return self._ALL_BODY_INDICES - if isinstance(body_ids, slice): - return wp.from_torch( - torch.arange(self.num_bodies, dtype=torch.int32, device=self.device)[body_ids], dtype=wp.int32 - ) - return body_ids + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_ENV_INDICES, device=self.device) - def _resolve_env_mask(self, env_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve environment mask to indices via torch.nonzero.""" - if env_mask is not None: - if isinstance(env_mask, wp.array): - env_mask = wp.to_torch(env_mask) - env_ids = torch.nonzero(env_mask)[:, 0] - else: - env_ids = self._ALL_ENV_INDICES - return env_ids + def _resolve_body_ids(self, body_ids=None, body_mask: wp.array | None = None) -> wp.array | torch.Tensor: + """Resolve body indices.""" + return resolve_ids(ids=body_ids, mask=body_mask, all_ids=self._ALL_BODY_INDICES, device=self.device) def _resolve_mask_writer_body_ids( self, body_mask: wp.array | None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None ) -> Sequence[int] | torch.Tensor | wp.array | slice: """Resolve the bodies of a mask writer, accepting the deprecated ``body_ids`` in place of ``body_mask``.""" if body_ids is None: - return self._resolve_body_mask(body_mask) + return self._resolve_body_ids(body_mask=body_mask) if body_mask is not None: raise ValueError("Pass either 'body_mask' or the deprecated 'body_ids', not both.") warnings.warn( @@ -1361,16 +1294,6 @@ def _resolve_mask_writer_body_ids( ) return body_ids - def _resolve_body_mask(self, body_mask: wp.array | None) -> torch.Tensor | wp.array: - """Resolve body mask to indices via torch.nonzero.""" - if body_mask is not None: - if isinstance(body_mask, wp.array): - body_mask = wp.to_torch(body_mask) - body_ids = torch.nonzero(body_mask)[:, 0] - else: - body_ids = self._ALL_BODY_INDICES - return body_ids - def _initialize_impl(self): # clear body names list to prevent double counting on re-initialization self._body_names_list.clear() diff --git a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py index e37a30f7a505..a2f5cb930c8f 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py @@ -16,6 +16,7 @@ from isaaclab.assets import AssetBase from isaaclab.utils import clone from isaaclab.utils.version import get_isaac_sim_version, has_kit +from isaaclab.utils.warp.utils import resolve_ids if TYPE_CHECKING: from isaacsim.robot.surface_gripper import GripperView @@ -187,10 +188,7 @@ def set_grippers_command_mask(self, states: torch.Tensor | wp.array, env_mask: t states: A tensor/array of floats representing the gripper command (full size). env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments). """ - if env_mask is not None: - env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.set_grippers_command_index(states, env_ids, full_data=True) def set_grippers_command(self, states: torch.Tensor, indices: torch.Tensor | None = None) -> None: @@ -274,10 +272,7 @@ def update_gripper_properties_mask( retry_interval: The retry interval. Shape (num_envs,). env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments). """ - if env_mask is not None: - env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.update_gripper_properties_index( max_grip_distance=max_grip_distance, coaxial_force_limit=coaxial_force_limit, @@ -407,10 +402,7 @@ def reset_mask(self, env_mask: torch.Tensor | None = None) -> None: Args: env_mask: Boolean mask of shape (num_envs,). Defaults to None (all environments). """ - if env_mask is not None: - env_ids = wp.from_torch(torch.argwhere(env_mask).view(-1).to(torch.int32), dtype=wp.int32) - else: - env_ids = self._ALL_INDICES + env_ids = self._resolve_env_ids(env_mask=env_mask) self.reset_index(env_ids) def reset(self, indices: torch.Tensor | None = None) -> None: @@ -552,24 +544,17 @@ def _process_cfg(self) -> None: Helper functions. """ - def _resolve_env_ids(self, env_ids) -> wp.array: + def _resolve_env_ids(self, env_ids=None, env_mask: wp.array | torch.Tensor | None = None) -> wp.array: """Resolve environment indices to a warp array. Args: env_ids: Environment indices. Can be None, a slice, a list, or a torch.Tensor. + env_mask: Boolean mask; takes precedence over ``env_ids``. Returns: A warp array of int32 indices. """ - if env_ids is None or env_ids == slice(None): - return self._ALL_INDICES - if isinstance(env_ids, slice): - return wp.from_torch(wp.to_torch(self._ALL_INDICES)[env_ids]) - elif isinstance(env_ids, list): - return wp.array(env_ids, dtype=wp.int32, device=self._device) - elif isinstance(env_ids, torch.Tensor): - return wp.from_torch(env_ids.to(torch.int32).contiguous(), dtype=wp.int32) - return env_ids + return resolve_ids(ids=env_ids, mask=env_mask, all_ids=self._ALL_INDICES, device=self._device, dtype=wp.int32) def parse_gripper_parameter( self, cfg_value: float | int | tuple | None, default_value: float | int | tuple | None, ndim: int = 0 diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 48344bd748c8..b33ecfa87ef9 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -1086,9 +1086,9 @@ def test_floating_articulation_root_writes_and_wrenches(articulation_scene: _Art actuator_reset = actuator.reset reset_env_ids = [] - def record_actuator_reset(env_ids=None) -> None: + def record_actuator_reset(env_ids=None, env_mask=None) -> None: reset_env_ids.append(env_ids) - actuator_reset(env_ids) + actuator_reset(env_ids, env_mask=env_mask) monkeypatch.setattr(actuator, "reset", record_actuator_reset) composers = (floating.instantaneous_wrench_composer, floating.permanent_wrench_composer) diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index cfd389fa99ef..4eeb4c246f8c 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -244,6 +244,22 @@ def test_deformable_nodal_state_writes(deformable_scene: _DeformableScene) -> No torch.testing.assert_close(raw_nodal_pos.reshape(nodal_pos.shape), nodal_pos) +@pytest.mark.isaacsim_ci +def test_deformable_nodal_state_mask_writes(deformable_scene: _DeformableScene) -> None: + """Mask writes take full-size data and reach only the environments the mask selects.""" + scene = deformable_scene + device = scene.device + env_mask = wp.array([False, True], dtype=wp.bool, device=device) + for deformable in (scene.volume, scene.surface): + initial_state = deformable.data.nodal_state_w.torch.clone() + state = initial_state.clone() + state[:, 0, :3] += torch.tensor([0.01, 0.02, 0.03], device=device) + state[:, 0, 3] = 0.25 + deformable.write_nodal_state_to_sim_mask(state, env_mask=env_mask) + expected = torch.cat((initial_state[:1], state[1:])) + torch.testing.assert_close(deformable.data.nodal_state_w.torch, expected) + + @pytest.mark.isaacsim_ci def test_volume_kinematic_targets_hold_selected_nodes(deformable_scene: _DeformableScene) -> None: """Kinematic targets of environment 0 hold its nodes while environment 1 keeps moving freely.""" diff --git a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py index 3e84d69abcec..066e1ffb6e5a 100644 --- a/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py +++ b/source/isaaclab_physx/test/assets/test_newton_actuators_physx.py @@ -296,9 +296,16 @@ def test_network_actuators_drive_haa_joints(rollouts: _Rollouts, robot: str) -> assert not torch.allclose(trace["applied_effort"][0][:, haa_ids], trace["applied_effort"][-1][:, haa_ids]) -def test_reset_clears_actuator_state_of_selected_envs(rollouts: _Rollouts) -> None: - """Resetting environment 0 clears the delay history of its DoFs only.""" +@pytest.mark.parametrize("selector", ["env_ids", "env_mask"]) +def test_reset_clears_actuator_state_of_selected_envs(rollouts: _Rollouts, selector: str) -> None: + """Resetting environment 0 by indices or by mask clears the delay history of its DoFs only.""" articulation = rollouts.robots["delayed"] + # Refill the delay history that a previous reset cleared. + sim = SimulationContext.instance() + for _ in range(3): + articulation.write_data_to_sim() + sim.step() + articulation.update(DT) adapter = articulation.newton_actuator_adapter delay_states = [ (actuator, state.delay_state) @@ -308,7 +315,10 @@ def test_reset_clears_actuator_state_of_selected_envs(rollouts: _Rollouts) -> No assert delay_states assert all((delay_state.num_pushes.numpy() > 0).all() for _, delay_state in delay_states) - articulation.reset(env_ids=torch.tensor([0], device=articulation.device)) + if selector == "env_ids": + articulation.reset(env_ids=torch.tensor([0], device=articulation.device)) + else: + articulation.reset(env_mask=wp.from_torch(torch.arange(NUM_ENVS, device=articulation.device) == 0)) for actuator, delay_state in delay_states: env_of_dof = actuator.indices.numpy() // adapter.num_joints pushes = delay_state.num_pushes.numpy() diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index 764e54ff941d..bdeb573cd776 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -644,17 +644,23 @@ def test_rigid_object_wrench_delivery_and_reset(scene: _Scene) -> None: > 0.99 ) - # A partial reset clears only the selected environment; a full reset deactivates the composers. + # A partial reset by indices or by mask clears only the selected environment; a full reset deactivates the + # composers. composers = (cube_object.instantaneous_wrench_composer, cube_object.permanent_wrench_composer) ones = torch.ones((_NUM_ENVS, 1, 3), device=device) - cube_object.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) - cube_object.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) - cube_object.reset(env_ids=torch.tensor([0], device=device)) - for composer in composers: - assert composer.active - for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): - assert torch.count_nonzero(buffer[0]) == 0 - assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() + selections = { + "env_ids": torch.tensor([0], device=device), + "env_mask": wp.from_torch(torch.arange(_NUM_ENVS, device=device) == 0), + } + for selector, selection in selections.items(): + cube_object.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + cube_object.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + cube_object.reset(**{selector: selection}) + for composer in composers: + assert composer.active + for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): + assert torch.count_nonzero(buffer[0]) == 0, selector + assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel(), selector cube_object.reset() for composer in composers: assert not composer.active @@ -834,17 +840,22 @@ def test_collection_wrench_delivery_and_reset(scene: _Scene) -> None: ) assert torch.all(heading[..., 1] > 0.99) - # A partial reset clears the external wrenches of the selected environments only. + # A partial reset by indices or by mask clears the external wrenches of the selected environments only. composers = (collection.instantaneous_wrench_composer, collection.permanent_wrench_composer) ones = torch.ones((_NUM_ENVS, _NUM_CUBES, 3), device=device) - collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) - collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) - collection.reset(env_ids=torch.tensor([0], device=device)) - for composer in composers: - assert composer.active - for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): - assert torch.count_nonzero(buffer[0]) == 0 - assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel() + selections = { + "env_ids": torch.tensor([0], device=device), + "env_mask": wp.from_torch(torch.arange(_NUM_ENVS, device=device) == 0), + } + for selector, selection in selections.items(): + collection.permanent_wrench_composer.set_forces_and_torques_index(forces=ones, torques=ones) + collection.instantaneous_wrench_composer.add_forces_and_torques_index(forces=ones, torques=ones) + collection.reset(**{selector: selection}) + for composer in composers: + assert composer.active + for buffer in (composer.out_force_b.torch, composer.out_torque_b.torch): + assert torch.count_nonzero(buffer[0]) == 0, selector + assert torch.count_nonzero(buffer[1:]) == buffer[1:].numel(), selector # A full reset clears every environment collection.reset() diff --git a/source/isaaclab_physx/test/sensors/test_frame_transformer.py b/source/isaaclab_physx/test/sensors/test_frame_transformer.py index 7ba9a23d3de4..50e52bfbc646 100644 --- a/source/isaaclab_physx/test/sensors/test_frame_transformer.py +++ b/source/isaaclab_physx/test/sensors/test_frame_transformer.py @@ -655,6 +655,8 @@ def _make_frame_transformer(use_recorded_launch: bool = True): sensor.cfg = SimpleNamespace(prim_path="/World/Source") sensor._device = device sensor._num_envs = 1 + sensor._ALL_ENV_MASK = wp.ones(1, dtype=wp.bool, device=device) + sensor._reset_mask = wp.zeros(1, dtype=wp.bool, device=device) sensor._num_target_frames = 1 sensor._frame_physx_view = transform_view sensor._backend_manager = SimpleNamespace(update_kinematics=lambda: None) diff --git a/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py b/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py index ea2a1b6df989..80c2956113de 100644 --- a/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py +++ b/source/isaaclab_physx/test/sensors/test_imu_pva_recorded_launch.py @@ -122,6 +122,8 @@ def _make_sensor(sensor_type: str, use_recorded_launch: bool = True): sensor._device = device sensor._backend_manager = SimpleNamespace(update_kinematics=lambda: None) sensor._num_envs = 2 + sensor._ALL_ENV_MASK = wp.ones(2, dtype=wp.bool, device=device) + sensor._reset_mask = wp.zeros(2, dtype=wp.bool, device=device) sensor._view = rigid_view sensor._offset_pos_b = wp.zeros(2, dtype=wp.vec3f, device=device) sensor._offset_quat_b = wp.array( diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst new file mode 100644 index 000000000000..0b608386bcdd --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.added.rst @@ -0,0 +1,5 @@ +* Added :class:`~isaaclab_tasks_experimental.core.reorient.config.shadow_hand.shadow_hand_warp_env.ShadowHandDirectWarpEnv`, + so ``Isaac-Reorient-Cube-Shadow-Direct`` runs under ``--frontend warp``. Like the stable task, it + commands the four finger tendons from the last four action columns. +* Added the ``Metrics/success_rate`` episode metric to the Warp in-hand reorientation environment, + matching the stable Direct task. diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.changed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.changed.rst new file mode 100644 index 000000000000..20484dbf7694 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.changed.rst @@ -0,0 +1,3 @@ +* Changed the Warp in-hand reorientation environment to raise :class:`NotImplementedError` for + configurations with ``asymmetric_obs`` enabled instead of silently omitting the critic + observations. diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.fixed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.fixed.rst new file mode 100644 index 000000000000..393a1e9d9837 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.fixed.rst @@ -0,0 +1,5 @@ +* Fixed the ``Isaac-Reorient-Cube-Allegro-Direct`` Warp environment measuring the object's + orientation error from the ``(y, z, w)`` quaternion components instead of ``(x, y, z)``, which + made its orientation reward, goal success, and goal resampling differ from the stable task. +* Fixed the ``Isaac-Reorient-Cube-Allegro-Direct`` Warp environment swapping the object's linear + and angular velocity in its full observation. diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.minor b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-task-cleanup-dex-warp-exp.minor new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst new file mode 100644 index 000000000000..918eff1c77e5 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-frontend-cleanup.changed.rst @@ -0,0 +1,4 @@ +* Changed the Cartpole, locomotion and reorient Warp direct environments to record their ``_apply_action``, + ``_get_dones``, ``_get_rewards``, ``_get_observations`` and ``_reset_idx`` methods as CUDA graphs with + :func:`~isaaclab_experimental.utils.captured`, since :class:`~isaaclab_experimental.envs.DirectRLEnvWarp` no + longer records them. diff --git a/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst new file mode 100644 index 000000000000..59b2cee59f4f --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/jichuanh-warp-rng-shared.changed.rst @@ -0,0 +1,2 @@ +* Changed the Cartpole, locomotion and reorient Warp environments to draw from the process-wide + :class:`~isaaclab.utils.seed.WarpRng` state instead of allocating and seeding their own. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py index 2a3ccdd47f0d..b43b54c01d3d 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/cartpole/cartpole_warp_env.py @@ -10,6 +10,9 @@ import warp as wp from isaaclab_experimental.envs import DirectRLEnvWarp from isaaclab_experimental.utils.warp.utils import wrap_to_pi +from isaaclab_experimental.utils.warp_capture import captured + +from isaaclab.utils.seed import WarpRng if TYPE_CHECKING: from isaaclab_tasks.core.cartpole.cartpole_direct_env_cfg import CartpoleEnvCfg @@ -174,16 +177,6 @@ def reset( state[env_index] = rng_state -@wp.kernel -def initialize_state( - state: wp.array(dtype=wp.uint32), - seed: wp.int32, -): - """Initialize each env's random number generator state from the seed.""" - env_index = wp.tid() - state[env_index] = wp.rand_init(seed, env_index) - - class CartpoleWarpEnv(DirectRLEnvWarp): cfg: CartpoleEnvCfg @@ -206,19 +199,6 @@ def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs self.observations = wp.zeros((self.num_envs), dtype=wp.vec4f, device=self.device) self.actions = wp.zeros((self.num_envs, self.cartpole.num_joints), dtype=wp.float32, device=self.device) self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.device) - self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.device) - - if self.cfg.seed is None: - self.cfg.seed = -1 - - wp.launch( - initialize_state, - dim=self.num_envs, - inputs=[ - self.states, - self.cfg.seed, - ], - ) # Bind torch buffers to warp buffers self.torch_obs_buf = wp.to_torch(self.observations) @@ -239,9 +219,11 @@ def _pre_physics_step(self, actions: wp.array) -> None: ], ) + @captured def _apply_action(self) -> None: self.cartpole.set_joint_effort_target_mask(target=self.actions) + @captured def _get_observations(self) -> dict: wp.launch( get_observations, @@ -258,6 +240,7 @@ def _get_observations(self) -> dict: ) return {"policy": self.torch_obs_buf} + @captured def _get_rewards(self) -> None: wp.launch( compute_rewards, @@ -278,6 +261,7 @@ def _get_rewards(self) -> None: ], ) + @captured def _get_dones(self) -> None: wp.launch( get_dones, @@ -294,6 +278,7 @@ def _get_dones(self) -> None: ], ) + @captured def _reset_idx(self, mask: wp.array | None = None) -> None: if mask is None: mask = self._ALL_ENV_MASK @@ -317,6 +302,6 @@ def _reset_idx(self, mask: wp.array | None = None) -> None: self.cfg.initial_pole_angle_range, self.cfg.initial_pole_velocity_range, mask, - self.states, + WarpRng.state, ], ) 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 index 0e7ac5ca0fcf..b9f61f1b4c05 100644 --- 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 @@ -40,11 +40,11 @@ def _joint_pos_target_l2_kernel( 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] + assert asset.data.joint_pos.warp.shape[1] == asset_cfg.joint_mask_wp.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], + inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask_wp, out, target], device=env.device, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py index dfcfad6e0f55..dad1560b9cf1 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/locomotion_env_warp.py @@ -7,8 +7,10 @@ import warp as wp from isaaclab_experimental.envs import DirectRLEnvWarp +from isaaclab_experimental.utils.warp_capture import captured from isaaclab.envs import DirectRLEnvCfg +from isaaclab.utils.seed import WarpRng from isaaclab.utils.string import resolve_matching_names_values @@ -416,15 +418,12 @@ def survival_rate( @wp.kernel -def initialize_state( +def initialize_targets( env_origins: wp.array(dtype=wp.vec3f), targets: wp.array(dtype=wp.vec3f), - state: wp.array(dtype=wp.uint32), - seed: wp.int32, ): - """Initialize each env's RNG state and place its target far along +x from the env origin.""" + """Place each env's target far along +x from the env origin.""" env_index = wp.tid() - state[env_index] = wp.rand_init(seed, env_index) targets[env_index] = env_origins[env_index] targets[env_index] += wp.static(wp.vec3f(1000.0, 0.0, 0.0)) @@ -480,7 +479,6 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs ) self.rewards = wp.zeros((self.num_envs), dtype=wp.float32, device=self.sim.device) self.actions = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) - self.states = wp.zeros((self.num_envs), dtype=wp.uint32, device=self.sim.device) self.potentials = wp.zeros(self.num_envs, dtype=wp.float32, device=self.sim.device) self.prev_potentials = wp.zeros_like(self.potentials) self.targets = wp.zeros((self.num_envs), dtype=wp.vec3f, device=self.sim.device) @@ -496,18 +494,13 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.env_origins = wp.from_torch(self.scene.env_origins, dtype=wp.vec3f) self.actions_mapped = wp.zeros((self.num_envs, self.robot.num_joints), dtype=wp.float32, device=self.sim.device) - # Initial states and targets - if self.cfg.seed is None: - self.cfg.seed = -1 - + # Initial targets wp.launch( - initialize_state, + initialize_targets, dim=self.num_envs, inputs=[ self.env_origins, self.targets, - self.states, - self.cfg.seed, ], ) @@ -533,6 +526,7 @@ def _pre_physics_step(self, actions: wp.array) -> None: inputs=[actions, self.actions_mapped, self.joint_gears, self.action_scale], ) + @captured def _apply_action(self) -> None: self.robot.set_joint_effort_target_mask(target=self.actions_mapped) @@ -576,6 +570,7 @@ def _compute_intermediate_values(self) -> None: ], ) + @captured def _get_observations(self) -> dict: wp.launch( observations, @@ -602,6 +597,7 @@ def _get_observations(self) -> dict: ) return {"policy": self.torch_obs_buf} + @captured def _get_rewards(self) -> None: wp.launch( compute_rewards, @@ -629,6 +625,7 @@ def _get_rewards(self) -> None: ], ) + @captured def _get_dones(self) -> None: self._compute_intermediate_values() @@ -646,6 +643,7 @@ def _get_dones(self) -> None: ], ) + @captured def _reset_idx(self, mask: wp.array | None = None): if mask is None: mask = self._ALL_ENV_MASK @@ -684,7 +682,7 @@ def _reset_idx(self, mask: wp.array | None = None): self.soft_joint_vel_limits, self.joint_pos, self.joint_vel, - self.states, + WarpRng.state, self.cfg.initial_joint_pos_range[0], self.cfg.initial_joint_pos_range[1], self.cfg.initial_joint_vel_range[0], diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py new file mode 100644 index 000000000000..f758624246e0 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/__init__.py @@ -0,0 +1,6 @@ +# 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 of the hand-specific in-hand reorientation tasks.""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/__init__.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/__init__.py new file mode 100644 index 000000000000..6abd1f868033 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/__init__.py @@ -0,0 +1,11 @@ +# 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 Shadow Hand reorientation task. + +``Isaac-Reorient-Cube-Shadow-Direct`` resolves +:class:`~isaaclab_tasks_experimental.core.reorient.config.shadow_hand.shadow_hand_warp_env.ShadowHandDirectWarpEnv` +by name from the stable env class; run the stable id with ``--frontend warp`` and ``presets=newton_mjwarp``. +""" diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/shadow_hand_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/shadow_hand_warp_env.py new file mode 100644 index 000000000000..349bfda33074 --- /dev/null +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/config/shadow_hand/shadow_hand_warp_env.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 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab_tasks.core.reorient.utils import resolve_actuated_tendons + +from ...reorient_warp_env import ReorientDirectWarpEnv, scale + +if TYPE_CHECKING: + from isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_env_cfg import ShadowHandEnvCfg + + +@wp.kernel +def apply_actions_to_tendon_targets( + # input + actions: wp.array2d(dtype=wp.float32), + first_tendon_action: wp.int32, + lower_limits: wp.array2d(dtype=wp.float32), + upper_limits: wp.array2d(dtype=wp.float32), + tendon_ids: wp.array(dtype=wp.int32), + # output + tendon_targets: wp.array2d(dtype=wp.float32), +): + """Scale tendon actions to each tendon's command range and clamp them to it, without smoothing.""" + env_id, i = wp.tid() + lower = lower_limits[env_id, i] + upper = upper_limits[env_id, i] + target = scale(actions[env_id, first_tendon_action + i], lower, upper) + tendon_targets[env_id, tendon_ids[i]] = wp.clamp(target, lower, upper) + + +class ShadowHandDirectWarpEnv(ReorientDirectWarpEnv): + """Warp twin of the stable Shadow Hand reorientation environment. + + Implements :class:`~isaaclab_tasks.core.reorient.config.shadow_hand.shadow_hand_direct_env.ShadowHandDirectEnv`: + four of this hand's twenty motors pull a tendon spanning a finger's middle and distal joints. + Actions are ordered joints first, then tendons. + """ + + cfg: ShadowHandEnvCfg + + def __init__(self, cfg: ShadowHandEnvCfg, render_mode: str | None = None, **kwargs): + super().__init__(cfg, render_mode, **kwargs) + tendon_ids, lower_limits, upper_limits = resolve_actuated_tendons( + self.hand, cfg.actuated_tendon_names, self.num_envs, self.device, cfg.actuated_tendon_position_limits + ) + self.num_actuated_tendons = len(tendon_ids) + self.actuated_tendon_ids = wp.array(tendon_ids, dtype=wp.int32, device=self.device) + self.actuated_tendon_mask = wp.array( + [tendon_id in tendon_ids for tendon_id in range(self.hand.num_fixed_tendons)], + dtype=wp.bool, + device=self.device, + ) + self.tendon_lower_limits = wp.from_torch(lower_limits) + self.tendon_upper_limits = wp.from_torch(upper_limits) + self.tendon_targets = wp.zeros( + (self.num_envs, self.hand.num_fixed_tendons), dtype=wp.float32, device=self.device + ) + + def _pre_physics_step(self, actions: wp.array) -> None: + super()._pre_physics_step(actions) + # Commanded here, outside the captured action stage: the articulation flags a buffered + # tendon target for submission on the host, which a graph replay would not repeat. + wp.launch( + apply_actions_to_tendon_targets, + dim=(self.num_envs, self.num_actuated_tendons), + inputs=[ + self.actions, + self.num_actuated_dofs, + self.tendon_lower_limits, + self.tendon_upper_limits, + self.actuated_tendon_ids, + self.tendon_targets, + ], + device=self.device, + ) + self.hand.set_fixed_tendon_position_target_mask( + target=self.tendon_targets, fixed_tendon_mask=self.actuated_tendon_mask + ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py index c2a34edc9da6..a25b9d547c07 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/reorient/reorient_warp_env.py @@ -6,29 +6,19 @@ from __future__ import annotations -from collections.abc import Sequence from typing import TYPE_CHECKING import torch import warp as wp from isaaclab_experimental.envs import DirectRLEnvWarp +from isaaclab_experimental.utils.warp_capture import captured + +from isaaclab.utils.seed import WarpRng if TYPE_CHECKING: from isaaclab_tasks.core.reorient.config.allegro_hand.allegro_hand_direct_env_cfg import AllegroHandEnvCfg -@wp.kernel -def initialize_rng_state( - # input - seed: wp.int32, - # output - state: wp.array(dtype=wp.uint32), -): - """Initialize each env's random number generator state from the seed.""" - env_id = wp.tid() - state[env_id] = wp.rand_init(seed, wp.int32(env_id)) - - @wp.kernel def apply_actions_to_targets( # input @@ -106,7 +96,7 @@ def reset_object( root_pose_w: wp.array(dtype=wp.transformf), root_vel_w: wp.array(dtype=wp.spatial_vectorf), ): - """Reset masked envs' object root pose with position noise and a randomized rotation, zeroing its velocity.""" + """Sample masked envs' object root pose with position noise and a randomized rotation, and a zero velocity.""" env_id = wp.tid() if env_mask[env_id]: nx = wp.randf(rng_state[env_id], wp.float32(-1.0), wp.float32(1.0)) @@ -126,7 +116,6 @@ def reset_object( rng_state[env_id] += wp.uint32(1) rot_w = randomize_rotation(rand0, rand1, x_unit_vec, y_unit_vec) - # The following should be equivalent, but consider using write_root_pose_to_sim and write_root_velocity_to_sim root_pose_w[env_id] = wp.transform(pos_w, rot_w) root_vel_w[env_id] = wp.spatial_vectorf( wp.float32(0.0), wp.float32(0.0), wp.float32(0.0), wp.float32(0.0), wp.float32(0.0), wp.float32(0.0) @@ -151,9 +140,8 @@ def reset_hand( joint_vel: wp.array2d(dtype=wp.float32), prev_targets: wp.array2d(dtype=wp.float32), cur_targets: wp.array2d(dtype=wp.float32), - hand_dof_targets: wp.array2d(dtype=wp.float32), ): - """Reset masked envs' hand joint positions and velocities with noise and initialize their position targets.""" + """Sample masked envs' reset joint positions and velocities with noise and initialize their position targets.""" env_id = wp.tid() if env_mask[env_id]: # Each env runs sequentially inside this kernel (avoids RNG races across DOFs). @@ -170,29 +158,47 @@ def reset_hand( rng_state[env_id] += wp.uint32(1) vel = default_joint_vel[env_id, dof_id] + reset_dof_vel_noise * dof_vel_noise - # The following lines should be equivalent to the following: - # self.hand.write_joint_state_to_sim(dof_pos, dof_vel, env_ids=env_ids) joint_pos[env_id, dof_id] = pos joint_vel[env_id, dof_id] = vel - prev_targets[env_id, dof_id] = pos cur_targets[env_id, dof_id] = pos - hand_dof_targets[env_id, dof_id] = pos @wp.kernel def reset_successes( # input env_mask: wp.array(dtype=wp.bool), - # output + # input/output successes: wp.array(dtype=wp.float32), + success_rate_stats: wp.array(dtype=wp.float32), ): - """Zero the consecutive-success counter for masked envs.""" + """Accumulate masked envs' finished-episode success rate into ``success_rate_stats``, then zero their counts. + + Reaching a goal draws a replacement, so an episode that reached ``n`` goals attempted ``n + 1``. + ``success_rate_stats`` must be zeroed before the launch; ``[0]`` receives the summed per-episode + rates and ``[1]`` the number of finished episodes. + """ env_id = wp.tid() if env_mask[env_id]: + goals = successes[env_id] + wp.atomic_add(success_rate_stats, 0, goals / (goals + wp.float32(1.0))) + wp.atomic_add(success_rate_stats, 1, wp.float32(1.0)) successes[env_id] = wp.float32(0.0) +@wp.kernel +def update_success_rate_from_stats( + # input + success_rate_stats: wp.array(dtype=wp.float32), + # output + success_rate: wp.array(dtype=wp.float32), +): + """Average the accumulated episode success rates; keep the last value when no episode finished.""" + # single-thread kernel (dim=1) + if success_rate_stats[1] > wp.float32(0.0): + success_rate[0] = success_rate_stats[0] / success_rate_stats[1] + + @wp.kernel def compute_intermediate_values( # input @@ -242,11 +248,17 @@ def get_dones( # input/output episode_length_buf: wp.array(dtype=wp.int32), # output + orientation_error: wp.array(dtype=wp.float32), + goal_reached: wp.array(dtype=wp.bool), out_of_reach: wp.array(dtype=wp.bool), time_out: wp.array(dtype=wp.bool), reset: wp.array(dtype=wp.bool), ): - """Flag object-fall and time-out/max-success termination, reset progress on goal, write the reset union.""" + """Evaluate goal success once per step, flag object-fall and time-out/max-success termination. + + The orientation error [rad] and success flags are written for the reward to reuse. Progress is + reset on reaching a goal when a consecutive-success cap is configured. + """ env_id = wp.tid() object_pos = wp.transform_get_translation(object_pose[env_id]) @@ -255,11 +267,14 @@ def get_dones( goal_dist = wp.length(object_pos - in_hand_pos[env_id]) out_of_reach[env_id] = goal_dist >= fall_dist + error = rotation_distance(object_rot, goal_rot[env_id]) + reached = error <= success_tolerance + orientation_error[env_id] = error + goal_reached[env_id] = reached + max_success_reached = False if max_consecutive_success > 0: - # Reset progress (episode length buf) on goal envs if max_consecutive_success > 0 - rot_dist = rotation_distance(object_rot, goal_rot[env_id]) - if wp.abs(rot_dist) <= success_tolerance: + if reached: episode_length_buf[env_id] = 0 max_success_reached = successes[env_id] >= wp.float32(max_consecutive_success) @@ -359,16 +374,15 @@ def compute_full_observations( observations[env_id, offset + 3] = obj_rot[3] offset += 4 - # spatial_vectorf layout: [0:3]=angular, [3:6]=linear - # torch reference order: linear (unscaled) first, then angular (scaled) - observations[env_id, offset + 0] = object_vels[env_id][3] - observations[env_id, offset + 1] = object_vels[env_id][4] - observations[env_id, offset + 2] = object_vels[env_id][5] + # root velocities are laid out [lin_vel, ang_vel]; only the angular part is scaled + observations[env_id, offset + 0] = object_vels[env_id][0] + observations[env_id, offset + 1] = object_vels[env_id][1] + observations[env_id, offset + 2] = object_vels[env_id][2] offset += 3 - observations[env_id, offset + 0] = vel_obs_scale * object_vels[env_id][0] - observations[env_id, offset + 1] = vel_obs_scale * object_vels[env_id][1] - observations[env_id, offset + 2] = vel_obs_scale * object_vels[env_id][2] + observations[env_id, offset + 0] = vel_obs_scale * object_vels[env_id][3] + observations[env_id, offset + 1] = vel_obs_scale * object_vels[env_id][4] + observations[env_id, offset + 2] = vel_obs_scale * object_vels[env_id][5] offset += 3 # goal @@ -438,13 +452,13 @@ def compute_rewards( reset_buf: wp.array(dtype=wp.bool), object_pose: wp.array(dtype=wp.transformf), target_pos: wp.array(dtype=wp.vec3f), - target_rot: wp.array(dtype=wp.quatf), + goal_reached: wp.array(dtype=wp.bool), + orientation_error: wp.array(dtype=wp.float32), dist_reward_scale: wp.float32, rot_reward_scale: wp.float32, rot_eps: wp.float32, actions: wp.array2d(dtype=wp.float32), action_penalty_scale: wp.float32, - success_tolerance: wp.float32, reach_goal_bonus: wp.float32, fall_dist: wp.float32, fall_penalty: wp.float32, @@ -461,13 +475,10 @@ def compute_rewards( env_id = wp.tid() obj_pos = wp.transform_get_translation(object_pose[env_id]) - obj_rot = wp.transform_get_rotation(object_pose[env_id]) - goal_dist = wp.length(obj_pos - target_pos[env_id]) - rot_dist = rotation_distance(obj_rot, target_rot[env_id]) dist_rew = goal_dist * dist_reward_scale - rot_rew = wp.float32(1.0) / (wp.abs(rot_dist) + rot_eps) * rot_reward_scale + rot_rew = rot_reward_scale / (orientation_error[env_id] + rot_eps) action_penalty = wp.float32(0.0) for i in range(action_dim): @@ -476,15 +487,11 @@ def compute_rewards( # Total reward is: position distance + orientation alignment + action regularization + success bonus + fall penalty reward = dist_rew + rot_rew + action_penalty * action_penalty_scale - # Find out which envs hit the goal and update successes count - reached = wp.abs(rot_dist) <= success_tolerance - goal_resets = reached or reset_goal_buf[env_id] + # a goal stays flagged until it is resampled, which happens later this step + goal_resets = goal_reached[env_id] or reset_goal_buf[env_id] reset_goal_buf[env_id] = goal_resets if goal_resets: successes[env_id] = successes[env_id] + wp.float32(1.0) - - # Success bonus: orientation is within `success_tolerance` of goal orientation - if goal_resets: reward = reward + reach_goal_bonus # Fall penalty: distance to the goal is larger than a threshold @@ -535,19 +542,20 @@ def randomize_rotation(rand0: wp.float32, rand1: wp.float32, x_axis: wp.vec3f, y @wp.func def rotation_distance(object_rot: wp.quatf, target_rot: wp.quatf) -> wp.float32: - # Orientation alignment for the cube in hand and goal cube + """Angle [rad] between two ``(x, y, z, w)`` orientations, as :func:`~isaaclab.utils.math.quat_error_magnitude`.""" quat_diff = object_rot * wp.quat_inverse(target_rot) - # Match Torch env convention: uses indices [1:4] for the vector part (see `rotation_distance` in Torch env). - v_norm = wp.sqrt(quat_diff[1] * quat_diff[1] + quat_diff[2] * quat_diff[2] + quat_diff[3] * quat_diff[3]) - v_norm = wp.min(v_norm, wp.float32(1.0)) - return wp.float32(2.0) * wp.asin(v_norm) + v_norm = wp.length(wp.vec3f(quat_diff[0], quat_diff[1], quat_diff[2])) + return wp.float32(2.0) * wp.asin(wp.min(v_norm, wp.float32(1.0))) class ReorientDirectWarpEnv(DirectRLEnvWarp): - cfg: AllegroHandEnvCfg # | ShadowHandEnvCfg + """Warp twin of :class:`~isaaclab_tasks.core.reorient.reorient_direct_env.ReorientDirectEnv`.""" + + cfg: AllegroHandEnvCfg - # def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | None = None, **kwargs): def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwargs): + if cfg.asymmetric_obs: + raise NotImplementedError("The Warp reorientation environment has no asymmetric critic observations.") super().__init__(cfg, render_mode, **kwargs) self.hand, self.object, self.goal_markers = [self.scene[name] for name in ("robot", "object", "goal_object")] @@ -564,13 +572,7 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa actuated_dof_indices.append(self.hand.joint_names.index(joint_name)) actuated_dof_indices.sort() self.num_actuated_dofs = len(actuated_dof_indices) - - # Warp index/mask helpers for kernels and articulation APIs. self.actuated_dof_indices = wp.array(actuated_dof_indices, dtype=wp.int32, device=self.device) - actuated_mask = [False] * self.num_hand_dofs - for idx in actuated_dof_indices: - actuated_mask[idx] = True - self.actuated_dof_mask = wp.array(actuated_mask, dtype=wp.bool, device=self.device) # finger bodies finger_bodies: list[int] = list() @@ -587,7 +589,6 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa # unit vectors self.x_unit_vec = wp.vec3f(1.0, 0.0, 0.0) 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) @@ -597,10 +598,19 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa # --------------------------------------------------------------------- # buffers for position targets - self.hand_dof_targets = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device) self.prev_targets = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device) self.cur_targets = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device) + # reset states sampled per env; the asset writers apply the masked rows + self.reset_joint_pos = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device) + self.reset_joint_vel = wp.zeros((self.num_envs, self.num_hand_dofs), dtype=wp.float32, device=self.device) + self.reset_object_pose = wp.zeros(self.num_envs, dtype=wp.transformf, device=self.device) + self.reset_object_vel = wp.zeros(self.num_envs, dtype=wp.spatial_vectorf, device=self.device) + + # per-step goal evaluation, written in `_get_dones` and reused by the reward + self.orientation_error = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) + self.goal_reached = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + # track goal resets self.reset_goal_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) # used to compare object position @@ -611,12 +621,15 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa self.goal_pos_w = wp.zeros(self.num_envs, dtype=wp.vec3f, device=self.device) # Initialize goal constants from Torch (avoid a one-off kernel launch). - default_root_pose = self.object.data.default_root_pose.torch.to(self.device) - in_hand_pos = default_root_pose[:, 0:3].clone() - in_hand_pos[:, 2] -= 0.04 + in_hand_pos = self.object.data.default_root_pose.torch[:, 0:3] + torch.tensor( + self.cfg.in_hand_pos_offset, device=self.device + ) self.in_hand_pos.assign(wp.from_torch(in_hand_pos, dtype=wp.vec3f)) - goal_pos = torch.tensor([-0.2, -0.45, 0.68], device=self.device, dtype=torch.float32).repeat((self.num_envs, 1)) + if self.cfg.goal_marker_position is None: + goal_pos = in_hand_pos + torch.tensor(self.cfg.goal_marker_offset, device=self.device) + else: + goal_pos = torch.tensor(self.cfg.goal_marker_position, device=self.device).repeat((self.num_envs, 1)) self.goal_pos.assign(wp.from_torch(goal_pos, dtype=wp.vec3f)) goal_rot = torch.zeros((self.num_envs, 4), device=self.device, dtype=torch.float32) @@ -629,6 +642,8 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa # track successes self.successes = wp.zeros(self.num_envs, dtype=wp.float32, device=self.device) self.consecutive_successes = wp.zeros(1, dtype=wp.float32, device=self.device) + self._success_rate_stats = wp.zeros(2, dtype=wp.float32, device=self.device) + self.success_rate = wp.zeros(1, dtype=wp.float32, device=self.device) # Persistent RL buffers (Warp). self.actions = wp.zeros((self.num_envs, self.cfg.action_space), dtype=wp.float32, device=self.device) @@ -647,20 +662,6 @@ def __init__(self, cfg: AllegroHandEnvCfg, render_mode: str | None = None, **kwa self.object_pose = wp.zeros(self.num_envs, dtype=wp.transformf, device=self.device) self.object_vels = wp.zeros(self.num_envs, dtype=wp.spatial_vectorf, device=self.device) - # RNG state (per-env) for randomizations in reset/goal resets. - self.rng_state = wp.zeros(self.num_envs, dtype=wp.uint32, device=self.device) - if self.cfg.seed is None: - self.cfg.seed = -1 - wp.launch( - initialize_rng_state, - dim=self.num_envs, - inputs=[ - self.cfg.seed, - self.rng_state, - ], - device=self.device, - ) - # --------------------------------------------------------------------- # Torch views / aliases # --------------------------------------------------------------------- @@ -676,6 +677,7 @@ def _pre_physics_step(self, actions: wp.array) -> None: # Store actions in a persistent Warp buffer (analogous to `actions.clone()` in the Torch env). wp.copy(self.actions, actions) + @captured def _apply_action(self) -> None: wp.launch( apply_actions_to_targets, @@ -692,17 +694,11 @@ def _apply_action(self) -> None: device=self.device, ) - # Apply position targets using mask method (CUDA graph safe). - # All joints are actuated for Allegro, so default masks (None = all) are correct. - self.hand.set_joint_position_target_mask(target=self.cur_targets) + # unactuated joints keep the targets their last reset wrote + self.hand.actuators.target_command.set_position_mask(value=self.cur_targets) + @captured def _get_observations(self) -> dict: - # if self.cfg.asymmetric_obs: - # self.fingertip_force_sensors = self.hand.root_physx_view.get_link_incoming_joint_force()[ - # :, self.finger_bodies - # ] - # NOTE: if re-enabled, `self.finger_bodies` holds public-order indices (from `body_names.index`) - # and must be translated to backend view order before indexing this `root_physx_view` array. if self.cfg.obs_type == "openai": self.compute_reduced_observations() elif self.cfg.obs_type == "full": @@ -711,10 +707,9 @@ def _get_observations(self) -> dict: raise ValueError(f"Unknown obs_type: {self.cfg.obs_type}") return {"policy": self.torch_obs_buf} + @captured def _get_rewards(self) -> None: # Clear reduction buffers before launching the reward kernel. - # wp.assign(self._num_resets, 0.0) - # wp.assign(self._finished_cons_successes, 0.0) self._num_resets.zero_() self._finished_cons_successes.zero_() wp.launch( @@ -724,13 +719,13 @@ def _get_rewards(self) -> None: self.reset_buf, self.object_pose, self.in_hand_pos, - self.goal_rot, + self.goal_reached, + self.orientation_error, self.cfg.dist_reward_scale, self.cfg.rot_reward_scale, self.cfg.rot_eps, self.actions, self.cfg.action_penalty_scale, - self.cfg.success_tolerance, self.cfg.reach_goal_bonus, self.cfg.fall_dist, self.cfg.fall_penalty, @@ -762,10 +757,9 @@ def _get_rewards(self) -> None: # .mean() cannot be called here as it causes problems on stream self.extras["log"]["consecutive_successes"] = wp.to_torch(self.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. - self._reset_target_pose(mask=self.reset_goal_buf) + self._reset_target_pose(self.reset_goal_buf) + @captured def _get_dones(self) -> None: self._compute_intermediate_values() @@ -782,6 +776,8 @@ def _get_dones(self) -> None: self.cfg.max_consecutive_success, self.successes, self._episode_length_buf_wp, + self.orientation_error, + self.goal_reached, self.reset_terminated, self.reset_time_outs, self.reset_buf, @@ -789,6 +785,7 @@ def _get_dones(self) -> None: device=self.device, ) + @captured def _reset_idx(self, mask: wp.array | None = None): if mask is None: mask = self._ALL_ENV_MASK @@ -796,10 +793,8 @@ def _reset_idx(self, mask: wp.array | None = None): # resets articulation and rigid body attributes super()._reset_idx(mask) - # reset goals - self._reset_target_pose(mask=mask) + self._reset_target_pose(mask) - # reset object wp.launch( reset_object, dim=self.num_envs, @@ -810,14 +805,15 @@ def _reset_idx(self, mask: wp.array | None = None): self.x_unit_vec, self.y_unit_vec, mask, - self.rng_state, - self.object.data.root_link_pose_w.warp, - self.object.data.root_com_vel_w.warp, + WarpRng.state, + self.reset_object_pose, + self.reset_object_vel, ], device=self.device, ) + self.object.write_root_pose_to_sim_mask(root_pose=self.reset_object_pose, env_mask=mask) + self.object.write_root_velocity_to_sim_mask(root_velocity=self.reset_object_vel, env_mask=mask) - # reset hand wp.launch( reset_hand, dim=self.num_envs, @@ -830,41 +826,36 @@ def _reset_idx(self, mask: wp.array | None = None): self.cfg.reset_dof_vel_noise, mask, self.num_hand_dofs, - self.rng_state, - self.hand.data.joint_pos.warp, - self.hand.data.joint_vel.warp, + WarpRng.state, + self.reset_joint_pos, + self.reset_joint_vel, self.prev_targets, self.cur_targets, - self.hand_dof_targets, ], device=self.device, ) + self.hand.actuators.target_command.set_position_mask(value=self.cur_targets, env_mask=mask) + self.hand.write_joint_position_to_sim_mask(position=self.reset_joint_pos, env_mask=mask) + self.hand.write_joint_velocity_to_sim_mask(velocity=self.reset_joint_vel, env_mask=mask) - self.hand.set_joint_position_target_mask(target=self.cur_targets, env_mask=mask) - + self._success_rate_stats.zero_() wp.launch( reset_successes, dim=self.num_envs, - inputs=[ - mask, - self.successes, - ], + inputs=[mask, self.successes, self._success_rate_stats], + device=self.device, + ) + wp.launch( + update_success_rate_from_stats, + dim=1, + inputs=[self._success_rate_stats, self.success_rate], device=self.device, ) + self.extras.setdefault("log", {})["Metrics/success_rate"] = wp.to_torch(self.success_rate) self._compute_intermediate_values() - def _reset_target_pose(self, env_ids: Sequence[int] | None = None, mask: wp.array | None = None): - # reset goal rotation - if mask is None: - if env_ids is None: - return - env_mask_list = [False] * self.num_envs - for env_id in env_ids: - env_mask_list[int(env_id)] = True - mask = wp.array(env_mask_list, dtype=wp.bool, device=self.device) - - # update goal pose and markers + def _reset_target_pose(self, mask: wp.array): wp.launch( reset_target_pose, dim=self.num_envs, @@ -874,7 +865,7 @@ def _reset_target_pose(self, env_ids: Sequence[int] | None = None, mask: wp.arra self.y_unit_vec, self.env_origins, self.goal_pos, - self.rng_state, + WarpRng.state, self.goal_rot, self.reset_goal_buf, self.goal_pos_w,