Fix lazy PhysX IMU and PVA acceleration

source/isaaclab_physx/benchmark/sensors/benchmark_imu_pva.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Benchmark the PhysX IMU and PVA sensor update paths.
+
+Usage:
+    ./isaaclab.sh -p source/isaaclab_physx/benchmark/sensors/benchmark_imu_pva.py \
+        --sensor imu --num_envs 4096
+    ./isaaclab.sh -p source/isaaclab_physx/benchmark/sensors/benchmark_imu_pva.py \
+        --sensor pva --num_envs 4096
+"""
+
+from __future__ import annotations
+
+import argparse
+
+from isaaclab.app import AppLauncher
+
+parser = argparse.ArgumentParser(description="Benchmark a PhysX IMU or PVA sensor update path.")
+parser.add_argument("--sensor", choices=("imu", "pva"), required=True, help="Sensor update path to benchmark.")
+parser.add_argument("--num_envs", type=int, default=4096, help="Number of environments and sensor instances.")
+parser.add_argument("--num_steps", type=int, default=500, help="Number of timed updates.")
+parser.add_argument("--warmup_steps", type=int, default=50, help="Number of untimed warm-up updates.")
+parser.add_argument("--label", type=str, default="current", help="Label printed with the benchmark results.")
+parser.add_argument(
+    "--disable_recorded_launch",
+    action="store_true",
+    help="Use cached PhysX views with ordinary eager Warp launches.",
+)
+AppLauncher.add_app_launcher_args(parser)
+args_cli = parser.parse_args()
+
+app_launcher = AppLauncher(args_cli)
+simulation_app = app_launcher.app
+
+"""Everything below follows application launch."""
+
+import statistics
+import time
+
+import torch
+import warp as wp
+
+import isaaclab.sim as sim_utils
+from isaaclab.assets import RigidObjectCfg
+from isaaclab.scene import InteractiveScene, InteractiveSceneCfg
+from isaaclab.sensors import ImuCfg, PvaCfg
+from isaaclab.utils.configclass import configclass
+
+
+@configclass
+class ImuPvaBenchmarkSceneCfg(InteractiveSceneCfg):
+    """One kinematic rigid body and one selected sensor per environment."""
+
+    body = RigidObjectCfg(
+        prim_path="{ENV_REGEX_NS}/Body",
+        spawn=sim_utils.CuboidCfg(
+            size=(0.1, 0.1, 0.1),
+            rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True, disable_gravity=True),
+        ),
+        init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.5)),
+    )
+    imu: ImuCfg | None = None
+    pva: PvaCfg | None = None
+
+
+def _percentile(samples: list[float], percentile: float) -> float:
+    """Return a linearly interpolated percentile for scalar samples."""
+    ordered = sorted(samples)
+    position = (len(ordered) - 1) * percentile
+    lower = int(position)
+    upper = min(lower + 1, len(ordered) - 1)
+    fraction = position - lower
+    return ordered[lower] + fraction * (ordered[upper] - ordered[lower])
+
+
+def main() -> None:
+    """Run the selected sensor benchmark and print latency statistics."""
+    sim_dt = 1.0 / 120.0
+    sim_cfg = sim_utils.SimulationCfg(dt=sim_dt, device=args_cli.device, gravity=(0.0, 0.0, -9.81))
+    sim = sim_utils.SimulationContext(sim_cfg)
+
+    scene_cfg = ImuPvaBenchmarkSceneCfg(
+        num_envs=args_cli.num_envs,
+        env_spacing=1.0,
+        lazy_sensor_update=True,
+    )
+    if args_cli.sensor == "imu":
+        scene_cfg.imu = ImuCfg(prim_path="{ENV_REGEX_NS}/Body")
+    else:
+        scene_cfg.pva = PvaCfg(prim_path="{ENV_REGEX_NS}/Body")
+
+    scene = InteractiveScene(scene_cfg)
+    sim.reset()
+    scene.reset()
+    sensor = scene[args_cli.sensor]
+
+    if args_cli.disable_recorded_launch:
+        sensor._use_recorded_launch = False
+    for _ in range(args_cli.warmup_steps):
+        sim.step(render=False)
+        sensor.update(sim_dt, force_recompute=True)
+    wp.synchronize_device(sim.device)
+
+    synchronized_ms: list[float] = []
+    submission_ms: list[float] = []
+    for _ in range(args_cli.num_steps):
+        sim.step(render=False)
+        wp.synchronize_device(sim.device)
+        start = time.perf_counter()
+        sensor.update(sim_dt, force_recompute=True)
+        submitted = time.perf_counter()
+        wp.synchronize_device(sim.device)
+        finished = time.perf_counter()
+        synchronized_ms.append((finished - start) * 1000.0)
+        submission_ms.append((submitted - start) * 1000.0)
+
+    lin_acc_b = sensor.data.lin_acc_b.torch
+    ang_vel_b = sensor.data.ang_vel_b.torch
+    finite_instances = int((torch.isfinite(lin_acc_b).all(dim=-1) & torch.isfinite(ang_vel_b).all(dim=-1)).sum())
+
+    print("-" * 80)
+    print(f"{args_cli.sensor.upper()} update benchmark (PhysX)")
+    print(f"  label                  : {args_cli.label}")
+    print(f"  device                 : {sim.device}")
+    print(f"  num_envs               : {args_cli.num_envs}")
+    print(f"  num_steps              : {args_cli.num_steps}")
+    print(f"  synchronized mean      : {statistics.mean(synchronized_ms):.3f} ms")
+    print(f"  synchronized p50       : {_percentile(synchronized_ms, 0.50):.3f} ms")
+    print(f"  synchronized p95       : {_percentile(synchronized_ms, 0.95):.3f} ms")
+    print(f"  submission mean        : {statistics.mean(submission_ms):.3f} ms")
+    print(f"  submission p50         : {_percentile(submission_ms, 0.50):.3f} ms")
+    print(f"  submission p95         : {_percentile(submission_ms, 0.95):.3f} ms")
+    print("-" * 80)
+    print(f"  finite sensor outputs  : {finite_instances} / {args_cli.num_envs}")
+
+
+if __name__ == "__main__":
+    main()
+    simulation_app.close()

source/isaaclab_physx/changelog.d/physx-imu-pva-record-launch.rst

@@ -0,0 +1,4 @@
+Fixed
+^^^^^
+
+* Fixed excessive PhysX IMU and PVA update overhead by reusing typed physics buffers and recorded kernel launches.

source/isaaclab_physx/isaaclab_physx/sensors/imu/imu.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import logging
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
@@ -22,6 +23,8 @@
 if TYPE_CHECKING:
     from isaaclab.sensors.imu import ImuCfg
 
+logger = logging.getLogger(__name__)
+
 
 class Imu(BaseImu):
     """The PhysX Inertial Measurement Unit (IMU) sensor.
