# 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
from __future__ import annotations
import logging
import re
from typing import TYPE_CHECKING
import numpy as np
import trimesh
import warp as wp
from pxr import Usd, UsdPhysics
import isaaclab.sim as sim_utils
from isaaclab import cloner
from isaaclab.sim.simulation_context import SimulationContext
from isaaclab.utils.mesh import PRIMITIVE_MESH_TYPES, create_trimesh_from_geom_mesh, create_trimesh_from_geom_shape
from isaaclab.utils.warp import ProxyArray, convert_to_warp_mesh
from isaaclab.utils.warp import kernels as warp_kernels
from .base_ray_caster import BaseRayCaster
from .kernels import copy_mesh_poses_to_table_kernel, fill_ray_hits_distance_inf_kernel
from .multi_mesh_ray_caster_data import MultiMeshRayCasterData
if TYPE_CHECKING:
from isaaclab.cloner import ClonePlan
from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg
logger = logging.getLogger(__name__)
def _matrix_from_quat_xyzw(quat: np.ndarray) -> np.ndarray:
"""Return a rotation matrix from an ``(x, y, z, w)`` quaternion."""
x, y, z, w = quat
two_s = 2.0 / np.dot(quat, quat)
return np.array(
[
[1.0 - two_s * (y * y + z * z), two_s * (x * y - z * w), two_s * (x * z + y * w)],
[two_s * (x * y + z * w), 1.0 - two_s * (x * x + z * z), two_s * (y * z - x * w)],
[two_s * (x * z - y * w), two_s * (y * z + x * w), 1.0 - two_s * (x * x + y * y)],
],
dtype=np.float64,
)
[docs]
class BaseMultiMeshRayCaster(BaseRayCaster):
"""A multi-mesh ray-casting sensor.
The ray-caster uses a set of rays to detect collisions with meshes in the scene. The rays are
defined in the sensor's local coordinate frame. The sensor can be configured to ray-cast against
a set of meshes with a given ray pattern.
The meshes are parsed from the list of primitive paths provided in the configuration. These are then
converted to warp meshes and stored in the :attr:`meshes` dictionary. The ray-caster then ray-casts
against these warp meshes using the ray pattern provided in the configuration.
Compared to the default RayCaster, the MultiMeshRayCaster provides additional functionality and flexibility as
an extension of the default RayCaster with the following enhancements:
- Raycasting against multiple target types : Supports primitive shapes (spheres, cubes, etc.) as well as arbitrary
meshes.
- Dynamic mesh tracking : Keeps track of specified meshes, enabling raycasting against moving parts
(e.g., robot links, articulated bodies, or dynamic obstacles).
- Memory-efficient caching : Avoids redundant memory usage by reusing mesh data across environments.
.. warning::
**Known limitation (multi-mesh closest-hit resolution):** When two meshes produce a
hit at the exact same distance for a given ray, the ``atomic_min`` + equality-check
pattern in the raycasting kernel is not fully thread-safe. The hit *position* is always
correct, but auxiliary outputs (normals, face IDs, mesh IDs) may originate from
different meshes for the affected ray. This requires an exact floating-point tie and is
rare in practice. See `warp#1058 <https://github.com/NVIDIA/warp/issues/1058>`_ for
upstream progress on a thread-safe ``atomic_min`` return value.
Example usage to raycast against the visual meshes of a robot (e.g. ANYmal):
.. code-block:: python
ray_caster_cfg = MultiMeshRayCasterCfg(
prim_path="{ENV_REGEX_NS}/Robot",
mesh_prim_paths=[
"/World/Ground",
MultiMeshRayCasterCfg.RaycastTargetCfg(prim_expr="{ENV_REGEX_NS}/Robot/LF_[^/]*/visuals"),
MultiMeshRayCasterCfg.RaycastTargetCfg(prim_expr="{ENV_REGEX_NS}/Robot/RF_[^/]*/visuals"),
MultiMeshRayCasterCfg.RaycastTargetCfg(prim_expr="{ENV_REGEX_NS}/Robot/LH_[^/]*/visuals"),
MultiMeshRayCasterCfg.RaycastTargetCfg(prim_expr="{ENV_REGEX_NS}/Robot/RH_[^/]*/visuals"),
MultiMeshRayCasterCfg.RaycastTargetCfg(prim_expr="{ENV_REGEX_NS}/Robot/base/visuals"),
],
ray_alignment="world",
pattern_cfg=patterns.GridPatternCfg(resolution=0.02, size=(2.5, 2.5), direction=(0, 0, -1)),
)
"""
cfg: MultiMeshRayCasterCfg
"""The configuration parameters."""
[docs]
def __init__(self, cfg: MultiMeshRayCasterCfg):
"""Initializes the ray-caster object.
Args:
cfg: The configuration parameters.
"""
super().__init__(cfg)
self._num_meshes_per_env: dict[str, int] = {}
self._raycast_targets_cfg: list[MultiMeshRayCasterCfg.RaycastTargetCfg] = []
for target in self.cfg.mesh_prim_paths:
if isinstance(target, str):
target_cfg = cfg.RaycastTargetCfg(prim_expr=target, track_mesh_transforms=False)
else:
target_cfg = target
target_cfg.prim_expr = cloner.expand_env_regex_ns(target_cfg.prim_expr)
self._raycast_targets_cfg.append(target_cfg)
self._data = MultiMeshRayCasterData()
def __str__(self) -> str:
"""Returns: A string containing information about the instance."""
return (
f"Ray-caster @ '{self.cfg.prim_path}': \n"
f"\tview type : {self._view.__class__}\n"
f"\tupdate period (s) : {self.cfg.update_period}\n"
f"\tnumber of meshes : {self._num_envs} x {sum(self._num_meshes_per_env.values())} \n"
f"\tnumber of sensors : {self._view_count}\n"
f"\tnumber of rays/sensor: {self.num_rays}\n"
f"\ttotal number of rays : {self.num_rays * self._view_count}"
)
"""
Properties
"""
@property
def data(self) -> MultiMeshRayCasterData:
self._update_outdated_buffers()
return self._data
"""
Implementation.
"""
def _initialize_warp_meshes(self):
"""Initialize mesh buffers from the ClonePlan when env-scoped, else from the stage."""
sim = SimulationContext.instance()
plan = sim.get_clone_plan() if sim is not None else None
target_records_by_expr = {}
dummy_mesh_id: int | None = None
self._mesh_views = []
# Build one per-env mesh list for each configured raycast target.
for target_cfg in self._raycast_targets_cfg:
records_per_env, dummy_mesh_id, tracked_target_exprs = self._build_mesh_records(
target_cfg, plan, dummy_mesh_id
)
self._num_meshes_per_env[target_cfg.prim_expr] = max(len(records) for records in records_per_env)
target_records_by_expr[target_cfg.prim_expr] = records_per_env
self._mesh_views.append(
self._create_tracked_target_view(tracked_target_exprs) if target_cfg.track_mesh_transforms else None
)
if dummy_mesh_id is None:
raise RuntimeError(
f"No meshes found for ray-casting! Please check the mesh prim paths: {self.cfg.mesh_prim_paths}"
)
total_meshes_per_env = sum(
self._num_meshes_per_env[target_cfg.prim_expr] for target_cfg in self._raycast_targets_cfg
)
mesh_ids = np.full((self._num_envs, total_meshes_per_env), dummy_mesh_id, dtype=np.uint64)
mesh_positions = np.full((self._num_envs, total_meshes_per_env, 3), 1.0e9, dtype=np.float32)
mesh_orientations = np.zeros((self._num_envs, total_meshes_per_env, 4), dtype=np.float32)
mesh_orientations[..., 3] = 1.0
mesh_offset = 0
for target_cfg in self._raycast_targets_cfg:
records_per_env = target_records_by_expr[target_cfg.prim_expr]
target_width = self._num_meshes_per_env[target_cfg.prim_expr]
for env_id, records in enumerate(records_per_env):
if not records:
continue
count = len(records)
record_mesh_ids, record_positions, record_orientations = zip(*records)
target_slice = slice(mesh_offset, mesh_offset + count)
mesh_ids[env_id, target_slice] = np.asarray(record_mesh_ids, dtype=np.uint64)
mesh_positions[env_id, target_slice] = np.asarray(record_positions, dtype=np.float32)
mesh_orientations[env_id, target_slice] = np.asarray(record_orientations, dtype=np.float32)
mesh_offset += target_width
self._mesh_ids_wp = wp.array2d(mesh_ids, dtype=wp.uint64, device=self.device)
self._mesh_positions_w = wp.array2d(mesh_positions, dtype=wp.vec3f, device=self.device)
self._mesh_orientations_w = wp.array2d(mesh_orientations, dtype=wp.quatf, device=self.device)
def _build_mesh_records(
self,
target_cfg: MultiMeshRayCasterCfg.RaycastTargetCfg,
plan: ClonePlan | None,
dummy_mesh_id: int | None,
):
"""Build mesh records for the target configuration."""
records_per_env = [[] for _ in range(self._num_envs)]
target_in_plan = False
tracked_target_exprs: list[str] = [target_cfg.prim_expr]
has_rigid_body_api = lambda p: p.HasAPI(UsdPhysics.RigidBodyAPI) # noqa: E731
# Prefer ClonePlan data for env-scoped targets; destination USD prims may not exist.
if plan is not None and target_cfg.track_mesh_transforms:
plan_tracked_target_exprs: list[str] = []
prim_expr = target_cfg.prim_expr
for source_root, destination_template, source_path, env_ids in cloner.query.iter_sources(plan, prim_expr):
target_in_plan = True
# Load meshes from the authored source entry.
source_pattern = re.compile(source_path)
source_prims = sim_utils.get_all_matching_child_prims(
source_root, lambda prim: source_pattern.fullmatch(prim.GetPath().pathString) is not None
)
if not source_prims:
raise RuntimeError(f"No ClonePlan source prims matched '{source_path}'.")
mesh_ids: list[int] = []
row_tracked_target_exprs: list[str] = []
for source_prim in source_prims:
source_prim_path = str(source_prim.GetPath())
# Use a bounded rigid-body ancestor when the match is below a body; otherwise
# enumerate rigid bodies under the match, including the matched prim itself.
owner_prim = None
rigid_body_records: list[tuple[Usd.Prim, Usd.Prim]]
if not source_prim.HasAPI(UsdPhysics.ArticulationRootAPI):
owner_prim = sim_utils.get_first_matching_ancestor_prim(source_prim_path, has_rigid_body_api)
owner_path = None if owner_prim is None else str(owner_prim.GetPath())
if owner_path and (owner_path == source_root or owner_path.startswith(f"{source_root}/")):
rigid_body_records = [(source_prim, owner_prim)]
else:
rigid_body_records = [
(p, p) for p in sim_utils.get_all_matching_child_prims(source_prim_path, has_rigid_body_api)
]
if not rigid_body_records:
raise RuntimeError(
f"Cannot track ClonePlan target '{target_cfg.prim_expr}' because source prim "
f"'{source_prim.GetPath()}' has no rigid-body ancestor or descendant."
)
for geometry_prim, owner_prim in rigid_body_records:
mesh_id = self._load_target_prim_warp_mesh(geometry_prim, target_cfg, reference_prim=owner_prim)
dummy_mesh_id = mesh_id if dummy_mesh_id is None else dummy_mesh_id
mesh_ids.append(mesh_id)
owner_path = str(owner_prim.GetPath())
if owner_path == source_root:
owner_suffix = ""
elif owner_path.startswith(source_root + "/"):
owner_suffix = owner_path[len(source_root) :]
else:
raise RuntimeError(
f"Tracked target owner '{owner_path}' is not under ClonePlan source root "
f"'{source_root}'."
)
row_tracked_target_exprs.append(destination_template.format("[^/]+") + owner_suffix)
if len(row_tracked_target_exprs) > len(plan_tracked_target_exprs):
plan_tracked_target_exprs = row_tracked_target_exprs
# Geometry is selected by ClonePlan; live pose is supplied by backend body/site views.
for env_id in env_ids:
for mesh_id in mesh_ids:
records_per_env[env_id].append((mesh_id, (1.0e9, 1.0e9, 1.0e9), (0.0, 0.0, 0.0, 1.0)))
if target_in_plan:
if not plan_tracked_target_exprs:
raise RuntimeError(
f"No tracked body expressions were resolved for target '{target_cfg.prim_expr}'."
)
return records_per_env, dummy_mesh_id, plan_tracked_target_exprs
# Fall back to authored USD prims for global targets and scenes without ClonePlan data.
target_prims = sim_utils.find_matching_prims(target_cfg.prim_expr)
if not target_prims:
raise RuntimeError(f"Failed to find a prim at path expression: {target_cfg.prim_expr}")
records = []
tracked_target_exprs = []
for target_prim in target_prims:
reference_prim = target_prim
if target_cfg.track_mesh_transforms:
while reference_prim and reference_prim.IsValid() and str(reference_prim.GetPath()) != "/":
if reference_prim.HasAPI(UsdPhysics.RigidBodyAPI):
break
reference_prim = reference_prim.GetParent()
if reference_prim is None or not reference_prim.IsValid() or not has_rigid_body_api(reference_prim):
raise RuntimeError(
f"Cannot track non-physics ray-cast target '{target_cfg.prim_expr}'. "
"Set track_mesh_transforms=False for static targets, or apply RigidBodyAPI to dynamic targets."
)
tracked_target_exprs.append(str(reference_prim.GetPath()))
mesh_id = self._load_target_prim_warp_mesh(target_prim, target_cfg, reference_prim=reference_prim)
dummy_mesh_id = mesh_id if dummy_mesh_id is None else dummy_mesh_id
records.append((mesh_id, *sim_utils.resolve_prim_pose(reference_prim)))
if len(records) == 1:
return [list(records) for _ in range(self._num_envs)], dummy_mesh_id, tracked_target_exprs
# Multiple USD matches are expected to be laid out evenly by environment.
if len(records) % self._num_envs != 0:
raise RuntimeError(
f"Target expression '{target_cfg.prim_expr}' matched {len(records)} mesh records, "
f"which cannot be evenly partitioned across {self._num_envs} environments."
)
n_meshes = len(records) // self._num_envs
records_per_env = [records[env_id * n_meshes : (env_id + 1) * n_meshes] for env_id in range(self._num_envs)]
return records_per_env, dummy_mesh_id, tracked_target_exprs
def _load_target_prim_warp_mesh(self, target_prim, target_cfg, reference_prim=None) -> int:
reference_prim = target_prim if reference_prim is None else reference_prim
prim_key = (f"{target_prim.GetPath()}@{reference_prim.GetPath()}", self._device)
if prim_key in BaseMultiMeshRayCaster.meshes:
return BaseMultiMeshRayCaster.meshes[prim_key].id
mesh_prims = sim_utils.get_all_matching_child_prims(
target_prim.GetPath(), lambda prim: prim.GetTypeName() in PRIMITIVE_MESH_TYPES + ["Mesh"]
)
if len(mesh_prims) == 0:
raise RuntimeError(
f"No mesh prims found at path: {target_prim.GetPath()} with supported types:"
f" {PRIMITIVE_MESH_TYPES + ['Mesh']}"
)
trimesh_meshes = []
for mesh_prim in mesh_prims:
if mesh_prim is None or not mesh_prim.IsValid():
raise RuntimeError(f"Invalid mesh prim path: {target_prim}")
if mesh_prim.GetTypeName() == "Mesh":
mesh = create_trimesh_from_geom_mesh(mesh_prim)
else:
mesh = create_trimesh_from_geom_shape(mesh_prim)
mesh.apply_scale(sim_utils.resolve_prim_scale(mesh_prim))
relative_pos, relative_quat = sim_utils.resolve_prim_pose(mesh_prim, reference_prim)
relative_pos = np.asarray(relative_pos, dtype=np.float64)
relative_quat = np.asarray(relative_quat, dtype=np.float64)
transform = np.eye(4)
transform[:3, :3] = _matrix_from_quat_xyzw(relative_quat)
transform[:3, 3] = relative_pos
mesh.apply_transform(transform)
trimesh_meshes.append(mesh)
if len(trimesh_meshes) == 1:
trimesh_mesh = trimesh_meshes[0]
elif target_cfg.merge_prim_meshes:
trimesh_mesh = trimesh.util.concatenate(trimesh_meshes)
else:
raise RuntimeError(
f"Multiple mesh prims found at path: {target_prim.GetPath()} but merging is disabled. Please"
" enable `merge_prim_meshes` in the configuration or specify each mesh separately."
)
wp_mesh = convert_to_warp_mesh(trimesh_mesh.vertices, trimesh_mesh.faces, device=self._device)
BaseMultiMeshRayCaster.meshes[prim_key] = wp_mesh
logger.info(
f"Read '{len(mesh_prims)}' mesh prims under path '{target_prim.GetPath()}' with"
f" {len(trimesh_mesh.vertices)} vertices and {len(trimesh_mesh.faces)} faces."
)
return wp_mesh.id
def _create_tracked_target_view(self, target_prim_paths: str | list[str]):
raise NotImplementedError("Tracked multi-mesh targets must be implemented by the active physics backend.")
def _initialize_rays_impl(self):
super()._initialize_rays_impl()
# Persistent buffer for tracking closest-hit distance across meshes (for atomic_min)
self._ray_distance_wp = wp.empty((self._view_count, self.num_rays), dtype=wp.float32, device=self._device)
if self.cfg.update_mesh_ids:
self._data.ray_mesh_ids = ProxyArray(
wp.zeros((self._view_count, self.num_rays), dtype=wp.int16, device=self._device)
)
else:
# Dummy 1×1 buffer so the kernel launch always has a valid array to bind
self._ray_mesh_id_wp = wp.empty((1, 1), dtype=wp.int16, device=self._device)
# Persistent dummy buffers for unused kernel outputs; allocated once to avoid per-step allocations.
self._dummy_normal_wp = wp.empty((1, 1), dtype=wp.vec3, device=self._device)
self._dummy_face_id_wp = wp.empty((1, 1), dtype=wp.int32, device=self._device)
def _update_mesh_transforms(self) -> None:
"""Update world-frame mesh positions and orientations for dynamically tracked targets.
Iterates over all tracked views and writes the current world poses into
the rectangular mesh pose buffers. Static (non-tracked) targets are
skipped; their initial poses were set during :meth:`_initialize_warp_meshes`.
"""
mesh_idx = 0
for view, target_cfg in zip(self._mesh_views, self._raycast_targets_cfg):
if not target_cfg.track_mesh_transforms:
mesh_idx += self._num_meshes_per_env[target_cfg.prim_expr]
continue
pos_w, ori_w = view.get_world_poses(None)
view_count = getattr(view, "count", pos_w.warp.shape[0])
meshes_per_env = view_count
if view_count != 1:
# Backend views return a flat list across envs; the mesh table is indexed per env.
meshes_per_env = view_count // self._num_envs
wp.launch(
copy_mesh_poses_to_table_kernel,
dim=(self._num_envs, meshes_per_env),
inputs=[
pos_w.warp,
ori_w.warp,
int(meshes_per_env),
int(mesh_idx),
bool(view_count == 1),
self._mesh_positions_w,
self._mesh_orientations_w,
],
device=self._device,
)
mesh_idx += self._num_meshes_per_env[target_cfg.prim_expr]
def _update_buffers_impl(self, env_mask: wp.array):
"""Fills the buffers of the sensor data."""
self._update_ray_infos(env_mask)
self._update_mesh_transforms()
# Fill output and distance buffers with inf for masked environments
wp.launch(
fill_ray_hits_distance_inf_kernel,
dim=(self._num_envs, self.num_rays),
inputs=[env_mask, False],
outputs=[self._data._ray_hits_w, self._ray_distance_wp, self._dummy_normal_wp],
device=self._device,
)
n_meshes = self._mesh_ids_wp.shape[1]
return_normal = False
return_face_id = False
write_mesh_ids = self.cfg.update_mesh_ids
# Ray-cast against all meshes; closest hit wins via atomic_min on ray_distance.
wp.launch(
warp_kernels.raycast_dynamic_meshes_kernel,
dim=(n_meshes, self._num_envs, self.num_rays),
inputs=[
env_mask,
self._mesh_ids_wp,
self._ray_starts_w,
self._ray_directions_w,
self._data._ray_hits_w,
self._ray_distance_wp,
self._dummy_normal_wp,
self._dummy_face_id_wp,
self._data.ray_mesh_ids.warp if self.cfg.update_mesh_ids else self._ray_mesh_id_wp,
self._mesh_positions_w,
self._mesh_orientations_w,
float(self.cfg.max_distance),
int(return_normal),
int(return_face_id),
int(write_mesh_ids),
],
device=self._device,
)
def _invalidate_initialize_callback(self, event):
"""Invalidates the scene elements."""
super()._invalidate_initialize_callback(event)
# clear mesh views so they are re-created on the next initialization
self._mesh_views = []