From 23007585fa2762a772907d5204598199579895cd Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 22:03:17 -0700 Subject: [PATCH 01/18] Skip wrench composer resets while inactive --- .../changelog.d/perf-hot-path-7b-newton.rst | 6 +++++ .../isaaclab/utils/wrench_composer.py | 4 +++- .../test/utils/test_wrench_composer.py | 22 +++++++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 source/isaaclab/changelog.d/perf-hot-path-7b-newton.rst diff --git a/source/isaaclab/changelog.d/perf-hot-path-7b-newton.rst b/source/isaaclab/changelog.d/perf-hot-path-7b-newton.rst new file mode 100644 index 000000000000..2a449b65f706 --- /dev/null +++ b/source/isaaclab/changelog.d/perf-hot-path-7b-newton.rst @@ -0,0 +1,6 @@ +Changed +^^^^^^^ + +* Changed :meth:`~isaaclab.utils.wrench_composer.WrenchComposer.reset` to return without launching work + when the composer is inactive, removing the per-step buffer clears for rigid objects and collections + that apply no external wrench. diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py index 4d859090d409..6331716cad5d 100644 --- a/source/isaaclab/isaaclab/utils/wrench_composer.py +++ b/source/isaaclab/isaaclab/utils/wrench_composer.py @@ -598,7 +598,7 @@ def reset( With no selection or ``env_ids=slice(None)``, zeros all seven buffers (5 input + 2 output) and clears all flags. Other ``env_ids`` or ``env_mask`` values perform a partial reset on the specified environments using - the reset kernels. + the reset kernels. An inactive composer already holds zero buffers, so any reset of it is a no-op. .. caution:: If both ``env_ids`` and ``env_mask`` are provided, ``env_mask`` takes precedence. @@ -606,6 +606,8 @@ def reset( env_ids: Environment indices. Defaults to None (all environments). env_mask: Environment mask. Defaults to None (all environments). """ + if not self._active: + return full_reset = env_mask is None and (env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None))) if full_reset: # Full reset: zero all 7 buffers diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index a86d990193e8..cdb76b332c1f 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -908,6 +908,28 @@ def test_full_reset_clears_active_flag(device: str, env_ids: slice | None): assert _get_wrench_without_pose_reads(composer)[2] is True +@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("selector", ["all", "env_ids", "env_mask"]) +def test_reset_of_inactive_composer_launches_no_work(device: str, selector: str): + """Resetting an inactive composer, as asset writes do every physics step, is a host-only no-op.""" + num_envs, num_bodies = 4, 2 + composer = WrenchComposer(create_mock_asset(num_envs, num_bodies, device)) + selection = { + "all": {}, + "env_ids": {"env_ids": [1, 3]}, + "env_mask": {"env_mask": wp.array([False, True, False, True], dtype=wp.bool, device=device)}, + }[selector] + + with patch.object(wp, "launch") as launch, patch.object(wp.array, "zero_") as zero: + composer.reset(**selection) + + launch.assert_not_called() + zero.assert_not_called() + assert not composer.active + assert not composer._dirty + np.testing.assert_array_equal(composer.out_force_b.warp.numpy(), np.zeros((num_envs, num_bodies, 3))) + + # ============================================================================ # Deprecated API Backward-Compatibility Tests # ============================================================================ From f0fb3fb5a31a3958d367ad68ea6df856f236a650 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 22:03:17 -0700 Subject: [PATCH 02/18] Remove redundant solver reset from Newton physics step --- .../isaaclab_newton/physics/mjwarp_manager.py | 8 ++++---- .../isaaclab_newton/physics/newton_manager.py | 3 +-- .../test/physics/test_newton_solver_reset.py | 9 +++++++++ 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_manager.py index 1b0bd4a4f0b9..0a4a18ef795b 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_manager.py @@ -116,10 +116,10 @@ def _reset_solver_internals(cls, world_mask: wp.array | None) -> None: With ``use_mujoco_cpu=True`` the solver owns a single global ``MjData`` and its reset path is not mask-aware — it clears the buffers for every - world. Since this hook fires on every step/forward boundary (usually - with an all-``False`` mask), the CPU path is gated on at least one - world actually being flagged so warm-starting is not defeated on every - step. + world. Since this hook can fire with an all-``False`` mask (e.g. on + every forward boundary after a graph-captured write), the CPU path is + gated on at least one world actually being flagged so warm-starting is + not defeated on every step. Args: world_mask: Per-world bool mask of shape ``(world_count + 1,)``. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 0d01ef42dc7a..7cd481d28e75 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -681,8 +681,6 @@ def step(cls) -> None: if sim is None or not sim.is_playing(): return - cls._reset_solver_internals_delegate(cls._world_reset_mask) - # Notify solver of model changes if cls._model_changes: with wp.ScopedDevice(PhysicsManager._device): @@ -694,6 +692,7 @@ def step(cls) -> None: NewtonManager._model_changes = set() # Reset-authored state and persistent solver resources must be ready before capture. + # forward() also resets solver internals for the worlds flagged since the last boundary. cls.forward() cfg = PhysicsManager._cfg device = PhysicsManager._device diff --git a/source/isaaclab_newton/test/physics/test_newton_solver_reset.py b/source/isaaclab_newton/test/physics/test_newton_solver_reset.py index fd4ac657832b..4d29c88918bc 100644 --- a/source/isaaclab_newton/test/physics/test_newton_solver_reset.py +++ b/source/isaaclab_newton/test/physics/test_newton_solver_reset.py @@ -138,3 +138,12 @@ def test_env_reset_clears_selected_mjwarp_solver_internals(device): torch.testing.assert_close(wp.to_torch(state.joint_qd), joint_qd_before) assert torch.count_nonzero(warm_start[0]).item() == 0 torch.testing.assert_close(warm_start[1], torch.full_like(warm_start[1], 29.0)) + + # A clean step has no flagged worlds, so it launches no solver reset. + with ( + patch.object(SimulationManager, "_simulate_full", classmethod(lambda cls: None)), + patch.object(SimulationManager, "_simulate_physics_only", classmethod(lambda cls: None)), + patch.object(SimulationManager, "_reset_solver_internals_delegate") as reset_delegate, + ): + sim.step(render=False) + reset_delegate.assert_not_called() From bc225bf2a9dffa934ebe3877cacd44d71a030416 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 22:03:29 -0700 Subject: [PATCH 03/18] Use elapsed time for Newton body COM acceleration --- .../assets/articulation/articulation_data.py | 4 +++- .../assets/rigid_object/rigid_object_data.py | 4 +++- .../rigid_object_collection_data.py | 4 +++- .../test/assets/test_articulation.py | 16 ++++++++++++++-- .../test/assets/test_rigid_object.py | 6 ++++++ .../test/assets/test_rigid_object_collection.py | 6 ++++++ 6 files changed, 35 insertions(+), 5 deletions(-) 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 3591606f8592..a629e8c22ee2 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -819,6 +819,8 @@ def body_com_acc_w(self) -> ProxyArray: All values are relative to the world. """ if self._body_com_acc_w.timestamp < self._sim_timestamp: + # Finite-difference over the elapsed time, which spans the decimation when Newton owns it. + time_elapsed = self._sim_timestamp - self._body_com_acc_w.timestamp body_ordering = self.body_ordering wp.launch( articulation_kernels.get_body_com_acc_from_body_com_vel_ordered, @@ -829,7 +831,7 @@ def body_com_acc_w(self) -> ProxyArray: self._previous_body_com_vel, body_ordering.user_to_backend if body_ordering is not None else None, body_ordering is not None, - SimulationManager.get_dt(), + time_elapsed, ], outputs=[self._body_com_acc_w.data], ) 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 ff8f88f9dda9..f76e2084b9f3 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 @@ -420,13 +420,15 @@ def body_com_acc_w(self) -> ProxyArray: This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. """ if self._body_com_acc_w.timestamp < self._sim_timestamp: + # Finite-difference over the elapsed time, which spans the decimation when Newton owns it. + time_elapsed = self._sim_timestamp - self._body_com_acc_w.timestamp wp.launch( shared_kernels.derive_body_acceleration_from_body_com_velocities, dim=(self._num_instances, 1), device=self.device, inputs=[ self._sim_bind_body_com_vel_w, - SimulationManager.get_dt(), + time_elapsed, self._previous_body_com_vel, ], outputs=[ diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py index 1b997e2a0758..8b75235bed6a 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection_data.py @@ -354,13 +354,15 @@ def body_com_acc_w(self) -> ProxyArray: This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. """ if self._body_com_acc_w.timestamp < self._sim_timestamp: + # Finite-difference over the elapsed time, which spans the decimation when Newton owns it. + time_elapsed = self._sim_timestamp - self._body_com_acc_w.timestamp wp.launch( shared_kernels.derive_body_acceleration_from_body_com_velocities, dim=(self.num_instances, self.num_bodies), device=self.device, inputs=[ self.body_com_vel_w.warp, - SimulationManager.get_dt(), + time_elapsed, self._previous_body_com_vel, ], outputs=[ diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index 95bbfc051039..c61f8e1ea09f 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -1026,6 +1026,7 @@ def test_newton_ordered_state_caches_invalidate_on_rebind( articulation.update(sim.cfg.dt) data = articulation.data + previous_body_com_vel = data._previous_body_com_vel.numpy().copy() primed_joint_vel_values = ( np.arange(np.prod(data._sim_bind_joint_vel.shape), dtype=np.float32).reshape(data._sim_bind_joint_vel.shape) + 100.0 @@ -1046,13 +1047,24 @@ def test_newton_ordered_state_caches_invalidate_on_rebind( # ``joint_vel`` shadow) observes the primed backend state. No-op under identity # ordering, where the getters alias the sim-bound arrays directly. data._refresh_user_order_state() - data.update(sim.cfg.dt) + # One update spanning two physics steps, as when Newton owns the decimation loop. + update_dt = 2 * sim.cfg.dt + data.update(update_dt) primed_joint_acc = data.joint_acc.warp.numpy().copy() primed_body_com_acc_w = data.body_com_acc_w.warp.numpy().copy() assert data._joint_acc.timestamp == data._sim_timestamp assert data._body_com_acc_w.timestamp == data._sim_timestamp assert np.any(primed_joint_acc != 0.0) - assert np.any(primed_body_com_acc_w != 0.0) + primed_body_user_to_backend = ( + np.asarray(articulation.body_ordering.user_to_backend_indices) + if articulation.body_ordering is not None + else np.arange(articulation.num_bodies) + ) + np.testing.assert_allclose( + primed_body_com_acc_w, + ((primed_body_com_vel_values - previous_body_com_vel) / update_dt)[:, primed_body_user_to_backend], + rtol=1e-5, + ) public_to_binding = { "joint_pos": "_sim_bind_joint_pos", diff --git a/source/isaaclab_newton/test/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index eadd4d00df64..6567aa7223e4 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -447,6 +447,12 @@ def test_gravity_vec_w_tracks_model_gravity(num_cubes, device): gravity[:, :, 2] = -9.81 torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) + # One update may span several physics steps, e.g. once per env step when Newton owns decimation. + sim.step() + sim.step() + cube_object.update(2 * sim.cfg.dt) + torch.testing.assert_close(cube_object.data.body_acc_w.torch, gravity) + # GRAVITY_VEC_W must share storage with Newton's per-env gravity array. model = SimulationManager.get_model() model_gravity_arr = model.gravity[: model.world_count] diff --git a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py index 8f1c4a41d620..d172b7295261 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -410,6 +410,12 @@ def test_gravity_vec_w_tracks_model_gravity(num_envs, num_cubes, device): gravity[..., 2] = -9.81 torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) + # One update may span several physics steps, e.g. once per env step when Newton owns decimation. + sim.step() + sim.step() + object_collection.update(2 * sim.cfg.dt) + torch.testing.assert_close(object_collection.data.body_com_acc_w.torch, gravity) + # GRAVITY_VEC_W must share storage with Newton's per-env gravity array. model = SimulationManager.get_model() model_gravity_arr = model.gravity[: model.world_count] From 5bd37c1e53a9d1b6aa8681c7796d41cc21d99b9a Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 22:03:29 -0700 Subject: [PATCH 04/18] Read back clamped joint defaults only when logged --- .../changelog.d/perf-hot-path-7b-newton.rst | 16 ++++++++ .../assets/articulation/articulation.py | 40 ++++++++++--------- .../test/assets/test_articulation.py | 30 ++++++++++++-- 3 files changed, 65 insertions(+), 21 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/perf-hot-path-7b-newton.rst diff --git a/source/isaaclab_newton/changelog.d/perf-hot-path-7b-newton.rst b/source/isaaclab_newton/changelog.d/perf-hot-path-7b-newton.rst new file mode 100644 index 000000000000..dc7f887bd292 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/perf-hot-path-7b-newton.rst @@ -0,0 +1,16 @@ +Changed +^^^^^^^ + +* Removed the redundant per-step solver-internal reset from :meth:`~isaaclab_newton.physics.NewtonManager.step`. + The reset still runs through :meth:`~isaaclab_newton.physics.NewtonManager.forward` for worlds flagged by + state writes. +* Skipped the device-to-host readback in joint position-limit writes of + :class:`~isaaclab_newton.assets.Articulation` when the clamped-default message would not be logged. + +Fixed +^^^^^ + +* Fixed ``body_com_acc_w`` of :class:`~isaaclab_newton.assets.Articulation`, + :class:`~isaaclab_newton.assets.RigidObject`, and :class:`~isaaclab_newton.assets.RigidObjectCollection` + dividing by the physics time step when an update spans several physics steps, such as when Newton owns + the decimation loop. The finite difference now uses the elapsed time since the previous update. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index e905d1adb0e7..93cf3044f6a5 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -1797,7 +1797,11 @@ def write_joint_position_limit_to_sim_index( """ env_ids = self._resolve_env_ids(env_ids) joint_ids = self._resolve_joint_ids(joint_ids) - clamped_defaults = wp.zeros(1, dtype=wp.int32, device=self.device) + # Count clamped defaults (a host sync) only when the resulting message would be logged. + log_level = logging.WARNING if warn_limit_violation else logging.INFO + report_clamping = logger.isEnabledFor(log_level) + if report_clamping: + self._clamped_default_count.zero_() if isinstance(limits, float): raise ValueError("Joint position limits must be a tensor or array, not a float.") self.assert_shape_and_dtype(limits, (env_ids.shape[0], joint_ids.shape[0]), wp.vec2f, "limits") @@ -1829,20 +1833,17 @@ def write_joint_position_limit_to_sim_index( self.data._sim_bind_joint_pos_limits_upper, self.data._soft_joint_pos_limits, self.data._default_joint_pos, - clamped_defaults, + self._clamped_default_count, ], device=self.device, ) self.data._joint_pos_limits.timestamp = self.data._sim_timestamp - if clamped_defaults.numpy()[0] > 0: - violation_message = ( + if report_clamping and self._clamped_default_count.numpy()[0] > 0: + logger.log( + log_level, "Some default joint positions are outside of the range of the new joint limits. Default joint positions" - " will be clamped to be within the new joint limits." + " will be clamped to be within the new joint limits.", ) - if warn_limit_violation: - logger.warning(violation_message) - else: - logger.info(violation_message) SimulationManager.add_model_change(ModelFlags.JOINT_DOF_PROPERTIES) def write_joint_position_limit_to_sim_mask( @@ -1871,7 +1872,11 @@ def write_joint_position_limit_to_sim_mask( """ env_mask = self._resolve_mask(env_mask, self._ALL_ENV_MASK) joint_mask = self._resolve_mask(joint_mask, self._ALL_JOINT_MASK) - clamped_defaults = wp.zeros(1, dtype=wp.int32, device=self.device) + # Count clamped defaults (a host sync) only when the resulting message would be logged. + log_level = logging.WARNING if warn_limit_violation else logging.INFO + report_clamping = logger.isEnabledFor(log_level) + if report_clamping: + self._clamped_default_count.zero_() if isinstance(limits, float): raise ValueError("Joint position limits must be a tensor or array, not a float.") self.assert_shape_and_dtype_mask(limits, (env_mask, joint_mask), wp.vec2f, "limits") @@ -1903,20 +1908,17 @@ def write_joint_position_limit_to_sim_mask( self.data._sim_bind_joint_pos_limits_upper, self.data._soft_joint_pos_limits, self.data._default_joint_pos, - clamped_defaults, + self._clamped_default_count, ], device=self.device, ) self.data._joint_pos_limits.timestamp = self.data._sim_timestamp - if clamped_defaults.numpy()[0] > 0: - violation_message = ( + if report_clamping and self._clamped_default_count.numpy()[0] > 0: + logger.log( + log_level, "Some default joint positions are outside of the range of the new joint limits. Default joint positions" - " will be clamped to be within the new joint limits." + " will be clamped to be within the new joint limits.", ) - if warn_limit_violation: - logger.warning(violation_message) - else: - logger.info(violation_message) SimulationManager.add_model_change(ModelFlags.JOINT_DOF_PROPERTIES) def write_joint_velocity_limit_to_sim_index( @@ -3451,6 +3453,8 @@ def _create_buffers(self): np.arange(self.num_spatial_tendons, dtype=np.int32), device=self.device ) self._ALL_SPATIAL_TENDON_MASK = wp.ones((self.num_spatial_tendons,), dtype=wp.bool, device=self.device) + # Scratch counter for default joint positions clamped by a joint position-limit write. + self._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=self.device) # Lazily-filled cache of backend-order collision-shape counts (see ``backend_num_shapes_per_body``). self._num_shapes_per_body_backend: list[int] | None = None diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index c61f8e1ea09f..f8532e29321f 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -28,10 +28,12 @@ """Rest everything follows.""" +import logging import sys from copy import copy, deepcopy from pathlib import Path from types import SimpleNamespace +from unittest.mock import patch import newton import numpy as np @@ -2412,14 +2414,15 @@ def test_setting_articulation_root_prim_path(sim, device, articulation_type, roo @pytest.mark.parametrize("device", test_devices()) @pytest.mark.parametrize("gravity_enabled", [False]) @pytest.mark.parametrize("articulation_type", ["anymal"]) -def test_write_joint_state_data_consistency(sim, num_articulations, device, gravity_enabled, articulation_type): +def test_write_joint_state_data_consistency(sim, num_articulations, device, gravity_enabled, articulation_type, caplog): """Joint limit and joint state writes update the joint buffers and refresh the body state without a step. This test verifies that: 1. Joint position limits are written and keep in-limit default joint positions 2. A partial joint state write with unsorted int64 selectors updates only the selected entries 3. A joint state write moves the bodies and refreshes the derived body poses and velocities - 4. Indexed joint limits that exclude a default joint position clamp it into the new limits + 4. Joint limits that exclude a default joint position clamp it into the new limits and report the + clamping at the requested log level, without a device readback when that level is disabled Args: sim: The simulation fixture @@ -2509,7 +2512,10 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) limits[..., 0] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * -0.1 limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * 0.1 - articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + articulation_logger = Articulation.__module__ + with caplog.at_level(logging.WARNING, logger=articulation_logger): + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) + assert any(record.levelno == logging.WARNING for record in caplog.records) # Check new limits are in place and the defaults are clamped into them torch.testing.assert_close(articulation.data.joint_pos_limits.torch[env_ids][:, joint_ids], limits) @@ -2519,6 +2525,24 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav ) assert torch.all(within_bounds) + # Info-level clamping reports are logged when enabled and skip the device readback otherwise. + articulation.data.default_joint_pos.torch.fill_(1.0) + full_limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + full_limits[..., 1] = 0.5 + caplog.clear() + with caplog.at_level(logging.INFO, logger=articulation_logger): + articulation.write_joint_position_limit_to_sim_mask(limits=full_limits, warn_limit_violation=False) + assert [record.levelno for record in caplog.records] == [logging.INFO] + articulation.data.default_joint_pos.torch.fill_(1.0) + caplog.clear() + with ( + caplog.at_level(logging.WARNING, logger=articulation_logger), + patch.object(wp.array, "numpy", side_effect=AssertionError("unexpected device readback")), + ): + articulation.write_joint_position_limit_to_sim_mask(limits=full_limits, warn_limit_violation=False) + assert not caplog.records + torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(full_limits[..., 1], 0.5)) + @pytest.mark.parametrize("selector_kind", ["index", "mask"]) @pytest.mark.parametrize("num_articulations", [2]) From e890cdfdb27f637c1ee6d5bf6b4e0a8ec557d2c8 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 22:05:38 -0700 Subject: [PATCH 05/18] Remove host syncs and redundant kernels from manager hot paths - Log command, termination, curriculum and success-rate episode metrics as device scalars instead of host floats. - Select heading-controlled envs in UniformVelocityCommand with torch.where instead of nonzero(). - Cache event range tensors by value, and make reset_root_state_uniform honor the ranges passed at call time. - Replace the DiffIK zero-quaternion guard with a sync-free torch.where. - Batch RewardManager per-term updates and resets over a term matrix. - Skip defensive copies around built-in noise and delay buffers. - Shift small CircularBuffer frames in two kernels regardless of length. - Check event modes via dict membership in the env step. --- .../changelog.d/perf-hot-path-7b-managers.rst | 24 ++++ .../isaaclab/envs/manager_based_env.py | 6 +- .../isaaclab/envs/manager_based_rl_env.py | 6 +- .../envs/mdp/actions/task_space_actions.py | 9 +- .../envs/mdp/commands/pose_2d_command.py | 6 +- .../envs/mdp/commands/pose_command.py | 6 +- .../envs/mdp/commands/velocity_command.py | 13 +-- source/isaaclab/isaaclab/envs/mdp/events.py | 65 ++++++----- source/isaaclab/isaaclab/envs/mdp/rewards.py | 2 +- .../isaaclab/managers/command_manager.py | 4 +- .../isaaclab/managers/curriculum_manager.py | 11 +- .../isaaclab/managers/observation_manager.py | 23 +++- .../isaaclab/managers/reward_manager.py | 47 ++++---- .../isaaclab/managers/termination_manager.py | 6 +- .../isaaclab/utils/buffers/circular_buffer.py | 18 +-- .../envs/test_diffik_jacobian_aliasing.py | 32 ++++++ .../test/envs/test_root_state_events.py | 103 ++++++++++++++++++ .../test/envs/test_velocity_command_term.py | 103 ++++++++++++++++++ .../test/managers/test_curriculum_manager.py | 62 +++++++++++ .../test/managers/test_observation_manager.py | 31 ++++++ .../test/managers/test_reward_manager.py | 54 +++++++++ .../test/managers/test_termination_manager.py | 36 ++++++ .../test/utils/test_circular_buffer.py | 25 +++-- 23 files changed, 575 insertions(+), 117 deletions(-) create mode 100644 source/isaaclab/changelog.d/perf-hot-path-7b-managers.rst create mode 100644 source/isaaclab/test/envs/test_root_state_events.py create mode 100644 source/isaaclab/test/envs/test_velocity_command_term.py create mode 100644 source/isaaclab/test/managers/test_curriculum_manager.py diff --git a/source/isaaclab/changelog.d/perf-hot-path-7b-managers.rst b/source/isaaclab/changelog.d/perf-hot-path-7b-managers.rst new file mode 100644 index 000000000000..973b1dd38895 --- /dev/null +++ b/source/isaaclab/changelog.d/perf-hot-path-7b-managers.rst @@ -0,0 +1,24 @@ +Changed +^^^^^^^ + +* Changed episode logging in :class:`~isaaclab.managers.CommandTerm`, + :class:`~isaaclab.managers.TerminationManager`, :class:`~isaaclab.managers.CurriculumManager`, + :class:`~isaaclab.envs.mdp.UniformPoseCommand`, :class:`~isaaclab.envs.mdp.UniformPose2dCommand` and + :class:`~isaaclab.envs.mdp.survival_success_rate` to report 0-d device tensors instead of Python floats, + matching :class:`~isaaclab.managers.RewardManager`, so resets no longer synchronize the stream. Code that + reads these ``extras["log"]`` entries as floats should call ``float()`` or ``.item()`` on them. +* Removed per-step host synchronizations from the heading control of + :class:`~isaaclab.envs.mdp.UniformVelocityCommand` and from + :class:`~isaaclab.envs.mdp.actions.task_space_actions.DifferentialInverseKinematicsAction`. +* Cached the device range tensors of the root-state, nodal-state, push and center-of-mass event terms by + value instead of uploading them on every call. +* Batched the per-term updates of :class:`~isaaclab.managers.RewardManager`, and skipped the defensive copy + before the built-in noise functions and after delay buffers in :class:`~isaaclab.managers.ObservationManager`. +* Shifted :class:`~isaaclab.utils.buffers.CircularBuffer` histories of small frames in two kernels + regardless of the history length. + +Fixed +^^^^^ + +* Fixed :class:`~isaaclab.envs.mdp.reset_root_state_uniform` ignoring the ``pose_range`` and ``velocity_range`` + passed at call time, including curriculum updates, in favor of the ranges present at construction. diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 778023f6a58b..b9c88d97f186 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -546,7 +546,7 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: self.recorder_manager.record_pre_step() # check if we need to do rendering within the physics loop - # note: uses cached property to avoid settings lookup every step + # note: evaluated once per step rather than once per physics substep is_rendering = self.sim.is_rendering # physics-owned decimation covers all substeps in one call @@ -562,7 +562,7 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: self.scene.update(dt=self.physics_dt * steps_per_call) # post-step: step interval event - if "interval" in self.event_manager.available_modes: + if "interval" in self.event_manager.active_terms: self.event_manager.apply(mode="interval", dt=self.step_dt) # advance video recorders (after render, before obs) @@ -631,7 +631,7 @@ def _reset_idx(self, env_ids: torch.Tensor | slice): self.scene.reset(env_ids) # apply events such as randomization for environments that need a reset - if "reset" in self.event_manager.available_modes: + if "reset" in self.event_manager.active_terms: env_step_count = self._sim_step_counter // self.cfg.decimation self.event_manager.apply(mode="reset", env_ids=env_ids, global_env_step_count=env_step_count) diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index 750c2af87b96..dfc59587562e 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -207,7 +207,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.recorder_manager.record_pre_step() # check if we need to do rendering within the physics loop - # note: uses cached property to avoid settings lookup every step + # note: evaluated once per step rather than once per physics substep is_rendering = self.sim.is_rendering # physics-owned decimation covers all substeps in one call @@ -269,7 +269,7 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: self.recorder_manager.record_post_reset(manual_reset_ids) self.command_manager.compute(dt=self.step_dt) - if "interval" in self.event_manager.available_modes: + if "interval" in self.event_manager.active_terms: self.event_manager.apply(mode="interval", dt=self.step_dt) # -- advance video recorders (after render and resets, before final obs) for recorder in self.video_recorders: @@ -387,7 +387,7 @@ def _reset_idx(self, env_ids: torch.Tensor | slice): # reset the internal buffers of the scene elements self.scene.reset(env_ids) # apply events such as randomizations for environments that need a reset - if "reset" in self.event_manager.available_modes: + if "reset" in self.event_manager.active_terms: env_step_count = self._sim_step_counter // self.cfg.decimation self.event_manager.apply(mode="reset", env_ids=env_ids, global_env_step_count=env_step_count) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index 3b077356cea5..cecfe3ef829a 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -210,11 +210,10 @@ def apply_actions(self): self._ik_controller.set_joint_pos_limits(limits[:, 0].clone(), limits[:, 1].clone()) self._limits_injected = True # compute the delta in joint-space - if ee_quat_curr.norm() != 0: - jacobian = self._compute_frame_jacobian() - joint_pos_des = self._ik_controller.compute(ee_pos_curr, ee_quat_curr, jacobian, joint_pos) - else: - joint_pos_des = joint_pos.clone() + jacobian = self._compute_frame_jacobian() + joint_pos_des = self._ik_controller.compute(ee_pos_curr, ee_quat_curr, jacobian, joint_pos) + # hold the joints while the frame pose is uninitialized (all-zero quaternions), without a host sync + joint_pos_des = torch.where(ee_quat_curr.norm() != 0, joint_pos_des, joint_pos) # set the joint position command self._asset.set_joint_position_target_index(target=joint_pos_des, joint_ids=self._joint_ids) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py index f1b0f747e341..ae9cd040bcdb 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_2d_command.py @@ -98,16 +98,14 @@ def _update_metrics(self): if self._track_success: self._succeeded |= self.metrics["error_pos"] < self.cfg.position_success_threshold - def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: extras = super().reset(env_ids) if self._track_success: if env_ids is None: env_ids = slice(None) # Write the unified ``Metrics/success_rate`` directly to env extras so it shares # a TensorBoard card with the same metric from other tasks. - self._env.extras.setdefault("log", {})["Metrics/success_rate"] = ( - self._succeeded[env_ids].float().mean().item() - ) + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._succeeded[env_ids].float().mean() index_fill_(self._succeeded, env_ids, False) return extras diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py index 74b329f2ed2a..89ef3e5a5258 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/pose_command.py @@ -151,16 +151,14 @@ def _compute_success(self, position_error: torch.Tensor, orientation_error: torc success[:] = False return success - def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: extras = super().reset(env_ids) if self._track_success: if env_ids is None: env_ids = slice(None) # Write the unified ``Metrics/success_rate`` directly to env extras so it shares # a TensorBoard card with the same metric from other tasks. - self._env.extras.setdefault("log", {})["Metrics/success_rate"] = ( - self._succeeded[env_ids].float().mean().item() - ) + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._succeeded[env_ids].float().mean() index_fill_(self._succeeded, env_ids, False) return extras diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index becf01280f75..c84aafd6ea4f 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -130,7 +130,7 @@ def _update_metrics(self): self._error_yaw_sum += error_yaw self._step_count += 1.0 - def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: # Finalize the just-ended episode's metrics into ``self.metrics`` BEFORE the base # class reads them. ``success_rate`` is per-env binary: the *episode-mean* error # is below both thresholds. Then super().reset() logs and zeros ``self.metrics``; @@ -181,17 +181,14 @@ def _update_command(self): """ # Compute angular velocity from heading direction if self.cfg.heading_command: - # resolve indices of heading envs - env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten() - # compute angular velocity - heading_error = math_utils.wrap_to_pi( - self.heading_target[env_ids] - self.robot.data.heading_w.torch[env_ids] - ) - self.vel_command_b[env_ids, 2] = torch.clip( + # compute angular velocity for all envs, then select heading envs without a host sync + heading_error = math_utils.wrap_to_pi(self.heading_target - self.robot.data.heading_w.torch) + heading_ang_vel = torch.clip( self.cfg.heading_control_stiffness * heading_error, min=self.cfg.ranges.ang_vel_z[0], max=self.cfg.ranges.ang_vel_z[1], ) + self.vel_command_b[:, 2] = torch.where(self.is_heading_env, heading_ang_vel, self.vel_command_b[:, 2]) # Enforce standing (i.e., zero velocity command) for standing envs index_fill_(self.vel_command_b, self.is_standing_env, 0.0) diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index eeac316a79ca..59fc271abaf7 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -14,7 +14,9 @@ from __future__ import annotations +import functools import logging +from collections.abc import Sequence from types import ModuleType from typing import TYPE_CHECKING, Literal @@ -38,6 +40,29 @@ logger = logging.getLogger(__name__) +@functools.lru_cache(maxsize=128) +def _cached_range_tensor(values: tuple[tuple[float, float], ...], device: str) -> torch.Tensor: + return torch.tensor(values, dtype=torch.float32, device=device) + + +def _range_tensor(ranges: dict[str, tuple[float, float]], keys: Sequence[str], device: str) -> torch.Tensor: + """Return the ``(min, max)`` bounds of ``keys`` as a device tensor of shape ``(len(keys), 2)``. + + The tensor is cached by value, so repeated calls avoid a host-to-device copy while range edits (for + example by a curriculum) still take effect. Callers must not modify the returned tensor. + + Args: + ranges: The ``(min, max)`` range for each key. Missing keys use ``(0.0, 0.0)``. + keys: The keys to read, in output row order. + device: The device of the returned tensor. + + Returns: + The range bounds. Shape is ``(len(keys), 2)``. + """ + values = tuple(tuple(float(v) for v in ranges.get(key, (0.0, 0.0))) for key in keys) + return _cached_range_tensor(values, str(device)) + + def randomize_rigid_body_scale( env: ManagerBasedEnv, env_ids: torch.Tensor | slice | None, @@ -522,8 +547,7 @@ def __call__( env_rows = env_ids[:, None] if not isinstance(env_ids, slice) and not isinstance(body_ids, slice) else env_ids # sample random CoM values - range_list = [com_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - ranges = torch.tensor(range_list, device=self.asset.device) + ranges = _range_tensor(com_range, ("x", "y", "z"), self.asset.device) num_envs = len(range(env.num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) rand_samples = math_utils.sample_uniform( ranges[:, 0], ranges[:, 1], (num_envs, 3), device=self.asset.device @@ -1257,8 +1281,7 @@ def push_by_setting_velocity( # velocities vel_w = asset.data.root_vel_w.torch[env_ids] # sample random velocities - range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - ranges = torch.tensor(range_list, device=asset.device) + ranges = _range_tensor(velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) vel_w += math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], vel_w.shape, device=asset.device) # set the velocities into the physics simulation asset.write_root_velocity_to_sim_index(root_velocity=vel_w, env_ids=env_ids) @@ -1277,20 +1300,8 @@ class reset_root_state_uniform(ManagerTermBase): The term takes a dictionary of pose and velocity ranges for each axis and rotation. The keys of the dictionary are ``x``, ``y``, ``z``, ``roll``, ``pitch``, and ``yaw``. The values are tuples of the form ``(min, max)``. If the dictionary does not contain a key, the position or velocity is set to zero for that axis. - - The range dictionaries are materialized as device tensors once at construction. """ - def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): - super().__init__(cfg, env) - keys = ("x", "y", "z", "roll", "pitch", "yaw") - pose_range = cfg.params.get("pose_range", {}) - velocity_range = cfg.params.get("velocity_range", {}) - self._pose_ranges = torch.tensor([tuple(pose_range.get(key, (0.0, 0.0))) for key in keys], device=env.device) - self._velocity_ranges = torch.tensor( - [tuple(velocity_range.get(key, (0.0, 0.0))) for key in keys], device=env.device - ) - def __call__( self, env: ManagerBasedEnv, @@ -1304,7 +1315,7 @@ def __call__( default_root_pose = asset.data.default_root_pose.torch[env_ids] default_root_vel = asset.data.default_root_vel.torch[env_ids] - ranges = self._pose_ranges + ranges = _range_tensor(pose_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) rand_samples = math_utils.sample_uniform( ranges[:, 0], ranges[:, 1], (default_root_pose.shape[0], 6), device=asset.device ) @@ -1313,7 +1324,7 @@ def __call__( orientations_delta = math_utils.quat_from_euler_xyz(rand_samples[:, 3], rand_samples[:, 4], rand_samples[:, 5]) orientations = math_utils.quat_mul(default_root_pose[:, 3:7], orientations_delta) # velocities - ranges = self._velocity_ranges + ranges = _range_tensor(velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) rand_samples = math_utils.sample_uniform( ranges[:, 0], ranges[:, 1], (default_root_pose.shape[0], 6), device=asset.device ) @@ -1356,8 +1367,7 @@ def reset_root_state_with_random_orientation( default_root_vel = asset.data.default_root_vel.torch[env_ids].clone() # poses - range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - ranges = torch.tensor(range_list, device=asset.device) + ranges = _range_tensor(pose_range, ("x", "y", "z"), asset.device) rand_samples = math_utils.sample_uniform( ranges[:, 0], ranges[:, 1], (default_root_pose.shape[0], 3), device=asset.device ) @@ -1366,8 +1376,7 @@ def reset_root_state_with_random_orientation( orientations = math_utils.random_orientation(default_root_pose.shape[0], device=asset.device) # velocities - range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - ranges = torch.tensor(range_list, device=asset.device) + ranges = _range_tensor(velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) rand_samples = math_utils.sample_uniform( ranges[:, 0], ranges[:, 1], (default_root_pose.shape[0], 6), device=asset.device ) @@ -1426,16 +1435,14 @@ def reset_root_state_from_terrain( positions = valid_positions[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids], ids] positions += asset.data.default_root_pose.torch[env_ids, :3] # sample random orientations - range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["roll", "pitch", "yaw"]] - ranges = torch.tensor(range_list, device=asset.device) + ranges = _range_tensor(pose_range, ("roll", "pitch", "yaw"), asset.device) rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (num_envs, 3), device=asset.device) # convert to quaternions orientations = math_utils.quat_from_euler_xyz(rand_samples[:, 0], rand_samples[:, 1], rand_samples[:, 2]) # sample random velocities - range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] - ranges = torch.tensor(range_list, device=asset.device) + ranges = _range_tensor(velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (num_envs, 6), device=asset.device) velocities = asset.data.default_root_vel.torch[env_ids] + rand_samples @@ -1699,8 +1706,7 @@ def reset_nodal_state_uniform( nodal_state = asset.data.default_nodal_state_w.torch[env_ids].clone() # position - range_list = [position_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - ranges = torch.tensor(range_list, device=asset.device) + ranges = _range_tensor(position_range, ("x", "y", "z"), asset.device) rand_samples = math_utils.sample_uniform( ranges[:, 0], ranges[:, 1], (nodal_state.shape[0], 1, 3), device=asset.device ) @@ -1708,8 +1714,7 @@ def reset_nodal_state_uniform( nodal_state[..., :3] += rand_samples # velocities - range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - ranges = torch.tensor(range_list, device=asset.device) + ranges = _range_tensor(velocity_range, ("x", "y", "z"), asset.device) rand_samples = math_utils.sample_uniform( ranges[:, 0], ranges[:, 1], (nodal_state.shape[0], 1, 3), device=asset.device ) diff --git a/source/isaaclab/isaaclab/envs/mdp/rewards.py b/source/isaaclab/isaaclab/envs/mdp/rewards.py index 0c96e700f76d..97c12b2b5e5f 100644 --- a/source/isaaclab/isaaclab/envs/mdp/rewards.py +++ b/source/isaaclab/isaaclab/envs/mdp/rewards.py @@ -88,7 +88,7 @@ class survival_success_rate(ManagerTermBase): def reset(self, env_ids: torch.Tensor) -> None: survived = self._env.termination_manager.time_outs[env_ids] - self._env.extras.setdefault("log", {})["Metrics/success_rate"] = survived.float().mean().item() + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = survived.float().mean() def __call__(self, env: ManagerBasedRLEnv) -> torch.Tensor: return torch.zeros(env.num_envs, device=env.device) diff --git a/source/isaaclab/isaaclab/managers/command_manager.py b/source/isaaclab/isaaclab/managers/command_manager.py index a2902beb206c..be467d0f22dc 100644 --- a/source/isaaclab/isaaclab/managers/command_manager.py +++ b/source/isaaclab/isaaclab/managers/command_manager.py @@ -112,7 +112,7 @@ def set_debug_vis(self, debug_vis: bool) -> bool: # return success return True - def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, float]: + def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, torch.Tensor]: """Reset the command generator and log metrics. This function resets the command counter and resamples the command. It should be called @@ -132,7 +132,7 @@ def reset(self, env_ids: Sequence[int] | slice | None = None) -> dict[str, float extras = {} for metric_name, metric_value in self.metrics.items(): # compute the mean metric value - extras[metric_name] = torch.mean(metric_value[env_ids]).item() + extras[metric_name] = torch.mean(metric_value[env_ids]) # reset the metric value index_fill_(metric_value, env_ids, 0.0) diff --git a/source/isaaclab/isaaclab/managers/curriculum_manager.py b/source/isaaclab/isaaclab/managers/curriculum_manager.py index 7ba550273384..bef9428024b4 100644 --- a/source/isaaclab/isaaclab/managers/curriculum_manager.py +++ b/source/isaaclab/isaaclab/managers/curriculum_manager.py @@ -90,7 +90,7 @@ def active_terms(self) -> list[str]: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float | torch.Tensor]: """Returns the current state of individual curriculum terms. Note: @@ -102,6 +102,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: Dictionary of curriculum terms and their states. """ extras = {} + # tensors are snapshotted on device: a host read would synchronize, and terms may update state in place for term_name, term_state in self._curriculum_state.items(): if term_state is not None: # deal with dict @@ -109,12 +110,12 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float]: # each key is a separate state to log for key, value in term_state.items(): if isinstance(value, torch.Tensor): - value = value.item() + value = value.detach().clone() extras[f"Curriculum/{term_name}/{key}"] = value else: # log directly if not a dict if isinstance(term_state, torch.Tensor): - term_state = term_state.item() + term_state = term_state.detach().clone() extras[f"Curriculum/{term_name}"] = term_state # reset all the curriculum terms for term_cfg in self._class_term_cfgs: @@ -161,12 +162,12 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc # each key is a separate state to log for key, value in term_state.items(): if isinstance(value, torch.Tensor): - value = value.item() + value = value.detach().clone() data.append(value) else: # log directly if not a dict if isinstance(term_state, torch.Tensor): - term_state = term_state.item() + term_state = term_state.detach().clone() data.append(term_state) terms.append((term_name, data)) diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index 1f7e1fbac070..692d91539d3e 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -16,7 +16,7 @@ from prettytable import PrettyTable from ..envs.utils.io_descriptors import _warn_io_descriptors_deprecated -from ..utils import instantiate, modifiers, noise, to_dict +from ..utils import callable_to_string, instantiate, modifiers, noise, to_dict from ..utils.buffers import CircularBuffer, DelayBuffer from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import ObservationGroupCfg, ObservationTermCfg @@ -24,6 +24,14 @@ if TYPE_CHECKING: from ..envs import ManagerBasedEnv +# Built-in noise functions that never modify their input, so the manager can skip the defensive copy. +# Configurations hold them as callables or as unresolved ``module:function`` strings. +_OUT_OF_PLACE_NOISE_FUNCS = tuple( + ref + for func in (noise.constant_noise, noise.gaussian_noise, noise.uniform_noise) + for ref in (func, callable_to_string(func)) +) + class ObservationManager(ManagerBase): """Manager for computing observation signals for a given world. @@ -443,8 +451,14 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch. else: obs = modifier.func(obs, **modifier.params) if isinstance(term_cfg.noise, noise.NoiseCfg): - obs = term_cfg.noise.func(obs.clone(), term_cfg.noise) - owned = False + if term_cfg.noise.func in _OUT_OF_PLACE_NOISE_FUNCS: + noisy_obs = term_cfg.noise.func(obs, term_cfg.noise) + # a new result is owned; a returned input keeps its ownership + owned = owned or noisy_obs is not obs + obs = noisy_obs + else: + obs = term_cfg.noise.func(obs.clone(), term_cfg.noise) + owned = False elif isinstance(term_cfg.noise, noise.NoiseModelCfg) and term_cfg.noise.func is not None: obs = term_cfg.noise.func(obs.clone()) owned = False @@ -462,7 +476,8 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch. obs = self._group_obs_term_delay_buffer[group_name][term_name].compute( obs, update_history=update_history ) - owned = False + # the delay buffer returns storage independent of its history + owned = True # Update the history buffer if observation term has history enabled if term_cfg.history_length > 0: circular_buffer = self._group_obs_term_history_buffer[group_name][term_name] diff --git a/source/isaaclab/isaaclab/managers/reward_manager.py b/source/isaaclab/isaaclab/managers/reward_manager.py index 8e4567d264c2..40f39db90245 100644 --- a/source/isaaclab/isaaclab/managers/reward_manager.py +++ b/source/isaaclab/isaaclab/managers/reward_manager.py @@ -56,10 +56,11 @@ def __init__(self, cfg: object, env: ManagerBasedRLEnv): # call the base class constructor (this will parse the terms config) super().__init__(cfg, env) - # prepare extra info to store individual reward term information - self._episode_sums = dict() - for term_name in self._term_names: - self._episode_sums[term_name] = torch.zeros(self.num_envs, dtype=torch.float, device=self.device) + # episodic sum of each term; the dictionary holds column views of one matrix for batched updates + self._episode_sum_buf = torch.zeros( + (self.num_envs, len(self._term_names)), dtype=torch.float, device=self.device + ) + self._episode_sums = {name: self._episode_sum_buf[:, i] for i, name in enumerate(self._term_names)} # create buffer for managing reward per environment self._reward_buf = torch.zeros(self.num_envs, dtype=torch.float, device=self.device) @@ -111,14 +112,11 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor] """ if env_ids is None: env_ids = slice(None) - extras = {} - for key in self._episode_sums.keys(): - # store information - # r_1 + r_2 + ... + r_n - episodic_sum_avg = torch.mean(self._episode_sums[key][env_ids]) - extras["Episode_Reward/" + key] = episodic_sum_avg / self._env.max_episode_length_s - # reset episodic sum - index_fill_(self._episode_sums[key], env_ids, 0.0) + # average of r_1 + r_2 + ... + r_n over the selected environments, for all terms at once + episodic_sum_avg = torch.mean(self._episode_sum_buf[env_ids], dim=0) / self._env.max_episode_length_s + extras = {"Episode_Reward/" + key: episodic_sum_avg[i] for i, key in enumerate(self._term_names)} + # reset episodic sums + index_fill_(self._episode_sum_buf, env_ids, 0.0) # reset all the reward terms for term_cfg in self._class_term_cfgs: term_cfg.func.reset(env_ids=env_ids) @@ -136,23 +134,22 @@ def compute(self, dt: float) -> torch.Tensor: Returns: The net reward signal of shape (num_envs,). """ - # reset computation - self._reward_buf[:] = 0.0 - # iterate over all the reward terms - for term_idx, (name, term_cfg) in enumerate(zip(self._term_names, self._term_cfgs)): + # compute the weighted reward of each term + for term_idx, term_cfg in enumerate(self._term_cfgs): # skip if weight is zero (kind of a micro-optimization) if term_cfg.weight == 0.0: self._step_reward[:, term_idx] = 0.0 continue - # compute term's value - value = term_cfg.func(self._env, **term_cfg.params) * term_cfg.weight * dt - # update total reward - self._reward_buf += value - # update episodic sum - self._episode_sums[name] += value - - # Update current reward for this step. - self._step_reward[:, term_idx] = value / dt + value = term_cfg.func(self._env, **term_cfg.params) + step_reward = self._step_reward[:, term_idx] + if isinstance(value, torch.Tensor) and value.shape == step_reward.shape: + torch.mul(value, term_cfg.weight, out=step_reward) + else: + # broadcast scalar values + step_reward[:] = value * term_cfg.weight + # update total reward and episodic sums with the time-step scaling + torch.sum(self._step_reward, dim=1, out=self._reward_buf).mul_(dt) + self._episode_sum_buf.add_(self._step_reward, alpha=dt) return self._reward_buf diff --git a/source/isaaclab/isaaclab/managers/termination_manager.py b/source/isaaclab/isaaclab/managers/termination_manager.py index 75f297eb314c..cb86c5792f7a 100644 --- a/source/isaaclab/isaaclab/managers/termination_manager.py +++ b/source/isaaclab/isaaclab/managers/termination_manager.py @@ -140,11 +140,11 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor] env_ids = slice(None) # add to episode dict extras = {} - # move to host once; per-element .item() would sync per term - last_episode_done_stats = self._last_episode_dones.float().mean(dim=0).cpu() + # log device scalars; a host read would synchronize the stream + last_episode_done_stats = self._last_episode_dones.float().mean(dim=0) for i, key in enumerate(self._term_names): # store information - extras["Episode_Termination/" + key] = last_episode_done_stats[i].item() + extras["Episode_Termination/" + key] = last_episode_done_stats[i] # reset all the termination terms for term_cfg in self._class_term_cfgs: term_cfg.func.reset(env_ids=env_ids) diff --git a/source/isaaclab/isaaclab/utils/buffers/circular_buffer.py b/source/isaaclab/isaaclab/utils/buffers/circular_buffer.py index 3417ac3e38f8..985a9d5ed2ab 100644 --- a/source/isaaclab/isaaclab/utils/buffers/circular_buffer.py +++ b/source/isaaclab/isaaclab/utils/buffers/circular_buffer.py @@ -9,6 +9,9 @@ from ..array import index_fill_ +# Frame size up to which one staged shift is faster than K - 1 slot copies (measured crossover is ~4 MiB). +_STAGED_SHIFT_MAX_FRAME_BYTES = 2 * 1024 * 1024 + class CircularBuffer: """Circular buffer for storing a history of batched tensor data. @@ -173,15 +176,14 @@ def append(self, data: torch.Tensor): if self._buffer is None: self._allocate_buffer(data) - # Shift slots so the newest write lands at the last K slot. Iterating front-to-back - # keeps adjacent-slot copies non-overlapping. Cheap at the typical frame-stack K=2-4. - if self._stack_dim_internal is None: - for i in range(self._max_len_int - 1): - self._buffer[i].copy_(self._buffer[i + 1]) - self._buffer[-1] = data + # Drop the oldest slot and write the newest at the last K slot. + k_pos = 0 if self._stack_dim_internal is None else self._stack_dim_internal + k = self._max_len_int + if k > 2 and data.numel() * data.element_size() <= _STAGED_SHIFT_MAX_FRAME_BYTES: + # small frames are launch bound: stage the overlapping shift in two kernels regardless of K + self._buffer.copy_(torch.cat((self._buffer.narrow(k_pos, 1, k - 1), data.unsqueeze(k_pos)), dim=k_pos)) else: - k_pos = self._stack_dim_internal - k = self._max_len_int + # large frames are bandwidth bound: front-to-back slot copies move the least data for i in range(k - 1): self._buffer.narrow(k_pos, i, 1).copy_(self._buffer.narrow(k_pos, i + 1, 1)) self._buffer.narrow(k_pos, k - 1, 1).copy_(data.unsqueeze(k_pos)) diff --git a/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py b/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py index f06ef4d368dd..17d5b3d88e9e 100644 --- a/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py +++ b/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py @@ -23,6 +23,7 @@ DifferentialInverseKinematicsAction, OperationalSpaceControllerAction, ) +from isaaclab.test.utils import test_devices from isaaclab.utils import math as math_utils pytestmark = pytest.mark.unit @@ -123,5 +124,36 @@ def test_body_offset_jacobian_uses_offset_in_root_frame(compute, returns_jacobia torch.testing.assert_close(result, expected) +@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize(("ee_quat", "ik_offset"), [((0.0, 0.0, 0.0, 1.0), 1.0), ((0.0, 0.0, 0.0, 0.0), 0.0)]) +def test_apply_actions_holds_joints_for_uninitialized_frame(device, ee_quat, ik_offset): + """IK targets are applied for a valid frame pose and joints hold for all-zero quaternions, without a sync.""" + joint_pos = torch.tensor([[0.1, 0.2], [0.3, 0.4]], device=device) + frame_pose = (torch.zeros(2, 3, device=device), torch.tensor([ee_quat] * 2, device=device)) + written = {} + stub = SimpleNamespace( + cfg=SimpleNamespace(controller=SimpleNamespace(joint_limit_avoidance_gain=0.0)), + _asset=SimpleNamespace( + data=SimpleNamespace(joint_pos=SimpleNamespace(torch=joint_pos)), + set_joint_position_target_index=lambda target, joint_ids: written.update(target=target), + ), + _joint_ids=slice(None), + _limits_injected=True, + _compute_frame_pose=lambda: frame_pose, + _compute_frame_jacobian=lambda: torch.zeros(2, 6, 2, device=device), + _ik_controller=SimpleNamespace(compute=lambda ee_pos, ee_quat, jacobian, joint_pos: joint_pos + 1.0), + ) + previous = torch.cuda.get_sync_debug_mode() if device.startswith("cuda") else None + if previous is not None: + torch.cuda.synchronize(device) + torch.cuda.set_sync_debug_mode("error") + try: + DifferentialInverseKinematicsAction.apply_actions(stub) + finally: + if previous is not None: + torch.cuda.set_sync_debug_mode(previous) + torch.testing.assert_close(written["target"], joint_pos + ik_offset) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/source/isaaclab/test/envs/test_root_state_events.py b/source/isaaclab/test/envs/test_root_state_events.py new file mode 100644 index 000000000000..a2e77eaefd3e --- /dev/null +++ b/source/isaaclab/test/envs/test_root_state_events.py @@ -0,0 +1,103 @@ +# 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 + +"""Unit tests for root-state event terms that sample from range dictionaries.""" + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from isaaclab.envs.mdp import push_by_setting_velocity, reset_root_state_uniform +from isaaclab.managers import EventTermCfg, SceneEntityCfg +from isaaclab.test.utils import test_devices + +pytestmark = pytest.mark.unit + + +@contextmanager +def _raise_on_cuda_sync(device: str): + """Raise if the enclosed code synchronizes a CUDA device.""" + if not device.startswith("cuda"): + yield + return + previous = torch.cuda.get_sync_debug_mode() + torch.cuda.synchronize(device) + torch.cuda.set_sync_debug_mode("error") + try: + yield + finally: + torch.cuda.set_sync_debug_mode(previous) + + +class _Scene(dict): + """Scene double exposing assets by name and the environment origins.""" + + def __init__(self, asset, env_origins: torch.Tensor): + super().__init__(robot=asset) + self.env_origins = env_origins + + +def _make_env(device: str, num_envs: int = 3) -> tuple[SimpleNamespace, SimpleNamespace]: + """Create an environment double whose robot records the written root state.""" + default_root_pose = torch.zeros(num_envs, 7, device=device) + default_root_pose[:, 2] = 0.4 + default_root_pose[:, 6] = 1.0 + asset = SimpleNamespace( + device=device, + data=SimpleNamespace( + default_root_pose=SimpleNamespace(torch=default_root_pose), + default_root_vel=SimpleNamespace(torch=torch.zeros(num_envs, 6, device=device)), + root_vel_w=SimpleNamespace(torch=torch.zeros(num_envs, 6, device=device)), + ), + written={}, + ) + asset.write_root_pose_to_sim_index = lambda root_pose, env_ids: asset.written.update(pose=root_pose) + asset.write_root_velocity_to_sim_index = lambda root_velocity, env_ids: asset.written.update(vel=root_velocity) + env_origins = torch.arange(num_envs * 3, dtype=torch.float, device=device).reshape(num_envs, 3) + env = SimpleNamespace(num_envs=num_envs, device=device, scene=_Scene(asset, env_origins)) + return env, asset + + +@pytest.mark.parametrize("device", test_devices()) +def test_reset_root_state_uniform_uses_call_ranges(device): + """The term samples from the ranges passed at call time, not the ones present at construction.""" + env, asset = _make_env(device) + asset_cfg = SceneEntityCfg("robot") + cfg = EventTermCfg( + func=reset_root_state_uniform, + mode="reset", + params={"pose_range": {}, "velocity_range": {}, "asset_cfg": asset_cfg}, + ) + term = reset_root_state_uniform(cfg, env) + env_ids = torch.arange(env.num_envs, device=device) + + term(env, env_ids, pose_range={"x": (1.0, 1.0), "z": (0.5, 0.5)}, velocity_range={"yaw": (2.0, 2.0)}) + + expected_pos = env.scene.env_origins + torch.tensor([1.0, 0.0, 0.9], device=device) + torch.testing.assert_close(asset.written["pose"][:, :3], expected_pos) + torch.testing.assert_close(asset.written["pose"][:, 3:], asset.data.default_root_pose.torch[:, 3:]) + expected_vel = torch.zeros(env.num_envs, 6, device=device) + expected_vel[:, 5] = 2.0 + torch.testing.assert_close(asset.written["vel"], expected_vel) + + +@pytest.mark.parametrize("device", test_devices()) +def test_push_by_setting_velocity_follows_range_edits_without_sync(device): + """Repeated ranges reuse device bounds without a sync, and in-place range edits still take effect.""" + env, asset = _make_env(device) + env_ids = torch.arange(env.num_envs, device=device) + velocity_range = {"x": (1.0, 1.0)} + push_by_setting_velocity(env, env_ids, velocity_range) + + with _raise_on_cuda_sync(device): + push_by_setting_velocity(env, env_ids, velocity_range) + torch.testing.assert_close(asset.written["vel"][:, 0], torch.ones(env.num_envs, device=device)) + + # a curriculum may edit the range dictionary in place + velocity_range["x"] = (-2.0, -2.0) + push_by_setting_velocity(env, env_ids, velocity_range) + torch.testing.assert_close(asset.written["vel"][:, 0], torch.full((env.num_envs,), -2.0, device=device)) diff --git a/source/isaaclab/test/envs/test_velocity_command_term.py b/source/isaaclab/test/envs/test_velocity_command_term.py new file mode 100644 index 000000000000..4802c42e707c --- /dev/null +++ b/source/isaaclab/test/envs/test_velocity_command_term.py @@ -0,0 +1,103 @@ +# 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 + +"""Unit tests for the uniform velocity command term.""" + +import math +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from isaaclab.envs.mdp import UniformVelocityCommand, UniformVelocityCommandCfg +from isaaclab.test.utils import test_devices + +pytestmark = pytest.mark.unit + + +@contextmanager +def _raise_on_cuda_sync(device: str): + """Raise if the enclosed code synchronizes a CUDA device.""" + if not device.startswith("cuda"): + yield + return + previous = torch.cuda.get_sync_debug_mode() + torch.cuda.synchronize(device) + torch.cuda.set_sync_debug_mode("error") + try: + yield + finally: + torch.cuda.set_sync_debug_mode(previous) + + +def _make_command(device: str, num_envs: int) -> tuple[UniformVelocityCommand, SimpleNamespace]: + """Create a heading-controlled velocity command on a stub robot.""" + robot_data = SimpleNamespace( + heading_w=SimpleNamespace(torch=torch.zeros(num_envs, device=device)), + root_lin_vel_b=SimpleNamespace(torch=torch.zeros(num_envs, 3, device=device)), + root_ang_vel_b=SimpleNamespace(torch=torch.zeros(num_envs, 3, device=device)), + ) + env = SimpleNamespace( + num_envs=num_envs, + device=device, + scene={"robot": SimpleNamespace(data=robot_data)}, + sim=MagicMock(), + extras={}, + ) + cfg = UniformVelocityCommandCfg( + asset_name="robot", + resampling_time_range=(10.0, 10.0), + heading_command=True, + heading_control_stiffness=0.8, + ranges=UniformVelocityCommandCfg.Ranges( + lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi) + ), + ) + return UniformVelocityCommand(cfg, env), robot_data + + +@pytest.mark.parametrize("device", test_devices()) +def test_heading_control_updates_only_heading_envs(device): + """Heading envs track the wrapped, clipped heading error without a sync; others keep their sampled yaw rate.""" + heading_target = [3.0, 0.5, -3.0, 0.2, 1.0, -0.1] + heading = [-3.0, 0.0, 3.0, 0.0, -1.0, 0.0] + is_heading_env = [True, True, True, False, True, False] + sampled_yaw_rate = [0.3, -0.4, 0.5, -0.6, 0.7, -0.8] + command, robot_data = _make_command(device, len(heading)) + command.heading_target[:] = torch.tensor(heading_target, device=device) + robot_data.heading_w.torch[:] = torch.tensor(heading, device=device) + command.is_heading_env[:] = torch.tensor(is_heading_env, device=device) + command.is_standing_env[:] = False + command.vel_command_b[:] = 0.25 + command.vel_command_b[:, 2] = torch.tensor(sampled_yaw_rate, device=device) + + with _raise_on_cuda_sync(device): + command._update_command() + + expected_yaw_rate = [] + for target, current, active, sampled in zip(heading_target, heading, is_heading_env, sampled_yaw_rate): + error = math.atan2(math.sin(target - current), math.cos(target - current)) + expected_yaw_rate.append(min(max(0.8 * error, -1.0), 1.0) if active else sampled) + torch.testing.assert_close(command.command[:, 2].cpu(), torch.tensor(expected_yaw_rate)) + torch.testing.assert_close(command.command[:, :2].cpu(), torch.full((len(heading), 2), 0.25)) + + +@pytest.mark.parametrize("device", test_devices()) +def test_reset_logs_metrics_without_host_reads(device): + """Reset logs episode metrics as device scalars without synchronizing the stream.""" + command, robot_data = _make_command(device, 4) + robot_data.root_lin_vel_b.torch[:2, 0] = 1.0 + command.compute(dt=0.1) + + with _raise_on_cuda_sync(device): + extras = command.reset() + + success_rate = command._env.extras["log"]["Metrics/success_rate"] + for value in (*extras.values(), success_rate): + assert isinstance(value, torch.Tensor) and value.ndim == 0 + assert set(extras) == {"error_vel_xy", "error_vel_yaw"} + torch.testing.assert_close(success_rate.cpu(), torch.tensor(0.5)) diff --git a/source/isaaclab/test/managers/test_curriculum_manager.py b/source/isaaclab/test/managers/test_curriculum_manager.py new file mode 100644 index 000000000000..90dd34ac5ec7 --- /dev/null +++ b/source/isaaclab/test/managers/test_curriculum_manager.py @@ -0,0 +1,62 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for the curriculum manager.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from isaaclab.managers import CurriculumManager, CurriculumTermCfg +from isaaclab.test.utils import test_devices + +pytestmark = pytest.mark.unit + + +def advance_level(env, env_ids): + """Advance a persistent level tensor in place and report it with a derived state.""" + env.level += 1.0 + return {"level": env.level, "half_level": env.level / 2.0, "name": "stage"} + + +def advance_scalar(env, env_ids): + """Report the persistent level tensor directly.""" + return env.level + + +@pytest.mark.parametrize("device", test_devices()) +def test_reset_logs_state_snapshots_without_host_reads(device): + """Reset logs tensor states as device snapshots, so later in-place term updates do not change them.""" + sim = MagicMock() + sim.is_playing.return_value = False + env = SimpleNamespace(num_envs=4, device=device, sim=sim, level=torch.zeros((), device=device)) + cfg = { + "dict_term": CurriculumTermCfg(func=advance_level), + "scalar_term": CurriculumTermCfg(func=advance_scalar), + } + manager = CurriculumManager(cfg, env) + manager.compute() + + previous = torch.cuda.get_sync_debug_mode() if device.startswith("cuda") else None + if previous is not None: + torch.cuda.synchronize(device) + torch.cuda.set_sync_debug_mode("error") + try: + extras = manager.reset() + finally: + if previous is not None: + torch.cuda.set_sync_debug_mode(previous) + manager.compute() + + assert extras["Curriculum/dict_term/name"] == "stage" + for key, expected in ( + ("Curriculum/dict_term/level", 1.0), + ("Curriculum/dict_term/half_level", 0.5), + ("Curriculum/scalar_term", 1.0), + ): + assert extras[key].device == torch.device(device) + assert extras[key].item() == expected diff --git a/source/isaaclab/test/managers/test_observation_manager.py b/source/isaaclab/test/managers/test_observation_manager.py index 27b76d300f67..47bc8357d481 100644 --- a/source/isaaclab/test/managers/test_observation_manager.py +++ b/source/isaaclab/test/managers/test_observation_manager.py @@ -371,6 +371,37 @@ def scratch_term(env, value): torch.testing.assert_close(result[:, 3:], torch.full_like(env.data.pos_w, 2.0)) +def in_place_noise(data: torch.Tensor, cfg: noise.NoiseCfg) -> torch.Tensor: + """Custom noise that modifies its input.""" + return data.add_(0.2) + + +@pytest.mark.parametrize( + "noise_cfg", + [ + noise.UniformNoiseCfg(n_min=-0.1, n_max=0.1), + noise.UniformNoiseCfg(n_min=0.0, n_max=0.0), + noise.GaussianNoiseCfg(func=noise.gaussian_noise, mean=0.1, std=0.2), + noise.ConstantNoiseCfg(bias=0.3, operation="abs"), + noise.NoiseCfg(func=in_place_noise), + ], + ids=["uniform", "zero_uniform", "gaussian_callable", "constant_abs", "custom_in_place"], +) +def test_noise_preserves_source_and_rng_stream(setup_env, noise_cfg): + """Noise, clipping and scaling leave the term's source storage intact and draw the expected samples.""" + env = setup_env + cfg = ObservationGroupCfg(enable_corruption=True) + cfg.position = ObservationTermCfg(func=pos_w_data, noise=noise_cfg, clip=(0.0, 0.5), scale=2.0) + manager = ObservationManager({"policy": cfg}, env) + source = env.data.pos_w.clone() + torch.manual_seed(0) + result = manager.compute()["policy"] + torch.manual_seed(0) + expected = noise_cfg.func(source.clone(), noise_cfg).clip(0.0, 0.5) * 2.0 + torch.testing.assert_close(env.data.pos_w, source) + torch.testing.assert_close(result, expected) + + def test_compute_with_2d_history(setup_env): env = setup_env """Test the observation computation with history buffers for 2D observations.""" diff --git a/source/isaaclab/test/managers/test_reward_manager.py b/source/isaaclab/test/managers/test_reward_manager.py index 56b1e1d1e7d4..f337cb092009 100644 --- a/source/isaaclab/test/managers/test_reward_manager.py +++ b/source/isaaclab/test/managers/test_reward_manager.py @@ -4,12 +4,14 @@ # SPDX-License-Identifier: BSD-3-Clause from collections import namedtuple +from types import SimpleNamespace from unittest.mock import MagicMock import pytest import torch from isaaclab.managers import RewardManager, RewardTermCfg +from isaaclab.test.utils import test_devices from isaaclab.utils import configclass pytestmark = pytest.mark.unit @@ -162,3 +164,55 @@ def test_invalid_reward_config(env): } with pytest.raises(ValueError): RewardManager(cfg, env) + + +def env_index_scaled(env, factor: float): + return factor * env.step * torch.arange(env.num_envs, dtype=torch.float, device=env.device) + + +def first_half_active(env): + return torch.arange(env.num_envs, device=env.device) < env.num_envs // 2 + + +@pytest.mark.parametrize("device", test_devices()) +def test_compute_and_reset_match_reference(device): + """Rewards, per-term step rewards, episodic sums and reset logs match a per-term reference.""" + num_envs, dt, max_episode_length_s = 8, 0.1, 2.0 + sim = MagicMock() + sim.is_playing.return_value = False + env = SimpleNamespace(num_envs=num_envs, device=device, sim=sim, max_episode_length_s=max_episode_length_s, step=0) + cfg = { + "scaled": RewardTermCfg(func=env_index_scaled, weight=-0.5, params={"factor": 2.0}), + "mask": RewardTermCfg(func=first_half_active, weight=3.0), + "scalar": RewardTermCfg(func=grilled_chicken, weight=10), + "disabled": RewardTermCfg(func=grilled_chicken_with_bbq, weight=0.0, params={"bbq": True}), + } + rew_man = RewardManager(cfg, env) + + env_index = torch.arange(num_envs, dtype=torch.float64) + expected_sums = {name: torch.zeros(num_envs, dtype=torch.float64) for name in cfg} + for step in range(1, 4): + env.step = step + expected_step = { + "scaled": -0.5 * 2.0 * step * env_index, + "mask": 3.0 * (env_index < num_envs // 2).double(), + "scalar": torch.full((num_envs,), 10.0, dtype=torch.float64), + "disabled": torch.zeros(num_envs, dtype=torch.float64), + } + rewards = rew_man.compute(dt=dt) + torch.testing.assert_close(rewards.cpu(), (sum(expected_step.values()) * dt).float()) + for index, name in enumerate(cfg): + expected_sums[name] += expected_step[name] * dt + torch.testing.assert_close(rew_man._step_reward[:, index].cpu(), expected_step[name].float()) + torch.testing.assert_close(rew_man._episode_sums[name].cpu(), expected_sums[name].float()) + + env_ids = torch.tensor([1, 6], device=device) + extras = rew_man.reset(env_ids) + for name in cfg: + value = extras["Episode_Reward/" + name] + assert value.ndim == 0 + expected = expected_sums[name][[1, 6]].mean() / max_episode_length_s + torch.testing.assert_close(value.cpu(), expected.float()) + # only the reset envs start a new episodic sum + expected_sums[name][[1, 6]] = 0.0 + torch.testing.assert_close(rew_man._episode_sums[name].cpu(), expected_sums[name].float()) diff --git a/source/isaaclab/test/managers/test_termination_manager.py b/source/isaaclab/test/managers/test_termination_manager.py index 6cb588444d9e..3f3d9c5ce00b 100644 --- a/source/isaaclab/test/managers/test_termination_manager.py +++ b/source/isaaclab/test/managers/test_termination_manager.py @@ -9,6 +9,7 @@ import torch from isaaclab.managers import TerminationManager, TerminationTermCfg +from isaaclab.test.utils import test_devices pytestmark = pytest.mark.unit @@ -120,3 +121,38 @@ def test_time_out_vs_terminated_split(env): out = tm.compute() assert torch.all(out) assert torch.all(tm.terminated) and torch.all(tm.time_outs) + + +def fail_first_quarter(env) -> torch.Tensor: + """Returns True for the first quarter of the envs.""" + return torch.arange(env.num_envs, device=env.device) < env.num_envs // 4 + + +@pytest.mark.parametrize("device", test_devices()) +def test_reset_logs_term_rates_without_host_reads(device): + """Reset logs each term's last-episode activation rate as a device scalar without synchronizing.""" + sim = MagicMock() + sim.is_playing.return_value = False + env = DummyEnv(num_envs=20, device=device, sim=sim) + cfg = { + "quarter": TerminationTermCfg(func=fail_first_quarter, time_out=False), + "term_3": TerminationTermCfg(func=fail_every_3_steps, time_out=True), + } + tm = TerminationManager(cfg, env) + env.counter = 1 + tm.compute() + + previous = torch.cuda.get_sync_debug_mode() if device.startswith("cuda") else None + if previous is not None: + torch.cuda.synchronize(device) + torch.cuda.set_sync_debug_mode("error") + try: + extras = tm.reset() + finally: + if previous is not None: + torch.cuda.set_sync_debug_mode(previous) + + assert set(extras) == {"Episode_Termination/quarter", "Episode_Termination/term_3"} + assert all(value.ndim == 0 and value.device == torch.device(device) for value in extras.values()) + assert extras["Episode_Termination/quarter"].item() == 0.25 + assert extras["Episode_Termination/term_3"].item() == 0.0 diff --git a/source/isaaclab/test/utils/test_circular_buffer.py b/source/isaaclab/test/utils/test_circular_buffer.py index 6563e8b89f4a..d72289dd6d70 100644 --- a/source/isaaclab/test/utils/test_circular_buffer.py +++ b/source/isaaclab/test/utils/test_circular_buffer.py @@ -257,19 +257,20 @@ def test_stack_dim_positive_index_equivalent_to_negative(): torch.testing.assert_close(buf_neg.stacked, buf_pos.stacked) -def test_stack_dim_ring_shift_after_overflow(): - """After K+1 frames, the oldest slot must be frame 1 (frame 0 evicted), newest = last.""" - B, H, W, C, K = 2, 4, 4, 3, 2 +@pytest.mark.parametrize( + "K,H,W", + [(2, 4, 4), (4, 4, 4), (4, 512, 256)], + ids=["two_slots", "small_frames", "large_frames"], +) +def test_stack_dim_ring_shift_after_overflow(K, H, W): + """After K+1 frames, the oldest frame is evicted and the slots hold frames 1..K in order.""" + B, C = 2, 3 buf = CircularBuffer(max_len=K, batch_size=B, device="cpu", stack_dim=-1) - f0 = torch.full((B, H, W, C), 0.0) - f1 = torch.full((B, H, W, C), 1.0) - f2 = torch.full((B, H, W, C), 2.0) - buf.append(f0) - buf.append(f1) - buf.append(f2) - stacked = buf.stacked # K=2, so slots are [f1, f2] - torch.testing.assert_close(stacked[..., :C], torch.full((B, H, W, C), 1.0)) - torch.testing.assert_close(stacked[..., C:], torch.full((B, H, W, C), 2.0)) + for value in range(K + 1): + buf.append(torch.full((B, H, W, C), float(value))) + stacked = buf.stacked + for slot in range(K): + torch.testing.assert_close(stacked[..., slot * C : (slot + 1) * C], torch.full((B, H, W, C), slot + 1.0)) def test_stack_dim_reset_clears_buffer(): From 92175d991e9f232f4e1ac5fbfafdcf2141139cf9 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 22:06:10 -0700 Subject: [PATCH 06/18] Trim sensor, marker, and renderer hot-path overhead - Skip ray-caster drift sampling on reset when drift ranges are zero and upload ray-cast drift ranges only when they change. - Stop eager per-step refreshes for debug-visualized sensors; debug callbacks now refresh outdated buffers lazily before drawing. - Return early from VisualizationMarkers.visualize when no backend is active. - Drop per-frame channel-compacting copies in the Isaac RTX renderer by indexing padded annotator tiles with their source channel count. - Remove the Newton Warp renderer read_output copy loop (outputs alias the camera buffers) and an unreachable PhysX ray-caster guard. --- .../changelog.d/perf-hot-path-7b-sensors.rst | 11 +++++ .../isaaclab/markers/visualization_markers.py | 13 +++--- .../sensors/ray_caster/base_ray_caster.py | 30 ++++++++++---- .../ray_caster/base_ray_caster_camera.py | 1 + .../isaaclab/isaaclab/sensors/sensor_base.py | 6 ++- .../markers/test_visualization_markers.py | 23 +++++++++++ .../test/sensors/test_ray_caster_sensor.py | 23 +++++++++++ .../isaaclab/test/sensors/test_sensor_base.py | 37 +++++++++++++++++ .../changelog.d/perf-hot-path-7b-sensors.rst | 4 ++ .../tacsl_sensor/visuotactile_sensor.py | 2 + .../changelog.d/perf-hot-path-7b-sensors.rst | 6 +++ .../renderers/newton_warp_renderer.py | 11 +---- .../sensors/contact_sensor/contact_sensor.py | 1 + .../isaaclab_newton/sensors/pva/pva.py | 1 + .../changelog.d/perf-hot-path-7b-sensors.rst | 5 +++ .../frame_transformer/frame_transformer.py | 1 + .../isaaclab_ov/sensors/pva/pva.py | 1 + .../changelog.d/perf-hot-path-7b-sensors.rst | 7 ++++ .../renderers/isaac_rtx_renderer.py | 24 +++++------ .../sensors/contact_sensor/contact_sensor.py | 1 + .../frame_transformer/frame_transformer.py | 1 + .../isaaclab_physx/sensors/pva/pva.py | 1 + .../sensors/ray_caster/ray_caster.py | 2 - .../test_isaac_rtx_renderer_contract.py | 40 +++++++++++++++++++ 24 files changed, 210 insertions(+), 42 deletions(-) create mode 100644 source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst create mode 100644 source/isaaclab_contrib/changelog.d/perf-hot-path-7b-sensors.rst create mode 100644 source/isaaclab_newton/changelog.d/perf-hot-path-7b-sensors.rst create mode 100644 source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst create mode 100644 source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst diff --git a/source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst new file mode 100644 index 000000000000..9136ed982c1b --- /dev/null +++ b/source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst @@ -0,0 +1,11 @@ +Changed +^^^^^^^ + +* Changed sensors with ``debug_vis=True`` to refresh their buffers lazily from the debug + visualization callback instead of on every :meth:`~isaaclab.sensors.SensorBase.update`. + Custom sensors must refresh outdated buffers (for example through ``data``) in their + ``_debug_vis_callback`` before reading internal data. +* Skipped ray-caster drift sampling on reset when ``drift_range`` and ``ray_cast_drift_range`` + are zero, and uploaded the ray-cast drift ranges to the device only when they change. +* Changed :meth:`~isaaclab.markers.VisualizationMarkers.visualize` to return without validating + or processing inputs when no visualizer backend is active, such as in headless runs. diff --git a/source/isaaclab/isaaclab/markers/visualization_markers.py b/source/isaaclab/isaaclab/markers/visualization_markers.py index 781305b57969..d2f8e9f55c44 100644 --- a/source/isaaclab/isaaclab/markers/visualization_markers.py +++ b/source/isaaclab/isaaclab/markers/visualization_markers.py @@ -159,9 +159,10 @@ def visualize( """Update markers in all initialized visualizer backends. .. note:: - If the markers are hidden, the function returns without updating - backend marker state. This avoids unnecessary work while debug - visualization is disabled. + If the markers are hidden or no visualizer backend is active (e.g. when + running headless), the function returns without validating inputs or + updating marker state. This avoids unnecessary work while nothing + displays the markers. Whenever updating the markers, the input arrays must have the same number of elements in the first dimension. Backends generally require @@ -206,9 +207,9 @@ def visualize( ValueError: When the function is called with all None arguments. """ self._ensure_backends_initialized() - # If markers are hidden, do not spend time normalizing or dispatching - # marker state to the active backends. - if not self.is_visible(): + # If no backend consumes markers (e.g. headless) or markers are hidden, do not spend + # time normalizing or dispatching marker state. + if not self._backends or not self.is_visible(): return norm_translations = self._to_tensor(translations, expected_width=3, name="translations") diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py index 2f1b7798ccbf..1f703a000329 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -69,6 +69,8 @@ def __init__(self, cfg: RayCasterCfg): BaseRayCaster._instance_count += 1 super().__init__(cfg) self._data = RayCasterData() + self._ray_cast_drift_range_list: tuple[tuple[float, float], ...] | None = None + self._ray_cast_drift_ranges: torch.Tensor | None = None def __str__(self) -> str: """Returns: A string containing information about the instance.""" @@ -104,6 +106,14 @@ def data(self) -> RayCasterData: def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): # reset the timers and counters super().reset(env_ids, env_mask) + # Drift buffers start at zero, so zero ranges (the defaults) need no resampling. + sample_drift = any(self.cfg.drift_range) + ray_cast_range_list = tuple( + tuple(self.cfg.ray_cast_drift_range.get(key, (0.0, 0.0))) for key in ("x", "y", "z") + ) + sample_ray_cast_drift = any(any(axis_range) for axis_range in ray_cast_range_list) + if not (sample_drift or sample_ray_cast_drift): + return # determine the selected batch size if env_ids is not None: num_envs_ids = len(range(self._view_count)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) @@ -114,14 +124,19 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None env_ids = slice(None) num_envs_ids = self._view_count # resample drift (uses torch views for indexing) - r = torch.empty(num_envs_ids, 3, device=self.device) - self.drift.torch[env_ids] = r.uniform_(*self.cfg.drift_range) + if sample_drift: + r = torch.empty(num_envs_ids, 3, device=self.device) + self.drift.torch[env_ids] = r.uniform_(*self.cfg.drift_range) # resample the ray cast drift - range_list = [self.cfg.ray_cast_drift_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - ranges = torch.tensor(range_list, device=self.device) - self.ray_cast_drift.torch[env_ids] = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (num_envs_ids, 3), device=self.device - ) + if sample_ray_cast_drift: + # Upload the ranges to the device only when the configuration changes. + if ray_cast_range_list != self._ray_cast_drift_range_list: + self._ray_cast_drift_range_list = ray_cast_range_list + self._ray_cast_drift_ranges = torch.tensor(ray_cast_range_list, device=self.device) + ranges = self._ray_cast_drift_ranges + self.ray_cast_drift.torch[env_ids] = math_utils.sample_uniform( + ranges[:, 0], ranges[:, 1], (num_envs_ids, 3), device=self.device + ) """ Implementation. @@ -324,6 +339,7 @@ def _set_debug_vis_impl(self, debug_vis: bool): def _debug_vis_callback(self, event): if self._data._ray_hits_w is None: return + self._update_outdated_buffers() # drop missed rays (inf) before visualizing viz_points = wp.to_torch(self._data._ray_hits_w).reshape(-1, 3) viz_points = viz_points[~torch.isinf(viz_points).any(dim=1)] diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster_camera.py index c5463f38376b..b8f5443d3510 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster_camera.py @@ -446,6 +446,7 @@ def _debug_vis_callback(self, event): # Debug visualization can be toggled before ray buffers are initialized. if not hasattr(self, "ray_hits_w"): return + self._update_outdated_buffers() # drop missed rays (inf) before visualizing viz_points = self.ray_hits_w.torch.reshape(-1, 3) viz_points = viz_points[~torch.isinf(viz_points).any(dim=-1)] diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py index ee06d9dfc1fc..00e443fa4753 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base.py @@ -206,8 +206,8 @@ def update(self, dt: float, force_recompute: bool = False): inputs=[self._is_outdated, self._timestamp, self._timestamp_last_update, dt, self.cfg.update_period], device=self._device, ) - # Update the buffers - if force_recompute or self._is_visualizing: + # Update the buffers; debug visualization refreshes lazily from its callback. + if force_recompute: self._update_outdated_buffers(force_recompute=force_recompute) @staticmethod @@ -316,6 +316,8 @@ def _debug_vis_callback(self, event): """Callback for debug visualization. This function calls the visualization objects and sets the data to visualize into them. + Sensor buffers are refreshed lazily, so implementations must refresh outdated buffers + (for example through :attr:`data`) before reading internal data. """ raise NotImplementedError(f"Debug visualization is not implemented for {self.__class__.__name__}.") diff --git a/source/isaaclab/test/markers/test_visualization_markers.py b/source/isaaclab/test/markers/test_visualization_markers.py index 01a4b9635dc5..c402d7929416 100644 --- a/source/isaaclab/test/markers/test_visualization_markers.py +++ b/source/isaaclab/test/markers/test_visualization_markers.py @@ -272,6 +272,7 @@ def test_first_visualize_defaults_to_first_prototype_when_count_matches_prototyp def test_usd_marker(sim): """Test with marker from a USD.""" + sim._has_offscreen_render = True # create a marker config = clone(FRAME_MARKER_CFG) config.prim_path = "/World/Visuals/test_frames" @@ -292,6 +293,7 @@ def test_usd_marker(sim): def test_visualization_skips_updates_when_invisible(sim): """When invisible, visualize should not update marker state.""" + sim._has_offscreen_render = True # create a marker config = clone(POSITION_GOAL_MARKER_CFG) config.prim_path = "/World/Visuals/test_protos" @@ -322,6 +324,27 @@ def test_visualization_skips_updates_when_invisible(sim): assert test_marker.count == 4 +def test_visualize_without_backend_skips_until_backend_is_active(sim): + """Without an active marker backend, visualize is a no-op; a backend enabled later receives updates.""" + config = VisualizationMarkersCfg( + prim_path="/World/Visuals/headless_marker", + markers={"test": sim_utils.SphereCfg(radius=0.1)}, + ) + test_marker = VisualizationMarkers(config) + assert test_marker._backends == [] + + # Mismatched inputs are not validated because nothing consumes them. + test_marker.visualize( + translations=torch.zeros(3, 3, device=sim.device), orientations=torch.zeros(2, 4, device=sim.device) + ) + assert test_marker.count == 1 + + sim._has_offscreen_render = True + test_marker.visualize(translations=torch.zeros(3, 3, device=sim.device)) + assert len(test_marker._backends) == 1 + assert test_marker.count == 3 + + def test_newton_marker_backend_registers_and_updates_state_without_frame_capture(sim): """Newton marker backend state should be registered and ready for Newton-family viewers.""" sim._visualizers.append(_FakeMarkerVisualizer(pumps_app_update=False)) diff --git a/source/isaaclab/test/sensors/test_ray_caster_sensor.py b/source/isaaclab/test/sensors/test_ray_caster_sensor.py index 40e9a1f346ab..84ef8e109521 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_sensor.py +++ b/source/isaaclab/test/sensors/test_ray_caster_sensor.py @@ -182,6 +182,7 @@ def test_ray_caster_reset_resamples_drift(sim_ground): sim_utils.create_prim("/World/Sensor", "Xform", translation=(0.0, 0.0, 2.0)) cfg = _ray_caster_cfg("/World/Sensor", "world") cfg.drift_range = (0.01, 0.05) # force non-zero drift + cfg.ray_cast_drift_range = {"x": (0.1, 0.2), "y": (-0.2, -0.1)} sensor = RayCaster(cfg) sim.reset() # sim.reset() initializes the sensor with zero drift; call sensor.reset() to resample @@ -216,6 +217,28 @@ def test_ray_caster_reset_resamples_drift(sim_ground): assert not torch.allclose(drift_after, drift_before), ( "reset() must resample drift; values must change from initial sample" ) + # Ray cast drift is sampled per axis; missing axes default to zero range. + ray_cast_drift = sensor.ray_cast_drift.torch + assert ((ray_cast_drift[:, 0] >= 0.1) & (ray_cast_drift[:, 0] <= 0.2)).all() + assert ((ray_cast_drift[:, 1] >= -0.2) & (ray_cast_drift[:, 1] <= -0.1)).all() + assert (ray_cast_drift[:, 2] == 0.0).all() + + +@pytest.mark.isaacsim_ci +def test_ray_caster_reset_skips_zero_drift_sampling(sim_ground): + """reset() with the default zero drift ranges leaves drift at zero without consuming random numbers.""" + sim = sim_ground + + sim_utils.create_prim("/World/Sensor", "Xform", translation=(0.0, 0.0, 2.0)) + sensor = RayCaster(_ray_caster_cfg("/World/Sensor", "world")) + sim.reset() + + get_rng_state = torch.cuda.get_rng_state if torch.device(sensor.device).type == "cuda" else torch.get_rng_state + rng_state = get_rng_state() + sensor.reset() + assert torch.equal(get_rng_state(), rng_state), "reset() must not sample drift when all drift ranges are zero" + assert (sensor.drift.torch == 0.0).all() + assert (sensor.ray_cast_drift.torch == 0.0).all() @pytest.mark.isaacsim_ci diff --git a/source/isaaclab/test/sensors/test_sensor_base.py b/source/isaaclab/test/sensors/test_sensor_base.py index 49d8608d3ca8..f0adba69327a 100644 --- a/source/isaaclab/test/sensors/test_sensor_base.py +++ b/source/isaaclab/test/sensors/test_sensor_base.py @@ -88,6 +88,21 @@ class DummySensorCfg(SensorBaseCfg): prim_path = "{ENV_REGEX_NS}/Cube/dummy_sensor" +class DummyVisSensor(DummySensor): + """Dummy sensor whose debug visualization records the counts it displays.""" + + def __init__(self, cfg): + self.visualized_counts = [] + super().__init__(cfg) + + def _set_debug_vis_impl(self, debug_vis: bool): + pass + + def _debug_vis_callback(self, event): + if self._is_initialized: + self.visualized_counts.append(self.data.count.clone()) + + def _populate_scene(): """""" @@ -307,6 +322,28 @@ def test_repeated_data_reads_are_graph_safe(create_dummy_sensor, device): wp.capture_launch(capture.graph) +@pytest.mark.parametrize("device", test_devices(DeviceScope.DEFAULT_CUDA)) +def test_debug_vis_refreshes_lazily_from_callback(create_dummy_sensor, device): + """Debug visualization does not force per-step refreshes; its callback sees current data.""" + sensor_cfg, sim, dt = create_dummy_sensor + sensor_cfg.debug_vis = True + sensor = DummyVisSensor(cfg=sensor_cfg) + sim.step() + sim.reset() + assert sensor._debug_vis_handle is not None + + backend_update_count = sensor.backend_update_count + for _ in range(3): + sensor.update(dt=dt) + assert sensor.backend_update_count == backend_update_count + + sim.vis_marker_registry.dispatch_callbacks() + assert sensor.backend_update_count == backend_update_count + 1 + torch.testing.assert_close( + sensor.visualized_counts[-1], torch.ones(sensor.num_instances, dtype=torch.int32, device=device) + ) + + @pytest.mark.parametrize("device", ("cpu",)) def test_rigid_body_ancestor_expr_trims_only_terminal_suffix(create_dummy_sensor, device): """Test that ancestor expression trimming keeps repeated path segments above the sensor.""" diff --git a/source/isaaclab_contrib/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab_contrib/changelog.d/perf-hot-path-7b-sensors.rst new file mode 100644 index 000000000000..1cd421cd175e --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/perf-hot-path-7b-sensors.rst @@ -0,0 +1,4 @@ +Changed +^^^^^^^ + +* Changed visuo-tactile sensor debug visualization to refresh outdated sensor buffers before drawing. diff --git a/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor.py b/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor.py index 9fc1a3c768e0..8f86d7fa073e 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor.py +++ b/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor.py @@ -905,6 +905,8 @@ def _debug_vis_callback(self, event): # Safety check - return if not properly initialized if not hasattr(self, "_tactile_visualizer") or self._tactile_visualizer is None: return + if self._is_initialized: + self._update_outdated_buffers() vis_points = None if self.cfg.visualize_sdf_closest_pts and hasattr(self, "debug_closest_points_wolrd"): diff --git a/source/isaaclab_newton/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab_newton/changelog.d/perf-hot-path-7b-sensors.rst new file mode 100644 index 000000000000..acb46b4bf75c --- /dev/null +++ b/source/isaaclab_newton/changelog.d/perf-hot-path-7b-sensors.rst @@ -0,0 +1,6 @@ +Changed +^^^^^^^ + +* Removed the unused per-frame output copy loop from + :class:`~isaaclab_newton.renderers.NewtonWarpRenderer`, whose outputs alias the camera buffers. +* Changed Newton contact and PVA debug visualization to refresh outdated sensor buffers before drawing. diff --git a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py index 5265417d3fca..849e4b99655c 100644 --- a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py +++ b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py @@ -656,17 +656,8 @@ def _launch_render(self, render_data: RenderData) -> None: render_data._convert_segmentation() def read_output(self, render_data: RenderData, camera_data: CameraData) -> None: - """Copy rendered outputs to the camera data buffers. + """Publish segmentation metadata; pixel outputs alias the camera buffers bound in :meth:`set_outputs`. See :meth:`~isaaclab.renderers.base_renderer.BaseRenderer.read_output`.""" - for output_name in camera_data.output: - if output_name == "rgb": - continue - image_data = render_data.get_output(output_name) - if image_data is not None: - output_wp = camera_data.output[output_name].warp - if image_data.ptr != output_wp.ptr: - wp.copy(output_wp, image_data) - # Publish the segmentation id-to-label metadata (idToLabels / idToSemantics) alongside the # pixel buffers. for output_name, info in render_data.segmentation_info().items(): diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py index 4f7ff5d70246..307ffaff1bc7 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py @@ -560,6 +560,7 @@ def _set_debug_vis_impl(self, debug_vis: bool): def _debug_vis_callback(self, event): if not self._is_initialized: return + self._update_outdated_buffers() sensing_transforms = wp.to_torch(self.contact_view.sensing_obj_transforms) positions = sensing_transforms.reshape(self._num_envs, self._num_sensors, 7)[..., :3] diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py index 14894e2f1fa2..269d6bca9bf1 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/pva/pva.py @@ -209,6 +209,7 @@ def _set_debug_vis_impl(self, debug_vis: bool): def _debug_vis_callback(self, event): if self._newton_model is None: return + self._update_outdated_buffers() # base position (offset upward for visibility) base_pos_w = self._data.pos_w.torch.clone() base_pos_w[:, 2] += 0.5 diff --git a/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst new file mode 100644 index 000000000000..42f6c9dbe3d6 --- /dev/null +++ b/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst @@ -0,0 +1,5 @@ +Changed +^^^^^^^ + +* Changed OVPhysX PVA and frame-transformer debug visualization to refresh outdated sensor + buffers before drawing. diff --git a/source/isaaclab_ov/isaaclab_ov/sensors/frame_transformer/frame_transformer.py b/source/isaaclab_ov/isaaclab_ov/sensors/frame_transformer/frame_transformer.py index 0dec1c3c1368..3c9d0c633135 100644 --- a/source/isaaclab_ov/isaaclab_ov/sensors/frame_transformer/frame_transformer.py +++ b/source/isaaclab_ov/isaaclab_ov/sensors/frame_transformer/frame_transformer.py @@ -477,6 +477,7 @@ def _set_debug_vis_impl(self, debug_vis: bool) -> None: def _debug_vis_callback(self, event) -> None: if not self.is_initialized or self._raw_transforms is None or not hasattr(self, "frame_visualizer"): return + self._update_outdated_buffers() # Convert warp -> torch at the boundary for visualization source_pos_w = wp.to_torch(self._data._source_pos_w) diff --git a/source/isaaclab_ov/isaaclab_ov/sensors/pva/pva.py b/source/isaaclab_ov/isaaclab_ov/sensors/pva/pva.py index 233866f61749..a29d15c4d9ac 100644 --- a/source/isaaclab_ov/isaaclab_ov/sensors/pva/pva.py +++ b/source/isaaclab_ov/isaaclab_ov/sensors/pva/pva.py @@ -272,6 +272,7 @@ def _debug_vis_callback(self, event): # safely return if the sensor has not been bound yet (matches the PhysX `_view is None` idiom) if self._root_view is None: return + self._update_outdated_buffers() # get marker location # -- base state (convert warp -> torch for visualization) base_pos_w = self._data.pos_w.torch.clone() diff --git a/source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst new file mode 100644 index 000000000000..f52bb69a5d5e --- /dev/null +++ b/source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst @@ -0,0 +1,7 @@ +Changed +^^^^^^^ + +* Removed the per-frame channel-compacting copies for normals, motion vectors, HDR color, and + simple-shading outputs in :class:`~isaaclab_physx.renderers.IsaacRtxRenderer`. +* Changed PhysX contact, PVA, and frame-transformer debug visualization to refresh outdated + sensor buffers before drawing. diff --git a/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py b/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py index bab4a9213b08..d0a50b556758 100644 --- a/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py +++ b/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py @@ -695,19 +695,14 @@ def tiling_grid_shape(): ptr=tiled_data_buffer.ptr, shape=(*tiled_data_buffer.shape, 4), dtype=wp.uint8, device=device ) - # For motion vectors, use specialized kernel that reads 4 channels but only writes 2 - # Note: Not doing this breaks the alignment of the data (check: https://github.com/isaac-sim/IsaacLab/issues/2003) - if data_type == "motion_vectors": - tiled_data_buffer = tiled_data_buffer[:, :, :2].contiguous() - - # For normals, we only require the first three channels of the tiled buffer - # Note: Not doing this breaks the alignment of the data (check: https://github.com/isaac-sim/IsaacLab/issues/4239) - if data_type == "normals": - tiled_data_buffer = tiled_data_buffer[:, :, :3].contiguous() - if data_type in SIMPLE_SHADING_MODES: - tiled_data_buffer = tiled_data_buffer[:, :, :3].contiguous() - if data_type == str(RenderBufferKind.RGB_HDR): - tiled_data_buffer = tiled_data_buffer[:, :, :3].contiguous() + tile_height, tile_width, num_channels = (int(dim) for dim in buf_wp.shape[1:]) + # Motion vectors, normals, HDR color, and simple shading annotators return 4 channels while the + # outputs keep only the leading ones. Index the source with its own channel count so the + # kernel, which copies only the destination channels, stays aligned without a compacting copy + # (see https://github.com/isaac-sim/IsaacLab/issues/2003 and #4239). + source_channels = num_channels + if data_type in ("motion_vectors", "normals", str(RenderBufferKind.RGB_HDR), *SIMPLE_SHADING_MODES): + source_channels = int(tiled_data_buffer.shape[2]) # ``reshape_tiled_image`` indexes the tiled buffer as # (num_tiles_y * height, num_tiles_x * width, channels), but annotators hand this data back @@ -718,12 +713,11 @@ def tiling_grid_shape(): # are ignored exactly as the previous flattened indexing ignored them. Keeping the view 3D # instead of 1D also keeps every dimension within Warp's per-dimension array size limit, so # large environment counts and camera resolutions no longer overflow a flattened dimension. - tile_height, tile_width, num_channels = (int(dim) for dim in buf_wp.shape[1:]) # ``tiled_source`` must outlive the view below: the view does not own the annotator memory. tiled_source = tiled_data_buffer tiled_data_buffer = wp.array( ptr=tiled_source.ptr, - shape=(num_tiles_y * tile_height, num_tiles_x * tile_width, num_channels), + shape=(num_tiles_y * tile_height, num_tiles_x * tile_width, source_channels), dtype=tiled_source.dtype, device=device, ) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py index d26d621c9f38..689a160e1eb6 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py @@ -701,6 +701,7 @@ def _debug_vis_callback(self, event): # note: this invalidity happens because of isaac sim view callbacks if self.body_physx_view is None: return + self._update_outdated_buffers() # Convert warp data to torch at the boundary for visualization net_forces_torch = self._data.net_normal_forces_w.torch # (N, B, 3) net_contact_force_w = torch.linalg.norm(net_forces_torch, dim=-1) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer.py b/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer.py index 9a9055e10c98..0401ab53868d 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer.py @@ -486,6 +486,7 @@ def _set_debug_vis_impl(self, debug_vis: bool): self.frame_visualizer.set_visibility(False) def _debug_vis_callback(self, event): + self._update_outdated_buffers() # Convert warp -> torch at the boundary for visualization source_pos_w = wp.to_torch(self._data._source_pos_w) source_quat_w = wp.to_torch(self._data._source_quat_w) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py b/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py index 7b52717ce34c..8e51737eea27 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py @@ -333,6 +333,7 @@ def _debug_vis_callback(self, event): # note: this invalidity happens because of isaac sim view callbacks if self._view is None: return + self._update_outdated_buffers() # get marker location # -- base state (convert warp -> torch for visualization) base_pos_w = self._data.pos_w.torch.clone() diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py index ca74b3562dfe..2aae7d9b2cf6 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py @@ -198,8 +198,6 @@ def _initialize_impl(self) -> None: def _get_view_transforms_wp(self) -> wp.array: """Refresh the PhysX transform buffer and return its cached typed Warp view.""" if self._transforms_prefetched: - if self._raw_transforms is None: - raise RuntimeError("RayCaster transforms were marked prefetched before a buffer was cached.") return self._raw_transforms if self._physx_body_view is None: diff --git a/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py b/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py index c7533bf16e06..f5c33f39a5cb 100644 --- a/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py +++ b/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py @@ -490,6 +490,46 @@ def test_render_treats_empty_annotator_frame_as_not_ready(monkeypatch, data_type launch.assert_not_called() +@pytest.mark.parametrize( + ("data_type", "num_channels"), + [("normals", 3), ("motion_vectors", 2), ("rgb_hdr", 3), ("simple_shading_full_mdl", 3)], +) +def test_render_keeps_leading_channels_of_padded_annotator_tiles(monkeypatch, data_type, num_channels): + """Four-channel annotator tiles are split per camera, keeping only the output's leading channels.""" + _install_omni_stubs(monkeypatch) + import isaaclab_physx.renderers.isaac_rtx_renderer as rtx_renderer + from isaaclab_physx.renderers.isaac_rtx_renderer_cfg import IsaacRtxRendererCfg + + # Four 2x3 cameras tiled on a 2x2 grid; every source value is distinct. + view_count, height, width = 4, 2, 3 + tiled = np.arange(2 * height * 2 * width * 4, dtype=np.float32).reshape(2 * height, 2 * width, 4) + annotator = MagicMock() + annotator.get_data.return_value = tiled + output_buffer = wp.zeros((view_count, height, width, num_channels), dtype=wp.float32, device="cpu") + render_data = SimpleNamespace( + annotators={data_type: annotator}, + output_data={data_type: SimpleNamespace(warp=output_buffer)}, + spec=SimpleNamespace(view_count=view_count, device="cpu", cfg=SimpleNamespace(width=width, height=height)), + renderer_info={}, + ppisp_pipeline=None, + _hdr_scratch_wp=None, + ) + renderer = rtx_renderer.IsaacRtxRenderer.__new__(rtx_renderer.IsaacRtxRenderer) + renderer.cfg = IsaacRtxRendererCfg() + + with patch.object(rtx_renderer, "ensure_isaac_rtx_render_update"): + renderer.render(render_data) + + expected = np.stack( + [ + tiled[row * height : (row + 1) * height, col * width : (col + 1) * width, :num_channels] + for row in range(2) + for col in range(2) + ] + ) + np.testing.assert_array_equal(output_buffer.numpy(), expected) + + @pytest.mark.parametrize( "states", [ From c8330862df804ffc59845b65061fb4373f02220d Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 22:06:51 -0700 Subject: [PATCH 07/18] Remove host syncs from task hot paths and RL wrappers Store direct-task joint, body, and tendon ids as device tensors, precompute static limits and default poses, use the actuator target-command API, log success metrics as 0-d tensors, compute Lift contact terms from one stacked norm, cache the ANYmal symmetry permutations and constant tensors, and send only marker indices for static success markers. --- .../changelog.d/perf-hot-path-7b-tasks.rst | 5 + .../isaaclab_rl/rl_games/rl_games.py | 1 - source/isaaclab_rl/isaaclab_rl/sb3.py | 5 +- .../perf-hot-path-7b-tasks.minor.rst | 26 ++++ .../core/cabinet/cabinet_direct_env.py | 28 ++-- .../core/cabinet/mdp/rewards.py | 5 +- .../core/fourbar_pole/mdp/rewards.py | 3 +- .../core/lift/mdp/commands/pose_commands.py | 43 +++--- .../isaaclab_tasks/core/lift/mdp/events.py | 19 ++- .../core/lift/mdp/observations.py | 17 +-- .../isaaclab_tasks/core/lift/mdp/rewards.py | 38 ++--- .../core/lift/mdp/terminations.py | 2 +- .../core/locomotion/locomotion_direct_env.py | 9 +- .../core/locomotion/mdp/observations.py | 13 +- .../core/reorient/mdp/rewards.py | 6 +- .../core/reorient/mdp/terminations.py | 2 +- .../core/reorient/reorient_direct_env.py | 66 ++++---- .../isaaclab_tasks/core/reorient/utils.py | 32 ++-- .../core/velocity/mdp/symmetry/anymal.py | 141 ++++++++---------- .../test/core/test_anymal_symmetry.py | 90 +++++++++++ .../test/core/test_lift_env_cfg.py | 69 ++++++++- .../test/core/test_reorient_utils.py | 21 ++- 22 files changed, 428 insertions(+), 213 deletions(-) create mode 100644 source/isaaclab_rl/changelog.d/perf-hot-path-7b-tasks.rst create mode 100644 source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst create mode 100644 source/isaaclab_tasks/test/core/test_anymal_symmetry.py diff --git a/source/isaaclab_rl/changelog.d/perf-hot-path-7b-tasks.rst b/source/isaaclab_rl/changelog.d/perf-hot-path-7b-tasks.rst new file mode 100644 index 000000000000..79fb7e693522 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/perf-hot-path-7b-tasks.rst @@ -0,0 +1,5 @@ +Changed +^^^^^^^ + +* Computed the Stable-Baselines3 wrapper's done flags from the host copies of the termination and + truncation flags, saving one device-to-host transfer per step. diff --git a/source/isaaclab_rl/isaaclab_rl/rl_games/rl_games.py b/source/isaaclab_rl/isaaclab_rl/rl_games/rl_games.py index 5acbe9f571c2..924b1317f6b7 100644 --- a/source/isaaclab_rl/isaaclab_rl/rl_games/rl_games.py +++ b/source/isaaclab_rl/isaaclab_rl/rl_games/rl_games.py @@ -310,7 +310,6 @@ def step(self, actions): # noqa: D102 # process observations and states obs_and_states = self._process_obs(obs_dict) # move buffers to rl-device - # note: we perform clone to prevent issues when rl-device and sim-device are the same. rew = rew.to(device=self._rl_device) dones = (terminated | truncated).to(device=self._rl_device) extras = { diff --git a/source/isaaclab_rl/isaaclab_rl/sb3.py b/source/isaaclab_rl/isaaclab_rl/sb3.py index 6348db5c0ceb..5eb5986c42f6 100644 --- a/source/isaaclab_rl/isaaclab_rl/sb3.py +++ b/source/isaaclab_rl/isaaclab_rl/sb3.py @@ -245,8 +245,6 @@ def step_async(self, actions): # noqa: D102 def step_wait(self) -> VecEnvStepReturn: # noqa: D102 # record step information obs_dict, rew, terminated, truncated, extras = self.env.step(self._async_actions) - # compute reset ids - dones = terminated | truncated # convert data types to numpy depending on backend # note: ManagerBasedRLEnv uses torch backend (by default). @@ -254,7 +252,8 @@ def step_wait(self) -> VecEnvStepReturn: # noqa: D102 rewards = rew.detach().cpu().numpy() terminated = terminated.detach().cpu().numpy() truncated = truncated.detach().cpu().numpy() - dones = dones.detach().cpu().numpy() + # compute reset ids on the host copies, which saves a device-to-host transfer + dones = terminated | truncated reset_ids = dones.nonzero()[0] diff --git a/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst b/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst new file mode 100644 index 000000000000..e7c34b2b55fb --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst @@ -0,0 +1,26 @@ +Changed +^^^^^^^ + +* Stored the joint, body, and tendon indices of the Reorient, Cabinet, and locomotion direct tasks as + device tensors, and precomputed their static joint limits and default arm pose, removing host-to-device + index uploads and CUDA synchronizations from every physics substep. The public + ``actuated_dof_indices``, ``finger_bodies``, ``finger_wrench_bodies``, ``arm_joint_ids``, and + ``finger_joint_ids`` attributes, and the indices returned by + :func:`~isaaclab_tasks.core.reorient.utils.resolve_actuated_tendons`, are now ``torch.long`` tensors + instead of lists. Call ``.tolist()`` where a list is needed. +* Switched the Reorient, Cabinet, and locomotion direct tasks from the deprecated articulation target + setters to ``articulation.actuators.target_command``. +* Logged the Cabinet, locomotion direct, Lift deformable and cable, and Fourbar pole success metrics as + 0-d device tensors instead of synchronizing with ``.item()`` on every reset. Call ``.item()`` where a + float is needed. +* Changed :meth:`~isaaclab_tasks.core.reorient.utils.EpisodeErrorRecorder.reset` to compute its statistics + without a host synchronization. Once any error has been recorded, it returns NaN statistics instead of an + empty dictionary when none of the selected environments has a sample. Check the values with + :func:`torch.isnan` instead of testing for missing keys. +* Stacked the Lift contact sensor forces so :func:`~isaaclab_tasks.core.lift.mdp.contacts` and + :func:`~isaaclab_tasks.core.lift.mdp.contact_count` take one norm per call instead of one per sensor. +* Precomputed the ANYmal symmetry augmentation as one cached column permutation and sign tensor per + transform, removing the per-minibatch index and sign uploads (about 1.5 ms to 0.15 ms per call). +* Sent only the marker indices to Lift success markers on static assets each step, cached the locomotion + walk-target offset and the Lift reset position bounds on the device, and read the Reorient goal directly + from the command buffers instead of concatenating the command on every access. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py index f760674092d8..e52f620fa3d4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py @@ -29,8 +29,11 @@ def __init__(self, cfg: CabinetDirectEnvCfg, render_mode: str | None = None, **k super().__init__(cfg, render_mode, **kwargs) self._robot, self._cabinet = self.scene["robot"], self.scene["cabinet"] - self.arm_joint_ids, _ = self._robot.find_joints(self.cfg.arm_joint_names) - self.finger_joint_ids, _ = self._robot.find_joints(self.cfg.finger_joint_names) + arm_joint_ids, _ = self._robot.find_joints(self.cfg.arm_joint_names) + finger_joint_ids, _ = self._robot.find_joints(self.cfg.finger_joint_names) + # device indices: indexing a CUDA tensor with a host list uploads it on every call + self.arm_joint_ids = torch.tensor(arm_joint_ids, dtype=torch.long, device=self.device) + self.finger_joint_ids = torch.tensor(finger_joint_ids, dtype=torch.long, device=self.device) self.ee_body_idx = self._robot.find_bodies(self.cfg.ee_body_name)[0][0] self.left_finger_body_idx = self._robot.find_bodies(self.cfg.left_finger_body_name)[0][0] self.right_finger_body_idx = self._robot.find_bodies(self.cfg.right_finger_body_name)[0][0] @@ -52,6 +55,8 @@ def __init__(self, cfg: CabinetDirectEnvCfg, render_mode: str | None = None, **k self.previous_actions = torch.zeros_like(self.actions) self.arm_joint_targets = torch.zeros((self.num_envs, len(self.arm_joint_ids)), device=self.device) + # the default arm pose is static, so its gather is hoisted out of the per-step action path + self._arm_default_joint_pos = self._robot.data.default_joint_pos.torch[:, self.arm_joint_ids] self.finger_joint_targets = torch.zeros((self.num_envs, len(self.finger_joint_ids)), device=self.device) # frame offsets, repeated for every environment @@ -89,8 +94,7 @@ def _pre_physics_step(self, actions: torch.Tensor) -> None: self.actions[:] = actions self.arm_joint_targets[:] = ( - self._robot.data.default_joint_pos.torch[:, self.arm_joint_ids] - + self.cfg.arm_action_scale * self.actions[:, : len(self.arm_joint_ids)] + self._arm_default_joint_pos + self.cfg.arm_action_scale * self.actions[:, : len(self.arm_joint_ids)] ) self.finger_joint_targets[:] = torch.where( self.actions[:, -1:] < 0.0, @@ -99,14 +103,9 @@ def _pre_physics_step(self, actions: torch.Tensor) -> None: ) def _apply_action(self) -> None: - self._robot.set_joint_position_target_index( - target=self.arm_joint_targets, - joint_ids=self.arm_joint_ids, - ) - self._robot.set_joint_position_target_index( - target=self.finger_joint_targets, - joint_ids=self.finger_joint_ids, - ) + target_command = self._robot.actuators.target_command + target_command.set_position_index(value=self.arm_joint_targets, joint_ids=self.arm_joint_ids) + target_command.set_position_index(value=self.finger_joint_targets, joint_ids=self.finger_joint_ids) def _get_observations(self) -> dict[str, torch.Tensor]: robot_joint_pos = self._robot.data.joint_pos.torch - self._robot.data.default_joint_pos.torch @@ -211,8 +210,9 @@ def _reset_idx(self, env_ids: Sequence[int] | None) -> None: env_ids = torch.arange(self.num_envs, device=self.device) log = self.extras.setdefault("log", {}) - log["Metrics/success_rate"] = self._episode_succeeded[env_ids].float().mean().item() - log["Metrics/drawer_pos"] = self._best_drawer_pos[env_ids].mean().item() + # 0-dim device tensors: no host sync on reset + log["Metrics/success_rate"] = self._episode_succeeded[env_ids].float().mean() + log["Metrics/drawer_pos"] = self._best_drawer_pos[env_ids].mean() for name, episode_sum in self._episode_reward_sums.items(): log[f"Episode_Reward/{name}"] = torch.mean(episode_sum[env_ids]) / self.max_episode_length_s index_fill_(episode_sum, env_ids, 0.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py index a52913fdc7c0..f22a7af233eb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py @@ -155,8 +155,9 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): def reset(self, env_ids: torch.Tensor): log = self._env.extras.setdefault("log", {}) - log["Metrics/success_rate"] = self.succeeded[env_ids].float().mean().item() - log["Metrics/drawer_pos"] = self.best_drawer_pos[env_ids].mean().item() + # 0-dim device tensors: no host sync on reset + log["Metrics/success_rate"] = self.succeeded[env_ids].float().mean() + log["Metrics/drawer_pos"] = self.best_drawer_pos[env_ids].mean() index_fill_(self.succeeded, env_ids, False) index_fill_(self.best_drawer_pos, env_ids, 0.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/mdp/rewards.py index 12f793d8f049..5eae41f9b1af 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/mdp/rewards.py @@ -35,7 +35,8 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): self._hold_steps = max(1, round(hold_time_s / env.step_dt)) def reset(self, env_ids: torch.Tensor): - self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._success[env_ids].mean().item() + # 0-dim device tensor: no host sync on reset + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._success[env_ids].mean() index_fill_(self._consecutive_upright, env_ids, 0.0) def __call__( diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/commands/pose_commands.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/commands/pose_commands.py index da05e6b61685..1143f5194456 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/commands/pose_commands.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/commands/pose_commands.py @@ -135,12 +135,7 @@ def _update_metrics(self): if not self.cfg.position_only: self.metrics["orientation_error"] = torch.linalg.norm(rot_error, dim=-1) success_id &= self.metrics["orientation_error"] < 0.5 - if self.success_vis_asset is not None: - self.success_visualizer.visualize( - self._get_success_vis_pos_w(), - marker_indices=success_id.int(), - environment_ids=self._env.scene._ALL_INDICES, - ) + self._visualize_success(success_id) def _get_success_vis_pos_w(self) -> torch.Tensor: """Return the success visualization positions in the world frame.""" @@ -148,6 +143,24 @@ def _get_success_vis_pos_w(self) -> torch.Tensor: return self._static_success_vis_pos_w return self.success_vis_asset.data.root_pos_w.torch + def _visualize_success(self, success: torch.Tensor) -> None: + """Show each environment's success marker. + + Args: + success: Per-environment success flags, shape ``(num_envs,)``. + """ + if self.success_vis_asset is None: + return + if self._static_success_vis_pos_w is not None: + # positions and environment ids were set at construction; resending them copies to the host + self.success_visualizer.visualize(marker_indices=success.int()) + return + self.success_visualizer.visualize( + self._get_success_vis_pos_w(), + marker_indices=success.int(), + environment_ids=self._env.scene._ALL_INDICES, + ) + def _resample_command(self, env_ids: Sequence[int]): # sample new pose targets # -- position @@ -248,15 +261,8 @@ def _update_metrics(self): com_w = self.object.data.root_pos_w.torch self.metrics["position_error"] = torch.linalg.norm(self.pose_command_w[:, :3] - com_w, dim=-1) - if self.success_vis_asset is None: - return # same success radius as the goal markers of the base class - success_id = (self.metrics["position_error"] < 0.05).int() - self.success_visualizer.visualize( - self._get_success_vis_pos_w(), - marker_indices=success_id, - environment_ids=self._env.scene._ALL_INDICES, - ) + self._visualize_success(self.metrics["position_error"] < 0.05) class CableUniformPoseCommand(ObjectUniformPoseCommand): @@ -288,14 +294,7 @@ def _update_metrics(self): segment_pos_w = self._segment_position_w() self.metrics["position_error"] = torch.linalg.norm(self.pose_command_w[:, :3] - segment_pos_w, dim=-1) - if self.success_vis_asset is None: - return - success_id = (self.metrics["position_error"] < 0.05).int() - self.success_visualizer.visualize( - self._get_success_vis_pos_w(), - marker_indices=success_id, - environment_ids=self._env.scene._ALL_INDICES, - ) + self._visualize_success(self.metrics["position_error"] < 0.05) def _debug_vis_callback(self, event): if not self.robot.is_initialized: diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py index c2a39d721d16..29743f6aecc1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py @@ -7,6 +7,7 @@ from __future__ import annotations +import functools from typing import TYPE_CHECKING import numpy as np @@ -72,6 +73,20 @@ def reset_joints_shared_offset( ) +def _xyz_ranges(ranges: dict[str, tuple[float, float]], device: str | torch.device) -> torch.Tensor: + """Per-axis ``(low, high)`` bounds [m] of the ``x``, ``y``, ``z`` keys as a device tensor, shape ``(3, 2)``. + + Missing axes are ``(0, 0)``. The tensor is cached so resets do not upload the constant bounds. + """ + bounds = tuple(tuple(ranges.get(axis, (0.0, 0.0))) for axis in ("x", "y", "z")) + return _cached_xyz_ranges(bounds, str(device)) + + +@functools.cache +def _cached_xyz_ranges(bounds: tuple[tuple[float, float], ...], device: str) -> torch.Tensor: + return torch.tensor(bounds, device=device) + + def reset_cable_state_uniform( env: ManagerBasedEnv, env_ids: torch.Tensor, @@ -82,7 +97,7 @@ def reset_cable_state_uniform( asset: CableObject = env.scene[asset_cfg.name] segment_pose = asset.data.default_segment_pose_w.torch[env_ids].clone() segment_velocity = asset.data.default_segment_velocity_w.torch[env_ids].clone() - ranges = torch.tensor([position_range.get(axis, (0.0, 0.0)) for axis in ("x", "y", "z")], device=asset.device) + ranges = _xyz_ranges(position_range, asset.device) offset = sample_uniform(ranges[:, 0], ranges[:, 1], (segment_pose.shape[0], 3), device=asset.device) segment_pose[..., :3] += offset.unsqueeze(1) asset.write_segment_pose_to_sim_index(segment_pose=segment_pose, env_ids=env_ids) @@ -822,7 +837,7 @@ def reset_deformable_over_support( deformable: DeformableObject = env.scene[asset_cfg.name] supports: tuple[RigidObject, RigidObject] = (env.scene[support_cfg[0].name], env.scene[support_cfg[1].name]) - ranges = torch.tensor([position_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z")], device=deformable.device) + ranges = _xyz_ranges(position_range, deformable.device) num_envs = len(range(env.num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) offset = sample_uniform(ranges[:, 0], ranges[:, 1], (num_envs, 3), device=deformable.device) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/observations.py index 315514474d44..95f02aed5cfc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/observations.py @@ -53,17 +53,8 @@ class body_state_b(ManagerTermBase): The state for each body is stacked horizontally as ``[position(3), quaternion(4)(xyzw), linvel(3), angvel(3)]`` and then concatenated over bodies. - - The body indices are baked to a device tensor at construction. """ - def __init__(self, cfg: ObservationTermCfg, env: ManagerBasedRLEnv): - super().__init__(cfg, env) - body_ids = cfg.params["body_asset_cfg"].body_ids - if isinstance(body_ids, list): - body_ids = torch.tensor(body_ids, dtype=torch.long, device=env.device) - self._body_ids = body_ids - def __call__( self, env: ManagerBasedRLEnv, @@ -86,10 +77,10 @@ def __call__( body_asset: Articulation = env.scene[body_asset_cfg.name] base_asset: Articulation = env.scene[base_asset_cfg.name] # world pose of the bodies, flattened over environments - body_pos_w = body_asset.data.body_pos_w.torch[:, self._body_ids] + body_pos_w = body_asset.data.body_pos_w.torch[:, body_asset_cfg.body_ids] num_bodies = body_pos_w.shape[1] body_pos_w = body_pos_w.reshape(-1, 3) - body_quat_w = body_asset.data.body_quat_w.torch[:, self._body_ids].reshape(-1, 4) + body_quat_w = body_asset.data.body_quat_w.torch[:, body_asset_cfg.body_ids].reshape(-1, 4) # world pose of the base frame, broadcast over the bodies root_pos_w = base_asset.data.root_link_pos_w.torch.unsqueeze(1).expand(-1, num_bodies, -1).reshape(-1, 3) root_quat_w = base_asset.data.root_link_quat_w.torch.unsqueeze(1).expand(-1, num_bodies, -1).reshape(-1, 4) @@ -99,8 +90,8 @@ def __call__( # and contact-response differences); pose-only states with observation history transfer across # physics backends, velocity states do not if include_vel: - body_lin_vel_w = body_asset.data.body_lin_vel_w.torch[:, self._body_ids].view(-1, 3) - body_ang_vel_w = body_asset.data.body_ang_vel_w.torch[:, self._body_ids].view(-1, 3) + body_lin_vel_w = body_asset.data.body_lin_vel_w.torch[:, body_asset_cfg.body_ids].view(-1, 3) + body_ang_vel_w = body_asset.data.body_ang_vel_w.torch[:, body_asset_cfg.body_ids].view(-1, 3) body_lin_vel_b = quat_apply_inverse(root_quat_w, body_lin_vel_w) body_ang_vel_b = quat_apply_inverse(root_quat_w, body_ang_vel_w) out = torch.cat((body_pos_b, body_quat_b, body_lin_vel_b, body_ang_vel_b), dim=1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py index eb30aeb47be4..4c68a4a1f54a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py @@ -67,14 +67,8 @@ def contacts(env: ManagerBasedRLEnv, threshold: float, thumb_name: str, finger_n Returns: Boolean tensor indicating good contact condition per environment. """ - thumb_mag = _contact_force_mag(env.scene.sensors[thumb_name], env.num_envs) - - any_finger_contact = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) - for finger_name in finger_names: - finger_mag = _contact_force_mag(env.scene.sensors[finger_name], env.num_envs) - any_finger_contact = any_finger_contact | (finger_mag > threshold) - - return (thumb_mag > threshold) & any_finger_contact + in_contact = _contact_force_mags(env, [thumb_name, *finger_names]) > threshold + return in_contact[:, 0] & in_contact[:, 1:].any(dim=1) def contact_count(env: ManagerBasedRLEnv, threshold: float, sensor_names: list[str]) -> torch.Tensor: @@ -91,11 +85,7 @@ def contact_count(env: ManagerBasedRLEnv, threshold: float, sensor_names: list[s Returns: Tensor of shape (num_envs,) with the count of sensors in contact per environment. """ - count = torch.zeros(env.num_envs, dtype=torch.float32, device=env.device) - - for sensor_name in sensor_names: - mag = _contact_force_mag(env.scene.sensors[sensor_name], env.num_envs) - count += (mag > threshold).float() + count = (_contact_force_mags(env, sensor_names) > threshold).float().sum(dim=1) return count / len(sensor_names) @@ -413,7 +403,8 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): def reset(self, env_ids: Sequence[int] | None = None) -> None: if env_ids is None: env_ids = slice(None) - self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._succeeded[env_ids].float().mean().item() + # 0-dim device tensor: no host sync on reset + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._succeeded[env_ids].float().mean() index_fill_(self._succeeded, env_ids, False) @@ -513,10 +504,21 @@ def cable_segment_goal_reached( return (distance < success_threshold).float() -def _contact_force_mag(sensor: ContactSensor, num_envs: int) -> torch.Tensor: - """Per-environment contact force magnitude [N] of a single-body, single-filter contact sensor.""" - force = sensor.data.normal_force_matrix_w.torch.view(num_envs, 3) - return torch.linalg.norm(force, dim=-1) +def _contact_force_mags(env: ManagerBasedRLEnv, sensor_names: Sequence[str]) -> torch.Tensor: + """Per-environment contact force magnitudes [N] of single-body, single-filter contact sensors. + + The forces are stacked so a single norm covers every sensor. + + Args: + env: The environment instance. + sensor_names: Names of the contact sensors in the scene. + + Returns: + Force magnitudes [N], shape ``(num_envs, len(sensor_names))``. + """ + sensors: list[ContactSensor] = [env.scene.sensors[name] for name in sensor_names] + forces = torch.stack([sensor.data.normal_force_matrix_w.torch.view(env.num_envs, 3) for sensor in sensors], dim=1) + return torch.linalg.norm(forces, dim=-1) def _deformable_com_goal_metrics( diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/terminations.py index 5f38e5d6ea9a..1eb5f0cd8cf7 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/terminations.py @@ -153,7 +153,7 @@ def joint_vel_out_of_sim_limit( ) -> torch.Tensor: """Terminate when joint velocities exceed solver limits [m/s or rad/s, depending on joint type].""" asset: Articulation = env.scene[asset_cfg.name] - joint_ids = asset_cfg.joint_ids if asset_cfg.joint_ids is not None else slice(None) + joint_ids = asset_cfg.joint_ids return torch.any( torch.abs(asset.data.joint_vel.torch[:, joint_ids]) > asset.data.joint_vel_limits.torch[:, joint_ids], dim=1 ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/locomotion_direct_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/locomotion_direct_env.py index e8efc3052826..14492b263154 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/locomotion_direct_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/locomotion_direct_env.py @@ -56,7 +56,9 @@ def __init__(self, cfg: AntEnvCfg | HumanoidEnvCfg, render_mode: str | None = No self.gear_ratio_scaled = self.joint_gears / torch.max(self.joint_gears) # resolve against the sensor's own body list: its ordering is backend-specific and does not # necessarily match the articulation's body ordering - self._feet_body_idx, _ = self.joint_wrench.find_bodies(self.cfg.feet_body_names) + feet_body_idx, _ = self.joint_wrench.find_bodies(self.cfg.feet_body_names) + # device indices: indexing a CUDA tensor with a host list uploads it on every call + self._feet_body_idx = torch.tensor(feet_body_idx, dtype=torch.long, device=self.device) # walk target, placed far enough away that the robot never reaches it self.targets = self.scene.env_origins + torch.tensor( @@ -69,7 +71,7 @@ def _pre_physics_step(self, actions: torch.Tensor) -> None: self.actions = actions.clone() # the action is clamped before scaling: unbounded joint efforts drive the solver to NaN forces = self.action_scale * self.joint_gears * torch.clamp(self.actions, -1.0, 1.0) - self.robot.set_joint_effort_target_index(target=forces) + self.robot.actuators.target_command.set_effort_index(value=forces) def _apply_action(self) -> None: # Joint effort targets persist in the articulation command buffer across decimation substeps. @@ -168,7 +170,8 @@ def _reset_idx(self, env_ids: Sequence[int] | None): # log survival success rate before resetting (survived = timed out without falling) survived = self.reset_time_outs[env_ids].float() - self.extras.setdefault("log", {})["Metrics/success_rate"] = survived.mean().item() + # 0-dim device tensor: no host sync on reset + self.extras.setdefault("log", {})["Metrics/success_rate"] = survived.mean() super()._reset_idx(env_ids) index_fill_(self.actions, env_ids, 0.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/mdp/observations.py index 60580dd2636b..5b9ae367a461 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/mdp/observations.py @@ -7,6 +7,7 @@ from __future__ import annotations +import functools from typing import TYPE_CHECKING import torch @@ -45,7 +46,17 @@ def walk_target_w(env: ManagerBasedEnv, target_pos: tuple[float, float, float]) The target is specified relative to the environment origin so that every robot walks along the same direction in its own frame, independently of where its environment sits in the grid. """ - return env.scene.env_origins + torch.tensor(target_pos, device=env.device) + return env.scene.env_origins + _walk_target_offset(tuple(target_pos), str(env.device)) + + +@functools.cache +def _walk_target_offset(target_pos: tuple[float, float, float], device: str) -> torch.Tensor: + """Walk target offset [m] as a device tensor. + + Cached per ``(target_pos, device)`` so the hot path does not upload the constant on every call; + ``device`` is a string so that it is hashable for the cache. + """ + return torch.tensor(target_pos, device=device) def base_heading_proj( diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/rewards.py index 05f244726fd1..b03ad147b48d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/rewards.py @@ -39,7 +39,7 @@ def success_bonus( command_term: ReorientCommand = env.command_manager.get_term(command_name) # obtain the goal orientation - goal_quat_w = command_term.command[:, 3:7] + goal_quat_w = command_term.quat_command_w # obtain the threshold for the orientation error threshold = command_term.cfg.orientation_success_threshold # calculate the orientation error @@ -65,7 +65,7 @@ def track_pos_l2( command_term: ReorientCommand = env.command_manager.get_term(command_name) # obtain the goal position - goal_pos_e = command_term.command[:, 0:3] + goal_pos_e = command_term.pos_command_e # obtain the object position in the environment frame object_pos_e = asset.data.root_pos_w.torch - env.scene.env_origins @@ -93,7 +93,7 @@ def track_orientation_inv_l2( command_term: ReorientCommand = env.command_manager.get_term(command_name) # obtain the goal orientation - goal_quat_w = command_term.command[:, 3:7] + goal_quat_w = command_term.quat_command_w # calculate the orientation error dtheta = math_utils.quat_error_magnitude(asset.data.root_quat_w.torch, goal_quat_w) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/terminations.py index f3eec876a1e9..cceb9705186d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/mdp/terminations.py @@ -42,6 +42,6 @@ def object_away_from_goal( # object pos asset_pos_e = asset.data.root_pos_w.torch - env.scene.env_origins - goal_pos_e = command_term.command[:, :3] + goal_pos_e = command_term.pos_command_e return torch.linalg.norm(asset_pos_e - goal_pos_e, ord=2, dim=1) > threshold diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py index 417de502387f..8b8a6266de3b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py @@ -45,26 +45,36 @@ def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | # -- robot introspection: joints, bodies, limits -- self.num_hand_dofs = self.hand.num_joints - self.actuated_dof_indices, _ = self.hand.find_joints(cfg.actuated_joint_names) - if len(self.actuated_dof_indices) != len(cfg.actuated_joint_names): + actuated_dof_indices, _ = self.hand.find_joints(cfg.actuated_joint_names) + if len(actuated_dof_indices) != len(cfg.actuated_joint_names): raise ValueError( - f"Expected {len(cfg.actuated_joint_names)} actuated joints, found {len(self.actuated_dof_indices)}." + f"Expected {len(cfg.actuated_joint_names)} actuated joints, found {len(actuated_dof_indices)}." ) - - self.finger_bodies, fingertip_body_names = self.hand.find_bodies(self.cfg.fingertip_body_names) - if len(self.finger_bodies) != len(self.cfg.fingertip_body_names): + # device indices: indexing a CUDA tensor with a host list uploads it on every call + self.actuated_dof_indices = torch.tensor(actuated_dof_indices, dtype=torch.long, device=self.device) + # a slice avoids the gather/scatter when every joint is actuated (``find_joints`` returns sorted ids) + all_actuated = actuated_dof_indices == list(range(self.num_hand_dofs)) + self._actuated_dof_ids = slice(None) if all_actuated else self.actuated_dof_indices + + finger_bodies, fingertip_body_names = self.hand.find_bodies(self.cfg.fingertip_body_names) + if len(finger_bodies) != len(self.cfg.fingertip_body_names): raise ValueError( - f"Expected {len(self.cfg.fingertip_body_names)} fingertip bodies, found {len(self.finger_bodies)}." + f"Expected {len(self.cfg.fingertip_body_names)} fingertip bodies, found {len(finger_bodies)}." ) - self.num_fingertips = len(self.finger_bodies) + self.finger_bodies = torch.tensor(finger_bodies, dtype=torch.long, device=self.device) + self.num_fingertips = len(finger_bodies) self.finger_wrench_bodies = [] if self._joint_wrench_sensor is not None: - for body_name in fingertip_body_names: - self.finger_wrench_bodies.append(self._joint_wrench_sensor.body_names.index(body_name)) - self.finger_wrench_bodies.sort() + finger_wrench_bodies = sorted( + self._joint_wrench_sensor.body_names.index(name) for name in fingertip_body_names + ) + self.finger_wrench_bodies = torch.tensor(finger_wrench_bodies, dtype=torch.long, device=self.device) joint_pos_limits = self.hand.data.joint_limits.torch.to(self.device) self.hand_dof_lower_limits = joint_pos_limits[..., 0] self.hand_dof_upper_limits = joint_pos_limits[..., 1] + # the joint limits are static, so the actuated subset is gathered once rather than every substep + self._actuated_lower_limits = self.hand_dof_lower_limits[:, self._actuated_dof_ids] + self._actuated_upper_limits = self.hand_dof_upper_limits[:, self._actuated_dof_ids] # -- actuation targets (EMA-smoothed joint position targets) -- self.prev_targets = torch.zeros((self.num_envs, self.num_hand_dofs), dtype=torch.float, device=self.device) @@ -94,7 +104,7 @@ def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | self.z_unit_tensor = torch.tensor([0, 0, 1], dtype=torch.float, device=self.device).repeat((self.num_envs, 1)) # -- articulation write handles -- - self._set_joint_pos_target = self.hand.set_joint_position_target_index + self._set_joint_pos_target = self.hand.actuators.target_command.set_position_index self._write_obj_root_pose = self.object.write_root_pose_to_sim_index self._write_obj_root_vel = self.object.write_root_velocity_to_sim_index self._write_hand_joint_pos = self.hand.write_joint_position_to_sim_index @@ -106,27 +116,19 @@ def _pre_physics_step(self, actions: torch.Tensor) -> None: def _apply_action(self) -> None: # Joint actions come first, matching the manager task's action-term order. A hand whose # motors also pull tendons consumes the remaining columns in its own subclass. + ids = self._actuated_dof_ids num_joint_actions = len(self.actuated_dof_indices) - self.cur_targets[:, self.actuated_dof_indices] = unscale_transform( - self.actions[:, :num_joint_actions], - self.hand_dof_lower_limits[:, self.actuated_dof_indices], - self.hand_dof_upper_limits[:, self.actuated_dof_indices], - ) - self.cur_targets[:, self.actuated_dof_indices] = ( - self.cfg.act_moving_average * self.cur_targets[:, self.actuated_dof_indices] - + (1.0 - self.cfg.act_moving_average) * self.prev_targets[:, self.actuated_dof_indices] + targets = unscale_transform( + self.actions[:, :num_joint_actions], self._actuated_lower_limits, self._actuated_upper_limits ) - self.cur_targets[:, self.actuated_dof_indices] = saturate( - self.cur_targets[:, self.actuated_dof_indices], - self.hand_dof_lower_limits[:, self.actuated_dof_indices], - self.hand_dof_upper_limits[:, self.actuated_dof_indices], + targets = ( + self.cfg.act_moving_average * targets + (1.0 - self.cfg.act_moving_average) * self.prev_targets[:, ids] ) + targets = saturate(targets, self._actuated_lower_limits, self._actuated_upper_limits) + self.cur_targets[:, ids] = targets + self.prev_targets[:, ids] = targets - self.prev_targets[:, self.actuated_dof_indices] = self.cur_targets[:, self.actuated_dof_indices] - - self._set_joint_pos_target( - target=self.cur_targets[:, self.actuated_dof_indices], joint_ids=self.actuated_dof_indices - ) + self._set_joint_pos_target(value=targets, joint_ids=ids) def _get_observations(self) -> dict: if self.cfg.asymmetric_obs: @@ -244,8 +246,8 @@ def _reset_idx(self, env_ids: Sequence[int]): self._reset_target_pose(env_ids) # reset object - object_default_pose = self.object.data.default_root_pose.torch.clone()[env_ids] - object_default_vel = self.object.data.default_root_vel.torch.clone()[env_ids] + object_default_pose = self.object.data.default_root_pose.torch[env_ids] + object_default_vel = self.object.data.default_root_vel.torch[env_ids] pos_noise = sample_uniform(-1.0, 1.0, (len(env_ids), 3), device=self.device) # global object positions object_default_pose[:, 0:3] = ( @@ -272,7 +274,7 @@ def _reset_idx(self, env_ids: Sequence[int]): self.prev_targets[env_ids] = dof_pos self.cur_targets[env_ids] = dof_pos - self._set_joint_pos_target(target=dof_pos, env_ids=env_ids) + self._set_joint_pos_target(value=dof_pos, env_ids=env_ids) self._write_hand_joint_pos(position=dof_pos, env_ids=env_ids) self._write_hand_joint_vel(velocity=dof_vel, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py index 2678b7ea265e..33c64ce12a49 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py @@ -41,6 +41,8 @@ def __init__(self, num_envs: int, device: str | torch.device): """ self.minimum_error = torch.full((num_envs,), torch.inf, device=device) self._has_sample = torch.zeros(num_envs, dtype=torch.bool, device=device) + # host-side flag: before the first update no episode has a sample, known without a device read + self._updated = False def update(self, error: torch.Tensor) -> None: """Record one error sample for every environment. @@ -58,6 +60,7 @@ def update(self, error: torch.Tensor) -> None: # force a host synchronization every step, so substitute-and-minimum instead torch.minimum(self.minimum_error, torch.where(finite, error, self.minimum_error), out=self.minimum_error) self._has_sample |= finite + self._updated = True def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> dict[str, torch.Tensor]: """Summarize and clear completed episodes. @@ -68,19 +71,21 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> Returns: Mean, median, and 90th-percentile episode-minimum errors as 0-dim device tensors, so logging them does not force a host - synchronization in the reset path. The result is empty when none of - the selected environments has a sample. + synchronization in the reset path. Environments without a sample + are excluded. The result is empty before the first :meth:`update`, + and its values are NaN when none of the selected environments has + a sample. """ if env_ids is None: env_ids = slice(None) - valid = self._has_sample[env_ids] - values = self.minimum_error[env_ids][valid] statistics = {} - if values.numel() > 0: + if self._updated: + # NaN marks environments without a sample; the nan-reductions skip them without a host sync + values = torch.where(self._has_sample[env_ids], self.minimum_error[env_ids], torch.nan) statistics = { - "mean": values.mean(), - "median": values.median(), - "p90": torch.quantile(values, 0.9), + "mean": values.nanmean(), + "median": values.nanmedian(), + "p90": torch.nanquantile(values, 0.9), } index_fill_(self.minimum_error, env_ids, torch.inf) index_fill_(self._has_sample, env_ids, False) @@ -255,7 +260,7 @@ def resolve_actuated_tendons( num_envs: int, device: str, position_limits: tuple[float, float], -) -> tuple[list[int], torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Resolve a hand's actuated tendons and the limits its command is unscaled against. A tendon has its own index space, so the indices come from :meth:`find_fixed_tendons` rather @@ -271,7 +276,8 @@ def resolve_actuated_tendons( position_limits: Lower and upper bound of the tendon command [m or rad]. Returns: - The tendon indices and the lower / upper limit tensors, shaped ``(num_envs, n_tendons)``. + The tendon indices as a device tensor, so per-step commands do not upload them, and the + lower / upper limit tensors, shaped ``(num_envs, n_tendons)``. Raises: ValueError: If a requested tendon is not present on the articulation. @@ -281,4 +287,8 @@ def resolve_actuated_tendons( raise ValueError(f"Expected {len(tendon_names)} actuated tendons, found {len(indices)}.") shape = (num_envs, len(indices)) lower, upper = position_limits - return indices, torch.full(shape, lower, device=device), torch.full(shape, upper, device=device) + return ( + torch.tensor(indices, dtype=torch.long, device=device), + torch.full(shape, lower, device=device), + torch.full(shape, upper, device=device), + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/symmetry/anymal.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/symmetry/anymal.py index 4a1bc3aa9b20..7c54f1b83bc2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/symmetry/anymal.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/symmetry/anymal.py @@ -7,6 +7,7 @@ from __future__ import annotations +import functools from typing import TYPE_CHECKING import torch @@ -104,29 +105,8 @@ def _transform_policy_obs_left_right(env: ManagerBasedRLEnv, obs: torch.Tensor) Returns: The transformed observation tensor with left-right symmetry applied. """ - # copy observation tensor - obs = obs.clone() - device = obs.device - # lin vel - obs[:, :3] = obs[:, :3] * torch.tensor([1, -1, 1], device=device) - # ang vel - obs[:, 3:6] = obs[:, 3:6] * torch.tensor([-1, 1, -1], device=device) - # projected gravity - obs[:, 6:9] = obs[:, 6:9] * torch.tensor([1, -1, 1], device=device) - # velocity command - obs[:, 9:12] = obs[:, 9:12] * torch.tensor([1, -1, -1], device=device) - # joint pos - obs[:, 12:24] = _switch_anymal_joints_left_right(obs[:, 12:24]) - # joint vel - obs[:, 24:36] = _switch_anymal_joints_left_right(obs[:, 24:36]) - # last actions - obs[:, 36:48] = _switch_anymal_joints_left_right(obs[:, 36:48]) - - # note: this is hard-coded for grid-pattern of ordering "xy" and size (1.6, 1.0) - if "height_scan" in env.observation_manager.active_terms["policy"]: - obs[:, 48:235] = obs[:, 48:235].view(-1, 11, 17).flip(dims=[1]).view(-1, 11 * 17) - - return obs + perm, sign = _policy_obs_symmetry("left_right", obs.shape[1], _has_height_scan(env), str(obs.device)) + return obs[:, perm] * sign def _transform_policy_obs_front_back(env: ManagerBasedRLEnv, obs: torch.Tensor) -> torch.Tensor: @@ -145,29 +125,8 @@ def _transform_policy_obs_front_back(env: ManagerBasedRLEnv, obs: torch.Tensor) Returns: The transformed observation tensor with front-back symmetry applied. """ - # copy observation tensor - obs = obs.clone() - device = obs.device - # lin vel - obs[:, :3] = obs[:, :3] * torch.tensor([-1, 1, 1], device=device) - # ang vel - obs[:, 3:6] = obs[:, 3:6] * torch.tensor([1, -1, -1], device=device) - # projected gravity - obs[:, 6:9] = obs[:, 6:9] * torch.tensor([-1, 1, 1], device=device) - # velocity command - obs[:, 9:12] = obs[:, 9:12] * torch.tensor([-1, 1, -1], device=device) - # joint pos - obs[:, 12:24] = _switch_anymal_joints_front_back(obs[:, 12:24]) - # joint vel - obs[:, 24:36] = _switch_anymal_joints_front_back(obs[:, 24:36]) - # last actions - obs[:, 36:48] = _switch_anymal_joints_front_back(obs[:, 36:48]) - - # note: this is hard-coded for grid-pattern of ordering "xy" and size (1.6, 1.0) - if "height_scan" in env.observation_manager.active_terms["policy"]: - obs[:, 48:235] = obs[:, 48:235].view(-1, 11, 17).flip(dims=[2]).view(-1, 11 * 17) - - return obs + perm, sign = _policy_obs_symmetry("front_back", obs.shape[1], _has_height_scan(env), str(obs.device)) + return obs[:, perm] * sign """ @@ -189,9 +148,8 @@ def _transform_actions_left_right(actions: torch.Tensor) -> torch.Tensor: Returns: The transformed actions tensor with left-right symmetry applied. """ - actions = actions.clone() - actions[:] = _switch_anymal_joints_left_right(actions[:]) - return actions + perm, sign = _joint_symmetry("left_right", str(actions.device)) + return actions[:, perm] * sign def _transform_actions_front_back(actions: torch.Tensor) -> torch.Tensor: @@ -208,9 +166,8 @@ def _transform_actions_front_back(actions: torch.Tensor) -> torch.Tensor: Returns: The transformed actions tensor with front-back symmetry applied. """ - actions = actions.clone() - actions[:] = _switch_anymal_joints_front_back(actions[:]) - return actions + perm, sign = _joint_symmetry("front_back", str(actions.device)) + return actions[:, perm] * sign """ @@ -229,32 +186,58 @@ def _transform_actions_front_back(actions: torch.Tensor) -> torch.Tensor: * LH = left hind --> [1, 5, 9] * RF = right front --> [2, 6, 10] * RH = right hind --> [3, 7, 11] -""" - -def _switch_anymal_joints_left_right(joint_data: torch.Tensor) -> torch.Tensor: - """Applies a left-right symmetry transformation to the joint data tensor.""" - joint_data_switched = torch.zeros_like(joint_data) - # left <-- right - joint_data_switched[..., [0, 4, 8, 1, 5, 9]] = joint_data[..., [2, 6, 10, 3, 7, 11]] - # right <-- left - joint_data_switched[..., [2, 6, 10, 3, 7, 11]] = joint_data[..., [0, 4, 8, 1, 5, 9]] - - # Flip the sign of the HAA joints - joint_data_switched[..., [0, 1, 2, 3]] *= -1.0 - - return joint_data_switched - - -def _switch_anymal_joints_front_back(joint_data: torch.Tensor) -> torch.Tensor: - """Applies a front-back symmetry transformation to the joint data tensor.""" - joint_data_switched = torch.zeros_like(joint_data) - # front <-- hind - joint_data_switched[..., [0, 4, 8, 2, 6, 10]] = joint_data[..., [1, 5, 9, 3, 7, 11]] - # hind <-- front - joint_data_switched[..., [1, 5, 9, 3, 7, 11]] = joint_data[..., [0, 4, 8, 2, 6, 10]] - - # Flip the sign of the HFE and KFE joints - joint_data_switched[..., 4:] *= -1 +Each transform is a column permutation and a sign per column: ``out[:, i] = sign[i] * x[:, perm[i]]``. +""" - return joint_data_switched +# signs of the base lin vel, ang vel, projected gravity, and velocity command columns +_BASE_SIGNS = { + "left_right": (1, -1, 1, -1, 1, -1, 1, -1, 1, 1, -1, -1), + "front_back": (-1, 1, 1, 1, -1, -1, -1, 1, 1, -1, 1, -1), +} +# source joint of each switched joint: left <-> right swaps LF/RF and LH/RH, front <-> back swaps LF/LH and RF/RH +_JOINT_PERMS = { + "left_right": (2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9), + "front_back": (1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10), +} +# left <-> right flips the HAA joints, front <-> back flips the HFE and KFE joints +_JOINT_SIGNS = { + "left_right": (-1,) * 4 + (1,) * 8, + "front_back": (1,) * 4 + (-1,) * 8, +} +# height-scan grid dimension flipped per transform, for the (11, 17) grid +_HEIGHT_SCAN_FLIP_DIM = {"left_right": 0, "front_back": 1} + + +def _has_height_scan(env: ManagerBasedRLEnv) -> bool: + """Whether the policy observation contains the height scan.""" + return "height_scan" in env.observation_manager.active_terms["policy"] + + +@functools.cache +def _joint_symmetry(kind: str, device: str) -> tuple[torch.Tensor, torch.Tensor]: + """Column permutation and signs of a transform over the 12 joints, cached per device.""" + return ( + torch.tensor(_JOINT_PERMS[kind], dtype=torch.long, device=device), + torch.tensor(_JOINT_SIGNS[kind], device=device), + ) + + +@functools.cache +def _policy_obs_symmetry(kind: str, num_obs: int, height_scan: bool, device: str) -> tuple[torch.Tensor, torch.Tensor]: + """Column permutation and signs of a transform over the policy observation, cached per device.""" + perm = list(range(num_obs)) + sign = [1] * num_obs + sign[:12] = _BASE_SIGNS[kind] + # joint positions, joint velocities, and last actions + for start in (12, 24, 36): + perm[start : start + 12] = [start + joint for joint in _JOINT_PERMS[kind]] + sign[start : start + 12] = _JOINT_SIGNS[kind] + # note: this is hard-coded for grid-pattern of ordering "xy" and size (1.6, 1.0) + if height_scan: + grid = torch.arange(48, 235).view(11, 17).flip(dims=[_HEIGHT_SCAN_FLIP_DIM[kind]]) + perm[48:235] = grid.flatten().tolist() + return ( + torch.tensor(perm, dtype=torch.long, device=device), + torch.tensor(sign, device=device), + ) diff --git a/source/isaaclab_tasks/test/core/test_anymal_symmetry.py b/source/isaaclab_tasks/test/core/test_anymal_symmetry.py new file mode 100644 index 000000000000..36c35ad3cbe3 --- /dev/null +++ b/source/isaaclab_tasks/test/core/test_anymal_symmetry.py @@ -0,0 +1,90 @@ +# Copyright (c) 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 ANYmal symmetry augmentation.""" + +from types import SimpleNamespace + +import pytest +import torch +from tensordict import TensorDict + +from isaaclab_tasks.core.velocity.mdp.symmetry import anymal + +# Reference: the index-assignment formulation the permutation tables must reproduce. +_LEFT = [0, 4, 8, 1, 5, 9] +_RIGHT = [2, 6, 10, 3, 7, 11] +_FRONT = [0, 4, 8, 2, 6, 10] +_HIND = [1, 5, 9, 3, 7, 11] + + +def _reference_joints(joint_data: torch.Tensor, kind: str) -> torch.Tensor: + first, second = (_LEFT, _RIGHT) if kind == "left_right" else (_FRONT, _HIND) + switched = torch.zeros_like(joint_data) + switched[..., first] = joint_data[..., second] + switched[..., second] = joint_data[..., first] + if kind == "left_right": + switched[..., [0, 1, 2, 3]] *= -1.0 + else: + switched[..., 4:] *= -1 + return switched + + +def _reference_obs(obs: torch.Tensor, kind: str, height_scan: bool) -> torch.Tensor: + signs = { + "left_right": ([1, -1, 1], [-1, 1, -1], [1, -1, 1], [1, -1, -1]), + "front_back": ([-1, 1, 1], [1, -1, -1], [-1, 1, 1], [-1, 1, -1]), + }[kind] + obs = obs.clone() + for block, sign in enumerate(signs): + obs[:, 3 * block : 3 * block + 3] = obs[:, 3 * block : 3 * block + 3] * torch.tensor(sign) + for start in (12, 24, 36): + obs[:, start : start + 12] = _reference_joints(obs[:, start : start + 12], kind) + if height_scan: + flip_dim = 1 if kind == "left_right" else 2 + obs[:, 48:235] = obs[:, 48:235].view(-1, 11, 17).flip(dims=[flip_dim]).view(-1, 11 * 17) + return obs + + +def _reference_augment(obs: torch.Tensor, actions: torch.Tensor, height_scan: bool): + left_right = _reference_obs(obs, "left_right", height_scan) + obs_aug = torch.cat( + ( + obs, + left_right, + _reference_obs(obs, "front_back", height_scan), + _reference_obs(left_right, "front_back", height_scan), + ) + ) + actions_left_right = _reference_joints(actions, "left_right") + actions_aug = torch.cat( + ( + actions, + actions_left_right, + _reference_joints(actions, "front_back"), + _reference_joints(actions_left_right, "front_back"), + ) + ) + return obs_aug, actions_aug + + +@pytest.mark.parametrize("height_scan", [False, True]) +def test_compute_symmetric_states_matches_reference(height_scan: bool): + """Verify the cached permutation and sign tables reproduce the index-assignment formulation.""" + num_envs = 5 + num_obs = 235 if height_scan else 48 + obs = torch.randn(num_envs, num_obs) + actions = torch.randn(num_envs, 12) + active_terms = ["base_lin_vel", "height_scan"] if height_scan else ["base_lin_vel"] + unwrapped = SimpleNamespace(observation_manager=SimpleNamespace(active_terms={"policy": active_terms})) + env = SimpleNamespace(unwrapped=unwrapped) + + obs_aug, actions_aug = anymal.compute_symmetric_states( + env, obs=TensorDict({"policy": obs}, batch_size=[num_envs]), actions=actions + ) + expected_obs, expected_actions = _reference_augment(obs, actions, height_scan) + + assert torch.equal(obs_aug["policy"], expected_obs) + assert torch.equal(actions_aug, expected_actions) diff --git a/source/isaaclab_tasks/test/core/test_lift_env_cfg.py b/source/isaaclab_tasks/test/core/test_lift_env_cfg.py index 69822a9e012d..ccecba79bfdc 100644 --- a/source/isaaclab_tasks/test/core/test_lift_env_cfg.py +++ b/source/isaaclab_tasks/test/core/test_lift_env_cfg.py @@ -13,6 +13,7 @@ from pxr import Usd +from isaaclab.assets import Asset from isaaclab.managers import CommandTerm, ObservationTermCfg, SceneEntityCfg from isaaclab.sim import select_usd_variants from isaaclab.utils.warp import ProxyArray @@ -116,9 +117,10 @@ def test_camera_normalization_is_stationary() -> None: assert torch.allclose(depth_obs.flatten(), torch.tanh(torch.tensor([0.0, 2.0]) / 2) - 0.5) -def test_lift_pose_markers_forward_environment_ids(monkeypatch: pytest.MonkeyPatch) -> None: - """Every lift pose, goal, and success marker should retain its environment ownership.""" - num_envs = 3 +def _make_pose_command( + monkeypatch: pytest.MonkeyPatch, num_envs: int, success_asset: object +) -> tuple[ObjectUniformPoseCommand, torch.Tensor]: + """Build a pose command around fake assets and spy markers; returns it with the shared root positions.""" environment_ids = torch.arange(num_envs) identity_quat = torch.zeros((num_envs, 4)) identity_quat[:, 3] = 1.0 @@ -139,7 +141,6 @@ def test_lift_pose_markers_forward_environment_ids(monkeypatch: pytest.MonkeyPat root_link_pose_w=SimpleNamespace(torch=root_pose_w), ) ) - success_asset = SimpleNamespace(data=SimpleNamespace(root_pos_w=SimpleNamespace(torch=root_pos_w))) scene = _FakeScene(environment_ids, robot=robot, object=object_asset, table=success_asset) env = SimpleNamespace(num_envs=num_envs, device="cpu", scene=scene) cfg = SimpleNamespace( @@ -161,8 +162,15 @@ def _initialize_command_term(command, command_cfg, command_env) -> None: monkeypatch.setattr(CommandTerm, "__init__", _initialize_command_term) monkeypatch.setattr(pose_commands, "VisualizationMarkers", _MarkerSpy) + return ObjectUniformPoseCommand(cfg, env), root_pos_w + - command = ObjectUniformPoseCommand(cfg, env) +def test_lift_pose_markers_forward_environment_ids(monkeypatch: pytest.MonkeyPatch) -> None: + """Every lift pose, goal, and success marker should retain its environment ownership.""" + num_envs = 3 + environment_ids = torch.arange(num_envs) + success_asset = SimpleNamespace(data=SimpleNamespace(root_pos_w=SimpleNamespace(torch=torch.zeros((num_envs, 3))))) + command, root_pos_w = _make_pose_command(monkeypatch, num_envs, success_asset) command._set_debug_vis_impl(True) command._debug_vis_callback(None) command.cfg.position_only = False @@ -184,6 +192,57 @@ def _initialize_command_term(command, command_cfg, command_env) -> None: assert torch.equal(kwargs["environment_ids"], environment_ids) +def test_lift_static_success_markers_update_only_marker_indices(monkeypatch: pytest.MonkeyPatch) -> None: + """Success markers on a static asset are placed once; per-step updates only switch the prototype.""" + num_envs = 3 + static_table = object.__new__(Asset) + static_table.cfg = SimpleNamespace(init_state=SimpleNamespace(pos=(0.5, 0.0, 0.0))) + command, root_pos_w = _make_pose_command(monkeypatch, num_envs, static_table) + # env 1 is at the goal, the others are not + command.pose_command_b[:, 0] = torch.tensor([1.0, 0.0, 1.0]) + + for update_metrics in ( + ObjectUniformPoseCommand._update_metrics, + DeformableUniformPoseCommand._update_metrics, + CableUniformPoseCommand._update_metrics, + ): + command._segment_position_w = lambda: root_pos_w + update_metrics(command) + + calls = command.success_visualizer.calls + placement_args, placement_kwargs = calls[0] + assert torch.equal(placement_args[0], torch.tensor([[0.5, 0.0, 0.0]] * num_envs)) + assert torch.equal(placement_kwargs["environment_ids"], torch.arange(num_envs)) + assert len(calls) == 4 + for args, kwargs in calls[1:]: + assert args == () + assert kwargs.keys() == {"marker_indices"} + assert torch.equal(kwargs["marker_indices"], torch.tensor([0, 1, 0], dtype=torch.int32)) + + +def test_lift_contact_terms_match_per_sensor_reference() -> None: + """Contact gating and counting from stacked sensor forces match a per-sensor evaluation.""" + num_envs, threshold = 64, 1.0 + names = ["thumb", "index", "middle", "ring"] + # contact forces [N] straddling the threshold, shape (num_envs, 1 body, 1 filter, 3) + forces = {name: 1.2 * torch.rand(num_envs, 1, 1, 3) for name in names} + sensors = { + name: SimpleNamespace(data=SimpleNamespace(normal_force_matrix_w=ProxyArray(wp.from_torch(force)))) + for name, force in forces.items() + } + env = SimpleNamespace(num_envs=num_envs, device="cpu", scene=SimpleNamespace(sensors=sensors)) + in_contact = { + name: torch.linalg.norm(force.view(num_envs, 3), dim=-1) > threshold for name, force in forces.items() + } + + expected_contacts = in_contact["thumb"] & (in_contact["index"] | in_contact["middle"] | in_contact["ring"]) + expected_count = sum(in_contact[name].float() for name in names) / len(names) + + assert torch.equal(mdp.contacts(env, threshold, "thumb", names[1:]), expected_contacts) + assert torch.equal(mdp.contact_count(env, threshold, names), expected_count) + assert not mdp.contacts(env, threshold, "thumb", []).any() + + def test_lift_point_cloud_markers_repeat_environment_ids_per_point() -> None: """Flattened point-cloud markers should retain env-major ownership.""" num_envs = 3 diff --git a/source/isaaclab_tasks/test/core/test_reorient_utils.py b/source/isaaclab_tasks/test/core/test_reorient_utils.py index 5eff5ca1e7df..0db609343021 100644 --- a/source/isaaclab_tasks/test/core/test_reorient_utils.py +++ b/source/isaaclab_tasks/test/core/test_reorient_utils.py @@ -68,6 +68,25 @@ def test_episode_error_recorder_skips_episodes_without_samples(): assert recorder.reset(torch.tensor([0, 1])) == {} +def test_episode_error_recorder_reset_excludes_episodes_without_samples(): + """Verify reset statistics cover only sampled episodes, and are NaN when none of them is.""" + recorder = reorient_utils.EpisodeErrorRecorder(num_envs=5, device="cpu") + nan = float("nan") + # envs 1 and 4 only see non-finite errors, so they have no sample + recorder.update(torch.tensor([0.3, nan, 0.1, 0.4, torch.inf])) + recorder.update(torch.tensor([0.2, nan, 0.5, 0.35, nan])) + + statistics = recorder.reset(torch.tensor([0, 1, 2, 3])) + sampled = torch.tensor([0.2, 0.1, 0.35]) + assert {k: v.item() for k, v in statistics.items()} == pytest.approx( + {"mean": sampled.mean().item(), "median": sampled.median().item(), "p90": torch.quantile(sampled, 0.9).item()} + ) + + statistics = recorder.reset(torch.tensor([4])) + assert set(statistics) == {"mean", "median", "p90"} + assert all(torch.isnan(value) for value in statistics.values()) + + def test_episode_error_recorder_update_matches_masked_indexing_reference(): """Verify the sync-free update equals boolean-mask indexing on NaN/inf/finite mixes.""" @@ -163,7 +182,7 @@ def test_resolve_actuated_tendons_returns_indices_in_action_order_and_limit_tens hand, ["rh_LFJ0", "rh_FFJ0"], num_envs=3, device="cpu", position_limits=(0.0, 2.0) ) - assert indices == [3, 0] + assert torch.equal(indices, torch.tensor([3, 0])) assert torch.equal(lower, torch.zeros(3, 2)) assert torch.equal(upper, torch.full((3, 2), 2.0)) From 0c5612a913f1fd9b2ed7bf19b3c8344aeaf15204 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 22:07:27 -0700 Subject: [PATCH 08/18] Use elapsed time for OVPhysX rigid-object COM acceleration --- source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst | 2 ++ .../isaaclab_ov/assets/rigid_object/rigid_object_data.py | 4 +++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst index 42f6c9dbe3d6..9c78cc792aed 100644 --- a/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst +++ b/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst @@ -3,3 +3,5 @@ Changed * Changed OVPhysX PVA and frame-transformer debug visualization to refresh outdated sensor buffers before drawing. +* Changed the OVPhysX rigid-object ``body_com_acc_w`` finite difference to use the elapsed time since the + previous update, matching joint accelerations and the Newton backend. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object_data.py b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object_data.py index f1eb37fd1a80..feaca89b506e 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object_data.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/rigid_object/rigid_object_data.py @@ -452,13 +452,15 @@ def body_com_acc_w(self) -> ProxyArray: if self._body_com_acc_w.timestamp < self._sim_timestamp: if self._previous_body_com_vel is None: self._previous_body_com_vel = wp.clone(self.body_com_vel_w.warp) + # Finite-difference over the elapsed time, matching joint accelerations. + time_elapsed = self._sim_timestamp - self._body_com_acc_w.timestamp wp.launch( shared_kernels.derive_body_acceleration_from_body_com_velocities, dim=(self._num_instances, 1), device=self.device, inputs=[ self.body_com_vel_w.warp, - SimulationManager.get_physics_dt(), + time_elapsed, self._previous_body_com_vel, ], outputs=[ From 282a75b0632d2f256e7ac76f67ac1bc3092bf7f3 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 27 Sep 2026 23:01:01 -0700 Subject: [PATCH 09/18] Trim redundant hot-path regression tests --- .../test/managers/test_observation_manager.py | 22 +++---- .../test/managers/test_reward_manager.py | 58 ++++++------------- .../test/utils/test_circular_buffer.py | 16 +++-- .../test/utils/test_wrench_composer.py | 16 +---- .../test/core/test_lift_env_cfg.py | 35 +++++------ 5 files changed, 53 insertions(+), 94 deletions(-) diff --git a/source/isaaclab/test/managers/test_observation_manager.py b/source/isaaclab/test/managers/test_observation_manager.py index 47bc8357d481..fba5862723a0 100644 --- a/source/isaaclab/test/managers/test_observation_manager.py +++ b/source/isaaclab/test/managers/test_observation_manager.py @@ -377,27 +377,27 @@ def in_place_noise(data: torch.Tensor, cfg: noise.NoiseCfg) -> torch.Tensor: @pytest.mark.parametrize( - "noise_cfg", + ("noise_cfg", "expected_delta"), [ - noise.UniformNoiseCfg(n_min=-0.1, n_max=0.1), - noise.UniformNoiseCfg(n_min=0.0, n_max=0.0), - noise.GaussianNoiseCfg(func=noise.gaussian_noise, mean=0.1, std=0.2), - noise.ConstantNoiseCfg(bias=0.3, operation="abs"), - noise.NoiseCfg(func=in_place_noise), + (noise.UniformNoiseCfg(n_min=-0.1, n_max=0.1), None), + (noise.UniformNoiseCfg(n_min=0.0, n_max=0.0), 0.0), + (noise.NoiseCfg(func=in_place_noise), 0.2), ], - ids=["uniform", "zero_uniform", "gaussian_callable", "constant_abs", "custom_in_place"], + ids=["out_of_place", "same_input", "in_place"], ) -def test_noise_preserves_source_and_rng_stream(setup_env, noise_cfg): - """Noise, clipping and scaling leave the term's source storage intact and draw the expected samples.""" +def test_noise_preserves_source_and_rng_stream(setup_env, noise_cfg, expected_delta): + """Noise ownership paths preserve source storage, clipping, scaling, and random samples.""" env = setup_env cfg = ObservationGroupCfg(enable_corruption=True) cfg.position = ObservationTermCfg(func=pos_w_data, noise=noise_cfg, clip=(0.0, 0.5), scale=2.0) manager = ObservationManager({"policy": cfg}, env) source = env.data.pos_w.clone() torch.manual_seed(0) - result = manager.compute()["policy"] + if expected_delta is None: + expected_delta = torch.rand_like(source) * 0.2 - 0.1 + expected = (source + expected_delta).clip(0.0, 0.5) * 2.0 torch.manual_seed(0) - expected = noise_cfg.func(source.clone(), noise_cfg).clip(0.0, 0.5) * 2.0 + result = manager.compute()["policy"] torch.testing.assert_close(env.data.pos_w, source) torch.testing.assert_close(result, expected) diff --git a/source/isaaclab/test/managers/test_reward_manager.py b/source/isaaclab/test/managers/test_reward_manager.py index f337cb092009..c56028bd9256 100644 --- a/source/isaaclab/test/managers/test_reward_manager.py +++ b/source/isaaclab/test/managers/test_reward_manager.py @@ -11,7 +11,6 @@ import torch from isaaclab.managers import RewardManager, RewardTermCfg -from isaaclab.test.utils import test_devices from isaaclab.utils import configclass pytestmark = pytest.mark.unit @@ -170,49 +169,30 @@ def env_index_scaled(env, factor: float): return factor * env.step * torch.arange(env.num_envs, dtype=torch.float, device=env.device) -def first_half_active(env): - return torch.arange(env.num_envs, device=env.device) < env.num_envs // 2 - - -@pytest.mark.parametrize("device", test_devices()) -def test_compute_and_reset_match_reference(device): - """Rewards, per-term step rewards, episodic sums and reset logs match a per-term reference.""" - num_envs, dt, max_episode_length_s = 8, 0.1, 2.0 +def test_compute_and_reset_match_reference(): + """A partial reset reports each term's accumulated reward and preserves other environments.""" + num_envs, dt, max_episode_length_s = 3, 0.25, 2.0 sim = MagicMock() sim.is_playing.return_value = False - env = SimpleNamespace(num_envs=num_envs, device=device, sim=sim, max_episode_length_s=max_episode_length_s, step=0) + env = SimpleNamespace(num_envs=num_envs, device="cpu", sim=sim, max_episode_length_s=max_episode_length_s, step=0) cfg = { - "scaled": RewardTermCfg(func=env_index_scaled, weight=-0.5, params={"factor": 2.0}), - "mask": RewardTermCfg(func=first_half_active, weight=3.0), - "scalar": RewardTermCfg(func=grilled_chicken, weight=10), + "scaled": RewardTermCfg(func=env_index_scaled, weight=-2.0, params={"factor": 1.0}), + "scalar": RewardTermCfg(func=grilled_chicken, weight=4.0), "disabled": RewardTermCfg(func=grilled_chicken_with_bbq, weight=0.0, params={"bbq": True}), } rew_man = RewardManager(cfg, env) - env_index = torch.arange(num_envs, dtype=torch.float64) - expected_sums = {name: torch.zeros(num_envs, dtype=torch.float64) for name in cfg} - for step in range(1, 4): + for step, expected in ((1, [1.0, 0.5, 0.0]), (2, [1.0, 0.0, -1.0])): env.step = step - expected_step = { - "scaled": -0.5 * 2.0 * step * env_index, - "mask": 3.0 * (env_index < num_envs // 2).double(), - "scalar": torch.full((num_envs,), 10.0, dtype=torch.float64), - "disabled": torch.zeros(num_envs, dtype=torch.float64), - } - rewards = rew_man.compute(dt=dt) - torch.testing.assert_close(rewards.cpu(), (sum(expected_step.values()) * dt).float()) - for index, name in enumerate(cfg): - expected_sums[name] += expected_step[name] * dt - torch.testing.assert_close(rew_man._step_reward[:, index].cpu(), expected_step[name].float()) - torch.testing.assert_close(rew_man._episode_sums[name].cpu(), expected_sums[name].float()) - - env_ids = torch.tensor([1, 6], device=device) - extras = rew_man.reset(env_ids) - for name in cfg: - value = extras["Episode_Reward/" + name] - assert value.ndim == 0 - expected = expected_sums[name][[1, 6]].mean() / max_episode_length_s - torch.testing.assert_close(value.cpu(), expected.float()) - # only the reset envs start a new episodic sum - expected_sums[name][[1, 6]] = 0.0 - torch.testing.assert_close(rew_man._episode_sums[name].cpu(), expected_sums[name].float()) + torch.testing.assert_close(rew_man.compute(dt=dt), torch.tensor(expected)) + + reset_env_2 = rew_man.reset(torch.tensor([2])) + assert {name: value.item() for name, value in reset_env_2.items()} == { + "Episode_Reward/scaled": -1.5, + "Episode_Reward/scalar": 1.0, + "Episode_Reward/disabled": 0.0, + } + reset_env_1 = rew_man.reset(torch.tensor([1])) + assert reset_env_1["Episode_Reward/scaled"].item() == -0.75 + assert reset_env_1["Episode_Reward/scalar"].item() == 1.0 + assert all(value.item() == 0.0 for value in rew_man.reset(torch.tensor([2])).values()) diff --git a/source/isaaclab/test/utils/test_circular_buffer.py b/source/isaaclab/test/utils/test_circular_buffer.py index d72289dd6d70..5b51f9fa5f9b 100644 --- a/source/isaaclab/test/utils/test_circular_buffer.py +++ b/source/isaaclab/test/utils/test_circular_buffer.py @@ -257,20 +257,18 @@ def test_stack_dim_positive_index_equivalent_to_negative(): torch.testing.assert_close(buf_neg.stacked, buf_pos.stacked) -@pytest.mark.parametrize( - "K,H,W", - [(2, 4, 4), (4, 4, 4), (4, 512, 256)], - ids=["two_slots", "small_frames", "large_frames"], -) -def test_stack_dim_ring_shift_after_overflow(K, H, W): +@pytest.mark.parametrize(("height", "width"), [(2, 3), (512, 256)], ids=["staged_shift", "large_frame_shift"]) +def test_stack_dim_ring_shift_after_overflow(height: int, width: int): """After K+1 frames, the oldest frame is evicted and the slots hold frames 1..K in order.""" - B, C = 2, 3 + B, C, K = 2, 3, 4 buf = CircularBuffer(max_len=K, batch_size=B, device="cpu", stack_dim=-1) for value in range(K + 1): - buf.append(torch.full((B, H, W, C), float(value))) + buf.append(torch.full((B, height, width, C), float(value))) stacked = buf.stacked for slot in range(K): - torch.testing.assert_close(stacked[..., slot * C : (slot + 1) * C], torch.full((B, H, W, C), slot + 1.0)) + torch.testing.assert_close( + stacked[..., slot * C : (slot + 1) * C], torch.full((B, height, width, C), slot + 1.0) + ) def test_stack_dim_reset_clears_buffer(): diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index cdb76b332c1f..4f6c86a54072 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -908,26 +908,16 @@ def test_full_reset_clears_active_flag(device: str, env_ids: slice | None): assert _get_wrench_without_pose_reads(composer)[2] is True -@pytest.mark.parametrize("device", test_devices()) -@pytest.mark.parametrize("selector", ["all", "env_ids", "env_mask"]) -def test_reset_of_inactive_composer_launches_no_work(device: str, selector: str): +def test_reset_of_inactive_composer_launches_no_work(): """Resetting an inactive composer, as asset writes do every physics step, is a host-only no-op.""" - num_envs, num_bodies = 4, 2 - composer = WrenchComposer(create_mock_asset(num_envs, num_bodies, device)) - selection = { - "all": {}, - "env_ids": {"env_ids": [1, 3]}, - "env_mask": {"env_mask": wp.array([False, True, False, True], dtype=wp.bool, device=device)}, - }[selector] + composer = WrenchComposer(create_mock_asset(2, 1, "cpu")) with patch.object(wp, "launch") as launch, patch.object(wp.array, "zero_") as zero: - composer.reset(**selection) + composer.reset(env_ids=[1]) launch.assert_not_called() zero.assert_not_called() assert not composer.active - assert not composer._dirty - np.testing.assert_array_equal(composer.out_force_b.warp.numpy(), np.zeros((num_envs, num_bodies, 3))) # ============================================================================ diff --git a/source/isaaclab_tasks/test/core/test_lift_env_cfg.py b/source/isaaclab_tasks/test/core/test_lift_env_cfg.py index 4f4bbc774e48..1479ca9c9e76 100644 --- a/source/isaaclab_tasks/test/core/test_lift_env_cfg.py +++ b/source/isaaclab_tasks/test/core/test_lift_env_cfg.py @@ -228,23 +228,17 @@ def test_lift_static_success_markers_update_only_marker_indices(monkeypatch: pyt num_envs = 3 static_table = object.__new__(Asset) static_table.cfg = SimpleNamespace(init_state=SimpleNamespace(pos=(0.5, 0.0, 0.0))) - command, root_pos_w = _make_pose_command(monkeypatch, num_envs, static_table) + command, _ = _make_pose_command(monkeypatch, num_envs, static_table) # env 1 is at the goal, the others are not command.pose_command_b[:, 0] = torch.tensor([1.0, 0.0, 1.0]) - for update_metrics in ( - ObjectUniformPoseCommand._update_metrics, - DeformableUniformPoseCommand._update_metrics, - CableUniformPoseCommand._update_metrics, - ): - command._segment_position_w = lambda: root_pos_w - update_metrics(command) + command._update_metrics() calls = command.success_visualizer.calls placement_args, placement_kwargs = calls[0] assert torch.equal(placement_args[0], torch.tensor([[0.5, 0.0, 0.0]] * num_envs)) assert torch.equal(placement_kwargs["environment_ids"], torch.arange(num_envs)) - assert len(calls) == 4 + assert len(calls) == 2 for args, kwargs in calls[1:]: assert args == () assert kwargs.keys() == {"marker_indices"} @@ -253,24 +247,21 @@ def test_lift_static_success_markers_update_only_marker_indices(monkeypatch: pyt def test_lift_contact_terms_match_per_sensor_reference() -> None: """Contact gating and counting from stacked sensor forces match a per-sensor evaluation.""" - num_envs, threshold = 64, 1.0 - names = ["thumb", "index", "middle", "ring"] - # contact forces [N] straddling the threshold, shape (num_envs, 1 body, 1 filter, 3) - forces = {name: 1.2 * torch.rand(num_envs, 1, 1, 3) for name in names} + num_envs, threshold = 3, 1.0 + names = ["thumb", "index", "middle"] + # The rows cover a finger without thumb, thumb with finger, and thumb alone. + magnitudes = {"thumb": [0.0, 2.0, 2.0], "index": [2.0, 0.0, 0.0], "middle": [0.0, 2.0, 0.0]} + forces = { + name: torch.tensor([[value, 0.0, 0.0] for value in values]).reshape(num_envs, 1, 1, 3) + for name, values in magnitudes.items() + } sensors = { name: SimpleNamespace(data=SimpleNamespace(normal_force_matrix_w=ProxyArray(wp.from_torch(force)))) for name, force in forces.items() } env = SimpleNamespace(num_envs=num_envs, device="cpu", scene=SimpleNamespace(sensors=sensors)) - in_contact = { - name: torch.linalg.norm(force.view(num_envs, 3), dim=-1) > threshold for name, force in forces.items() - } - - expected_contacts = in_contact["thumb"] & (in_contact["index"] | in_contact["middle"] | in_contact["ring"]) - expected_count = sum(in_contact[name].float() for name in names) / len(names) - - assert torch.equal(mdp.contacts(env, threshold, "thumb", names[1:]), expected_contacts) - assert torch.equal(mdp.contact_count(env, threshold, names), expected_count) + assert torch.equal(mdp.contacts(env, threshold, "thumb", names[1:]), torch.tensor([False, True, False])) + torch.testing.assert_close(mdp.contact_count(env, threshold, names), torch.tensor([1 / 3, 2 / 3, 1 / 3])) assert not mdp.contacts(env, threshold, "thumb", []).any() From ec656283e52f0772bfa63a32e491427305bcb5d5 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Tue, 29 Sep 2026 09:22:12 -0700 Subject: [PATCH 10/18] Address hot-path review feedback and consolidate tests --- .../changelog.d/perf-hot-path-7b-sensors.rst | 5 ++ .../isaaclab/managers/curriculum_manager.py | 4 +- .../isaaclab/managers/observation_manager.py | 16 +++-- .../sensors/ray_caster/base_ray_caster.py | 8 ++- .../envs/test_diffik_jacobian_aliasing.py | 4 +- .../test/envs/test_root_state_events.py | 6 +- .../test/envs/test_velocity_command_term.py | 6 +- .../test/managers/test_curriculum_manager.py | 7 +- .../test/managers/test_observation_manager.py | 5 ++ .../test/managers/test_reward_manager.py | 64 +++++++------------ .../test/managers/test_termination_manager.py | 4 +- .../markers/test_visualization_markers.py | 33 ++++------ .../test/sensors/test_ray_caster_sensor.py | 37 ++++++----- .../isaaclab/test/sensors/test_sensor_base.py | 33 +++------- .../test/utils/test_wrench_composer.py | 15 ++--- .../changelog.d/perf-hot-path-7b-sensors.rst | 3 + .../assets/articulation/articulation.py | 29 +++++---- .../test/assets/test_articulation.py | 19 +++++- .../changelog.d/perf-hot-path-7b-sensors.rst | 3 + .../assets/articulation/articulation.py | 15 +++-- .../test/assets/test_articulation.py | 19 +++++- .../perf-hot-path-7b-tasks.minor.rst | 4 +- .../core/lift/mdp/commands/pose_commands.py | 43 +++++++------ .../isaaclab_tasks/core/lift/mdp/events.py | 2 +- .../test/core/test_lift_env_cfg.py | 38 ++++------- .../test/core/test_reorient_utils.py | 11 +--- 26 files changed, 221 insertions(+), 212 deletions(-) diff --git a/source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst index 9136ed982c1b..80edccd2e498 100644 --- a/source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst +++ b/source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst @@ -9,3 +9,8 @@ Changed are zero, and uploaded the ray-cast drift ranges to the device only when they change. * Changed :meth:`~isaaclab.markers.VisualizationMarkers.visualize` to return without validating or processing inputs when no visualizer backend is active, such as in headless runs. + +Fixed +^^^^^ + +* Cleared previously sampled ray-caster drift on reset when its configured range was changed to zero. diff --git a/source/isaaclab/isaaclab/managers/curriculum_manager.py b/source/isaaclab/isaaclab/managers/curriculum_manager.py index bef9428024b4..13d3e346b5cb 100644 --- a/source/isaaclab/isaaclab/managers/curriculum_manager.py +++ b/source/isaaclab/isaaclab/managers/curriculum_manager.py @@ -162,12 +162,12 @@ def get_active_iterable_terms(self, env_idx: int) -> Sequence[tuple[str, Sequenc # each key is a separate state to log for key, value in term_state.items(): if isinstance(value, torch.Tensor): - value = value.detach().clone() + value = value.item() data.append(value) else: # log directly if not a dict if isinstance(term_state, torch.Tensor): - term_state = term_state.detach().clone() + term_state = term_state.item() data.append(term_state) terms.append((term_name, data)) diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index 692d91539d3e..04ae6b69c3e1 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -8,7 +8,7 @@ from __future__ import annotations import inspect -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import TYPE_CHECKING import numpy as np @@ -425,10 +425,10 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch. group_term_names = self._group_obs_term_names[group_name] group_obs = dict.fromkeys(group_term_names, None) term_cfgs = self._group_obs_term_cfgs[group_name] - obs_terms = zip(group_term_names, term_cfgs) + obs_terms = zip(group_term_names, term_cfgs, self._group_obs_term_safe_noise_funcs[group_name]) # evaluate terms: compute, add noise, clip, scale, custom modifiers - for term_name, term_cfg in obs_terms: + for term_name, term_cfg, safe_noise_func in obs_terms: output_spec = self._term_output_specs.get((group_name, term_name)) owned = output_spec is not None if output_spec is not None: @@ -451,7 +451,7 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch. else: obs = modifier.func(obs, **modifier.params) if isinstance(term_cfg.noise, noise.NoiseCfg): - if term_cfg.noise.func in _OUT_OF_PLACE_NOISE_FUNCS: + if term_cfg.noise.func is safe_noise_func: noisy_obs = term_cfg.noise.func(obs, term_cfg.noise) # a new result is owned; a returned input keeps its ownership owned = owned or noisy_obs is not obs @@ -565,6 +565,7 @@ def _prepare_terms(self): self._group_obs_term_names: dict[str, list[str]] = {} self._group_obs_term_dim: dict[str, list[tuple[int, ...]]] = {} self._group_obs_term_cfgs: dict[str, list[ObservationTermCfg]] = {} + self._group_obs_term_safe_noise_funcs: dict[str, list[Callable | str | None]] = {} self._group_obs_class_term_cfgs: dict[str, list[ObservationTermCfg]] = {} self._group_obs_concatenate: dict[str, bool] = {} self._group_obs_concatenate_dim: dict[str, int] = {} @@ -604,6 +605,7 @@ def _prepare_terms(self): self._group_obs_term_names[group_name] = [] self._group_obs_term_dim[group_name] = [] self._group_obs_term_cfgs[group_name] = [] + self._group_obs_term_safe_noise_funcs[group_name] = [] self._group_obs_class_term_cfgs[group_name] = [] # history buffers @@ -664,6 +666,12 @@ def _prepare_terms(self): # add term config to list self._group_obs_term_names[group_name].append(term_name) self._group_obs_term_cfgs[group_name].append(term_cfg) + # Keep custom noise isolated even if its function is replaced after setup. + self._group_obs_term_safe_noise_funcs[group_name].append( + term_cfg.noise.func + if isinstance(term_cfg.noise, noise.NoiseCfg) and term_cfg.noise.func in _OUT_OF_PLACE_NOISE_FUNCS + else None + ) # call function the first time to fill up dimensions obs_dims = self._prepare_term_output(group_name, term_name, term_cfg) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py index 1f703a000329..1a98a5c41fd9 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -69,6 +69,7 @@ def __init__(self, cfg: RayCasterCfg): BaseRayCaster._instance_count += 1 super().__init__(cfg) self._data = RayCasterData() + self._drift_sampled = False self._ray_cast_drift_range_list: tuple[tuple[float, float], ...] | None = None self._ray_cast_drift_ranges: torch.Tensor | None = None @@ -112,7 +113,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None tuple(self.cfg.ray_cast_drift_range.get(key, (0.0, 0.0))) for key in ("x", "y", "z") ) sample_ray_cast_drift = any(any(axis_range) for axis_range in ray_cast_range_list) - if not (sample_drift or sample_ray_cast_drift): + self._drift_sampled |= sample_drift or sample_ray_cast_drift + if not self._drift_sampled: return # determine the selected batch size if env_ids is not None: @@ -127,6 +129,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None if sample_drift: r = torch.empty(num_envs_ids, 3, device=self.device) self.drift.torch[env_ids] = r.uniform_(*self.cfg.drift_range) + else: + self.drift.torch[env_ids] = 0.0 # resample the ray cast drift if sample_ray_cast_drift: # Upload the ranges to the device only when the configuration changes. @@ -137,6 +141,8 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None self.ray_cast_drift.torch[env_ids] = math_utils.sample_uniform( ranges[:, 0], ranges[:, 1], (num_envs_ids, 3), device=self.device ) + else: + self.ray_cast_drift.torch[env_ids] = 0.0 """ Implementation. diff --git a/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py b/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py index 17d5b3d88e9e..0707cf582b6f 100644 --- a/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py +++ b/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py @@ -23,7 +23,7 @@ DifferentialInverseKinematicsAction, OperationalSpaceControllerAction, ) -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils import math as math_utils pytestmark = pytest.mark.unit @@ -124,7 +124,7 @@ def test_body_offset_jacobian_uses_offset_in_root_frame(compute, returns_jacobia torch.testing.assert_close(result, expected) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.parametrize(("ee_quat", "ik_offset"), [((0.0, 0.0, 0.0, 1.0), 1.0), ((0.0, 0.0, 0.0, 0.0), 0.0)]) def test_apply_actions_holds_joints_for_uninitialized_frame(device, ee_quat, ik_offset): """IK targets are applied for a valid frame pose and joints hold for all-zero quaternions, without a sync.""" diff --git a/source/isaaclab/test/envs/test_root_state_events.py b/source/isaaclab/test/envs/test_root_state_events.py index a2e77eaefd3e..8f51fc55aea5 100644 --- a/source/isaaclab/test/envs/test_root_state_events.py +++ b/source/isaaclab/test/envs/test_root_state_events.py @@ -13,7 +13,7 @@ from isaaclab.envs.mdp import push_by_setting_velocity, reset_root_state_uniform from isaaclab.managers import EventTermCfg, SceneEntityCfg -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices pytestmark = pytest.mark.unit @@ -62,7 +62,7 @@ def _make_env(device: str, num_envs: int = 3) -> tuple[SimpleNamespace, SimpleNa return env, asset -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_reset_root_state_uniform_uses_call_ranges(device): """The term samples from the ranges passed at call time, not the ones present at construction.""" env, asset = _make_env(device) @@ -85,7 +85,7 @@ def test_reset_root_state_uniform_uses_call_ranges(device): torch.testing.assert_close(asset.written["vel"], expected_vel) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_push_by_setting_velocity_follows_range_edits_without_sync(device): """Repeated ranges reuse device bounds without a sync, and in-place range edits still take effect.""" env, asset = _make_env(device) diff --git a/source/isaaclab/test/envs/test_velocity_command_term.py b/source/isaaclab/test/envs/test_velocity_command_term.py index 4802c42e707c..b24b8ef4fedf 100644 --- a/source/isaaclab/test/envs/test_velocity_command_term.py +++ b/source/isaaclab/test/envs/test_velocity_command_term.py @@ -14,7 +14,7 @@ import torch from isaaclab.envs.mdp import UniformVelocityCommand, UniformVelocityCommandCfg -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices pytestmark = pytest.mark.unit @@ -60,7 +60,7 @@ def _make_command(device: str, num_envs: int) -> tuple[UniformVelocityCommand, S return UniformVelocityCommand(cfg, env), robot_data -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_heading_control_updates_only_heading_envs(device): """Heading envs track the wrapped, clipped heading error without a sync; others keep their sampled yaw rate.""" heading_target = [3.0, 0.5, -3.0, 0.2, 1.0, -0.1] @@ -86,7 +86,7 @@ def test_heading_control_updates_only_heading_envs(device): torch.testing.assert_close(command.command[:, :2].cpu(), torch.full((len(heading), 2), 0.25)) -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_reset_logs_metrics_without_host_reads(device): """Reset logs episode metrics as device scalars without synchronizing the stream.""" command, robot_data = _make_command(device, 4) diff --git a/source/isaaclab/test/managers/test_curriculum_manager.py b/source/isaaclab/test/managers/test_curriculum_manager.py index 90dd34ac5ec7..051eb729aefc 100644 --- a/source/isaaclab/test/managers/test_curriculum_manager.py +++ b/source/isaaclab/test/managers/test_curriculum_manager.py @@ -12,7 +12,7 @@ import torch from isaaclab.managers import CurriculumManager, CurriculumTermCfg -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices pytestmark = pytest.mark.unit @@ -28,7 +28,7 @@ def advance_scalar(env, env_ids): return env.level -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_reset_logs_state_snapshots_without_host_reads(device): """Reset logs tensor states as device snapshots, so later in-place term updates do not change them.""" sim = MagicMock() @@ -60,3 +60,6 @@ def test_reset_logs_state_snapshots_without_host_reads(device): ): assert extras[key].device == torch.device(device) assert extras[key].item() == expected + + assert manager.get_active_iterable_terms(0) == [("dict_term", [2.0, 1.0, "stage"]), ("scalar_term", [2.0])] + assert isinstance(manager.get_active_iterable_terms(0)[1][1][0], float) diff --git a/source/isaaclab/test/managers/test_observation_manager.py b/source/isaaclab/test/managers/test_observation_manager.py index fba5862723a0..ac9780fe571e 100644 --- a/source/isaaclab/test/managers/test_observation_manager.py +++ b/source/isaaclab/test/managers/test_observation_manager.py @@ -401,6 +401,11 @@ def test_noise_preserves_source_and_rng_stream(setup_env, noise_cfg, expected_de torch.testing.assert_close(env.data.pos_w, source) torch.testing.assert_close(result, expected) + manager.cfg["policy"].position.noise.func = in_place_noise + result = manager.compute()["policy"] + torch.testing.assert_close(env.data.pos_w, source) + torch.testing.assert_close(result, (source + 0.2).clip(0.0, 0.5) * 2.0) + def test_compute_with_2d_history(setup_env): env = setup_env diff --git a/source/isaaclab/test/managers/test_reward_manager.py b/source/isaaclab/test/managers/test_reward_manager.py index c56028bd9256..e7d3f910521d 100644 --- a/source/isaaclab/test/managers/test_reward_manager.py +++ b/source/isaaclab/test/managers/test_reward_manager.py @@ -91,20 +91,33 @@ class MyRewardManagerAnnotatedCfg: assert rew_man_from_dict._term_cfgs == rew_man_from_cfg._term_cfgs -def test_compute(env): - """Test the computation of reward.""" +def test_compute(): + """A partial reset reports each term's accumulated reward and preserves other environments.""" + num_envs, dt, max_episode_length_s = 3, 0.25, 2.0 + sim = MagicMock() + sim.is_playing.return_value = False + env = SimpleNamespace(num_envs=num_envs, device="cpu", sim=sim, max_episode_length_s=max_episode_length_s, step=0) cfg = { - "term_1": RewardTermCfg(func=grilled_chicken, weight=10), - "term_2": RewardTermCfg(func=grilled_chicken_with_curry, weight=0.0, params={"hot": False}), + "scaled": RewardTermCfg(func=env_index_scaled, weight=-2.0, params={"factor": 1.0}), + "scalar": RewardTermCfg(func=grilled_chicken, weight=4.0), + "disabled": RewardTermCfg(func=grilled_chicken_with_bbq, weight=0.0, params={"bbq": True}), } rew_man = RewardManager(cfg, env) - # compute expected reward - expected_reward = cfg["term_1"].weight * env.dt - # compute reward using manager - rewards = rew_man.compute(dt=env.dt) - # check the reward for environment index 0 - assert float(rewards[0]) == expected_reward - assert tuple(rewards.shape) == (env.num_envs,) + + for step, expected in ((1, [1.0, 0.5, 0.0]), (2, [1.0, 0.0, -1.0])): + env.step = step + torch.testing.assert_close(rew_man.compute(dt=dt), torch.tensor(expected)) + + reset_env_2 = rew_man.reset(torch.tensor([2])) + assert {name: value.item() for name, value in reset_env_2.items()} == { + "Episode_Reward/scaled": -1.5, + "Episode_Reward/scalar": 1.0, + "Episode_Reward/disabled": 0.0, + } + reset_env_1 = rew_man.reset(torch.tensor([1])) + assert reset_env_1["Episode_Reward/scaled"].item() == -0.75 + assert reset_env_1["Episode_Reward/scalar"].item() == 1.0 + assert all(value.item() == 0.0 for value in rew_man.reset(torch.tensor([2])).values()) def test_config_empty(env): @@ -167,32 +180,3 @@ def test_invalid_reward_config(env): def env_index_scaled(env, factor: float): return factor * env.step * torch.arange(env.num_envs, dtype=torch.float, device=env.device) - - -def test_compute_and_reset_match_reference(): - """A partial reset reports each term's accumulated reward and preserves other environments.""" - num_envs, dt, max_episode_length_s = 3, 0.25, 2.0 - sim = MagicMock() - sim.is_playing.return_value = False - env = SimpleNamespace(num_envs=num_envs, device="cpu", sim=sim, max_episode_length_s=max_episode_length_s, step=0) - cfg = { - "scaled": RewardTermCfg(func=env_index_scaled, weight=-2.0, params={"factor": 1.0}), - "scalar": RewardTermCfg(func=grilled_chicken, weight=4.0), - "disabled": RewardTermCfg(func=grilled_chicken_with_bbq, weight=0.0, params={"bbq": True}), - } - rew_man = RewardManager(cfg, env) - - for step, expected in ((1, [1.0, 0.5, 0.0]), (2, [1.0, 0.0, -1.0])): - env.step = step - torch.testing.assert_close(rew_man.compute(dt=dt), torch.tensor(expected)) - - reset_env_2 = rew_man.reset(torch.tensor([2])) - assert {name: value.item() for name, value in reset_env_2.items()} == { - "Episode_Reward/scaled": -1.5, - "Episode_Reward/scalar": 1.0, - "Episode_Reward/disabled": 0.0, - } - reset_env_1 = rew_man.reset(torch.tensor([1])) - assert reset_env_1["Episode_Reward/scaled"].item() == -0.75 - assert reset_env_1["Episode_Reward/scalar"].item() == 1.0 - assert all(value.item() == 0.0 for value in rew_man.reset(torch.tensor([2])).values()) diff --git a/source/isaaclab/test/managers/test_termination_manager.py b/source/isaaclab/test/managers/test_termination_manager.py index 3f3d9c5ce00b..fb967e10225f 100644 --- a/source/isaaclab/test/managers/test_termination_manager.py +++ b/source/isaaclab/test/managers/test_termination_manager.py @@ -9,7 +9,7 @@ import torch from isaaclab.managers import TerminationManager, TerminationTermCfg -from isaaclab.test.utils import test_devices +from isaaclab.test.utils import DeviceScope, test_devices pytestmark = pytest.mark.unit @@ -128,7 +128,7 @@ def fail_first_quarter(env) -> torch.Tensor: return torch.arange(env.num_envs, device=env.device) < env.num_envs // 4 -@pytest.mark.parametrize("device", test_devices()) +@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) def test_reset_logs_term_rates_without_host_reads(device): """Reset logs each term's last-episode activation rate as a device scalar without synchronizing.""" sim = MagicMock() diff --git a/source/isaaclab/test/markers/test_visualization_markers.py b/source/isaaclab/test/markers/test_visualization_markers.py index 56ec605a65b3..d5bd320214b2 100644 --- a/source/isaaclab/test/markers/test_visualization_markers.py +++ b/source/isaaclab/test/markers/test_visualization_markers.py @@ -288,11 +288,21 @@ def test_usd_marker(sim): def test_visualization_skips_updates_when_invisible(sim): """When invisible, visualize should not update marker state.""" - sim._has_offscreen_render = True # create a marker config = clone(POSITION_GOAL_MARKER_CFG) config.prim_path = "/World/Visuals/test_protos" test_marker = VisualizationMarkers(config) + assert test_marker._backends == [] + initial_count = test_marker.count + # Unused inputs are ignored until a backend becomes active. + test_marker.visualize( + translations=torch.zeros(3, 3, device=sim.device), orientations=torch.zeros(2, 4, device=sim.device) + ) + assert test_marker.count == initial_count + sim._has_offscreen_render = True + test_marker.visualize(translations=torch.zeros(3, 3, device=sim.device)) + assert len(test_marker._backends) == 1 + assert test_marker.count == 3 # play the simulation sim.reset() @@ -319,27 +329,6 @@ def test_visualization_skips_updates_when_invisible(sim): assert test_marker.count == 4 -def test_visualize_without_backend_skips_until_backend_is_active(sim): - """Without an active marker backend, visualize is a no-op; a backend enabled later receives updates.""" - config = VisualizationMarkersCfg( - prim_path="/World/Visuals/headless_marker", - markers={"test": sim_utils.SphereCfg(radius=0.1)}, - ) - test_marker = VisualizationMarkers(config) - assert test_marker._backends == [] - - # Mismatched inputs are not validated because nothing consumes them. - test_marker.visualize( - translations=torch.zeros(3, 3, device=sim.device), orientations=torch.zeros(2, 4, device=sim.device) - ) - assert test_marker.count == 1 - - sim._has_offscreen_render = True - test_marker.visualize(translations=torch.zeros(3, 3, device=sim.device)) - assert len(test_marker._backends) == 1 - assert test_marker.count == 3 - - def test_newton_marker_backend_registers_and_updates_state_without_frame_capture(sim): """Newton marker backend state should be registered and ready for Newton-family viewers.""" sim._visualizers.append(_FakeMarkerVisualizer(pumps_app_update=False)) diff --git a/source/isaaclab/test/sensors/test_ray_caster_sensor.py b/source/isaaclab/test/sensors/test_ray_caster_sensor.py index 10b03830be1c..f170432de192 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_sensor.py +++ b/source/isaaclab/test/sensors/test_ray_caster_sensor.py @@ -181,19 +181,24 @@ def test_ray_caster_reset_resamples_drift(sim_ground): sim_utils.create_prim("/World/Sensor", "Xform", translation=(0.0, 0.0, 2.0)) cfg = _ray_caster_cfg("/World/Sensor", "world") - cfg.drift_range = (0.01, 0.05) # force non-zero drift - cfg.ray_cast_drift_range = {"x": (0.1, 0.2), "y": (-0.2, -0.1)} sensor = RayCaster(cfg) sim.reset() - # sim.reset() initializes the sensor with zero drift; call sensor.reset() to resample - # from the configured drift_range before we capture the baseline. + get_rng_state = torch.cuda.get_rng_state if torch.device(sensor.device).type == "cuda" else torch.get_rng_state + rng_state = get_rng_state() + sensor.reset() + assert torch.equal(get_rng_state(), rng_state) + assert (sensor.drift.torch == 0.0).all() + assert (sensor.ray_cast_drift.torch == 0.0).all() + + sensor.cfg.drift_range = (0.01, 0.05) + sensor.cfg.ray_cast_drift_range = {"x": (0.1, 0.2), "y": (-0.2, -0.1)} sensor.reset() dt = 0.01 sensor.update(dt) drift_before = sensor.drift.clone() # (1, 3) torch tensor - lo, hi = cfg.drift_range + lo, hi = sensor.cfg.drift_range # After sensor.reset(), drift should be within the configured range assert drift_before.shape == (1, 3), f"Drift shape should be (1, 3), got {drift_before.shape}" @@ -223,21 +228,15 @@ def test_ray_caster_reset_resamples_drift(sim_ground): assert ((ray_cast_drift[:, 1] >= -0.2) & (ray_cast_drift[:, 1] <= -0.1)).all() assert (ray_cast_drift[:, 2] == 0.0).all() - -@pytest.mark.isaacsim_ci -def test_ray_caster_reset_skips_zero_drift_sampling(sim_ground): - """reset() with the default zero drift ranges leaves drift at zero without consuming random numbers.""" - sim = sim_ground - - sim_utils.create_prim("/World/Sensor", "Xform", translation=(0.0, 0.0, 2.0)) - sensor = RayCaster(_ray_caster_cfg("/World/Sensor", "world")) - sim.reset() - - get_rng_state = torch.cuda.get_rng_state if torch.device(sensor.device).type == "cuda" else torch.get_rng_state - rng_state = get_rng_state() - sensor.reset() - assert torch.equal(get_rng_state(), rng_state), "reset() must not sample drift when all drift ranges are zero" + # Disabling either drift clears its old sample. + sensor.cfg.drift_range = (0.0, 0.0) + sensor.reset(env_ids=[0]) assert (sensor.drift.torch == 0.0).all() + assert (sensor.ray_cast_drift.torch[:, 0] >= 0.1).all() + sensor.cfg.ray_cast_drift_range = {} + rng_state = get_rng_state() + sensor.reset(env_ids=[0]) + assert torch.equal(get_rng_state(), rng_state) assert (sensor.ray_cast_drift.torch == 0.0).all() diff --git a/source/isaaclab/test/sensors/test_sensor_base.py b/source/isaaclab/test/sensors/test_sensor_base.py index c16f988de967..9a229dccdd00 100644 --- a/source/isaaclab/test/sensors/test_sensor_base.py +++ b/source/isaaclab/test/sensors/test_sensor_base.py @@ -186,7 +186,8 @@ def test_sensor_update_rate(create_dummy_sensor, device): """ sensor_cfg, sim, dt = create_dummy_sensor sensor_cfg.update_period = 2 * dt - sensor = DummySensor(cfg=sensor_cfg) + sensor_cfg.debug_vis = True + sensor = DummyVisSensor(cfg=sensor_cfg) # Play sim sim.step() @@ -206,6 +207,14 @@ def test_sensor_update_rate(create_dummy_sensor, device): ) expected_value += i % 2 + backend_update_count = sensor.backend_update_count + for _ in range(3): + sensor.update(dt=dt) + assert sensor.backend_update_count == backend_update_count + sim.vis_marker_registry.dispatch_callbacks() + assert sensor.backend_update_count == backend_update_count + 1 + torch.testing.assert_close(sensor.visualized_counts[-1], sensor.data.count) + @pytest.mark.parametrize("device", test_devices(DeviceScope.CPU)) def test_sensor_reset(create_dummy_sensor, device): @@ -315,28 +324,6 @@ def test_repeated_data_reads_are_graph_safe(create_dummy_sensor, device): wp.capture_launch(capture.graph) -@pytest.mark.parametrize("device", test_devices(DeviceScope.DEFAULT_CUDA)) -def test_debug_vis_refreshes_lazily_from_callback(create_dummy_sensor, device): - """Debug visualization does not force per-step refreshes; its callback sees current data.""" - sensor_cfg, sim, dt = create_dummy_sensor - sensor_cfg.debug_vis = True - sensor = DummyVisSensor(cfg=sensor_cfg) - sim.step() - sim.reset() - assert sensor._debug_vis_handle is not None - - backend_update_count = sensor.backend_update_count - for _ in range(3): - sensor.update(dt=dt) - assert sensor.backend_update_count == backend_update_count - - sim.vis_marker_registry.dispatch_callbacks() - assert sensor.backend_update_count == backend_update_count + 1 - torch.testing.assert_close( - sensor.visualized_counts[-1], torch.ones(sensor.num_instances, dtype=torch.int32, device=device) - ) - - @pytest.mark.parametrize("device", ("cpu",)) def test_rigid_body_ancestor_expr_trims_only_terminal_suffix(create_dummy_sensor, device): """Test that ancestor expression trimming keeps repeated path segments above the sensor.""" diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index 4f6c86a54072..e7fed7cf8880 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -902,16 +902,6 @@ def test_full_reset_clears_active_flag(device: str, env_ids: slice | None): assert is_global is False np.testing.assert_array_equal(force.numpy(), np.zeros_like(forces_np)) np.testing.assert_array_equal(torque.numpy(), np.zeros_like(forces_np)) - composer.add_forces_and_torques_index( - forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), is_global=True - ) - assert _get_wrench_without_pose_reads(composer)[2] is True - - -def test_reset_of_inactive_composer_launches_no_work(): - """Resetting an inactive composer, as asset writes do every physics step, is a host-only no-op.""" - composer = WrenchComposer(create_mock_asset(2, 1, "cpu")) - with patch.object(wp, "launch") as launch, patch.object(wp.array, "zero_") as zero: composer.reset(env_ids=[1]) @@ -919,6 +909,11 @@ def test_reset_of_inactive_composer_launches_no_work(): zero.assert_not_called() assert not composer.active + composer.add_forces_and_torques_index( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), is_global=True + ) + assert _get_wrench_without_pose_reads(composer)[2] is True + # ============================================================================ # Deprecated API Backward-Compatibility Tests diff --git a/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst index 9c78cc792aed..68a7df5d784b 100644 --- a/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst +++ b/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst @@ -5,3 +5,6 @@ Changed buffers before drawing. * Changed the OVPhysX rigid-object ``body_com_acc_w`` finite difference to use the elapsed time since the previous update, matching joint accelerations and the Newton backend. + +* Skipped the joint-limit clamping counter readback when its logging level is disabled and reused + the counter buffer across writes. diff --git a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py index 8fc7852c6875..a1af330993c1 100644 --- a/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py +++ b/source/isaaclab_ov/isaaclab_ov/assets/articulation/articulation.py @@ -1518,7 +1518,10 @@ def write_joint_position_limit_to_sim_index( device=self._device, ) # Clamp default_joint_pos to the new limits and refresh soft_joint_pos_limits. - clamped_count = wp.zeros(1, dtype=wp.int32, device=self._device) + log_level = logging.WARNING if warn_limit_violation else logging.INFO + report_clamping = logger.isEnabledFor(log_level) + if report_clamping: + self._clamped_default_count.zero_() wp.launch( clamp_default_joint_pos_and_update_soft_limits_index_kernel(env_ids, joint_ids), dim=(env_ids.shape[0], joint_ids.shape[0]), @@ -1531,19 +1534,16 @@ def write_joint_position_limit_to_sim_index( outputs=[ self._data._default_joint_pos, self._data._soft_joint_pos_limits, - clamped_count, + self._clamped_default_count, ], device=self._device, ) - if clamped_count.numpy()[0] > 0: + if report_clamping and self._clamped_default_count.numpy()[0] > 0: violation_message = ( "Some default joint positions are outside of the range of the new joint limits. Default joint" " positions will be clamped to be within the new joint limits." ) - if warn_limit_violation: - logger.warning(violation_message) - else: - logger.info(violation_message) + logger.log(log_level, violation_message) # Stage to pinned-host CPU: flatten the vec2f buffer to float32 view. self._push_joint_property( TT.DOF_LIMIT, @@ -1613,7 +1613,10 @@ def write_joint_position_limit_to_sim_mask( device=self._device, ) # Clamp default_joint_pos to the new limits and refresh soft_joint_pos_limits. - clamped_count = wp.zeros(1, dtype=wp.int32, device=self._device) + log_level = logging.WARNING if warn_limit_violation else logging.INFO + report_clamping = logger.isEnabledFor(log_level) + if report_clamping: + self._clamped_default_count.zero_() wp.launch( clamp_default_joint_pos_and_update_soft_limits_mask, dim=(self._num_instances, self._num_joints), @@ -1626,19 +1629,16 @@ def write_joint_position_limit_to_sim_mask( outputs=[ self._data._default_joint_pos, self._data._soft_joint_pos_limits, - clamped_count, + self._clamped_default_count, ], device=self._device, ) - if clamped_count.numpy()[0] > 0: + if report_clamping and self._clamped_default_count.numpy()[0] > 0: violation_message = ( "Some default joint positions are outside of the range of the new joint limits. Default joint" " positions will be clamped to be within the new joint limits." ) - if warn_limit_violation: - logger.warning(violation_message) - else: - logger.info(violation_message) + logger.log(log_level, violation_message) self._push_joint_property( TT.DOF_LIMIT, self._data._joint_pos_limits.data, @@ -4080,6 +4080,7 @@ def _create_buffers(self) -> None: self._ALL_INDICES = wp.array(np.arange(N, dtype=np.int32), device=device) self._ALL_BODY_INDICES = wp.array(np.arange(B, dtype=np.int32), device=device) self._ALL_JOINT_INDICES = wp.array(np.arange(J, dtype=np.int32), device=device) + self._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=self._device) self._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(FT, dtype=np.int32), device=device) self._ALL_SPATIAL_TENDON_INDICES = wp.array(np.arange(ST, dtype=np.int32), device=device) self._sim_env_ids = wp.empty(N, dtype=wp.int32, device=device) diff --git a/source/isaaclab_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index f075d4462872..6f477f924b3b 100644 --- a/source/isaaclab_ov/test/assets/test_articulation.py +++ b/source/isaaclab_ov/test/assets/test_articulation.py @@ -51,6 +51,7 @@ from __future__ import annotations import importlib +import logging import sys from pathlib import Path from types import SimpleNamespace @@ -1665,7 +1666,7 @@ def test_out_of_range_default_joint_vel(sim, device): @pytest.mark.parametrize("num_articulations", [2]) @pytest.mark.parametrize("device", test_devices()) @pytest.mark.parametrize("add_ground_plane", [True]) -def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): +def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, caplog): """Test write_joint_limits_to_sim API and when default pos falls outside of the new limits. This test verifies that: @@ -1736,6 +1737,22 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): ) assert torch.all(within_bounds) + # Logging level changes must not change clamping or reuse an old violation count. + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 1] = 0.5 + articulation_logger = type(articulation).__module__ + for level in (logging.WARNING, logging.INFO): + articulation.data.default_joint_pos.torch.fill_(1.0) + caplog.clear() + with caplog.at_level(level, logger=articulation_logger): + articulation.write_joint_position_limit_to_sim_index(limits=limits, warn_limit_violation=False) + assert bool(caplog.records) == (level == logging.INFO) + torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(limits[..., 1], 0.5)) + caplog.clear() + with caplog.at_level(logging.INFO, logger=articulation_logger): + articulation.write_joint_position_limit_to_sim_mask(limits=limits, warn_limit_violation=False) + assert not caplog.records + @pytest.mark.parametrize("num_articulations", [2]) @pytest.mark.parametrize("device", test_devices()) diff --git a/source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst b/source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst index f52bb69a5d5e..4f16c7a5ac44 100644 --- a/source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst +++ b/source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst @@ -5,3 +5,6 @@ Changed simple-shading outputs in :class:`~isaaclab_physx.renderers.IsaacRtxRenderer`. * Changed PhysX contact, PVA, and frame-transformer debug visualization to refresh outdated sensor buffers before drawing. + +* Skipped the joint-limit clamping counter readback when its logging level is disabled and reused + the counter buffer across writes. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 27c87503269f..c023134d3a95 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -1603,7 +1603,10 @@ def write_joint_position_limit_to_sim_index( if env_ids.shape[0] == 0 or joint_ids.shape[0] == 0: return - clamped_defaults = wp.zeros(1, dtype=wp.int32, device=self.device) + log_level = logging.WARNING if warn_limit_violation else logging.INFO + report_clamping = logger.isEnabledFor(log_level) + if report_clamping: + self._clamped_default_count.zero_() sim_env_ids = self._sim_env_ids_view(env_ids.shape[0]) # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. # Note: we are doing a single launch for faster performance. Prior versions would do this in multiple launches. @@ -1623,21 +1626,18 @@ def write_joint_position_limit_to_sim_index( self.data._joint_pos_limits, self.data._soft_joint_pos_limits, self.data._default_joint_pos, - clamped_defaults, + self._clamped_default_count, sim_env_ids, ], device=self.device, ) # Log a warning if the default joint positions are outside of the new limits. - if clamped_defaults.numpy()[0] > 0: + if report_clamping and self._clamped_default_count.numpy()[0] > 0: violation_message = ( "Some default joint positions are outside of the range of the new joint limits. Default joint positions" " will be clamped to be within the new joint limits." ) - if warn_limit_violation: - logger.warning(violation_message) - else: - logger.info(violation_message) + logger.log(log_level, violation_message) # Set into simulation, note that when updating "model" properties with PhysX we need to do it on CPU. cpu_env_ids = self._get_cpu_env_ids(env_ids, sim_env_ids) joint_pos_limits_backend = self._get_backend_ordered_joint_buffer( @@ -3971,6 +3971,7 @@ def _resolve_joint_dof_signs(self) -> tuple[int, ...]: def _create_buffers(self): self._ALL_INDICES = wp.array(np.arange(self.num_instances, dtype=np.int32), device=self.device) self._ALL_JOINT_INDICES = wp.array(np.arange(self.num_joints, dtype=np.int32), device=self.device) + self._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=self.device) self._ALL_BODY_INDICES = wp.array(np.arange(self.num_bodies, dtype=np.int32), device=self.device) self._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(self.num_fixed_tendons, dtype=np.int32), device=self.device) self._ALL_SPATIAL_TENDON_INDICES = wp.array( diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 7078c528911e..e0a81ef3aa97 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -22,6 +22,7 @@ launch_test_simulation() +import logging import sys from pathlib import Path @@ -879,7 +880,7 @@ def test_out_of_range_default_joint_vel(sim, device): @pytest.mark.parametrize("num_articulations", [2]) @pytest.mark.parametrize("device", test_devices()) @pytest.mark.parametrize("add_ground_plane", [True]) -def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): +def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane, caplog): """Test write_joint_limits_to_sim API and when default pos falls outside of the new limits. This test verifies that: @@ -950,6 +951,22 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): ) assert torch.all(within_bounds) + # Logging level changes must not change clamping or reuse an old violation count. + limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) + limits[..., 1] = 0.5 + articulation_logger = type(articulation).__module__ + for level in (logging.WARNING, logging.INFO): + articulation.data.default_joint_pos.torch.fill_(1.0) + caplog.clear() + with caplog.at_level(level, logger=articulation_logger): + articulation.write_joint_position_limit_to_sim_index(limits=limits, warn_limit_violation=False) + assert bool(caplog.records) == (level == logging.INFO) + torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(limits[..., 1], 0.5)) + caplog.clear() + with caplog.at_level(logging.INFO, logger=articulation_logger): + articulation.write_joint_position_limit_to_sim_mask(limits=limits, warn_limit_violation=False) + assert not caplog.records + @pytest.mark.parametrize("num_articulations", [2]) @pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) diff --git a/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst b/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst index e7c34b2b55fb..d1ed38a9dde0 100644 --- a/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst +++ b/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst @@ -21,6 +21,6 @@ Changed :func:`~isaaclab_tasks.core.lift.mdp.contact_count` take one norm per call instead of one per sensor. * Precomputed the ANYmal symmetry augmentation as one cached column permutation and sign tensor per transform, removing the per-minibatch index and sign uploads (about 1.5 ms to 0.15 ms per call). -* Sent only the marker indices to Lift success markers on static assets each step, cached the locomotion - walk-target offset and the Lift reset position bounds on the device, and read the Reorient goal directly +* Cached the locomotion walk-target offset and the Lift reset position bounds on the device, + and read the Reorient goal directly from the command buffers instead of concatenating the command on every access. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/commands/pose_commands.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/commands/pose_commands.py index 1143f5194456..da05e6b61685 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/commands/pose_commands.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/commands/pose_commands.py @@ -135,7 +135,12 @@ def _update_metrics(self): if not self.cfg.position_only: self.metrics["orientation_error"] = torch.linalg.norm(rot_error, dim=-1) success_id &= self.metrics["orientation_error"] < 0.5 - self._visualize_success(success_id) + if self.success_vis_asset is not None: + self.success_visualizer.visualize( + self._get_success_vis_pos_w(), + marker_indices=success_id.int(), + environment_ids=self._env.scene._ALL_INDICES, + ) def _get_success_vis_pos_w(self) -> torch.Tensor: """Return the success visualization positions in the world frame.""" @@ -143,24 +148,6 @@ def _get_success_vis_pos_w(self) -> torch.Tensor: return self._static_success_vis_pos_w return self.success_vis_asset.data.root_pos_w.torch - def _visualize_success(self, success: torch.Tensor) -> None: - """Show each environment's success marker. - - Args: - success: Per-environment success flags, shape ``(num_envs,)``. - """ - if self.success_vis_asset is None: - return - if self._static_success_vis_pos_w is not None: - # positions and environment ids were set at construction; resending them copies to the host - self.success_visualizer.visualize(marker_indices=success.int()) - return - self.success_visualizer.visualize( - self._get_success_vis_pos_w(), - marker_indices=success.int(), - environment_ids=self._env.scene._ALL_INDICES, - ) - def _resample_command(self, env_ids: Sequence[int]): # sample new pose targets # -- position @@ -261,8 +248,15 @@ def _update_metrics(self): com_w = self.object.data.root_pos_w.torch self.metrics["position_error"] = torch.linalg.norm(self.pose_command_w[:, :3] - com_w, dim=-1) + if self.success_vis_asset is None: + return # same success radius as the goal markers of the base class - self._visualize_success(self.metrics["position_error"] < 0.05) + success_id = (self.metrics["position_error"] < 0.05).int() + self.success_visualizer.visualize( + self._get_success_vis_pos_w(), + marker_indices=success_id, + environment_ids=self._env.scene._ALL_INDICES, + ) class CableUniformPoseCommand(ObjectUniformPoseCommand): @@ -294,7 +288,14 @@ def _update_metrics(self): segment_pos_w = self._segment_position_w() self.metrics["position_error"] = torch.linalg.norm(self.pose_command_w[:, :3] - segment_pos_w, dim=-1) - self._visualize_success(self.metrics["position_error"] < 0.05) + if self.success_vis_asset is None: + return + success_id = (self.metrics["position_error"] < 0.05).int() + self.success_visualizer.visualize( + self._get_success_vis_pos_w(), + marker_indices=success_id, + environment_ids=self._env.scene._ALL_INDICES, + ) def _debug_vis_callback(self, event): if not self.robot.is_initialized: diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py index 29743f6aecc1..aeef012499d7 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py @@ -82,7 +82,7 @@ def _xyz_ranges(ranges: dict[str, tuple[float, float]], device: str | torch.devi return _cached_xyz_ranges(bounds, str(device)) -@functools.cache +@functools.lru_cache(maxsize=128) def _cached_xyz_ranges(bounds: tuple[tuple[float, float], ...], device: str) -> torch.Tensor: return torch.tensor(bounds, device=device) diff --git a/source/isaaclab_tasks/test/core/test_lift_env_cfg.py b/source/isaaclab_tasks/test/core/test_lift_env_cfg.py index 1479ca9c9e76..b6c867a592fd 100644 --- a/source/isaaclab_tasks/test/core/test_lift_env_cfg.py +++ b/source/isaaclab_tasks/test/core/test_lift_env_cfg.py @@ -196,24 +196,33 @@ def _initialize_command_term(command, command_cfg, command_env) -> None: return ObjectUniformPoseCommand(cfg, env), root_pos_w -def test_lift_pose_markers_forward_environment_ids(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("static", [False, True]) +def test_lift_pose_markers_forward_environment_ids(monkeypatch: pytest.MonkeyPatch, static: bool) -> None: """Every lift pose, goal, and success marker should retain its environment ownership.""" num_envs = 3 environment_ids = torch.arange(num_envs) success_asset = SimpleNamespace(data=SimpleNamespace(root_pos_w=SimpleNamespace(torch=torch.zeros((num_envs, 3))))) + if static: + success_asset = object.__new__(Asset) + success_asset.cfg = SimpleNamespace(init_state=SimpleNamespace(pos=(0.5, 0.0, 0.0))) command, root_pos_w = _make_pose_command(monkeypatch, num_envs, success_asset) command._set_debug_vis_impl(True) command._debug_vis_callback(None) command.cfg.position_only = False command._debug_vis_callback(None) + # Discard construction-time placement, as happens before a backend is active. + command.success_visualizer.calls.clear() + command.cfg.position_only = True + command.pose_command_b[:, 0] = torch.tensor([1.0, 0.0, 1.0]) command._update_metrics() + assert torch.equal(command.success_visualizer.calls[0][1]["marker_indices"], torch.tensor([0, 1, 0])) DeformableUniformPoseCommand._update_metrics(command) command._segment_position_w = lambda: root_pos_w CableUniformPoseCommand._update_metrics(command) CableUniformPoseCommand._debug_vis_callback(command, None) expected_call_counts = { - command.success_visualizer: 4, + command.success_visualizer: 3, command.goal_visualizer: 3, command.curr_visualizer: 3, } @@ -221,28 +230,9 @@ def test_lift_pose_markers_forward_environment_ids(monkeypatch: pytest.MonkeyPat assert len(visualizer.calls) == expected_count for _, kwargs in visualizer.calls: assert torch.equal(kwargs["environment_ids"], environment_ids) - - -def test_lift_static_success_markers_update_only_marker_indices(monkeypatch: pytest.MonkeyPatch) -> None: - """Success markers on a static asset are placed once; per-step updates only switch the prototype.""" - num_envs = 3 - static_table = object.__new__(Asset) - static_table.cfg = SimpleNamespace(init_state=SimpleNamespace(pos=(0.5, 0.0, 0.0))) - command, _ = _make_pose_command(monkeypatch, num_envs, static_table) - # env 1 is at the goal, the others are not - command.pose_command_b[:, 0] = torch.tensor([1.0, 0.0, 1.0]) - - command._update_metrics() - - calls = command.success_visualizer.calls - placement_args, placement_kwargs = calls[0] - assert torch.equal(placement_args[0], torch.tensor([[0.5, 0.0, 0.0]] * num_envs)) - assert torch.equal(placement_kwargs["environment_ids"], torch.arange(num_envs)) - assert len(calls) == 2 - for args, kwargs in calls[1:]: - assert args == () - assert kwargs.keys() == {"marker_indices"} - assert torch.equal(kwargs["marker_indices"], torch.tensor([0, 1, 0], dtype=torch.int32)) + if static: + for args, _ in command.success_visualizer.calls: + torch.testing.assert_close(args[0], torch.tensor([[0.5, 0.0, 0.0]] * num_envs)) def test_lift_contact_terms_match_per_sensor_reference() -> None: diff --git a/source/isaaclab_tasks/test/core/test_reorient_utils.py b/source/isaaclab_tasks/test/core/test_reorient_utils.py index 0db609343021..f13d6ecc7ab6 100644 --- a/source/isaaclab_tasks/test/core/test_reorient_utils.py +++ b/source/isaaclab_tasks/test/core/test_reorient_utils.py @@ -62,15 +62,10 @@ def test_episode_error_recorder_reports_threshold_independent_statistics(): def test_episode_error_recorder_skips_episodes_without_samples(): - """Verify initial resets do not emit non-finite diagnostic values.""" - recorder = reorient_utils.EpisodeErrorRecorder(num_envs=2, device="cpu") - - assert recorder.reset(torch.tensor([0, 1])) == {} - - -def test_episode_error_recorder_reset_excludes_episodes_without_samples(): - """Verify reset statistics cover only sampled episodes, and are NaN when none of them is.""" + """Initial resets are empty; later statistics exclude unsampled episodes.""" recorder = reorient_utils.EpisodeErrorRecorder(num_envs=5, device="cpu") + assert recorder.reset(torch.arange(5)) == {} + nan = float("nan") # envs 1 and 4 only see non-finite errors, so they have no sample recorder.update(torch.tensor([0.3, nan, 0.1, 0.4, torch.inf])) From a71c9af888f5e4871330ec1f273fae70480191ff Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Tue, 29 Sep 2026 10:06:38 -0700 Subject: [PATCH 11/18] Centralize sampling from named event ranges --- ...st => perf-hot-path-7b-managers.minor.rst} | 7 ++ source/isaaclab/isaaclab/envs/mdp/events.py | 75 ++++++------------- source/isaaclab/isaaclab/utils/math.py | 44 +++++++++++ .../test/envs/test_root_state_events.py | 15 +++- .../perf-hot-path-7b-tasks.minor.rst | 3 +- .../isaaclab_tasks/core/lift/mdp/events.py | 23 +----- 6 files changed, 90 insertions(+), 77 deletions(-) rename source/isaaclab/changelog.d/{perf-hot-path-7b-managers.rst => perf-hot-path-7b-managers.minor.rst} (85%) diff --git a/source/isaaclab/changelog.d/perf-hot-path-7b-managers.rst b/source/isaaclab/changelog.d/perf-hot-path-7b-managers.minor.rst similarity index 85% rename from source/isaaclab/changelog.d/perf-hot-path-7b-managers.rst rename to source/isaaclab/changelog.d/perf-hot-path-7b-managers.minor.rst index 973b1dd38895..9625e65a8249 100644 --- a/source/isaaclab/changelog.d/perf-hot-path-7b-managers.rst +++ b/source/isaaclab/changelog.d/perf-hot-path-7b-managers.minor.rst @@ -1,3 +1,10 @@ +Added +^^^^^ + +* Added :func:`~isaaclab.utils.math.sample_uniform_from_ranges` to sample named components with + shared, bounded caching of device bounds. Core and Lift event terms now use this sampler instead + of maintaining separate tensor-cache helpers. + Changed ^^^^^^^ diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 59fc271abaf7..59309db504e7 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -14,9 +14,7 @@ from __future__ import annotations -import functools import logging -from collections.abc import Sequence from types import ModuleType from typing import TYPE_CHECKING, Literal @@ -40,29 +38,6 @@ logger = logging.getLogger(__name__) -@functools.lru_cache(maxsize=128) -def _cached_range_tensor(values: tuple[tuple[float, float], ...], device: str) -> torch.Tensor: - return torch.tensor(values, dtype=torch.float32, device=device) - - -def _range_tensor(ranges: dict[str, tuple[float, float]], keys: Sequence[str], device: str) -> torch.Tensor: - """Return the ``(min, max)`` bounds of ``keys`` as a device tensor of shape ``(len(keys), 2)``. - - The tensor is cached by value, so repeated calls avoid a host-to-device copy while range edits (for - example by a curriculum) still take effect. Callers must not modify the returned tensor. - - Args: - ranges: The ``(min, max)`` range for each key. Missing keys use ``(0.0, 0.0)``. - keys: The keys to read, in output row order. - device: The device of the returned tensor. - - Returns: - The range bounds. Shape is ``(len(keys), 2)``. - """ - values = tuple(tuple(float(v) for v in ranges.get(key, (0.0, 0.0))) for key in keys) - return _cached_range_tensor(values, str(device)) - - def randomize_rigid_body_scale( env: ManagerBasedEnv, env_ids: torch.Tensor | slice | None, @@ -547,10 +522,9 @@ def __call__( env_rows = env_ids[:, None] if not isinstance(env_ids, slice) and not isinstance(body_ids, slice) else env_ids # sample random CoM values - ranges = _range_tensor(com_range, ("x", "y", "z"), self.asset.device) num_envs = len(range(env.num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) - rand_samples = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (num_envs, 3), device=self.asset.device + rand_samples = math_utils.sample_uniform_from_ranges( + com_range, ("x", "y", "z"), num_envs, device=self.asset.device ).unsqueeze(1) # start from defaults and add random offsets @@ -1281,8 +1255,9 @@ def push_by_setting_velocity( # velocities vel_w = asset.data.root_vel_w.torch[env_ids] # sample random velocities - ranges = _range_tensor(velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) - vel_w += math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], vel_w.shape, device=asset.device) + vel_w += math_utils.sample_uniform_from_ranges( + velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), vel_w.shape[:-1], device=asset.device + ) # set the velocities into the physics simulation asset.write_root_velocity_to_sim_index(root_velocity=vel_w, env_ids=env_ids) @@ -1315,18 +1290,16 @@ def __call__( default_root_pose = asset.data.default_root_pose.torch[env_ids] default_root_vel = asset.data.default_root_vel.torch[env_ids] - ranges = _range_tensor(pose_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) - rand_samples = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (default_root_pose.shape[0], 6), device=asset.device + rand_samples = math_utils.sample_uniform_from_ranges( + pose_range, ("x", "y", "z", "roll", "pitch", "yaw"), default_root_pose.shape[0], device=asset.device ) positions = default_root_pose[:, 0:3] + env.scene.env_origins[env_ids] + rand_samples[:, 0:3] orientations_delta = math_utils.quat_from_euler_xyz(rand_samples[:, 3], rand_samples[:, 4], rand_samples[:, 5]) orientations = math_utils.quat_mul(default_root_pose[:, 3:7], orientations_delta) # velocities - ranges = _range_tensor(velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) - rand_samples = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (default_root_pose.shape[0], 6), device=asset.device + rand_samples = math_utils.sample_uniform_from_ranges( + velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), default_root_pose.shape[0], device=asset.device ) velocities = default_root_vel + rand_samples @@ -1367,18 +1340,16 @@ def reset_root_state_with_random_orientation( default_root_vel = asset.data.default_root_vel.torch[env_ids].clone() # poses - ranges = _range_tensor(pose_range, ("x", "y", "z"), asset.device) - rand_samples = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (default_root_pose.shape[0], 3), device=asset.device + rand_samples = math_utils.sample_uniform_from_ranges( + pose_range, ("x", "y", "z"), default_root_pose.shape[0], device=asset.device ) positions = default_root_pose[:, 0:3] + env.scene.env_origins[env_ids] + rand_samples orientations = math_utils.random_orientation(default_root_pose.shape[0], device=asset.device) # velocities - ranges = _range_tensor(velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) - rand_samples = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (default_root_pose.shape[0], 6), device=asset.device + rand_samples = math_utils.sample_uniform_from_ranges( + velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), default_root_pose.shape[0], device=asset.device ) velocities = default_root_vel + rand_samples @@ -1435,15 +1406,17 @@ def reset_root_state_from_terrain( positions = valid_positions[terrain.terrain_levels[env_ids], terrain.terrain_types[env_ids], ids] positions += asset.data.default_root_pose.torch[env_ids, :3] # sample random orientations - ranges = _range_tensor(pose_range, ("roll", "pitch", "yaw"), asset.device) - rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (num_envs, 3), device=asset.device) + rand_samples = math_utils.sample_uniform_from_ranges( + pose_range, ("roll", "pitch", "yaw"), num_envs, device=asset.device + ) # convert to quaternions orientations = math_utils.quat_from_euler_xyz(rand_samples[:, 0], rand_samples[:, 1], rand_samples[:, 2]) # sample random velocities - ranges = _range_tensor(velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), asset.device) - rand_samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (num_envs, 6), device=asset.device) + rand_samples = math_utils.sample_uniform_from_ranges( + velocity_range, ("x", "y", "z", "roll", "pitch", "yaw"), num_envs, device=asset.device + ) velocities = asset.data.default_root_vel.torch[env_ids] + rand_samples @@ -1706,17 +1679,15 @@ def reset_nodal_state_uniform( nodal_state = asset.data.default_nodal_state_w.torch[env_ids].clone() # position - ranges = _range_tensor(position_range, ("x", "y", "z"), asset.device) - rand_samples = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (nodal_state.shape[0], 1, 3), device=asset.device + rand_samples = math_utils.sample_uniform_from_ranges( + position_range, ("x", "y", "z"), (nodal_state.shape[0], 1), device=asset.device ) nodal_state[..., :3] += rand_samples # velocities - ranges = _range_tensor(velocity_range, ("x", "y", "z"), asset.device) - rand_samples = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (nodal_state.shape[0], 1, 3), device=asset.device + rand_samples = math_utils.sample_uniform_from_ranges( + velocity_range, ("x", "y", "z"), (nodal_state.shape[0], 1), device=asset.device ) nodal_state[..., 3:] += rand_samples diff --git a/source/isaaclab/isaaclab/utils/math.py b/source/isaaclab/isaaclab/utils/math.py index eff9e22d576d..c3366a3d8f81 100644 --- a/source/isaaclab/isaaclab/utils/math.py +++ b/source/isaaclab/isaaclab/utils/math.py @@ -8,8 +8,10 @@ # needed to import for allowing type-hinting: torch.Tensor | np.ndarray from __future__ import annotations +import functools import logging import math +from collections.abc import Sequence from typing import Literal import numpy as np @@ -1400,6 +1402,40 @@ def sample_uniform( return torch.rand(*size, device=device) * (upper - lower) + lower +def sample_uniform_from_ranges( + ranges: dict[str, tuple[float, float]], + keys: Sequence[str], + size: int | tuple[int, ...], + device: str | torch.device, +) -> torch.Tensor: + """Sample independent components uniformly from named ranges. + + Missing keys use zero bounds. Device bounds are cached by value, so edits to + ``ranges`` take effect on the next call without repeated uploads of unchanged bounds. + Only bounds are cached; each call draws fresh samples in the order given by ``keys``. + + Args: + ranges: Lower and upper bounds for each named component. + keys: Component names, in output order. + size: Batch shape, excluding the component dimension. + device: Device to create the samples on. + + Returns: + Sampled tensor of shape ``(*size, len(keys))`` using the default floating-point dtype. + """ + if isinstance(size, int): + size = (size,) + device = torch.device(device) + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + bounds = [] + for key in keys: + lower, upper = ranges.get(key, (0.0, 0.0)) + bounds.append((float(lower), float(upper))) + limits = _cached_uniform_bounds(tuple(bounds), device, torch.get_default_dtype()) + return sample_uniform(limits[:, 0], limits[:, 1], (*size, len(keys)), device) + + def sample_log_uniform( lower: torch.Tensor | float, upper: torch.Tensor | float, size: int | tuple[int, ...], device: str ) -> torch.Tensor: @@ -1975,3 +2011,11 @@ def generate_random_transformation_matrix(pos_boundary: float = 1, rot_boundary: T[:3, 3] = translation return T + + +@functools.lru_cache(maxsize=128) +def _cached_uniform_bounds( + bounds: tuple[tuple[float, float], ...], device: torch.device, dtype: torch.dtype +) -> torch.Tensor: + """Keep shared device bounds private so callers cannot mutate cached tensors.""" + return torch.tensor(bounds, dtype=dtype, device=device).reshape(-1, 2) diff --git a/source/isaaclab/test/envs/test_root_state_events.py b/source/isaaclab/test/envs/test_root_state_events.py index 8f51fc55aea5..5784f0354a72 100644 --- a/source/isaaclab/test/envs/test_root_state_events.py +++ b/source/isaaclab/test/envs/test_root_state_events.py @@ -90,12 +90,19 @@ def test_push_by_setting_velocity_follows_range_edits_without_sync(device): """Repeated ranges reuse device bounds without a sync, and in-place range edits still take effect.""" env, asset = _make_env(device) env_ids = torch.arange(env.num_envs, device=device) - velocity_range = {"x": (1.0, 1.0)} + # Dictionary insertion order differs from the command's component order. + velocity_range = {"yaw": (-0.2, 0.8), "x": (-2.0, 1.0)} push_by_setting_velocity(env, env_ids, velocity_range) - with _raise_on_cuda_sync(device): - push_by_setting_velocity(env, env_ids, velocity_range) - torch.testing.assert_close(asset.written["vel"][:, 0], torch.ones(env.num_envs, device=device)) + lower = torch.tensor([-2.0, 0.0, 0.0, 0.0, 0.0, -0.2], device=device) + upper = torch.tensor([1.0, 0.0, 0.0, 0.0, 0.0, 0.8], device=device) + torch.manual_seed(7) + expected = [torch.rand(env.num_envs, 6, device=device) * (upper - lower) + lower for _ in range(2)] + torch.manual_seed(7) + for sample in expected: + with _raise_on_cuda_sync(device): + push_by_setting_velocity(env, env_ids, velocity_range) + torch.testing.assert_close(asset.written["vel"], sample, rtol=0, atol=0) # a curriculum may edit the range dictionary in place velocity_range["x"] = (-2.0, -2.0) diff --git a/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst b/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst index d1ed38a9dde0..12978fae6808 100644 --- a/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst +++ b/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst @@ -21,6 +21,7 @@ Changed :func:`~isaaclab_tasks.core.lift.mdp.contact_count` take one norm per call instead of one per sensor. * Precomputed the ANYmal symmetry augmentation as one cached column permutation and sign tensor per transform, removing the per-minibatch index and sign uploads (about 1.5 ms to 0.15 ms per call). -* Cached the locomotion walk-target offset and the Lift reset position bounds on the device, +* Cached the locomotion walk-target offset and sampled Lift reset offsets through the shared + :func:`~isaaclab.utils.math.sample_uniform_from_ranges` helper, and read the Reorient goal directly from the command buffers instead of concatenating the command on every access. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py index aeef012499d7..2949b509ce94 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py @@ -7,7 +7,6 @@ from __future__ import annotations -import functools from typing import TYPE_CHECKING import numpy as np @@ -19,7 +18,7 @@ from isaaclab import cloner from isaaclab.managers import EventTermCfg, ManagerTermBase, ManagerTermBaseCfg, SceneEntityCfg from isaaclab.utils import instantiate -from isaaclab.utils.math import quat_apply, random_orientation, sample_uniform +from isaaclab.utils.math import quat_apply, random_orientation, sample_uniform, sample_uniform_from_ranges from isaaclab_tasks.utils.success_monitor import SuccessMonitor, SuccessMonitorCfg @@ -73,20 +72,6 @@ def reset_joints_shared_offset( ) -def _xyz_ranges(ranges: dict[str, tuple[float, float]], device: str | torch.device) -> torch.Tensor: - """Per-axis ``(low, high)`` bounds [m] of the ``x``, ``y``, ``z`` keys as a device tensor, shape ``(3, 2)``. - - Missing axes are ``(0, 0)``. The tensor is cached so resets do not upload the constant bounds. - """ - bounds = tuple(tuple(ranges.get(axis, (0.0, 0.0))) for axis in ("x", "y", "z")) - return _cached_xyz_ranges(bounds, str(device)) - - -@functools.lru_cache(maxsize=128) -def _cached_xyz_ranges(bounds: tuple[tuple[float, float], ...], device: str) -> torch.Tensor: - return torch.tensor(bounds, device=device) - - def reset_cable_state_uniform( env: ManagerBasedEnv, env_ids: torch.Tensor, @@ -97,8 +82,7 @@ def reset_cable_state_uniform( asset: CableObject = env.scene[asset_cfg.name] segment_pose = asset.data.default_segment_pose_w.torch[env_ids].clone() segment_velocity = asset.data.default_segment_velocity_w.torch[env_ids].clone() - ranges = _xyz_ranges(position_range, asset.device) - offset = sample_uniform(ranges[:, 0], ranges[:, 1], (segment_pose.shape[0], 3), device=asset.device) + offset = sample_uniform_from_ranges(position_range, ("x", "y", "z"), segment_pose.shape[0], device=asset.device) segment_pose[..., :3] += offset.unsqueeze(1) asset.write_segment_pose_to_sim_index(segment_pose=segment_pose, env_ids=env_ids) asset.write_segment_velocity_to_sim_index(segment_velocity=segment_velocity, env_ids=env_ids) @@ -837,9 +821,8 @@ def reset_deformable_over_support( deformable: DeformableObject = env.scene[asset_cfg.name] supports: tuple[RigidObject, RigidObject] = (env.scene[support_cfg[0].name], env.scene[support_cfg[1].name]) - ranges = _xyz_ranges(position_range, deformable.device) num_envs = len(range(env.num_envs)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) - offset = sample_uniform(ranges[:, 0], ranges[:, 1], (num_envs, 3), device=deformable.device) + offset = sample_uniform_from_ranges(position_range, ("x", "y", "z"), num_envs, device=deformable.device) nodal_state = deformable.data.default_nodal_state_w.torch[env_ids].clone() nodal_state[..., :3] += offset.unsqueeze(1) From e8a0e12bf8b88b442ad044d117d9b518a7385e65 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Tue, 29 Sep 2026 17:35:34 -0700 Subject: [PATCH 12/18] Require noise callbacks to preserve their inputs --- ...st => perf-hot-path-7b-managers.major.rst} | 10 ++++- .../isaaclab/managers/manager_term_cfg.py | 6 ++- .../isaaclab/managers/observation_manager.py | 38 ++++--------------- .../isaaclab/utils/noise/noise_cfg.py | 8 +++- .../isaaclab/utils/noise/noise_model.py | 4 +- .../test/managers/test_observation_manager.py | 18 +++++---- source/isaaclab/test/utils/test_noise.py | 6 +++ 7 files changed, 48 insertions(+), 42 deletions(-) rename source/isaaclab/changelog.d/{perf-hot-path-7b-managers.minor.rst => perf-hot-path-7b-managers.major.rst} (70%) diff --git a/source/isaaclab/changelog.d/perf-hot-path-7b-managers.minor.rst b/source/isaaclab/changelog.d/perf-hot-path-7b-managers.major.rst similarity index 70% rename from source/isaaclab/changelog.d/perf-hot-path-7b-managers.minor.rst rename to source/isaaclab/changelog.d/perf-hot-path-7b-managers.major.rst index 9625e65a8249..b268d5c9a022 100644 --- a/source/isaaclab/changelog.d/perf-hot-path-7b-managers.minor.rst +++ b/source/isaaclab/changelog.d/perf-hot-path-7b-managers.major.rst @@ -8,6 +8,14 @@ Added Changed ^^^^^^^ +* **Breaking:** Required noise functions and models configured through + :class:`~isaaclab.utils.noise.NoiseCfg` and :class:`~isaaclab.utils.noise.NoiseModelCfg` + to leave their input unchanged. :class:`~isaaclab.managers.ObservationManager` now calls + them without a defensive input copy. Custom callbacks using in-place operations must + clone their input first (for example, ``data.clone().add_(bias)``) or use out-of-place + operations (``data + bias``). Returning the unchanged input or a view remains supported; + the manager still protects borrowed outputs during subsequent processing. + * Changed episode logging in :class:`~isaaclab.managers.CommandTerm`, :class:`~isaaclab.managers.TerminationManager`, :class:`~isaaclab.managers.CurriculumManager`, :class:`~isaaclab.envs.mdp.UniformPoseCommand`, :class:`~isaaclab.envs.mdp.UniformPose2dCommand` and @@ -20,7 +28,7 @@ Changed * Cached the device range tensors of the root-state, nodal-state, push and center-of-mass event terms by value instead of uploading them on every call. * Batched the per-term updates of :class:`~isaaclab.managers.RewardManager`, and skipped the defensive copy - before the built-in noise functions and after delay buffers in :class:`~isaaclab.managers.ObservationManager`. + before noise callbacks and after delay buffers in :class:`~isaaclab.managers.ObservationManager`. * Shifted :class:`~isaaclab.utils.buffers.CircularBuffer` histories of small frames in two kernels regardless of the history length. diff --git a/source/isaaclab/isaaclab/managers/manager_term_cfg.py b/source/isaaclab/isaaclab/managers/manager_term_cfg.py index 0a1dbeeac4c3..d5db5cb447ac 100644 --- a/source/isaaclab/isaaclab/managers/manager_term_cfg.py +++ b/source/isaaclab/isaaclab/managers/manager_term_cfg.py @@ -171,7 +171,11 @@ class ObservationTermCfg(ManagerTermBaseCfg): """ noise: NoiseCfg | NoiseModelCfg | None = None - """The noise to add to the observation. Defaults to None, in which case no noise is added.""" + """The noise to add to the observation. Defaults to None, in which case no noise is added. + + Noise callbacks must not modify their input. Use out-of-place operations or clone the + input inside the callback before modifying it. + """ clip: tuple[float, float] | None = None """The clipping range for the observation after adding noise. Defaults to None, diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index 04ae6b69c3e1..134ec65d3e63 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -8,7 +8,7 @@ from __future__ import annotations import inspect -from collections.abc import Callable, Sequence +from collections.abc import Sequence from typing import TYPE_CHECKING import numpy as np @@ -16,7 +16,7 @@ from prettytable import PrettyTable from ..envs.utils.io_descriptors import _warn_io_descriptors_deprecated -from ..utils import callable_to_string, instantiate, modifiers, noise, to_dict +from ..utils import instantiate, modifiers, noise, to_dict from ..utils.buffers import CircularBuffer, DelayBuffer from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import ObservationGroupCfg, ObservationTermCfg @@ -24,14 +24,6 @@ if TYPE_CHECKING: from ..envs import ManagerBasedEnv -# Built-in noise functions that never modify their input, so the manager can skip the defensive copy. -# Configurations hold them as callables or as unresolved ``module:function`` strings. -_OUT_OF_PLACE_NOISE_FUNCS = tuple( - ref - for func in (noise.constant_noise, noise.gaussian_noise, noise.uniform_noise) - for ref in (func, callable_to_string(func)) -) - class ObservationManager(ManagerBase): """Manager for computing observation signals for a given world. @@ -425,10 +417,10 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch. group_term_names = self._group_obs_term_names[group_name] group_obs = dict.fromkeys(group_term_names, None) term_cfgs = self._group_obs_term_cfgs[group_name] - obs_terms = zip(group_term_names, term_cfgs, self._group_obs_term_safe_noise_funcs[group_name]) + obs_terms = zip(group_term_names, term_cfgs) # evaluate terms: compute, add noise, clip, scale, custom modifiers - for term_name, term_cfg, safe_noise_func in obs_terms: + for term_name, term_cfg in obs_terms: output_spec = self._term_output_specs.get((group_name, term_name)) owned = output_spec is not None if output_spec is not None: @@ -450,17 +442,12 @@ def compute_group(self, group_name: str, update_history: bool = False) -> torch. obs = modifier.func(obs) else: obs = modifier.func(obs, **modifier.params) + # Noise callbacks must not modify their input, but may return borrowed storage. if isinstance(term_cfg.noise, noise.NoiseCfg): - if term_cfg.noise.func is safe_noise_func: - noisy_obs = term_cfg.noise.func(obs, term_cfg.noise) - # a new result is owned; a returned input keeps its ownership - owned = owned or noisy_obs is not obs - obs = noisy_obs - else: - obs = term_cfg.noise.func(obs.clone(), term_cfg.noise) - owned = False + obs = term_cfg.noise.func(obs, term_cfg.noise) + owned = False elif isinstance(term_cfg.noise, noise.NoiseModelCfg) and term_cfg.noise.func is not None: - obs = term_cfg.noise.func(obs.clone()) + obs = term_cfg.noise.func(obs) owned = False if term_cfg.clip: obs = ( @@ -565,7 +552,6 @@ def _prepare_terms(self): self._group_obs_term_names: dict[str, list[str]] = {} self._group_obs_term_dim: dict[str, list[tuple[int, ...]]] = {} self._group_obs_term_cfgs: dict[str, list[ObservationTermCfg]] = {} - self._group_obs_term_safe_noise_funcs: dict[str, list[Callable | str | None]] = {} self._group_obs_class_term_cfgs: dict[str, list[ObservationTermCfg]] = {} self._group_obs_concatenate: dict[str, bool] = {} self._group_obs_concatenate_dim: dict[str, int] = {} @@ -605,7 +591,6 @@ def _prepare_terms(self): self._group_obs_term_names[group_name] = [] self._group_obs_term_dim[group_name] = [] self._group_obs_term_cfgs[group_name] = [] - self._group_obs_term_safe_noise_funcs[group_name] = [] self._group_obs_class_term_cfgs[group_name] = [] # history buffers @@ -666,13 +651,6 @@ def _prepare_terms(self): # add term config to list self._group_obs_term_names[group_name].append(term_name) self._group_obs_term_cfgs[group_name].append(term_cfg) - # Keep custom noise isolated even if its function is replaced after setup. - self._group_obs_term_safe_noise_funcs[group_name].append( - term_cfg.noise.func - if isinstance(term_cfg.noise, noise.NoiseCfg) and term_cfg.noise.func in _OUT_OF_PLACE_NOISE_FUNCS - else None - ) - # call function the first time to fill up dimensions obs_dims = self._prepare_term_output(group_name, term_name, term_cfg) diff --git a/source/isaaclab/isaaclab/utils/noise/noise_cfg.py b/source/isaaclab/isaaclab/utils/noise/noise_cfg.py index da76e38b7212..5ba3824dd8a1 100644 --- a/source/isaaclab/isaaclab/utils/noise/noise_cfg.py +++ b/source/isaaclab/isaaclab/utils/noise/noise_cfg.py @@ -25,7 +25,9 @@ class NoiseCfg: """The function to be called for applying the noise. Note: - The shape of the input and output tensors must be the same. + The shape of the input and output tensors must be the same. The function must not + modify its input, including through a view. Callers do not provide a defensive copy; + use out-of-place operations or clone the input before modifying it. """ operation: Literal["add", "scale", "abs"] = "add" """The operation to apply the noise on the data. Defaults to "add".""" @@ -84,7 +86,9 @@ class NoiseModelCfg: """Function or callable class used by this noise model. The function must take a single `torch.Tensor` (the batch of observations) as input - and return a `torch.Tensor` of the same shape with noise applied. + and return a `torch.Tensor` of the same shape with noise applied. It must not modify + its input, including through a view. Use out-of-place operations or clone the input + before modifying it; callers do not provide a defensive copy. It also supports `callable classes `_, i.e. classes that implement the ``__call__()`` method. In this case, the class should inherit from the diff --git a/source/isaaclab/isaaclab/utils/noise/noise_model.py b/source/isaaclab/isaaclab/utils/noise/noise_model.py index f3ad251bd508..ff150d304ea3 100644 --- a/source/isaaclab/isaaclab/utils/noise/noise_model.py +++ b/source/isaaclab/isaaclab/utils/noise/noise_model.py @@ -120,7 +120,9 @@ def reset(self, env_ids: Sequence[int] | None = None): pass def __call__(self, data: torch.Tensor) -> torch.Tensor: - """Apply the noise to the data. + """Apply the noise without modifying the input data. + + Implementations must use out-of-place operations or clone the input before modifying it. Args: data: The data to apply the noise to. Shape is (num_envs, ...). diff --git a/source/isaaclab/test/managers/test_observation_manager.py b/source/isaaclab/test/managers/test_observation_manager.py index ac9780fe571e..8bd2f3dfc7c8 100644 --- a/source/isaaclab/test/managers/test_observation_manager.py +++ b/source/isaaclab/test/managers/test_observation_manager.py @@ -371,9 +371,9 @@ def scratch_term(env, value): torch.testing.assert_close(result[:, 3:], torch.full_like(env.data.pos_w, 2.0)) -def in_place_noise(data: torch.Tensor, cfg: noise.NoiseCfg) -> torch.Tensor: - """Custom noise that modifies its input.""" - return data.add_(0.2) +def custom_noise(data: torch.Tensor, cfg: noise.NoiseCfg) -> torch.Tensor: + """Custom noise owns any storage it modifies; absolute mode returns an unchanged view.""" + return data.view_as(data) if cfg.operation == "abs" else data.clone().add_(0.2) @pytest.mark.parametrize( @@ -381,9 +381,11 @@ def in_place_noise(data: torch.Tensor, cfg: noise.NoiseCfg) -> torch.Tensor: [ (noise.UniformNoiseCfg(n_min=-0.1, n_max=0.1), None), (noise.UniformNoiseCfg(n_min=0.0, n_max=0.0), 0.0), - (noise.NoiseCfg(func=in_place_noise), 0.2), + (noise.NoiseCfg(func=custom_noise), 0.2), + (noise.NoiseCfg(func=custom_noise, operation="abs"), 0.0), + (noise.NoiseModelCfg(noise_cfg=noise.NoiseCfg(func=custom_noise, operation="abs")), 0.0), ], - ids=["out_of_place", "same_input", "in_place"], + ids=["out_of_place", "same_input", "custom", "view", "model_view"], ) def test_noise_preserves_source_and_rng_stream(setup_env, noise_cfg, expected_delta): """Noise ownership paths preserve source storage, clipping, scaling, and random samples.""" @@ -401,10 +403,12 @@ def test_noise_preserves_source_and_rng_stream(setup_env, noise_cfg, expected_de torch.testing.assert_close(env.data.pos_w, source) torch.testing.assert_close(result, expected) - manager.cfg["policy"].position.noise.func = in_place_noise + # Without post-processing, the returned observation must still own its storage. + manager.cfg["policy"].position.clip = None + manager.cfg["policy"].position.scale = None result = manager.compute()["policy"] + result.zero_() torch.testing.assert_close(env.data.pos_w, source) - torch.testing.assert_close(result, (source + 0.2).clip(0.0, 0.5) * 2.0) def test_compute_with_2d_history(setup_env): diff --git a/source/isaaclab/test/utils/test_noise.py b/source/isaaclab/test/utils/test_noise.py index bbcac7872d81..5c8033a84156 100644 --- a/source/isaaclab/test/utils/test_noise.py +++ b/source/isaaclab/test/utils/test_noise.py @@ -31,9 +31,11 @@ def test_gaussian_noise(device, noise_device, op): # create noise config noise_cfg = noise.GaussianNoiseCfg(std=std, mean=mean, operation=op) + source = data.clone() for _ in range(2): # apply noise noisy_data = noise_cfg.func(data, cfg=noise_cfg) + torch.testing.assert_close(data, source) # calculate resulting noise compared to original data set if op == "add": std_result, mean_result = torch.std_mean(noisy_data - data, dim=0) @@ -61,9 +63,11 @@ def test_uniform_noise(device, noise_device, op): # create noise config noise_cfg = noise.UniformNoiseCfg(n_max=n_max, n_min=n_min, operation=op) + source = data.clone() for _ in range(2): # apply noise noisy_data = noise_cfg.func(data, cfg=noise_cfg) + torch.testing.assert_close(data, source) # calculate resulting noise compared to original data set if op == "add": min_result, _ = torch.min(noisy_data - data, dim=0) @@ -93,9 +97,11 @@ def test_constant_noise(device, noise_device, op): # create noise config noise_cfg = noise.ConstantNoiseCfg(bias=bias, operation=op) + source = data.clone() for _ in range(2): # apply noise noisy_data = noise_cfg.func(data, cfg=noise_cfg) + torch.testing.assert_close(data, source) # calculate resulting noise compared to original data set if op == "add": bias_result = noisy_data - data From cd13a329eda973faf8d25ea31aaf977fc4b5341c Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Tue, 29 Sep 2026 17:37:06 -0700 Subject: [PATCH 13/18] Keep noise configuration docs concise --- .../isaaclab/managers/manager_term_cfg.py | 6 +----- .../isaaclab/isaaclab/utils/noise/noise_cfg.py | 16 ++++------------ 2 files changed, 5 insertions(+), 17 deletions(-) diff --git a/source/isaaclab/isaaclab/managers/manager_term_cfg.py b/source/isaaclab/isaaclab/managers/manager_term_cfg.py index d5db5cb447ac..0a1dbeeac4c3 100644 --- a/source/isaaclab/isaaclab/managers/manager_term_cfg.py +++ b/source/isaaclab/isaaclab/managers/manager_term_cfg.py @@ -171,11 +171,7 @@ class ObservationTermCfg(ManagerTermBaseCfg): """ noise: NoiseCfg | NoiseModelCfg | None = None - """The noise to add to the observation. Defaults to None, in which case no noise is added. - - Noise callbacks must not modify their input. Use out-of-place operations or clone the - input inside the callback before modifying it. - """ + """The noise to add to the observation. Defaults to None, in which case no noise is added.""" clip: tuple[float, float] | None = None """The clipping range for the observation after adding noise. Defaults to None, diff --git a/source/isaaclab/isaaclab/utils/noise/noise_cfg.py b/source/isaaclab/isaaclab/utils/noise/noise_cfg.py index 5ba3824dd8a1..63c7329c0d53 100644 --- a/source/isaaclab/isaaclab/utils/noise/noise_cfg.py +++ b/source/isaaclab/isaaclab/utils/noise/noise_cfg.py @@ -19,16 +19,10 @@ @configclass class NoiseCfg: - """Base configuration for a noise term.""" + """Base configuration for noise that does not modify its input.""" func: Callable[[torch.Tensor, NoiseCfg], torch.Tensor] = MISSING - """The function to be called for applying the noise. - - Note: - The shape of the input and output tensors must be the same. The function must not - modify its input, including through a view. Callers do not provide a defensive copy; - use out-of-place operations or clone the input before modifying it. - """ + """Function that applies noise without changing its input or shape.""" operation: Literal["add", "scale", "abs"] = "add" """The operation to apply the noise on the data. Defaults to "add".""" @@ -74,7 +68,7 @@ class GaussianNoiseCfg(NoiseCfg): @configclass class NoiseModelCfg: - """Configuration for a noise model.""" + """Configuration for a noise model that does not modify its input.""" class_type: type[NoiseModel] | str = "{DIR}.noise_model:NoiseModel" """The class type of the noise model.""" @@ -86,9 +80,7 @@ class NoiseModelCfg: """Function or callable class used by this noise model. The function must take a single `torch.Tensor` (the batch of observations) as input - and return a `torch.Tensor` of the same shape with noise applied. It must not modify - its input, including through a view. Use out-of-place operations or clone the input - before modifying it; callers do not provide a defensive copy. + and return a `torch.Tensor` of the same shape with noise applied. It also supports `callable classes `_, i.e. classes that implement the ``__call__()`` method. In this case, the class should inherit from the From e81da34587de1bab9e30f1d80c94e027b42b59fc Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Tue, 29 Sep 2026 17:42:31 -0700 Subject: [PATCH 14/18] Trim redundant host-sync comments --- source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py | 1 - source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py | 1 - source/isaaclab/isaaclab/managers/curriculum_manager.py | 1 - source/isaaclab/isaaclab/managers/termination_manager.py | 1 - source/isaaclab_rl/isaaclab_rl/sb3.py | 1 - .../isaaclab_tasks/core/cabinet/cabinet_direct_env.py | 2 -- .../isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py | 1 - .../isaaclab_tasks/core/fourbar_pole/mdp/rewards.py | 1 - source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py | 1 - .../isaaclab_tasks/core/locomotion/locomotion_direct_env.py | 2 -- .../isaaclab_tasks/core/reorient/reorient_direct_env.py | 2 -- source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py | 2 -- 12 files changed, 16 deletions(-) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index cecfe3ef829a..6d039894c528 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -212,7 +212,6 @@ def apply_actions(self): # compute the delta in joint-space jacobian = self._compute_frame_jacobian() joint_pos_des = self._ik_controller.compute(ee_pos_curr, ee_quat_curr, jacobian, joint_pos) - # hold the joints while the frame pose is uninitialized (all-zero quaternions), without a host sync joint_pos_des = torch.where(ee_quat_curr.norm() != 0, joint_pos_des, joint_pos) # set the joint position command self._asset.set_joint_position_target_index(target=joint_pos_des, joint_ids=self._joint_ids) diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index c84aafd6ea4f..f580e87d3052 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -181,7 +181,6 @@ def _update_command(self): """ # Compute angular velocity from heading direction if self.cfg.heading_command: - # compute angular velocity for all envs, then select heading envs without a host sync heading_error = math_utils.wrap_to_pi(self.heading_target - self.robot.data.heading_w.torch) heading_ang_vel = torch.clip( self.cfg.heading_control_stiffness * heading_error, diff --git a/source/isaaclab/isaaclab/managers/curriculum_manager.py b/source/isaaclab/isaaclab/managers/curriculum_manager.py index 13d3e346b5cb..76db6d0d7b58 100644 --- a/source/isaaclab/isaaclab/managers/curriculum_manager.py +++ b/source/isaaclab/isaaclab/managers/curriculum_manager.py @@ -102,7 +102,6 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, float | torch Dictionary of curriculum terms and their states. """ extras = {} - # tensors are snapshotted on device: a host read would synchronize, and terms may update state in place for term_name, term_state in self._curriculum_state.items(): if term_state is not None: # deal with dict diff --git a/source/isaaclab/isaaclab/managers/termination_manager.py b/source/isaaclab/isaaclab/managers/termination_manager.py index cb86c5792f7a..dad362207acf 100644 --- a/source/isaaclab/isaaclab/managers/termination_manager.py +++ b/source/isaaclab/isaaclab/managers/termination_manager.py @@ -140,7 +140,6 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor] env_ids = slice(None) # add to episode dict extras = {} - # log device scalars; a host read would synchronize the stream last_episode_done_stats = self._last_episode_dones.float().mean(dim=0) for i, key in enumerate(self._term_names): # store information diff --git a/source/isaaclab_rl/isaaclab_rl/sb3.py b/source/isaaclab_rl/isaaclab_rl/sb3.py index 5eb5986c42f6..b8403920d92e 100644 --- a/source/isaaclab_rl/isaaclab_rl/sb3.py +++ b/source/isaaclab_rl/isaaclab_rl/sb3.py @@ -252,7 +252,6 @@ def step_wait(self) -> VecEnvStepReturn: # noqa: D102 rewards = rew.detach().cpu().numpy() terminated = terminated.detach().cpu().numpy() truncated = truncated.detach().cpu().numpy() - # compute reset ids on the host copies, which saves a device-to-host transfer dones = terminated | truncated reset_ids = dones.nonzero()[0] diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py index e52f620fa3d4..37ed3f5e781d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py @@ -31,7 +31,6 @@ def __init__(self, cfg: CabinetDirectEnvCfg, render_mode: str | None = None, **k self._robot, self._cabinet = self.scene["robot"], self.scene["cabinet"] arm_joint_ids, _ = self._robot.find_joints(self.cfg.arm_joint_names) finger_joint_ids, _ = self._robot.find_joints(self.cfg.finger_joint_names) - # device indices: indexing a CUDA tensor with a host list uploads it on every call self.arm_joint_ids = torch.tensor(arm_joint_ids, dtype=torch.long, device=self.device) self.finger_joint_ids = torch.tensor(finger_joint_ids, dtype=torch.long, device=self.device) self.ee_body_idx = self._robot.find_bodies(self.cfg.ee_body_name)[0][0] @@ -210,7 +209,6 @@ def _reset_idx(self, env_ids: Sequence[int] | None) -> None: env_ids = torch.arange(self.num_envs, device=self.device) log = self.extras.setdefault("log", {}) - # 0-dim device tensors: no host sync on reset log["Metrics/success_rate"] = self._episode_succeeded[env_ids].float().mean() log["Metrics/drawer_pos"] = self._best_drawer_pos[env_ids].mean() for name, episode_sum in self._episode_reward_sums.items(): diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py index f22a7af233eb..6cb144bb4f5c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/mdp/rewards.py @@ -155,7 +155,6 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): def reset(self, env_ids: torch.Tensor): log = self._env.extras.setdefault("log", {}) - # 0-dim device tensors: no host sync on reset log["Metrics/success_rate"] = self.succeeded[env_ids].float().mean() log["Metrics/drawer_pos"] = self.best_drawer_pos[env_ids].mean() index_fill_(self.succeeded, env_ids, False) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/mdp/rewards.py index 5eae41f9b1af..092fa0e7f74a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/fourbar_pole/mdp/rewards.py @@ -35,7 +35,6 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): self._hold_steps = max(1, round(hold_time_s / env.step_dt)) def reset(self, env_ids: torch.Tensor): - # 0-dim device tensor: no host sync on reset self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._success[env_ids].mean() index_fill_(self._consecutive_upright, env_ids, 0.0) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py index 4c68a4a1f54a..421960f64836 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py @@ -403,7 +403,6 @@ def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): def reset(self, env_ids: Sequence[int] | None = None) -> None: if env_ids is None: env_ids = slice(None) - # 0-dim device tensor: no host sync on reset self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._succeeded[env_ids].float().mean() index_fill_(self._succeeded, env_ids, False) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/locomotion_direct_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/locomotion_direct_env.py index 14492b263154..7002d8fdf798 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/locomotion_direct_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/locomotion/locomotion_direct_env.py @@ -57,7 +57,6 @@ def __init__(self, cfg: AntEnvCfg | HumanoidEnvCfg, render_mode: str | None = No # resolve against the sensor's own body list: its ordering is backend-specific and does not # necessarily match the articulation's body ordering feet_body_idx, _ = self.joint_wrench.find_bodies(self.cfg.feet_body_names) - # device indices: indexing a CUDA tensor with a host list uploads it on every call self._feet_body_idx = torch.tensor(feet_body_idx, dtype=torch.long, device=self.device) # walk target, placed far enough away that the robot never reaches it @@ -170,7 +169,6 @@ def _reset_idx(self, env_ids: Sequence[int] | None): # log survival success rate before resetting (survived = timed out without falling) survived = self.reset_time_outs[env_ids].float() - # 0-dim device tensor: no host sync on reset self.extras.setdefault("log", {})["Metrics/success_rate"] = survived.mean() super()._reset_idx(env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py index 8b8a6266de3b..20a84724c8b4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/reorient_direct_env.py @@ -50,7 +50,6 @@ def __init__(self, cfg: AllegroHandEnvCfg | ShadowHandEnvCfg, render_mode: str | raise ValueError( f"Expected {len(cfg.actuated_joint_names)} actuated joints, found {len(actuated_dof_indices)}." ) - # device indices: indexing a CUDA tensor with a host list uploads it on every call self.actuated_dof_indices = torch.tensor(actuated_dof_indices, dtype=torch.long, device=self.device) # a slice avoids the gather/scatter when every joint is actuated (``find_joints`` returns sorted ids) all_actuated = actuated_dof_indices == list(range(self.num_hand_dofs)) @@ -234,7 +233,6 @@ def _reset_idx(self, env_ids: Sequence[int]): # Per-attempt success rate, matching the manager workflow's metric of the same # name: reaching a goal draws a replacement, so exactly one goal is outstanding # when the episode ends and it presented ``successes + 1``. - # 0-dim device tensor: avoids a host sync here; consumers read it at logging cadence goals = self.successes[env_ids] self.extras.setdefault("log", {})["Metrics/success_rate"] = (goals / (goals + 1.0)).mean() for statistic, value in self._orientation_error.reset(env_ids).items(): diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py index 33c64ce12a49..dfb9b06a621a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py @@ -41,7 +41,6 @@ def __init__(self, num_envs: int, device: str | torch.device): """ self.minimum_error = torch.full((num_envs,), torch.inf, device=device) self._has_sample = torch.zeros(num_envs, dtype=torch.bool, device=device) - # host-side flag: before the first update no episode has a sample, known without a device read self._updated = False def update(self, error: torch.Tensor) -> None: @@ -80,7 +79,6 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> env_ids = slice(None) statistics = {} if self._updated: - # NaN marks environments without a sample; the nan-reductions skip them without a host sync values = torch.where(self._has_sample[env_ids], self.minimum_error[env_ids], torch.nan) statistics = { "mean": values.nanmean(), From ea03b6542177a6cbc83233494c960f6990d01cda Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Tue, 29 Sep 2026 18:41:03 -0700 Subject: [PATCH 15/18] Remove sync-focused curriculum and velocity command tests --- .../test/envs/test_velocity_command_term.py | 103 ------------------ .../test/managers/test_curriculum_manager.py | 65 ----------- 2 files changed, 168 deletions(-) delete mode 100644 source/isaaclab/test/envs/test_velocity_command_term.py delete mode 100644 source/isaaclab/test/managers/test_curriculum_manager.py diff --git a/source/isaaclab/test/envs/test_velocity_command_term.py b/source/isaaclab/test/envs/test_velocity_command_term.py deleted file mode 100644 index b24b8ef4fedf..000000000000 --- a/source/isaaclab/test/envs/test_velocity_command_term.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Unit tests for the uniform velocity command term.""" - -import math -from contextlib import contextmanager -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -import torch - -from isaaclab.envs.mdp import UniformVelocityCommand, UniformVelocityCommandCfg -from isaaclab.test.utils import DeviceScope, test_devices - -pytestmark = pytest.mark.unit - - -@contextmanager -def _raise_on_cuda_sync(device: str): - """Raise if the enclosed code synchronizes a CUDA device.""" - if not device.startswith("cuda"): - yield - return - previous = torch.cuda.get_sync_debug_mode() - torch.cuda.synchronize(device) - torch.cuda.set_sync_debug_mode("error") - try: - yield - finally: - torch.cuda.set_sync_debug_mode(previous) - - -def _make_command(device: str, num_envs: int) -> tuple[UniformVelocityCommand, SimpleNamespace]: - """Create a heading-controlled velocity command on a stub robot.""" - robot_data = SimpleNamespace( - heading_w=SimpleNamespace(torch=torch.zeros(num_envs, device=device)), - root_lin_vel_b=SimpleNamespace(torch=torch.zeros(num_envs, 3, device=device)), - root_ang_vel_b=SimpleNamespace(torch=torch.zeros(num_envs, 3, device=device)), - ) - env = SimpleNamespace( - num_envs=num_envs, - device=device, - scene={"robot": SimpleNamespace(data=robot_data)}, - sim=MagicMock(), - extras={}, - ) - cfg = UniformVelocityCommandCfg( - asset_name="robot", - resampling_time_range=(10.0, 10.0), - heading_command=True, - heading_control_stiffness=0.8, - ranges=UniformVelocityCommandCfg.Ranges( - lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi) - ), - ) - return UniformVelocityCommand(cfg, env), robot_data - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_heading_control_updates_only_heading_envs(device): - """Heading envs track the wrapped, clipped heading error without a sync; others keep their sampled yaw rate.""" - heading_target = [3.0, 0.5, -3.0, 0.2, 1.0, -0.1] - heading = [-3.0, 0.0, 3.0, 0.0, -1.0, 0.0] - is_heading_env = [True, True, True, False, True, False] - sampled_yaw_rate = [0.3, -0.4, 0.5, -0.6, 0.7, -0.8] - command, robot_data = _make_command(device, len(heading)) - command.heading_target[:] = torch.tensor(heading_target, device=device) - robot_data.heading_w.torch[:] = torch.tensor(heading, device=device) - command.is_heading_env[:] = torch.tensor(is_heading_env, device=device) - command.is_standing_env[:] = False - command.vel_command_b[:] = 0.25 - command.vel_command_b[:, 2] = torch.tensor(sampled_yaw_rate, device=device) - - with _raise_on_cuda_sync(device): - command._update_command() - - expected_yaw_rate = [] - for target, current, active, sampled in zip(heading_target, heading, is_heading_env, sampled_yaw_rate): - error = math.atan2(math.sin(target - current), math.cos(target - current)) - expected_yaw_rate.append(min(max(0.8 * error, -1.0), 1.0) if active else sampled) - torch.testing.assert_close(command.command[:, 2].cpu(), torch.tensor(expected_yaw_rate)) - torch.testing.assert_close(command.command[:, :2].cpu(), torch.full((len(heading), 2), 0.25)) - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_reset_logs_metrics_without_host_reads(device): - """Reset logs episode metrics as device scalars without synchronizing the stream.""" - command, robot_data = _make_command(device, 4) - robot_data.root_lin_vel_b.torch[:2, 0] = 1.0 - command.compute(dt=0.1) - - with _raise_on_cuda_sync(device): - extras = command.reset() - - success_rate = command._env.extras["log"]["Metrics/success_rate"] - for value in (*extras.values(), success_rate): - assert isinstance(value, torch.Tensor) and value.ndim == 0 - assert set(extras) == {"error_vel_xy", "error_vel_yaw"} - torch.testing.assert_close(success_rate.cpu(), torch.tensor(0.5)) diff --git a/source/isaaclab/test/managers/test_curriculum_manager.py b/source/isaaclab/test/managers/test_curriculum_manager.py deleted file mode 100644 index 051eb729aefc..000000000000 --- a/source/isaaclab/test/managers/test_curriculum_manager.py +++ /dev/null @@ -1,65 +0,0 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Unit tests for the curriculum manager.""" - -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -import torch - -from isaaclab.managers import CurriculumManager, CurriculumTermCfg -from isaaclab.test.utils import DeviceScope, test_devices - -pytestmark = pytest.mark.unit - - -def advance_level(env, env_ids): - """Advance a persistent level tensor in place and report it with a derived state.""" - env.level += 1.0 - return {"level": env.level, "half_level": env.level / 2.0, "name": "stage"} - - -def advance_scalar(env, env_ids): - """Report the persistent level tensor directly.""" - return env.level - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_reset_logs_state_snapshots_without_host_reads(device): - """Reset logs tensor states as device snapshots, so later in-place term updates do not change them.""" - sim = MagicMock() - sim.is_playing.return_value = False - env = SimpleNamespace(num_envs=4, device=device, sim=sim, level=torch.zeros((), device=device)) - cfg = { - "dict_term": CurriculumTermCfg(func=advance_level), - "scalar_term": CurriculumTermCfg(func=advance_scalar), - } - manager = CurriculumManager(cfg, env) - manager.compute() - - previous = torch.cuda.get_sync_debug_mode() if device.startswith("cuda") else None - if previous is not None: - torch.cuda.synchronize(device) - torch.cuda.set_sync_debug_mode("error") - try: - extras = manager.reset() - finally: - if previous is not None: - torch.cuda.set_sync_debug_mode(previous) - manager.compute() - - assert extras["Curriculum/dict_term/name"] == "stage" - for key, expected in ( - ("Curriculum/dict_term/level", 1.0), - ("Curriculum/dict_term/half_level", 0.5), - ("Curriculum/scalar_term", 1.0), - ): - assert extras[key].device == torch.device(device) - assert extras[key].item() == expected - - assert manager.get_active_iterable_terms(0) == [("dict_term", [2.0, 1.0, "stage"]), ("scalar_term", [2.0])] - assert isinstance(manager.get_active_iterable_terms(0)[1][1][0], float) From dc82bb91305f7145a354c643184ca4feb4925981 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Tue, 29 Sep 2026 18:44:47 -0700 Subject: [PATCH 16/18] Trim sync, RNG, and private-state test assertions --- .../envs/test_diffik_jacobian_aliasing.py | 17 ++------ .../test/envs/test_root_state_events.py | 41 ++++--------------- .../test/managers/test_observation_manager.py | 3 +- .../test/managers/test_termination_manager.py | 36 ---------------- .../markers/test_visualization_markers.py | 2 - .../test/sensors/test_ray_caster_sensor.py | 5 --- .../test/utils/test_wrench_composer.py | 7 ---- .../test/assets/test_articulation.py | 10 ++--- .../test/physics/test_newton_solver_reset.py | 9 ---- .../test_isaac_rtx_renderer_contract.py | 2 +- 10 files changed, 18 insertions(+), 114 deletions(-) diff --git a/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py b/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py index 0707cf582b6f..4193f3e90fc8 100644 --- a/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py +++ b/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py @@ -23,7 +23,6 @@ DifferentialInverseKinematicsAction, OperationalSpaceControllerAction, ) -from isaaclab.test.utils import DeviceScope, test_devices from isaaclab.utils import math as math_utils pytestmark = pytest.mark.unit @@ -124,10 +123,10 @@ def test_body_offset_jacobian_uses_offset_in_root_frame(compute, returns_jacobia torch.testing.assert_close(result, expected) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) @pytest.mark.parametrize(("ee_quat", "ik_offset"), [((0.0, 0.0, 0.0, 1.0), 1.0), ((0.0, 0.0, 0.0, 0.0), 0.0)]) -def test_apply_actions_holds_joints_for_uninitialized_frame(device, ee_quat, ik_offset): - """IK targets are applied for a valid frame pose and joints hold for all-zero quaternions, without a sync.""" +def test_apply_actions_holds_joints_for_uninitialized_frame(ee_quat, ik_offset): + """IK targets are applied for a valid frame pose and joints hold for all-zero quaternions.""" + device = "cpu" joint_pos = torch.tensor([[0.1, 0.2], [0.3, 0.4]], device=device) frame_pose = (torch.zeros(2, 3, device=device), torch.tensor([ee_quat] * 2, device=device)) written = {} @@ -143,15 +142,7 @@ def test_apply_actions_holds_joints_for_uninitialized_frame(device, ee_quat, ik_ _compute_frame_jacobian=lambda: torch.zeros(2, 6, 2, device=device), _ik_controller=SimpleNamespace(compute=lambda ee_pos, ee_quat, jacobian, joint_pos: joint_pos + 1.0), ) - previous = torch.cuda.get_sync_debug_mode() if device.startswith("cuda") else None - if previous is not None: - torch.cuda.synchronize(device) - torch.cuda.set_sync_debug_mode("error") - try: - DifferentialInverseKinematicsAction.apply_actions(stub) - finally: - if previous is not None: - torch.cuda.set_sync_debug_mode(previous) + DifferentialInverseKinematicsAction.apply_actions(stub) torch.testing.assert_close(written["target"], joint_pos + ik_offset) diff --git a/source/isaaclab/test/envs/test_root_state_events.py b/source/isaaclab/test/envs/test_root_state_events.py index 5784f0354a72..fc1f76b21239 100644 --- a/source/isaaclab/test/envs/test_root_state_events.py +++ b/source/isaaclab/test/envs/test_root_state_events.py @@ -5,7 +5,6 @@ """Unit tests for root-state event terms that sample from range dictionaries.""" -from contextlib import contextmanager from types import SimpleNamespace import pytest @@ -13,26 +12,10 @@ from isaaclab.envs.mdp import push_by_setting_velocity, reset_root_state_uniform from isaaclab.managers import EventTermCfg, SceneEntityCfg -from isaaclab.test.utils import DeviceScope, test_devices pytestmark = pytest.mark.unit -@contextmanager -def _raise_on_cuda_sync(device: str): - """Raise if the enclosed code synchronizes a CUDA device.""" - if not device.startswith("cuda"): - yield - return - previous = torch.cuda.get_sync_debug_mode() - torch.cuda.synchronize(device) - torch.cuda.set_sync_debug_mode("error") - try: - yield - finally: - torch.cuda.set_sync_debug_mode(previous) - - class _Scene(dict): """Scene double exposing assets by name and the environment origins.""" @@ -62,9 +45,9 @@ def _make_env(device: str, num_envs: int = 3) -> tuple[SimpleNamespace, SimpleNa return env, asset -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_reset_root_state_uniform_uses_call_ranges(device): +def test_reset_root_state_uniform_uses_call_ranges(): """The term samples from the ranges passed at call time, not the ones present at construction.""" + device = "cpu" env, asset = _make_env(device) asset_cfg = SceneEntityCfg("robot") cfg = EventTermCfg( @@ -85,24 +68,18 @@ def test_reset_root_state_uniform_uses_call_ranges(device): torch.testing.assert_close(asset.written["vel"], expected_vel) -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_push_by_setting_velocity_follows_range_edits_without_sync(device): - """Repeated ranges reuse device bounds without a sync, and in-place range edits still take effect.""" +def test_push_by_setting_velocity_follows_range_edits(): + """Ranges are matched to components by name, and in-place range edits still take effect.""" + device = "cpu" env, asset = _make_env(device) env_ids = torch.arange(env.num_envs, device=device) # Dictionary insertion order differs from the command's component order. velocity_range = {"yaw": (-0.2, 0.8), "x": (-2.0, 1.0)} push_by_setting_velocity(env, env_ids, velocity_range) - - lower = torch.tensor([-2.0, 0.0, 0.0, 0.0, 0.0, -0.2], device=device) - upper = torch.tensor([1.0, 0.0, 0.0, 0.0, 0.0, 0.8], device=device) - torch.manual_seed(7) - expected = [torch.rand(env.num_envs, 6, device=device) * (upper - lower) + lower for _ in range(2)] - torch.manual_seed(7) - for sample in expected: - with _raise_on_cuda_sync(device): - push_by_setting_velocity(env, env_ids, velocity_range) - torch.testing.assert_close(asset.written["vel"], sample, rtol=0, atol=0) + velocity = asset.written["vel"] + assert ((velocity[:, 0] >= -2.0) & (velocity[:, 0] <= 1.0)).all() + assert ((velocity[:, 5] >= -0.2) & (velocity[:, 5] <= 0.8)).all() + assert (velocity[:, 1:5] == 0.0).all() # a curriculum may edit the range dictionary in place velocity_range["x"] = (-2.0, -2.0) diff --git a/source/isaaclab/test/managers/test_observation_manager.py b/source/isaaclab/test/managers/test_observation_manager.py index 8bd2f3dfc7c8..7eee886b257f 100644 --- a/source/isaaclab/test/managers/test_observation_manager.py +++ b/source/isaaclab/test/managers/test_observation_manager.py @@ -380,12 +380,11 @@ def custom_noise(data: torch.Tensor, cfg: noise.NoiseCfg) -> torch.Tensor: ("noise_cfg", "expected_delta"), [ (noise.UniformNoiseCfg(n_min=-0.1, n_max=0.1), None), - (noise.UniformNoiseCfg(n_min=0.0, n_max=0.0), 0.0), (noise.NoiseCfg(func=custom_noise), 0.2), (noise.NoiseCfg(func=custom_noise, operation="abs"), 0.0), (noise.NoiseModelCfg(noise_cfg=noise.NoiseCfg(func=custom_noise, operation="abs")), 0.0), ], - ids=["out_of_place", "same_input", "custom", "view", "model_view"], + ids=["out_of_place", "custom", "view", "model_view"], ) def test_noise_preserves_source_and_rng_stream(setup_env, noise_cfg, expected_delta): """Noise ownership paths preserve source storage, clipping, scaling, and random samples.""" diff --git a/source/isaaclab/test/managers/test_termination_manager.py b/source/isaaclab/test/managers/test_termination_manager.py index fb967e10225f..6cb588444d9e 100644 --- a/source/isaaclab/test/managers/test_termination_manager.py +++ b/source/isaaclab/test/managers/test_termination_manager.py @@ -9,7 +9,6 @@ import torch from isaaclab.managers import TerminationManager, TerminationTermCfg -from isaaclab.test.utils import DeviceScope, test_devices pytestmark = pytest.mark.unit @@ -121,38 +120,3 @@ def test_time_out_vs_terminated_split(env): out = tm.compute() assert torch.all(out) assert torch.all(tm.terminated) and torch.all(tm.time_outs) - - -def fail_first_quarter(env) -> torch.Tensor: - """Returns True for the first quarter of the envs.""" - return torch.arange(env.num_envs, device=env.device) < env.num_envs // 4 - - -@pytest.mark.parametrize("device", test_devices(DeviceScope.CUDA)) -def test_reset_logs_term_rates_without_host_reads(device): - """Reset logs each term's last-episode activation rate as a device scalar without synchronizing.""" - sim = MagicMock() - sim.is_playing.return_value = False - env = DummyEnv(num_envs=20, device=device, sim=sim) - cfg = { - "quarter": TerminationTermCfg(func=fail_first_quarter, time_out=False), - "term_3": TerminationTermCfg(func=fail_every_3_steps, time_out=True), - } - tm = TerminationManager(cfg, env) - env.counter = 1 - tm.compute() - - previous = torch.cuda.get_sync_debug_mode() if device.startswith("cuda") else None - if previous is not None: - torch.cuda.synchronize(device) - torch.cuda.set_sync_debug_mode("error") - try: - extras = tm.reset() - finally: - if previous is not None: - torch.cuda.set_sync_debug_mode(previous) - - assert set(extras) == {"Episode_Termination/quarter", "Episode_Termination/term_3"} - assert all(value.ndim == 0 and value.device == torch.device(device) for value in extras.values()) - assert extras["Episode_Termination/quarter"].item() == 0.25 - assert extras["Episode_Termination/term_3"].item() == 0.0 diff --git a/source/isaaclab/test/markers/test_visualization_markers.py b/source/isaaclab/test/markers/test_visualization_markers.py index d5bd320214b2..d1721d0c3ac4 100644 --- a/source/isaaclab/test/markers/test_visualization_markers.py +++ b/source/isaaclab/test/markers/test_visualization_markers.py @@ -292,7 +292,6 @@ def test_visualization_skips_updates_when_invisible(sim): config = clone(POSITION_GOAL_MARKER_CFG) config.prim_path = "/World/Visuals/test_protos" test_marker = VisualizationMarkers(config) - assert test_marker._backends == [] initial_count = test_marker.count # Unused inputs are ignored until a backend becomes active. test_marker.visualize( @@ -301,7 +300,6 @@ def test_visualization_skips_updates_when_invisible(sim): assert test_marker.count == initial_count sim._has_offscreen_render = True test_marker.visualize(translations=torch.zeros(3, 3, device=sim.device)) - assert len(test_marker._backends) == 1 assert test_marker.count == 3 # play the simulation diff --git a/source/isaaclab/test/sensors/test_ray_caster_sensor.py b/source/isaaclab/test/sensors/test_ray_caster_sensor.py index f170432de192..2d973c705df9 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_sensor.py +++ b/source/isaaclab/test/sensors/test_ray_caster_sensor.py @@ -183,10 +183,7 @@ def test_ray_caster_reset_resamples_drift(sim_ground): cfg = _ray_caster_cfg("/World/Sensor", "world") sensor = RayCaster(cfg) sim.reset() - get_rng_state = torch.cuda.get_rng_state if torch.device(sensor.device).type == "cuda" else torch.get_rng_state - rng_state = get_rng_state() sensor.reset() - assert torch.equal(get_rng_state(), rng_state) assert (sensor.drift.torch == 0.0).all() assert (sensor.ray_cast_drift.torch == 0.0).all() @@ -234,9 +231,7 @@ def test_ray_caster_reset_resamples_drift(sim_ground): assert (sensor.drift.torch == 0.0).all() assert (sensor.ray_cast_drift.torch[:, 0] >= 0.1).all() sensor.cfg.ray_cast_drift_range = {} - rng_state = get_rng_state() sensor.reset(env_ids=[0]) - assert torch.equal(get_rng_state(), rng_state) assert (sensor.ray_cast_drift.torch == 0.0).all() diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index e7fed7cf8880..a86d990193e8 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -902,13 +902,6 @@ def test_full_reset_clears_active_flag(device: str, env_ids: slice | None): assert is_global is False np.testing.assert_array_equal(force.numpy(), np.zeros_like(forces_np)) np.testing.assert_array_equal(torque.numpy(), np.zeros_like(forces_np)) - with patch.object(wp, "launch") as launch, patch.object(wp.array, "zero_") as zero: - composer.reset(env_ids=[1]) - - launch.assert_not_called() - zero.assert_not_called() - assert not composer.active - composer.add_forces_and_torques_index( forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), is_global=True ) diff --git a/source/isaaclab_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index 2640a65dd836..2bb728645b39 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -2416,7 +2416,7 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav 2. A partial joint state write with unsorted int64 selectors updates only the selected entries 3. A joint state write moves the bodies and refreshes the derived body poses and velocities 4. Joint limits that exclude a default joint position clamp it into the new limits and report the - clamping at the requested log level, without a device readback when that level is disabled + clamping at the requested log level Args: sim: The simulation fixture @@ -2519,7 +2519,7 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav ) assert torch.all(within_bounds) - # Info-level clamping reports are logged when enabled and skip the device readback otherwise. + # Info-level clamping reports are logged only when that level is enabled. articulation.data.default_joint_pos.torch.fill_(1.0) full_limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) full_limits[..., 1] = 0.5 @@ -2529,10 +2529,7 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav assert [record.levelno for record in caplog.records] == [logging.INFO] articulation.data.default_joint_pos.torch.fill_(1.0) caplog.clear() - with ( - caplog.at_level(logging.WARNING, logger=articulation_logger), - patch.object(wp.array, "numpy", side_effect=AssertionError("unexpected device readback")), - ): + with caplog.at_level(logging.WARNING, logger=articulation_logger): articulation.write_joint_position_limit_to_sim_mask(limits=full_limits, warn_limit_violation=False) assert not caplog.records torch.testing.assert_close(articulation.data.default_joint_pos.torch, torch.full_like(full_limits[..., 1], 0.5)) @@ -2613,7 +2610,6 @@ def test_body_q_consistent_after_root_write(num_articulations, device, articulat active, then patches ``_simulate_physics_only`` to capture body_q at the moment collide() is called and asserts it matches joint_q. """ - from unittest.mock import patch sim_cfg = SimulationCfg( dt=1 / 200, diff --git a/source/isaaclab_newton/test/physics/test_newton_solver_reset.py b/source/isaaclab_newton/test/physics/test_newton_solver_reset.py index 381d92ea6bd4..c6353c54a9bb 100644 --- a/source/isaaclab_newton/test/physics/test_newton_solver_reset.py +++ b/source/isaaclab_newton/test/physics/test_newton_solver_reset.py @@ -137,12 +137,3 @@ def test_env_reset_clears_selected_mjwarp_solver_internals(device): torch.testing.assert_close(wp.to_torch(state.joint_qd), joint_qd_before) assert torch.count_nonzero(warm_start[0]).item() == 0 torch.testing.assert_close(warm_start[1], torch.full_like(warm_start[1], 29.0)) - - # A clean step has no flagged worlds, so it launches no solver reset. - with ( - patch.object(SimulationManager, "_simulate_full", classmethod(lambda cls: None)), - patch.object(SimulationManager, "_simulate_physics_only", classmethod(lambda cls: None)), - patch.object(SimulationManager, "_reset_solver_internals_delegate") as reset_delegate, - ): - sim.step(render=False) - reset_delegate.assert_not_called() diff --git a/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py b/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py index bffed8b3c704..0a75262bc7f3 100644 --- a/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py +++ b/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py @@ -488,7 +488,7 @@ def test_render_treats_empty_annotator_frame_as_not_ready(monkeypatch, data_type @pytest.mark.parametrize( ("data_type", "num_channels"), - [("normals", 3), ("motion_vectors", 2), ("rgb_hdr", 3), ("simple_shading_full_mdl", 3)], + [("normals", 3), ("motion_vectors", 2)], ) def test_render_keeps_leading_channels_of_padded_annotator_tiles(monkeypatch, data_type, num_channels): """Four-channel annotator tiles are split per camera, keeping only the output's leading channels.""" From 1ca362a7da4bba6cc58041b24041e1761315998d Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Wed, 30 Sep 2026 09:24:07 -0700 Subject: [PATCH 17/18] Restore contact force magnitude helper name --- .../isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py index 421960f64836..c0d7376b02fa 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/rewards.py @@ -67,7 +67,7 @@ def contacts(env: ManagerBasedRLEnv, threshold: float, thumb_name: str, finger_n Returns: Boolean tensor indicating good contact condition per environment. """ - in_contact = _contact_force_mags(env, [thumb_name, *finger_names]) > threshold + in_contact = _contact_force_mag(env, [thumb_name, *finger_names]) > threshold return in_contact[:, 0] & in_contact[:, 1:].any(dim=1) @@ -85,7 +85,7 @@ def contact_count(env: ManagerBasedRLEnv, threshold: float, sensor_names: list[s Returns: Tensor of shape (num_envs,) with the count of sensors in contact per environment. """ - count = (_contact_force_mags(env, sensor_names) > threshold).float().sum(dim=1) + count = (_contact_force_mag(env, sensor_names) > threshold).float().sum(dim=1) return count / len(sensor_names) @@ -503,7 +503,7 @@ def cable_segment_goal_reached( return (distance < success_threshold).float() -def _contact_force_mags(env: ManagerBasedRLEnv, sensor_names: Sequence[str]) -> torch.Tensor: +def _contact_force_mag(env: ManagerBasedRLEnv, sensor_names: Sequence[str]) -> torch.Tensor: """Per-environment contact force magnitudes [N] of single-body, single-filter contact sensors. The forces are stacked so a single norm covers every sensor. From e9d4425877c2d926aa5fe241f9651332e98ef81e Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Wed, 30 Sep 2026 10:22:39 -0700 Subject: [PATCH 18/18] Fix hot-path PR test failures and benchmark cache permissions --- .github/workflows/build.yaml | 2 ++ source/isaaclab/changelog.d/fix-pr8105-ci.skip | 0 source/isaaclab/test/assets/_articulation_iface_test_utils.py | 2 ++ source/isaaclab/test/managers/test_observation_manager.py | 2 +- source/isaaclab/test/sensors/test_sensor_base.py | 2 +- .../changelog.d/fix-pr8105-handover-tendons.rst | 4 ++++ .../isaaclab_tasks/core/handover/handover_env.py | 4 ++-- 7 files changed, 12 insertions(+), 4 deletions(-) create mode 100644 source/isaaclab/changelog.d/fix-pr8105-ci.skip create mode 100644 source/isaaclab_tasks/changelog.d/fix-pr8105-handover-tendons.rst diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 03d84b31059d..acc2ad693b38 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -1025,6 +1025,8 @@ jobs: run: | set -uo pipefail install -d -m 0777 "${GITHUB_WORKSPACE}/jit-cache/warp" "${GITHUB_WORKSPACE}/jit-cache/nv" + # Restored version directories and kernels must also be writable by the container user. + chmod -R a+rwX "${GITHUB_WORKSPACE}/jit-cache" || exit 1 container="" cleanup_container() { diff --git a/source/isaaclab/changelog.d/fix-pr8105-ci.skip b/source/isaaclab/changelog.d/fix-pr8105-ci.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab/test/assets/_articulation_iface_test_utils.py b/source/isaaclab/test/assets/_articulation_iface_test_utils.py index 25d10a44ee6e..d852c5f13e35 100644 --- a/source/isaaclab/test/assets/_articulation_iface_test_utils.py +++ b/source/isaaclab/test/assets/_articulation_iface_test_utils.py @@ -115,6 +115,7 @@ def create_physx_articulation( object.__setattr__(articulation, "_root_view", mock_view) object.__setattr__(articulation, "_device", device) + articulation._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=device) # We can't call the initialize method here, because we don't have a good mock for the actuators yet. # We need to set the _data attribute manually. @@ -392,6 +393,7 @@ def create_newton_articulation( object.__setattr__(articulation, "_device", device) object.__setattr__(articulation, "_data", data) object.__setattr__(articulation, "_test_simulation_manager", mock_manager) + articulation._clamped_default_count = wp.zeros(1, dtype=wp.int32, device=device) # the solver builds this adapter; the shell has no model, so it stays absent object.__setattr__(articulation, "_fixed_tendon_control", None) diff --git a/source/isaaclab/test/managers/test_observation_manager.py b/source/isaaclab/test/managers/test_observation_manager.py index 7eee886b257f..fe00e34d64fd 100644 --- a/source/isaaclab/test/managers/test_observation_manager.py +++ b/source/isaaclab/test/managers/test_observation_manager.py @@ -392,7 +392,7 @@ def test_noise_preserves_source_and_rng_stream(setup_env, noise_cfg, expected_de cfg = ObservationGroupCfg(enable_corruption=True) cfg.position = ObservationTermCfg(func=pos_w_data, noise=noise_cfg, clip=(0.0, 0.5), scale=2.0) manager = ObservationManager({"policy": cfg}, env) - source = env.data.pos_w.clone() + source = torch.clone(env.data.pos_w) torch.manual_seed(0) if expected_delta is None: expected_delta = torch.rand_like(source) * 0.2 - 0.1 diff --git a/source/isaaclab/test/sensors/test_sensor_base.py b/source/isaaclab/test/sensors/test_sensor_base.py index 9a229dccdd00..235569656fe6 100644 --- a/source/isaaclab/test/sensors/test_sensor_base.py +++ b/source/isaaclab/test/sensors/test_sensor_base.py @@ -93,7 +93,7 @@ def _set_debug_vis_impl(self, debug_vis: bool): def _debug_vis_callback(self, event): if self._is_initialized: - self.visualized_counts.append(self.data.count.clone()) + self.visualized_counts.append(torch.clone(self.data.count)) def _populate_scene(): diff --git a/source/isaaclab_tasks/changelog.d/fix-pr8105-handover-tendons.rst b/source/isaaclab_tasks/changelog.d/fix-pr8105-handover-tendons.rst new file mode 100644 index 000000000000..c868a8384953 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/fix-pr8105-handover-tendons.rst @@ -0,0 +1,4 @@ +Fixed +^^^^^ + +* Fixed direct handover action application with device-resident tendon indices. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env.py index 54434498e466..3b0eb2965f83 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/handover/handover_env.py @@ -70,7 +70,7 @@ def __init__(self, cfg: HandoverEnvCfg, render_mode: str | None = None, **kwargs # Motors that pull a tendon rather than drive a joint. Both hands are the same model, so # one index set serves both. - self.actuated_tendon_indices: list[int] = [] + self.actuated_tendon_indices = torch.empty(0, dtype=torch.long, device=self.device) if cfg.actuated_tendon_names: self.actuated_tendon_indices, self.tendon_lower_limits, self.tendon_upper_limits = resolve_actuated_tendons( self.right_hand, @@ -146,7 +146,7 @@ def _apply_hand_action( prev_targets[:, idx] = targets hand.set_joint_position_target_index(target=targets, joint_ids=idx) - if self.actuated_tendon_indices: + if len(self.actuated_tendon_indices) > 0: # No moving average on the tendon target: the manager task's action term applies none, # and the two task variants have to stay comparable. # saturate like the joint target above: a Gaussian policy samples past the action range,