@@ -62,6 +65,12 @@ def __init__(self, cfg: ImuCfg):
         super().__init__(cfg)
         self._data = ImuData()
         self._rigid_parent_expr: str | None = None
+        self._raw_transforms: wp.array | None = None
+        self._raw_velocities: wp.array | None = None
+        self._raw_coms: wp.array | None = None
+        self._update_cmd: wp.Launch | None = None
+        self._update_env_mask: wp.array | None = None
+        self._use_recorded_launch: bool = False
 
     def __str__(self) -> str:
         """Returns: A string containing information about the instance."""
@@ -105,10 +114,6 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None
             device=self._device,
         )
 
-    def update(self, dt: float, force_recompute: bool = False):
-        self._dt = dt
-        super().update(dt, force_recompute)
-
     """
     Implementation.
     """
@@ -148,32 +153,77 @@ def _initialize_impl(self):
             self._offset_pos_b = wp.from_torch(composed_p.contiguous(), dtype=wp.vec3f)
             self._offset_quat_b = wp.from_torch(composed_q.contiguous(), dtype=wp.quatf)
 
+        self._use_recorded_launch = wp.get_device(self._device).is_cuda
+
     def _update_buffers_impl(self, env_mask: wp.array | None = None):
         """Fills the buffers of the sensor data."""
         env_mask = self._resolve_indices_and_mask(None, env_mask)
 
