From f6364827e5fc00bfa17b6a01321445ae66622b0f Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 01:39:55 -0700 Subject: [PATCH 1/5] Add Warp launch replay cache Provide an owner-scoped path for pointer-stable kernels that cannot stay inside CUDA graph regions. Preserve eager A/B mode and validate recorded storage contracts in debug builds. --- .../warp-launch-record-replay.minor.rst | 5 + .../isaaclab/isaaclab/utils/warp/__init__.pyi | 2 + .../isaaclab/utils/warp/warp_launch_cache.py | 343 ++++++++++++++++++ .../test/utils/warp/test_warp_launch_cache.py | 254 +++++++++++++ 4 files changed, 604 insertions(+) create mode 100644 source/isaaclab/changelog.d/warp-launch-record-replay.minor.rst create mode 100644 source/isaaclab/isaaclab/utils/warp/warp_launch_cache.py create mode 100644 source/isaaclab/test/utils/warp/test_warp_launch_cache.py diff --git a/source/isaaclab/changelog.d/warp-launch-record-replay.minor.rst b/source/isaaclab/changelog.d/warp-launch-record-replay.minor.rst new file mode 100644 index 000000000000..2ca617fdb126 --- /dev/null +++ b/source/isaaclab/changelog.d/warp-launch-record-replay.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added :class:`~isaaclab.utils.warp.WarpLaunchCache` for low-overhead replay + of pointer-stable Warp kernel launches outside CUDA graphs. diff --git a/source/isaaclab/isaaclab/utils/warp/__init__.pyi b/source/isaaclab/isaaclab/utils/warp/__init__.pyi index 7be5ce3ed1c8..fb401ac1a128 100644 --- a/source/isaaclab/isaaclab/utils/warp/__init__.pyi +++ b/source/isaaclab/isaaclab/utils/warp/__init__.pyi @@ -6,6 +6,7 @@ __all__ = [ "ParticleMeshCounter", "ProxyArray", + "WarpLaunchCache", "convert_to_warp_mesh", "make_box_region_mesh", "make_frustum_region_mesh", @@ -17,3 +18,4 @@ __all__ = [ from .ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh, raycast_single_mesh from .particle_mesh import ParticleMeshCounter, make_box_region_mesh, make_frustum_region_mesh from .proxy_array import ProxyArray +from .warp_launch_cache import WarpLaunchCache diff --git a/source/isaaclab/isaaclab/utils/warp/warp_launch_cache.py b/source/isaaclab/isaaclab/utils/warp/warp_launch_cache.py new file mode 100644 index 000000000000..1d5cb3eedfc5 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/warp/warp_launch_cache.py @@ -0,0 +1,343 @@ +# 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 + +"""Cached Warp kernel launches for low-overhead record-and-replay execution.""" + +from __future__ import annotations + +import operator +import os +import struct +from collections.abc import Hashable, Sequence +from dataclasses import dataclass +from typing import Any + +import warp as wp + +_WARP_LAUNCH_MODE_ENV = "ISAACLAB_WARP_LAUNCH_MODE" +_WARP_LAUNCH_DEBUG_ENV = "ISAACLAB_WARP_LAUNCH_DEBUG" + + +class WarpLaunchCache: + """Launch Warp kernels eagerly or replay a command recorded on first use. + + This utility is the per-kernel counterpart to a CUDA graph cache. Each + ``(kernel, site)`` pair is recorded independently, so pointer-stable kernels + can replay even when neighboring work cannot be captured or recorded. The + launch interface contains only the arguments used by persistent Isaac Lab + kernels. + + Recorded arguments are static. Their storage, layout, and scalar values must + remain unchanged for the lifetime of the cached command. Use a distinct + ``site`` when one owner invokes the same kernel with multiple persistent + argument sets. Leave launches with changing argument storage on + :func:`warp.launch`. + + The launch dimension is checked before replay. Debug mode additionally checks + the static-argument contract on every replay. It is meant for tests and + validation; normal replay only performs the command lookup, an optional + dimension update, and :meth:`warp.Launch.launch`. + + Recorded launches may be replayed inside CUDA graph capture. CUDA graph + nodes retain their own launch data, so discarding the Python command is + safe after outstanding work is drained. Do not replay an old graph after + replacing the recorded argument storage it references. + + Args: + mode: ``"replay"`` to cache commands or ``"eager"`` to call + :func:`warp.launch` directly. If omitted, reads + ``ISAACLAB_WARP_LAUNCH_MODE`` and defaults to ``"replay"``. + debug: Whether to validate recorded arguments on every replay. If + omitted, reads ``ISAACLAB_WARP_LAUNCH_DEBUG`` and defaults to + ``False``. + device: Fixed device for the cache owner. + enabled: Whether replay is enabled. A disabled cache always launches + eagerly. + """ + + @dataclass(frozen=True) + class _ArrayDescriptor: + """Pointer and layout properties of one array argument.""" + + ptr: int + grad_ptr: int + device: str + dtype: str + shape: tuple[int, ...] + strides: tuple[int, ...] + + @dataclass + class _Entry: + """One recorded command and its static launch contract.""" + + command: wp.Launch + site: Hashable | None + dim: tuple[int, ...] + argument_counts: tuple[int, int] + argument_tokens: tuple[tuple[Any, ...], ...] | None + argument_owners: tuple[Any, ...] + + def __init__( + self, + mode: str | None = None, + debug: bool | None = None, + *, + device: str | wp.Device, + enabled: bool = True, + ): + if not enabled: + mode = "eager" + elif mode is None: + mode = os.environ.get(_WARP_LAUNCH_MODE_ENV, "replay") + mode = mode.lower() + if mode not in {"eager", "replay"}: + raise ValueError(f"Invalid Warp launch mode {mode!r}. Set {_WARP_LAUNCH_MODE_ENV} to 'eager' or 'replay'.") + + if debug is None: + debug = self._read_debug_environment() + + self._mode = mode + self._debug = debug + self._device = wp.get_device(device) + self._entries: dict[Hashable, WarpLaunchCache._Entry] = {} + + def launch( + self, + kernel: wp.Kernel, + dim: int | Sequence[int], + inputs: Sequence[Any] = (), + outputs: Sequence[Any] = (), + *, + site: Hashable | None = None, + ) -> None: + """Launch a Warp kernel eagerly or through a cached command. + + Args: + kernel: Warp kernel to launch. + dim: Number or shape of threads to launch. A zero-sized launch is a + no-op and is not cached. + inputs: Forward input arguments. + outputs: Forward output arguments. + site: Optional stable identity for one persistent argument set. + + Raises: + TypeError: If ``site`` is not hashable. + ValueError: If the launch dimension is invalid. + RuntimeError: In debug mode, if a recorded static argument changes. + """ + if self._mode == "eager": + wp.launch( + kernel, + dim=dim, + inputs=inputs, + outputs=outputs, + device=self._device, + ) + return + + logical_key: Hashable = kernel if site is None else (kernel, site) + try: + entry = self._entries.get(logical_key) + except TypeError as exc: + raise TypeError(f"Warp launch site must be hashable, got {type(site).__name__}.") from exc + + if entry is not None: + if isinstance(dim, int): + dim_matches = len(entry.dim) == 1 and dim == entry.dim[0] + else: + dim_matches = isinstance(dim, tuple) and dim == entry.dim + if not dim_matches: + normalized_dim, total_dim = self._normalize_dim(dim) + if total_dim == 0: + return + if normalized_dim != entry.dim: + entry.command.set_dim(normalized_dim) + entry.dim = normalized_dim + if self._debug: + self._validate_entry(entry, inputs, outputs) + entry.command.launch() + return + + normalized_dim, total_dim = self._normalize_dim(dim) + if total_dim == 0: + return + command = wp.launch( + kernel, + dim=normalized_dim, + inputs=inputs, + outputs=outputs, + device=self._device, + record_cmd=True, + ) + if command is None: + raise RuntimeError(f"Warp did not return a Launch while recording {logical_key!r} on {self._device}.") + + arguments = (*inputs, *outputs) + entry = self._Entry( + command=command, + site=site, + dim=normalized_dim, + argument_counts=(len(inputs), len(outputs)), + argument_tokens=self._argument_tokens(arguments) if self._debug else None, + argument_owners=arguments, + ) + self._entries[logical_key] = entry + command.launch() + + __call__ = launch + + def invalidate(self, site: Hashable | None = None) -> None: + """Discard recorded commands. + + Args: + site: Optional site to invalidate across kernels. If omitted, every + command is discarded. + + Outstanding work on the owner device is synchronized before the cached + commands' retained argument owners are released. An empty cache does not + synchronize. + """ + if site is None: + self.reset() + return + stale_keys = [key for key, entry in self._entries.items() if self._sites_match(entry.site, site)] + if stale_keys: + wp.synchronize_device(self._device) + for key in stale_keys: + del self._entries[key] + + def reset(self) -> None: + """Drain dependent work and discard every recorded command. + + The owner device is synchronized only when the cache contains recorded + commands. Disabled and otherwise empty caches return immediately. + """ + if self._entries: + wp.synchronize_device(self._device) + self._entries.clear() + + @staticmethod + def _sites_match(recorded: Hashable | None, requested: Hashable) -> bool: + """Compare site keys without forcing array-like equality to a boolean.""" + if recorded is requested: + return True + try: + if hash(recorded) != hash(requested): + return False + return bool(recorded == requested) + except (RuntimeError, TypeError, ValueError): + return False + + def _validate_entry( + self, + entry: WarpLaunchCache._Entry, + inputs: Sequence[Any], + outputs: Sequence[Any], + ) -> None: + """Validate the static command contract in debug mode.""" + counts = (len(inputs), len(outputs)) + if counts != entry.argument_counts: + raise RuntimeError( + f"Warp launch argument counts changed: recorded {entry.argument_counts}, received {counts}. " + "Use a distinct site or invalidate the cache." + ) + arguments = (*inputs, *outputs) + tokens = self._argument_tokens(arguments) + if tokens != entry.argument_tokens: + mismatch = next( + index for index, pair in enumerate(zip(entry.argument_tokens or (), tokens)) if pair[0] != pair[1] + ) + raise RuntimeError( + f"Warp launch static argument {mismatch} changed. Recorded commands require persistent argument " + "storage and values; use a distinct site, invalidate the cache, or leave this launch eager." + ) + + @staticmethod + def _normalize_dim(dim: int | Sequence[int]) -> tuple[tuple[int, ...], int]: + """Normalize Warp launch dimensions and compute their product.""" + if isinstance(dim, int): + normalized = (dim,) + else: + try: + normalized = tuple(operator.index(extent) for extent in dim) + except TypeError: + normalized = (operator.index(dim),) + + if len(normalized) > 4: + raise ValueError(f"Warp launch dimensions must have at most 4 axes, got {len(normalized)}.") + + total = 1 + for extent in normalized: + if extent < 0: + raise ValueError(f"Warp launch dimensions must be non-negative, got {normalized}.") + total *= extent + return normalized, total + + @classmethod + def _argument_tokens(cls, arguments: Sequence[Any]) -> tuple[tuple[Any, ...], ...]: + """Snapshot static arguments for debug validation.""" + return tuple(cls._argument_token(value) for value in arguments) + + @classmethod + def _argument_token(cls, value: Any) -> tuple[Any, ...]: + """Return a stable pointer, layout, or value token for one argument.""" + descriptor = cls._array_descriptor_or_none(value) + if descriptor is not None: + return ("array", descriptor) + if type(value) is float: + return ("float", struct.pack("=d", value)) + if type(value) is complex: + return ("complex", struct.pack("=dd", value.real, value.imag)) + if value is None or isinstance(value, (bool, int, str, bytes)): + return ("value", type(value), value) + if isinstance(value, tuple): + return ("tuple", type(value), tuple(cls._argument_token(item) for item in value)) + if hasattr(type(value), "_wp_scalar_type_"): + try: + return ("warp_value", type(value), tuple(value)) + except TypeError: + pass + if hasattr(value, "value"): + return ("value_attribute", type(value), value.value) + return ("object", type(value), id(value)) + + @classmethod + def _array_descriptor_or_none(cls, value: Any) -> WarpLaunchCache._ArrayDescriptor | None: + """Return a descriptor for a Warp array, ProxyArray, or Torch tensor.""" + array = value if isinstance(value, wp.array) else getattr(value, "warp", None) + if isinstance(array, wp.array): + grad = array.grad + return cls._ArrayDescriptor( + ptr=0 if array.ptr is None else int(array.ptr), + grad_ptr=0 if grad is None or grad.ptr is None else int(grad.ptr), + device=str(array.device), + dtype=str(array.dtype), + shape=tuple(int(extent) for extent in array.shape), + strides=tuple(int(stride) for stride in array.strides), + ) + + data_ptr = getattr(value, "data_ptr", None) + stride = getattr(value, "stride", None) + element_size = getattr(value, "element_size", None) + if callable(data_ptr) and callable(stride) and callable(element_size) and hasattr(value, "shape"): + return cls._ArrayDescriptor( + ptr=int(data_ptr()), + grad_ptr=0, + device=str(getattr(value, "device", "cpu")), + dtype=str(getattr(value, "dtype", type(value))), + shape=tuple(int(extent) for extent in value.shape), + strides=tuple(int(value_stride) * int(element_size()) for value_stride in stride()), + ) + return None + + @staticmethod + def _read_debug_environment() -> bool: + """Read the boolean debug switch from the environment.""" + value = os.environ.get(_WARP_LAUNCH_DEBUG_ENV, "0").lower() + if value in {"0", "false", "no", "off"}: + return False + if value in {"1", "true", "yes", "on"}: + return True + raise ValueError(f"Invalid {_WARP_LAUNCH_DEBUG_ENV} value {value!r}; use 0/1, false/true, no/yes, or off/on.") diff --git a/source/isaaclab/test/utils/warp/test_warp_launch_cache.py b/source/isaaclab/test/utils/warp/test_warp_launch_cache.py new file mode 100644 index 000000000000..2b7fd1230850 --- /dev/null +++ b/source/isaaclab/test/utils/warp/test_warp_launch_cache.py @@ -0,0 +1,254 @@ +# 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 WarpLaunchCache record-and-replay behavior.""" + +from __future__ import annotations + +import os +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import torch +import warp as wp + +from isaaclab.utils.warp import ProxyArray, WarpLaunchCache + + +@wp.kernel +def _affine( + source: wp.array(dtype=wp.float32), + scale: float, + output: wp.array(dtype=wp.float32), +): + """Apply a scalar multiplier to an array.""" + index = wp.tid() + output[index] = source[index] * scale + + +@wp.kernel +def _increment(output: wp.array(dtype=wp.float32)): + """Increment every launched output element.""" + index = wp.tid() + output[index] = output[index] + 1.0 + + +@wp.kernel +def _write_index(output: wp.array(dtype=wp.float32)): + """Write one-based thread indices to an output array.""" + index = wp.tid() + output[index] = wp.float32(index + 1) + + +class TestWarpLaunchCache(unittest.TestCase): + """Tests for :class:`WarpLaunchCache`.""" + + @classmethod + def setUpClass(cls): + wp.init() + if not wp.is_cuda_available(): + raise unittest.SkipTest("CUDA is required for Warp launch replay tests.") + + def setUp(self): + self.device = wp.get_device("cuda:0") + self.cache = WarpLaunchCache(device=self.device) + + def tearDown(self): + wp.synchronize_device(self.device) + + def test_static_launch_records_once_then_replays(self): + """The first call should record and execute, then later calls should replay.""" + output = wp.zeros(8, dtype=wp.float32, device=self.device) + + self.cache.launch(_increment, dim=8, inputs=[output]) + self.cache.launch(_increment, dim=8, inputs=[output]) + + self.assertEqual(len(self.cache._entries), 1) + np.testing.assert_array_equal(output.numpy(), np.full(8, 2.0, dtype=np.float32)) + + def test_fixed_device_cache_is_callable_from_its_owner(self): + """An owner should retain and call one fixed-device launcher.""" + owner = SimpleNamespace(_warp_launch=WarpLaunchCache(device=self.device)) + output = wp.zeros(4, dtype=wp.float32, device=self.device) + + owner._warp_launch(_increment, dim=4, inputs=[output]) + owner._warp_launch(_increment, dim=4, inputs=[output]) + + np.testing.assert_array_equal(output.numpy(), np.full(4, 2.0, dtype=np.float32)) + + def test_sites_distinguish_persistent_argument_sets(self): + """One kernel should own independent commands for distinct stable sites.""" + output_a = wp.zeros(4, dtype=wp.float32, device=self.device) + output_b = wp.zeros(4, dtype=wp.float32, device=self.device) + + self.cache(_increment, dim=4, inputs=[output_a], site=output_a) + self.cache(_increment, dim=4, inputs=[output_b], site=output_b) + self.cache(_increment, dim=4, inputs=[output_a], site=output_a) + self.cache(_increment, dim=4, inputs=[output_b], site=output_b) + + self.assertEqual(len(self.cache._entries), 2) + np.testing.assert_array_equal(output_a.numpy(), np.full(4, 2.0, dtype=np.float32)) + np.testing.assert_array_equal(output_b.numpy(), np.full(4, 2.0, dtype=np.float32)) + + def test_dim_change_updates_recorded_bounds(self): + """A changed launch dimension should update the cached command bounds.""" + output = wp.zeros(8, dtype=wp.float32, device=self.device) + + self.cache(_write_index, dim=4, inputs=[output]) + self.cache(_write_index, dim=8, inputs=[output]) + + np.testing.assert_array_equal(output.numpy(), np.arange(1, 9, dtype=np.float32)) + + def test_eager_environment_mode_uses_current_arguments(self): + """The eager switch should bypass recording and use each call's arguments.""" + source = wp.full(4, value=2.0, dtype=wp.float32, device=self.device) + output = wp.zeros(4, dtype=wp.float32, device=self.device) + + with patch.dict(os.environ, {"ISAACLAB_WARP_LAUNCH_MODE": "eager"}): + cache = WarpLaunchCache(device=self.device) + cache(_affine, dim=4, inputs=[source, 2.0], outputs=[output]) + cache(_affine, dim=4, inputs=[source, 5.0], outputs=[output]) + + self.assertFalse(cache._entries) + np.testing.assert_array_equal(output.numpy(), np.full(4, 10.0, dtype=np.float32)) + + def test_disabled_cache_launches_eagerly(self): + """A disabled owner cache should remain a direct-launch fallback.""" + cache = WarpLaunchCache(device=self.device, enabled=False) + source = wp.full(4, value=2.0, dtype=wp.float32, device=self.device) + output = wp.zeros(4, dtype=wp.float32, device=self.device) + + cache(_affine, dim=4, inputs=[source, 2.0], outputs=[output]) + cache(_affine, dim=4, inputs=[source, 3.0], outputs=[output]) + + self.assertFalse(cache._entries) + np.testing.assert_array_equal(output.numpy(), np.full(4, 6.0, dtype=np.float32)) + + def test_debug_mode_rejects_changed_static_array(self): + """Debug mode should detect replacement of persistent array storage.""" + source_a = wp.ones(4, dtype=wp.float32, device=self.device) + source_b = wp.ones(4, dtype=wp.float32, device=self.device) + output = wp.zeros(4, dtype=wp.float32, device=self.device) + cache = WarpLaunchCache(device=self.device, debug=True) + + cache(_affine, dim=4, inputs=[source_a, 1.0], outputs=[output]) + with self.assertRaisesRegex(RuntimeError, "static argument 0 changed"): + cache(_affine, dim=4, inputs=[source_b, 1.0], outputs=[output]) + + def test_debug_mode_rejects_repointed_torch_tensor(self): + """Debug mode should detect a tensor object repointed to new storage.""" + source = torch.ones(4, dtype=torch.float32, device=str(self.device)) + replacement = torch.full_like(source, 3.0) + output = wp.zeros(4, dtype=wp.float32, device=self.device) + cache = WarpLaunchCache(device=self.device, debug=True) + + cache(_affine, dim=4, inputs=[source, 2.0], outputs=[output]) + source.set_(replacement) + + with self.assertRaisesRegex(RuntimeError, "static argument 0 changed"): + cache(_affine, dim=4, inputs=[source, 2.0], outputs=[output]) + + def test_debug_mode_rejects_changed_static_scalar(self): + """Debug mode should reject a changed recorded scalar value.""" + source = wp.ones(4, dtype=wp.float32, device=self.device) + output = wp.zeros(4, dtype=wp.float32, device=self.device) + cache = WarpLaunchCache(device=self.device, debug=True) + + cache(_affine, dim=4, inputs=[source, 1.0], outputs=[output]) + with self.assertRaisesRegex(RuntimeError, "static argument 1 changed"): + cache(_affine, dim=4, inputs=[source, 2.0], outputs=[output]) + + def test_debug_mode_accepts_new_proxy_for_same_array(self): + """Debug mode should compare ProxyArray storage instead of wrapper identity.""" + source = wp.ones(4, dtype=wp.float32, device=self.device) + output = wp.zeros(4, dtype=wp.float32, device=self.device) + cache = WarpLaunchCache(device=self.device, debug=True) + + cache(_affine, dim=4, inputs=[ProxyArray(source), 2.0], outputs=[output]) + cache(_affine, dim=4, inputs=[ProxyArray(source), 2.0], outputs=[output]) + + np.testing.assert_array_equal(output.numpy(), np.full(4, 2.0, dtype=np.float32)) + + def test_zero_dim_is_noop_and_does_not_populate_cache(self): + """A zero-sized launch should neither execute nor reserve a command.""" + output = wp.zeros(4, dtype=wp.float32, device=self.device) + + self.cache(_increment, dim=(4, 0), inputs=[output]) + + self.assertFalse(self.cache._entries) + np.testing.assert_array_equal(output.numpy(), np.zeros(4, dtype=np.float32)) + + def test_fixed_device_fast_path_skips_device_lookup(self): + """A fixed-device cache should not resolve its device after construction.""" + output = wp.zeros(4, dtype=wp.float32, device=self.device) + self.cache(_increment, dim=4, inputs=[output]) + + with patch.object(wp, "get_device", side_effect=AssertionError("unexpected device lookup")): + self.cache(_increment, dim=4, inputs=[output]) + + np.testing.assert_array_equal(output.numpy(), np.full(4, 2.0, dtype=np.float32)) + + def test_invalidate_site_and_reset(self): + """Invalidation should drop one stable site or every recorded command.""" + output_a = wp.zeros(1, dtype=wp.float32, device=self.device) + output_b = wp.zeros(1, dtype=wp.float32, device=self.device) + self.cache(_increment, dim=1, inputs=[output_a], site="a") + self.cache(_increment, dim=1, inputs=[output_b], site="b") + + self.cache.invalidate(site="a") + self.assertEqual(len(self.cache._entries), 1) + self.cache.reset() + self.assertFalse(self.cache._entries) + + def test_invalidate_accepts_tensor_site_identity(self): + """Invalidation should accept a tensor site without evaluating elementwise equality.""" + output = wp.zeros(1, dtype=wp.float32, device=self.device) + site = torch.arange(2, device=str(self.device)) + self.cache(_increment, dim=1, inputs=[output], site=site) + + self.cache.invalidate(site=site) + + self.assertFalse(self.cache._entries) + + def test_reset_synchronizes_only_nonempty_cache_device(self): + """Reset should drain recorded commands without syncing an empty cache.""" + empty_cache = WarpLaunchCache(device=self.device, enabled=False) + with patch.object(wp, "synchronize_device", wraps=wp.synchronize_device) as synchronize: + empty_cache.reset() + synchronize.assert_not_called() + + output = wp.zeros(1, dtype=wp.float32, device=self.device) + self.cache(_increment, dim=1, inputs=[output]) + self.cache.reset() + synchronize.assert_called_once_with(self.device) + + def test_replay_can_be_recorded_inside_cuda_capture(self): + """A warmed command replay should compose with CUDA graph capture.""" + output = wp.zeros(4, dtype=wp.float32, device=self.device) + self.cache(_increment, dim=4, inputs=[output]) + + with wp.ScopedCapture(device=self.device) as capture: + self.cache(_increment, dim=4, inputs=[output]) + + wp.capture_launch(capture.graph) + np.testing.assert_array_equal(output.numpy(), np.full(4, 2.0, dtype=np.float32)) + + def test_captured_graph_survives_launch_cache_reset(self): + """A graph should retain a captured replay after its Python command is reset.""" + output = wp.zeros(4, dtype=wp.float32, device=self.device) + self.cache(_increment, dim=4, inputs=[output]) + + with wp.ScopedCapture(device=self.device) as capture: + self.cache(_increment, dim=4, inputs=[output]) + + self.cache.reset() + wp.capture_launch(capture.graph) + np.testing.assert_array_equal(output.numpy(), np.full(4, 2.0, dtype=np.float32)) + + +if __name__ == "__main__": + unittest.main() From 24040adee2907c85347ba7a0f1da52e801c13a84 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 01:40:05 -0700 Subject: [PATCH 2/5] Use replay in Warp frontend hot paths Route persistent direct-task, MDP, and manager launches through environment-owned caches so uncapturable stages avoid eager marshalling on every kernel. Keep transient allocations eager. --- .../warp-launch-record-replay.minor.rst | 11 + .../envs/direct_rl_env_warp.py | 25 +- .../envs/manager_based_env_warp.py | 30 +++ .../envs/manager_based_rl_env_warp.py | 41 ++- .../envs/mdp/actions/joint_actions.py | 12 +- .../isaaclab_experimental/envs/mdp/events.py | 196 ++++++++------ .../envs/mdp/observations.py | 69 ++--- .../isaaclab_experimental/envs/mdp/rewards.py | 123 ++++----- .../envs/mdp/terminations.py | 24 +- .../managers/action_manager.py | 12 +- .../managers/command_manager.py | 127 +++++++-- .../managers/event_manager.py | 30 ++- .../managers/observation_manager.py | 12 +- .../managers/reward_manager.py | 31 ++- .../managers/termination_manager.py | 15 +- .../utils/manager_call_switch.py | 16 +- .../utils/warp_graph_cache.py | 63 ++++- .../test/envs/mdp/parity_helpers.py | 4 + .../test/envs/mdp/test_actions_warp_parity.py | 39 +++ .../test/envs/mdp/test_events_warp_parity.py | 33 +-- .../envs/mdp/test_observations_warp_parity.py | 23 ++ .../test/envs/mdp/test_rewards_warp_parity.py | 28 ++ .../envs/mdp/test_terminations_warp_parity.py | 15 ++ .../envs/test_manager_based_rl_env_warp.py | 250 ++++++++++++++++++ .../test/utils/test_manager_call_switch.py | 37 ++- .../test/utils/test_warp_graph_cache.py | 51 ++++ .../warp-launch-record-replay.minor.rst | 6 + .../direct/cartpole/cartpole_warp_env.py | 12 +- .../inhand_manipulation_warp_env.py | 49 ++-- .../direct/locomotion/locomotion_env_warp.py | 22 +- .../classic/cartpole/mdp/rewards.py | 8 +- .../classic/humanoid/mdp/observations.py | 24 +- .../classic/humanoid/mdp/rewards.py | 36 +-- .../locomotion/velocity/mdp/rewards.py | 85 ++---- .../locomotion/velocity/mdp/terminations.py | 37 ++- .../manipulation/reach/mdp/rewards.py | 44 ++- .../test/test_cartpole_warp_rewards.py | 30 +++ 37 files changed, 1155 insertions(+), 515 deletions(-) create mode 100644 source/isaaclab_experimental/changelog.d/warp-launch-record-replay.minor.rst create mode 100644 source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py create mode 100644 source/isaaclab_tasks_experimental/changelog.d/warp-launch-record-replay.minor.rst create mode 100644 source/isaaclab_tasks_experimental/test/test_cartpole_warp_rewards.py diff --git a/source/isaaclab_experimental/changelog.d/warp-launch-record-replay.minor.rst b/source/isaaclab_experimental/changelog.d/warp-launch-record-replay.minor.rst new file mode 100644 index 000000000000..15ab0ba9862d --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/warp-launch-record-replay.minor.rst @@ -0,0 +1,11 @@ +Added +^^^^^ + +* Added recorded-launch replay to reusable Warp manager and MDP kernels, with + eager and replay modes for performance comparison. + +Fixed +^^^^^ + +* Fixed Warp manager-based locomotion environments to expose persistent Warp + views of stable command-manager outputs. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py index d69b13a8feec..9e0b76e0eacd 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 @@ -31,11 +31,13 @@ # from isaaclab.envs.ui import ViewportCameraController from isaaclab.managers import EventManager +from isaaclab.physics import PhysicsEvent from isaaclab.sim import SimulationContext from isaaclab.sim.utils import use_stage from isaaclab.utils.noise import NoiseModel from isaaclab.utils.seed import configure_seed from isaaclab.utils.timer import Timer +from isaaclab.utils.warp import WarpLaunchCache from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache @@ -245,6 +247,14 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # Warp CUDA graph cache for capture-or-replay self._graph_cache = WarpGraphCache() + # Warp launch cache for low-overhead execution outside captured stages + self._warp_launch = WarpLaunchCache(device=self.device) + self._warp_cache_physics_ready_handle = self.sim.physics_manager.register_callback( + self._reset_warp_caches_after_physics_ready, + PhysicsEvent.PHYSICS_READY, + order=30, + name="direct_rl_env_warp_cache_rebind", + ) # setup the action and observation spaces for Gym self._configure_gym_env_spaces() @@ -486,7 +496,7 @@ def step_warp_action(self) -> None: def _step_warp_end_pre(self) -> None: """Capturable portion before write_data_to_sim (pure warp kernels).""" - wp.launch( + self._warp_launch.launch( add_to_env, dim=self.num_envs, inputs=[ @@ -614,6 +624,11 @@ def render(self, recompute: bool = False) -> np.ndarray | None: def close(self): """Cleanup for the environment.""" if not self._is_closed: + if self._warp_cache_physics_ready_handle is not None: + self._warp_cache_physics_ready_handle.deregister() + self._warp_cache_physics_ready_handle = None + self._reset_warp_caches_after_physics_ready(None) + # close entities related to the environment # note: this is order-sensitive to avoid any dangling references if self.cfg.events: @@ -639,6 +654,11 @@ def close(self): # update closing status self._is_closed = True + def _reset_warp_caches_after_physics_ready(self, _payload: Any) -> None: + """Drop graphs and launches recorded against previous simulation storage.""" + self._graph_cache.invalidate() + self._warp_launch.reset() + """ Operations - Debug Visualization. """ @@ -740,13 +760,14 @@ def _reset_idx(self, mask: wp.array | None = None): # self._observation_noise_model.reset(env_ids) # reset the episode length buffer - wp.launch( + self._warp_launch.launch( zero_mask_int32, dim=self.num_envs, inputs=[ mask, self._episode_length_buf_wp, ], + site=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 2b4791d22a16..b9335b7e905e 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 @@ -30,12 +30,14 @@ from isaaclab.envs.manager_based_env_cfg import ManagerBasedEnvCfg from isaaclab.envs.ui import ViewportCameraController from isaaclab.envs.utils.io_descriptors import export_articulations_data, export_scene_data +from isaaclab.physics import PhysicsEvent from isaaclab.sim import SimulationContext from isaaclab.sim.utils import use_stage from isaaclab.ui.widgets import ManagerLiveVisualizer from isaaclab.utils.seed import configure_seed from isaaclab.utils.timer import Timer from isaaclab.utils.version import has_kit +from isaaclab.utils.warp import WarpLaunchCache from isaaclab_experimental.envs.interactive_scene_warp import InteractiveSceneWarp as InteractiveScene from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch @@ -103,6 +105,15 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # raise RuntimeError("Simulation context already exists. Cannot create a new one.") self.sim: SimulationContext = SimulationContext.instance() + # Owner-scoped Warp launcher for eager or recorded execution. + self._warp_launch = WarpLaunchCache(device=self.device) + self._warp_cache_physics_ready_handle = self.sim.physics_manager.register_callback( + self._reset_warp_caches_after_physics_ready, + PhysicsEvent.PHYSICS_READY, + order=30, + name="manager_based_env_warp_cache_rebind", + ) + # make sure torch is running on the correct device if "cuda" in self.device: torch.cuda.set_device(self.device) @@ -592,6 +603,11 @@ def seed(seed: int = -1) -> int: def close(self): """Cleanup for the environment.""" if not self._is_closed: + if self._warp_cache_physics_ready_handle is not None: + self._warp_cache_physics_ready_handle.deregister() + self._warp_cache_physics_ready_handle = None + self._reset_warp_caches_after_physics_ready(None) + # destructor is order-sensitive del self.viewport_camera_controller del self.action_manager @@ -609,6 +625,20 @@ def close(self): # update closing status self._is_closed = True + def invalidate_wp_graphs(self) -> None: + """Invalidate cached Warp graphs and recorded launches. + + Call this if launch topology, term shapes, or static argument storage + changes. The method synchronizes device work before releasing recorded + argument owners. + """ + self._reset_warp_caches_after_physics_ready(None) + + def _reset_warp_caches_after_physics_ready(self, _payload: Any) -> None: + """Drop graphs and launches recorded against previous simulation storage.""" + self._manager_call_switch.invalidate_graphs() + self._warp_launch.reset() + """ Helper functions. """ diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py b/source/isaaclab_experimental/isaaclab_experimental/envs/manager_based_rl_env_warp.py index 2e1894e54003..a9e12b6c9f53 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 @@ -32,6 +32,7 @@ from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode from isaaclab_experimental.utils.torch_utils import clone_obs_buffer +from ..managers.command_manager import _CommandViewCache from .manager_based_env_warp import ManagerBasedEnvWarp DEBUG_TIMERS = os.environ.get("DEBUG_TIMERS", "0") == "1" @@ -74,6 +75,35 @@ class ManagerBasedRLEnvWarp(ManagerBasedEnvWarp, gym.Env): """ + class _CommandManagerWarpView(CommandManager): + """Stable command manager with persistent Warp views of its Torch commands.""" + + def __init__(self, cfg: object, env: ManagerBasedRLEnvWarp): + super().__init__(cfg, env) + self._command_views = _CommandViewCache(self.active_terms, self.get_command) + + def compute(self, dt: float): + """Update commands and refresh staging buffers for computed command properties.""" + super().compute(dt) + self._command_views.refresh() + + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: + """Reset commands and refresh staging buffers before observations are computed.""" + extras = super().reset(env_ids=env_ids) + self._command_views.refresh() + return extras + + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32): + """Return the persistent Warp view of a command. + + Args: + name: Name of the command term. + + Returns: + Persistent command array with shape ``(num_envs, command_dim)``. + """ + return self._command_views.get(name) + is_vector_env: ClassVar[bool] = True """Whether the environment is a vectorized environment.""" metadata: ClassVar[dict[str, Any]] = { @@ -156,7 +186,7 @@ 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) - self.command_manager = CommandManager(self.cfg.commands, self) + self.command_manager = self._CommandManagerWarpView(self.cfg.commands, self) print("[INFO] Command Manager: ", self.command_manager) # call the parent class to load the managers for observations and actions. @@ -200,13 +230,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() @@ -269,8 +292,6 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: warp_call={"fn": self.scene.write_data_to_sim}, timer=DEBUG_TIMER_STEP, ) - - # simulate with Timer(name="simulate", msg="Newton simulation step took:", enable=DEBUG_TIMER_STEP, time_unit="us"): self.sim.step(render=False) self.recorder_manager.record_post_physics_decimation_step() diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/actions/joint_actions.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/actions/joint_actions.py index f379f66c1900..8abc477a6b06 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/actions/joint_actions.py @@ -231,8 +231,8 @@ def IO_descriptor(self) -> GenericActionIODescriptor: """ def process_actions(self, actions: wp.array, action_offset: int = 0): - wp.launch( - kernel=_process_joint_actions_kernel, + self._env._warp_launch.launch( + _process_joint_actions_kernel, dim=(self.num_envs, self.action_dim), inputs=[ actions, @@ -243,7 +243,7 @@ def process_actions(self, actions: wp.array, action_offset: int = 0): self._raw_actions, self._processed_actions, ], - device=self.device, + site=(self, int(action_offset)), ) def reset(self, env_mask: wp.array | None = None) -> None: @@ -251,11 +251,11 @@ def reset(self, env_mask: wp.array | None = None) -> None: if env_mask is None: self._raw_actions.fill_(0.0) return - wp.launch( - kernel=zero_masked_2d, + self._env._warp_launch.launch( + zero_masked_2d, dim=(self.num_envs, self.action_dim), inputs=[env_mask, self._raw_actions], - device=self.device, + site=(self, env_mask), ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py index d9e61041a7ad..abe860b89f15 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/events.py @@ -34,6 +34,16 @@ if TYPE_CHECKING: from isaaclab.assets import Articulation + +def _get_event_scratch(env) -> dict[tuple[object, ...], tuple[wp.array, ...]]: + """Return the persistent Warp scratch arrays owned by an environment.""" + scratch = getattr(env, "_warp_mdp_event_scratch", None) + if scratch is None: + scratch = {} + env._warp_mdp_event_scratch = scratch + return scratch + + # --------------------------------------------------------------------------- # Randomize rigid body center of mass # --------------------------------------------------------------------------- @@ -79,26 +89,22 @@ def randomize_rigid_body_com( """ 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 + ranges = tuple(tuple(float(value) for value in com_range.get(key, (0.0, 0.0))) for key in ["x", "y", "z"]) + com_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) + com_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) - wp.launch( - kernel=_randomize_com_kernel, + env._warp_launch.launch( + _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, + com_lo, + com_hi, ], - device=env.device, + site=(asset_cfg.body_ids_wp, ranges, env_mask), ) # Notify the solver that inertial properties changed (COM position affects inertia). @@ -156,36 +162,37 @@ def apply_external_force_torque( 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 + force_range = (float(force_range[0]), float(force_range[1])) + torque_range = (float(torque_range[0]), float(torque_range[1])) + + scratch = _get_event_scratch(env) + scratch_key = (apply_external_force_torque, asset_cfg.name, force_range, torque_range) + if scratch_key not in scratch: + scratch[scratch_key] = ( + wp.zeros((env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device), + wp.zeros((env.num_envs, asset.num_bodies), dtype=wp.vec3f, device=env.device), ) - apply_external_force_torque._is_warmed_up = True + scratch_forces, scratch_torques = scratch[scratch_key] - wp.launch( - kernel=_apply_external_force_torque_kernel, + env._warp_launch.launch( + _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, + scratch_forces, + scratch_torques, force_range[0], force_range[1], torque_range[0], torque_range[1], ], - device=env.device, + site=(scratch_forces, env_mask, force_range, torque_range), ) asset.permanent_wrench_composer.set_forces_and_torques_mask( - forces=apply_external_force_torque._scratch_forces, - torques=apply_external_force_torque._scratch_torques, + forces=scratch_forces, + torques=scratch_torques, env_mask=env_mask, ) @@ -237,33 +244,38 @@ def push_by_setting_velocity( """ asset: Articulation = env.scene[asset_cfg.name] - # First-call: allocate scratch and pre-parse constant range arguments. - if not getattr(push_by_setting_velocity, "_is_warmed_up", False): - push_by_setting_velocity._scratch_vel = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) - r = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - push_by_setting_velocity._lin_lo = wp.vec3f(r[0][0], r[1][0], r[2][0]) - push_by_setting_velocity._lin_hi = wp.vec3f(r[0][1], r[1][1], r[2][1]) - push_by_setting_velocity._ang_lo = wp.vec3f(r[3][0], r[4][0], r[5][0]) - push_by_setting_velocity._ang_hi = wp.vec3f(r[3][1], r[4][1], r[5][1]) - push_by_setting_velocity._is_warmed_up = True - - wp.launch( - kernel=_push_by_setting_velocity_kernel, + scratch = _get_event_scratch(env) + ranges = tuple( + tuple(float(value) for value in velocity_range.get(key, (0.0, 0.0))) + for key in ["x", "y", "z", "roll", "pitch", "yaw"] + ) + scratch_key = (push_by_setting_velocity, asset_cfg.name, *ranges) + if scratch_key not in scratch: + scratch[scratch_key] = (wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device),) + (scratch_vel,) = scratch[scratch_key] + + lin_lo = wp.vec3f(ranges[0][0], ranges[1][0], ranges[2][0]) + lin_hi = wp.vec3f(ranges[0][1], ranges[1][1], ranges[2][1]) + ang_lo = wp.vec3f(ranges[3][0], ranges[4][0], ranges[5][0]) + ang_hi = wp.vec3f(ranges[3][1], ranges[4][1], ranges[5][1]) + + env._warp_launch.launch( + _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, + scratch_vel, + lin_lo, + lin_hi, + ang_lo, + ang_hi, ], - device=env.device, + site=(scratch_vel, env_mask, ranges), ) - asset.write_root_velocity_to_sim_mask(root_velocity=push_by_setting_velocity._scratch_vel, env_mask=env_mask) + asset.write_root_velocity_to_sim_mask(root_velocity=scratch_vel, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -348,26 +360,34 @@ def reset_root_state_uniform( """ asset: Articulation = env.scene[asset_cfg.name] - # First-call: allocate scratch and pre-parse range dicts. - if not getattr(reset_root_state_uniform, "_is_warmed_up", False): - reset_root_state_uniform._scratch_pose = wp.zeros((env.num_envs,), dtype=wp.transformf, device=env.device) - reset_root_state_uniform._scratch_vel = wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device) - # Pre-parse pose_range dict - p = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - reset_root_state_uniform._pos_lo = wp.vec3f(p[0][0], p[1][0], p[2][0]) - reset_root_state_uniform._pos_hi = wp.vec3f(p[0][1], p[1][1], p[2][1]) - reset_root_state_uniform._rot_lo = wp.vec3f(p[3][0], p[4][0], p[5][0]) - reset_root_state_uniform._rot_hi = wp.vec3f(p[3][1], p[4][1], p[5][1]) - # Pre-parse velocity_range dict - v = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - reset_root_state_uniform._vel_lin_lo = wp.vec3f(v[0][0], v[1][0], v[2][0]) - reset_root_state_uniform._vel_lin_hi = wp.vec3f(v[0][1], v[1][1], v[2][1]) - reset_root_state_uniform._vel_ang_lo = wp.vec3f(v[3][0], v[4][0], v[5][0]) - reset_root_state_uniform._vel_ang_hi = wp.vec3f(v[3][1], v[4][1], v[5][1]) - reset_root_state_uniform._is_warmed_up = True - - wp.launch( - kernel=_reset_root_state_uniform_kernel, + scratch = _get_event_scratch(env) + pose_ranges = tuple( + tuple(float(value) for value in pose_range.get(key, (0.0, 0.0))) + for key in ["x", "y", "z", "roll", "pitch", "yaw"] + ) + velocity_ranges = tuple( + tuple(float(value) for value in velocity_range.get(key, (0.0, 0.0))) + for key in ["x", "y", "z", "roll", "pitch", "yaw"] + ) + scratch_key = (reset_root_state_uniform, asset_cfg.name, *pose_ranges, *velocity_ranges) + if scratch_key not in scratch: + scratch[scratch_key] = ( + wp.zeros((env.num_envs,), dtype=wp.transformf, device=env.device), + wp.zeros((env.num_envs,), dtype=wp.spatial_vectorf, device=env.device), + ) + scratch_pose, scratch_vel = scratch[scratch_key] + + pos_lo = wp.vec3f(pose_ranges[0][0], pose_ranges[1][0], pose_ranges[2][0]) + pos_hi = wp.vec3f(pose_ranges[0][1], pose_ranges[1][1], pose_ranges[2][1]) + rot_lo = wp.vec3f(pose_ranges[3][0], pose_ranges[4][0], pose_ranges[5][0]) + rot_hi = wp.vec3f(pose_ranges[3][1], pose_ranges[4][1], pose_ranges[5][1]) + vel_lin_lo = wp.vec3f(velocity_ranges[0][0], velocity_ranges[1][0], velocity_ranges[2][0]) + vel_lin_hi = wp.vec3f(velocity_ranges[0][1], velocity_ranges[1][1], velocity_ranges[2][1]) + vel_ang_lo = wp.vec3f(velocity_ranges[3][0], velocity_ranges[4][0], velocity_ranges[5][0]) + vel_ang_hi = wp.vec3f(velocity_ranges[3][1], velocity_ranges[4][1], velocity_ranges[5][1]) + + env._warp_launch.launch( + _reset_root_state_uniform_kernel, dim=env.num_envs, inputs=[ env_mask, @@ -375,22 +395,22 @@ def reset_root_state_uniform( 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, + scratch_pose, + scratch_vel, + pos_lo, + pos_hi, + rot_lo, + rot_hi, + vel_lin_lo, + vel_lin_hi, + vel_ang_lo, + vel_ang_hi, ], - device=env.device, + site=(scratch_pose, env_mask, pose_ranges, velocity_ranges), ) - 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) + asset.write_root_pose_to_sim_mask(root_pose=scratch_pose, env_mask=env_mask) + asset.write_root_velocity_to_sim_mask(root_velocity=scratch_vel, env_mask=env_mask) # --------------------------------------------------------------------------- @@ -456,6 +476,8 @@ def reset_joints_by_offset( via `isaaclab_experimental.envs.mdp`. """ asset: Articulation = env.scene[asset_cfg.name] + position_range = (float(position_range[0]), float(position_range[1])) + velocity_range = (float(velocity_range[0]), float(velocity_range[1])) # Assume cfg params are already resolved by the manager stack (Warp-first workflow). if asset_cfg.joint_ids_wp is None: @@ -470,8 +492,8 @@ def reset_joints_by_offset( "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment." ) - wp.launch( - kernel=_reset_joints_by_offset_kernel, + env._warp_launch.launch( + _reset_joints_by_offset_kernel, dim=env.num_envs, inputs=[ env_mask, @@ -488,7 +510,7 @@ def reset_joints_by_offset( float(velocity_range[0]), float(velocity_range[1]), ], - device=env.device, + site=(asset_cfg.joint_ids_wp, position_range, velocity_range, env_mask), ) # Sync derived buffers (_previous_joint_vel, joint_acc) for reset envs. @@ -548,6 +570,8 @@ def reset_joints_by_scale( ): """Warp-first reset of joint state by scaling defaults with random factors.""" asset: Articulation = env.scene[asset_cfg.name] + position_range = (float(position_range[0]), float(position_range[1])) + velocity_range = (float(velocity_range[0]), float(velocity_range[1])) if asset_cfg.joint_ids_wp is None: raise ValueError( @@ -561,8 +585,8 @@ def reset_joints_by_scale( "Use ManagerBasedEnvWarp or ManagerBasedRLEnvWarp as the base environment." ) - wp.launch( - kernel=_reset_joints_by_scale_kernel, + env._warp_launch.launch( + _reset_joints_by_scale_kernel, dim=env.num_envs, inputs=[ env_mask, @@ -579,7 +603,7 @@ def reset_joints_by_scale( float(velocity_range[0]), float(velocity_range[1]), ], - device=env.device, + site=(asset_cfg.joint_ids_wp, position_range, velocity_range, env_mask), ) # Sync derived buffers (_previous_joint_vel, joint_acc) for reset envs. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py index 808231c8faf6..a5ea20bd2ab6 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/observations.py @@ -90,11 +90,11 @@ def _base_pos_z_kernel( def base_pos_z(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Root height in the simulation world frame.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_base_pos_z_kernel, + env._warp_launch.launch( + _base_pos_z_kernel, dim=env.num_envs, inputs=[asset.data.root_pos_w.warp, out], - device=env.device, + site=out, ) @@ -125,11 +125,11 @@ def _base_lin_vel_kernel( def base_lin_vel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Root linear velocity in the asset's root frame.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_base_lin_vel_kernel, + env._warp_launch.launch( + _base_lin_vel_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, out], - device=env.device, + site=out, ) @@ -152,11 +152,11 @@ def _base_ang_vel_kernel( def base_ang_vel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Root angular velocity in the asset's root frame.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_base_ang_vel_kernel, + env._warp_launch.launch( + _base_ang_vel_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, out], - device=env.device, + site=out, ) @@ -179,11 +179,11 @@ def _projected_gravity_kernel( def projected_gravity(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Gravity projection on the asset's root frame.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_projected_gravity_kernel, + env._warp_launch.launch( + _projected_gravity_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, asset.data.GRAVITY_VEC_W.warp, out], - device=env.device, + site=out, ) @@ -204,11 +204,11 @@ def joint_pos(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntity "SceneEntityCfg.joint_ids_wp is required for subset joint observations in Warp-first observations. " "Pass `asset_cfg` via term cfg params so it is resolved at manager init." ) - wp.launch( - kernel=_joint_gather_kernel, + env._warp_launch.launch( + _joint_gather_kernel, dim=(env.num_envs, out.shape[1]), inputs=[asset.data.joint_pos.warp, joint_ids_wp, out], - device=env.device, + site=("joint_pos", out), ) @@ -240,11 +240,11 @@ def joint_pos_rel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEn "SceneEntityCfg.joint_ids_wp is required for subset joint observations in Warp-first observations. " "Pass `asset_cfg` via term cfg params so it is resolved at manager init." ) - wp.launch( - kernel=_joint_rel_gather_kernel, + env._warp_launch.launch( + _joint_rel_gather_kernel, dim=(env.num_envs, out.shape[1]), inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, joint_ids_wp, out], - device=env.device, + site=("joint_pos_rel", out), ) @@ -274,11 +274,11 @@ def joint_pos_limit_normalized(env: ManagerBasedEnv, out, asset_cfg: SceneEntity "SceneEntityCfg.joint_ids_wp is required for subset joint observations in Warp-first observations. " "Pass `asset_cfg` via term cfg params so it is resolved at manager init." ) - wp.launch( - kernel=_joint_pos_limit_normalized_kernel, + env._warp_launch.launch( + _joint_pos_limit_normalized_kernel, dim=(env.num_envs, out.shape[1]), inputs=[asset.data.joint_pos.warp, asset.data.soft_joint_pos_limits.warp, joint_ids_wp, out], - device=env.device, + site=out, ) @@ -294,11 +294,11 @@ def joint_vel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntity "SceneEntityCfg.joint_ids_wp is required for subset joint observations in Warp-first observations. " "Pass `asset_cfg` via term cfg params so it is resolved at manager init." ) - wp.launch( - kernel=_joint_gather_kernel, + env._warp_launch.launch( + _joint_gather_kernel, dim=(env.num_envs, out.shape[1]), inputs=[asset.data.joint_vel.warp, joint_ids_wp, out], - device=env.device, + site=("joint_vel", out), ) @@ -318,11 +318,11 @@ def joint_vel_rel(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEn "SceneEntityCfg.joint_ids_wp is required for subset joint observations in Warp-first observations. " "Pass `asset_cfg` via term cfg params so it is resolved at manager init." ) - wp.launch( - kernel=_joint_rel_gather_kernel, + env._warp_launch.launch( + _joint_rel_gather_kernel, dim=(env.num_envs, out.shape[1]), inputs=[asset.data.joint_vel.warp, asset.data.default_joint_vel.warp, joint_ids_wp, out], - device=env.device, + site=("joint_vel_rel", out), ) @@ -353,17 +353,6 @@ def generated_commands(env: ManagerBasedEnv, out, command_name: str) -> None: """The generated command from the command manager. Writes into ``out``. Warp-first override of :func:`isaaclab.envs.mdp.observations.generated_commands`. - Uses ``wp.from_torch`` to create a zero-copy warp view of the command tensor on first call. + Uses the command manager's persistent Warp view. """ - # TODO(warp-migration): Cross-manager access (observation → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - fn = generated_commands - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - fn._cmd_wp = cmd - else: - fn._cmd_wp = wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True - wp.copy(out, fn._cmd_wp) + wp.copy(out, env.command_manager.get_command_wp(command_name)) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index 338d27c99c5e..efd7666cea18 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -44,11 +44,11 @@ def _is_alive_kernel(terminated: wp.array(dtype=wp.bool), out: wp.array(dtype=wp def is_alive(env: ManagerBasedRLEnv, out: wp.array(dtype=wp.float32)) -> None: """Reward for being alive. Writes into ``out`` (shape: (num_envs,)).""" - wp.launch( - kernel=_is_alive_kernel, + env._warp_launch.launch( + _is_alive_kernel, dim=env.num_envs, inputs=[env.termination_manager.terminated_wp, out], - device=env.device, + site=out, ) @@ -60,11 +60,11 @@ def _is_terminated_kernel(terminated: wp.array(dtype=wp.bool), out: wp.array(dty def is_terminated(env: ManagerBasedRLEnv, out) -> None: """Penalize terminated episodes. Writes into ``out``.""" - wp.launch( - kernel=_is_terminated_kernel, + env._warp_launch.launch( + _is_terminated_kernel, dim=env.num_envs, inputs=[env.termination_manager.terminated_wp, out], - device=env.device, + site=out, ) @@ -95,11 +95,11 @@ def _lin_vel_z_l2_kernel( def lin_vel_z_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Penalize z-axis base linear velocity using L2 squared kernel.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_lin_vel_z_l2_kernel, + env._warp_launch.launch( + _lin_vel_z_l2_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, out], - device=env.device, + site=out, ) @@ -117,11 +117,11 @@ def _ang_vel_xy_l2_kernel( def ang_vel_xy_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Penalize xy-axis base angular velocity using L2 squared kernel.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_ang_vel_xy_l2_kernel, + env._warp_launch.launch( + _ang_vel_xy_l2_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, out], - device=env.device, + site=out, ) @@ -141,11 +141,11 @@ def _flat_orientation_l2_kernel( def flat_orientation_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Penalize non-flat base orientation using L2 squared kernel.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_flat_orientation_l2_kernel, + env._warp_launch.launch( + _flat_orientation_l2_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, asset.data.GRAVITY_VEC_W.warp, out], - device=env.device, + site=out, ) @@ -171,11 +171,11 @@ def _sum_sq_masked_kernel( def joint_torques_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Penalize joint torques applied on the articulation using L2 squared kernel.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_sum_sq_masked_kernel, + env._warp_launch.launch( + _sum_sq_masked_kernel, dim=env.num_envs, inputs=[asset.data.applied_torque.warp, asset_cfg.joint_mask, out], - device=env.device, + site=("joint_torques_l2", out), ) @@ -195,33 +195,33 @@ def _sum_abs_masked_kernel( def joint_vel_l1(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg) -> None: """Penalize joint velocities on the articulation using an L1-kernel. Writes into ``out``.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_sum_abs_masked_kernel, + env._warp_launch.launch( + _sum_abs_masked_kernel, dim=env.num_envs, inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask, out], - device=env.device, + site=out, ) def joint_vel_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Penalize joint velocities on the articulation using L2 squared kernel.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_sum_sq_masked_kernel, + env._warp_launch.launch( + _sum_sq_masked_kernel, dim=env.num_envs, inputs=[asset.data.joint_vel.warp, asset_cfg.joint_mask, out], - device=env.device, + site=("joint_vel_l2", out), ) def joint_acc_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Penalize joint accelerations on the articulation using L2 squared kernel.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_sum_sq_masked_kernel, + env._warp_launch.launch( + _sum_sq_masked_kernel, dim=env.num_envs, inputs=[asset.data.joint_acc.warp, asset_cfg.joint_mask, out], - device=env.device, + site=("joint_acc_l2", out), ) @@ -244,11 +244,11 @@ def _sum_abs_diff_masked_kernel( def joint_deviation_l1(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Penalize joint positions that deviate from the default one.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_sum_abs_diff_masked_kernel, + env._warp_launch.launch( + _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], - device=env.device, + site=out, ) @@ -282,11 +282,11 @@ def _joint_pos_limits_kernel( def joint_pos_limits(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Penalize joint positions if they cross the soft limits.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_joint_pos_limits_kernel, + env._warp_launch.launch( + _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], - device=env.device, + site=out, ) @@ -312,11 +312,11 @@ def _sum_sq_diff_2d_kernel( def action_rate_l2(env: ManagerBasedRLEnv, out) -> None: """Penalize the rate of change of the actions using L2 squared kernel.""" - wp.launch( - kernel=_sum_sq_diff_2d_kernel, + env._warp_launch.launch( + _sum_sq_diff_2d_kernel, dim=env.num_envs, inputs=[env.action_manager.action, env.action_manager.prev_action, out], - device=env.device, + site=out, ) @@ -332,11 +332,11 @@ def _sum_sq_2d_kernel(x: wp.array(dtype=wp.float32, ndim=2), out: wp.array(dtype def action_l2(env: ManagerBasedRLEnv, out) -> None: """Penalize the actions using L2 squared kernel.""" - wp.launch( - kernel=_sum_sq_2d_kernel, + env._warp_launch.launch( + _sum_sq_2d_kernel, dim=env.num_envs, inputs=[env.action_manager.action, out], - device=env.device, + site=out, ) @@ -374,11 +374,11 @@ def undesired_contacts(env: ManagerBasedRLEnv, out, threshold: float, sensor_cfg Warp-first override of :func:`isaaclab.envs.mdp.rewards.undesired_contacts`. """ contact_sensor = env.scene.sensors[sensor_cfg.name] - wp.launch( - kernel=_undesired_contacts_kernel, + env._warp_launch.launch( + _undesired_contacts_kernel, dim=env.num_envs, inputs=[contact_sensor.data.net_forces_w_history.warp, sensor_cfg.body_ids_wp, threshold, out], - device=env.device, + site=(out, float(threshold)), ) @@ -415,28 +415,18 @@ def track_lin_vel_xy_exp( Warp-first override of :func:`isaaclab.envs.mdp.rewards.track_lin_vel_xy_exp`. """ asset: Articulation = env.scene[asset_cfg.name] - # cache the warp view of the command tensor on first call (zero-copy) - # TODO(warp-migration): Cross-manager access (reward → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - if not getattr(track_lin_vel_xy_exp, "_is_warmed_up", False) or track_lin_vel_xy_exp._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - track_lin_vel_xy_exp._cmd_wp = cmd - else: - track_lin_vel_xy_exp._cmd_wp = wp.from_torch(cmd) - track_lin_vel_xy_exp._cmd_name = command_name - track_lin_vel_xy_exp._is_warmed_up = True - wp.launch( - kernel=_track_lin_vel_xy_exp_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _track_lin_vel_xy_exp_kernel, dim=env.num_envs, inputs=[ asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, - track_lin_vel_xy_exp._cmd_wp, + command, 1.0 / (std * std), out, ], - device=env.device, + site=(out, command_name, float(std)), ) @@ -466,26 +456,17 @@ def track_ang_vel_z_exp( Warp-first override of :func:`isaaclab.envs.mdp.rewards.track_ang_vel_z_exp`. """ asset: Articulation = env.scene[asset_cfg.name] - # TODO(warp-migration): Cross-manager access (reward → command). Replace with direct - # warp getter once all managers are guaranteed to be warp-native. - if not getattr(track_ang_vel_z_exp, "_is_warmed_up", False) or track_ang_vel_z_exp._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - if isinstance(cmd, wp.array): - track_ang_vel_z_exp._cmd_wp = cmd - else: - track_ang_vel_z_exp._cmd_wp = wp.from_torch(cmd) - track_ang_vel_z_exp._cmd_name = command_name - track_ang_vel_z_exp._is_warmed_up = True - wp.launch( - kernel=_track_ang_vel_z_exp_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _track_ang_vel_z_exp_kernel, dim=env.num_envs, inputs=[ asset.data.root_link_pose_w.warp, asset.data.root_com_vel_w.warp, - track_ang_vel_z_exp._cmd_wp, + command, 2, 1.0 / (std * std), out, ], - device=env.device, + site=(out, command_name, float(std)), ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py index d309a70538bb..9e586d3a52f4 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py @@ -41,11 +41,11 @@ def _time_out_kernel( def time_out(env: ManagerBasedRLEnv, out) -> None: """Terminate the episode when episode length exceeds the maximum episode length.""" - wp.launch( - kernel=_time_out_kernel, + env._warp_launch.launch( + _time_out_kernel, dim=env.num_envs, inputs=[env._episode_length_buf_wp, env.max_episode_length, out], - device=env.device, + site=(out, int(env.max_episode_length)), ) @@ -69,11 +69,11 @@ def root_height_below_minimum( ) -> None: """Terminate when the asset's root height is below the minimum height.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_root_height_below_min_kernel, + env._warp_launch.launch( + _root_height_below_min_kernel, dim=env.num_envs, inputs=[asset.data.root_pos_w.warp, minimum_height, out], - device=env.device, + site=(out, float(minimum_height)), ) @@ -113,11 +113,11 @@ def joint_pos_out_of_manual_limit( f"joint_mask length ({asset_cfg.joint_mask.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, + env._warp_launch.launch( + _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], - device=env.device, + site=(out, float(bounds[0]), float(bounds[1])), ) @@ -152,9 +152,9 @@ def illegal_contact(env: ManagerBasedRLEnv, out, threshold: float, sensor_cfg: S Warp-first override of :func:`isaaclab.envs.mdp.terminations.illegal_contact`. """ contact_sensor = env.scene.sensors[sensor_cfg.name] - wp.launch( - kernel=_illegal_contact_kernel, + env._warp_launch.launch( + _illegal_contact_kernel, dim=env.num_envs, inputs=[contact_sensor.data.net_forces_w_history.warp, sensor_cfg.body_ids_wp, threshold, out], - device=env.device, + site=(out, float(threshold)), ) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py index 672af82bc8ab..d766a5199119 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/action_manager.py @@ -387,17 +387,17 @@ def reset( self._prev_action.fill_(0.0) self._action.fill_(0.0) else: - wp.launch( - kernel=_zero_masked_2d, + self._env._warp_launch.launch( + _zero_masked_2d, dim=(self.num_envs, self.total_action_dim), inputs=[env_mask, self._prev_action], - device=self.device, + site=(self._prev_action, env_mask), ) - wp.launch( - kernel=_zero_masked_2d, + self._env._warp_launch.launch( + _zero_masked_2d, dim=(self.num_envs, self.total_action_dim), inputs=[env_mask, self._action], - device=self.device, + site=(self._action, env_mask), ) # reset all action terms diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py index 8b4e8f83dbe7..263f688d3d36 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/command_manager.py @@ -10,7 +10,7 @@ import inspect import weakref from abc import abstractmethod -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING import torch @@ -225,31 +225,36 @@ def reset( # compute selected count and reset scale self._reset_count_wp.zero_() self._reset_scale_wp.zero_() - wp.launch(kernel=count_masked, dim=self.num_envs, inputs=[env_mask, self._reset_count_wp], device=self.device) - wp.launch( - kernel=compute_reset_scale, + self._env._warp_launch.launch( + count_masked, + dim=self.num_envs, + inputs=[env_mask, self._reset_count_wp], + site=(self, env_mask), + ) + self._env._warp_launch.launch( + compute_reset_scale, dim=1, inputs=[self._reset_count_wp, 1.0, self._reset_scale_wp], - device=self.device, + site=self, ) # update pre-allocated reset extras and clear selected metric rows for metric_name, metric_value_wp in self.metrics.items(): out_mean_wp = self._reset_metric_mean_wp[metric_name] out_mean_wp.zero_() - wp.launch( - kernel=_sum_and_zero_masked, + self._env._warp_launch.launch( + _sum_and_zero_masked, dim=self.num_envs, inputs=[env_mask, self._reset_scale_wp, metric_value_wp, out_mean_wp], - device=self.device, + site=(out_mean_wp, env_mask), ) # set the command counter to zero - wp.launch( - kernel=_zero_counter_masked, + self._env._warp_launch.launch( + _zero_counter_masked, dim=self.num_envs, inputs=[env_mask, self.command_counter_wp], - device=self.device, + site=(self, env_mask), ) # resample the command self._resample(env_mask=env_mask) @@ -288,11 +293,11 @@ def compute(self, dt: float): # update the metrics based on current state self._update_metrics() # reduce the time left before resampling and build resample mask - wp.launch( - kernel=_step_time_left_and_build_resample_mask, + self._env._warp_launch.launch( + _step_time_left_and_build_resample_mask, dim=self.num_envs, inputs=[self.time_left_wp, float(dt), self._resample_mask_wp], - device=self.device, + site=(self, float(dt)), ) # resample masked envs self._resample(env_mask=self._resample_mask_wp) @@ -320,8 +325,8 @@ def _resample(self, env_mask: wp.array): raise RuntimeError("Environment rng_state_wp is not initialized.") # resample time-left and increment command-counter for masked envs - wp.launch( - kernel=_resample_time_left_and_increment_counter, + self._env._warp_launch.launch( + _resample_time_left_and_increment_counter, dim=self.num_envs, inputs=[ env_mask, @@ -331,7 +336,12 @@ def _resample(self, env_mask: wp.array): float(self.cfg.resampling_time_range[0]), float(self.cfg.resampling_time_range[1]), ], - device=self.device, + site=( + self, + env_mask, + float(self.cfg.resampling_time_range[0]), + float(self.cfg.resampling_time_range[1]), + ), ) # resample command values for masked envs self._resample_command(env_mask) @@ -372,6 +382,73 @@ def _debug_vis_callback(self, event): raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") +class _CommandViewCache: + """Maintain pointer-stable Warp command views across replacing properties.""" + + def __init__(self, names: Sequence[str], get_command: Callable[[str], torch.Tensor | wp.array]): + self._names = frozenset(names) + self._get_command = get_command + self._owners: dict[str, torch.Tensor | wp.array] = {} + self._sources: dict[str, torch.Tensor | wp.array] = {} + self.views: dict[str, wp.array(dtype=wp.float32)] = {} + + def _create(self, name: str) -> None: + """Create the persistent view for one command on first use.""" + if name not in self._names: + raise KeyError(f"Unknown command term {name!r}.") + command = self._get_command(name) + if isinstance(command, torch.Tensor): + if command.dtype != torch.float32: + raise TypeError(f"Command {name!r} must use torch.float32, got {command.dtype}.") + # Keep a distinct tensor object so ``command.set_(...)`` cannot + # repoint the storage retained by the persistent Warp view. + owner = command.detach() + self._owners[name] = owner + self.views[name] = wp.from_torch(owner) + elif isinstance(command, wp.array): + if command.dtype != wp.float32: + raise TypeError(f"Command {name!r} must use wp.float32, got {command.dtype}.") + self._owners[name] = command + self.views[name] = command + else: + raise TypeError(f"Command {name!r} must return a Torch tensor or Warp array, got {type(command).__name__}.") + + def refresh(self) -> None: + """Copy replaced command properties into their original persistent buffers.""" + for name, owner in self._owners.items(): + current = self._get_command(name) + if isinstance(owner, torch.Tensor) and isinstance(current, torch.Tensor): + owner_specialization = (owner.shape, owner.dtype, owner.device, owner.stride()) + current_specialization = (current.shape, current.dtype, current.device, current.stride()) + if owner_specialization != current_specialization: + raise RuntimeError( + f"Command {name!r} changed tensor specialization from " + f"{owner_specialization} to {current_specialization}." + ) + if owner.data_ptr() != current.data_ptr(): + owner.copy_(current) + self._sources[name] = current + elif isinstance(owner, wp.array) and isinstance(current, wp.array): + owner_specialization = (owner.shape, owner.dtype, owner.device, owner.strides) + current_specialization = (current.shape, current.dtype, current.device, current.strides) + if owner_specialization != current_specialization: + raise RuntimeError( + f"Command {name!r} changed array specialization from " + f"{owner_specialization} to {current_specialization}." + ) + if owner.ptr != current.ptr: + wp.copy(owner, current) + self._sources[name] = current + else: + raise TypeError(f"Command {name!r} changed from {type(owner).__name__} to {type(current).__name__}.") + + def get(self, name: str) -> wp.array(dtype=wp.float32): + """Return the pointer-stable Warp view for a command term.""" + if name not in self.views: + self._create(name) + return self.views[name] + + class CommandManager(ManagerBase): """Manager for generating commands. @@ -403,6 +480,8 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): super().__init__(cfg, env) # store the commands self._commands = dict() + self._command_views = _CommandViewCache(self.active_terms, lambda name: self._terms[name].command) + self._commands_wp = self._command_views.views if self.cfg: self.cfg.debug_vis = False for term in self._terms.values(): @@ -526,6 +605,8 @@ def reset( # reset the command term term.reset(env_mask=env_mask) + self._command_views.refresh() + return self._reset_extras def compute(self, dt: float): @@ -541,6 +622,7 @@ def compute(self, dt: float): for term in self._terms.values(): # compute term's value term.compute(dt) + self._command_views.refresh() def get_command(self, name: str) -> torch.Tensor: """Returns the command for the specified command term. @@ -556,6 +638,17 @@ def get_command(self, name: str) -> torch.Tensor: return wp.to_torch(command) return command + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32): + """Return the persistent Warp view of a command. + + Args: + name: Name of the command term. + + Returns: + Command array with shape ``(num_envs, command_dim)``. + """ + return self._command_views.get(name) + def get_term(self, name: str) -> CommandTerm: """Returns the command term with the specified name. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py index ab8a0264ea9b..9ecec06bc65a 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/event_manager.py @@ -230,6 +230,7 @@ def set_term_cfg(self, term_name: str, cfg: EventTermCfg): break if not term_found: raise ValueError(f"Event term '{term_name}' not found.") + self._env.invalidate_wp_graphs() def get_term_cfg(self, term_name: str) -> EventTermCfg: for mode, terms in self._mode_term_names.items(): @@ -266,8 +267,8 @@ def reset( if term_cfg.is_global_time or not term_cfg.resample_interval_on_reset: continue lower, upper = self._interval_term_ranges[i] - wp.launch( - kernel=_interval_reset_selected, + self._env._warp_launch.launch( + _interval_reset_selected, dim=self.num_envs, inputs=[ env_mask, @@ -276,7 +277,7 @@ def reset( float(lower), float(upper), ], - device=self.device, + site=(self._interval_term_time_left_wp[i], env_mask, float(lower), float(upper)), ) return {} @@ -342,8 +343,8 @@ def _apply_interval(self, dt: float) -> None: "EventManager._apply_interval: _interval_global_rng_state_wp is not initialized." ) # update scalar time_left and scalar flag (mask is a broadcast view of the flag) - wp.launch( - kernel=_interval_step_global, + self._env._warp_launch.launch( + _interval_step_global, dim=1, inputs=[ self._interval_term_time_left_wp[i], @@ -353,12 +354,12 @@ def _apply_interval(self, dt: float) -> None: float(lower), float(upper), ], - device=self.device, + site=(self._interval_term_time_left_wp[i], float(dt), float(lower), float(upper)), ) term_cfg.func(self._env, self._scratch_interval_trigger_mask_view_wp, **term_cfg.params) else: - wp.launch( - kernel=_interval_step_per_env, + self._env._warp_launch.launch( + _interval_step_per_env, dim=self.num_envs, inputs=[ self._interval_term_time_left_wp[i], @@ -368,7 +369,7 @@ def _apply_interval(self, dt: float) -> None: float(lower), float(upper), ], - device=self.device, + site=(self._interval_term_time_left_wp[i], float(dt), float(lower), float(upper)), ) term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params) @@ -379,8 +380,8 @@ def _apply_reset(self, env_mask_wp: wp.array, global_env_step_count_wp: wp.array # 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, + self._env._warp_launch.launch( + _reset_compute_valid_mask, dim=self.num_envs, inputs=[ env_mask_wp, @@ -390,7 +391,12 @@ def _apply_reset(self, env_mask_wp: wp.array, global_env_step_count_wp: wp.array global_env_step_count_wp, int(min_step_count), ], - device=self.device, + site=( + self._reset_term_last_triggered_step_wp[index], + env_mask_wp, + global_env_step_count_wp, + min_step_count, + ), ) term_cfg.func(self._env, self._scratch_term_mask_wp, **term_cfg.params) diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py index 74d0a9a955be..1e9fd6b7f73c 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/observation_manager.py @@ -509,18 +509,18 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch. # clip then scale (stable semantics); implementation may use Warp kernels if term_cfg.clip is not None: - wp.launch( - kernel=_apply_clip, + self._env._warp_launch.launch( + _apply_clip, dim=self.num_envs, inputs=[term_cfg.out_wp, float(term_cfg.clip[0]), float(term_cfg.clip[1])], - device=self.device, + site=(term_cfg.out_wp, float(term_cfg.clip[0]), float(term_cfg.clip[1])), ) if term_cfg.scale is not None: - wp.launch( - kernel=_apply_scale, + self._env._warp_launch.launch( + _apply_scale, dim=self.num_envs, inputs=[term_cfg.out_wp, term_cfg.scale_wp], - device=self.device, + site=term_cfg.out_wp, ) # TODO(jichuanh): This is not migrated yet. Need revisit. diff --git a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py index 67dcbc055c2f..20823d0181c2 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/reward_manager.py @@ -255,20 +255,25 @@ def reset( self._reset_count_wp.zero_() self._reset_scale_wp.zero_() - wp.launch(kernel=count_masked, dim=self.num_envs, inputs=[env_mask, self._reset_count_wp], device=self.device) - wp.launch( - kernel=compute_reset_scale, + self._env._warp_launch.launch( + count_masked, + dim=self.num_envs, + inputs=[env_mask, self._reset_count_wp], + site=(self, env_mask), + ) + self._env._warp_launch.launch( + compute_reset_scale, dim=1, inputs=[self._reset_count_wp, float(self._env.max_episode_length_s), self._reset_scale_wp], - device=self.device, + site=(self, float(self._env.max_episode_length_s)), ) if self._num_terms > 0: - wp.launch( - kernel=_sum_and_zero_masked, + self._env._warp_launch.launch( + _sum_and_zero_masked, dim=(self._num_terms, self.num_envs), inputs=[env_mask, self._reset_scale_wp, self._episode_sums_wp, self._episode_sum_avg_wp], - device=self.device, + site=(self, env_mask), ) # reset all the reward terms @@ -291,11 +296,10 @@ def compute(self, dt: float) -> torch.Tensor: """ # 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, + self._env._warp_launch.launch( + _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: @@ -307,8 +311,8 @@ def compute(self, dt: float) -> torch.Tensor: term_cfg.func(self._env, term_cfg.out, **term_cfg.params) # update total reward, episodic sums and step rewards in a single kernel launch - wp.launch( - kernel=_reward_finalize, + self._env._warp_launch.launch( + _reward_finalize, dim=self.num_envs, inputs=[ self._term_outs_wp, @@ -318,7 +322,7 @@ def compute(self, dt: float) -> torch.Tensor: self._episode_sums_wp, self._step_reward_wp, ], - device=self.device, + site=float(dt), ) return self._reward_tensor_view @@ -347,6 +351,7 @@ 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) + self._env.invalidate_wp_graphs() 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/termination_manager.py b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py index 8a768651b292..03e31febafce 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py +++ b/source/isaaclab_experimental/isaaclab_experimental/managers/termination_manager.py @@ -259,11 +259,10 @@ def reset( 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( - kernel=_termination_reset_mean_all_2d, + self._env._warp_launch.launch( + _termination_reset_mean_all_2d, dim=(self.num_envs, len(self._term_names)), inputs=[self._last_episode_dones_wp, self._term_done_avg_wp], - device=self.device, ) for term_cfg in self._class_term_cfgs: term_cfg.func.reset(env_mask=env_mask) @@ -281,11 +280,10 @@ def compute(self) -> torch.Tensor: The combined termination signal of shape (num_envs,). """ # reset computation (Warp buffers) in a single kernel launch - wp.launch( - kernel=_termination_pre_compute_reset, + self._env._warp_launch.launch( + _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) @@ -293,8 +291,8 @@ def compute(self) -> torch.Tensor: term_cfg.func(self._env, term_cfg.out, **term_cfg.params) # finalize dones and update last-episode term flags (single kernel launch) - wp.launch( - kernel=_termination_finalize, + self._env._warp_launch.launch( + _termination_finalize, dim=self.num_envs, inputs=[ self._term_dones_wp, @@ -304,7 +302,6 @@ def compute(self) -> torch.Tensor: self._dones_wp, self._last_episode_dones_wp, ], - device=self.device, ) return self._dones_tensor_view diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py index 89924c789bc3..d392f9db2642 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/manager_call_switch.py @@ -22,7 +22,8 @@ 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_NOT_CAPTURED`` (1): Call Warp-compatible implementations without enclosing the manager stage in a + CUDA graph. Individual pointer-stable kernel calls may still use recorded-launch replay. * ``WARP_CAPTURED`` (2): Call Warp implementations with CUDA graph capture/replay. """ @@ -35,8 +36,9 @@ 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. + stable Python, uncaptured Warp, or a captured Warp CUDA graph. The uncaptured + path delegates individual kernel launch policy to owner-scoped launch caches. + Optionally wraps each call in a :class:`Timer` context for profiling. """ DEFAULT_CONFIG: dict[str, int] = {"default": 2} @@ -121,7 +123,7 @@ def call_stage( Args: stage: Stage identifier in the form ``"ManagerName_function_name"``. - warp_call: Call spec for the warp path (eager or captured). + warp_call: Call spec for the uncaptured or graph-captured Warp path. 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). @@ -202,7 +204,11 @@ def register_manager_capturability(self, manager_name: str, capturable: bool) -> # ------------------------------------------------------------------ def _run_call(self, call: dict[str, Any]) -> Any: - """Execute a single call spec eagerly.""" + """Execute a call without stage-level graph capture. + + The called Warp implementation may replay its own pointer-stable kernel + commands through an owner-scoped launch cache. + """ return call["fn"](*call.get("args", ()), **call.get("kwargs", {})) def _wp_capture_or_launch(self, stage: str, call: dict[str, Any]) -> Any: diff --git a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py index d446be28d099..f5884838c824 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py +++ b/source/isaaclab_experimental/isaaclab_experimental/utils/warp_graph_cache.py @@ -5,24 +5,28 @@ """Warp CUDA graph capture-or-replay utility.""" +import os from collections.abc import Callable from typing import Any import warp as wp +_WARP_GRAPH_MODE_ENV = "ISAACLAB_WARP_GRAPH_MODE" + class WarpGraphCache: - """Caches Warp CUDA graphs by stage name: captures on first call, replays after. + """Cache Warp CUDA graphs by stage name, or execute stages eagerly. - 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 + In capture mode, the first call for a stage executes one eager warm-up + before graph capture. This lets one-time initialization code (memory + allocations, Torch dtype casts, ``hasattr`` guards, and similar work) + 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. + The return value from the capture run is cached and returned on every later + replay, ensuring captured stages return the same references, such as tensor + views. In eager mode, every call executes the callable and returns that + invocation's result without caching. Usage:: @@ -30,9 +34,20 @@ class 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) + + Args: + mode: Execution mode. ``"capture"`` caches CUDA graphs and ``"eager"`` + calls each stage directly. If omitted, the value is read from + ``ISAACLAB_WARP_GRAPH_MODE`` and defaults to ``"capture"``. """ - def __init__(self): + def __init__(self, mode: str | None = None): + if mode is None: + mode = os.environ.get(_WARP_GRAPH_MODE_ENV, "capture") + mode = mode.lower() + if mode not in {"capture", "eager"}: + raise ValueError(f"Invalid Warp graph mode {mode!r}. Set {_WARP_GRAPH_MODE_ENV} to 'capture' or 'eager'.") + self._mode = mode self._graphs: dict[str, Any] = {} self._results: dict[str, Any] = {} @@ -43,7 +58,7 @@ def capture_or_replay( args: tuple = (), kwargs: dict[str, Any] | None = None, ) -> Any: - """Capture *fn* into a CUDA graph on the first call, then replay. + """Execute a callable eagerly or capture it once and replay its graph. Args: stage: Unique name identifying this captured scope. @@ -54,10 +69,13 @@ def capture_or_replay( kwargs: Keyword arguments forwarded to *fn*. Defaults to ``None``. Returns: - The cached return value from the first (capture) invocation. + In eager mode, the current invocation's return value. In capture + mode, the cached return value from the capture invocation. """ if kwargs is None: kwargs = {} + if self._mode == "eager": + return fn(*args, **kwargs) graph = self._graphs.get(stage) if graph is not None: wp.capture_launch(graph) @@ -72,10 +90,31 @@ def capture_or_replay( return result def invalidate(self, stage: str | None = None) -> None: - """Drop cached graph(s). If *stage* is ``None``, drop all.""" + """Drain and drop one captured graph or every captured graph. + + Args: + stage: Stage to invalidate. If ``None``, invalidate every stage. + + Devices are synchronized only when the selected stage owns a captured + graph. Eager and otherwise empty caches return without synchronizing. + """ if stage is None: + graphs = tuple(self._graphs.values()) + self._synchronize_graphs(graphs) self._graphs.clear() self._results.clear() else: + graph = self._graphs.get(stage) + graphs = () if graph is None else (graph,) + self._synchronize_graphs(graphs) self._graphs.pop(stage, None) self._results.pop(stage, None) + + @staticmethod + def _synchronize_graphs(graphs: tuple[Any, ...]) -> None: + """Synchronize each device referenced by a collection of graphs.""" + devices: dict[str, Any] = {} + for graph in graphs: + devices[str(graph.device)] = graph.device + for device in devices.values(): + wp.synchronize_device(device) diff --git a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py index 4ceec2ecb60b..30005b1a8f16 100644 --- a/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py +++ b/source/isaaclab_experimental/test/envs/mdp/parity_helpers.py @@ -533,11 +533,15 @@ class MockCommandManager: def __init__(self, command_tensor: torch.Tensor, cmd_term: MockCommandTerm): self._cmd = command_tensor + self._cmd_wp = wp.from_torch(command_tensor) self._term = cmd_term def get_command(self, name: str) -> torch.Tensor: return self._cmd + def get_command_wp(self, name: str) -> wp.array(dtype=wp.float32): + return self._cmd_wp + def get_term(self, name: str): return self._term diff --git a/source/isaaclab_experimental/test/envs/mdp/test_actions_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_actions_warp_parity.py index f2db4f64fbe9..764824556ef3 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_actions_warp_parity.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_actions_warp_parity.py @@ -12,6 +12,8 @@ import torch import warp as wp +from isaaclab.utils.warp import WarpLaunchCache + wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") @@ -42,6 +44,7 @@ def __init__(self, asset): self.scene = MockScene({"robot": asset}, env_origins=None) self.num_envs = NUM_ENVS self.device = DEVICE + self._warp_launch = WarpLaunchCache(device=DEVICE) # ============================================================================ @@ -167,6 +170,42 @@ def test_joint_action_with_scale(self, env, asset, actions_wp): expected = raw * 2.5 assert_close(term.processed_actions, expected) + def test_joint_action_replay_reads_mutated_persistent_action(self, env): + """Replay should consume updated values from the persistent action buffer.""" + cfg = JointEffortActionCfg(asset_name="robot", joint_names=[".*"]) + term = JointEffortAction(cfg, env) + first = wp.full((NUM_ENVS, 2 * NUM_JOINTS), value=1.0, dtype=wp.float32, device=DEVICE) + second_torch = torch.cat( + ( + torch.full((NUM_ENVS, NUM_JOINTS), 2.0, device=DEVICE), + torch.full((NUM_ENVS, NUM_JOINTS), 7.0, device=DEVICE), + ), + dim=1, + ) + + term.process_actions(first, action_offset=NUM_JOINTS) + wp.copy(first, wp.from_torch(second_torch)) + term.process_actions(first, action_offset=NUM_JOINTS) + + assert_close(term.raw_actions, torch.full((NUM_ENVS, NUM_JOINTS), 7.0, device=DEVICE)) + + def test_joint_action_replay_updates_reset_mask_pointer(self, env): + """Replay should apply the current reset mask when its allocation changes.""" + cfg = JointEffortActionCfg(asset_name="robot", joint_names=[".*"]) + term = JointEffortAction(cfg, env) + actions = wp.ones((NUM_ENVS, NUM_JOINTS), dtype=wp.float32, device=DEVICE) + first_mask = wp.array([index < NUM_ENVS // 2 for index in range(NUM_ENVS)], dtype=wp.bool, device=DEVICE) + second_mask = wp.array([index >= NUM_ENVS // 2 for index in range(NUM_ENVS)], dtype=wp.bool, device=DEVICE) + + term.process_actions(actions) + term.reset(env_mask=first_mask) + term.process_actions(actions) + term.reset(env_mask=second_mask) + + raw = wp.to_torch(term.raw_actions) + assert_close(raw[: NUM_ENVS // 2], torch.ones((NUM_ENVS // 2, NUM_JOINTS), device=DEVICE)) + assert_close(raw[NUM_ENVS // 2 :], torch.zeros((NUM_ENVS // 2, NUM_JOINTS), device=DEVICE)) + # ============================================================================ # Mathematical parity tests: warp processed_actions == raw * scale + offset diff --git a/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_events_warp_parity.py index 24b9eb022665..b0efc9773bdd 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 @@ -12,6 +12,8 @@ import torch import warp as wp +from isaaclab.utils.warp import WarpLaunchCache + # Skip entire module if no CUDA device available wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") @@ -37,25 +39,6 @@ # ============================================================================ -@pytest.fixture(autouse=True) -def _clear_function_caches(): - """Clear first-call caches on warp MDP functions so each test starts fresh. - - Functions that cache warp views via the ``hasattr`` pattern need clearing - between tests to avoid stale references from prior fixtures. - """ - yield - for fn in ( - warp_evt.push_by_setting_velocity, - warp_evt.apply_external_force_torque, - warp_evt.reset_root_state_uniform, - warp_evt.randomize_rigid_body_com, - ): - for attr in list(vars(fn)): - if attr.startswith("_"): - delattr(fn, attr) - - @pytest.fixture() def art_data(): return MockArticulationData(NUM_ENVS, NUM_JOINTS, DEVICE) @@ -99,6 +82,7 @@ class _Env: env.action_manager = MockActionManagerWarp(action_wp[0], action_wp[1]) env.num_envs = NUM_ENVS env.device = DEVICE + env._warp_launch = WarpLaunchCache(device=DEVICE) env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 @@ -288,7 +272,9 @@ def test_push_by_setting_velocity(self, warp_env, art_data, all_joints_cfg): wp.capture_launch(cap.graph) wp.synchronize() - scratch = wp.to_torch(warp_evt.push_by_setting_velocity._scratch_vel) + range_values = tuple(zero_range[key] for key in ("x", "y", "z", "roll", "pitch", "yaw")) + (scratch_vel,) = warp_env._warp_mdp_event_scratch[(warp_evt.push_by_setting_velocity, "robot", *range_values)] + scratch = wp.to_torch(scratch_vel) expected = torch.tensor([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], device=DEVICE).expand(NUM_ENVS, -1) assert_close(scratch, expected) @@ -305,8 +291,11 @@ def test_apply_external_force_torque(self, warp_env, art_data, all_joints_cfg): wp.capture_launch(cap.graph) wp.synchronize() - forces = wp.to_torch(warp_evt.apply_external_force_torque._scratch_forces) - torques = wp.to_torch(warp_evt.apply_external_force_torque._scratch_torques) + scratch_forces, scratch_torques = warp_env._warp_mdp_event_scratch[ + (warp_evt.apply_external_force_torque, "robot", (0.0, 0.0), (0.0, 0.0)) + ] + forces = wp.to_torch(scratch_forces) + torques = wp.to_torch(scratch_torques) assert_close(forces, torch.zeros_like(forces)) assert_close(torques, torch.zeros_like(torques)) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_observations_warp_parity.py index 9f1eea23bf8d..460736d75703 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.warp import WarpLaunchCache + # Skip entire module if no CUDA device available wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") @@ -127,6 +129,7 @@ class _Env: env.action_manager = MockActionManagerWarp(action_wp[0], action_wp[1]) env.num_envs = NUM_ENVS env.device = DEVICE + env._warp_launch = WarpLaunchCache(device=DEVICE) env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 @@ -165,6 +168,7 @@ class _Env: env.command_manager = MockCommandManager(cmd_tensor, cmd_term) env.num_envs = NUM_ENVS env.device = DEVICE + env._warp_launch = WarpLaunchCache(device=DEVICE) env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length = 500 @@ -283,6 +287,25 @@ def test_joint_vel_subset(self, warp_env, stable_env, subset_cfg): assert_close(actual, expected) assert_close(actual_cap, expected) + @pytest.mark.parametrize( + ("first_fn", "second_fn", "stable_second_fn"), + [ + (warp_obs.joint_pos, warp_obs.joint_vel, stable_obs.joint_vel), + (warp_obs.joint_pos_rel, warp_obs.joint_vel_rel, stable_obs.joint_vel_rel), + ], + ) + def test_joint_helpers_reuse_output( + self, warp_env, stable_env, all_joints_cfg, first_fn, second_fn, stable_second_fn + ): + """Distinct public helpers sharing a kernel may safely reuse one output buffer.""" + warp_env._warp_launch = WarpLaunchCache(device=DEVICE, debug=False) + out = wp.empty((NUM_ENVS, NUM_JOINTS), dtype=wp.float32, device=DEVICE) + + first_fn(warp_env, out, asset_cfg=all_joints_cfg) + second_fn(warp_env, out, asset_cfg=all_joints_cfg) + + assert_close(wp.to_torch(out).clone(), stable_second_fn(stable_env, asset_cfg=all_joints_cfg)) + # -- Normalized joint position ---------------------------------------------- def test_joint_pos_limit_normalized(self, warp_env, stable_env, all_joints_cfg): diff --git a/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_rewards_warp_parity.py index 69cabcd14a36..b5edf0583053 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.warp import WarpLaunchCache + # Skip entire module if no CUDA device available wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") @@ -138,6 +140,7 @@ class _Env: env.termination_manager = term_mgr env.num_envs = NUM_ENVS env.device = DEVICE + env._warp_launch = WarpLaunchCache(device=DEVICE) env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 @@ -177,6 +180,7 @@ class _Env: env.command_manager = MockCommandManager(cmd_tensor, cmd_term) env.num_envs = NUM_ENVS env.device = DEVICE + env._warp_launch = WarpLaunchCache(device=DEVICE) env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length = 500 @@ -299,6 +303,17 @@ def test_joint_torques_l2(self, warp_env, stable_env, all_joints_cfg): assert_close(actual, expected) assert_close(actual_cap, expected) + def test_sum_squared_helpers_reuse_output(self, warp_env, stable_env, all_joints_cfg): + """Distinct public reductions sharing a kernel may safely reuse one output buffer.""" + warp_env._warp_launch = WarpLaunchCache(device=DEVICE, debug=False) + out = wp.empty((NUM_ENVS,), dtype=wp.float32, device=DEVICE) + + warp_rew.joint_torques_l2(warp_env, out, asset_cfg=all_joints_cfg) + warp_rew.joint_vel_l2(warp_env, out, asset_cfg=all_joints_cfg) + warp_rew.joint_acc_l2(warp_env, out, asset_cfg=all_joints_cfg) + + assert_close(wp.to_torch(out).clone(), stable_rew.joint_acc_l2(stable_env, asset_cfg=all_joints_cfg)) + # -- Joint L1 penalties (masked) -------------------------------------------- def test_joint_vel_l1(self, warp_env, stable_env, all_joints_cfg): @@ -368,6 +383,19 @@ def test_track_lin_vel_xy_exp(self, warp_env_bodies, stable_env_bodies, body_cfg assert_close(actual, expected) assert_close(actual_cap, expected) + def test_track_lin_vel_xy_exp_rerecords_changed_std(self, warp_env_bodies, stable_env_bodies): + """A changed scalar term parameter records a command with the new value.""" + cfg = MockBodyCfg("robot") + cfg.joint_ids = list(range(NUM_JOINTS)) + warp_env_bodies._warp_launch = WarpLaunchCache(device=DEVICE, debug=False) + out = wp.empty((NUM_ENVS,), dtype=wp.float32, device=DEVICE) + + warp_rew.track_lin_vel_xy_exp(warp_env_bodies, out, std=0.25, command_name="vel", asset_cfg=cfg) + warp_rew.track_lin_vel_xy_exp(warp_env_bodies, out, std=0.75, command_name="vel", asset_cfg=cfg) + + expected = stable_rew.track_lin_vel_xy_exp(stable_env_bodies, std=0.75, command_name="vel", asset_cfg=cfg) + assert_close(wp.to_torch(out).clone(), expected) + def test_track_ang_vel_z_exp(self, warp_env_bodies, stable_env_bodies, body_cfg): cfg = MockBodyCfg("robot") cfg.joint_ids = list(range(NUM_JOINTS)) diff --git a/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py b/source/isaaclab_experimental/test/envs/mdp/test_terminations_warp_parity.py index d82339e01228..2f77ed8a3be8 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 @@ -12,6 +12,8 @@ import torch import warp as wp +from isaaclab.utils.warp import WarpLaunchCache + # Skip entire module if no CUDA device available wp.init() pytestmark = pytest.mark.skipif(not wp.is_cuda_available(), reason="CUDA device required") @@ -131,6 +133,7 @@ class _Env: env.action_manager = MockActionManagerWarp(action_wp[0], action_wp[1]) env.num_envs = NUM_ENVS env.device = DEVICE + env._warp_launch = WarpLaunchCache(device=DEVICE) env.episode_length_buf = episode_length_buf env.step_dt = 0.02 env.max_episode_length_s = 10.0 @@ -169,6 +172,7 @@ class _Env: env.command_manager = MockCommandManager(cmd_tensor, cmd_term) env.num_envs = NUM_ENVS env.device = DEVICE + env._warp_launch = WarpLaunchCache(device=DEVICE) env.episode_length_buf = episode_length_buf env._episode_length_buf_wp = wp.from_torch(episode_length_buf) env.step_dt = 0.02 @@ -230,6 +234,17 @@ def test_root_height_below_minimum(self, warp_env, stable_env, all_joints_cfg): assert_equal(actual, expected) assert_equal(actual_cap, expected) + def test_root_height_rerecords_changed_minimum(self, warp_env, stable_env, all_joints_cfg): + """A changed scalar term parameter records a command with the new value.""" + warp_env._warp_launch = WarpLaunchCache(device=DEVICE, debug=False) + out = wp.empty((NUM_ENVS,), dtype=wp.bool, device=DEVICE) + + warp_term.root_height_below_minimum(warp_env, out, minimum_height=-100.0, asset_cfg=all_joints_cfg) + warp_term.root_height_below_minimum(warp_env, out, minimum_height=100.0, asset_cfg=all_joints_cfg) + + expected = stable_term.root_height_below_minimum(stable_env, minimum_height=100.0, asset_cfg=all_joints_cfg) + assert_equal(wp.to_torch(out).clone(), expected) + def test_joint_pos_out_of_manual_limit(self, warp_env, stable_env, all_joints_cfg): cfg = all_joints_cfg bounds = (-1.0, 1.0) diff --git a/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py new file mode 100644 index 000000000000..1d0a660a7d78 --- /dev/null +++ b/source/isaaclab_experimental/test/envs/test_manager_based_rl_env_warp.py @@ -0,0 +1,250 @@ +# 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 experimental Warp manager-based RL environment.""" + +from types import SimpleNamespace + +import pytest +import torch +import warp as wp +from isaaclab_experimental.envs.direct_rl_env_warp import DirectRLEnvWarp +from isaaclab_experimental.envs.manager_based_env_warp import ManagerBasedEnvWarp +from isaaclab_experimental.envs.manager_based_rl_env_warp import ManagerBasedRLEnvWarp +from isaaclab_experimental.managers.command_manager import CommandManager as ExperimentalCommandManager +from isaaclab_experimental.managers.command_manager import _CommandViewCache +from isaaclab_experimental.managers.reward_manager import RewardManager + +from isaaclab.managers import CommandManager +from isaaclab.utils.warp import WarpLaunchCache + + +@pytest.mark.parametrize("env_type", [ManagerBasedEnvWarp, DirectRLEnvWarp]) +def test_physics_ready_destroys_graphs_before_resetting_launch_cache(env_type, monkeypatch): + """A full rebind should release graphs before delegating the cache drain.""" + calls = [] + env = object.__new__(env_type) + env.sim = SimpleNamespace(device="cuda:0") + env._warp_launch = SimpleNamespace(reset=lambda: calls.append("launch_reset")) + if env_type is ManagerBasedEnvWarp: + env._manager_call_switch = SimpleNamespace(invalidate_graphs=lambda: calls.append("graph_invalidate")) + else: + env._graph_cache = SimpleNamespace(invalidate=lambda: calls.append("graph_invalidate")) + monkeypatch.setattr(wp, "synchronize_device", lambda device: pytest.fail(f"unexpected direct sync on {device}")) + + env._reset_warp_caches_after_physics_ready(None) + + assert calls == ["graph_invalidate", "launch_reset"] + + +def test_public_graph_invalidation_also_resets_recorded_launches(monkeypatch): + """Topology changes should invalidate graphs before delegating the cache drain.""" + calls = [] + env = object.__new__(ManagerBasedRLEnvWarp) + env.sim = SimpleNamespace(device="cuda:0") + env._manager_call_switch = SimpleNamespace(invalidate_graphs=lambda: calls.append("graph_invalidate")) + env._warp_launch = SimpleNamespace(reset=lambda: calls.append("launch_reset")) + monkeypatch.setattr(wp, "synchronize_device", lambda device: pytest.fail(f"unexpected direct sync on {device}")) + + env.invalidate_wp_graphs() + + assert calls == ["graph_invalidate", "launch_reset"] + + +def test_environment_uses_warp_command_manager(): + """The Warp environment should instantiate the manager that exposes persistent Warp commands.""" + manager_type = ManagerBasedRLEnvWarp._CommandManagerWarpView + assert issubclass(manager_type, CommandManager) + assert hasattr(manager_type, "get_command_wp") + assert "_CommandManagerWarpView" in ManagerBasedRLEnvWarp.load_managers.__code__.co_names + + +def test_warp_command_manager_stages_computed_commands(monkeypatch): + """Command properties that allocate should be copied into a pointer-stable staging buffer.""" + + class AllocatingCommandTerm: + def __init__(self): + self.position = torch.zeros((2, 2)) + self.heading = torch.zeros((2, 1)) + + @property + def command(self) -> torch.Tensor: + return torch.cat((self.position, self.heading), dim=1) + + def compute(self, dt: float): + self.position.add_(dt) + self.heading.add_(2.0 * dt) + + def reset(self, env_ids=None) -> dict: + self.position.fill_(3.0) + self.heading.fill_(5.0) + return {} + + term = AllocatingCommandTerm() + + def initialize_command_manager(manager, _cfg, _env): + manager._terms = {"pose": term} + manager._resolve_terms_handle = None + + monkeypatch.setattr(CommandManager, "__init__", initialize_command_manager) + manager = ManagerBasedRLEnvWarp._CommandManagerWarpView(None, None) + command_wp = manager.get_command_wp("pose") + + manager.compute(0.5) + + assert manager.get_command_wp("pose") is command_wp + torch.testing.assert_close(torch.asarray(command_wp), term.command) + + manager.reset() + + assert manager.get_command_wp("pose") is command_wp + torch.testing.assert_close(torch.asarray(command_wp), term.command) + + +def test_experimental_command_manager_refreshes_allocating_commands(): + """The exported manager should refresh the same persistent view after compute and reset.""" + + class AllocatingCommandTerm: + def __init__(self): + self.position = torch.zeros((2, 2)) + self.heading = torch.zeros((2, 1)) + + @property + def command(self) -> torch.Tensor: + return torch.cat((self.position, self.heading), dim=1) + + def compute(self, dt: float): + self.position.add_(dt) + self.heading.add_(2.0 * dt) + + def reset(self, env_mask: wp.array): + del env_mask + self.position.fill_(4.0) + self.heading.fill_(6.0) + + term = AllocatingCommandTerm() + manager = object.__new__(ExperimentalCommandManager) + manager._terms = {"pose": term} + manager._command_views = _CommandViewCache(["pose"], lambda name: manager._terms[name].command) + manager._commands_wp = manager._command_views.views + manager._reset_extras = {} + manager._env = SimpleNamespace(num_envs=2, device="cpu") + command_wp = manager.get_command_wp("pose") + + manager.compute(0.5) + assert manager.get_command_wp("pose") is command_wp + torch.testing.assert_close(torch.asarray(command_wp), term.command) + + manager.reset(env_mask=wp.ones(2, dtype=wp.bool, device="cpu")) + assert manager.get_command_wp("pose") is command_wp + torch.testing.assert_close(torch.asarray(command_wp), term.command) + + +def test_command_view_cache_refreshes_late_pointer_replacements(): + """A command that starts pointer-stable should remain current after replacing its tensor.""" + + class ReplacingCommandTerm: + def __init__(self): + self.command = torch.zeros((2, 3)) + + term = ReplacingCommandTerm() + views = _CommandViewCache(["pose"], lambda name: term.command) + command_wp = views.get("pose") + + term.command = torch.full_like(term.command, 7.0) + views.refresh() + + assert views.get("pose") is command_wp + torch.testing.assert_close(torch.asarray(command_wp), term.command) + + +def test_command_view_cache_refreshes_in_place_storage_replacements(): + """A tensor that repoints itself should not leave its persistent Warp view stale.""" + command = torch.zeros((2, 3)) + views = _CommandViewCache(["pose"], lambda _name: command) + command_wp = views.get("pose") + + command.set_(torch.full_like(command, 7.0)) + views.refresh() + + assert views.get("pose") is command_wp + torch.testing.assert_close(torch.asarray(command_wp), command) + + +def test_command_view_cache_rejects_layout_replacements(): + """Rebinding to the same storage with a new layout should fail clearly.""" + command = torch.arange(4, dtype=torch.float32).reshape(2, 2) + views = _CommandViewCache(["pose"], lambda _name: command) + views.get("pose") + + command.set_(command.T) + + with pytest.raises(RuntimeError, match="changed tensor specialization"): + views.refresh() + + +def test_command_view_cache_defers_unused_command_properties(): + """Constructing a manager should not evaluate a command until its Warp view is requested.""" + calls = 0 + + def get_null_command(name: str): + nonlocal calls + calls += 1 + raise RuntimeError(f"{name} does not generate a command") + + views = _CommandViewCache(["null"], get_null_command) + + assert calls == 0 + with pytest.raises(RuntimeError, match="does not generate a command"): + views.get("null") + assert calls == 1 + + +@pytest.mark.parametrize( + "command", + [ + torch.zeros((2, 3), dtype=torch.float64), + wp.zeros((2, 3), dtype=wp.int32, device="cpu"), + ], +) +def test_command_view_cache_rejects_non_float32_commands(command): + """The public float32 Warp-view contract should reject incompatible commands.""" + views = _CommandViewCache(["command"], lambda _name: command) + + with pytest.raises(TypeError, match="must use .*float32"): + views.get("command") + + +def test_reward_manager_replay_updates_dt(): + """Reward finalization should use the current time step on every replay.""" + + def fill_reward(_env, out: wp.array(dtype=wp.float32)): + out.fill_(2.0) + + num_envs = 4 + env = SimpleNamespace(num_envs=num_envs, device="cpu", _warp_launch=WarpLaunchCache(device="cpu")) + manager = object.__new__(RewardManager) + manager._env = env + manager._num_terms = 1 + manager._term_outs_wp = wp.zeros((1, num_envs), dtype=wp.float32, device="cpu") + term_out = wp.array( + ptr=manager._term_outs_wp.ptr, + dtype=wp.float32, + shape=(num_envs,), + strides=(manager._term_outs_wp.strides[1],), + device="cpu", + ) + manager._term_cfgs = [SimpleNamespace(weight=3.0, func=fill_reward, out=term_out, params={})] + manager._reward_wp = wp.zeros(num_envs, dtype=wp.float32, device="cpu") + manager._step_reward_wp = wp.zeros((num_envs, 1), dtype=wp.float32, device="cpu") + manager._term_weights_wp = wp.array([3.0], dtype=wp.float32, device="cpu") + manager._episode_sums_wp = wp.zeros((1, num_envs), dtype=wp.float32, device="cpu") + manager._reward_tensor_view = wp.to_torch(manager._reward_wp) + + reward_half_second = manager.compute(0.5).clone() + reward_quarter_second = manager.compute(0.25).clone() + + torch.testing.assert_close(reward_half_second, torch.full((num_envs,), 3.0)) + torch.testing.assert_close(reward_quarter_second, torch.full((num_envs,), 1.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 index b4745939eaa0..5e8094df370b 100644 --- a/source/isaaclab_experimental/test/utils/test_manager_call_switch.py +++ b/source/isaaclab_experimental/test/utils/test_manager_call_switch.py @@ -14,6 +14,8 @@ import warp as wp from isaaclab_experimental.utils.manager_call_switch import ManagerCallMode, ManagerCallSwitch +from isaaclab.utils.warp import WarpLaunchCache + @wp.kernel def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)): @@ -270,6 +272,33 @@ def warp_fn(a, b, key=None): ) self.assertEqual(received, {"a": 1, "b": 2, "key": "val"}) + def test_uncapturable_manager_still_replays_neighboring_launch(self): + """A manager downgrade should preserve launch-level replay for its other terms.""" + switch = ManagerCallSwitch(cfg_source={"default": 2}) + switch.register_manager_capturability("RewardManager", capturable=False) + cache = WarpLaunchCache(device="cpu") + source = wp.full(4, value=1.0, dtype=wp.float32, device="cpu") + output = wp.zeros(4, dtype=wp.float32, device="cpu") + uncapturable_call_count = 0 + + def manager_stage(): + nonlocal uncapturable_call_count + uncapturable_call_count += 1 + cache.launch(_add_one, dim=4, inputs=[source, output]) + + switch.call_stage(stage="RewardManager_compute", warp_call={"fn": manager_stage}) + command = next(iter(cache._entries.values())).command + self.assertEqual(output.numpy().tolist(), [2.0] * 4) + + source.fill_(5.0) + switch.call_stage(stage="RewardManager_compute", warp_call={"fn": manager_stage}) + + self.assertEqual(switch.get_mode_for_manager("RewardManager"), ManagerCallMode.WARP_NOT_CAPTURED) + self.assertEqual(uncapturable_call_count, 2) + self.assertEqual(len(cache._entries), 1) + self.assertIs(next(iter(cache._entries.values())).command, command) + self.assertEqual(output.numpy().tolist(), [6.0] * 4) + class TestStageDispatchCaptured(unittest.TestCase): """Tests for WARP_CAPTURED mode (requires GPU).""" @@ -362,8 +391,8 @@ def fn_with_guard(): for val in result2.numpy(): self.assertAlmostEqual(val, 6.0, places=5) - def test_warp_eager_no_warmup(self): - """WARP_NOT_CAPTURED mode should call fn exactly once per call (no warm-up).""" + def test_warp_uncaptured_has_no_stage_warmup(self): + """WARP_NOT_CAPTURED should call the stage once while allowing launch-level replay inside it.""" switch = ManagerCallSwitch(cfg_source={"default": 1}) call_count = [0] @@ -371,10 +400,10 @@ def warp_fn(): call_count[0] += 1 switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 1, "Eager mode should call fn exactly once") + self.assertEqual(call_count[0], 1, "Uncaptured mode should call the stage exactly once") switch.call_stage(stage="RewardManager_compute", warp_call={"fn": warp_fn}) - self.assertEqual(call_count[0], 2, "Eager mode should call fn again each time") + self.assertEqual(call_count[0], 2, "Uncaptured mode should call the stage again each time") # ====================================================================== diff --git a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py index 1b9c10a76895..8cbaaf558fc6 100644 --- a/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py +++ b/source/isaaclab_experimental/test/utils/test_warp_graph_cache.py @@ -8,10 +8,13 @@ from __future__ import annotations import unittest +from unittest.mock import patch import warp as wp from isaaclab_experimental.utils.warp_graph_cache import WarpGraphCache +from isaaclab.utils.warp import WarpLaunchCache + @wp.kernel def _add_one(a: wp.array(dtype=wp.float32), b: wp.array(dtype=wp.float32)): @@ -223,6 +226,54 @@ def test_invalidate_nonexistent_stage_is_noop(self): """Invalidating a stage that was never captured should not raise.""" self.cache.invalidate("nonexistent") # should not raise + def test_graph_invalidation_drains_when_launch_cache_is_eager(self): + """A graph should drain even when the neighboring launch cache has no commands.""" + launch_cache = WarpLaunchCache(device=self.device, enabled=False) + src = wp.zeros(4, 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) + + self.cache.capture_or_replay("stage", my_fn) + graph_device = self.cache._graphs["stage"].device + + with patch.object(wp, "synchronize_device", wraps=wp.synchronize_device) as synchronize: + self.cache.invalidate() + launch_cache.reset() + + synchronize.assert_called_once_with(graph_device) + self.assertFalse(self.cache._graphs) + self.assertFalse(launch_cache._entries) + + def test_empty_graph_and_launch_caches_do_not_synchronize(self): + """Empty eager caches should not introduce a lifecycle synchronization.""" + graph_cache = WarpGraphCache(mode="eager") + launch_cache = WarpLaunchCache(device=self.device, enabled=False) + + with patch.object(wp, "synchronize_device", side_effect=AssertionError("unexpected sync")): + graph_cache.invalidate() + launch_cache.reset() + + def test_eager_mode_invokes_function_on_every_call(self): + """Eager mode should bypass graph capture without changing call results.""" + cache = WarpGraphCache(mode="eager") + call_count = 0 + + def counted_fn(value): + nonlocal call_count + call_count += 1 + return value + + self.assertEqual(cache.capture_or_replay("stage", counted_fn, args=(1,)), 1) + self.assertEqual(cache.capture_or_replay("stage", counted_fn, args=(2,)), 2) + self.assertEqual(call_count, 2) + + def test_invalid_mode_is_rejected(self): + """An unknown graph execution mode should fail during construction.""" + with self.assertRaisesRegex(ValueError, "Invalid Warp graph mode"): + WarpGraphCache(mode="invalid") + # ------------------------------------------------------------------ # Args / kwargs forwarding # ------------------------------------------------------------------ diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-launch-record-replay.minor.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-launch-record-replay.minor.rst new file mode 100644 index 000000000000..945b017676b6 --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-launch-record-replay.minor.rst @@ -0,0 +1,6 @@ +Added +^^^^^ + +* Added recorded launch replay to reusable Cartpole, locomotion, and in-hand + manipulation Warp task kernels. Set + ``ISAACLAB_WARP_LAUNCH_MODE=eager`` to restore eager launches for comparison. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py index 14f776c95b0b..3703cb692dcb 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/cartpole/cartpole_warp_env.py @@ -221,7 +221,7 @@ def _setup_scene(self) -> None: def _pre_physics_step(self, actions: wp.array) -> None: wp.launch( - update_actions, + kernel=update_actions, dim=self.num_envs, inputs=[ actions, @@ -229,13 +229,14 @@ def _pre_physics_step(self, actions: wp.array) -> None: self.action_scale, self._cart_dof_idx[0], ], + device=self.device, ) def _apply_action(self) -> None: self.cartpole.set_joint_effort_target_mask(target=self.actions) def _get_observations(self) -> dict: - wp.launch( + self._warp_launch.launch( get_observations, dim=self.num_envs, inputs=[ @@ -249,7 +250,7 @@ def _get_observations(self) -> dict: return {"policy": self.torch_obs_buf} def _get_rewards(self) -> None: - wp.launch( + self._warp_launch.launch( compute_rewards, dim=self.num_envs, inputs=[ @@ -268,7 +269,7 @@ def _get_rewards(self) -> None: ) def _get_dones(self) -> None: - wp.launch( + self._warp_launch.launch( get_dones, dim=self.num_envs, inputs=[ @@ -290,7 +291,7 @@ def _reset_idx(self, mask: wp.array | None = None) -> None: super()._reset_idx(mask) - wp.launch( + self._warp_launch.launch( reset, dim=self.num_envs, inputs=[ @@ -304,4 +305,5 @@ def _reset_idx(self, mask: wp.array | None = None) -> None: mask, self.states, ], + site=mask, ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py index 2082d470ff9b..b648120230ad 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/inhand_manipulation/inhand_manipulation_warp_env.py @@ -596,6 +596,7 @@ def __init__(self, cfg: AllegroHandWarpEnvCfg, render_mode: str | None = None, * # track goal resets self.reset_goal_buf = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) + self._target_pose_env_mask = wp.zeros(self.num_envs, dtype=wp.bool, device=self.device) # used to compare object position self.in_hand_pos = wp.zeros(self.num_envs, dtype=wp.vec3f, device=self.device) # default goal positions @@ -690,7 +691,7 @@ def _pre_physics_step(self, actions: wp.array) -> None: wp.copy(self.actions, actions) def _apply_action(self) -> None: - wp.launch( + self._warp_launch.launch( apply_actions_to_targets, dim=(self.num_envs, self.num_actuated_dofs), inputs=[ @@ -702,7 +703,6 @@ def _apply_action(self) -> None: self.prev_targets, self.cur_targets, ], - device=self.device, ) # Apply position targets using mask method (CUDA graph safe). @@ -728,7 +728,7 @@ def _get_rewards(self) -> None: # wp.assign(self._finished_cons_successes, 0.0) self._num_resets.zero_() self._finished_cons_successes.zero_() - wp.launch( + self._warp_launch.launch( compute_rewards, dim=self.num_envs, inputs=[ @@ -752,11 +752,10 @@ def _get_rewards(self) -> None: self._finished_cons_successes, self.rewards, ], - device=self.device, ) # A separate kernel is needed as Warp does not support thread synchronization for reductions. - wp.launch( + self._warp_launch.launch( update_consecutive_successes_from_stats, dim=1, inputs=[ @@ -765,7 +764,6 @@ def _get_rewards(self) -> None: self.cfg.av_factor, self.consecutive_successes, ], - device=self.device, ) if "log" not in self.extras: @@ -780,7 +778,7 @@ def _get_rewards(self) -> None: def _get_dones(self) -> None: self._compute_intermediate_values() - wp.launch( + self._warp_launch.launch( get_dones, dim=self.num_envs, inputs=[ @@ -797,7 +795,6 @@ def _get_dones(self) -> None: self.reset_time_outs, self.reset_buf, ], - device=self.device, ) def _reset_idx(self, mask: wp.array | None = None): @@ -811,7 +808,7 @@ def _reset_idx(self, mask: wp.array | None = None): self._reset_target_pose(mask=mask) # reset object - wp.launch( + self._warp_launch.launch( reset_object, dim=self.num_envs, inputs=[ @@ -825,11 +822,11 @@ def _reset_idx(self, mask: wp.array | None = None): self.object.data.root_link_pose_w.warp, self.object.data.root_com_vel_w.warp, ], - device=self.device, + site=mask, ) # reset hand - wp.launch( + self._warp_launch.launch( reset_hand, dim=self.num_envs, inputs=[ @@ -848,19 +845,19 @@ def _reset_idx(self, mask: wp.array | None = None): self.cur_targets, self.hand_dof_targets, ], - device=self.device, + site=mask, ) self.hand.set_joint_position_target_mask(target=self.cur_targets, env_mask=mask) - wp.launch( + self._warp_launch.launch( reset_successes, dim=self.num_envs, inputs=[ mask, self.successes, ], - device=self.device, + site=mask, ) self._compute_intermediate_values() @@ -873,10 +870,11 @@ def _reset_target_pose(self, env_ids: Sequence[int] | None = None, mask: wp.arra 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) + self._target_pose_env_mask.assign(wp.array(env_mask_list, dtype=wp.bool, device=self.device)) + mask = self._target_pose_env_mask # update goal pose and markers - wp.launch( + self._warp_launch.launch( reset_target_pose, dim=self.num_envs, inputs=[ @@ -890,7 +888,7 @@ def _reset_target_pose(self, env_ids: Sequence[int] | None = None, mask: wp.arra self.reset_goal_buf, self.goal_pos_w, ], - device=self.device, + site=mask, ) def _post_step_visualize(self) -> None: @@ -899,7 +897,7 @@ def _post_step_visualize(self) -> None: def _compute_intermediate_values(self): # data for hand/object (Warp version of the Torch env's `_compute_intermediate_values`) - wp.launch( + self._warp_launch.launch( compute_intermediate_values, dim=self.num_envs, inputs=[ @@ -917,7 +915,6 @@ def _compute_intermediate_values(self): self.object_pose, self.object_vels, ], - device=self.device, ) def compute_reduced_observations(self): @@ -925,7 +922,7 @@ def compute_reduced_observations(self): # Fingertip positions # Object Position, but not orientation # Relative target orientation - wp.launch( + self._warp_launch.launch( compute_reduced_observations, dim=self.num_envs, inputs=[ @@ -937,19 +934,18 @@ def compute_reduced_observations(self): self.cfg.action_space, self.observations, ], - device=self.device, ) # Warp-native non-finite sanitization + print-once. - wp.launch( + self._warp_launch.launch( sanitize_and_print_once, dim=(self.num_envs * self.cfg.observation_space), inputs=[self.observations.flatten(), self.obs_nonfinite_flag], - device=self.device, + site="reduced_observations_sanitize", ) self.obs_nonfinite_flag.zero_() def compute_full_observations(self): - wp.launch( + self._warp_launch.launch( compute_full_observations, dim=self.num_envs, inputs=[ @@ -971,13 +967,12 @@ def compute_full_observations(self): self.cfg.action_space, self.observations, ], - device=self.device, ) # Warp-native non-finite sanitization + print-once. - wp.launch( + self._warp_launch.launch( sanitize_and_print_once, dim=(self.num_envs * self.cfg.observation_space), inputs=[self.observations.flatten(), self.obs_nonfinite_flag], - device=self.device, + site="full_observations_sanitize", ) self.obs_nonfinite_flag.zero_() diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py index 62400c4d48fe..c253414d1c94 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/direct/locomotion/locomotion_env_warp.py @@ -428,17 +428,17 @@ def _setup_scene(self) -> None: def _pre_physics_step(self, actions: wp.array) -> None: self.actions.assign(actions) - wp.launch( + self._warp_launch.launch( update_actions, dim=(self.num_envs, self.robot.num_joints), - inputs=[actions, self.actions_mapped, self.joint_gears, self.action_scale], + inputs=[self.actions, self.actions_mapped, self.joint_gears, self.action_scale], ) def _apply_action(self) -> None: self.robot.set_joint_effort_target_mask(target=self.actions_mapped) def _compute_intermediate_values(self) -> None: - wp.launch( + self._warp_launch.launch( compute_heading_and_up, dim=self.num_envs, inputs=[ @@ -455,7 +455,7 @@ def _compute_intermediate_values(self) -> None: ], ) - wp.launch( + self._warp_launch.launch( compute_rot, dim=self.num_envs, inputs=[ @@ -467,7 +467,7 @@ def _compute_intermediate_values(self) -> None: self.angle_to_target, ], ) - wp.launch( + self._warp_launch.launch( scale_dof_pos, dim=(self.num_envs, self.robot.num_joints), inputs=[ @@ -478,7 +478,7 @@ def _compute_intermediate_values(self) -> None: ) def _get_observations(self) -> dict: - wp.launch( + self._warp_launch.launch( observations, dim=self.num_envs, inputs=[ @@ -500,7 +500,7 @@ def _get_observations(self) -> dict: return {"policy": self.torch_obs_buf} def _get_rewards(self) -> None: - wp.launch( + self._warp_launch.launch( compute_rewards, dim=self.num_envs, inputs=[ @@ -527,7 +527,7 @@ def _get_rewards(self) -> None: def _get_dones(self) -> None: self._compute_intermediate_values() - wp.launch( + self._warp_launch.launch( get_dones, dim=self.num_envs, inputs=[ @@ -547,7 +547,7 @@ def _reset_idx(self, mask: wp.array | None = None): super()._reset_idx(mask) - wp.launch( + self._warp_launch.launch( reset_root, dim=self.num_envs, inputs=[ @@ -562,8 +562,9 @@ def _reset_idx(self, mask: wp.array | None = None): self.root_vel_w, mask, ], + site=mask, ) - wp.launch( + self._warp_launch.launch( reset_joints, dim=(self.num_envs, self.robot.num_joints), inputs=[ @@ -573,6 +574,7 @@ def _reset_idx(self, mask: wp.array | None = None): self.joint_vel, mask, ], + site=mask, ) self._compute_intermediate_values() diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py index 2678ef03de65..ee6449ab76e8 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/cartpole/mdp/rewards.py @@ -37,9 +37,9 @@ def joint_pos_target_l2(env: ManagerBasedRLEnv, out, target: float, asset_cfg: S """Penalize joint position deviation from a target value. Writes into ``out``.""" asset: Articulation = env.scene[asset_cfg.name] assert asset.data.joint_pos.warp.shape[1] == asset_cfg.joint_mask.shape[0] - wp.launch( - kernel=_joint_pos_target_l2_kernel, + env._warp_launch.launch( + _joint_pos_target_l2_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, out, target], - device=env.device, + inputs=[asset.data.joint_pos.warp, asset_cfg.joint_mask, out, float(target)], + site=(out, float(target)), ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py index 4d9ab39b4ff7..9488110ab48d 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/observations.py @@ -51,11 +51,11 @@ def _base_yaw_roll_kernel( def base_yaw_roll(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Yaw and roll of the base in the simulation world frame. Shape: (num_envs, 2).""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_base_yaw_roll_kernel, + env._warp_launch.launch( + _base_yaw_roll_kernel, dim=env.num_envs, inputs=[asset.data.root_quat_w.warp, out], - device=env.device, + site=out, ) @@ -79,11 +79,11 @@ def _base_up_proj_kernel( def base_up_proj(env: ManagerBasedEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> None: """Projection of the base up vector onto the world up vector. Shape: (num_envs, 1).""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_base_up_proj_kernel, + env._warp_launch.launch( + _base_up_proj_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, asset.data.GRAVITY_VEC_W.warp, out], - device=env.device, + site=out, ) @@ -124,8 +124,8 @@ def base_heading_proj( ) -> None: """Dot product between the base forward direction and direction to target. Shape: (num_envs, 1).""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_base_heading_proj_kernel, + env._warp_launch.launch( + _base_heading_proj_kernel, dim=env.num_envs, inputs=[ asset.data.root_pos_w.warp, @@ -135,7 +135,7 @@ def base_heading_proj( target_pos[2], out, ], - device=env.device, + site=(out, tuple(target_pos)), ) @@ -177,9 +177,9 @@ def base_angle_to_target( ) -> None: """Angle between the base forward vector and the vector to the target. Shape: (num_envs, 1).""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_base_angle_to_target_kernel, + env._warp_launch.launch( + _base_angle_to_target_kernel, dim=env.num_envs, inputs=[asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, target_pos[0], target_pos[1], out], - device=env.device, + site=(out, tuple(target_pos)), ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py index a090d3a16f69..b1829e415608 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/classic/humanoid/mdp/rewards.py @@ -57,11 +57,11 @@ def upright_posture_bonus( ) -> None: """Reward for maintaining an upright posture. Writes 1.0 if up_proj > threshold, else 0.0.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_upright_posture_bonus_kernel, + env._warp_launch.launch( + _upright_posture_bonus_kernel, dim=env.num_envs, inputs=[asset.data.root_link_pose_w.warp, asset.data.GRAVITY_VEC_W.warp, threshold, out], - device=env.device, + site=(out, float(threshold)), ) @@ -99,8 +99,8 @@ def move_to_target_bonus( ) -> None: """Reward for heading towards the target.""" asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_move_to_target_bonus_kernel, + env._warp_launch.launch( + _move_to_target_bonus_kernel, dim=env.num_envs, inputs=[ asset.data.root_pos_w.warp, @@ -110,7 +110,7 @@ def move_to_target_bonus( threshold, out, ], - device=env.device, + site=(out, float(threshold), tuple(target_pos)), ) @@ -174,8 +174,8 @@ def reset(self, env_mask: wp.array | None = None) -> None: return asset: Articulation = self._env.scene["robot"] inv_dt = 1.0 / self._env.step_dt - wp.launch( - kernel=_progress_reward_reset_kernel, + self._env._warp_launch.launch( + _progress_reward_reset_kernel, dim=self.num_envs, inputs=[ env_mask, @@ -186,7 +186,7 @@ def reset(self, env_mask: wp.array | None = None) -> None: inv_dt, self.potentials, ], - device=self.device, + site=(self, env_mask, tuple(self._target_pos), float(inv_dt)), ) def __call__( @@ -198,11 +198,11 @@ def __call__( ) -> None: asset: Articulation = env.scene[asset_cfg.name] inv_dt = 1.0 / env.step_dt - wp.launch( - kernel=_progress_reward_kernel, + env._warp_launch.launch( + _progress_reward_kernel, dim=env.num_envs, inputs=[asset.data.root_pos_w.warp, target_pos[0], target_pos[1], inv_dt, self.potentials, out], - device=env.device, + site=(self, out, tuple(target_pos), float(inv_dt)), ) @@ -259,8 +259,8 @@ def __call__( asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), ) -> None: asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_joint_pos_limits_penalty_ratio_kernel, + env._warp_launch.launch( + _joint_pos_limits_penalty_ratio_kernel, dim=env.num_envs, inputs=[ asset.data.joint_pos.warp, @@ -270,7 +270,7 @@ def __call__( 1.0 / (1.0 - threshold), out, ], - device=env.device, + site=(self, out, float(threshold)), ) @@ -313,9 +313,9 @@ def __call__( asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), ) -> None: asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_power_consumption_kernel, + env._warp_launch.launch( + _power_consumption_kernel, dim=env.num_envs, inputs=[env.action_manager.action, asset.data.joint_vel.warp, self._gear_ratio_scaled_wp, out], - device=env.device, + site=(self, out), ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py index d7e11c517c7d..7452561e8d45 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/rewards.py @@ -6,8 +6,7 @@ """Warp-first reward functions for the velocity locomotion environment. All functions follow the ``func(env, out, **params) -> None`` signature. -Cross-manager torch tensors (contact sensor, commands) are cached as zero-copy -warp views on first call via ``wp.from_torch``. +Cross-manager commands use the manager's persistent Warp views. """ from __future__ import annotations @@ -50,28 +49,22 @@ def _feet_air_time_kernel( def feet_air_time(env: ManagerBasedRLEnv, out, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float) -> None: """Reward long steps taken by the feet using L2-kernel.""" - fn = feet_air_time contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - # Cache command bridge (persistent pointer) - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True + command = env.command_manager.get_command_wp(command_name) # Newton contact sensor returns persistent wp.arrays — use directly, no wp.from_torch needed first_contact = contact_sensor.compute_first_contact(env.step_dt) - wp.launch( - kernel=_feet_air_time_kernel, + env._warp_launch.launch( + _feet_air_time_kernel, dim=env.num_envs, inputs=[ contact_sensor.data.last_air_time.warp, first_contact, sensor_cfg.body_ids_wp, - fn._cmd_wp, + command, threshold, out, ], - device=env.device, + site=(out, command_name, float(threshold)), ) @@ -114,25 +107,20 @@ def _feet_air_time_positive_biped_kernel( def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg) -> None: """Reward long steps taken by the feet for bipeds.""" - fn = feet_air_time_positive_biped contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True - wp.launch( - kernel=_feet_air_time_positive_biped_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _feet_air_time_positive_biped_kernel, dim=env.num_envs, inputs=[ contact_sensor.data.current_air_time.warp, contact_sensor.data.current_contact_time.warp, sensor_cfg.body_ids_wp, - fn._cmd_wp, + command, threshold, out, ], - device=env.device, + site=(out, command_name, float(threshold)), ) @@ -172,8 +160,8 @@ def feet_slide(env, out, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = """Penalize feet sliding.""" contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] asset = env.scene[asset_cfg.name] - wp.launch( - kernel=_feet_slide_kernel, + env._warp_launch.launch( + _feet_slide_kernel, dim=env.num_envs, inputs=[ asset.data.body_lin_vel_w.warp, @@ -182,7 +170,7 @@ def feet_slide(env, out, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = contact_sensor.data.net_forces_w_history.warp.shape[1], out, ], - device=env.device, + site=(out, int(contact_sensor.data.net_forces_w_history.warp.shape[1])), ) @@ -224,18 +212,13 @@ def track_lin_vel_xy_yaw_frame_exp( env, out, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame.""" - fn = track_lin_vel_xy_yaw_frame_exp asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True - wp.launch( - kernel=_track_lin_vel_xy_yaw_frame_exp_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _track_lin_vel_xy_yaw_frame_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w.warp, asset.data.root_lin_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], - device=env.device, + inputs=[asset.data.root_quat_w.warp, asset.data.root_lin_vel_w.warp, command, 1.0 / (std * std), out], + site=(out, command_name, float(std)), ) @@ -260,18 +243,13 @@ def track_ang_vel_z_world_exp( env, out, command_name: str, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Reward tracking of angular velocity commands (yaw) in world frame.""" - fn = track_ang_vel_z_world_exp asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True - wp.launch( - kernel=_track_ang_vel_z_world_exp_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _track_ang_vel_z_world_exp_kernel, dim=env.num_envs, - inputs=[asset.data.root_ang_vel_w.warp, fn._cmd_wp, 1.0 / (std * std), out], - device=env.device, + inputs=[asset.data.root_ang_vel_w.warp, command, 1.0 / (std * std), out], + site=(out, command_name, float(std)), ) @@ -304,16 +282,11 @@ def stand_still_joint_deviation_l1( env, out, command_name: str, command_threshold: float = 0.06, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> None: """Penalize offsets from the default joint positions when the command is very small.""" - fn = stand_still_joint_deviation_l1 asset = env.scene[asset_cfg.name] - if not getattr(fn, "_is_warmed_up", False) or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - fn._is_warmed_up = True - wp.launch( - kernel=_stand_still_joint_deviation_l1_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _stand_still_joint_deviation_l1_kernel, dim=env.num_envs, - inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, fn._cmd_wp, command_threshold, out], - device=env.device, + inputs=[asset.data.joint_pos.warp, asset.data.default_joint_pos.warp, command, command_threshold, out], + site=(out, command_name, float(command_threshold)), ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py index 4e0c32c58e64..d3c2ce897ede 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/locomotion/velocity/mdp/terminations.py @@ -36,31 +36,24 @@ def terrain_out_of_bounds( env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), distance_buffer: float = 3.0 ) -> None: """Terminate when the actor moves too close to the edge of the terrain.""" - fn = terrain_out_of_bounds - if not getattr(fn, "_is_warmed_up", False): - terrain_type = env.scene.cfg.terrain.terrain_type - if terrain_type == "plane": - fn._is_plane = True - elif terrain_type == "generator": - fn._is_plane = False - terrain_gen_cfg = env.scene.terrain.cfg.terrain_generator - grid_width, grid_length = terrain_gen_cfg.size - n_rows, n_cols = terrain_gen_cfg.num_rows, terrain_gen_cfg.num_cols - border_width = terrain_gen_cfg.border_width - fn._half_width = 0.5 * (n_rows * grid_width + 2 * border_width) - fn._half_height = 0.5 * (n_cols * grid_length + 2 * border_width) - else: - raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") - fn._is_warmed_up = True - - if fn._is_plane: + terrain_type = env.scene.cfg.terrain.terrain_type + if terrain_type == "plane": out.zero_() return + if terrain_type != "generator": + raise ValueError("Received unsupported terrain type, must be either 'plane' or 'generator'.") + + terrain_gen_cfg = env.scene.terrain.cfg.terrain_generator + grid_width, grid_length = terrain_gen_cfg.size + n_rows, n_cols = terrain_gen_cfg.num_rows, terrain_gen_cfg.num_cols + border_width = terrain_gen_cfg.border_width + half_width = 0.5 * (n_rows * grid_width + 2 * border_width) + half_height = 0.5 * (n_cols * grid_length + 2 * border_width) asset: Articulation = env.scene[asset_cfg.name] - wp.launch( - kernel=_terrain_out_of_bounds_kernel, + env._warp_launch.launch( + _terrain_out_of_bounds_kernel, dim=env.num_envs, - inputs=[asset.data.root_pos_w.warp, fn._half_width, fn._half_height, distance_buffer, out], - device=env.device, + inputs=[asset.data.root_pos_w.warp, half_width, half_height, distance_buffer, out], + site=(out, float(half_width), float(half_height), float(distance_buffer)), ) diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py index 7435752528d6..d5f984d3ea0c 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/manager_based/manipulation/reach/mdp/rewards.py @@ -6,7 +6,7 @@ """Warp-first reward terms for the reach task. All functions follow the ``func(env, out, **params) -> None`` signature. -Command tensors are cached as zero-copy warp views on first call. +Commands use the manager's persistent Warp views. """ from __future__ import annotations @@ -50,24 +50,20 @@ def _position_command_error_kernel( def position_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: """Penalize tracking of the position error using L2-norm.""" - fn = position_command_error asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - wp.launch( - kernel=_position_command_error_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _position_command_error_kernel, dim=env.num_envs, inputs=[ asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, asset.data.body_pos_w.warp, - fn._cmd_wp, + command, asset_cfg.body_ids[0], out, ], - device=env.device, + site=(out, command_name, int(asset_cfg.body_ids[0])), ) @@ -101,25 +97,21 @@ def position_command_error_tanh( env: ManagerBasedRLEnv, out, std: float, command_name: str, asset_cfg: SceneEntityCfg ) -> None: """Reward tracking of the position using the tanh kernel.""" - fn = position_command_error_tanh asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - wp.launch( - kernel=_position_command_error_tanh_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _position_command_error_tanh_kernel, dim=env.num_envs, inputs=[ asset.data.root_pos_w.warp, asset.data.root_quat_w.warp, asset.data.body_pos_w.warp, - fn._cmd_wp, + command, asset_cfg.body_ids[0], 1.0 / std, out, ], - device=env.device, + site=(out, command_name, int(asset_cfg.body_ids[0]), float(std)), ) @@ -152,15 +144,11 @@ def _orientation_command_error_kernel( def orientation_command_error(env: ManagerBasedRLEnv, out, command_name: str, asset_cfg: SceneEntityCfg) -> None: """Penalize tracking orientation error using shortest path.""" - fn = orientation_command_error asset: Articulation = env.scene[asset_cfg.name] - if not hasattr(fn, "_cmd_wp") or fn._cmd_name != command_name: - cmd = env.command_manager.get_command(command_name) - fn._cmd_wp = cmd if isinstance(cmd, wp.array) else wp.from_torch(cmd) - fn._cmd_name = command_name - wp.launch( - kernel=_orientation_command_error_kernel, + command = env.command_manager.get_command_wp(command_name) + env._warp_launch.launch( + _orientation_command_error_kernel, dim=env.num_envs, - inputs=[asset.data.root_quat_w.warp, asset.data.body_quat_w.warp, fn._cmd_wp, asset_cfg.body_ids[0], out], - device=env.device, + inputs=[asset.data.root_quat_w.warp, asset.data.body_quat_w.warp, command, asset_cfg.body_ids[0], out], + site=(out, command_name, int(asset_cfg.body_ids[0])), ) diff --git a/source/isaaclab_tasks_experimental/test/test_cartpole_warp_rewards.py b/source/isaaclab_tasks_experimental/test/test_cartpole_warp_rewards.py new file mode 100644 index 000000000000..6de237d1ecd6 --- /dev/null +++ b/source/isaaclab_tasks_experimental/test/test_cartpole_warp_rewards.py @@ -0,0 +1,30 @@ +# 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 Cartpole Warp reward launch replay.""" + +from types import SimpleNamespace + +import numpy as np +import warp as wp +from isaaclab_tasks_experimental.manager_based.classic.cartpole.mdp.rewards import joint_pos_target_l2 + +from isaaclab.utils.warp import WarpLaunchCache + + +def test_joint_pos_target_l2_replay_selects_target_site(): + """A changed reward target should use a command recorded for that static value.""" + joint_pos = wp.array([[0.0], [2.0]], dtype=wp.float32, device="cpu") + joint_mask = wp.array([True], dtype=wp.bool, device="cpu") + output = wp.zeros(2, dtype=wp.float32, device="cpu") + asset = SimpleNamespace(data=SimpleNamespace(joint_pos=SimpleNamespace(warp=joint_pos))) + env = SimpleNamespace(scene={"robot": asset}, num_envs=2, device="cpu", _warp_launch=WarpLaunchCache(device="cpu")) + asset_cfg = SimpleNamespace(name="robot", joint_mask=joint_mask) + + joint_pos_target_l2(env, output, target=0.0, asset_cfg=asset_cfg) + np.testing.assert_allclose(output.numpy(), np.array([0.0, 4.0], dtype=np.float32), atol=2.0e-6) + + joint_pos_target_l2(env, output, target=1.0, asset_cfg=asset_cfg) + np.testing.assert_allclose(output.numpy(), np.array([1.0, 1.0], dtype=np.float32), atol=2.0e-6) From 41239d542f790963303b3a93570d743302c251cc Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 01:40:13 -0700 Subject: [PATCH 3/5] Add opt-in Newton launch replay Gate replay for validated pointer-stable articulation, rigid-object, derived-state, and actuator paths. Leave the feature disabled by default while broader workload coverage matures. --- .../warp-launch-record-replay.minor.rst | 6 + .../isaaclab_newton/actuators/adapter.py | 17 +- .../assets/articulation/articulation.py | 73 +++-- .../assets/articulation/articulation_data.py | 110 +++---- .../assets/rigid_object/rigid_object.py | 13 +- .../assets/rigid_object/rigid_object_data.py | 77 +++-- .../isaaclab_newton/physics/newton_manager.py | 32 +- .../physics/newton_manager_cfg.py | 10 + .../test_articulation_warp_launch_cache.py | 281 ++++++++++++++++++ .../test_warp_launch_cache_hot_paths.py | 258 ++++++++++++++++ 10 files changed, 764 insertions(+), 113 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/warp-launch-record-replay.minor.rst create mode 100644 source/isaaclab_newton/test/assets/test_articulation_warp_launch_cache.py create mode 100644 source/isaaclab_newton/test/assets/test_warp_launch_cache_hot_paths.py diff --git a/source/isaaclab_newton/changelog.d/warp-launch-record-replay.minor.rst b/source/isaaclab_newton/changelog.d/warp-launch-record-replay.minor.rst new file mode 100644 index 000000000000..b1e5c54c53ec --- /dev/null +++ b/source/isaaclab_newton/changelog.d/warp-launch-record-replay.minor.rst @@ -0,0 +1,6 @@ +Added +^^^^^ + +* Added :attr:`~isaaclab_newton.physics.NewtonCfg.use_warp_launch_cache` to + enable recorded replay for pointer-stable articulation target publication, + lazy derived-state, and actuator adapter kernels outside the solver graph. diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py index 689232b81e60..667fd7a4dc52 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py +++ b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py @@ -26,6 +26,8 @@ import warp as wp from newton.actuators import Actuator, Clamping, Delay +from isaaclab.utils.warp import WarpLaunchCache + from .kernels import ( build_implicit_dof_mask, build_per_dof_env_mask_kernel, @@ -85,6 +87,8 @@ def __init__( num_joints: int, dof_offset: int, device: str, + *, + warp_launch: WarpLaunchCache | None = None, ): self.actuators = actuators self.num_joints = num_joints @@ -92,6 +96,7 @@ def __init__( self._num_envs = num_envs self._dof_offset = dof_offset self._device = device + self._warp_launch = warp_launch if warp_launch is not None else WarpLaunchCache(device=device, enabled=False) # Collect the set of local DOFs covered by some actuator. Only the # env-0 slice of each actuator's flat ``indices`` array is needed — @@ -157,10 +162,11 @@ def step(self, sim_state: Any, sim_control: Any, dt: float) -> None: # Zero before scatter-add (actuators accumulate into this buffer). self._computed_effort.zero_() for act in self.actuators: - wp.launch( + self._warp_launch.launch( zero_at_indices_kernel, dim=act.indices.shape[0], inputs=[sim_control.joint_f, act.indices], + site=act, ) for act, sa, sb in zip(self.actuators, self._states_a, self._states_b): act.step(sim_state, sim_control, sa, sb, dt=dt) @@ -325,6 +331,15 @@ def from_usd( ) return cls(actuators, num_envs, num_joints, dof_offset=0, device=device) + def _invalidate_launch_cache(self) -> None: + """Release recorded launches after the owning Newton graph is destroyed. + + Callers clear the Newton graph reference before invoking this method. Device + synchronization drains any recorded command work before retained argument + owners are released. Disabled caches return without synchronizing. + """ + self._warp_launch.reset() + # --------------------------------------------------------------------------- # Per-articulation initial-gain snapshot — consumed by diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index e0c26072c1da..9ad1f9e39228 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -36,6 +36,7 @@ from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions from isaaclab.utils.version import get_isaac_sim_version, has_kit +from isaaclab.utils.warp import WarpLaunchCache from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -133,6 +134,11 @@ def __init__(self, cfg: ArticulationCfg): sim_ctx = SimulationContext.instance() self._sim_cfg = sim_ctx.cfg if sim_ctx is not None else None + physics_cfg = self._sim_cfg.physics if self._sim_cfg is not None else None + self._warp_launch = WarpLaunchCache( + device=sim_ctx.device if sim_ctx is not None else None, + enabled=getattr(physics_cfg, "use_warp_launch_cache", False), + ) """ Properties @@ -340,7 +346,7 @@ def write_data_to_sim(self): # joint_effort. Under identity ordering ``_joint_backend_to_user_map`` # returns the identity arange, so this single gather also serves the # no-reorder case without a dedicated per-buffer copy. - wp.launch( + self._warp_launch.launch( ordering_kernels.reorder_joint_targets_user_to_backend, dim=(self.num_instances, self.num_joints), inputs=[ @@ -359,7 +365,7 @@ def write_data_to_sim(self): self.data._sim_bind_joint_velocity_target, self.data._sim_bind_joint_act, ], - device=self.device, + site="newton_joint_targets", ) else: # Standard Lab actuator path @@ -369,7 +375,7 @@ def write_data_to_sim(self): # write_joint_act is off, so it is never indexed. Under identity # ordering ``_joint_backend_to_user_map`` returns the identity arange, # so this single gather also serves the no-reorder case. - wp.launch( + self._warp_launch.launch( ordering_kernels.reorder_joint_targets_user_to_backend, dim=(self.num_instances, self.num_joints), inputs=[ @@ -388,7 +394,7 @@ def write_data_to_sim(self): self.data._sim_bind_joint_velocity_target, self.data._sim_bind_joint_effort, ], - device=self.device, + site="lab_joint_targets", ) def update(self, dt: float): @@ -3392,7 +3398,7 @@ def has_articulation_root_api(prim) -> bool: SimulationManager.get_physics_sim_view().append(self._root_view) # container for data access - self._data = ArticulationData(self.root_view, self.device) + self._data = ArticulationData(self.root_view, self.device, warp_launch=self._warp_launch) # Register callback to rebind simulation data after a full reset (model/state recreation). self._physics_ready_handle = SimulationManager.register_callback( @@ -3538,6 +3544,7 @@ def _process_cfg(self): def _invalidate_initialize_callback(self, event): """Invalidates the scene elements.""" + self._warp_launch.invalidate() # call parent super()._invalidate_initialize_callback(event) self._root_view = None @@ -3892,33 +3899,51 @@ def _apply_actuator_model(self): joint_vel=self._data.joint_vel.torch[:, actuator.joint_indices], ) # update targets (these are set into the simulation) - joint_indices = actuator.joint_indices - if actuator.joint_indices == slice(None) or actuator.joint_indices is None: + actuator_joint_indices = actuator.joint_indices + is_full_joint_group = actuator_joint_indices is None or ( + isinstance(actuator_joint_indices, slice) and actuator_joint_indices == slice(None) + ) + joint_indices = actuator_joint_indices + if is_full_joint_group: joint_indices = self._ALL_JOINT_INDICES if hasattr(actuator, "gear_ratio"): gear_ratio = actuator.gear_ratio else: gear_ratio = None - wp.launch( - articulation_kernels.update_targets, - dim=(self.num_instances, joint_indices.shape[0]), - inputs=[ - control_action.joint_positions, - control_action.joint_velocities, - control_action.joint_efforts, - joint_indices, - ], - outputs=[ - self._joint_pos_target_sim, - self._joint_vel_target_sim, - self._joint_effort_target_sim, - ], - device=self.device, - ) + launch_dim = (self.num_instances, joint_indices.shape[0]) + target_inputs = [ + control_action.joint_positions, + control_action.joint_velocities, + control_action.joint_efforts, + joint_indices, + ] + target_outputs = [ + self._joint_pos_target_sim, + self._joint_vel_target_sim, + self._joint_effort_target_sim, + ] + if getattr(actuator, "is_implicit_model", False) and is_full_joint_group: + self._warp_launch.launch( + articulation_kernels.update_targets, + dim=launch_dim, + inputs=target_inputs, + outputs=target_outputs, + site=actuator, + ) + else: + # Explicit models may replace target tensors, while partial-group + # indexing allocates fresh Torch tensors on every call. + wp.launch( + articulation_kernels.update_targets, + dim=launch_dim, + inputs=target_inputs, + outputs=target_outputs, + device=self.device, + ) # update state of the actuator model wp.launch( articulation_kernels.update_actuator_state_model, - dim=(self.num_instances, joint_indices.shape[0]), + dim=launch_dim, inputs=[ actuator.computed_effort, actuator.applied_effort, diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py index 4c27091dd3cb..cd782c3cfaef 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -16,7 +16,7 @@ from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.buffers import reset_timestamps -from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp import ProxyArray, WarpLaunchCache from isaaclab.utils.warp.utils import capture_unsafe from isaaclab_newton.assets import kernels as shared_kernels @@ -58,18 +58,26 @@ class ArticulationData(BaseArticulationData): __backend_name__: str = "newton" """The name of the backend for the articulation data.""" - def __init__(self, root_view: ArticulationView, device: str): + def __init__( + self, + root_view: ArticulationView, + device: str, + *, + warp_launch: WarpLaunchCache | None = None, + ): """Initializes the articulation data. Args: root_view: The root articulation view. device: The device used for processing. + warp_launch: Optional owner launcher for recorded Warp launches. """ super().__init__(root_view, device) # Set the root articulation view # note: this is stored as a weak reference to avoid circular references between the asset class # and the data container. This is important to avoid memory leaks. self._root_view: ArticulationView = weakref.proxy(root_view) + self._warp_launch = warp_launch if warp_launch is not None else WarpLaunchCache(device=device, enabled=False) # Set initial time stamp self._sim_timestamp = 0.0 @@ -678,18 +686,17 @@ def root_link_vel_w(self) -> ProxyArray: relative to the world. """ if self._root_link_vel_w.timestamp < self._sim_timestamp: - wp.launch( + inputs = [ + self.root_com_vel_w.warp, + self.root_link_pose_w.warp, + self._sim_bind_body_com_pos_b, + ] + outputs = [self._root_link_vel_w.data] + self._warp_launch.launch( shared_kernels.get_root_link_vel_from_root_com_vel, dim=self._num_instances, - inputs=[ - self.root_com_vel_w.warp, - self.root_link_pose_w.warp, - self._sim_bind_body_com_pos_b, - ], - outputs=[ - self._root_link_vel_w.data, - ], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._root_link_vel_w.timestamp = self._sim_timestamp @@ -707,17 +714,13 @@ def root_com_pose_w(self) -> ProxyArray: """ if self._root_com_pose_w.timestamp < self._sim_timestamp: # apply local transform to center of mass frame - wp.launch( + inputs = [self.root_link_pose_w.warp, self._sim_bind_body_com_pos_b] + outputs = [self._root_com_pose_w.data] + self._warp_launch.launch( shared_kernels.get_root_com_pose_from_root_link_pose, dim=self._num_instances, - inputs=[ - self.root_link_pose_w.warp, - self._sim_bind_body_com_pos_b, - ], - outputs=[ - self._root_com_pose_w.data, - ], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._root_com_pose_w.timestamp = self._sim_timestamp @@ -783,18 +786,17 @@ def body_link_vel_w(self) -> ProxyArray: relative to the world. """ if self._body_link_vel_w.timestamp < self._sim_timestamp: - wp.launch( + inputs = [ + self.body_com_vel_w.warp, + self.body_link_pose_w.warp, + self.body_com_pos_b.warp, + ] + outputs = [self._body_link_vel_w.data] + self._warp_launch.launch( shared_kernels.get_body_link_vel_from_body_com_vel, dim=(self._num_instances, self._num_bodies), - inputs=[ - self.body_com_vel_w.warp, - self.body_link_pose_w.warp, - self.body_com_pos_b.warp, - ], - outputs=[ - self._body_link_vel_w.data, - ], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._body_link_vel_w.timestamp = self._sim_timestamp @@ -813,17 +815,13 @@ def body_com_pose_w(self) -> ProxyArray: """ self._ensure_fk_fresh() if self._body_com_pose_w.timestamp < self._sim_timestamp: - wp.launch( + inputs = [self.body_link_pose_w.warp, self.body_com_pos_b.warp] + outputs = [self._body_com_pose_w.data] + self._warp_launch.launch( shared_kernels.get_body_com_pose_from_body_link_pose, dim=(self._num_instances, self._num_bodies), - inputs=[ - self.body_link_pose_w.warp, - self.body_com_pos_b.warp, - ], - outputs=[ - self._body_com_pose_w.data, - ], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._body_com_pose_w.timestamp = self._sim_timestamp @@ -1104,12 +1102,13 @@ def projected_gravity_b(self) -> ProxyArray: Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ if self._projected_gravity_b.timestamp < self._sim_timestamp: - wp.launch( + inputs = [self.GRAVITY_VEC_W.warp, self.root_link_quat_w.warp] + outputs = [self._projected_gravity_b.data] + self._warp_launch.launch( shared_kernels.projected_gravity_b_kernel, dim=self._num_instances, - inputs=[self.GRAVITY_VEC_W.warp, self.root_link_quat_w.warp], - outputs=[self._projected_gravity_b.data], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._projected_gravity_b.timestamp = self._sim_timestamp return self._projected_gravity_b_ta @@ -1126,12 +1125,13 @@ def heading_w(self) -> ProxyArray: frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ if self._heading_w.timestamp < self._sim_timestamp: - wp.launch( + inputs = [self.FORWARD_VEC_B.warp, self.root_link_quat_w.warp] + outputs = [self._heading_w.data] + self._warp_launch.launch( shared_kernels.root_heading_w, dim=self._num_instances, - inputs=[self.FORWARD_VEC_B.warp, self.root_link_quat_w.warp], - outputs=[self._heading_w.data], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._heading_w.timestamp = self._sim_timestamp return self._heading_w_ta @@ -1152,12 +1152,13 @@ def root_link_lin_vel_b(self) -> ProxyArray: ) self._root_link_lin_vel_b_ta = ProxyArray(self._root_link_lin_vel_b.data) if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: - wp.launch( + inputs = [self.root_link_lin_vel_w.warp, self.root_link_quat_w.warp] + outputs = [self._root_link_lin_vel_b.data] + self._warp_launch.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_link_lin_vel_w.warp, self.root_link_quat_w.warp], - outputs=[self._root_link_lin_vel_b.data], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._root_link_lin_vel_b.timestamp = self._sim_timestamp return self._root_link_lin_vel_b_ta @@ -1524,6 +1525,11 @@ def _create_simulation_bindings(self) -> None: .. caution:: This is possible if and only if the properties that we access are strided from newton and not indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. """ + # A full simulation reset replaces state/model pointers after the Newton + # graph has been destroyed. Drain recorded work, then release commands + # and owners recorded against the previous storage. + self._warp_launch.reset() + # Short-hand for the number of instances, number of links, and number of joints. self._num_instances = self._root_view.count self._num_joints = self._root_view.joint_dof_count diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py index 3265ad2896dd..13156e6d6d76 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 @@ -23,6 +23,7 @@ from isaaclab.physics import PhysicsEvent from isaaclab.sim.utils.queries import resolve_matching_prims_from_source from isaaclab.utils.version import has_kit +from isaaclab.utils.warp import WarpLaunchCache from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -61,11 +62,20 @@ def __init__(self, cfg: RigidObjectCfg): Args: cfg: A configuration instance. """ + from isaaclab.sim import SimulationContext # noqa: PLC0415 + super().__init__(cfg) if has_kit(): queue_usd_replication(cfg) queue_newton_physics_replication(cfg) + sim_ctx = SimulationContext.instance() + physics_cfg = sim_ctx.cfg.physics if sim_ctx is not None else None + self._warp_launch = WarpLaunchCache( + device=sim_ctx.device if sim_ctx is not None else None, + enabled=getattr(physics_cfg, "use_warp_launch_cache", False), + ) + """ Properties """ @@ -1031,7 +1041,7 @@ def has_rigid_body_api(prim) -> bool: ) # container for data access - self._data = RigidObjectData(self.root_view, self.device) + self._data = RigidObjectData(self.root_view, self.device, warp_launch=self._warp_launch) # Register callback to rebind simulation data after a full reset (model/state recreation). self._physics_ready_handle = SimulationManager.register_callback( @@ -1127,6 +1137,7 @@ def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | def _invalidate_initialize_callback(self, event): """Invalidates the scene elements.""" + self._warp_launch.invalidate() # call parent super()._invalidate_initialize_callback(event) # set all existing views to None to invalidate them diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py index c568427bac4c..d1f22b536d8f 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py @@ -15,7 +15,7 @@ from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer from isaaclab.utils.buffers import reset_timestamps -from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp import ProxyArray, WarpLaunchCache from isaaclab.utils.warp.utils import capture_unsafe from isaaclab_newton.assets import kernels as shared_kernels @@ -61,18 +61,26 @@ class RigidObjectData(BaseRigidObjectData): __backend_name__: str = "newton" """The name of the backend for the rigid object data.""" - def __init__(self, root_view: ArticulationView, device: str): + def __init__( + self, + root_view: ArticulationView, + device: str, + *, + warp_launch: WarpLaunchCache | None = None, + ): """Initializes the rigid object data. Args: root_view: The root rigid body view. device: The device used for processing. + warp_launch: Optional owner launcher for recorded Warp launches. """ super().__init__(root_view, device) # Set the root rigid body view # note: this is stored as a weak reference to avoid circular references between the asset class # and the data container. This is important to avoid memory leaks. self._root_view: ArticulationView = weakref.proxy(root_view) + self._warp_launch = warp_launch if warp_launch is not None else WarpLaunchCache(device=device, enabled=False) # Set initial time stamp self._sim_timestamp = 0.0 @@ -280,18 +288,17 @@ def root_link_vel_w(self) -> ProxyArray: """ if self._root_link_vel_w.timestamp < self._sim_timestamp: # read the CoM velocity and compute link velocity - wp.launch( + inputs = [ + self.root_com_vel_w.warp, + self.root_link_pose_w.warp, + self.body_com_pos_b.warp, + ] + outputs = [self._root_link_vel_w.data] + self._warp_launch.launch( shared_kernels.get_root_link_vel_from_root_com_vel, dim=self._num_instances, - inputs=[ - self.root_com_vel_w.warp, - self.root_link_pose_w.warp, - self.body_com_pos_b.warp, - ], - outputs=[ - self._root_link_vel_w.data, - ], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._root_link_vel_w.timestamp = self._sim_timestamp @@ -308,17 +315,13 @@ def root_com_pose_w(self) -> ProxyArray: """ if self._root_com_pose_w.timestamp < self._sim_timestamp: # apply local transform to center of mass frame - wp.launch( + inputs = [self.root_link_pose_w.warp, self.body_com_pos_b.warp] + outputs = [self._root_com_pose_w.data] + self._warp_launch.launch( shared_kernels.get_root_com_pose_from_root_link_pose, dim=self._num_instances, - inputs=[ - self.root_link_pose_w.warp, - self.body_com_pos_b.warp, - ], - outputs=[ - self._root_com_pose_w.data, - ], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._root_com_pose_w.timestamp = self._sim_timestamp @@ -486,12 +489,13 @@ def projected_gravity_b(self) -> ProxyArray: Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). """ if self._projected_gravity_b.timestamp < self._sim_timestamp: - wp.launch( + inputs = [self.GRAVITY_VEC_W.warp, self.root_link_quat_w.warp] + outputs = [self._projected_gravity_b.data] + self._warp_launch.launch( shared_kernels.projected_gravity_b_kernel, dim=self._num_instances, - inputs=[self.GRAVITY_VEC_W.warp, self.root_link_quat_w.warp], - outputs=[self._projected_gravity_b.data], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._projected_gravity_b.timestamp = self._sim_timestamp return self._projected_gravity_b_ta @@ -508,12 +512,13 @@ def heading_w(self) -> ProxyArray: frame is along x-direction, i.e. :math:`(1, 0, 0)`. """ if self._heading_w.timestamp < self._sim_timestamp: - wp.launch( + inputs = [self.FORWARD_VEC_B.warp, self.root_link_quat_w.warp] + outputs = [self._heading_w.data] + self._warp_launch.launch( shared_kernels.root_heading_w, dim=self._num_instances, - inputs=[self.FORWARD_VEC_B.warp, self.root_link_quat_w.warp], - outputs=[self._heading_w.data], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._heading_w.timestamp = self._sim_timestamp return self._heading_w_ta @@ -533,12 +538,13 @@ def root_link_lin_vel_b(self) -> ProxyArray: ) self._root_link_lin_vel_b_ta = ProxyArray(self._root_link_lin_vel_b.data) if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: - wp.launch( + inputs = [self.root_link_lin_vel_w.warp, self.root_link_quat_w.warp] + outputs = [self._root_link_lin_vel_b.data] + self._warp_launch.launch( shared_kernels.quat_apply_inverse_1D_kernel, dim=self._num_instances, - inputs=[self.root_link_lin_vel_w.warp, self.root_link_quat_w.warp], - outputs=[self._root_link_lin_vel_b.data], - device=self.device, + inputs=inputs, + outputs=outputs, ) self._root_link_lin_vel_b.timestamp = self._sim_timestamp return self._root_link_lin_vel_b_ta @@ -872,6 +878,11 @@ def _create_simulation_bindings(self) -> None: .. caution:: This is possible if and only if the properties that we access are strided from newton and not indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. """ + # A full simulation reset replaces state/model pointers after the Newton + # graph has been destroyed. Drain recorded work, then release commands + # and owners recorded against the previous storage. + self._warp_launch.reset() + # Short-hand for the number of instances, number of links, and number of joints. self._num_instances = self._root_view.count self._num_bodies = self._root_view.link_count diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index f3ec45149c9c..090d1b9ced6a 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -73,6 +73,7 @@ def _paused_gc(): from isaaclab.utils import checked_apply from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.timer import Timer +from isaaclab.utils.warp import WarpLaunchCache from isaaclab_newton.cloner.newton_clone_utils import replicate_builder_mapping from isaaclab_newton.physics.visualization_builder import build_visualization_builder_from_stage_envs @@ -303,6 +304,7 @@ def provides_implicit_damping(cls) -> bool: # Newton actuator adapter (owns actuators and double-buffered states) _adapter: NewtonActuatorAdapter | None = None + _warp_launch: WarpLaunchCache | None = None # In-graph hooks invoked after the actuator step and before the solver # substeps, in registration order. Multiple articulations register their # implicit-DOF telemetry / FF-routing kernels here. @@ -787,6 +789,13 @@ def step(cls) -> None: @classmethod def close(cls) -> None: """Clean up Newton physics resources.""" + # Asset STOP callbacks invalidate recorded commands. A solver graph may + # contain those commands, so release the graph and drain device work + # before PhysicsManager dispatches STOP. + NewtonManager._graph = None + NewtonManager._graph_capture_pending = False + if PhysicsManager._device is not None: + wp.synchronize_device(PhysicsManager._device) super().close() cls.clear() @@ -829,6 +838,10 @@ def clear(cls): """Clear all Newton-specific state (callbacks cleared by super().close()).""" if cls._cubric is not None and cls._cubric_adapter is not None: cls._cubric.release_adapter(cls._cubric_adapter) + # Recorded commands may be embedded in the graph and retain pointers + # owned by the adapter cache. Destroy the graph before that cache. + NewtonManager._graph = None + NewtonManager._graph_capture_pending = False NewtonManager._cubric = None NewtonManager._cubric_adapter = None NewtonManager._cubric_bound_fabric_id = None @@ -848,7 +861,10 @@ def clear(cls): NewtonManager._newton_frame_transform_sensors = [] NewtonManager._newton_imu_sensors = [] NewtonManager._report_contacts = False + if NewtonManager._adapter is not None: + NewtonManager._adapter._invalidate_launch_cache() NewtonManager._adapter = None + NewtonManager._warp_launch = None NewtonManager._post_actuator_callbacks = [] NewtonManager._post_step_callbacks = [] # Set by an articulation that took the ``use_newton_actuators=True`` @@ -860,8 +876,6 @@ def clear(cls): # Per-world reset masks NewtonManager._world_reset_mask = None NewtonManager._fk_reset_mask = None - NewtonManager._graph = None - NewtonManager._graph_capture_pending = False NewtonManager._newton_stage_path = None NewtonManager._usdrt_stage = None NewtonManager._transforms_dirty = False @@ -1237,6 +1251,11 @@ def start_simulation(cls) -> None: Note: Collision pipeline is initialized later in initialize_solver() after we determine whether the solver needs external collision detection. """ + # A previous graph may embed recorded adapter launches. Destroy it + # before replacing model/state storage or the owning adapter cache. + NewtonManager._graph = None + NewtonManager._graph_capture_pending = False + logger.debug(f"Builder: {cls._builder}") cls._drain_stale_cuda_error() @@ -1281,7 +1300,10 @@ def start_simulation(cls) -> None: # articulation after this point. Assign through the explicit base # class so external readers (which import ``NewtonManager`` directly) # observe the canonical state regardless of which subclass is active. + if NewtonManager._adapter is not None: + NewtonManager._adapter._invalidate_launch_cache() NewtonManager._adapter = None + NewtonManager._warp_launch = None NewtonManager._use_newton_actuators_active = False # Allocate per-world reset masks (used by all solvers for masked FK, and by Kamino for masked reset). @@ -2162,12 +2184,18 @@ def activate_newton_actuator_path(cls) -> None: from isaaclab_newton.actuators import NewtonActuatorAdapter # noqa: PLC0415 dofs_per_env = cls._model.joint_dof_count // cls._num_envs + cfg = PhysicsManager._cfg + NewtonManager._warp_launch = WarpLaunchCache( + device=PhysicsManager._device, + enabled=getattr(cfg, "use_warp_launch_cache", False), + ) NewtonManager._adapter = NewtonActuatorAdapter( actuators=list(cls._model.actuators), num_envs=cls._num_envs, num_joints=dofs_per_env, dof_offset=0, device=PhysicsManager._device, + warp_launch=NewtonManager._warp_launch, ) cls._adapter.finalize(cls._control) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py index 019b1b0bdced..dd6632a00450 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py @@ -112,6 +112,16 @@ class NewtonCfg(PhysicsCfg): If set to False, the simulation performance will be severely degraded. """ + use_warp_launch_cache: bool = False + """Whether to replay recorded Warp launches outside the Newton solver graph. + + When enabled, cached :class:`warp.Launch` commands reduce Python launch + preparation for pointer-stable articulation target publication, eager + Newton-actuator stepping, and capture-unsafe lazy derived-state properties. + It does not make Python timestamp branches, Torch actuator computation, or + the complete Isaac Lab actuator publication path CUDA-graph-safe. + """ + solver_cfg: NewtonSolverCfg | None = None """Solver configuration. If None (default), MJWarpSolverCfg is used by default.""" diff --git a/source/isaaclab_newton/test/assets/test_articulation_warp_launch_cache.py b/source/isaaclab_newton/test/assets/test_articulation_warp_launch_cache.py new file mode 100644 index 000000000000..8cca8760df8c --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_articulation_warp_launch_cache.py @@ -0,0 +1,281 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for Warp launch handling at the Newton Isaac Lab actuator boundary.""" + +from types import SimpleNamespace + +import isaaclab_newton.assets.articulation.articulation as articulation_module +import isaaclab_newton.assets.rigid_object.rigid_object as rigid_object_module +import pytest +import torch +import warp as wp +from isaaclab_newton.assets.articulation.articulation import Articulation +from isaaclab_newton.assets.rigid_object.rigid_object import RigidObject +from isaaclab_newton.physics import NewtonCfg + +from isaaclab.sim import SimulationContext +from isaaclab.utils.types import ArticulationActions +from isaaclab.utils.warp import WarpLaunchCache + + +class _AllocatingActuator: + """Minimal explicit actuator that returns new Torch allocations per call.""" + + def __init__(self, num_envs: int, num_joints: int, device: str): + self.is_implicit_model = False + self.joint_indices = slice(None) + self.gear_ratio = None + self.velocity_limit = torch.full((num_envs, num_joints), 10.0, device=device) + self.computed_effort = torch.zeros((num_envs, num_joints), device=device) + self.applied_effort = torch.zeros_like(self.computed_effort) + self.output_pointers: list[tuple[int, int, int, int]] = [] + + def compute( + self, + control_action: ArticulationActions, + joint_pos: torch.Tensor, + joint_vel: torch.Tensor, + ) -> ArticulationActions: + del joint_pos, joint_vel + self.computed_effort = control_action.joint_efforts + 1.0 + self.applied_effort = self.computed_effort * 2.0 + if self.output_pointers: + self.gear_ratio = self.computed_effort + 10.0 + control_action.joint_positions = self.applied_effort + 100.0 + else: + self.gear_ratio = None + control_action.joint_positions = None + self.velocity_limit = self.computed_effort + 20.0 + self.output_pointers.append( + ( + self.computed_effort.data_ptr(), + self.applied_effort.data_ptr(), + 0 if self.gear_ratio is None else self.gear_ratio.data_ptr(), + self.velocity_limit.data_ptr(), + ) + ) + control_action.joint_velocities = None + control_action.joint_efforts = self.applied_effort + return control_action + + +class _AllocatingImplicitActuator: + """Minimal implicit actuator with persistent targets and allocating telemetry.""" + + is_implicit_model = True + + def __init__(self, num_envs: int, num_joints: int, device: str): + self.joint_indices = slice(None) + self.gear_ratio = None + self.velocity_limit = torch.full((num_envs, num_joints), 10.0, device=device) + self.computed_effort = torch.zeros((num_envs, num_joints), device=device) + self.applied_effort = torch.zeros_like(self.computed_effort) + self.output_pointers: list[tuple[int, int]] = [] + + def compute( + self, + control_action: ArticulationActions, + joint_pos: torch.Tensor, + joint_vel: torch.Tensor, + ) -> ArticulationActions: + del joint_pos, joint_vel + self.computed_effort = control_action.joint_efforts + 1.0 + self.applied_effort = self.computed_effort * 2.0 + self.output_pointers.append((self.computed_effort.data_ptr(), self.applied_effort.data_ptr())) + return control_action + + +def _make_articulation_stub( + actuator: _AllocatingActuator | _AllocatingImplicitActuator, + num_envs: int, + num_joints: int, + device: str, +) -> tuple[SimpleNamespace, torch.Tensor]: + """Create the minimum articulation state needed by the actuator publication path.""" + shape = (num_envs, num_joints) + articulation = SimpleNamespace(device=device, num_instances=num_envs) + articulation._ALL_JOINT_INDICES = wp.array(range(num_joints), dtype=wp.int32, device=device) + articulation._joint_pos_target_sim = wp.zeros(shape, dtype=wp.float32, device=device) + articulation._joint_vel_target_sim = wp.zeros(shape, dtype=wp.float32, device=device) + articulation._joint_effort_target_sim = wp.zeros(shape, dtype=wp.float32, device=device) + articulation._warp_launch = WarpLaunchCache(mode="replay", debug=True, device=device) + articulation.actuators = {"test": actuator} + + joint_effort_target = torch.ones(shape, device=device) + articulation._data = SimpleNamespace( + joint_pos_target=SimpleNamespace(torch=torch.zeros(shape, device=device)), + joint_vel_target=SimpleNamespace(torch=torch.zeros(shape, device=device)), + joint_effort_target=SimpleNamespace(torch=joint_effort_target), + joint_pos=SimpleNamespace(torch=torch.zeros(shape, device=device)), + joint_vel=SimpleNamespace(torch=torch.zeros(shape, device=device)), + computed_torque=wp.zeros(shape, dtype=wp.float32, device=device), + applied_torque=wp.zeros(shape, dtype=wp.float32, device=device), + gear_ratio=wp.zeros(shape, dtype=wp.float32, device=device), + soft_joint_vel_limits=wp.zeros(shape, dtype=wp.float32, device=device), + ) + return articulation, joint_effort_target + + +@pytest.fixture(scope="module") +def cuda_device() -> str: + """Return a CUDA device or skip the module when CUDA is unavailable.""" + wp.init() + if not wp.is_cuda_available(): + pytest.skip("CUDA is required for Warp launch replay tests.") + return "cuda:0" + + +def test_newton_launch_cache_is_disabled_by_default(): + """Newton should preserve eager actuator-boundary launches unless opted in.""" + assert NewtonCfg().use_warp_launch_cache is False + + +@pytest.mark.parametrize( + ("asset_cls", "asset_module"), + [(Articulation, articulation_module), (RigidObject, rigid_object_module)], +) +def test_asset_launch_cache_uses_sim_device_before_asset_initialization( + asset_cls: type, + asset_module, + monkeypatch: pytest.MonkeyPatch, +): + """Asset construction should not read ``asset.device`` before the ready callback sets it.""" + sim_context = SimpleNamespace( + cfg=SimpleNamespace(physics=SimpleNamespace(use_warp_launch_cache=True)), + device="cpu", + ) + monkeypatch.setattr(SimulationContext, "instance", classmethod(lambda cls: sim_context)) + monkeypatch.setattr(asset_cls.__mro__[1], "__init__", lambda self, cfg: setattr(self, "cfg", cfg)) + monkeypatch.setattr(asset_cls, "_clear_callbacks", lambda self: None) + monkeypatch.setattr(asset_module, "has_kit", lambda: False) + monkeypatch.setattr(asset_module, "queue_newton_physics_replication", lambda cfg: None) + + asset = object.__new__(asset_cls) + asset_cls.__init__(asset, SimpleNamespace()) + + assert not hasattr(asset, "_device") + assert asset._warp_launch._device == wp.get_device("cpu") + assert asset._warp_launch._mode == "replay" + del asset + + +def test_actuator_publication_accepts_reallocated_torch_outputs(cuda_device: str): + """Eager publication should use current actuator tensors after their pointers change.""" + num_envs = 8 + num_joints = 3 + actuator = _AllocatingActuator(num_envs, num_joints, cuda_device) + articulation, joint_effort_target = _make_articulation_stub(actuator, num_envs, num_joints, cuda_device) + shape = (num_envs, num_joints) + + Articulation._apply_actuator_model(articulation) + joint_effort_target.fill_(3.0) + Articulation._apply_actuator_model(articulation) + wp.synchronize_device(cuda_device) + + assert actuator.output_pointers[0] != actuator.output_pointers[1] + assert len(articulation._warp_launch._entries) == 0 + torch.testing.assert_close( + wp.to_torch(articulation._joint_pos_target_sim), torch.full(shape, 108.0, device=cuda_device) + ) + torch.testing.assert_close( + wp.to_torch(articulation._joint_effort_target_sim), torch.full(shape, 8.0, device=cuda_device) + ) + torch.testing.assert_close( + wp.to_torch(articulation._data.computed_torque), torch.full(shape, 4.0, device=cuda_device) + ) + torch.testing.assert_close( + wp.to_torch(articulation._data.applied_torque), torch.full(shape, 8.0, device=cuda_device) + ) + torch.testing.assert_close(wp.to_torch(articulation._data.gear_ratio), torch.full(shape, 14.0, device=cuda_device)) + torch.testing.assert_close( + wp.to_torch(articulation._data.soft_joint_vel_limits), torch.full(shape, 24.0, device=cuda_device) + ) + + +def test_full_slice_implicit_targets_use_recorded_launch(cuda_device: str): + """Persistent full-group implicit targets should replay while allocating telemetry stays eager.""" + num_envs = 8 + num_joints = 3 + shape = (num_envs, num_joints) + actuator = _AllocatingImplicitActuator(num_envs, num_joints, cuda_device) + articulation, joint_effort_target = _make_articulation_stub(actuator, num_envs, num_joints, cuda_device) + + Articulation._apply_actuator_model(articulation) + joint_effort_target.fill_(3.0) + Articulation._apply_actuator_model(articulation) + wp.synchronize_device(cuda_device) + + assert actuator.output_pointers[0] != actuator.output_pointers[1] + assert len(articulation._warp_launch._entries) == 1 + torch.testing.assert_close( + wp.to_torch(articulation._joint_effort_target_sim), torch.full(shape, 3.0, device=cuda_device) + ) + torch.testing.assert_close( + wp.to_torch(articulation._data.computed_torque), torch.full(shape, 4.0, device=cuda_device) + ) + torch.testing.assert_close( + wp.to_torch(articulation._data.applied_torque), torch.full(shape, 8.0, device=cuda_device) + ) + + +@pytest.mark.parametrize("newton_actuator_path", [False, True]) +def test_joint_target_reorder_uses_recorded_launch(cuda_device: str, newton_actuator_path: bool): + """Persistent Lab and Newton joint-target reorders should replay current in-place values.""" + num_envs = 4 + num_joints = 3 + shape = (num_envs, num_joints) + source_effort = wp.ones(shape, dtype=wp.float32, device=cuda_device) + source_pos = wp.full(shape, 2.0, dtype=wp.float32, device=cuda_device) + source_vel = wp.full(shape, 3.0, dtype=wp.float32, device=cuda_device) + backend_effort = wp.zeros(shape, dtype=wp.float32, device=cuda_device) + backend_pos = wp.zeros(shape, dtype=wp.float32, device=cuda_device) + backend_vel = wp.zeros(shape, dtype=wp.float32, device=cuda_device) + backend_act = wp.zeros(shape, dtype=wp.float32, device=cuda_device) + backend_to_user = wp.array([2, 0, 1], dtype=wp.int32, device=cuda_device) + + articulation = SimpleNamespace( + _has_newton_actuators=newton_actuator_path, + _has_implicit_actuators=True, + _instantaneous_wrench_composer=SimpleNamespace(active=False), + _permanent_wrench_composer=SimpleNamespace(active=False), + _joint_effort_target_sim=source_effort, + _joint_pos_target_sim=source_pos, + _joint_vel_target_sim=source_vel, + _joint_backend_to_user_map=lambda: backend_to_user, + _apply_actuator_model=lambda: None, + _warp_launch=WarpLaunchCache(mode="replay", debug=True, device=cuda_device), + _data=SimpleNamespace( + _joint_effort_target=source_effort, + _joint_pos_target=source_pos, + _joint_vel_target=source_vel, + ), + data=SimpleNamespace( + _sim_bind_joint_effort=backend_effort, + _sim_bind_joint_position_target=backend_pos, + _sim_bind_joint_velocity_target=backend_vel, + _sim_bind_joint_act=backend_act, + ), + device=cuda_device, + num_instances=num_envs, + num_joints=num_joints, + ) + + Articulation.write_data_to_sim(articulation) + source_effort.fill_(4.0) + source_pos.fill_(5.0) + source_vel.fill_(6.0) + Articulation.write_data_to_sim(articulation) + wp.synchronize_device(cuda_device) + + assert len(articulation._warp_launch._entries) == 1 + entry = next(iter(articulation._warp_launch._entries.values())) + expected_site = "newton_joint_targets" if newton_actuator_path else "lab_joint_targets" + assert entry.site == expected_site + torch.testing.assert_close(wp.to_torch(backend_effort), torch.full(shape, 4.0, device=cuda_device)) + torch.testing.assert_close(wp.to_torch(backend_pos), torch.full(shape, 5.0, device=cuda_device)) + torch.testing.assert_close(wp.to_torch(backend_vel), torch.full(shape, 6.0, device=cuda_device)) + if newton_actuator_path: + torch.testing.assert_close(wp.to_torch(backend_act), torch.full(shape, 4.0, device=cuda_device)) diff --git a/source/isaaclab_newton/test/assets/test_warp_launch_cache_hot_paths.py b/source/isaaclab_newton/test/assets/test_warp_launch_cache_hot_paths.py new file mode 100644 index 000000000000..a6e452bc6a67 --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_warp_launch_cache_hot_paths.py @@ -0,0 +1,258 @@ +# 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 recorded replay on capture-unsafe Newton hot paths.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import warp as wp +from isaaclab_newton.actuators import NewtonActuatorAdapter +from isaaclab_newton.assets import ArticulationData, RigidObjectData +from isaaclab_newton.physics import NewtonManager + +from isaaclab.physics import PhysicsManager +from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.warp import ProxyArray, WarpLaunchCache + + +class _NoOpActuator: + """Minimal Newton actuator for adapter launch tests.""" + + def __init__(self, indices: wp.array): + self.indices = indices + self.control_computed_output_attr = None + + def state(self): + return None + + def step(self, sim_state, sim_control, state_in, state_out, dt: float) -> None: + del sim_state, sim_control, state_in, state_out, dt + + def is_graphable(self) -> bool: + return False + + +class _RebindReset(RuntimeError): + """Sentinel raised when the rebind hook resets its launch cache.""" + + +class _ResetProbe: + """Cache probe that stops a simulation rebind after observing reset.""" + + def __init__(self, events: list[tuple[str, str | None]]): + self.events = events + + def reset(self) -> None: + self.events.append(("reset", None)) + raise _RebindReset + + +@pytest.fixture(scope="module") +def cuda_device() -> str: + """Return a CUDA device or skip the module when CUDA is unavailable.""" + wp.init() + if not wp.is_cuda_available(): + pytest.skip("CUDA is required for Warp launch replay tests.") + return "cuda:0" + + +def _spatial_vectors(num_envs: int, value: float, device: str) -> wp.array: + """Create constant spatial vectors for a derived-state test.""" + values = np.full((num_envs, 6), value, dtype=np.float32) + return wp.array(values, dtype=wp.spatial_vectorf, device=device) + + +def _make_root_velocity_data(data_cls: type, source: wp.array, warp_launch: WarpLaunchCache, device: str): + """Build the fields used by a data class' lazy root-link velocity property.""" + num_envs = source.shape[0] + root_pose = np.zeros((num_envs, 7), dtype=np.float32) + root_pose[:, 6] = 1.0 + + data = object.__new__(data_cls) + data.device = device + data._warp_launch = warp_launch + data._num_instances = num_envs + data._sim_timestamp = 1.0 + data._root_com_vel_w_ta = ProxyArray(source) + data._root_link_pose_w_ta = ProxyArray(wp.array(root_pose, dtype=wp.transformf, device=device)) + data._body_com_pos_b_ta = ProxyArray(wp.zeros((num_envs, 1), dtype=wp.vec3f, device=device)) + data._sim_bind_body_com_pos_b = data._body_com_pos_b_ta.warp + data._root_link_vel_w = TimestampedBuffer( + shape=(num_envs,), + dtype=wp.spatial_vectorf, + device=device, + ) + data._root_link_vel_w_ta = ProxyArray(data._root_link_vel_w.data) + return data + + +@pytest.mark.parametrize("data_cls", [ArticulationData, RigidObjectData]) +def test_lazy_derived_replay_refreshes_stale_timestamp_and_rebound_source(data_cls: type, cuda_device: str): + """Replay should refresh stale data, while invalidation should admit a rebound source pointer.""" + cache = WarpLaunchCache(mode="replay", debug=True, device=cuda_device) + source = _spatial_vectors(8, 1.0, cuda_device) + data = _make_root_velocity_data(data_cls, source, cache, cuda_device) + + first = data.root_link_vel_w + wp.synchronize_device(cuda_device) + np.testing.assert_allclose(first.warp.numpy(), source.numpy()) + + source.fill_(2.0) + data._sim_timestamp = 2.0 + second = data.root_link_vel_w + wp.synchronize_device(cuda_device) + np.testing.assert_allclose(second.warp.numpy(), source.numpy()) + + rebound_source = _spatial_vectors(8, 3.0, cuda_device) + data._root_com_vel_w_ta = ProxyArray(rebound_source) + cache.invalidate() + data._sim_timestamp = 3.0 + third = data.root_link_vel_w + wp.synchronize_device(cuda_device) + np.testing.assert_allclose(third.warp.numpy(), rebound_source.numpy()) + + +@pytest.mark.parametrize("data_cls", [ArticulationData, RigidObjectData]) +def test_simulation_rebind_delegates_cache_drain(data_cls: type, monkeypatch: pytest.MonkeyPatch): + """Simulation rebinding should delegate conditional synchronization to the cache.""" + events: list[tuple[str, str | None]] = [] + data = object.__new__(data_cls) + data.device = "cuda:0" + data._warp_launch = _ResetProbe(events) + monkeypatch.setattr(wp, "synchronize_device", lambda device: pytest.fail(f"unexpected direct sync on {device}")) + + with pytest.raises(_RebindReset): + data._create_simulation_bindings() + + assert events == [("reset", None)] + + +def test_adapter_teardown_delegates_cache_drain(monkeypatch: pytest.MonkeyPatch): + """Adapter teardown should delegate conditional synchronization to the cache.""" + events: list[tuple[str, str | None]] = [] + adapter = object.__new__(NewtonActuatorAdapter) + adapter._device = "cuda:0" + adapter._warp_launch = _ResetProbe(events) + monkeypatch.setattr(wp, "synchronize_device", lambda device: pytest.fail(f"unexpected direct sync on {device}")) + + with pytest.raises(_RebindReset): + adapter._invalidate_launch_cache() + + assert events == [("reset", None)] + + +def test_newton_close_releases_graph_before_asset_stop_callbacks(monkeypatch: pytest.MonkeyPatch): + """Newton shutdown should release graphs and drain work before asset caches are invalidated.""" + events: list[tuple[str, object]] = [] + graph = object() + saved_graph = NewtonManager._graph + saved_capture_pending = NewtonManager._graph_capture_pending + saved_device = PhysicsManager._device + try: + NewtonManager._graph = graph + NewtonManager._graph_capture_pending = True + PhysicsManager._device = "cuda:0" + monkeypatch.setattr( + wp, + "synchronize_device", + lambda device: events.append(("synchronize", device)), + ) + monkeypatch.setattr( + PhysicsManager, + "close", + classmethod(lambda cls: events.append(("stop", NewtonManager._graph))), + ) + monkeypatch.setattr( + NewtonManager, + "clear", + classmethod(lambda cls: events.append(("clear", NewtonManager._graph))), + ) + + NewtonManager.close() + + assert NewtonManager._graph is None + assert NewtonManager._graph_capture_pending is False + assert events == [("synchronize", "cuda:0"), ("stop", None), ("clear", None)] + finally: + NewtonManager._graph = saved_graph + NewtonManager._graph_capture_pending = saved_capture_pending + PhysicsManager._device = saved_device + + +@pytest.mark.parametrize("use_cache", [False, True]) +def test_adapter_step_zeros_current_effort_with_eager_or_replay(use_cache: bool, cuda_device: str): + """Adapter zeroing should preserve eager behavior and replay correctly when enabled.""" + indices = wp.array([1, 3], dtype=wp.uint32, device=cuda_device) + actuator = _NoOpActuator(indices) + warp_launch = WarpLaunchCache(device=cuda_device, enabled=use_cache, debug=True) + adapter = NewtonActuatorAdapter( + [actuator], + num_envs=1, + num_joints=4, + dof_offset=0, + device=cuda_device, + warp_launch=warp_launch, + ) + control = SimpleNamespace(joint_f=wp.full(4, 7.0, dtype=wp.float32, device=cuda_device)) + + adapter.step(None, control, 0.01) + wp.synchronize_device(cuda_device) + np.testing.assert_allclose(control.joint_f.numpy(), [7.0, 0.0, 7.0, 0.0]) + + control.joint_f.fill_(9.0) + adapter.step(None, control, 0.01) + wp.synchronize_device(cuda_device) + np.testing.assert_allclose(control.joint_f.numpy(), [9.0, 0.0, 9.0, 0.0]) + + assert bool(warp_launch._entries) is use_cache + adapter._invalidate_launch_cache() + assert not warp_launch._entries + + +@pytest.mark.parametrize("use_cache", [False, True]) +def test_newton_manager_flag_selects_adapter_launch_mode(use_cache: bool, cuda_device: str): + """The Newton manager should share an eager or replay launcher with its adapter.""" + had_active_flag = hasattr(NewtonManager, "_use_newton_actuators_active") + saved_state = { + "adapter": NewtonManager._adapter, + "warp_launch": NewtonManager._warp_launch, + "control": NewtonManager._control, + "device": PhysicsManager._device, + "cfg": PhysicsManager._cfg, + "model": NewtonManager._model, + "num_envs": NewtonManager._num_envs, + "use_newton_actuators_active": getattr(NewtonManager, "_use_newton_actuators_active", False), + } + try: + actuator = _NoOpActuator(wp.array([1, 3], dtype=wp.uint32, device=cuda_device)) + PhysicsManager._cfg = SimpleNamespace(use_warp_launch_cache=use_cache) + PhysicsManager._device = cuda_device + NewtonManager._model = SimpleNamespace(actuators=[actuator], joint_dof_count=4) + NewtonManager._num_envs = 1 + NewtonManager._control = SimpleNamespace() + NewtonManager._adapter = None + NewtonManager._use_newton_actuators_active = False + + NewtonManager.activate_newton_actuator_path() + + assert NewtonManager._adapter is not None + assert NewtonManager._adapter._warp_launch is NewtonManager._warp_launch + assert NewtonManager._warp_launch._mode == ("replay" if use_cache else "eager") + finally: + if NewtonManager._adapter is not None: + NewtonManager._adapter._invalidate_launch_cache() + NewtonManager._adapter = saved_state["adapter"] + NewtonManager._warp_launch = saved_state["warp_launch"] + NewtonManager._control = saved_state["control"] + PhysicsManager._device = saved_state["device"] + PhysicsManager._cfg = saved_state["cfg"] + NewtonManager._model = saved_state["model"] + NewtonManager._num_envs = saved_state["num_envs"] + if had_active_flag: + NewtonManager._use_newton_actuators_active = saved_state["use_newton_actuators_active"] + else: + del NewtonManager._use_newton_actuators_active From fbe99e6833c0dcfd6b9228119340578a167f9913 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 01:40:22 -0700 Subject: [PATCH 4/5] Add Warp launch mode benchmarks Add a reproducible mixed-kernel matrix and completed-step timing controls to compare eager launch, recorded replay, dynamic setters, and graph composition across launch and thread scales. --- scripts/benchmarks/runtime.py | 31 +- .../benchmarks/test/test_warp_launch_modes.py | 128 +++ scripts/benchmarks/warp_launch_modes.py | 836 ++++++++++++++++++ .../isaaclab/test/benchmark/stepping.py | 129 ++- .../isaaclab/test/benchmark/test_stepping.py | 47 +- 5 files changed, 1156 insertions(+), 15 deletions(-) create mode 100644 scripts/benchmarks/test/test_warp_launch_modes.py create mode 100644 scripts/benchmarks/warp_launch_modes.py diff --git a/scripts/benchmarks/runtime.py b/scripts/benchmarks/runtime.py index e22875097655..7c2f2198f7e9 100644 --- a/scripts/benchmarks/runtime.py +++ b/scripts/benchmarks/runtime.py @@ -42,6 +42,17 @@ def _parse_args(argv: list[str]) -> tuple[argparse.Namespace, list[str]]: parser.add_argument("--task", type=str, required=True, help="Gym task id to benchmark.") parser.add_argument("--num_envs", type=int, default=None, help="Number of parallel environments.") parser.add_argument("--num_frames", type=int, default=100, help="Number of environment steps to benchmark.") + parser.add_argument("--warmup_frames", type=int, default=0, help="Untimed environment steps before measurement.") + parser.add_argument( + "--synchronize_steps", + action="store_true", + help="Synchronize CUDA before and after each step to measure completed-step latency.", + ) + parser.add_argument( + "--reuse_action_buffer", + action="store_true", + help="Reuse one action allocation and randomize its contents in-place between steps.", + ) parser.add_argument("--seed", type=int, default=None, help="Environment seed.") parser.add_argument("--output_path", type=str, default=".", help="Directory to write the output JSON.") parser.add_argument( @@ -123,6 +134,9 @@ def run(argv: list[str]) -> None: {"name": "task", "data": args.task}, {"name": "num_envs", "data": args.num_envs}, {"name": "num_frames", "data": args.num_frames}, + {"name": "warmup_frames", "data": args.warmup_frames}, + {"name": "synchronize_steps", "data": args.synchronize_steps}, + {"name": "reuse_action_buffer", "data": args.reuse_action_buffer}, {"name": "presets", "data": ",".join(cfg.presets)}, ] }, @@ -135,24 +149,31 @@ def run(argv: list[str]) -> None: num_envs = env.unwrapped.num_envs with BenchmarkMonitor(benchmark, interval=1.0): - step_times_s = stepping.run_runtime_loop(env, args.num_frames) + timing = stepping.measure_runtime_loop( + env, + args.num_frames, + warmup_frames=args.warmup_frames, + synchronize_steps=args.synchronize_steps, + reuse_action_buffer=args.reuse_action_buffer, + ) benchmark.update_manual_recorders() startup = StartupTime( app_launch=(app_t1 - app_t0) / 1e9, env_creation=(env_t1 - env_t0) / 1e9, - first_step=(step_times_s[0] if step_times_s else 0.0), + first_step=timing.first_step_s, python_imports=(imports_t1 - imports_t0) / 1e9, task_config=(task_config_t1 - task_config_t0) / 1e9, ) - fps = [num_envs / t for t in step_times_s if t > 0] + step_times_s = timing.step_times_s + instantaneous_fps = [num_envs / t for t in step_times_s if t > 0] runtime = builders.build_runtime( startup_time_s=startup, iteration_times_s=step_times_s, - collection_fps=fps, - total_fps=fps, + collection_fps=instantaneous_fps, + total_fps=instantaneous_fps, steps_per_iteration=num_envs, ) diff --git a/scripts/benchmarks/test/test_warp_launch_modes.py b/scripts/benchmarks/test/test_warp_launch_modes.py new file mode 100644 index 000000000000..7736f952900c --- /dev/null +++ b/scripts/benchmarks/test/test_warp_launch_modes.py @@ -0,0 +1,128 @@ +# 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 + +"""Lightweight tests for the mixed Warp launch-mode benchmark.""" + +import argparse +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +ROOT = Path(__file__).resolve().parents[3] +RUNNER = ROOT / "scripts" / "benchmarks" / "warp_launch_modes.py" + + +@pytest.fixture(scope="module") +def modes_module(): + """Load the standalone benchmark as a module without running its CLI.""" + spec = importlib.util.spec_from_file_location("warp_launch_modes_test_module", RUNNER) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_dry_run_builds_deterministic_cycles_and_launch_counts(modes_module, capsys: pytest.CaptureFixture[str]): + """The dry-run matrix should preserve the four-kernel cycle structure.""" + arguments = [ + "--dry_run", + "--threads", + "256,1024", + "--stage_repeats", + "1,4", + "--modes", + "eager,cache_eager", + "--case_order", + "randomized", + "--case_seed", + "17", + ] + assert modes_module.run(arguments) == 0 + first_output = capsys.readouterr().out + assert modes_module.run(arguments) == 0 + second_output = capsys.readouterr().out + + assert first_output == second_output + assert "Matrix cases: 8" in first_output + assert "eager__threads_256__nodes_4" in first_output + assert "cache_eager__threads_1024__nodes_16" in first_output + + +@pytest.mark.parametrize( + ("mode", "total_executions", "expected_action_checksum"), + [ + ("eager", 3, 2.0), + ("bare_dynamic_normal", 2, 2.0), + ("bare_dynamic_normal", 3, -2.0), + ("bare_dynamic_ctype_constructed", 3, -2.0), + ("bare_dynamic_ctype_prepacked", 3, -2.0), + ], +) +def test_semantic_expectations_track_dynamic_action_bank( + modes_module, mode: str, total_executions: int, expected_action_checksum: float +): + """Analytical checksums should account for alternating dynamic action pointers.""" + case = modes_module._Case(mode=mode, threads=8, stage_repeats=1) + expected = modes_module._semantic_expectations(case, total_executions) + + assert expected == { + "action_checksum": expected_action_checksum, + "observation_checksum": 4.8, + "reward_checksum": 7.896, + "reset_count": 0, + } + + +def test_mode_parser_exposes_explicit_no_conversion_variants(modes_module): + """Descriptor construction and prepacking must remain distinct CLI modes.""" + modes = modes_module._mode_list("bare_dynamic_ctype_constructed,bare_dynamic_ctype_prepacked,cache_eager") + assert modes == ( + "bare_dynamic_ctype_constructed", + "bare_dynamic_ctype_prepacked", + "cache_eager", + ) + with pytest.raises(argparse.ArgumentTypeError, match="unknown modes"): + modes_module._mode_list("bare_dynamic_ctype") + with pytest.raises(argparse.ArgumentTypeError, match="unknown modes"): + modes_module._mode_list("cache_dynamic_changed") + + +def test_default_case_order_is_randomized(modes_module): + """The benchmark should randomize cases by default to limit time-order bias.""" + assert modes_module._parse_args([]).case_order == "randomized" + + +def test_dynamic_probe_rejects_a_noop_setter(modes_module): + """Two same-sign outputs must not pass the explicit dynamic-pointer probe.""" + case = modes_module._Case(mode="bare_dynamic_normal", threads=8, stage_repeats=1) + + modes_module._validate_dynamic_probe(case, (-2.0, 2.0)) + with pytest.raises(RuntimeError, match="Dynamic action probe failed"): + modes_module._validate_dynamic_probe(case, (2.0, 2.0)) + + +def test_graph_activity_selection_prefers_graph_events(modes_module): + """Graph modes should prefer one graph event when child kernel events are also available.""" + case = modes_module._Case(mode="graph_eager", threads=8, stage_repeats=1) + kernel_results = [SimpleNamespace(elapsed=0.25) for _ in range(case.nodes_per_step)] + graph_results = [SimpleNamespace(elapsed=0.5)] + + selected = modes_module._select_gpu_activity(case, kernel_results, graph_results) + + assert selected == (1.0, 0.5, 0.5, "graph", 1) + + +def test_graph_activity_selection_accepts_child_kernel_events(modes_module): + """Graph modes should fall back to child kernels when no graph event is reported.""" + case = modes_module._Case(mode="graph_replay", threads=8, stage_repeats=1) + kernel_results = [SimpleNamespace(elapsed=0.25) for _ in range(case.nodes_per_step)] + + selected = modes_module._select_gpu_activity(case, kernel_results, []) + + assert selected == (1.0, None, 1.0, "kernel", case.nodes_per_step) diff --git a/scripts/benchmarks/warp_launch_modes.py b/scripts/benchmarks/warp_launch_modes.py new file mode 100644 index 000000000000..0dca666d9f1b --- /dev/null +++ b/scripts/benchmarks/warp_launch_modes.py @@ -0,0 +1,836 @@ +# 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 + +"""Benchmark Warp launch modes on a mixed production-like kernel sequence. + +The sequence models the recurring action, observation, reward, and termination +kernel roles in the experimental Warp Cartpole task. It compares eager launches, +bare recorded commands, typed and documented no-conversion command updates, +the Isaac Lab ``WarpLaunchCache`` wrapper, and CUDA graphs. One-time +recording/capture costs, host submission, synchronized wall time, and GPU +activity are reported separately. + +Example:: + + ./isaaclab.sh -p scripts/benchmarks/warp_launch_modes.py \ + --device cuda:0 --threads 256,1024,4096,16384 \ + --stage_repeats 1,4,16,64 --output_dir results/warp_launch_modes +""" + +from __future__ import annotations + +import argparse +import csv +import dataclasses +import datetime +import hashlib +import itertools +import json +import math +import os +import pathlib +import platform +import random +import subprocess +import sys +import time +from collections.abc import Callable, Sequence +from typing import Any + +import warp as wp + +_REPOSITORY_ROOT = pathlib.Path(__file__).resolve().parents[2] +_ISAACLAB_SOURCE = _REPOSITORY_ROOT / "source" / "isaaclab" +if str(_ISAACLAB_SOURCE) not in sys.path: + sys.path.insert(0, str(_ISAACLAB_SOURCE)) + +from isaaclab.utils.warp import WarpLaunchCache # noqa: E402 + + +@wp.kernel +def update_actions( + input_actions: wp.array2d(dtype=wp.float32), + actions: wp.array2d(dtype=wp.float32), + action_scale: wp.float32, + cart_dof_idx: wp.int32, +): + """Apply one persistent action buffer to a Cartpole-shaped output.""" + env_index = wp.tid() + actions[env_index, cart_dof_idx] = action_scale * input_actions[env_index, 0] + + +@wp.kernel +def get_observations( + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + cart_dof_idx: wp.int32, + pole_dof_idx: wp.int32, + observations: wp.array(dtype=wp.vec4f), +): + """Gather four Cartpole-shaped state values.""" + env_index = wp.tid() + observations[env_index][0] = joint_pos[env_index, pole_dof_idx] + observations[env_index][1] = joint_vel[env_index, pole_dof_idx] + observations[env_index][2] = joint_pos[env_index, cart_dof_idx] + observations[env_index][3] = joint_vel[env_index, cart_dof_idx] + + +@wp.kernel +def compute_rewards( + rew_scale_alive: wp.float32, + rew_scale_terminated: wp.float32, + rew_scale_pole_pos: wp.float32, + rew_scale_cart_vel: wp.float32, + rew_scale_pole_vel: wp.float32, + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + cart_dof_idx: wp.int32, + pole_dof_idx: wp.int32, + reset_terminated: wp.array(dtype=wp.bool), + reward: wp.array(dtype=wp.float32), +): + """Compute a representative multi-input reward kernel.""" + env_index = wp.tid() + alive = wp.where(reset_terminated[env_index], wp.float32(0.0), rew_scale_alive) + terminated = wp.where(reset_terminated[env_index], rew_scale_terminated, wp.float32(0.0)) + pole_pos = joint_pos[env_index, pole_dof_idx] + reward[env_index] = ( + alive + + terminated + + rew_scale_pole_pos * pole_pos * pole_pos + + rew_scale_cart_vel * wp.abs(joint_vel[env_index, cart_dof_idx]) + + rew_scale_pole_vel * wp.abs(joint_vel[env_index, pole_dof_idx]) + ) + + +@wp.kernel +def get_dones( + joint_pos: wp.array2d(dtype=wp.float32), + episode_length: wp.array(dtype=wp.int32), + cart_dof_idx: wp.int32, + pole_dof_idx: wp.int32, + max_episode_length: wp.int32, + max_cart_pos: wp.float32, + out_of_bounds: wp.array(dtype=wp.bool), + time_out: wp.array(dtype=wp.bool), + reset: wp.array(dtype=wp.bool), +): + """Compute representative termination flags.""" + env_index = wp.tid() + out_of_bounds[env_index] = (wp.abs(joint_pos[env_index, cart_dof_idx]) > max_cart_pos) or ( + wp.abs(joint_pos[env_index, pole_dof_idx]) > wp.pi / 2.0 + ) + time_out[env_index] = episode_length[env_index] >= max_episode_length - 1 + reset[env_index] = out_of_bounds[env_index] or time_out[env_index] + + +_BASE_NODE_COUNT = 4 +_ALL_MODES = ( + "eager", + "cache_eager", + "bare_replay", + "bare_dynamic_normal", + "bare_dynamic_ctype_constructed", + "bare_dynamic_ctype_prepacked", + "cache_static", + "graph_eager", + "graph_replay", +) +_CSV_FIELDS = ( + "timestamp_utc", + "case_id", + "status", + "message", + "sample_index", + "mode", + "threads", + "stage_repeats", + "cycles_per_step", + "nodes_per_step", + "dynamic_nodes_per_step", + "dynamic_updates_per_step", + "setter_method", + "warmup_executions", + "record_setup_ms", + "capture_setup_ms", + "first_execution_batch_ms", + "cpu_submit_batch_ms", + "cpu_submit_us_per_node", + "synchronized_wall_batch_ms", + "synchronized_wall_us_per_node", + "gpu_kernel_batch_ms", + "gpu_graph_batch_ms", + "gpu_activity_batch_ms", + "gpu_activity_us_per_node", + "gpu_activity_type", + "gpu_activity_count", + "action_checksum", + "observation_checksum", + "reward_checksum", + "reset_count", + "dynamic_probe_first_action_checksum", + "dynamic_probe_second_action_checksum", + "semantic_validation", + "device", +) + + +@dataclasses.dataclass(frozen=True) +class _Case: + """One point in the mixed-sequence benchmark matrix.""" + + mode: str + threads: int + stage_repeats: int + + @property + def nodes_per_step(self) -> int: + return _BASE_NODE_COUNT * self.stage_repeats + + @property + def case_id(self) -> str: + return f"{self.mode}__threads_{self.threads}__nodes_{self.nodes_per_step}" + + +@dataclasses.dataclass +class _Node: + """One retained kernel call in the production-like sequence.""" + + kernel: wp.Kernel + arguments: list[Any] + dynamic_index: int | None = None + + +@dataclasses.dataclass +class _Workload: + """Buffers and nodes retained for one benchmark case.""" + + nodes: list[_Node] + action_banks: tuple[wp.array, wp.array] + actions: wp.array + reward: wp.array + observations: wp.array + reset_mask: wp.array + + +@dataclasses.dataclass +class _PreparedMode: + """Prepared steady-state callback and one-time setup metadata.""" + + execute: Callable[[], None] + record_setup_ms: float + capture_setup_ms: float + dynamic_nodes_per_step: int + dynamic_updates_per_step: int + setter_method: str + + +def run(argv: Sequence[str] | None = None) -> int: + """Run the selected matrix and return a process exit code.""" + args = _parse_args(argv) + cases = [ + _Case(mode=mode, threads=threads, stage_repeats=stage_repeats) + for mode, threads, stage_repeats in itertools.product(args.modes, args.threads, args.stage_repeats) + ] + if args.case_order == "randomized": + random.Random(args.case_seed).shuffle(cases) + if args.dry_run: + for case in cases: + print(case.case_id) + print(f"Matrix cases: {len(cases)}") + return 0 + + wp.init() + device = wp.get_device(args.device) + if not device.is_cuda: + raise RuntimeError("This benchmark requires a CUDA device because two modes use CUDA graphs.") + + started_utc = _utc_now() + run_name = args.run_name or f"warp_launch_modes_{started_utc.replace('-', '').replace(':', '')[:15]}Z" + _validate_run_name(run_name) + output_dir = pathlib.Path(args.output_dir).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + csv_path = output_dir / f"{run_name}.csv" + json_path = output_dir / f"{run_name}.json" + if csv_path.exists() or json_path.exists(): + raise FileExistsError(f"Refusing to overwrite existing benchmark output for run {run_name!r}.") + + kernel_source = pathlib.Path(__file__).resolve() + cache_source = pathlib.Path(sys.modules[WarpLaunchCache.__module__].__file__).resolve() + manifest: dict[str, Any] = { + "schema_version": 1, + "status": "running", + "started_utc": started_utc, + "completed_utc": None, + "command": [sys.executable, str(pathlib.Path(__file__).resolve()), *sys.argv[1:]], + "arguments": vars(args), + "matrix": {"total_cases": len(cases), "completed_cases": 0, "failed_cases": 0}, + "software": { + "python": platform.python_version(), + "warp": wp.__version__, + "warp_path": str(pathlib.Path(wp.__file__).resolve()), + "git_commit": _git_commit(), + }, + "sources": { + "benchmark": _file_metadata(pathlib.Path(__file__)), + "warp_launch_cache": _file_metadata(cache_source), + }, + "device": _device_metadata(device), + "kernel_sequence": { + "source": str(kernel_source), + "source_sha256": hashlib.sha256(kernel_source.read_bytes()).hexdigest(), + "base_nodes": ["update_actions", "get_observations", "compute_rewards", "get_dones"], + }, + "metric_notes": { + "cpu_submit_batch_ms": "Host time to submit one complete sequence; trailing synchronization excluded.", + "synchronized_wall_batch_ms": "Wall time with device synchronization before and after the sequence.", + "gpu_activity_batch_ms": "Warp CUDA activity attributed to kernels or the graph during the sequence.", + "record_setup_ms": ( + "Host time to record every command/replay-cache entry; replay-cache setup also submits its first call." + ), + "capture_setup_ms": "Host wall time to construct the CUDA graph after kernel precompilation.", + "dynamic_updates_per_step": ( + "Actual setter calls expected in steady state; unchanged cache tokens are elided." + ), + "bare_dynamic_ctype_constructed": ( + "Constructs an array descriptor inside every update, then uses Warp's documented " + "set_param_at_index_from_ctype no-conversion API." + ), + "bare_dynamic_ctype_prepacked": ( + "Reuses retained array descriptors constructed before steady-state timing, then uses Warp's " + "documented set_param_at_index_from_ctype no-conversion API." + ), + }, + "rows": [], + } + _write_outputs(csv_path, json_path, manifest, []) + + _precompile(device) + rows: list[dict[str, Any]] = [] + failed_cases = 0 + for case_index, case in enumerate(cases, start=1): + print(f"[{case_index}/{len(cases)}] {case.case_id}", flush=True) + try: + case_rows = _run_case(case, args.warmup_steps, args.warmup_seconds, args.repeats, device) + except Exception as exc: + failed_cases += 1 + case_rows = [ + _make_row( + case, + status="failed", + message=f"{type(exc).__name__}: {exc}", + sample_index=None, + device=device.alias, + ) + ] + print(f"FAILED {case.case_id}: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True) + rows.extend(case_rows) + manifest["rows"] = rows + manifest["matrix"]["completed_cases"] = case_index - failed_cases + manifest["matrix"]["failed_cases"] = failed_cases + _write_outputs(csv_path, json_path, manifest, rows) + + manifest["status"] = "completed" if failed_cases == 0 else "completed_with_failures" + manifest["completed_utc"] = _utc_now() + _write_outputs(csv_path, json_path, manifest, rows) + print(f"Cases: {len(cases) - failed_cases} completed, {failed_cases} failed") + print(f"CSV: {csv_path}") + return 0 if failed_cases == 0 else 1 + + +def _parse_args(argv: Sequence[str] | None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Benchmark a mixed sequence of recurring Warp Cartpole kernels.") + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--threads", type=_positive_int_list, default=(256, 1024, 4096, 16384)) + parser.add_argument("--stage_repeats", type=_positive_int_list, default=(1, 4, 16, 64)) + parser.add_argument("--modes", type=_mode_list, default=_ALL_MODES) + parser.add_argument("--warmup_steps", type=_nonnegative_int, default=5) + parser.add_argument( + "--warmup_seconds", + type=_nonnegative_float, + default=1.0, + help="Minimum synchronized warm-up duration per case to stabilize adaptive clocks.", + ) + parser.add_argument("--repeats", type=_positive_int, default=10) + parser.add_argument("--case_order", choices=("grouped", "randomized"), default="randomized") + parser.add_argument("--case_seed", type=int, default=271828) + parser.add_argument("--output_dir", default="results/warp_launch_modes") + parser.add_argument("--run_name", default=None) + parser.add_argument("--dry_run", action="store_true") + return parser.parse_args(argv) + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def _nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("value must be nonnegative") + return parsed + + +def _nonnegative_float(value: str) -> float: + parsed = float(value) + if parsed < 0.0: + raise argparse.ArgumentTypeError("value must be nonnegative") + return parsed + + +def _positive_int_list(value: str) -> tuple[int, ...]: + parsed = tuple(dict.fromkeys(_positive_int(item.strip()) for item in value.split(",") if item.strip())) + if not parsed: + raise argparse.ArgumentTypeError("at least one positive integer is required") + return parsed + + +def _mode_list(value: str) -> tuple[str, ...]: + parsed = tuple(dict.fromkeys(item.strip() for item in value.split(",") if item.strip())) + unknown = [mode for mode in parsed if mode not in _ALL_MODES] + if unknown: + raise argparse.ArgumentTypeError(f"unknown modes {unknown}; choose from {', '.join(_ALL_MODES)}") + if not parsed: + raise argparse.ArgumentTypeError("at least one mode is required") + return parsed + + +def _precompile(device: wp.Device) -> None: + """Compile every production kernel before setup timing.""" + workload = _make_workload(1, 1, device) + for node in workload.nodes: + wp.launch(node.kernel, dim=1, inputs=node.arguments, device=device) + wp.synchronize_device(device) + + +def _make_workload(threads: int, stage_repeats: int, device: wp.Device) -> _Workload: + """Allocate Cartpole-shaped buffers and construct the repeated node sequence.""" + joint_pos = wp.full((threads, 2), 0.1, dtype=wp.float32, device=device) + joint_vel = wp.full((threads, 2), 0.2, dtype=wp.float32, device=device) + action_banks = ( + wp.full((threads, 1), 0.25, dtype=wp.float32, device=device), + wp.full((threads, 1), -0.25, dtype=wp.float32, device=device), + ) + actions = wp.zeros((threads, 2), dtype=wp.float32, device=device) + observations = wp.zeros(threads, dtype=wp.vec4f, device=device) + episode_length = wp.zeros(threads, dtype=wp.int32, device=device) + out_of_bounds = wp.zeros(threads, dtype=wp.bool, device=device) + time_out = wp.zeros(threads, dtype=wp.bool, device=device) + reset_mask = wp.zeros(threads, dtype=wp.bool, device=device) + reward = wp.zeros(threads, dtype=wp.float32, device=device) + + nodes: list[_Node] = [] + for _ in range(stage_repeats): + nodes.extend( + ( + _Node( + update_actions, + [action_banks[0], actions, 1.0, 0], + dynamic_index=0, + ), + _Node( + get_observations, + [joint_pos, joint_vel, 0, 1, observations], + ), + _Node( + compute_rewards, + [1.0, -2.0, -1.0, -0.01, -0.005, joint_pos, joint_vel, 0, 1, out_of_bounds, reward], + ), + _Node( + get_dones, + [joint_pos, episode_length, 0, 1, 500, 3.0, out_of_bounds, time_out, reset_mask], + ), + ) + ) + return _Workload( + nodes=nodes, + action_banks=action_banks, + actions=actions, + reward=reward, + observations=observations, + reset_mask=reset_mask, + ) + + +def _record_commands(workload: _Workload, threads: int, device: wp.Device) -> tuple[list[wp.Launch], float]: + wp.synchronize_device(device) + start_ns = time.perf_counter_ns() + commands = [ + wp.launch(node.kernel, dim=threads, inputs=node.arguments, device=device, record_cmd=True) + for node in workload.nodes + ] + record_setup_ms = (time.perf_counter_ns() - start_ns) / 1.0e6 + if any(command is None for command in commands): + raise RuntimeError("wp.launch(record_cmd=True) did not return a command for every sequence node.") + return commands, record_setup_ms # type: ignore[return-value] + + +def _prepare_mode(case: _Case, workload: _Workload, device: wp.Device) -> _PreparedMode: + """Prepare one execution path and time its one-time setup.""" + record_setup_ms = 0.0 + capture_setup_ms = 0.0 + dynamic_nodes = sum(node.dynamic_index is not None for node in workload.nodes) + + if case.mode == "eager": + + def execute() -> None: + for node in workload.nodes: + wp.launch(node.kernel, dim=case.threads, inputs=node.arguments, device=device) + + return _PreparedMode(execute, 0.0, 0.0, 0, 0, "none") + + if case.mode in { + "bare_replay", + "bare_dynamic_normal", + "bare_dynamic_ctype_constructed", + "bare_dynamic_ctype_prepacked", + "graph_replay", + }: + commands, record_setup_ms = _record_commands(workload, case.threads, device) + else: + commands = [] + + if case.mode == "bare_replay": + + def execute() -> None: + for command in commands: + command.launch() + + return _PreparedMode(execute, record_setup_ms, 0.0, 0, 0, "none") + + if case.mode in { + "bare_dynamic_normal", + "bare_dynamic_ctype_constructed", + "bare_dynamic_ctype_prepacked", + }: + execution_index = 0 + action_descriptors = tuple(action.__ctype__() for action in workload.action_banks) + + def execute() -> None: + nonlocal execution_index + action_index = (execution_index + 1) & 1 + action = workload.action_banks[action_index] + for node, command in zip(workload.nodes, commands): + if node.dynamic_index is not None: + if case.mode == "bare_dynamic_normal": + command.set_param_at_index(node.dynamic_index, action) + elif case.mode == "bare_dynamic_ctype_constructed": + command.set_param_at_index_from_ctype(node.dynamic_index, action.__ctype__()) + else: + command.set_param_at_index_from_ctype(node.dynamic_index, action_descriptors[action_index]) + command.launch() + execution_index += 1 + + setter = { + "bare_dynamic_normal": "set_param_at_index", + "bare_dynamic_ctype_constructed": "set_param_at_index_from_ctype + array.__ctype__", + "bare_dynamic_ctype_prepacked": "set_param_at_index_from_ctype (prepacked)", + }[case.mode] + return _PreparedMode(execute, record_setup_ms, 0.0, dynamic_nodes, dynamic_nodes, setter) + + if case.mode in {"cache_eager", "cache_static"}: + cache = WarpLaunchCache( + mode="eager" if case.mode == "cache_eager" else "replay", + debug=False, + device=device, + ) + if case.mode != "cache_eager": + wp.synchronize_device(device) + start_ns = time.perf_counter_ns() + for node in workload.nodes: + cache.launch( + node.kernel, + dim=case.threads, + inputs=node.arguments, + ) + record_setup_ms = (time.perf_counter_ns() - start_ns) / 1.0e6 + + def execute() -> None: + for node in workload.nodes: + cache.launch( + node.kernel, + dim=case.threads, + inputs=node.arguments, + ) + + return _PreparedMode(execute, record_setup_ms, 0.0, 0, 0, "none") + + if case.mode in {"graph_eager", "graph_replay"}: + wp.synchronize_device(device) + start_ns = time.perf_counter_ns() + with wp.ScopedCapture(device=device) as capture: + if case.mode == "graph_eager": + for node in workload.nodes: + wp.launch(node.kernel, dim=case.threads, inputs=node.arguments, device=device) + else: + for command in commands: + command.launch() + capture_setup_ms = (time.perf_counter_ns() - start_ns) / 1.0e6 + graph = capture.graph + + def execute() -> None: + wp.capture_launch(graph) + + return _PreparedMode(execute, record_setup_ms, capture_setup_ms, 0, 0, "none") + + raise ValueError(f"Unknown mode: {case.mode}") + + +def _run_case( + case: _Case, + warmup_steps: int, + warmup_seconds: float, + repeats: int, + device: wp.Device, +) -> list[dict[str, Any]]: + workload = _make_workload(case.threads, case.stage_repeats, device) + prepared = _prepare_mode(case, workload, device) + + with wp.ScopedTimer("first_execution", print=False, synchronize=True) as first_timer: + prepared.execute() + warmup_executions = _warm_up(prepared.execute, warmup_steps, warmup_seconds, device) + + rows: list[dict[str, Any]] = [] + for sample_index in range(repeats): + wp.synchronize_device(device) + submit_start_ns = time.perf_counter_ns() + prepared.execute() + cpu_submit_batch_ms = (time.perf_counter_ns() - submit_start_ns) / 1.0e6 + wp.synchronize_device(device) + + with wp.ScopedTimer( + "synchronized_sequence", + print=False, + synchronize=True, + cuda_filter=wp.TIMING_KERNEL | wp.TIMING_GRAPH, + ) as synchronized_timer: + prepared.execute() + + kernel_results = [timing for timing in synchronized_timer.timing_results if timing.filter & wp.TIMING_KERNEL] + graph_results = [timing for timing in synchronized_timer.timing_results if timing.filter & wp.TIMING_GRAPH] + ( + gpu_kernel_batch_ms, + gpu_graph_batch_ms, + gpu_activity_batch_ms, + gpu_activity_type, + gpu_activity_count, + ) = _select_gpu_activity(case, kernel_results, graph_results) + expected_activity_count = ( + 1 if case.mode.startswith("graph_") and gpu_activity_type == "graph" else case.nodes_per_step + ) + if gpu_activity_count != expected_activity_count: + raise RuntimeError( + f"GPU activity count mismatch for {case.case_id}: measured {gpu_activity_count}, " + f"expected {expected_activity_count}." + ) + + rows.append( + _make_row( + case, + status="ok", + sample_index=sample_index, + dynamic_nodes_per_step=prepared.dynamic_nodes_per_step, + dynamic_updates_per_step=prepared.dynamic_updates_per_step, + setter_method=prepared.setter_method, + warmup_executions=warmup_executions, + record_setup_ms=prepared.record_setup_ms, + capture_setup_ms=prepared.capture_setup_ms, + first_execution_batch_ms=first_timer.elapsed, + cpu_submit_batch_ms=cpu_submit_batch_ms, + cpu_submit_us_per_node=cpu_submit_batch_ms * 1_000.0 / case.nodes_per_step, + synchronized_wall_batch_ms=synchronized_timer.elapsed, + synchronized_wall_us_per_node=synchronized_timer.elapsed * 1_000.0 / case.nodes_per_step, + gpu_kernel_batch_ms=gpu_kernel_batch_ms, + gpu_graph_batch_ms=gpu_graph_batch_ms, + gpu_activity_batch_ms=gpu_activity_batch_ms, + gpu_activity_us_per_node=( + gpu_activity_batch_ms * 1_000.0 / case.nodes_per_step if gpu_activity_batch_ms is not None else None + ), + gpu_activity_type=gpu_activity_type, + gpu_activity_count=gpu_activity_count, + device=device.alias, + ) + ) + + action_checksum = float(workload.actions.numpy().sum()) + observation_checksum = float(workload.observations.numpy().sum()) + reward_checksum = float(workload.reward.numpy().sum()) + reset_count = int(workload.reset_mask.numpy().sum()) + total_executions = 1 + warmup_executions + 2 * repeats + expected = _semantic_expectations(case, total_executions) + actual = { + "action_checksum": action_checksum, + "observation_checksum": observation_checksum, + "reward_checksum": reward_checksum, + "reset_count": reset_count, + } + if not all(math.isfinite(value) for value in (action_checksum, observation_checksum, reward_checksum)): + raise RuntimeError(f"Non-finite semantic checksum: {actual}") + if ( + not math.isclose(action_checksum, expected["action_checksum"], rel_tol=1.0e-5, abs_tol=1.0e-5) + or not math.isclose(observation_checksum, expected["observation_checksum"], rel_tol=1.0e-5, abs_tol=1.0e-5) + or not math.isclose(reward_checksum, expected["reward_checksum"], rel_tol=1.0e-5, abs_tol=1.0e-5) + or reset_count != expected["reset_count"] + ): + raise RuntimeError(f"Semantic validation failed: actual={actual}, expected={expected}") + + dynamic_probe_checksums: tuple[float | None, float | None] = (None, None) + if _is_changed_dynamic_mode(case.mode): + probe_checksums = [] + for _ in range(2): + prepared.execute() + wp.synchronize_device(device) + probe_checksums.append(float(workload.actions.numpy().sum())) + _validate_dynamic_probe(case, probe_checksums) + dynamic_probe_checksums = (probe_checksums[0], probe_checksums[1]) + for row in rows: + row.update(actual) + row["dynamic_probe_first_action_checksum"] = dynamic_probe_checksums[0] + row["dynamic_probe_second_action_checksum"] = dynamic_probe_checksums[1] + row["semantic_validation"] = "passed" + return rows + + +def _warm_up(execute: Callable[[], None], minimum_steps: int, minimum_seconds: float, device: wp.Device) -> int: + """Warm up a case until both count and elapsed-time thresholds are met.""" + executions = 0 + start = time.perf_counter() + while executions < minimum_steps or time.perf_counter() - start < minimum_seconds: + execute() + executions += 1 + if executions % 64 == 0: + wp.synchronize_device(device) + wp.synchronize_device(device) + return executions + + +def _select_gpu_activity( + case: _Case, + kernel_results: Sequence[Any], + graph_results: Sequence[Any], +) -> tuple[float | None, float | None, float | None, str, int]: + """Select the GPU activity representation appropriate for an execution mode.""" + gpu_kernel_batch_ms = sum(timing.elapsed for timing in kernel_results) if kernel_results else None + gpu_graph_batch_ms = sum(timing.elapsed for timing in graph_results) if graph_results else None + if case.mode.startswith("graph_") and gpu_graph_batch_ms is not None: + return gpu_kernel_batch_ms, gpu_graph_batch_ms, gpu_graph_batch_ms, "graph", len(graph_results) + if gpu_kernel_batch_ms is not None: + return gpu_kernel_batch_ms, gpu_graph_batch_ms, gpu_kernel_batch_ms, "kernel", len(kernel_results) + if gpu_graph_batch_ms is not None: + return gpu_kernel_batch_ms, gpu_graph_batch_ms, gpu_graph_batch_ms, "graph", len(graph_results) + return None, None, None, "none", 0 + + +def _semantic_expectations(case: _Case, total_executions: int) -> dict[str, float | int]: + """Return analytical checksums for the final state of a benchmark case.""" + final_action_sign = -1.0 if _is_changed_dynamic_mode(case.mode) and total_executions % 2 else 1.0 + return { + "action_checksum": final_action_sign * 0.25 * case.threads, + "observation_checksum": 0.6 * case.threads, + "reward_checksum": 0.987 * case.threads, + "reset_count": 0, + } + + +def _is_changed_dynamic_mode(mode: str) -> bool: + """Return whether an execution mode alternates its dynamic action pointer.""" + return mode in { + "bare_dynamic_normal", + "bare_dynamic_ctype_constructed", + "bare_dynamic_ctype_prepacked", + } + + +def _validate_dynamic_probe(case: _Case, checksums: Sequence[float]) -> None: + """Validate that two explicit post-timing executions used opposite action banks.""" + expected_magnitude = 0.25 * case.threads + if len(checksums) != 2 or not ( + math.isclose(checksums[0], -checksums[1], rel_tol=1.0e-5, abs_tol=1.0e-5) + and math.isclose(abs(checksums[0]), expected_magnitude, rel_tol=1.0e-5, abs_tol=1.0e-5) + ): + raise RuntimeError( + f"Dynamic action probe failed for {case.case_id}: measured {checksums}, " + f"expected alternating +/-{expected_magnitude}." + ) + + +def _make_row(case: _Case, **updates: Any) -> dict[str, Any]: + row: dict[str, Any] = dict.fromkeys(_CSV_FIELDS) + row.update( + { + "timestamp_utc": _utc_now(), + "case_id": case.case_id, + "mode": case.mode, + "threads": case.threads, + "stage_repeats": case.stage_repeats, + "cycles_per_step": case.stage_repeats, + "nodes_per_step": case.nodes_per_step, + } + ) + row.update(updates) + return row + + +def _write_outputs( + csv_path: pathlib.Path, + json_path: pathlib.Path, + manifest: dict[str, Any], + rows: list[dict[str, Any]], +) -> None: + csv_temporary = csv_path.with_name(f".{csv_path.name}.{os.getpid()}.tmp") + json_temporary = json_path.with_name(f".{json_path.name}.{os.getpid()}.tmp") + with csv_temporary.open("w", newline="", encoding="utf-8") as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=_CSV_FIELDS) + writer.writeheader() + writer.writerows(rows) + os.replace(csv_temporary, csv_path) + with json_temporary.open("w", encoding="utf-8") as json_file: + json.dump(manifest, json_file, indent=2, sort_keys=True, allow_nan=False) + json_file.write("\n") + os.replace(json_temporary, json_path) + + +def _device_metadata(device: wp.Device) -> dict[str, Any]: + return { + "alias": device.alias, + "name": device.name, + "arch": getattr(device, "arch", None), + "ordinal": getattr(device, "ordinal", None), + "uuid": str(getattr(device, "uuid", "")), + "is_cuda": device.is_cuda, + } + + +def _git_commit() -> str | None: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=_REPOSITORY_ROOT, + check=False, + capture_output=True, + text=True, + ) + return result.stdout.strip() if result.returncode == 0 else None + + +def _file_metadata(path: pathlib.Path) -> dict[str, str]: + """Return an absolute source path and its content digest.""" + path = path.resolve() + return {"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()} + + +def _validate_run_name(run_name: str) -> None: + if not run_name or pathlib.Path(run_name).name != run_name: + raise ValueError("run_name must be a non-empty filename stem without path separators") + + +def _utc_now() -> str: + return datetime.datetime.now(datetime.UTC).isoformat() + + +if __name__ == "__main__": + raise SystemExit(run(sys.argv[1:])) diff --git a/source/isaaclab/isaaclab/test/benchmark/stepping.py b/source/isaaclab/isaaclab/test/benchmark/stepping.py index 9536d626765d..735c0e087ecd 100644 --- a/source/isaaclab/isaaclab/test/benchmark/stepping.py +++ b/source/isaaclab/isaaclab/test/benchmark/stepping.py @@ -13,6 +13,7 @@ from __future__ import annotations import time +from dataclasses import dataclass from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -59,32 +60,142 @@ def sample_random_actions(env) -> torch.Tensor | dict[str, torch.Tensor]: return 2.0 * torch.rand(u.num_envs, u.single_action_space.shape[0], device=u.device) - 1.0 -def run_runtime_loop(env, num_frames: int) -> list[float]: - """Step the environment ``num_frames`` times and record per-step wall times [s]. +@dataclass(frozen=True) +class RuntimeLoopTiming: + """Timing samples collected while stepping an environment.""" - Calls ``env.reset()`` once before the loop, then on each frame samples - random actions via :func:`sample_random_actions`, steps the environment, - and records the elapsed wall-clock time for that step. + first_step_s: float + """Wall time of the first environment step [s].""" + + step_times_s: list[float] + """Environment step wall times collected after the requested warm-up [s].""" + + +def measure_runtime_loop( + env, + num_frames: int, + *, + warmup_frames: int = 0, + synchronize_steps: bool = False, + reuse_action_buffer: bool = False, +) -> RuntimeLoopTiming: + """Step an environment and measure startup and per-step wall times. + + Calls ``env.reset()`` once, runs untimed warm-up frames, then records the + requested frames. Action sampling stays outside the timed region. Optional + CUDA synchronization measures completed-step latency rather than host + submission latency. Optional action-buffer reuse changes values in-place so + pointer-sensitive frontends can be measured independently from allocation. Args: env: A Gym-compatible environment. num_frames: Number of environment steps to run. + warmup_frames: Number of untimed environment steps before measurement. + synchronize_steps: Whether to synchronize the environment CUDA device + immediately before and after every measured step. + reuse_action_buffer: Whether to reuse one action allocation and randomize + its contents in-place between steps. Returns: - A list of length ``num_frames`` containing per-step wall times [s]. + The first-step time and a list of length ``num_frames`` containing + per-step wall times after the requested warm-up [s]. With no warm-up, + the first measured sample is also the reported first step. """ - env.reset() + if num_frames < 0: + raise ValueError(f"num_frames must be non-negative, got {num_frames}.") + if warmup_frames < 0: + raise ValueError(f"warmup_frames must be non-negative, got {warmup_frames}.") + env.reset() + reusable_actions = sample_random_actions(env) if reuse_action_buffer else None step_times: list[float] = [] + def next_actions(): + if reusable_actions is None: + return sample_random_actions(env) + _randomize_actions_in_place(reusable_actions) + return reusable_actions + + first_step_s = 0.0 + for warmup_index in range(warmup_frames): + actions = next_actions() + if warmup_index == 0: + if synchronize_steps: + _synchronize_env_device(env) + t0 = time.perf_counter_ns() + env.step(actions) + if synchronize_steps: + _synchronize_env_device(env) + t1 = time.perf_counter_ns() + first_step_s = (t1 - t0) / 1e9 + else: + env.step(actions) + if synchronize_steps: + _synchronize_env_device(env) + for _ in range(num_frames): - actions = sample_random_actions(env) + actions = next_actions() + if synchronize_steps: + _synchronize_env_device(env) t0 = time.perf_counter_ns() env.step(actions) + if synchronize_steps: + _synchronize_env_device(env) t1 = time.perf_counter_ns() step_times.append((t1 - t0) / 1e9) + if first_step_s == 0.0 and len(step_times) == 1: + first_step_s = step_times[0] + + return RuntimeLoopTiming(first_step_s=first_step_s, step_times_s=step_times) + + +def run_runtime_loop( + env, + num_frames: int, + *, + warmup_frames: int = 0, + synchronize_steps: bool = False, + reuse_action_buffer: bool = False, +) -> list[float]: + """Step an environment and return per-step wall times after warm-up [s]. + + This compatibility wrapper delegates to :func:`measure_runtime_loop`. + + Args: + env: A Gym-compatible environment. + num_frames: Number of environment steps to measure. + warmup_frames: Number of untimed environment steps before measurement. + synchronize_steps: Whether to synchronize the environment CUDA device + immediately before and after every measured step. + reuse_action_buffer: Whether to reuse one action allocation and randomize + its contents in-place between steps. + + Returns: + A list of length ``num_frames`` containing per-step wall times [s]. + """ + return measure_runtime_loop( + env, + num_frames, + warmup_frames=warmup_frames, + synchronize_steps=synchronize_steps, + reuse_action_buffer=reuse_action_buffer, + ).step_times_s + + +def _randomize_actions_in_place(actions: torch.Tensor | dict[str, torch.Tensor]) -> None: + """Fill reusable single-agent or multi-agent action tensors in-place.""" + tensors = actions.values() if isinstance(actions, dict) else (actions,) + for tensor in tensors: + tensor.uniform_(-1.0, 1.0) + + +def _synchronize_env_device(env) -> None: + """Synchronize the environment CUDA device when it uses CUDA.""" + import torch # noqa: PLC0415 - return step_times + device = torch.device(env.unwrapped.device) + if device.type == "cuda": + torch.cuda.synchronize(device) def _extract_success(extras) -> float | None: diff --git a/source/isaaclab/test/benchmark/test_stepping.py b/source/isaaclab/test/benchmark/test_stepping.py index d2c53117230c..93fe0c707e9d 100644 --- a/source/isaaclab/test/benchmark/test_stepping.py +++ b/source/isaaclab/test/benchmark/test_stepping.py @@ -6,9 +6,11 @@ """Tests for the runtime stepping helpers.""" import numpy as np +import pytest import torch -from isaaclab.test.benchmark.stepping import run_runtime_loop, sample_random_actions +from isaaclab.test.benchmark import stepping +from isaaclab.test.benchmark.stepping import measure_runtime_loop, run_runtime_loop, sample_random_actions class _Space: @@ -26,12 +28,14 @@ def __init__(self): self.unwrapped = _Env._U() self.reset_called = False self.steps = 0 + self.action_ids = [] def reset(self): self.reset_called = True def step(self, actions): self.steps += 1 + self.action_ids.append(id(actions)) return (None, None, None, {}) @@ -49,6 +53,47 @@ def test_run_runtime_loop_steps_and_times(): assert len(times) == 5 and all(t >= 0.0 for t in times) +def test_run_runtime_loop_excludes_warmup_and_reuses_actions(): + env = _Env() + times = run_runtime_loop(env, num_frames=3, warmup_frames=2, reuse_action_buffer=True) + assert env.steps == 5 + assert len(times) == 3 + assert len(set(env.action_ids)) == 1 + + +def test_measure_runtime_loop_reports_first_step_separately(): + env = _Env() + timing = measure_runtime_loop(env, num_frames=3, warmup_frames=2) + assert timing.first_step_s >= 0.0 + assert len(timing.step_times_s) == 3 + assert env.steps == 5 + + +def test_measure_runtime_loop_excludes_startup_action_sampling_and_synchronizes(monkeypatch): + """The startup sample should time one completed step without action-generation work.""" + events = [] + clock = iter((100, 200)) + env = _Env() + + monkeypatch.setattr(stepping, "sample_random_actions", lambda _env: events.append("sample") or torch.zeros((4, 3))) + monkeypatch.setattr(stepping, "_synchronize_env_device", lambda _env: events.append("sync")) + monkeypatch.setattr(stepping.time, "perf_counter_ns", lambda: events.append("clock") or next(clock)) + original_step = env.step + env.step = lambda actions: events.append("step") or original_step(actions) + + timing = measure_runtime_loop(env, num_frames=0, warmup_frames=1, synchronize_steps=True) + + assert timing.first_step_s == pytest.approx(1e-7) + assert events[:6] == ["sample", "sync", "clock", "step", "sync", "clock"] + + +def test_run_runtime_loop_rejects_negative_frame_counts(): + env = _Env() + for kwargs in ({"num_frames": -1}, {"num_frames": 1, "warmup_frames": -1}): + with pytest.raises(ValueError): + run_runtime_loop(env, **kwargs) + + class _MASpace: def __init__(self, n): self._n = n From f75ec2ba91d9a393a00b4d19dc53125648dda8c3 Mon Sep 17 00:00:00 2001 From: jichuanh Date: Thu, 16 Jul 2026 18:35:58 -0700 Subject: [PATCH 5/5] Invalidate Warp caches after curriculum changes Runtime manager-term updates can replace scalar or callable inputs that are baked into captured graphs and recorded launches. Invalidate the owner caches only when a curriculum actually writes a new value. --- .../warp-launch-record-replay.minor.rst | 5 +++ .../isaaclab/isaaclab/envs/mdp/curriculums.py | 13 ++++++++ .../envs/test_modify_env_param_curr_term.py | 33 +++++++++++++++++++ 3 files changed, 51 insertions(+) diff --git a/source/isaaclab/changelog.d/warp-launch-record-replay.minor.rst b/source/isaaclab/changelog.d/warp-launch-record-replay.minor.rst index 2ca617fdb126..e0350325258e 100644 --- a/source/isaaclab/changelog.d/warp-launch-record-replay.minor.rst +++ b/source/isaaclab/changelog.d/warp-launch-record-replay.minor.rst @@ -3,3 +3,8 @@ Added * Added :class:`~isaaclab.utils.warp.WarpLaunchCache` for low-overhead replay of pointer-stable Warp kernel launches outside CUDA graphs. + +Fixed +^^^^^ + +* Fixed runtime manager-term updates to invalidate cached Warp work. diff --git a/source/isaaclab/isaaclab/envs/mdp/curriculums.py b/source/isaaclab/isaaclab/envs/mdp/curriculums.py index 8438e5bec925..d48505736152 100644 --- a/source/isaaclab/isaaclab/envs/mdp/curriculums.py +++ b/source/isaaclab/isaaclab/envs/mdp/curriculums.py @@ -294,3 +294,16 @@ def __init__(self, cfg, env): super().__init__(cfg, env) # overwrite the simplified address with the full manager path self._address = self._address.replace("s.", "_manager.cfg.", 1) + + def _process_accessors(self, root: ManagerBasedRLEnv, path: str) -> tuple[callable, callable]: + """Build accessors that invalidate cached Warp work after a term changes.""" + get_value, set_value = super()._process_accessors(root, path) + invalidate_wp_graphs = getattr(root, "invalidate_wp_graphs", None) + if not callable(invalidate_wp_graphs): + return get_value, set_value + + def set_value_and_invalidate(value): + set_value(value) + invalidate_wp_graphs() + + return get_value, set_value_and_invalidate diff --git a/source/isaaclab/test/envs/test_modify_env_param_curr_term.py b/source/isaaclab/test/envs/test_modify_env_param_curr_term.py index d4b4c1297662..3bb7316d8592 100644 --- a/source/isaaclab/test/envs/test_modify_env_param_curr_term.py +++ b/source/isaaclab/test/envs/test_modify_env_param_curr_term.py @@ -5,6 +5,8 @@ """Test texture randomization in the cartpole scene using pytest.""" +from types import SimpleNamespace + from isaaclab.app import AppLauncher # launch omniverse app @@ -32,6 +34,37 @@ def replace_value(env, env_id, data, value, num_steps): return mdp.modify_env_param.NO_CHANGE +def test_modify_term_cfg_invalidates_warp_caches_on_change(): + """Ensure changed term parameters invalidate cached Warp work exactly once.""" + + class DummyEnv: + def __init__(self): + term_cfg = SimpleNamespace(params={"scale": 1.0}) + self.reward_manager = SimpleNamespace(cfg=SimpleNamespace(term=term_cfg)) + self.invalidations = 0 + + def invalidate_wp_graphs(self): + self.invalidations += 1 + + env = DummyEnv() + cfg = CurrTerm( + func=mdp.modify_term_cfg, + params={ + "address": "rewards.term.params.scale", + "modify_fn": lambda env, env_ids, data: 2.0, + }, + ) + term = mdp.modify_term_cfg(cfg, env) + + term(env, [], **cfg.params) + assert env.reward_manager.cfg.term.params["scale"] == 2.0 + assert env.invalidations == 1 + + cfg.params["modify_fn"] = lambda env, env_ids, data: mdp.modify_term_cfg.NO_CHANGE + term(env, [], **cfg.params) + assert env.invalidations == 1 + + @configclass class CurriculumsCfg: modify_observation_joint_pos = CurrTerm(