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/changelog.d/perf-hot-path-7b-managers.major.rst b/source/isaaclab/changelog.d/perf-hot-path-7b-managers.major.rst new file mode 100644 index 000000000000..b268d5c9a022 --- /dev/null +++ b/source/isaaclab/changelog.d/perf-hot-path-7b-managers.major.rst @@ -0,0 +1,39 @@ +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 +^^^^^^^ + +* **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 + :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 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. + +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/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/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..80edccd2e498 --- /dev/null +++ b/source/isaaclab/changelog.d/perf-hot-path-7b-sensors.rst @@ -0,0 +1,16 @@ +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. + +Fixed +^^^^^ + +* Cleared previously sampled ray-caster drift on reset when its configured range was changed to zero. 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..6d039894c528 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,9 @@ 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) + 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..f580e87d3052 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,13 @@ 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( + 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..59309db504e7 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -522,11 +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 - 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) 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 @@ -1257,9 +1255,9 @@ 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) - 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) @@ -1277,20 +1275,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,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 = self._pose_ranges - 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 = self._velocity_ranges - 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 @@ -1356,20 +1340,16 @@ 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) - 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 - 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) - 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 @@ -1426,17 +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 - range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["roll", "pitch", "yaw"]] - ranges = torch.tensor(range_list, device=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 - 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) - 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 @@ -1699,19 +1679,15 @@ 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) - 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 - range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - ranges = torch.tensor(range_list, device=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/envs/mdp/rewards.py b/source/isaaclab/isaaclab/envs/mdp/rewards.py index 127c8cd3d9b4..7fcae5494399 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..76db6d0d7b58 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: @@ -109,12 +109,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: diff --git a/source/isaaclab/isaaclab/managers/observation_manager.py b/source/isaaclab/isaaclab/managers/observation_manager.py index 1f7e1fbac070..134ec65d3e63 100644 --- a/source/isaaclab/isaaclab/managers/observation_manager.py +++ b/source/isaaclab/isaaclab/managers/observation_manager.py @@ -442,11 +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): - obs = term_cfg.noise.func(obs.clone(), term_cfg.noise) + 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 = ( @@ -462,7 +463,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] @@ -649,7 +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) - # 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/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..dad362207acf 100644 --- a/source/isaaclab/isaaclab/managers/termination_manager.py +++ b/source/isaaclab/isaaclab/managers/termination_manager.py @@ -140,11 +140,10 @@ 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() + 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/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..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,9 @@ 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 def __str__(self) -> str: """Returns: A string containing information about the instance.""" @@ -104,6 +107,15 @@ 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) + 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: num_envs_ids = len(range(self._view_count)[env_ids]) if isinstance(env_ids, slice) else len(env_ids) @@ -114,14 +126,23 @@ 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) + else: + self.drift.torch[env_ids] = 0.0 # 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 + ) + else: + self.ray_cast_drift.torch[env_ids] = 0.0 """ Implementation. @@ -324,6 +345,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/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/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/isaaclab/utils/noise/noise_cfg.py b/source/isaaclab/isaaclab/utils/noise/noise_cfg.py index da76e38b7212..63c7329c0d53 100644 --- a/source/isaaclab/isaaclab/utils/noise/noise_cfg.py +++ b/source/isaaclab/isaaclab/utils/noise/noise_cfg.py @@ -19,14 +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. - """ + """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".""" @@ -72,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.""" 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/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/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/envs/test_diffik_jacobian_aliasing.py b/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py index f06ef4d368dd..4193f3e90fc8 100644 --- a/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py +++ b/source/isaaclab/test/envs/test_diffik_jacobian_aliasing.py @@ -123,5 +123,28 @@ def test_body_offset_jacobian_uses_offset_in_root_frame(compute, returns_jacobia torch.testing.assert_close(result, expected) +@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(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 = {} + 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), + ) + DifferentialInverseKinematicsAction.apply_actions(stub) + 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..fc1f76b21239 --- /dev/null +++ b/source/isaaclab/test/envs/test_root_state_events.py @@ -0,0 +1,87 @@ +# 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 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 + +pytestmark = pytest.mark.unit + + +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 + + +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( + 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) + + +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) + 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) + 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/managers/test_observation_manager.py b/source/isaaclab/test/managers/test_observation_manager.py index 27b76d300f67..fe00e34d64fd 100644 --- a/source/isaaclab/test/managers/test_observation_manager.py +++ b/source/isaaclab/test/managers/test_observation_manager.py @@ -371,6 +371,45 @@ def scratch_term(env, value): torch.testing.assert_close(result[:, 3:], torch.full_like(env.data.pos_w, 2.0)) +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( + ("noise_cfg", "expected_delta"), + [ + (noise.UniformNoiseCfg(n_min=-0.1, n_max=0.1), None), + (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", "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.""" + 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 = 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 + expected = (source + expected_delta).clip(0.0, 0.5) * 2.0 + torch.manual_seed(0) + result = manager.compute()["policy"] + torch.testing.assert_close(env.data.pos_w, source) + torch.testing.assert_close(result, expected) + + # 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) + + 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..e7d3f910521d 100644 --- a/source/isaaclab/test/managers/test_reward_manager.py +++ b/source/isaaclab/test/managers/test_reward_manager.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause from collections import namedtuple +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -90,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): @@ -162,3 +176,7 @@ 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) diff --git a/source/isaaclab/test/markers/test_visualization_markers.py b/source/isaaclab/test/markers/test_visualization_markers.py index 7529e08595ae..d1721d0c3ac4 100644 --- a/source/isaaclab/test/markers/test_visualization_markers.py +++ b/source/isaaclab/test/markers/test_visualization_markers.py @@ -267,6 +267,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" @@ -291,6 +292,15 @@ 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) + 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 test_marker.count == 3 # play the simulation sim.reset() diff --git a/source/isaaclab/test/sensors/test_ray_caster_sensor.py b/source/isaaclab/test/sensors/test_ray_caster_sensor.py index 18c63863f311..2d973c705df9 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_sensor.py +++ b/source/isaaclab/test/sensors/test_ray_caster_sensor.py @@ -181,18 +181,21 @@ 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 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. + sensor.reset() + 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}" @@ -216,6 +219,20 @@ 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() + + # 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 = {} + sensor.reset(env_ids=[0]) + 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 cfcb70d33cfb..235569656fe6 100644 --- a/source/isaaclab/test/sensors/test_sensor_base.py +++ b/source/isaaclab/test/sensors/test_sensor_base.py @@ -81,6 +81,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(torch.clone(self.data.count)) + + def _populate_scene(): """""" @@ -171,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() @@ -191,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): diff --git a/source/isaaclab/test/utils/test_circular_buffer.py b/source/isaaclab/test/utils/test_circular_buffer.py index 6563e8b89f4a..5b51f9fa5f9b 100644 --- a/source/isaaclab/test/utils/test_circular_buffer.py +++ b/source/isaaclab/test/utils/test_circular_buffer.py @@ -257,19 +257,18 @@ 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(("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, K = 2, 3, 4 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, 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, height, width, C), slot + 1.0) + ) def test_stack_dim_reset_clears_buffer(): 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 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-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/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/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/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/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 ef081c173a70..9642989db9e0 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -714,8 +714,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): @@ -727,6 +725,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/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_newton/test/assets/test_articulation.py b/source/isaaclab_newton/test/assets/test_articulation.py index f0f4b5da6dbc..2bb728645b39 100644 --- a/source/isaaclab_newton/test/assets/test_articulation.py +++ b/source/isaaclab_newton/test/assets/test_articulation.py @@ -22,10 +22,12 @@ launch_test_simulation() +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 @@ -1020,6 +1022,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 @@ -1040,13 +1043,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", @@ -2394,14 +2408,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 Args: sim: The simulation fixture @@ -2491,7 +2506,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) @@ -2501,6 +2519,21 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav ) assert torch.all(within_bounds) + # 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 + 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): + 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]) @@ -2577,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/assets/test_rigid_object.py b/source/isaaclab_newton/test/assets/test_rigid_object.py index 39a756a60b12..a49149b30455 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object.py @@ -440,6 +440,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 8724ad8ac51d..563266cab4fa 100644 --- a/source/isaaclab_newton/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_newton/test/assets/test_rigid_object_collection.py @@ -403,6 +403,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] 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..68a7df5d784b --- /dev/null +++ b/source/isaaclab_ov/changelog.d/perf-hot-path-7b-sensors.rst @@ -0,0 +1,10 @@ +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. + +* 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 b1137fe2ff93..8bc71c184231 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, @@ -4081,6 +4081,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/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=[ 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_ov/test/assets/test_articulation.py b/source/isaaclab_ov/test/assets/test_articulation.py index 40cea80c376a..f95e7aa082a4 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 @@ -1763,7 +1764,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: @@ -1834,6 +1835,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 new file mode 100644 index 000000000000..4f16c7a5ac44 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/perf-hot-path-7b-sensors.rst @@ -0,0 +1,10 @@ +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. + +* 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/isaaclab_physx/renderers/isaac_rtx_renderer.py b/source/isaaclab_physx/isaaclab_physx/renderers/isaac_rtx_renderer.py index df6b40130987..c49d3f51d3ca 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 1aeb71a2daeb..083db94eaf65 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/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_physx/test/renderers/test_isaac_rtx_renderer_contract.py b/source/isaaclab_physx/test/renderers/test_isaac_rtx_renderer_contract.py index 319ce055b8cb..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 @@ -486,6 +486,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)], +) +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", [ 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..b8403920d92e 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,7 @@ 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() + dones = terminated | truncated reset_ids = dones.nonzero()[0] 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/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..12978fae6808 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/perf-hot-path-7b-tasks.minor.rst @@ -0,0 +1,27 @@ +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). +* 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/cabinet/cabinet_direct_env.py b/source/isaaclab_tasks/isaaclab_tasks/core/cabinet/cabinet_direct_env.py index f760674092d8..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 @@ -29,8 +29,10 @@ 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) + 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 +54,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 +93,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 +102,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 +209,8 @@ 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() + 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..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,8 +155,8 @@ 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() + 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..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,7 @@ 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() + 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/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, 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..2949b509ce94 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/lift/mdp/events.py @@ -18,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 @@ -82,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 = torch.tensor([position_range.get(axis, (0.0, 0.0)) for axis in ("x", "y", "z")], device=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) @@ -822,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 = torch.tensor([position_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z")], device=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) 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..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,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_mag(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_mag(env, sensor_names) > threshold).float().sum(dim=1) return count / len(sensor_names) @@ -413,7 +403,7 @@ 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() + self._env.extras.setdefault("log", {})["Metrics/success_rate"] = self._succeeded[env_ids].float().mean() index_fill_(self._succeeded, env_ids, False) @@ -513,10 +503,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_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. + + 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..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 @@ -56,7 +56,8 @@ 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) + 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 +70,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 +169,7 @@ 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() + 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..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 @@ -45,26 +45,35 @@ 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.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 - 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): + 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 +103,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 +115,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: @@ -232,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(): @@ -244,8 +244,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 +272,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..dfb9b06a621a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/reorient/utils.py @@ -41,6 +41,7 @@ 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) + self._updated = False def update(self, error: torch.Tensor) -> None: """Record one error sample for every environment. @@ -58,6 +59,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 +70,20 @@ 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: + 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 +258,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 +274,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 +285,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 05f206cf6aa8..5aadcb08de89 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 @@ -146,9 +147,10 @@ def test_camera_normalization_is_stationary(data_type: str) -> None: torch.testing.assert_close(output.flatten(), expected) -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 @@ -169,7 +171,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( @@ -191,20 +192,36 @@ 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) + +@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, } @@ -212,6 +229,29 @@ def _initialize_command_term(command, command_cfg, command_env) -> None: assert len(visualizer.calls) == expected_count for _, kwargs in visualizer.calls: assert torch.equal(kwargs["environment_ids"], environment_ids) + 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: + """Contact gating and counting from stacked sensor forces match a per-sensor evaluation.""" + 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)) + 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() def test_lift_point_cloud_markers_repeat_environment_ids_per_point() -> None: diff --git a/source/isaaclab_tasks/test/core/test_reorient_utils.py b/source/isaaclab_tasks/test/core/test_reorient_utils.py index 5eff5ca1e7df..f13d6ecc7ab6 100644 --- a/source/isaaclab_tasks/test/core/test_reorient_utils.py +++ b/source/isaaclab_tasks/test/core/test_reorient_utils.py @@ -62,10 +62,24 @@ 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") + """Initial resets are empty; later statistics exclude unsampled episodes.""" + recorder = reorient_utils.EpisodeErrorRecorder(num_envs=5, device="cpu") + assert recorder.reset(torch.arange(5)) == {} - assert recorder.reset(torch.tensor([0, 1])) == {} + 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(): @@ -163,7 +177,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))