fixup! [FEATURE] Speed up and add noise options for tactile sensors: hysteresis, dead taxels, probe gain

Author: duburcqa

genesis/engine/sensors/kinematic_tactile.py

@@ -24,7 +24,7 @@
     triangle_face_normal,
 )
 
-from .raycaster import ensure_solver_bvhs, update_solver_bvhs
+from .raycaster import RaycastContext
 
 from .base_sensor import RigidSensorMetadataMixin, RigidSensorMixin, SimpleSensor, SimpleSensorMetadata
 from .probe import (
@@ -914,12 +914,13 @@ def build(self):
         # SDF activation is the default fallback; the idempotent call keeps it active even in mixed-mode scenes.
         self._shared_metadata.solver.collider.activate_sdf()
 
-        # Last-value-wins propagation of contact_depth_query. Build the shared collision BVH lazily on raycast opt-in.
+        # Last-value-wins propagation of contact_depth_query. In raycast mode, activate the shared RaycastContext so it
+        # builds its BVH (only on raycast opt-in); SDF-mode scenes never pay for it.
         mode = self._options.contact_depth_query
         if mode is not None:
             self._shared_metadata.contact_depth_query = mode
         if self._shared_metadata.contact_depth_query == "raycast":
-            self._shared_metadata.collision_bvh = ensure_solver_bvhs(self._manager._sim)
+            self._shared_context.activate()
 
 
 @dataclass
@@ -939,7 +940,7 @@ class ContactDepthProbeSensor(
     KinematicTactileSensorMixin[ContactDepthProbeMetadata],
     ProbesWithNormalSensorMixin[ContactDepthProbeMetadata],
     RigidSensorMixin[ContactDepthProbeMetadata],
-    SimpleSensor[ContactDepthProbeOptions, None, ContactDepthProbeMetadata, tuple],
+    SimpleSensor[ContactDepthProbeOptions, RaycastContext, ContactDepthProbeMetadata, tuple],
 ):
     """
     Returns contact depth in meters per probe.
@@ -965,7 +966,7 @@ def _get_cache_dtype(cls) -> torch.dtype:
     @classmethod
     def _update_current_timestep_data(
         cls,
-        shared_context: None,
+        shared_context: RaycastContext,
         shared_metadata: ContactDepthProbeMetadata,
         current_ground_truth_data_T: torch.Tensor,
         ground_truth_data_timeline: "TensorRingBuffer | None",
@@ -1013,8 +1014,7 @@ def _update_current_timestep_data(
                 solver.collider._collider_state,
                 shared_metadata.sensor_candidate_geom_mask,
             )
-            update_solver_bvhs(shared_metadata.collision_bvh)
-            rigid_entry = next(e for e in shared_metadata.collision_bvh if e.raycast_mask is None)
+            rigid_entry = next(e for e in shared_context.bvh_contexts if e.raycast_mask is None)
             _kernel_contact_depth_probe_bvh(
                 shared_metadata.probe_positions,
                 shared_metadata.probe_local_normal,
@@ -1059,7 +1059,9 @@ class ContactProbeMetadata(ContactDepthProbeMetadata):
     release_threshold_row: torch.Tensor = make_tensor_field((0,))
 
 
-class ContactProbeSensor(ContactDepthProbeSensor, SimpleSensor[ContactProbeOptions, None, ContactProbeMetadata, tuple]):
+class ContactProbeSensor(
+    ContactDepthProbeSensor, SimpleSensor[ContactProbeOptions, RaycastContext, ContactProbeMetadata, tuple]
+):
     """
     Returns boolean contact per probe with optional Schmitt-trigger hysteresis. Shares the depth-probe kernel.
 
@@ -1321,7 +1323,7 @@ class KinematicTaxelSensor(
     KinematicTactileSensorMixin[KinematicTaxelMetadata],
     ProbesWithNormalSensorMixin[KinematicTaxelMetadata],
     RigidSensorMixin[KinematicTaxelMetadata],
-    SimpleSensor[KinematicTaxelOptions, None, KinematicTaxelMetadata, KinematicTaxelReturnType],
+    SimpleSensor[KinematicTaxelOptions, RaycastContext, KinematicTaxelMetadata, KinematicTaxelReturnType],
 ):
     """Kinematic taxels: spring-damper force and torque per probe from contact geometry and relative motion."""
 
@@ -1409,7 +1411,7 @@ def _get_cache_dtype(cls) -> torch.dtype:
     @classmethod
     def _update_current_timestep_data(
         cls,
-        shared_context: None,
+        shared_context: RaycastContext,
         shared_metadata: KinematicTaxelMetadata,
         current_ground_truth_data_T: torch.Tensor,
         ground_truth_data_timeline: "TensorRingBuffer | None",
@@ -1472,8 +1474,7 @@ def _update_current_timestep_data(
                 solver.collider._collider_state,
                 shared_metadata.sensor_candidate_geom_mask,
             )
-            update_solver_bvhs(shared_metadata.collision_bvh)
-            rigid_entry = next(e for e in shared_metadata.collision_bvh if e.raycast_mask is None)
+            rigid_entry = next(e for e in shared_context.bvh_contexts if e.raycast_mask is None)
             _kernel_kinematic_taxel_bvh(
                 shared_metadata.probe_positions,
                 shared_metadata.probe_local_normal,

genesis/engine/sensors/point_cloud_tactile.py

@@ -29,7 +29,7 @@
     func_noised_probe_radius,
     get_measured_bufs,
 )
-from .raycaster import ensure_solver_bvhs, update_solver_bvhs
+from .raycaster import RaycastContext
 from .tactile_shared import (
     BVH_LEAF_SIZE,
     BVH_STACK_SIZE,
@@ -1754,7 +1754,7 @@ class ElastomerTaxelSensor(
     PointCloudTactileSensorMixin[ElastomerTaxelSensorMetadata],
     ProbesWithNormalSensorMixin[ElastomerTaxelSensorMetadata],
     RigidSensorMixin[ElastomerTaxelSensorMetadata],
-    SimpleSensor[ElastomerTaxelSensorOptions, None, ElastomerTaxelSensorMetadata],
+    SimpleSensor[ElastomerTaxelSensorOptions, RaycastContext, ElastomerTaxelSensorMetadata],
 ):
     def __init__(
         self,
@@ -1796,14 +1796,14 @@ def build(self):
             gs.raise_exception("ElastomerTaxel must be attached to a rigid link with collision geometry.")
 
         # Last-value-wins propagation of contact_depth_query. Skip activate_sdf only when this sensor explicitly opts
-        # into raycast -- with last-wins, mixed-mode scenes may end up with SDF activated unused (harmless).
+        # into raycast - with last-wins, mixed-mode scenes may end up with SDF activated unused (harmless).
         mode = self._options.contact_depth_query
         if mode is not None:
             self._shared_metadata.contact_depth_query = mode
         if mode != "raycast":
             solver.collider.activate_sdf()
         if self._shared_metadata.contact_depth_query == "raycast":
-            self._shared_metadata.collision_bvh = ensure_solver_bvhs(self._manager._sim)
+            self._shared_context.activate()
 
         elastomer_geom_start_row = self._shared_metadata.elastomer_geom_idx.shape[0]
         elastomer_geom_idx, elastomer_geom_active_envs_mask = _collect_collision_geom_idx(
@@ -2021,7 +2021,7 @@ def _apply_transform(
     @classmethod
     def _update_current_timestep_data(
         cls,
-        shared_context: None,
+        shared_context: RaycastContext,
         shared_metadata: ElastomerTaxelSensorMetadata,
         current_ground_truth_data_T: torch.Tensor,
         ground_truth_data_timeline: "TensorRingBuffer | None",
@@ -2051,8 +2051,7 @@ def _update_current_timestep_data(
                 shared_metadata.probe_depth_buf,
             )
         else:
-            update_solver_bvhs(shared_metadata.collision_bvh)
-            rigid_entry = next(e for e in shared_metadata.collision_bvh if e.raycast_mask is None)
+            rigid_entry = next(e for e in shared_context.bvh_contexts if e.raycast_mask is None)
             _kernel_elastomer_probe_depth_bvh(
                 shared_metadata.probe_positions,
                 shared_metadata.probe_sensor_idx,
@@ -2136,7 +2135,7 @@ def _update_current_timestep_data(
                     shared_metadata.surface_candidate_buf,
                 )
             else:
-                rigid_entry = next(e for e in shared_metadata.collision_bvh if e.raycast_mask is None)
+                rigid_entry = next(e for e in shared_context.bvh_contexts if e.raycast_mask is None)
                 _kernel_elastomer_surface_state_via_global_bvh(
                     shared_metadata.links_idx,
                     shared_metadata.sensor_elastomer_geom_start,

genesis/engine/sensors/raycaster.py

@@ -215,78 +215,6 @@ def destroy(self):
         self._bvh_contexts.clear()
 
 
-# Per-process cache of collision/visual BVH lists keyed by ``id(sim)``, used by tactile-probe sensors in
-# ``contact_depth_query="raycast"`` mode to lazily build (only on raycast opt-in) and share one BVH per simulator.
-# Their mode is a runtime option, so the BVH cannot be declared as a static ``RaycastContext``; it is built on demand.
-_SIM_BVHS: dict[int, list[BVHContext]] = {}
-
-
-def ensure_solver_bvhs(sim) -> list[BVHContext]:
-    """
-    Idempotently build the per-(solver, mesh-type) collision/visual BVHs for ``sim`` and return the shared list.
-
-    Both ``RaycasterSensor`` and tactile-probe sensors in ``contact_depth_query="raycast"`` mode reuse the same list
-    so they share the BVH state for a given sim.
-    """
-    sim_key = id(sim)
-    existing = _SIM_BVHS.get(sim_key)
-    if existing is not None:
-        return existing
-
-    solver_bvhs: list[BVHContext] = []
-    for solver in (sim.rigid_solver, sim.kinematic_solver):
-        if not solver.is_active:
-            continue
-        n_envs = solver._B
-        if isinstance(solver, RigidSolver):
-            n_faces = solver.faces_info.geom_idx.shape[0]
-            aabb = AABB(n_batches=n_envs, n_aabbs=n_faces)
-            bvh = LBVH(aabb, max_n_query_result_per_aabb=0, n_radix_sort_groups=64)
-            solver_bvhs.append(BVHContext(solver, bvh, aabb, None))
-        n_vfaces = solver.vfaces_info.vgeom_idx.shape[0]
-        if n_vfaces > 0:
-            mask = RaycastContext._compute_visual_raycast_mask(solver)
-            if mask.any():
-                aabb = AABB(n_batches=n_envs, n_aabbs=n_vfaces)
-                bvh = LBVH(aabb, max_n_query_result_per_aabb=0, n_radix_sort_groups=64)
-                solver_bvhs.append(BVHContext(solver, bvh, aabb, mask))
-
-    _SIM_BVHS[sim_key] = solver_bvhs
-    return solver_bvhs
-
-
-def update_solver_bvhs(solver_bvhs: list[BVHContext]) -> None:
-    """Rebuild ``solver_bvhs`` from current scene state. Cheap dedup is intentionally NOT applied here: callers can
-    invoke this between scene mutations (e.g. ``set_pos``) without an intervening simulator substep, so any cross-call
-    deduplication would mask real geometry changes."""
-    for entry in solver_bvhs:
-        if entry.raycast_mask is None:
-            kernel_update_verts_and_aabbs(
-                geoms_info=entry.solver.geoms_info,
-                geoms_state=entry.solver.geoms_state,
-                verts_info=entry.solver.verts_info,
-                faces_info=entry.solver.faces_info,
-                free_verts_state=entry.solver.free_verts_state,
-                fixed_verts_state=entry.solver.fixed_verts_state,
-                links_info=entry.solver.links_info,
-                static_rigid_sim_config=entry.solver._static_rigid_sim_config,
-                aabb_state=entry.aabb,
-            )
-            entry.bvh.build()
-        else:
-            entry.solver.update_forward_pos()
-            entry.solver.update_vgeoms()
-            kernel_update_visual_aabbs(
-                vverts_info=entry.solver.vverts_info,
-                vverts_state=entry.solver.vverts_state,
-                vfaces_info=entry.solver.vfaces_info,
-                vgeoms_state=entry.solver.vgeoms_state,
-                face_mask=entry.raycast_mask,
-                aabb_state=entry.aabb,
-            )
-            entry.bvh.build()
-
-
 @dataclass
 class RaycasterSharedMetadata(KinematicSensorMetadataMixin, SimpleSensorMetadata):
     # The BVHs cast against each frame live on the shared ``RaycastContext`` (one per active solver per mesh type),

genesis/engine/sensors/tactile_shared.py

@@ -11,7 +11,6 @@
 from genesis.utils.misc import concat_with_tensor, make_tensor_field
 
 if TYPE_CHECKING:
-    from genesis.engine.sensors.raycaster import _SolverBVH
     from genesis.utils.ring_buffer import TensorRingBuffer
 
 
@@ -442,18 +441,18 @@ class ContactDepthQueryMetadataMixin:
     """
     Shared per-sensor-class state for the contact-depth query backend.
 
-    ``contact_depth_query`` is the resolved mode for every sensor of this class -- ``"sdf"`` or ``"raycast"``. Each
+    ``contact_depth_query`` is the resolved mode for every sensor of this class - ``"sdf"`` or ``"raycast"``. Each
     sensor's ``build()`` overwrites the field with its own option when non-``None`` (last-value-wins); ``None`` at
     update time falls back to ``"sdf"``.
 
-    When mode is ``"raycast"``, ``collision_bvh`` references the per-sim shared BVH list built by
-    ``ensure_solver_bvhs``. ``sensor_candidate_geom_mask`` is a per-(env, sensor, geom) bool gate -- scattered per
-    step from the contact prefilter (KinematicTactile family) or once at build from ``sensor_track_geom_idx``
-    (ElastomerTaxel) -- so BVH leaves whose ``faces_info.geom_idx`` falls outside the mask are skipped.
+    When mode is ``"raycast"``, the collision/visual BVHs come from the shared ``RaycastContext`` (the sensor's
+    ``shared_context``), built lazily on raycast opt-in and refreshed once per step by ``SensorManager``.
+    ``sensor_candidate_geom_mask`` is a per-(env, sensor, geom) bool gate - scattered per step from the contact
+    prefilter (KinematicTactile family) or once at build from ``sensor_track_geom_idx`` (ElastomerTaxel) - so BVH
+    leaves whose ``faces_info.geom_idx`` falls outside the mask are skipped.
     """
 
     contact_depth_query: str | None = None
-    collision_bvh: "list[_SolverBVH] | None" = None
     sensor_candidate_geom_mask: torch.Tensor = make_tensor_field((0, 0, 0), dtype_factory=lambda: gs.tc_bool)