-        transforms = self._view.get_transforms().view(wp.transformf)
-        velocities = self._view.get_velocities().view(wp.spatial_vectorf)
-        wp.copy(self._coms_buffer, self._view.get_coms().view(wp.transformf))
-
-        wp.launch(
+        # Refresh the PhysX buffers every update, but create their typed Warp views only once:
+        # the getters lazily allocate their output buffers and refresh the same memory in place
+        # on every call, so the cached views (and the recorded launch that consumes them) stay
+        # valid. A re-backed buffer would silently freeze the sensor data, so fail loudly.
+        transforms = self._view.get_transforms()
+        velocities = self._view.get_velocities()
+        coms = self._view.get_coms()
+        if self._raw_transforms is None:
+            self._raw_transforms = transforms.view(wp.transformf)
+            self._raw_velocities = velocities.view(wp.spatial_vectorf)
+            self._raw_coms = coms.view(wp.transformf)
+        elif (
+            transforms.ptr != self._raw_transforms.ptr
+            or velocities.ptr != self._raw_velocities.ptr
+            or coms.ptr != self._raw_coms.ptr
+        ):
+            raise RuntimeError(
+                f"A PhysX rigid body buffer of the sensor at '{self.cfg.prim_path}' was re-allocated"
+                " after its warp view was cached. The cached views and the recorded launch require"
+                " pointer-stable buffers refreshed in place."
+            )
+        wp.copy(self._coms_buffer, self._raw_coms)
+
+        if self._use_recorded_launch:
+            if self._update_cmd is None:
+                try:
+                    self._update_cmd = self._launch_update(env_mask, record_cmd=True)
+                    self._update_env_mask = env_mask
+                except Exception as exc:
+                    self._use_recorded_launch = False
+                    logger.warning(
+                        f"Failed to record the update of the IMU at '{self.cfg.prim_path}'."
+                        f" Falling back to eager kernel launches. Reason: {exc}"
+                    )
+            if self._update_cmd is not None:
+                if env_mask is not self._update_env_mask:
+                    self._update_cmd.set_param_by_name("env_mask", env_mask)
+                    self._update_env_mask = env_mask
+                self._update_cmd.launch()
+                return
+
+        self._launch_update(env_mask)
+
+    def _launch_update(self, env_mask: wp.array, record_cmd: bool = False) -> wp.Launch | None:
+        """Launch or record the kernel that updates the IMU data."""
+
+        return wp.launch(
             imu_update_kernel,
             dim=self._num_envs,
             inputs=[
                 env_mask,
-                transforms,
-                velocities,
+                self._raw_transforms,
+                self._raw_velocities,
                 self._coms_buffer,
                 self._offset_pos_b,
                 self._offset_quat_b,
                 self._gravity_bias_w,
-                1.0 / self._dt,
                 self._timestamp,
+                self._timestamp_last_update,
                 self._prev_lin_vel_w,
                 self._data._ang_vel_b,
                 self._data._lin_acc_b,
             ],
             device=self._device,
+            record_cmd=record_cmd,
         )
 
     def _initialize_buffers_impl(self):
@@ -188,3 +238,13 @@ def _initialize_buffers_impl(self):
         self._offset_quat_b = wp.from_torch(offset_quat_torch.contiguous(), dtype=wp.quatf)
 
         self._coms_buffer = wp.zeros(self._view.count, dtype=wp.transformf, device=self._device)
+
+    def _invalidate_initialize_callback(self, event):
+        """Invalidate the sensor and release cached PhysX and launch state."""
+        super()._invalidate_initialize_callback(event)
+        self._view = None
+        self._raw_transforms = None
+        self._raw_velocities = None
+        self._raw_coms = None
+        self._update_cmd = None
+        self._update_env_mask = None

source/isaaclab_physx/isaaclab_physx/sensors/imu/kernels.py

@@ -16,8 +16,8 @@ def imu_update_kernel(
     offset_pos_b: wp.array(dtype=wp.vec3f),
     offset_quat_b: wp.array(dtype=wp.quatf),
     gravity_bias_w: wp.array(dtype=wp.vec3f),
-    inv_dt: wp.float32,
     timestamp: wp.array(dtype=wp.float32),
+    timestamp_last_update: wp.array(dtype=wp.float32),
     # inputs / outputs
     prev_lin_vel_w: wp.array(dtype=wp.vec3f),
     # outputs
@@ -34,8 +34,8 @@ def imu_update_kernel(
         offset_pos_b: Offset positions of the sensors.
         offset_quat_b: Offset quaternions of the sensors.
         gravity_bias_w: Gravity bias in the world frame.
-        inv_dt: Inverse of the time step.
         timestamp: Timestamp of the environment.
+        timestamp_last_update: Timestamp of the previous sensor sample.
         prev_lin_vel_w: Previous linear velocity in the world frame.
         out_ang_vel_b: Output angular velocity in the body frame.
         out_lin_acc_b: Output linear acceleration in the body frame.
@@ -49,6 +49,11 @@ def imu_update_kernel(
     if timestamp[idx] == 0.0:
         return
 
+    elapsed_time = timestamp[idx] - timestamp_last_update[idx]
+    if elapsed_time <= 0.0:
+        return
+    inv_dt = 1.0 / elapsed_time
+
     body_quat = wp.transform_get_rotation(transforms[idx])
 
     lin_vel_w = wp.spatial_top(velocities[idx])

source/isaaclab_physx/isaaclab_physx/sensors/pva/kernels.py

@@ -16,8 +16,8 @@ def pva_update_kernel(
     offset_pos_b: wp.array(dtype=wp.vec3f),
     offset_quat_b: wp.array(dtype=wp.quatf),
     gravity_vec_w: wp.array(dtype=wp.vec3f),
-    inv_dt: wp.float32,
     timestamp: wp.array(dtype=wp.float32),
+    timestamp_last_update: wp.array(dtype=wp.float32),
     # inputs / outputs
     prev_lin_vel_w: wp.array(dtype=wp.vec3f),
     prev_ang_vel_w: wp.array(dtype=wp.vec3f),
@@ -40,8 +40,8 @@ def pva_update_kernel(
         offset_pos_b: Offset positions of the sensors.
         offset_quat_b: Offset quaternions of the sensors.
         gravity_vec_w: Gravity direction unit vector in the world frame.
-        inv_dt: Inverse of the time step.
         timestamp: Timestamp of the environment.
+        timestamp_last_update: Timestamp of the previous sensor sample.
         prev_lin_vel_w: Previous linear velocity in the world frame.
         prev_ang_vel_w: Previous angular velocity in the world frame.
         out_pos_w: Output position in the world frame.
@@ -61,6 +61,11 @@ def pva_update_kernel(
     if timestamp[idx] == 0.0:
         return
 
+    elapsed_time = timestamp[idx] - timestamp_last_update[idx]
+    if elapsed_time <= 0.0:
+        return
+    inv_dt = 1.0 / elapsed_time
+
     body_pos = wp.transform_get_translation(transforms[idx])
     body_quat = wp.transform_get_rotation(transforms[idx])