diff --git a/source/isaaclab/changelog.d/octi-raycaster-backend-split.minor.rst b/source/isaaclab/changelog.d/octi-raycaster-backend-split.minor.rst new file mode 100644 index 00000000000..96c6b1d58ec --- /dev/null +++ b/source/isaaclab/changelog.d/octi-raycaster-backend-split.minor.rst @@ -0,0 +1,26 @@ +Added +^^^^^ + +* Added :class:`~isaaclab.sensors.ray_caster.BaseRayCaster`, + :class:`~isaaclab.sensors.ray_caster.BaseRayCasterCamera`, + :class:`~isaaclab.sensors.ray_caster.BaseMultiMeshRayCaster`, and + :class:`~isaaclab.sensors.ray_caster.BaseMultiMeshRayCasterCamera` + carrying the backend-agnostic ray-caster logic. Backend subclasses + override only the body-tracker and target-mesh-tracker hooks. + +Changed +^^^^^^^ + +* **Breaking:** Changed :class:`~isaaclab.sensors.camera.CameraData` + camera-owned buffers to :class:`~isaaclab.utils.warp.ProxyArray`. + Access torch tensor operations through the explicit ``.torch`` view. +* :class:`~isaaclab.sensors.ray_caster.RayCaster`, + :class:`~isaaclab.sensors.ray_caster.RayCasterCamera`, + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCaster`, and + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCasterCamera` are now + :class:`~isaaclab.utils.backend_utils.FactoryBase` shims dispatching + to PhysX / Newton implementations. Cfg surface and runtime semantics + unchanged. +* Changed ray-caster camera update paths to keep pose, ray, depth, normal, + and mesh-id buffers Warp-owned internally, while exposing public camera + outputs through :class:`~isaaclab.utils.warp.ProxyArray`. diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/__init__.pyi b/source/isaaclab/isaaclab/sensors/ray_caster/__init__.pyi index 44b76de2fa2..034c0ff84dd 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/__init__.pyi +++ b/source/isaaclab/isaaclab/sensors/ray_caster/__init__.pyi @@ -4,6 +4,10 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "BaseMultiMeshRayCaster", + "BaseMultiMeshRayCasterCamera", + "BaseRayCaster", + "BaseRayCasterCamera", "MultiMeshRayCaster", "MultiMeshRayCasterCamera", "MultiMeshRayCasterCameraCfg", @@ -18,6 +22,10 @@ __all__ = [ "patterns", ] +from .base_multi_mesh_ray_caster import BaseMultiMeshRayCaster +from .base_multi_mesh_ray_caster_camera import BaseMultiMeshRayCasterCamera +from .base_ray_caster import BaseRayCaster +from .base_ray_caster_camera import BaseRayCasterCamera from .multi_mesh_ray_caster import MultiMeshRayCaster from .multi_mesh_ray_caster_camera import MultiMeshRayCasterCamera from .multi_mesh_ray_caster_camera_cfg import MultiMeshRayCasterCameraCfg diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster.py new file mode 100644 index 00000000000..43920af5833 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster.py @@ -0,0 +1,473 @@ +# 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 UsdPhysics + +import isaaclab.sim as sim_utils +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, + ) + + +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 `_ 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.""" + + 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 = target_cfg.prim_expr.format(ENV_REGEX_NS="/World/envs/env_.*") + 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] + + # Prefer ClonePlan data for env-scoped targets; destination USD prims may not exist. + if plan is not None and target_cfg.track_mesh_transforms: + target_path = re.sub(r"env_\.\*", "env_0", target_cfg.prim_expr) + plan_tracked_target_exprs: list[str] = [] + for row, (source_root, destination_template) in enumerate(zip(plan.sources, plan.destinations)): + if "{}" not in destination_template: + continue + + dest_path = destination_template.format(0) + suffix = target_path.removeprefix(dest_path) + if suffix == target_path or (suffix and not suffix.startswith("/")): + continue + + target_in_plan = True + env_ids = plan.clone_mask[row].nonzero(as_tuple=False).squeeze(-1) + if env_ids.numel() == 0: + continue + + # Load meshes from the authored source row. + source_prims = sim_utils.find_matching_prims(source_root + suffix) + if not source_prims: + raise RuntimeError(f"No ClonePlan source prims matched '{source_root + suffix}'.") + + mesh_ids: list[int] = [] + row_tracked_target_exprs: list[str] = [] + for source_prim in source_prims: + owner_prim = source_prim + while owner_prim and owner_prim.IsValid() and str(owner_prim.GetPath()) != "/": + if owner_prim.HasAPI(UsdPhysics.RigidBodyAPI): + break + owner_prim = owner_prim.GetParent() + if owner_prim is None or not owner_prim.IsValid() or not owner_prim.HasAPI(UsdPhysics.RigidBodyAPI): + raise RuntimeError( + f"Cannot track ClonePlan target '{target_cfg.prim_expr}' because source prim " + f"'{source_prim.GetPath()}' has no rigid-body ancestor." + ) + mesh_id = self._load_target_prim_warp_mesh(source_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 '{source_root}'." + ) + row_tracked_target_exprs.append(destination_template.replace("{}", ".*") + 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.tolist(): + 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 reference_prim.HasAPI(UsdPhysics.RigidBodyAPI) + ): + 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 + pos, quat = sim_utils.resolve_prim_pose(reference_prim) + pos = (float(pos[0]), float(pos[1]), float(pos[2])) + quat = (float(quat[0]), float(quat[1]), float(quat[2]), float(quat[3])) + records.append((mesh_id, pos, quat)) + + 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 = [] diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster_camera.py new file mode 100644 index 00000000000..cc9f660edc6 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster_camera.py @@ -0,0 +1,262 @@ +# 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 + +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp import kernels as warp_kernels + +from . import kernels as ray_caster_kernels +from .base_multi_mesh_ray_caster import BaseMultiMeshRayCaster +from .base_ray_caster_camera import BaseRayCasterCamera +from .multi_mesh_ray_caster_camera_data import MultiMeshRayCasterCameraData + +if TYPE_CHECKING: + from .multi_mesh_ray_caster_camera_cfg import MultiMeshRayCasterCameraCfg + + +class BaseMultiMeshRayCasterCamera(BaseRayCasterCamera, BaseMultiMeshRayCaster): + """A multi-mesh ray-casting camera sensor. + + The ray-caster camera uses a set of rays to get the distances to meshes in the scene. The rays are + defined in the sensor's local coordinate frame. The sensor has the same interface as the + :class:`isaaclab.sensors.Camera` that implements the camera class through USD camera prims. + However, this class provides a faster image generation. The sensor converts meshes from the list of + primitive paths provided in the configuration to Warp meshes. The camera then ray-casts against these + Warp meshes only. + + Currently, only the following annotators are supported: + + - ``"distance_to_camera"``: An image containing the distance to camera optical center. + - ``"distance_to_image_plane"``: An image containing distances of 3D points from camera plane along camera's z-axis. + - ``"normals"``: An image containing the local surface normal vectors at each pixel. + """ + + cfg: MultiMeshRayCasterCameraCfg + """The configuration parameters.""" + + def __init__(self, cfg: MultiMeshRayCasterCameraCfg): + """Initializes the camera object. + + Args: + cfg: The configuration parameters. + + Raises: + ValueError: If the provided data types are not supported by the ray-caster camera. + """ + self._check_supported_data_types(cfg) + # initialize base class + BaseMultiMeshRayCaster.__init__(self, cfg) + # create empty variables for storing output data + self._data = MultiMeshRayCasterCameraData() + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + return ( + f"Multi-Mesh Ray-Caster-Camera @ '{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}\n" + f"\timage shape : {self.image_shape}" + ) + + def _initialize_warp_meshes(self): + # The camera MRO would pick the single-mesh camera path; use the multi-mesh setup. + BaseMultiMeshRayCaster._initialize_warp_meshes(self) + + def _create_buffers(self): + super()._create_buffers() + self._data.image_mesh_ids = ProxyArray( + wp.zeros((self._num_envs, *self.image_shape, 1), dtype=wp.int16, device=self.device) + ) + + def _initialize_rays_impl(self): + # NOTE: This method intentionally does NOT call super()._initialize_rays_impl() through the MRO + # chain. The intermediate classes (RayCasterCamera, MultiMeshRayCaster) use different internal + # buffer names and orderings that are incompatible with the camera's full init path: + # - RayCasterCamera creates single-mesh ray buffers (_ray_distance, _ray_normal_w, etc.) + # - MultiMeshRayCaster creates _ray_distance_wp / _ray_mesh_id_wp for multi-mesh use + # The camera replaces all of these with its own camera-named equivalents below. + # If either parent class gains new shared buffers, they must be added here explicitly. + + # Camera-specific bookkeeping buffers + self._frame_wp = wp.zeros(self._view.count, dtype=wp.int64, device=self._device) + self._frame = wp.to_torch(self._frame_wp) + + # Build camera output buffers (intrinsics, image data, etc.) + self._create_buffers() + self._compute_intrinsic_matrices() + + # Compute local ray starts/directions from the camera pattern (torch, init-time only) + ray_starts_local, ray_directions_local = self.cfg.pattern_cfg.func( + self.cfg.pattern_cfg, self._data.intrinsic_matrices.torch, self._device + ) + self.num_rays = ray_directions_local.shape[1] + + # Store local (sensor-frame) ray arrays as ProxyArrays for Warp kernel dispatch. + self.ray_starts = ProxyArray(wp.from_torch(ray_starts_local.contiguous(), dtype=wp.vec3f)) + self.ray_directions = ProxyArray(wp.from_torch(ray_directions_local.contiguous(), dtype=wp.vec3f)) + + # Camera-frame offset and drift buffers. + self._create_offset_buffers(self._view.count) + + # World-frame ray buffers. + self._ray_starts_w = wp.empty((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + self._ray_directions_w = wp.empty((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + + # Ray hit positions as a warp array; expose a ProxyArray for debug visualisation. + self.ray_hits_w = ProxyArray(wp.empty((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device)) + + # Per-ray closest-hit distance for atomic_min across meshes + self._ray_distance_cam_wp = wp.empty((self._view.count, self.num_rays), dtype=wp.float32, device=self._device) + + # Optional normal buffer (always allocated; filled only when "normals" is requested) + self._ray_normal_w = wp.empty((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + + # Mesh-id buffers from MultiMeshRayCaster._initialize_rays_impl + if self.cfg.update_mesh_ids: + self._ray_mesh_id_wp = wp.zeros((self._view.count, self.num_rays), dtype=wp.int16, device=self._device) + else: + self._ray_mesh_id_wp = wp.empty((1, 1), dtype=wp.int16, device=self._device) + + # Dummy face-id buffer (not used by camera but required by kernel signature) + self._ray_face_id_wp = wp.empty((1, 1), dtype=wp.int32, device=self._device) + + def _update_ray_infos(self, env_mask: wp.array): + """Updates camera poses and world-frame ray buffers for masked environments. + + Args: + env_mask: Boolean mask selecting which environments to update. Shape is (num_envs,). + """ + transforms = self._get_view_transforms_wp() + wp.launch( + ray_caster_kernels.update_ray_caster_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + transforms, + env_mask, + self._offset_pos_wp, + self._offset_quat_wp, + self.drift.warp, + self.ray_cast_drift.warp, + self.ray_starts.warp, + self.ray_directions.warp, + int(ray_caster_kernels.ALIGNMENT_BASE), + ], + outputs=[ + self._data.pos_w.warp, + self._data.quat_w_world.warp, + self._ray_starts_w, + self._ray_directions_w, + ], + device=self._device, + ) + + def _update_buffers_impl(self, env_mask: wp.array): + """Fills the buffers of the sensor data.""" + self._update_ray_infos(env_mask) + + # Increment frame count for updated environments + self._update_frame(env_mask, frame_op=1) + + self._update_mesh_transforms() + + return_normal = "normals" in self.cfg.data_types + + # Fill ray hit, distance, and optional normal buffers with inf for masked environments. + wp.launch( + ray_caster_kernels.fill_ray_hits_distance_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, return_normal], + outputs=[self.ray_hits_w.warp, self._ray_distance_cam_wp, self._ray_normal_w], + device=self._device, + ) + + n_meshes = self._mesh_ids_wp.shape[1] + + # 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.ray_hits_w.warp, + self._ray_distance_cam_wp, + self._ray_normal_w, + self._ray_face_id_wp, + self._ray_mesh_id_wp, + self._mesh_positions_w, + self._mesh_orientations_w, + float(ray_caster_kernels.CAMERA_RAYCAST_MAX_DIST), + int(return_normal), + int(False), + int(self.cfg.update_mesh_ids), + ], + device=self._device, + ) + + if "distance_to_image_plane" in self.cfg.data_types: + wp.launch( + ray_caster_kernels.compute_distance_to_image_plane_to_image_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + env_mask, + self._data.quat_w_world.warp, + self._ray_distance_cam_wp, + self._ray_directions_w, + int(self.image_shape[1]), + bool(self._depth_clip_enabled), + float(self.cfg.max_distance), + self._depth_clip_fill_value, + ], + outputs=[ + self._data.output["distance_to_image_plane"].warp, + ], + device=self._device, + ) + + if "distance_to_camera" in self.cfg.data_types: + wp.launch( + ray_caster_kernels.copy_float2d_to_image1_depth_clipped_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + env_mask, + self._ray_distance_cam_wp, + int(self.image_shape[1]), + bool(self._depth_clip_enabled), + float(self.cfg.max_distance), + self._depth_clip_fill_value, + ], + outputs=[ + self._data.output["distance_to_camera"].warp, + ], + device=self._device, + ) + + if return_normal: + wp.launch( + ray_caster_kernels.copy_vec3_2d_to_image3_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, self._ray_normal_w, int(self.image_shape[1]), self._data.output["normals"].warp], + device=self._device, + ) + + if self.cfg.update_mesh_ids: + wp.launch( + ray_caster_kernels.copy_int16_2d_to_image1_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, self._ray_mesh_id_wp, int(self.image_shape[1]), self._data.image_mesh_ids.warp], + device=self._device, + ) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py new file mode 100644 index 00000000000..27af7fa8b0b --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -0,0 +1,360 @@ +# 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 +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, ClassVar + +import numpy as np +import torch +import warp as wp + +from pxr import Gf, Usd, UsdGeom + +import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils +from isaaclab.markers import VisualizationMarkers +from isaaclab.terrains.trimesh.utils import make_plane +from isaaclab.utils.warp import ProxyArray, convert_to_warp_mesh +from isaaclab.utils.warp.kernels import raycast_mesh_masked_kernel + +from ..sensor_base import SensorBase +from . import kernels as ray_caster_kernels +from .ray_caster_data import RayCasterData + +if TYPE_CHECKING: + from .ray_caster_cfg import RayCasterCfg + +logger = logging.getLogger(__name__) + + +class BaseRayCaster(SensorBase): + """A 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. + + .. note:: + Currently, only static meshes are supported. Extending the warp mesh to support dynamic meshes + is a work in progress. + """ + + cfg: RayCasterCfg + """The configuration parameters.""" + + meshes: ClassVar[dict[tuple[str, str], wp.Mesh]] = {} + """A dictionary to store warp meshes for raycasting, shared across all instances. + + The keys are ``(prim_path, device)`` tuples and values are the corresponding warp Mesh objects. Meshes are + created lazily for the sensor's active device, not eagerly for every device. Including the device in the key + prevents a mesh created on one device (e.g. CPU) from being reused by a kernel running on a different device + (e.g. CUDA) when multiple simulation contexts or tests use different devices in the same Python process.""" + _instance_count: ClassVar[int] = 0 + """A counter to track the number of RayCaster instances, used to manage class variable lifecycle.""" + + def __init__(self, cfg: RayCasterCfg): + """Initializes the ray-caster object. + + Args: + cfg: The configuration parameters. + """ + BaseRayCaster._instance_count += 1 + super().__init__(cfg) + # Resolve physics-body paths and spawn the sensor Xform child if needed. + self._requested_prim_path = self.cfg.prim_path + self._resolve_and_spawn("raycaster") + self._data = RayCasterData() + + 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 : {len(BaseRayCaster.meshes)}\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 num_instances(self) -> int: + return self._view_count + + @property + def data(self) -> RayCasterData: + # update sensors if needed + self._update_outdated_buffers() + # return the data + return self._data + + """ + Operations. + """ + + 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) + # resolve to indices for torch indexing + if env_ids is not None: + num_envs_ids = len(env_ids) + elif env_mask is not None: + env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) + num_envs_ids = len(env_ids) + else: + 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) + # 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 + ) + + """ + Implementation. + """ + + def _initialize_impl(self): + super()._initialize_impl() + self._initialize_pose_tracking() + if not hasattr(self, "_view_count"): + view: Any = self._view + self._view_count = view.count + + # Resolve alignment mode to integer constant for kernel dispatch + alignment_map = {"world": 0, "yaw": 1, "base": 2} + if self.cfg.ray_alignment not in alignment_map: + raise RuntimeError(f"Unsupported ray_alignment type: {self.cfg.ray_alignment}.") + self._alignment_mode = alignment_map[self.cfg.ray_alignment] + + # load the meshes by parsing the stage + self._initialize_warp_meshes() + self._initialize_rays_impl() + + def _initialize_pose_tracking(self) -> None: + """Initialize backend-specific sensor pose tracking. + + Backend subclasses must set ``_view_count`` and provide transforms + through either ``_view.get_world_poses(indices=None)`` or an override of + :meth:`_get_view_transforms_wp`. They must also set ``_offset_pos_wp`` + and ``_offset_quat_wp`` to the sensor-frame offset relative to the + tracked backend body/site. + """ + raise NotImplementedError(f"{self.__class__.__name__} must initialize backend pose tracking.") + + def _initialize_warp_meshes(self): + # check number of mesh prims provided + if len(self.cfg.mesh_prim_paths) != 1: + raise NotImplementedError( + f"RayCaster currently only supports one mesh prim. Received: {len(self.cfg.mesh_prim_paths)}" + ) + + # read prims to ray-cast + for mesh_prim_path in self.cfg.mesh_prim_paths: + mesh_key = (mesh_prim_path, self._device) + if mesh_key in BaseRayCaster.meshes: + continue + + mesh_prim = sim_utils.get_first_matching_child_prim( + mesh_prim_path, lambda prim: prim.GetTypeName() == "Plane" + ) + if mesh_prim is None: + mesh_prim = sim_utils.get_first_matching_child_prim( + mesh_prim_path, lambda prim: prim.GetTypeName() == "Mesh" + ) + if mesh_prim is None or not mesh_prim.IsValid(): + raise RuntimeError(f"Invalid mesh prim path: {mesh_prim_path}") + mesh_prim = UsdGeom.Mesh(mesh_prim) + points = np.asarray(mesh_prim.GetPointsAttr().Get()) + xformable = UsdGeom.Xformable(mesh_prim) + world_transform: Gf.Matrix4d = xformable.ComputeLocalToWorldTransform(Usd.TimeCode.Default()) + transform_matrix = np.array(world_transform).T + points = np.matmul(points, transform_matrix[:3, :3].T) + points += transform_matrix[:3, 3] + indices = np.asarray(mesh_prim.GetFaceVertexIndicesAttr().Get()) + wp_mesh = convert_to_warp_mesh(points, indices, device=self._device) + logger.info( + f"Read mesh prim: {mesh_prim.GetPath()} with {len(points)} vertices and {len(indices)} faces." + ) + else: + mesh = make_plane(size=(2e6, 2e6), height=0.0, center_zero=True) + wp_mesh = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=self._device) + logger.info(f"Created infinite plane mesh prim: {mesh_prim.GetPath()}.") + BaseRayCaster.meshes[mesh_key] = wp_mesh + + if all((path, self._device) not in BaseRayCaster.meshes for path in self.cfg.mesh_prim_paths): + raise RuntimeError(f"No meshes found for ray-casting! Please check paths: {self.cfg.mesh_prim_paths}") + + def _initialize_rays_impl(self): + # Compute ray starts and directions from pattern (torch, init-time only) + ray_starts_torch, ray_directions_torch = self.cfg.pattern_cfg.func(self.cfg.pattern_cfg, self._device) + self.num_rays = len(ray_directions_torch) + + # Apply sensor offset rotation/position to local ray pattern + offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device) + offset_quat = torch.tensor(list(self.cfg.offset.rot), device=self._device) + ray_directions_torch = math_utils.quat_apply( + offset_quat.repeat(len(ray_directions_torch), 1), ray_directions_torch + ) + ray_starts_torch += offset_pos + + # Repeat for each environment + ray_starts_torch = ray_starts_torch.repeat(self._view_count, 1, 1).contiguous() + ray_directions_torch = ray_directions_torch.repeat(self._view_count, 1, 1).contiguous() + + # Keep public aliases warp-first; kernels use the underlying Warp arrays. + self.ray_starts = ProxyArray(wp.from_torch(ray_starts_torch, dtype=wp.vec3f)) + self.ray_directions = ProxyArray(wp.from_torch(ray_directions_torch, dtype=wp.vec3f)) + + # Drift buffers are warp-first; reset uses explicit .torch views for sampling. + self.drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device)) + self.ray_cast_drift = ProxyArray(wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device)) + + # World-frame ray buffers + self._ray_starts_w = wp.empty((self._view_count, self.num_rays), dtype=wp.vec3f, device=self._device) + self._ray_directions_w = wp.empty((self._view_count, self.num_rays), dtype=wp.vec3f, device=self._device) + + # Data buffers + self._data.create_buffers(self._view_count, self.num_rays, self._device) + + # Dummy distance/normal buffers required by the merged raycast_mesh_masked_kernel signature. + # Sized (1, 1) even though the kernel is launched at (num_envs, num_rays): the kernel only + # writes to these buffers when return_distance==1 or return_normal==1 respectively, and + # RayCaster always passes 0 for both flags. If those flags are ever enabled here, these + # buffers must be resized to (num_envs, num_rays) to avoid an out-of-bounds write. + self._dummy_ray_distance = wp.empty((1, 1), dtype=wp.float32, device=self._device) + self._dummy_ray_normal = wp.empty((1, 1), dtype=wp.vec3f, device=self._device) + + def _get_view_transforms_wp(self) -> wp.array: + """Get world transforms from the frame view as a warp array of ``wp.transformf``. + + Returns: + Warp array of ``wp.transformf`` with shape ``(num_envs,)``. Layout is + ``(tx, ty, tz, qx, qy, qz, qw)`` per element, matching the quaternion + convention returned by the backend pose tracker. + """ + pos_w, quat_w = self._view.get_world_poses() + pos_torch = pos_w.torch.reshape(-1, 3) + quat_torch = quat_w.torch.reshape(-1, 4) + poses = torch.cat([pos_torch, quat_torch], dim=-1).contiguous() + return wp.from_torch(poses).view(wp.transformf) + + def _update_ray_infos(self, env_mask: wp.array): + """Updates sensor poses and ray world-frame buffers via a single warp kernel.""" + transforms = self._get_view_transforms_wp() + + wp.launch( + ray_caster_kernels.update_ray_caster_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + transforms, + env_mask, + self._offset_pos_wp, + self._offset_quat_wp, + self.drift.warp, + self.ray_cast_drift.warp, + self.ray_starts.warp, + self.ray_directions.warp, + self._alignment_mode, + ], + outputs=[ + self._data._pos_w, + self._data._quat_w, + self._ray_starts_w, + self._ray_directions_w, + ], + device=self._device, + ) + + def _update_buffers_impl(self, env_mask: wp.array): + """Fills the buffers of the sensor data.""" + self._update_ray_infos(env_mask) + + # Fill ray hits with inf before raycasting + wp.launch( + ray_caster_kernels.fill_vec3_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, wp.inf, self._data._ray_hits_w], + device=self._device, + ) + + # Ray-cast against the mesh + wp.launch( + raycast_mesh_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + BaseRayCaster.meshes[(self.cfg.mesh_prim_paths[0], self._device)].id, + env_mask, + self._ray_starts_w, + self._ray_directions_w, + float(self.cfg.max_distance), + int(False), # return_distance: not needed by RayCaster + int(False), # return_normal: not needed by RayCaster + self._data._ray_hits_w, + self._dummy_ray_distance, + self._dummy_ray_normal, + ], + device=self._device, + ) + + # Apply vertical drift to ray hits + wp.launch( + ray_caster_kernels.apply_z_drift_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, self.ray_cast_drift.warp, self._data._ray_hits_w], + device=self._device, + ) + + def _set_debug_vis_impl(self, debug_vis: bool): + if debug_vis: + if not hasattr(self, "ray_visualizer"): + self.ray_visualizer = VisualizationMarkers(self.cfg.visualizer_cfg) + self.ray_visualizer.set_visibility(True) + else: + if hasattr(self, "ray_visualizer"): + self.ray_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + if self._data._ray_hits_w is None: + return + ray_hits_torch = wp.to_torch(self._data._ray_hits_w) + # remove possible inf values + viz_points = ray_hits_torch.reshape(-1, 3) + viz_points = viz_points[~torch.any(torch.isinf(viz_points), dim=1)] + + # if no points to visualize, skip + if viz_points.shape[0] == 0: + return + + self.ray_visualizer.visualize(viz_points) + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + super()._invalidate_initialize_callback(event) + self._view = None + + def __del__(self): + BaseRayCaster._instance_count -= 1 + if BaseRayCaster._instance_count == 0: + BaseRayCaster.meshes.clear() 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 new file mode 100644 index 00000000000..98a82bbac29 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster_camera.py @@ -0,0 +1,587 @@ +# 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 +from collections.abc import Sequence +from typing import TYPE_CHECKING, ClassVar, Literal + +import torch +import warp as wp + +from pxr import UsdGeom + +import isaaclab.utils.math as math_utils +from isaaclab.sensors.camera import CameraData +from isaaclab.utils.warp import ProxyArray +from isaaclab.utils.warp.kernels import raycast_mesh_masked_kernel + +from ..sensor_base import SensorBase +from . import kernels as ray_caster_kernels +from .base_ray_caster import BaseRayCaster + +if TYPE_CHECKING: + from .ray_caster_camera_cfg import RayCasterCameraCfg + +logger = logging.getLogger(__name__) + + +class BaseRayCasterCamera(BaseRayCaster): + """A ray-casting camera sensor. + + The ray-caster camera uses a set of rays to get the distances to meshes in the scene. The rays are + defined in the sensor's local coordinate frame. The sensor has the same interface as the + :class:`isaaclab.sensors.Camera` that implements the camera class through USD camera prims. + However, this class provides a faster image generation. The sensor converts meshes from the list of + primitive paths provided in the configuration to Warp meshes. The camera then ray-casts against these + Warp meshes only. + + Currently, only the following annotators are supported: + + - ``"distance_to_camera"``: An image containing the distance to camera optical center. + - ``"distance_to_image_plane"``: An image containing distances of 3D points from camera plane along camera's z-axis. + - ``"normals"``: An image containing the local surface normal vectors at each pixel. + + .. note:: + Currently, only static meshes are supported. Extending the warp mesh to support dynamic meshes + is a work in progress. + """ + + cfg: RayCasterCameraCfg + """The configuration parameters.""" + UNSUPPORTED_TYPES: ClassVar[set[str]] = { + "rgb", + "instance_id_segmentation", + "instance_id_segmentation_fast", + "instance_segmentation", + "instance_segmentation_fast", + "semantic_segmentation", + "skeleton_data", + "motion_vectors", + "bounding_box_2d_tight", + "bounding_box_2d_tight_fast", + "bounding_box_2d_loose", + "bounding_box_2d_loose_fast", + "bounding_box_3d", + "bounding_box_3d_fast", + } + """A set of sensor types that are not supported by the ray-caster camera.""" + + def __init__(self, cfg: RayCasterCameraCfg): + """Initializes the camera object. + + Args: + cfg: The configuration parameters. + + Raises: + ValueError: If the provided data types are not supported by the ray-caster camera. + """ + # perform check on supported data types + self._check_supported_data_types(cfg) + # initialize base class + super().__init__(cfg) + # create empty variables for storing output data + self._data = CameraData() + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + return ( + f"Ray-Caster-Camera @ '{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 : {len(BaseRayCaster.meshes)}\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}\n" + f"\timage shape : {self.image_shape}" + ) + + @property + def data(self) -> CameraData: + # update sensors if needed + self._update_outdated_buffers() + # return the data + return self._data + + @property + def image_shape(self) -> tuple[int, int]: + """A tuple containing (height, width) of the camera sensor.""" + return (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width) + + @property + def frame(self) -> torch.tensor: + """Frame number when the measurement took place.""" + return self._frame + + def set_intrinsic_matrices( + self, matrices: torch.Tensor, focal_length: float = 1.0, env_ids: Sequence[int] | None = None + ): + """Set the intrinsic matrix of the camera. + + Args: + matrices: The intrinsic matrices for the camera. Shape is (N, 3, 3). + focal_length: Focal length to use when computing aperture values (in cm). Defaults to 1.0. + env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. + """ + # resolve env_ids + if env_ids is None: + env_ids = slice(None) + # save new intrinsic matrices and focal length + self._data.intrinsic_matrices.torch[env_ids] = matrices.to(self._device) + self._focal_length = focal_length + # recompute ray directions + ray_starts_torch = self.ray_starts.torch if hasattr(self.ray_starts, "torch") else self.ray_starts + ray_directions_torch = ( + self.ray_directions.torch if hasattr(self.ray_directions, "torch") else self.ray_directions + ) + ray_starts_torch[env_ids], ray_directions_torch[env_ids] = self.cfg.pattern_cfg.func( + self.cfg.pattern_cfg, self._data.intrinsic_matrices.torch[env_ids], self._device + ) + # Refresh warp views of local ray buffers; .contiguous() may produce a copy so we store + # the contiguous tensors explicitly to prevent GC while the warp views are alive. + if hasattr(self, "_ray_starts_local"): + self._ray_starts_contiguous = ray_starts_torch.contiguous() + self._ray_directions_contiguous = ray_directions_torch.contiguous() + self._ray_starts_local = wp.from_torch(self._ray_starts_contiguous, dtype=wp.vec3f) + self._ray_directions_local = wp.from_torch(self._ray_directions_contiguous, dtype=wp.vec3f) + + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): + env_mask = self._resolve_indices_and_mask(env_ids, env_mask) + # reset the timestamps + SensorBase.reset(self, env_mask=env_mask) + # reset the data through the same Warp path used by updates. + # note: this recomputation is useful if one performs events such as randomizations on the camera poses. + self._update_ray_infos(env_mask) + self._update_frame(env_mask, frame_op=2) + + def set_world_poses( + self, + positions: torch.Tensor | None = None, + orientations: torch.Tensor | None = None, + env_ids: Sequence[int] | None = None, + convention: Literal["opengl", "ros", "world"] = "ros", + ): + """Set the pose of the camera w.r.t. the world frame using specified convention. + + Since different fields use different conventions for camera orientations, the method allows users to + set the camera poses in the specified convention. Possible conventions are: + + - :obj:`"opengl"` - forward axis: -Z - up axis +Y - Offset is applied in the OpenGL (Usd.Camera) convention + - :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention + - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention + + See :meth:`isaaclab.utils.math.convert_camera_frame_orientation_convention` for more details + on the conventions. + + Args: + positions: The cartesian coordinates (in meters). Shape is (N, 3). + Defaults to None, in which case the camera position in not changed. + orientations: The quaternion orientation in (x, y, z, w). Shape is (N, 4). + Defaults to None, in which case the camera orientation in not changed. + env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. + convention: The convention in which the poses are fed. Defaults to "ros". + + Raises: + RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. + """ + # resolve env_ids for compact source arrays and env mask for output refresh + if env_ids is None or isinstance(env_ids, slice): + env_ids_wp = self._empty_env_ids_wp + env_mask = self._ALL_ENV_MASK + count = self._num_envs + use_env_ids = False + else: + self._target_env_ids_torch = torch.as_tensor(env_ids, dtype=torch.int32, device=self._device).reshape(-1) + env_ids_wp = wp.from_torch(self._target_env_ids_torch, dtype=wp.int32) + self._reset_mask.zero_() + self._reset_mask_torch[self._target_env_ids_torch.to(dtype=torch.long)] = True + env_mask = self._reset_mask + count = self._target_env_ids_torch.numel() + use_env_ids = True + + target_positions_wp = self._offset_pos_wp + if positions is not None: + positions = torch.as_tensor(positions, dtype=torch.float32, device=self._device).reshape(-1, 3) + if positions.shape[0] == 1 and count != 1: + positions = positions.expand(count, -1) + elif positions.shape[0] != count: + raise ValueError(f"Expected {count} camera positions, got {positions.shape[0]}.") + self._target_positions_torch = positions.contiguous() + target_positions_wp = wp.from_torch(self._target_positions_torch, dtype=wp.vec3f) + + target_quats_wp = self._offset_quat_wp + if orientations is not None: + # convert rotation matrix from input convention to world + quat_w_set = math_utils.convert_camera_frame_orientation_convention( + torch.as_tensor(orientations, dtype=torch.float32, device=self._device), + origin=convention, + target="world", + ) + quat_w_set = quat_w_set.reshape(-1, 4) + if quat_w_set.shape[0] == 1 and count != 1: + quat_w_set = quat_w_set.expand(count, -1) + elif quat_w_set.shape[0] != count: + raise ValueError(f"Expected {count} camera orientations, got {quat_w_set.shape[0]}.") + self._target_quats_torch = quat_w_set.contiguous() + target_quats_wp = wp.from_torch(self._target_quats_torch, dtype=wp.quatf) + + wp.launch( + ray_caster_kernels.update_camera_offsets_kernel, + dim=count, + inputs=[ + self._get_view_transforms_wp(), + env_ids_wp, + target_positions_wp, + target_quats_wp, + use_env_ids, + positions is not None, + orientations is not None, + self._offset_pos_wp, + self._offset_quat_wp, + ], + device=self._device, + ) + + # update the data through the same Warp path used by normal sensor updates + self._update_ray_infos(env_mask) + + def set_world_poses_from_view( + self, eyes: torch.Tensor, targets: torch.Tensor, env_ids: Sequence[int] | None = None + ): + """Set the poses of the camera from the eye position and look-at target position. + + Args: + eyes: The positions of the camera's eye. Shape is (N, 3). + targets: The target locations to look at. Shape is (N, 3). + env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. + + Raises: + RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. + NotImplementedError: If the stage up-axis is not "Y" or "Z". + """ + # get up axis of current stage + up_axis = UsdGeom.GetStageUpAxis(self.stage) + # camera position and rotation in opengl convention + orientations = math_utils.quat_from_matrix( + math_utils.create_rotation_matrix_from_view(eyes, targets, up_axis=up_axis, device=self._device) + ) + self.set_world_poses(eyes, orientations, env_ids, convention="opengl") + + def _create_offset_buffers(self, count: int): + """Create Warp-owned camera offset and drift buffers.""" + quat_w = math_utils.convert_camera_frame_orientation_convention( + torch.tensor([self.cfg.offset.rot], device=self._device), origin=self.cfg.offset.convention, target="world" + ) + offset_pos = [tuple(float(v) for v in self.cfg.offset.pos)] * count + offset_quat = [tuple(float(v) for v in quat_w.squeeze(0).tolist())] * count + self._offset_pos_wp = wp.array(offset_pos, dtype=wp.vec3f, device=self._device) + self._offset_quat_wp = wp.array(offset_quat, dtype=wp.quatf, device=self._device) + self.drift = ProxyArray(wp.zeros(count, dtype=wp.vec3f, device=self._device)) + self.ray_cast_drift = ProxyArray(wp.zeros(count, dtype=wp.vec3f, device=self._device)) + self._empty_env_ids_wp = wp.empty(0, dtype=wp.int32, device=self._device) + + def _initialize_rays_impl(self): + # Frame count is updated by Warp kernels; expose a torch view for the public API. + self._frame_wp = wp.zeros(self._view.count, dtype=wp.int64, device=self._device) + self._frame = wp.to_torch(self._frame_wp) + # create buffers + self._create_buffers() + # compute intrinsic matrices + self._compute_intrinsic_matrices() + # compute ray starts and directions + self.ray_starts, self.ray_directions = self.cfg.pattern_cfg.func( + self.cfg.pattern_cfg, self._data.intrinsic_matrices.torch, self._device + ) + self.num_rays = self.ray_directions.shape[1] + + # Offset/drift buffers are warp-primary so kernels always see current values without re-wrapping. + self._create_offset_buffers(self._view.count) + + # Warp buffers for world-frame rays (used by update kernel) + self._ray_starts_w = wp.empty((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + self._ray_directions_w = wp.empty((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + + # Warp views for ray_starts and ray_directions (from torch tensors returned by pattern_cfg.func) + # These are (num_envs, num_rays, 3) torch tensors; wrap as warp vec3f arrays. + # Store contiguous tensors explicitly so they are not garbage-collected while the + # warp views are alive (mirrors the pattern in RayCaster._initialize_impl). + self._ray_starts_contiguous = self.ray_starts.contiguous() + self._ray_directions_contiguous = self.ray_directions.contiguous() + self._ray_starts_local = wp.from_torch(self._ray_starts_contiguous, dtype=wp.vec3f) + self._ray_directions_local = wp.from_torch(self._ray_directions_contiguous, dtype=wp.vec3f) + + # Intermediate warp buffers for ray results (filled with inf before each raycasting step) + self._ray_distance_wp = wp.empty((self._view.count, self.num_rays), dtype=wp.float32, device=self._device) + if "normals" in self.cfg.data_types: + self._ray_normal_w = wp.empty((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) + else: + self._ray_normal_w = wp.empty((1, 1), dtype=wp.vec3f, device=self._device) + + # Ray hit buffer used by raycasting and debug visualization. + self.ray_hits_w = ProxyArray(wp.empty((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device)) + + def _update_ray_infos(self, env_mask: wp.array): + """Updates camera poses and world-frame ray buffers via a single warp kernel.""" + transforms = self._get_view_transforms_wp() + wp.launch( + ray_caster_kernels.update_ray_caster_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + transforms, + env_mask, + self._offset_pos_wp, + self._offset_quat_wp, + self.drift.warp, + self.ray_cast_drift.warp, + self._ray_starts_local, + self._ray_directions_local, + int(ray_caster_kernels.ALIGNMENT_BASE), + ], + outputs=[ + self._data.pos_w.warp, + self._data.quat_w_world.warp, + self._ray_starts_w, + self._ray_directions_w, + ], + device=self._device, + ) + + def _update_frame(self, env_mask: wp.array, frame_op: int): + """Update frame counters for masked environments.""" + wp.launch( + ray_caster_kernels.update_frame_masked_kernel, + dim=self._num_envs, + inputs=[env_mask, frame_op, self._frame_wp], + device=self._device, + ) + + def _update_buffers_impl(self, env_mask: wp.array): + """Fills the buffers of the sensor data.""" + # increment frame count + self._update_frame(env_mask, frame_op=1) + + self._update_ray_infos(env_mask) + + # Determine whether to compute normals. + need_normal = int("normals" in self.cfg.data_types) + + # Fill ray hit, distance, and optional normal buffers with inf before raycasting. + wp.launch( + ray_caster_kernels.fill_ray_hits_distance_inf_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, bool(need_normal)], + outputs=[self.ray_hits_w.warp, self._ray_distance_wp, self._ray_normal_w], + device=self._device, + ) + + # Ray-cast against the mesh; use a large upper-bound max_dist so depth clipping + # can be applied per-data-type afterwards (matching the original behaviour). + wp.launch( + raycast_mesh_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + BaseRayCaster.meshes[(self.cfg.mesh_prim_paths[0], self._device)].id, + env_mask, + self._ray_starts_w, + self._ray_directions_w, + float(ray_caster_kernels.CAMERA_RAYCAST_MAX_DIST), + int(True), # return_distance: always needed for depth output + need_normal, + self.ray_hits_w.warp, + self._ray_distance_wp, + self._ray_normal_w, + ], + device=self._device, + ) + + if "distance_to_image_plane" in self.cfg.data_types: + wp.launch( + ray_caster_kernels.compute_distance_to_image_plane_to_image_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + env_mask, + self._data.quat_w_world.warp, + self._ray_distance_wp, + self._ray_directions_w, + int(self.image_shape[1]), + bool(self._depth_clip_enabled), + float(self.cfg.max_distance), + self._depth_clip_fill_value, + ], + outputs=[ + self._data.output["distance_to_image_plane"].warp, + ], + device=self._device, + ) + + if "distance_to_camera" in self.cfg.data_types: + wp.launch( + ray_caster_kernels.copy_float2d_to_image1_depth_clipped_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + env_mask, + self._ray_distance_wp, + int(self.image_shape[1]), + bool(self._depth_clip_enabled), + float(self.cfg.max_distance), + self._depth_clip_fill_value, + ], + outputs=[ + self._data.output["distance_to_camera"].warp, + ], + device=self._device, + ) + + if "normals" in self.cfg.data_types: + wp.launch( + ray_caster_kernels.copy_vec3_2d_to_image3_masked_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[env_mask, self._ray_normal_w, int(self.image_shape[1]), self._data.output["normals"].warp], + device=self._device, + ) + + def _debug_vis_callback(self, event): + # Debug visualization can be toggled before ray buffers are initialized. + if not hasattr(self, "ray_hits_w"): + return + # filter out missed rays (inf values) before visualizing + ray_hits_flat = self.ray_hits_w.torch.reshape(-1, 3) + valid_mask = ~torch.isinf(ray_hits_flat).any(dim=-1) + viz_points = ray_hits_flat[valid_mask] + # if no valid hits, skip + if viz_points.shape[0] == 0: + return + self.ray_visualizer.visualize(viz_points) + + def _check_supported_data_types(self, cfg: RayCasterCameraCfg): + """Checks if the data types are supported by the ray-caster camera.""" + # check if there is any intersection in unsupported types + # reason: we cannot obtain this data from simplified warp-based ray caster + common_elements = set(cfg.data_types) & BaseRayCasterCamera.UNSUPPORTED_TYPES + if common_elements: + raise ValueError( + f"RayCasterCamera class does not support the following sensor types: {common_elements}." + "\n\tThis is because these sensor types cannot be obtained in a fast way using ''warp''." + "\n\tHint: If you need to work with these sensor types, we recommend using the USD camera" + " interface from the isaaclab.sensors.camera module." + ) + + def _create_buffers(self): + """Create buffers for storing data.""" + self._depth_clip_enabled = True + if self.cfg.depth_clipping_behavior == "none": + self._depth_clip_enabled = False + self._depth_clip_fill_value = 0.0 + elif self.cfg.depth_clipping_behavior == "max": + self._depth_clip_fill_value = float(self.cfg.max_distance) + elif self.cfg.depth_clipping_behavior == "zero": + self._depth_clip_fill_value = 0.0 + else: + raise ValueError( + f"Unknown depth_clipping_behavior: {self.cfg.depth_clipping_behavior!r}." + " Valid values are 'max', 'zero', and 'none'." + ) + # create the data object + # -- pose of the cameras + self._data.create_buffers(self._view.count, self._device) + # -- intrinsic matrix + self._data.intrinsic_matrices.torch[:, 2, 2] = 1.0 + self._data.image_shape = self.image_shape + # -- output data + # create the buffers to store the annotator data. + self._data._output = {} + self._data.info = {name: None for name in self.cfg.data_types} + for name in self.cfg.data_types: + if name in ["distance_to_image_plane", "distance_to_camera"]: + shape = (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width, 1) + elif name in ["normals"]: + shape = (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width, 3) + else: + raise ValueError(f"Received unknown data type: {name}. Please check the configuration.") + # allocate tensor to store the data + self._data.output[name] = ProxyArray( + wp.zeros((self._view.count, *shape), dtype=wp.float32, device=self._device) + ) + + def _compute_intrinsic_matrices(self): + """Computes the intrinsic matrices for the camera based on the config provided.""" + # get the sensor properties + pattern_cfg = self.cfg.pattern_cfg + + # check if vertical aperture is provided + # if not then it is auto-computed based on the aspect ratio to preserve squared pixels + if pattern_cfg.vertical_aperture is None: + pattern_cfg.vertical_aperture = pattern_cfg.horizontal_aperture * pattern_cfg.height / pattern_cfg.width + + # compute the intrinsic matrix + f_x = pattern_cfg.width * pattern_cfg.focal_length / pattern_cfg.horizontal_aperture + f_y = pattern_cfg.height * pattern_cfg.focal_length / pattern_cfg.vertical_aperture + c_x = pattern_cfg.horizontal_aperture_offset * f_x + pattern_cfg.width / 2 + c_y = pattern_cfg.vertical_aperture_offset * f_y + pattern_cfg.height / 2 + # allocate the intrinsic matrices + self._data.intrinsic_matrices.torch[:, 0, 0] = f_x + self._data.intrinsic_matrices.torch[:, 0, 2] = c_x + self._data.intrinsic_matrices.torch[:, 1, 1] = f_y + self._data.intrinsic_matrices.torch[:, 1, 2] = c_y + + # save focal length + self._focal_length = pattern_cfg.focal_length + + def _compute_view_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]: + """Obtains the pose of the view the camera is attached to in the world frame. + + .. deprecated v2.3.1: + This function will be removed in a future release. Call + ``self._view.get_world_poses(indices)`` directly instead. The returned + ProxyArray pair exposes ``.warp`` and ``.torch`` accessors. + + Returns: + A tuple of the position (in meters) and quaternion (x, y, z, w). + + + """ + logger.warning( + "The function '_compute_view_world_poses' is deprecated." + " Call 'self._view.get_world_poses(indices)' directly instead." + ) + + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None + pos_w, quat_w = self._view.get_world_poses(indices) + return pos_w.torch.clone(), quat_w.torch.clone() + + def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]: + """Computes the pose of the camera in the world frame. + + This function applies the offset pose to the pose of the view the camera is attached to. + + .. deprecated v2.3.1: + This function will be removed in a future release. Instead, use the code block below: + + .. code-block:: python + + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) + pos_w, quat_w = self._view.get_world_poses(indices) + # The returned ProxyArray pair exposes .warp and .torch accessors + pos_w, quat_w = pos_w.torch.clone(), quat_w.torch.clone() + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + ) + + Returns: + A tuple of the position (in meters) and quaternion (x, y, z, w) in "world" convention. + """ + logger.warning( + "The function '_compute_camera_world_poses' is deprecated." + " Call 'self._view.get_world_poses(indices)' and 'math_utils.combine_frame_transforms' directly instead." + ) + + indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None + pos_w, quat_w = self._view.get_world_poses(indices) + offset_pos = wp.to_torch(self._offset_pos_wp) + offset_quat = wp.to_torch(self._offset_quat_wp) + return math_utils.combine_frame_transforms( + pos_w.torch.clone(), quat_w.torch.clone(), offset_pos[env_ids], offset_quat[env_ids] + ) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py index 98c54ea7141..4bcd10aba33 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/kernels.py @@ -171,91 +171,226 @@ def apply_z_drift_kernel( @wp.kernel(enable_backward=False) -def fill_float2d_masked_kernel( +def fill_ray_hits_distance_inf_kernel( # input env_mask: wp.array(dtype=wp.bool), - val: wp.float32, + fill_normals: bool, # output - data: wp.array2d(dtype=wp.float32), + ray_hits: wp.array2d(dtype=wp.vec3f), + ray_distance: wp.array2d(dtype=wp.float32), + ray_normals: wp.array2d(dtype=wp.vec3f), ): - """Fill a 2D float32 array with a given value for masked environments. + """Fill ray hit, distance, and optionally normal buffers with infinity for masked environments. Launch with dim=(num_envs, num_rays). Args: env_mask: Boolean mask for which environments to update. Shape is (num_envs,). - val: Value to fill with. - data: Array to fill. Shape is (num_envs, num_rays). + fill_normals: Whether to fill ``ray_normals``. + ray_hits: Ray hit positions to fill with ``wp.inf``. Shape is (num_envs, num_rays). + ray_distance: Ray distances to fill with ``wp.inf``. Shape is (num_envs, num_rays). + ray_normals: Ray normals to fill with ``wp.inf`` when requested. Shape is (num_envs, num_rays). """ env, ray = wp.tid() if not env_mask[env]: return - data[env, ray] = val + inf_vec = wp.vec3f(wp.inf, wp.inf, wp.inf) + ray_hits[env, ray] = inf_vec + ray_distance[env, ray] = wp.inf + if fill_normals: + ray_normals[env, ray] = inf_vec @wp.kernel(enable_backward=False) -def compute_distance_to_image_plane_masked_kernel( +def update_frame_masked_kernel( # input env_mask: wp.array(dtype=wp.bool), - quat_w: wp.array(dtype=wp.quatf), - ray_distance: wp.array2d(dtype=wp.float32), - ray_directions_w: wp.array2d(dtype=wp.vec3f), + frame_op: int, # output - distance_to_image_plane: wp.array2d(dtype=wp.float32), + frame: wp.array(dtype=wp.int64), ): - """Compute distance-to-image-plane from ray depth and direction for masked environments. + """Update frame counters for masked environments. - The distance to the image plane is the signed projection of the hit displacement - (``ray_distance * ray_direction_w``) onto the camera forward axis (+X in world convention). - This equals the x-component of the hit vector in the camera frame. + ``frame_op`` uses 1 for increment and 2 for reset. + """ + env = wp.tid() + if not env_mask[env]: + return + if frame_op == 1: + frame[env] = frame[env] + wp.int64(1) + elif frame_op == 2: + frame[env] = wp.int64(0) - Launch with dim=(num_envs, num_rays). - Args: - env_mask: Boolean mask for which environments to update. Shape is (num_envs,). - quat_w: Camera orientation in world frame (x, y, z, w). Shape is (num_envs,). - ray_distance: Per-ray hit distances [m]. Shape is (num_envs, num_rays). - Contains inf for missed rays. - ray_directions_w: World-frame unit ray directions. Shape is (num_envs, num_rays). - distance_to_image_plane: Output distance-to-image-plane [m]. Shape is (num_envs, num_rays). +@wp.kernel(enable_backward=False) +def update_camera_offsets_kernel( + # input + transforms: wp.array(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + target_positions: wp.array(dtype=wp.vec3f), + target_quats: wp.array(dtype=wp.quatf), + use_env_ids: bool, + update_position: bool, + update_orientation: bool, + # output + offset_pos: wp.array(dtype=wp.vec3f), + offset_quat: wp.array(dtype=wp.quatf), +): + """Update camera-frame offsets from target world poses. + + Launch with ``dim=count`` where ``count`` is either the number of selected + environments or all environments. ``target_positions`` and ``target_quats`` + are compact arrays indexed by the launch id. """ + src_id = wp.tid() + env_id = src_id + if use_env_ids: + env_id = env_ids[src_id] + + view_transform = transforms[env_id] + view_pos = wp.transform_get_translation(view_transform) + view_quat = wp.transform_get_rotation(view_transform) + + if update_position: + offset_pos[env_id] = wp.quat_rotate_inv(view_quat, target_positions[src_id] - view_pos) + if update_orientation: + offset_quat[env_id] = wp.quat_inverse(view_quat) * target_quats[src_id] + + +@wp.kernel(enable_backward=False) +def copy_float2d_to_image1_depth_clipped_masked_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + src: wp.array2d(dtype=wp.float32), + width: int, + clip_depth: bool, + max_dist: wp.float32, + fill_val: wp.float32, + # output + dst: wp.array4d(dtype=wp.float32), +): + """Copy a flat float buffer to ``(N, H, W, 1)`` camera output with optional depth clipping.""" env, ray = wp.tid() if not env_mask[env]: return + value = src[env, ray] + if clip_depth and (value > max_dist or wp.isnan(value)): + value = fill_val + row = ray // width + col = ray - row * width + dst[env, row, col, 0] = value - depth = ray_distance[env, ray] - dir_w = ray_directions_w[env, ray] - # displacement vector in world frame - disp_w = wp.vec3f(depth * dir_w[0], depth * dir_w[1], depth * dir_w[2]) - # rotate into camera frame (quat_rotate_inv applies q^-1 * v * q) - disp_cam = wp.quat_rotate_inv(quat_w[env], disp_w) - # x-component is the forward (depth) axis of the camera in world convention - distance_to_image_plane[env, ray] = disp_cam[0] + +@wp.kernel(enable_backward=False) +def copy_vec3_2d_to_image3_masked_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + src: wp.array2d(dtype=wp.vec3f), + width: int, + # output + dst: wp.array4d(dtype=wp.float32), +): + """Copy a flat per-ray vec3 buffer to ``(N, H, W, 3)`` camera output.""" + env, ray = wp.tid() + if not env_mask[env]: + return + row = ray // width + col = ray - row * width + value = src[env, ray] + dst[env, row, col, 0] = value[0] + dst[env, row, col, 1] = value[1] + dst[env, row, col, 2] = value[2] @wp.kernel(enable_backward=False) -def apply_depth_clipping_masked_kernel( +def copy_int16_2d_to_image1_masked_kernel( # input env_mask: wp.array(dtype=wp.bool), - max_dist: wp.float32, - fill_val: wp.float32, + src: wp.array2d(dtype=wp.int16), + width: int, # output - depth: wp.array2d(dtype=wp.float32), + dst: wp.array4d(dtype=wp.int16), ): - """Clip depth values in-place, replacing values above max_dist or NaN with fill_val. + """Copy a flat per-ray int16 buffer to ``(N, H, W, 1)`` camera output.""" + env, ray = wp.tid() + if not env_mask[env]: + return + row = ray // width + col = ray - row * width + dst[env, row, col, 0] = src[env, ray] - Launch with dim=(num_envs, num_rays). - Args: - env_mask: Boolean mask for which environments to update. Shape is (num_envs,). - max_dist: Maximum depth threshold [m]. - fill_val: Replacement value [m] written for depths exceeding max_dist or NaN. - Pass ``max_dist`` for "max" clipping or ``0.0`` for "zero" clipping. - depth: Depth buffer to clip in-place. Shape is (num_envs, num_rays). - """ +@wp.kernel(enable_backward=False) +def copy_mesh_poses_to_table_kernel( + # input + positions_src: wp.array(dtype=wp.vec3f), + orientations_src: wp.array(dtype=wp.quatf), + meshes_per_env: int, + mesh_offset: int, + broadcast_single_source: bool, + # output + positions_dst: wp.array2d(dtype=wp.vec3f), + orientations_dst: wp.array2d(dtype=wp.quatf), +): + """Copy flat tracked-mesh poses into the rectangular per-env mesh table.""" + env, local_mesh = wp.tid() + src_index = local_mesh + if not broadcast_single_source: + src_index = env * meshes_per_env + local_mesh + dst_index = mesh_offset + local_mesh + positions_dst[env, dst_index] = positions_src[src_index] + orientations_dst[env, dst_index] = orientations_src[src_index] + + +@wp.kernel(enable_backward=False) +def copy_mesh_transforms_to_table_kernel( + # input + transforms_src: wp.array(dtype=wp.transformf), + meshes_per_env: int, + mesh_offset: int, + broadcast_single_source: bool, + # output + positions_dst: wp.array2d(dtype=wp.vec3f), + orientations_dst: wp.array2d(dtype=wp.quatf), +): + """Copy flat tracked-mesh transforms into the rectangular per-env mesh table.""" + env, local_mesh = wp.tid() + src_index = local_mesh + if not broadcast_single_source: + src_index = env * meshes_per_env + local_mesh + dst_index = mesh_offset + local_mesh + xform = transforms_src[src_index] + positions_dst[env, dst_index] = wp.transform_get_translation(xform) + orientations_dst[env, dst_index] = wp.transform_get_rotation(xform) + + +@wp.kernel(enable_backward=False) +def compute_distance_to_image_plane_to_image_masked_kernel( + # input + env_mask: wp.array(dtype=wp.bool), + quat_w: wp.array(dtype=wp.quatf), + ray_distance: wp.array2d(dtype=wp.float32), + ray_directions_w: wp.array2d(dtype=wp.vec3f), + width: int, + clip_depth: bool, + max_dist: wp.float32, + fill_val: wp.float32, + # output + dst: wp.array4d(dtype=wp.float32), +): + """Compute distance-to-image-plane, optionally clip it, and write camera output.""" env, ray = wp.tid() if not env_mask[env]: return - val = depth[env, ray] - if val > max_dist or wp.isnan(val): - depth[env, ray] = fill_val + + depth = ray_distance[env, ray] + dir_w = ray_directions_w[env, ray] + disp_w = wp.vec3f(depth * dir_w[0], depth * dir_w[1], depth * dir_w[2]) + disp_cam = wp.quat_rotate_inv(quat_w[env], disp_w) + value = disp_cam[0] + if clip_depth and (value > max_dist or wp.isnan(value)): + value = fill_val + + row = ray // width + col = ray - row * width + dst[env, row, col, 0] = value diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py index f9d0493a0f6..8e61c9ffb63 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py @@ -5,441 +5,12 @@ from __future__ import annotations -import logging -import re -from typing import TYPE_CHECKING, ClassVar +from isaaclab.utils.backend_utils import FactoryBase -import numpy as np -import torch -import trimesh -import warp as wp +from .base_multi_mesh_ray_caster import BaseMultiMeshRayCaster -import isaaclab.sim as sim_utils -from isaaclab.sim.views import BaseFrameView, FrameView -from isaaclab.utils.math import matrix_from_quat -from isaaclab.utils.mesh import PRIMITIVE_MESH_TYPES, create_trimesh_from_geom_mesh, create_trimesh_from_geom_shape -from isaaclab.utils.warp import convert_to_warp_mesh -from isaaclab.utils.warp import kernels as warp_kernels -from .kernels import fill_float2d_masked_kernel, fill_vec3_inf_kernel -from .multi_mesh_ray_caster_data import MultiMeshRayCasterData -from .ray_caster import RayCaster +class MultiMeshRayCaster(FactoryBase, BaseMultiMeshRayCaster): + """Backend-dispatching multi-mesh ray-caster sensor.""" -if TYPE_CHECKING: - from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg - -logger = logging.getLogger(__name__) - - -class MultiMeshRayCaster(RayCaster): - """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 `_ 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.""" - - mesh_views: ClassVar[dict[str, BaseFrameView]] = {} - """A dictionary to store mesh views for raycasting, shared across all instances. - - The keys correspond to the prim path for the mesh views, and values are the corresponding view objects. - """ - - 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): - self._raycast_targets_cfg.append(cfg.RaycastTargetCfg(prim_expr=target, track_mesh_transforms=False)) - else: - self._raycast_targets_cfg.append(target) - - for cfg in self._raycast_targets_cfg: - cfg.prim_expr = cfg.prim_expr.format(ENV_REGEX_NS="/World/envs/env_.*") - - 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): - """Parse mesh prim expressions, build (or reuse) Warp meshes, and cache per-env mesh IDs. - - High-level steps (per target expression): - - 1. Resolve matching prims by regex/path expression. - 2. Collect supported mesh child prims; merge into a single mesh if configured. - 3. Deduplicate identical vertex buffers (exact match) to avoid uploading duplicates to Warp. - 4. Partition mesh IDs per environment or mark as globally shared. - 5. Optionally create physics views (articulation / rigid body / fallback XForm) and cache local offsets. - - Exceptions: - Raises a RuntimeError if: - - - No prims match the provided expression. - - No supported mesh prims are found under a matched prim. - - Multiple mesh prims are found but merging is disabled. - - """ - multi_mesh_ids: dict[str, list[list[int]]] = {} - for target_cfg in self._raycast_targets_cfg: - target_prim_path = target_cfg.prim_expr - # check if mesh already casted into warp mesh and skip if so. - if target_prim_path in multi_mesh_ids: - logger.warning( - f"Mesh at target prim path '{target_prim_path}' already exists in the mesh cache. Duplicate entries" - " in `mesh_prim_paths`? This mesh will be skipped." - ) - continue - - target_prims = sim_utils.find_matching_prims(target_prim_path) - if len(target_prims) == 0: - raise RuntimeError(f"Failed to find a prim at path expression: {target_prim_path}") - - is_global_prim = len(target_prims) == 1 - - loaded_vertices: list[np.ndarray | None] = [] - wp_mesh_ids = [] - - for target_prim in target_prims: - if target_cfg.is_shared and len(wp_mesh_ids) > 0: - # Verify if this mesh has already been registered in an earlier environment. - # Note, this check may fail, if the prim path is not following the env_.* pattern - # Which (worst case) leads to parsing the mesh and skipping registering it at a later stage - curr_prim_base_path = re.sub(r"env_\d+", "env_0", str(target_prim.GetPath())) - base_key = (curr_prim_base_path, self._device) - if base_key in MultiMeshRayCaster.meshes: - MultiMeshRayCaster.meshes[(str(target_prim.GetPath()), self._device)] = ( - MultiMeshRayCaster.meshes[base_key] - ) - prim_key = (str(target_prim.GetPath()), self._device) - if prim_key in MultiMeshRayCaster.meshes: - wp_mesh_ids.append(MultiMeshRayCaster.meshes[prim_key].id) - loaded_vertices.append(None) - continue - - 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: - warn_msg = ( - f"No mesh prims found at path: {target_prim.GetPath()} with supported types:" - f" {PRIMITIVE_MESH_TYPES + ['Mesh']}" - " Skipping this target." - ) - for prim in sim_utils.get_all_matching_child_prims(target_prim.GetPath(), lambda prim: True): - warn_msg += f"\n - Available prim '{prim.GetPath()}' of type '{prim.GetTypeName()}'" - logger.warning(warn_msg) - continue - - 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) - scale = sim_utils.resolve_prim_scale(mesh_prim) - mesh.apply_scale(scale) - - relative_pos, relative_quat = sim_utils.resolve_prim_pose(mesh_prim, target_prim) - relative_pos = torch.tensor(relative_pos, dtype=torch.float32) - relative_quat = torch.tensor(relative_quat, dtype=torch.float32) - - rotation = matrix_from_quat(relative_quat) - transform = np.eye(4) - transform[:3, :3] = rotation.numpy() - transform[:3, 3] = relative_pos.numpy() - 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." - ) - - registered_idx = _registered_points_idx(trimesh_mesh.vertices, loaded_vertices) - if registered_idx != -1 and self.cfg.reference_meshes: - logger.info("Found a duplicate mesh, only reference the mesh.") - loaded_vertices.append(None) - wp_mesh_ids.append(wp_mesh_ids[registered_idx]) - else: - loaded_vertices.append(trimesh_mesh.vertices) - wp_mesh = convert_to_warp_mesh(trimesh_mesh.vertices, trimesh_mesh.faces, device=self._device) - MultiMeshRayCaster.meshes[(str(target_prim.GetPath()), self._device)] = wp_mesh - wp_mesh_ids.append(wp_mesh.id) - - if registered_idx != -1: - logger.info(f"Found duplicate mesh for mesh prims under path '{target_prim.GetPath()}'.") - else: - 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." - ) - - if is_global_prim: - multi_mesh_ids[target_prim_path] = [wp_mesh_ids] * self._num_envs - self._num_meshes_per_env[target_prim_path] = len(wp_mesh_ids) - else: - multi_mesh_ids[target_prim_path] = [] - mesh_idx = 0 - n_meshes_per_env = len(wp_mesh_ids) // self._num_envs - self._num_meshes_per_env[target_prim_path] = n_meshes_per_env - for _ in range(self._num_envs): - multi_mesh_ids[target_prim_path].append(wp_mesh_ids[mesh_idx : mesh_idx + n_meshes_per_env]) - mesh_idx += n_meshes_per_env - - if target_cfg.track_mesh_transforms: - MultiMeshRayCaster.mesh_views[target_prim_path] = FrameView( - target_prim_path, device=self._device, stage=self.stage - ) - - if all([target_cfg.prim_expr not in multi_mesh_ids for target_cfg in self._raycast_targets_cfg]): - raise RuntimeError( - f"No meshes found for ray-casting! Please check the mesh prim paths: {self.cfg.mesh_prim_paths}" - ) - - total_n_meshes_per_env = sum(self._num_meshes_per_env.values()) - self._mesh_positions_w = wp.zeros((self._num_envs, total_n_meshes_per_env), dtype=wp.vec3, device=self.device) - self._mesh_orientations_w = wp.zeros( - (self._num_envs, total_n_meshes_per_env), dtype=wp.quat, device=self.device - ) - # Zero-copy torch views for writing from physics view results (torch tensors) - self._mesh_positions_w_torch = wp.to_torch(self._mesh_positions_w) - self._mesh_orientations_w_torch = wp.to_torch(self._mesh_orientations_w) - - mesh_idx = 0 - for target_cfg in self._raycast_targets_cfg: - n_meshes = self._num_meshes_per_env[target_cfg.prim_expr] - - pos_w, ori_w = [], [] - for prim in sim_utils.find_matching_prims(target_cfg.prim_expr): - translation, quat = sim_utils.resolve_prim_pose(prim) - pos_w.append(translation) - ori_w.append(quat) - pos_w = torch.tensor(pos_w, device=self.device, dtype=torch.float32).view(-1, n_meshes, 3) - ori_w = torch.tensor(ori_w, device=self.device, dtype=torch.float32).view(-1, n_meshes, 4) - - self._mesh_positions_w_torch[:, mesh_idx : mesh_idx + n_meshes] = pos_w - self._mesh_orientations_w_torch[:, mesh_idx : mesh_idx + n_meshes] = ori_w - mesh_idx += n_meshes - - multi_mesh_ids_flattened = [] - for env_idx in range(self._num_envs): - meshes_in_env = [] - for target_cfg in self._raycast_targets_cfg: - meshes_in_env.extend(multi_mesh_ids[target_cfg.prim_expr][env_idx]) - multi_mesh_ids_flattened.append(meshes_in_env) - - self._mesh_views = [ - self.mesh_views[target_cfg.prim_expr] if target_cfg.track_mesh_transforms else None - for target_cfg in self._raycast_targets_cfg - ] - - self._mesh_ids_wp = wp.array2d(multi_mesh_ids_flattened, dtype=wp.uint64, device=self.device) - - def _initialize_rays_impl(self): - super()._initialize_rays_impl() - # Persistent buffer for tracking closest-hit distance across meshes (for atomic_min) - self._ray_distance_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.float32, device=self._device) - if self.cfg.update_mesh_ids: - self._ray_mesh_id_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.int16, device=self._device) - # Zero-copy torch view with the trailing dim expected by consumers of ray_mesh_ids - self._data.ray_mesh_ids = wp.to_torch(self._ray_mesh_id_w).unsqueeze(-1) - else: - # Dummy 1×1 buffer so the kernel launch always has a valid array to bind - self._ray_mesh_id_w = 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_w = wp.empty((1, 1), dtype=wp.vec3, device=self._device) - self._dummy_face_id_w = 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 - ``_mesh_positions_w_torch`` and ``_mesh_orientations_w_torch``. 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 - - # update position of the target meshes - pos_w, ori_w = view.get_world_poses(None) - pos_w, ori_w = pos_w.torch, ori_w.torch - pos_w = pos_w.squeeze(0) if len(pos_w.shape) == 3 else pos_w - ori_w = ori_w.squeeze(0) if len(ori_w.shape) == 3 else ori_w - - count = view.count - if count != 1: - count = count // self._num_envs - pos_w = pos_w.view(self._num_envs, count, 3) - ori_w = ori_w.view(self._num_envs, count, 4) - - self._mesh_positions_w_torch[:, mesh_idx : mesh_idx + count] = pos_w - self._mesh_orientations_w_torch[:, mesh_idx : mesh_idx + count] = ori_w - mesh_idx += count - - 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() - - n_meshes = self._mesh_ids_wp.shape[1] - - # Fill output and distance buffers with inf for masked environments - wp.launch( - fill_vec3_inf_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._data._ray_hits_w], - device=self._device, - ) - wp.launch( - fill_float2d_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._ray_distance_w], - device=self._device, - ) - - # 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_w, - self._dummy_normal_w, - self._dummy_face_id_w, - self._ray_mesh_id_w, - self._mesh_positions_w, - self._mesh_orientations_w, - float(self.cfg.max_distance), - int(False), - int(False), - int(self.cfg.update_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 - MultiMeshRayCaster.mesh_views.clear() - - def __del__(self): - super().__del__() - if RayCaster._instance_count == 0: - MultiMeshRayCaster.mesh_views.clear() - - -""" -Helper functions -""" - - -def _registered_points_idx(points: np.ndarray, registered_points: list[np.ndarray | None]) -> int: - """Check if the points are already registered in the list of registered points. - - Args: - points: The points to check. - registered_points: The list of registered points. - - Returns: - The index of the registered points if found, otherwise -1. - """ - for idx, reg_points in enumerate(registered_points): - if reg_points is None: - continue - if reg_points.shape == points.shape and (reg_points == points).all(): - return idx - return -1 + _backend_class_names = {"physx": "MultiMeshRayCaster", "newton": "MultiMeshRayCaster"} diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py index a868d17c784..e201e61d1f2 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py @@ -5,286 +5,12 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from isaaclab.utils.backend_utils import FactoryBase -import torch -import warp as wp +from .base_multi_mesh_ray_caster_camera import BaseMultiMeshRayCasterCamera -import isaaclab.utils.math as math_utils -from isaaclab.utils.warp import kernels as warp_kernels -from .kernels import ( - CAMERA_RAYCAST_MAX_DIST, - compute_distance_to_image_plane_masked_kernel, - fill_float2d_masked_kernel, - fill_vec3_inf_kernel, -) -from .multi_mesh_ray_caster import MultiMeshRayCaster -from .multi_mesh_ray_caster_camera_data import MultiMeshRayCasterCameraData -from .ray_caster_camera import RayCasterCamera +class MultiMeshRayCasterCamera(FactoryBase, BaseMultiMeshRayCasterCamera): + """Backend-dispatching multi-mesh ray-caster camera sensor.""" -if TYPE_CHECKING: - from .multi_mesh_ray_caster_camera_cfg import MultiMeshRayCasterCameraCfg - - -class MultiMeshRayCasterCamera(RayCasterCamera, MultiMeshRayCaster): - """A multi-mesh ray-casting camera sensor. - - The ray-caster camera uses a set of rays to get the distances to meshes in the scene. The rays are - defined in the sensor's local coordinate frame. The sensor has the same interface as the - :class:`isaaclab.sensors.Camera` that implements the camera class through USD camera prims. - However, this class provides a faster image generation. The sensor converts meshes from the list of - primitive paths provided in the configuration to Warp meshes. The camera then ray-casts against these - Warp meshes only. - - Currently, only the following annotators are supported: - - - ``"distance_to_camera"``: An image containing the distance to camera optical center. - - ``"distance_to_image_plane"``: An image containing distances of 3D points from camera plane along camera's z-axis. - - ``"normals"``: An image containing the local surface normal vectors at each pixel. - """ - - cfg: MultiMeshRayCasterCameraCfg - """The configuration parameters.""" - - def __init__(self, cfg: MultiMeshRayCasterCameraCfg): - """Initializes the camera object. - - Args: - cfg: The configuration parameters. - - Raises: - ValueError: If the provided data types are not supported by the ray-caster camera. - """ - self._check_supported_data_types(cfg) - # initialize base class - MultiMeshRayCaster.__init__(self, cfg) - # create empty variables for storing output data - self._data = MultiMeshRayCasterCameraData() - - def __str__(self) -> str: - """Returns: A string containing information about the instance.""" - return ( - f"Multi-Mesh Ray-Caster-Camera @ '{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 : {len(MultiMeshRayCaster.meshes)}\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}\n" - f"\timage shape : {self.image_shape}" - ) - - """ - Implementation. - """ - - def _initialize_warp_meshes(self): - MultiMeshRayCaster._initialize_warp_meshes(self) - - def _create_buffers(self): - super()._create_buffers() - self._data.image_mesh_ids = torch.zeros( - self._num_envs, *self.image_shape, 1, device=self.device, dtype=torch.int16 - ) - - def _initialize_rays_impl(self): - # NOTE: This method intentionally does NOT call super()._initialize_rays_impl() through the MRO - # chain. The intermediate classes (RayCasterCamera, MultiMeshRayCaster) use different internal - # buffer names and orderings that are incompatible with the camera's full init path: - # - RayCasterCamera creates single-mesh ray buffers (_ray_distance, _ray_normal_w, etc.) - # - MultiMeshRayCaster creates _ray_distance_w / _ray_mesh_id_w for multi-mesh use - # The camera replaces all of these with its own camera-named equivalents below. - # If either parent class gains new shared buffers, they must be added here explicitly. - - # Camera-specific bookkeeping buffers - self._ALL_INDICES = torch.arange(self._view.count, device=self._device, dtype=torch.long) - self._frame = torch.zeros(self._view.count, device=self._device, dtype=torch.long) - - # Build camera output buffers (intrinsics, image data, etc.) - self._create_buffers() - self._compute_intrinsic_matrices() - - # Compute local ray starts/directions from the camera pattern (torch, init-time only) - ray_starts_local, ray_directions_local = self.cfg.pattern_cfg.func( - self.cfg.pattern_cfg, self._data.intrinsic_matrices, self._device - ) - self.num_rays = ray_directions_local.shape[1] - - # Store local (sensor-frame) ray arrays as torch tensors for per-env camera-convention rotation - self.ray_starts = ray_starts_local - self.ray_directions = ray_directions_local - - # Camera-frame offset: convert from cfg convention to world convention - quat_offset = math_utils.convert_camera_frame_orientation_convention( - torch.tensor([self.cfg.offset.rot], device=self._device), - origin=self.cfg.offset.convention, - target="world", - ) - self._offset_quat = quat_offset.repeat(self._view.count, 1) - self._offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device).repeat(self._view.count, 1) - - # Warp-backed camera orientation buffer for warp kernel calls; - # updated from self._data.quat_w_world in _update_ray_infos. - self._quat_w_wp = wp.zeros(self._view.count, dtype=wp.quatf, device=self._device) - self._quat_w_wp_torch = wp.to_torch(self._quat_w_wp) - - # Warp buffer for distance_to_image_plane output (if requested) - if "distance_to_image_plane" in self.cfg.data_types: - self._distance_to_image_plane_wp = wp.zeros( - (self._view.count, self.num_rays), dtype=wp.float32, device=self._device - ) - - # World-frame ray buffers: allocate as warp arrays first, then create zero-copy torch views. - # Keeping warp arrays as primary storage avoids lifetime issues when passing to kernels. - self._ray_starts_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - self._ray_directions_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - # Zero-copy torch views used for indexing and post-processing - self._ray_starts_w_torch = wp.to_torch(self._ray_starts_w) - self._ray_directions_w_torch = wp.to_torch(self._ray_directions_w) - - # Ray hit positions as a warp array; expose a torch view for debug visualisation - self._ray_hits_w_cam = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - self.ray_hits_w = wp.to_torch(self._ray_hits_w_cam) - - # Per-ray closest-hit distance for atomic_min across meshes - self._ray_distance_cam_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.float32, device=self._device) - - # Optional normal buffer (always allocated; filled only when "normals" is requested) - self._ray_normal_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - - # Mesh-id buffers from MultiMeshRayCaster._initialize_rays_impl - if self.cfg.update_mesh_ids: - self._ray_mesh_id_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.int16, device=self._device) - self._data.ray_mesh_ids = wp.to_torch(self._ray_mesh_id_w).unsqueeze(-1) - else: - self._ray_mesh_id_w = wp.empty((1, 1), dtype=wp.int16, device=self._device) - - # Dummy face-id buffer (not used by camera but required by kernel signature) - self._ray_face_id_w = wp.empty((1, 1), dtype=wp.int32, device=self._device) - - def _update_ray_infos(self, env_mask: wp.array): - """Updates camera poses and world-frame ray buffers for masked environments. - - Args: - env_mask: Boolean mask selecting which environments to update. Shape is (num_envs,). - """ - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - if len(env_ids) == 0: - return - - # Compute camera world poses by composing view pose with sensor offset - indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) - pos_w, quat_w = self._view.get_world_poses(indices) - pos_w, quat_w = pos_w.torch, quat_w.torch - pos_w, quat_w = math_utils.combine_frame_transforms( - pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] - ) - # Store camera pose in CameraData (torch tensors) and warp-backed orientation buffer - self._data.pos_w.torch[env_ids] = pos_w - self._data.quat_w_world.torch[env_ids] = quat_w - self._quat_w_wp_torch[env_ids] = quat_w - - # Rotate local ray starts and directions into world frame using full camera orientation - quat_w_repeated = quat_w.repeat(1, self.num_rays).reshape(-1, 4) - ray_starts_local = self.ray_starts[env_ids].reshape(-1, 3) - ray_dirs_local = self.ray_directions[env_ids].reshape(-1, 3) - - ray_starts_world = math_utils.quat_apply(quat_w_repeated, ray_starts_local).reshape( - len(env_ids), self.num_rays, 3 - ) - ray_starts_world += pos_w.unsqueeze(1) - ray_dirs_world = math_utils.quat_apply(quat_w_repeated, ray_dirs_local).reshape(len(env_ids), self.num_rays, 3) - - # Write back into the warp-backed buffers via zero-copy torch views - self._ray_starts_w_torch[env_ids] = ray_starts_world - self._ray_directions_w_torch[env_ids] = ray_dirs_world - - def _update_buffers_impl(self, env_mask: wp.array): - """Fills the buffers of the sensor data.""" - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - if len(env_ids) == 0: - return - - self._update_ray_infos(env_mask) - - # Increment frame count for updated environments - self._frame[env_ids] += 1 - - self._update_mesh_transforms() - - n_meshes = self._mesh_ids_wp.shape[1] - return_normal = "normals" in self.cfg.data_types - - # Fill ray hit and distance buffers with inf for masked environments - wp.launch( - fill_vec3_inf_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._ray_hits_w_cam], - device=self._device, - ) - wp.launch( - fill_float2d_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._ray_distance_cam_w], - device=self._device, - ) - if return_normal: - wp.launch( - fill_vec3_inf_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._ray_normal_w], - device=self._device, - ) - - # 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._ray_hits_w_cam, - self._ray_distance_cam_w, - self._ray_normal_w, - self._ray_face_id_w, - self._ray_mesh_id_w, - self._mesh_positions_w, - self._mesh_orientations_w, - float(CAMERA_RAYCAST_MAX_DIST), - int(return_normal), - int(False), - int(self.cfg.update_mesh_ids), - ], - device=self._device, - ) - - if "distance_to_image_plane" in self.cfg.data_types: - wp.launch( - compute_distance_to_image_plane_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, self._quat_w_wp, self._ray_distance_cam_w, self._ray_directions_w], - outputs=[self._distance_to_image_plane_wp], - device=self._device, - ) - # Apply depth clipping on the intermediate buffer (leaves _ray_distance_cam_w unmodified) - self._apply_depth_clipping(env_mask, self._distance_to_image_plane_wp) - d2ip_torch = wp.to_torch(self._distance_to_image_plane_wp) - self._data.output["distance_to_image_plane"][env_ids] = d2ip_torch[env_ids].view(-1, *self.image_shape, 1) - - if "distance_to_camera" in self.cfg.data_types: - # d2ip (if requested) was computed before this block so _ray_distance_cam_w is still unclipped. - self._apply_depth_clipping(env_mask, self._ray_distance_cam_w) - ray_dist_torch = wp.to_torch(self._ray_distance_cam_w) - self._data.output["distance_to_camera"][env_ids] = ray_dist_torch[env_ids].view(-1, *self.image_shape, 1) - - if return_normal: - ray_normal_torch = wp.to_torch(self._ray_normal_w) - self._data.output["normals"][env_ids] = ray_normal_torch[env_ids].view(-1, *self.image_shape, 3) - - if self.cfg.update_mesh_ids: - self._data.image_mesh_ids[env_ids] = wp.to_torch(self._ray_mesh_id_w)[env_ids].view( - -1, *self.image_shape, 1 - ) + _backend_class_names = {"physx": "MultiMeshRayCasterCamera", "newton": "MultiMeshRayCasterCamera"} diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py index 21338f0a061..687185c3e0b 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py @@ -5,9 +5,8 @@ """Data container for the multi-mesh ray-cast camera sensor.""" -import torch - from isaaclab.sensors.camera import CameraData +from isaaclab.utils.warp import ProxyArray class MultiMeshRayCasterCameraData(CameraData): @@ -19,7 +18,7 @@ class MultiMeshRayCasterCameraData(CameraData): warp-native :class:`RayCasterData`. """ - image_mesh_ids: torch.Tensor = None + image_mesh_ids: ProxyArray = None """The mesh ids of the image pixels. Shape is (N, H, W, 1), where N is the number of sensors, H and W are the height and width of the image, diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py index b9ae187591b..331dcb4af79 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py @@ -6,7 +6,7 @@ """Data container for the multi-mesh ray-cast sensor.""" -import torch +from isaaclab.utils.warp import ProxyArray from .ray_caster_data import RayCasterData @@ -14,7 +14,7 @@ class MultiMeshRayCasterData(RayCasterData): """Data container for the multi-mesh ray-cast sensor.""" - ray_mesh_ids: torch.Tensor = None + ray_mesh_ids: ProxyArray = None """The mesh ids of the ray hits. Shape is (N, B, 1), where N is the number of sensors, B is the number of rays diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py index ab8ca8b7d4c..4b21117c0c2 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py @@ -5,370 +5,12 @@ from __future__ import annotations -import logging -from collections.abc import Sequence -from typing import TYPE_CHECKING, ClassVar +from isaaclab.utils.backend_utils import FactoryBase -import numpy as np -import torch -import warp as wp +from .base_ray_caster import BaseRayCaster -from pxr import Gf, Usd, UsdGeom -import isaaclab.sim as sim_utils -import isaaclab.utils.math as math_utils -from isaaclab.markers import VisualizationMarkers -from isaaclab.sim.views import FrameView -from isaaclab.terrains.trimesh.utils import make_plane -from isaaclab.utils.warp import convert_to_warp_mesh -from isaaclab.utils.warp.kernels import raycast_mesh_masked_kernel +class RayCaster(FactoryBase, BaseRayCaster): + """Backend-dispatching ray-caster sensor.""" -from ..sensor_base import SensorBase -from .kernels import ( - apply_z_drift_kernel, - fill_vec3_inf_kernel, - update_ray_caster_kernel, -) -from .ray_caster_data import RayCasterData - -if TYPE_CHECKING: - from .ray_caster_cfg import RayCasterCfg - -logger = logging.getLogger(__name__) - - -class RayCaster(SensorBase): - """A 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. - - .. note:: - Currently, only static meshes are supported. Extending the warp mesh to support dynamic meshes - is a work in progress. - """ - - cfg: RayCasterCfg - """The configuration parameters.""" - - meshes: ClassVar[dict[tuple[str, str], wp.Mesh]] = {} - """A dictionary to store warp meshes for raycasting, shared across all instances. - - The keys are ``(prim_path, device)`` tuples and values are the corresponding warp Mesh objects. - Including the device in the key prevents a mesh created on one device (e.g. CPU) from being - reused by a kernel running on a different device (e.g. CUDA).""" - _instance_count: ClassVar[int] = 0 - """A counter to track the number of RayCaster instances, used to manage class variable lifecycle.""" - - def __init__(self, cfg: RayCasterCfg): - """Initializes the ray-caster object. - - Args: - cfg: The configuration parameters. - """ - RayCaster._instance_count += 1 - super().__init__(cfg) - # Resolve physics-body paths and spawn the sensor Xform child if needed. - self._resolve_and_spawn("raycaster") - self._data = RayCasterData() - - 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 : {len(RayCaster.meshes)}\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 num_instances(self) -> int: - return self._view.count - - @property - def data(self) -> RayCasterData: - # update sensors if needed - self._update_outdated_buffers() - # return the data - return self._data - - """ - Operations. - """ - - 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) - # resolve to indices for torch indexing - if env_ids is not None: - num_envs_ids = len(env_ids) - elif env_mask is not None: - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - num_envs_ids = len(env_ids) - else: - 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[env_ids] = r.uniform_(*self.cfg.drift_range) - # resample the ray cast drift - range_list = [self.cfg.ray_cast_drift_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z"]] - ranges = torch.tensor(range_list, device=self.device) - self.ray_cast_drift[env_ids] = math_utils.sample_uniform( - ranges[:, 0], ranges[:, 1], (num_envs_ids, 3), device=self.device - ) - - """ - Implementation. - """ - - def _initialize_impl(self): - super()._initialize_impl() - # Build a FrameView over the sensor prim paths. The FrameView tracks the spawned - # (non-physics) Xform directly, so no physics-body redirect or offset resolution - # is needed at runtime — the world pose returned already includes any offset - # baked into the prim's local transform. - self._view = FrameView(self.cfg.prim_path, device=self._device, stage=self.stage) - - # Per-env identity offsets (kept for kernel ABI compatibility): the sensor frame is - # already the FrameView's tracked prim, so no additional view-to-sensor offset applies. - self._offset_pos_wp = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) - identity_quat = torch.zeros(self._view.count, 4, device=self._device) - identity_quat[:, 3] = 1.0 - self._offset_quat_contiguous = identity_quat.contiguous() - self._offset_quat_wp = wp.from_torch(self._offset_quat_contiguous, dtype=wp.quatf) - - # Resolve alignment mode to integer constant for kernel dispatch - alignment_map = {"world": 0, "yaw": 1, "base": 2} - if self.cfg.ray_alignment not in alignment_map: - raise RuntimeError(f"Unsupported ray_alignment type: {self.cfg.ray_alignment}.") - self._alignment_mode = alignment_map[self.cfg.ray_alignment] - - # load the meshes by parsing the stage - self._initialize_warp_meshes() - self._initialize_rays_impl() - - def _initialize_warp_meshes(self): - # check number of mesh prims provided - if len(self.cfg.mesh_prim_paths) != 1: - raise NotImplementedError( - f"RayCaster currently only supports one mesh prim. Received: {len(self.cfg.mesh_prim_paths)}" - ) - - # read prims to ray-cast - for mesh_prim_path in self.cfg.mesh_prim_paths: - mesh_key = (mesh_prim_path, self._device) - if mesh_key in RayCaster.meshes: - continue - - mesh_prim = sim_utils.get_first_matching_child_prim( - mesh_prim_path, lambda prim: prim.GetTypeName() == "Plane" - ) - if mesh_prim is None: - mesh_prim = sim_utils.get_first_matching_child_prim( - mesh_prim_path, lambda prim: prim.GetTypeName() == "Mesh" - ) - if mesh_prim is None or not mesh_prim.IsValid(): - raise RuntimeError(f"Invalid mesh prim path: {mesh_prim_path}") - mesh_prim = UsdGeom.Mesh(mesh_prim) - points = np.asarray(mesh_prim.GetPointsAttr().Get()) - xformable = UsdGeom.Xformable(mesh_prim) - world_transform: Gf.Matrix4d = xformable.ComputeLocalToWorldTransform(Usd.TimeCode.Default()) - transform_matrix = np.array(world_transform).T - points = np.matmul(points, transform_matrix[:3, :3].T) - points += transform_matrix[:3, 3] - indices = np.asarray(mesh_prim.GetFaceVertexIndicesAttr().Get()) - wp_mesh = convert_to_warp_mesh(points, indices, device=self._device) - logger.info( - f"Read mesh prim: {mesh_prim.GetPath()} with {len(points)} vertices and {len(indices)} faces." - ) - else: - mesh = make_plane(size=(2e6, 2e6), height=0.0, center_zero=True) - wp_mesh = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=self._device) - logger.info(f"Created infinite plane mesh prim: {mesh_prim.GetPath()}.") - RayCaster.meshes[mesh_key] = wp_mesh - - if all((mesh_prim_path, self._device) not in RayCaster.meshes for mesh_prim_path in self.cfg.mesh_prim_paths): - raise RuntimeError( - f"No meshes found for ray-casting! Please check the mesh prim paths: {self.cfg.mesh_prim_paths}" - ) - - def _initialize_rays_impl(self): - # Compute ray starts and directions from pattern (torch, init-time only) - ray_starts_torch, ray_directions_torch = self.cfg.pattern_cfg.func(self.cfg.pattern_cfg, self._device) - self.num_rays = len(ray_directions_torch) - - # Apply sensor offset rotation/position to local ray pattern - offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device) - offset_quat = torch.tensor(list(self.cfg.offset.rot), device=self._device) - ray_directions_torch = math_utils.quat_apply( - offset_quat.repeat(len(ray_directions_torch), 1), ray_directions_torch - ) - ray_starts_torch += offset_pos - - # Repeat for each environment - ray_starts_torch = ray_starts_torch.repeat(self._view.count, 1, 1) - ray_directions_torch = ray_directions_torch.repeat(self._view.count, 1, 1) - - # Create warp arrays from the init-time torch data - # The warp arrays own the memory; torch views provide backward-compat indexing - self._ray_starts_local = wp.from_torch(ray_starts_torch.contiguous(), dtype=wp.vec3f) - self._ray_directions_local = wp.from_torch(ray_directions_torch.contiguous(), dtype=wp.vec3f) - - # Torch views (same attribute names as before for subclass compatibility) - self.ray_starts = wp.to_torch(self._ray_starts_local) - self.ray_directions = wp.to_torch(self._ray_directions_local) - - # Drift buffers (warp-owned, torch views for reset indexing) - self._drift = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) - self._ray_cast_drift = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) - self.drift = wp.to_torch(self._drift) - self.ray_cast_drift = wp.to_torch(self._ray_cast_drift) - - # World-frame ray buffers - self._ray_starts_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - self._ray_directions_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - - # Torch views for subclass compatibility - self._ray_starts_w_torch = wp.to_torch(self._ray_starts_w) - self._ray_directions_w_torch = wp.to_torch(self._ray_directions_w) - - # Data buffers - self._data.create_buffers(self._view.count, self.num_rays, self._device) - - # Dummy distance/normal buffers required by the merged raycast_mesh_masked_kernel signature. - # Sized (1, 1) even though the kernel is launched at (num_envs, num_rays): the kernel only - # writes to these buffers when return_distance==1 or return_normal==1 respectively, and - # RayCaster always passes 0 for both flags. If those flags are ever enabled here, these - # buffers must be resized to (num_envs, num_rays) to avoid an out-of-bounds write. - self._dummy_ray_distance = wp.empty((1, 1), dtype=wp.float32, device=self._device) - self._dummy_ray_normal = wp.empty((1, 1), dtype=wp.vec3f, device=self._device) - - def _get_view_transforms_wp(self) -> wp.array: - """Get world transforms from the frame view as a warp array of ``wp.transformf``. - - Returns: - Warp array of ``wp.transformf`` with shape ``(num_envs,)``. Layout is - ``(tx, ty, tz, qx, qy, qz, qw)`` per element, matching the quaternion - convention returned by :class:`~isaaclab.sim.views.FrameView`. - """ - pos_w, quat_w = self._view.get_world_poses() - pos_torch = pos_w.torch.reshape(-1, 3) - quat_torch = quat_w.torch.reshape(-1, 4) - poses = torch.cat([pos_torch, quat_torch], dim=-1).contiguous() - return wp.from_torch(poses).view(wp.transformf) - - def _update_ray_infos(self, env_mask: wp.array): - """Updates sensor poses and ray world-frame buffers via a single warp kernel.""" - transforms = self._get_view_transforms_wp() - - wp.launch( - update_ray_caster_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[ - transforms, - env_mask, - self._offset_pos_wp, - self._offset_quat_wp, - self._drift, - self._ray_cast_drift, - self._ray_starts_local, - self._ray_directions_local, - self._alignment_mode, - ], - outputs=[ - self._data._pos_w, - self._data._quat_w, - self._ray_starts_w, - self._ray_directions_w, - ], - device=self._device, - ) - - def _update_buffers_impl(self, env_mask: wp.array): - """Fills the buffers of the sensor data.""" - self._update_ray_infos(env_mask) - - # Fill ray hits with inf before raycasting - wp.launch( - fill_vec3_inf_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._data._ray_hits_w], - device=self._device, - ) - - # Ray-cast against the mesh - wp.launch( - raycast_mesh_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[ - RayCaster.meshes[(self.cfg.mesh_prim_paths[0], self._device)].id, - env_mask, - self._ray_starts_w, - self._ray_directions_w, - float(self.cfg.max_distance), - int(False), # return_distance: not needed by RayCaster - int(False), # return_normal: not needed by RayCaster - self._data._ray_hits_w, - self._dummy_ray_distance, - self._dummy_ray_normal, - ], - device=self._device, - ) - - # Apply vertical drift to ray hits - wp.launch( - apply_z_drift_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, self._ray_cast_drift, self._data._ray_hits_w], - device=self._device, - ) - - def _set_debug_vis_impl(self, debug_vis: bool): - if debug_vis: - if not hasattr(self, "ray_visualizer"): - self.ray_visualizer = VisualizationMarkers(self.cfg.visualizer_cfg) - self.ray_visualizer.set_visibility(True) - else: - if hasattr(self, "ray_visualizer"): - self.ray_visualizer.set_visibility(False) - - def _debug_vis_callback(self, event): - if self._data._ray_hits_w is None: - return - ray_hits_torch = wp.to_torch(self._data._ray_hits_w) - # remove possible inf values - viz_points = ray_hits_torch.reshape(-1, 3) - viz_points = viz_points[~torch.any(torch.isinf(viz_points), dim=1)] - - # if no points to visualize, skip - if viz_points.shape[0] == 0: - return - - self.ray_visualizer.visualize(viz_points) - - """ - Internal simulation callbacks. - """ - - def _invalidate_initialize_callback(self, event): - """Invalidates the scene elements.""" - super()._invalidate_initialize_callback(event) - self._view = None - - def __del__(self): - RayCaster._instance_count -= 1 - if RayCaster._instance_count == 0: - RayCaster.meshes.clear() + _backend_class_names = {"physx": "RayCaster", "newton": "RayCaster"} diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py index 650dfee54ac..13765d23d44 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py @@ -5,647 +5,12 @@ from __future__ import annotations -import logging -from collections.abc import Sequence -from typing import TYPE_CHECKING, ClassVar, Literal +from isaaclab.utils.backend_utils import FactoryBase -import torch -import warp as wp +from .base_ray_caster_camera import BaseRayCasterCamera -from pxr import UsdGeom -import isaaclab.utils.math as math_utils -from isaaclab.sensors.camera import CameraData -from isaaclab.utils.warp import ProxyArray -from isaaclab.utils.warp.kernels import raycast_mesh_masked_kernel +class RayCasterCamera(FactoryBase, BaseRayCasterCamera): + """Backend-dispatching ray-caster camera sensor.""" -from .kernels import ( - ALIGNMENT_BASE, - CAMERA_RAYCAST_MAX_DIST, - apply_depth_clipping_masked_kernel, - compute_distance_to_image_plane_masked_kernel, - fill_float2d_masked_kernel, - fill_vec3_inf_kernel, - update_ray_caster_kernel, -) -from .ray_caster import RayCaster - -if TYPE_CHECKING: - from .ray_caster_camera_cfg import RayCasterCameraCfg - -# import logger -logger = logging.getLogger(__name__) - - -class RayCasterCamera(RayCaster): - """A ray-casting camera sensor. - - The ray-caster camera uses a set of rays to get the distances to meshes in the scene. The rays are - defined in the sensor's local coordinate frame. The sensor has the same interface as the - :class:`isaaclab.sensors.Camera` that implements the camera class through USD camera prims. - However, this class provides a faster image generation. The sensor converts meshes from the list of - primitive paths provided in the configuration to Warp meshes. The camera then ray-casts against these - Warp meshes only. - - Currently, only the following annotators are supported: - - - ``"distance_to_camera"``: An image containing the distance to camera optical center. - - ``"distance_to_image_plane"``: An image containing distances of 3D points from camera plane along camera's z-axis. - - ``"normals"``: An image containing the local surface normal vectors at each pixel. - - .. note:: - Currently, only static meshes are supported. Extending the warp mesh to support dynamic meshes - is a work in progress. - """ - - cfg: RayCasterCameraCfg - """The configuration parameters.""" - UNSUPPORTED_TYPES: ClassVar[set[str]] = { - "rgb", - "instance_id_segmentation", - "instance_id_segmentation_fast", - "instance_segmentation", - "instance_segmentation_fast", - "semantic_segmentation", - "skeleton_data", - "motion_vectors", - "bounding_box_2d_tight", - "bounding_box_2d_tight_fast", - "bounding_box_2d_loose", - "bounding_box_2d_loose_fast", - "bounding_box_3d", - "bounding_box_3d_fast", - } - """A set of sensor types that are not supported by the ray-caster camera.""" - - def __init__(self, cfg: RayCasterCameraCfg): - """Initializes the camera object. - - Args: - cfg: The configuration parameters. - - Raises: - ValueError: If the provided data types are not supported by the ray-caster camera. - """ - # perform check on supported data types - self._check_supported_data_types(cfg) - # initialize base class - super().__init__(cfg) - # create empty variables for storing output data - self._data = CameraData() - - def __str__(self) -> str: - """Returns: A string containing information about the instance.""" - return ( - f"Ray-Caster-Camera @ '{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 : {len(RayCaster.meshes)}\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}\n" - f"\timage shape : {self.image_shape}" - ) - - """ - Properties - """ - - @property - def data(self) -> CameraData: - # update sensors if needed - self._update_outdated_buffers() - # return the data - return self._data - - @property - def image_shape(self) -> tuple[int, int]: - """A tuple containing (height, width) of the camera sensor.""" - return (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width) - - @property - def frame(self) -> torch.tensor: - """Frame number when the measurement took place.""" - return self._frame - - """ - Operations. - """ - - def set_intrinsic_matrices( - self, matrices: torch.Tensor, focal_length: float = 1.0, env_ids: Sequence[int] | None = None - ): - """Set the intrinsic matrix of the camera. - - Args: - matrices: The intrinsic matrices for the camera. Shape is (N, 3, 3). - focal_length: Focal length to use when computing aperture values (in cm). Defaults to 1.0. - env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. - """ - # resolve env_ids - if env_ids is None: - env_ids = slice(None) - # save new intrinsic matrices and focal length - self._data.intrinsic_matrices[env_ids] = matrices.to(self._device) - self._focal_length = focal_length - # recompute ray directions - self.ray_starts[env_ids], self.ray_directions[env_ids] = self.cfg.pattern_cfg.func( - self.cfg.pattern_cfg, self._data.intrinsic_matrices[env_ids], self._device - ) - # Refresh warp views of local ray buffers; .contiguous() may produce a copy so we store - # the contiguous tensors explicitly to prevent GC while the warp views are alive. - if hasattr(self, "_ray_starts_local"): - self._ray_starts_contiguous = self.ray_starts.contiguous() - self._ray_directions_contiguous = self.ray_directions.contiguous() - self._ray_starts_local = wp.from_torch(self._ray_starts_contiguous, dtype=wp.vec3f) - self._ray_directions_local = wp.from_torch(self._ray_directions_contiguous, dtype=wp.vec3f) - - def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): - # reset the timestamps - super().reset(env_ids, env_mask) - # resolve to indices for torch indexing - if env_ids is None and env_mask is not None: - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - elif env_ids is None or isinstance(env_ids, slice): - env_ids = self._ALL_INDICES - if not isinstance(env_ids, torch.Tensor): - env_ids = torch.tensor(env_ids, dtype=torch.long, device=self._device) - # reset the data - # note: this recomputation is useful if one performs events such as randomizations on the camera poses. - indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_w, quat_w = self._view.get_world_poses(indices) - pos_w, quat_w = math_utils.combine_frame_transforms( - pos_w.torch.clone(), quat_w.torch.clone(), self._offset_pos[env_ids], self._offset_quat[env_ids] - ) - self._data.pos_w[env_ids] = pos_w - self._data.quat_w_world[env_ids] = quat_w - # Reset the frame count - self._frame[env_ids] = 0 - - def set_world_poses( - self, - positions: torch.Tensor | None = None, - orientations: torch.Tensor | None = None, - env_ids: Sequence[int] | None = None, - convention: Literal["opengl", "ros", "world"] = "ros", - ): - """Set the pose of the camera w.r.t. the world frame using specified convention. - - Since different fields use different conventions for camera orientations, the method allows users to - set the camera poses in the specified convention. Possible conventions are: - - - :obj:`"opengl"` - forward axis: -Z - up axis +Y - Offset is applied in the OpenGL (Usd.Camera) convention - - :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention - - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention - - See :meth:`isaaclab.utils.math.convert_camera_frame_orientation_convention` for more details - on the conventions. - - Args: - positions: The cartesian coordinates (in meters). Shape is (N, 3). - Defaults to None, in which case the camera position in not changed. - orientations: The quaternion orientation in (x, y, z, w). Shape is (N, 4). - Defaults to None, in which case the camera orientation in not changed. - env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. - convention: The convention in which the poses are fed. Defaults to "ros". - - Raises: - RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. - """ - # resolve env_ids - if env_ids is None or isinstance(env_ids, slice): - env_ids = self._ALL_INDICES - - # get current positions - indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_w, quat_w = self._view.get_world_poses(indices) - pos_w_torch = pos_w.torch - quat_w_torch = quat_w.torch - if positions is not None: - # transform to camera frame - pos_offset_world_frame = positions - pos_w_torch - self._offset_pos[env_ids] = math_utils.quat_apply(math_utils.quat_inv(quat_w_torch), pos_offset_world_frame) - if orientations is not None: - # convert rotation matrix from input convention to world - quat_w_set = math_utils.convert_camera_frame_orientation_convention( - orientations, origin=convention, target="world" - ) - self._offset_quat[env_ids] = math_utils.quat_mul(math_utils.quat_inv(quat_w_torch), quat_w_set) - - # update the data - pos_w2, quat_w2 = self._view.get_world_poses(indices) - pos_w_out, quat_w_out = math_utils.combine_frame_transforms( - pos_w2.torch.clone(), quat_w2.torch.clone(), self._offset_pos[env_ids], self._offset_quat[env_ids] - ) - self._data.pos_w[env_ids] = pos_w_out - self._data.quat_w_world[env_ids] = quat_w_out - - def set_world_poses_from_view( - self, eyes: torch.Tensor, targets: torch.Tensor, env_ids: Sequence[int] | None = None - ): - """Set the poses of the camera from the eye position and look-at target position. - - Args: - eyes: The positions of the camera's eye. Shape is (N, 3). - targets: The target locations to look at. Shape is (N, 3). - env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices. - - Raises: - RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. - NotImplementedError: If the stage up-axis is not "Y" or "Z". - ValueError: If every eye position equals its target (look-at direction undefined for the - whole batch). When only some rows are degenerate, those rows are skipped and the - remaining poses are still applied; a warning is logged. - """ - # resolve env_ids to a tensor up front so we can index it during partial-failure filtering - if env_ids is None: - env_ids = self._ALL_INDICES - if not isinstance(env_ids, torch.Tensor): - env_ids = torch.tensor(env_ids, dtype=torch.long, device=self._device) - # get up axis of current stage - up_axis = UsdGeom.GetStageUpAxis(self.stage) - # camera position and rotation in opengl convention; degenerate rows (eye == target) come back as NaN - rotation_matrix = math_utils.create_rotation_matrix_from_view( - eyes, targets, up_axis=up_axis, device=self._device - ) - valid_indices = (~torch.isnan(rotation_matrix).any(dim=(-2, -1))).nonzero(as_tuple=True)[0] - n_valid = valid_indices.numel() - n_total = rotation_matrix.shape[0] - if n_valid == 0: - raise ValueError("look-at is undefined: every eye position equals its target") - if n_valid < n_total: - logger.warning( - "set_world_poses_from_view: skipping %d pose(s) where eye equals target", - n_total - n_valid, - ) - rotation_matrix = rotation_matrix.index_select(0, valid_indices) - eyes = eyes.index_select(0, valid_indices) - env_ids = env_ids.index_select(0, valid_indices) - orientations = math_utils.quat_from_matrix(rotation_matrix) - self.set_world_poses(eyes, orientations, env_ids, convention="opengl") - - """ - Implementation. - """ - - def _initialize_rays_impl(self): - # Create all indices buffer - self._ALL_INDICES = torch.arange(self._view.count, device=self._device, dtype=torch.long) - # Create frame count buffer - self._frame = torch.zeros(self._view.count, device=self._device, dtype=torch.long) - # create buffers - self._create_buffers() - # compute intrinsic matrices - self._compute_intrinsic_matrices() - # compute ray starts and directions - self.ray_starts, self.ray_directions = self.cfg.pattern_cfg.func( - self.cfg.pattern_cfg, self._data.intrinsic_matrices, self._device - ) - self.num_rays = self.ray_directions.shape[1] - - # Offset buffers: warp-primary so the kernel always sees the current values without re-wrapping. - # Zero-copy torch views (_offset_pos, _offset_quat) are used by set_world_poses for indexed writes. - self._offset_pos_wp = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) - self._offset_quat_wp = wp.zeros(self._view.count, dtype=wp.quatf, device=self._device) - self._offset_pos = wp.to_torch(self._offset_pos_wp) - self._offset_quat = wp.to_torch(self._offset_quat_wp) - # Initialize from config - quat_w = math_utils.convert_camera_frame_orientation_convention( - torch.tensor([self.cfg.offset.rot], device=self._device), origin=self.cfg.offset.convention, target="world" - ) - self._offset_pos[:] = torch.tensor(list(self.cfg.offset.pos), device=self._device) - self._offset_quat[:] = quat_w - - # Warp buffers for world-frame rays (used by update kernel) - self._ray_starts_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - self._ray_directions_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - - # Warp views for ray_starts and ray_directions (from torch tensors returned by pattern_cfg.func) - # These are (num_envs, num_rays, 3) torch tensors; wrap as warp vec3f arrays. - # Store contiguous tensors explicitly so they are not garbage-collected while the - # warp views are alive (mirrors the pattern in RayCaster._initialize_impl). - self._ray_starts_contiguous = self.ray_starts.contiguous() - self._ray_directions_contiguous = self.ray_directions.contiguous() - self._ray_starts_local = wp.from_torch(self._ray_starts_contiguous, dtype=wp.vec3f) - self._ray_directions_local = wp.from_torch(self._ray_directions_contiguous, dtype=wp.vec3f) - - # Wrap the torch drift buffers (created in _create_buffers) as warp arrays (zero-copy). - # Cameras do not apply positional drift, so these remain zero. - self._drift_contiguous = self.drift.contiguous() - self._ray_cast_drift_contiguous = self.ray_cast_drift.contiguous() - self._drift = wp.from_torch(self._drift_contiguous, dtype=wp.vec3f) - self._ray_cast_drift = wp.from_torch(self._ray_cast_drift_contiguous, dtype=wp.vec3f) - - # Warp buffers for camera pose outputs - self._pos_w_wp = wp.zeros(self._view.count, dtype=wp.vec3f, device=self._device) - self._quat_w_wp = wp.zeros(self._view.count, dtype=wp.quatf, device=self._device) - - # Intermediate warp buffers for ray results (filled with inf before each raycasting step) - self._ray_distance = wp.zeros((self._view.count, self.num_rays), dtype=wp.float32, device=self._device) - if "normals" in self.cfg.data_types: - self._ray_normal_w = wp.zeros((self._view.count, self.num_rays), dtype=wp.vec3f, device=self._device) - else: - self._ray_normal_w = wp.zeros((1, 1), dtype=wp.vec3f, device=self._device) - - if "distance_to_image_plane" in self.cfg.data_types: - self._distance_to_image_plane_wp = wp.zeros( - (self._view.count, self.num_rays), dtype=wp.float32, device=self._device - ) - - # Torch buffer for ray hits (used by debug visualizer) - self.ray_hits_w = torch.full((self._view.count, self.num_rays, 3), float("inf"), device=self._device) - # Warp view of ray_hits_w - self._ray_hits_w_wp = wp.from_torch(self.ray_hits_w.contiguous(), dtype=wp.vec3f) - - # Cache zero-copy torch views of warp output buffers to avoid per-step wrapper allocation. - self._pos_w_torch = wp.to_torch(self._pos_w_wp) - self._quat_w_torch = wp.to_torch(self._quat_w_wp) - self._ray_distance_torch = wp.to_torch(self._ray_distance) - if "distance_to_image_plane" in self.cfg.data_types: - self._distance_to_image_plane_torch = wp.to_torch(self._distance_to_image_plane_wp) - if "normals" in self.cfg.data_types: - self._ray_normal_w_torch = wp.to_torch(self._ray_normal_w) - - def _update_buffers_impl(self, env_mask: wp.array): - """Fills the buffers of the sensor data.""" - # Convert mask to indices for torch-indexed writes - env_ids = wp.to_torch(env_mask).nonzero(as_tuple=False).squeeze(-1) - if len(env_ids) == 0: - return - # increment frame count - self._frame[env_ids] += 1 - - # Update world-frame ray starts/directions and camera pose via warp kernel. - # Camera always uses ALIGNMENT_BASE (full orientation) and zero drift. - transforms = self._get_view_transforms_wp() - wp.launch( - update_ray_caster_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[ - transforms, - env_mask, - self._offset_pos_wp, - self._offset_quat_wp, - self._drift, - self._ray_cast_drift, - self._ray_starts_local, - self._ray_directions_local, - int(ALIGNMENT_BASE), - ], - outputs=[ - self._pos_w_wp, - self._quat_w_wp, - self._ray_starts_w, - self._ray_directions_w, - ], - device=self._device, - ) - - # Write camera pose to CameraData (torch tensors) - self._data.pos_w[env_ids] = self._pos_w_torch[env_ids] - self._data.quat_w_world[env_ids] = self._quat_w_torch[env_ids] - - # Fill ray hit positions with inf before raycasting - wp.launch( - fill_vec3_inf_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._ray_hits_w_wp], - device=self._device, - ) - - # Fill ray distance with inf before raycasting - wp.launch( - fill_float2d_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._ray_distance], - device=self._device, - ) - - # Determine whether to compute normals - need_normal = int("normals" in self.cfg.data_types) - if need_normal: - # Fill normal buffer with inf before raycasting - wp.launch( - fill_vec3_inf_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float("inf"), self._ray_normal_w], - device=self._device, - ) - - # Ray-cast against the mesh; use a large upper-bound max_dist so depth clipping - # can be applied per-data-type afterwards (matching the original behaviour). - wp.launch( - raycast_mesh_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[ - RayCaster.meshes[(self.cfg.mesh_prim_paths[0], self._device)].id, - env_mask, - self._ray_starts_w, - self._ray_directions_w, - float(CAMERA_RAYCAST_MAX_DIST), - int(True), # return_distance: always needed for depth output - need_normal, - self._ray_hits_w_wp, - self._ray_distance, - self._ray_normal_w, - ], - device=self._device, - ) - - # Compute distance_to_image_plane using a warp kernel - if "distance_to_image_plane" in self.cfg.data_types: - wp.launch( - compute_distance_to_image_plane_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[ - env_mask, - self._quat_w_wp, - self._ray_distance, - self._ray_directions_w, - ], - outputs=[ - self._distance_to_image_plane_wp, - ], - device=self._device, - ) - # Apply depth clipping on the intermediate buffer (leaves _ray_distance unmodified) - self._apply_depth_clipping(env_mask, self._distance_to_image_plane_wp) - self._data.output["distance_to_image_plane"][env_ids] = self._distance_to_image_plane_torch[env_ids].view( - -1, *self.image_shape, 1 - ) - - if "distance_to_camera" in self.cfg.data_types: - # d2ip (if requested) was computed before this block so _ray_distance is still unclipped. - self._apply_depth_clipping(env_mask, self._ray_distance) - self._data.output["distance_to_camera"][env_ids] = self._ray_distance_torch[env_ids].view( - -1, *self.image_shape, 1 - ) - - if "normals" in self.cfg.data_types: - self._data.output["normals"][env_ids] = self._ray_normal_w_torch[env_ids].view(-1, *self.image_shape, 3) - - def _debug_vis_callback(self, event): - # in case it crashes be safe - if not hasattr(self, "ray_hits_w"): - return - # filter out missed rays (inf values) before visualizing - ray_hits_flat = self.ray_hits_w.reshape(-1, 3) - valid_mask = ~torch.isinf(ray_hits_flat).any(dim=-1) - viz_points = ray_hits_flat[valid_mask] - # if no valid hits, skip - if viz_points.shape[0] == 0: - return - self.ray_visualizer.visualize(viz_points) - - """ - Private Helpers - """ - - def _apply_depth_clipping(self, env_mask: wp.array, depth: wp.array) -> None: - """Apply depth clipping in-place on a warp float32 buffer. - - Uses :attr:`cfg.depth_clipping_behavior` to determine the fill value: - ``"max"`` replaces out-of-range and NaN values with :attr:`cfg.max_distance`; - ``"zero"`` replaces them with 0. No-op when behavior is ``"none"``. - - Args: - env_mask: Boolean mask selecting which environments to update. Shape is (num_envs,). - depth: Warp 2-D float32 buffer to clip in-place. Shape is (num_envs, num_rays). - """ - if self.cfg.depth_clipping_behavior == "max": - wp.launch( - apply_depth_clipping_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float(self.cfg.max_distance), float(self.cfg.max_distance), depth], - device=self._device, - ) - elif self.cfg.depth_clipping_behavior == "zero": - wp.launch( - apply_depth_clipping_masked_kernel, - dim=(self._num_envs, self.num_rays), - inputs=[env_mask, float(self.cfg.max_distance), float(0.0), depth], - device=self._device, - ) - elif self.cfg.depth_clipping_behavior == "none": - pass # no clipping: inf values remain as-is - else: - raise ValueError( - f"Unknown depth_clipping_behavior: {self.cfg.depth_clipping_behavior!r}." - " Valid values are 'max', 'zero', and 'none'." - ) - - def _check_supported_data_types(self, cfg: RayCasterCameraCfg): - """Checks if the data types are supported by the ray-caster camera.""" - # check if there is any intersection in unsupported types - # reason: we cannot obtain this data from simplified warp-based ray caster - common_elements = set(cfg.data_types) & RayCasterCamera.UNSUPPORTED_TYPES - if common_elements: - raise ValueError( - f"RayCasterCamera class does not support the following sensor types: {common_elements}." - "\n\tThis is because these sensor types cannot be obtained in a fast way using ''warp''." - "\n\tHint: If you need to work with these sensor types, we recommend using the USD camera" - " interface from the isaaclab.sensors.camera module." - ) - - def _create_buffers(self): - """Create buffers for storing data.""" - # prepare drift (kept as torch tensors so subclasses may use torch indexing) - self.drift = torch.zeros(self._view.count, 3, device=self.device) - self.ray_cast_drift = torch.zeros(self._view.count, 3, device=self.device) - # create the data object - # -- pose and intrinsics as warp-backed ProxyArrays - device_str = self._device if isinstance(self._device, str) else str(self._device) - self._data.create_buffers(self._view.count, device_str) - self._data.intrinsic_matrices.torch[:, 2, 2] = 1.0 - self._data.image_shape = self.image_shape - # -- output data as warp-backed ProxyArrays - output = {} - self._data.info = {name: None for name in self.cfg.data_types} - for name in self.cfg.data_types: - if name in ["distance_to_image_plane", "distance_to_camera"]: - shape = (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width, 1) - elif name in ["normals"]: - shape = (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width, 3) - else: - raise ValueError(f"Received unknown data type: {name}. Please check the configuration.") - wp_arr = wp.zeros((self._view.count, *shape), dtype=wp.float32, device=device_str) - output[name] = ProxyArray(wp_arr) - self._data._output = output - - def _compute_intrinsic_matrices(self): - """Computes the intrinsic matrices for the camera based on the config provided.""" - # get the sensor properties - pattern_cfg = self.cfg.pattern_cfg - - # check if vertical aperture is provided - # if not then it is auto-computed based on the aspect ratio to preserve squared pixels - if pattern_cfg.vertical_aperture is None: - pattern_cfg.vertical_aperture = pattern_cfg.horizontal_aperture * pattern_cfg.height / pattern_cfg.width - - # compute the intrinsic matrix - f_x = pattern_cfg.width * pattern_cfg.focal_length / pattern_cfg.horizontal_aperture - f_y = pattern_cfg.height * pattern_cfg.focal_length / pattern_cfg.vertical_aperture - c_x = pattern_cfg.horizontal_aperture_offset * f_x + pattern_cfg.width / 2 - c_y = pattern_cfg.vertical_aperture_offset * f_y + pattern_cfg.height / 2 - # allocate the intrinsic matrices - self._data.intrinsic_matrices[:, 0, 0] = f_x - self._data.intrinsic_matrices[:, 0, 2] = c_x - self._data.intrinsic_matrices[:, 1, 1] = f_y - self._data.intrinsic_matrices[:, 1, 2] = c_y - - # save focal length - self._focal_length = pattern_cfg.focal_length - - def _compute_view_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]: - """Obtains the pose of the view the camera is attached to in the world frame. - - .. deprecated v2.3.1: - This function will be removed in a future release. Call - ``self._view.get_world_poses(indices)`` directly instead. The returned - ProxyArray pair exposes ``.warp`` and ``.torch`` accessors. - - Returns: - A tuple of the position (in meters) and quaternion (x, y, z, w). - - - """ - logger.warning( - "The function '_compute_view_world_poses' is deprecated." - " Call 'self._view.get_world_poses(indices)' directly instead." - ) - - indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_w, quat_w = self._view.get_world_poses(indices) - return pos_w.torch.clone(), quat_w.torch.clone() - - def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]: - """Computes the pose of the camera in the world frame. - - This function applies the offset pose to the pose of the view the camera is attached to. - - .. deprecated v2.3.1: - This function will be removed in a future release. Instead, use the code block below: - - .. code-block:: python - - indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) - pos_w, quat_w = self._view.get_world_poses(indices) - # The returned ProxyArray pair exposes .warp and .torch accessors - pos_w, quat_w = pos_w.torch.clone(), quat_w.torch.clone() - pos_w, quat_w = math_utils.combine_frame_transforms( - pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] - ) - - Returns: - A tuple of the position (in meters) and quaternion (x, y, z, w) in "world" convention. - """ - logger.warning( - "The function '_compute_camera_world_poses' is deprecated." - " Call 'self._view.get_world_poses(indices)' and 'math_utils.combine_frame_transforms' directly instead." - ) - - indices = wp.from_torch(env_ids.to(dtype=torch.int32), dtype=wp.int32) if env_ids is not None else None - pos_w, quat_w = self._view.get_world_poses(indices) - return math_utils.combine_frame_transforms( - pos_w.torch.clone(), quat_w.torch.clone(), self._offset_pos[env_ids], self._offset_quat[env_ids] - ) + _backend_class_names = {"physx": "RayCasterCamera", "newton": "RayCasterCamera"} diff --git a/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py index eed77643127..5823e6c79f8 100644 --- a/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py +++ b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py @@ -96,16 +96,36 @@ def design_scene(sim: SimulationContext, num_envs: int = 2048): f"/World/envs/env_0/object_{i}", object, translation=(0.0 + random.random(), 0.0 + random.random(), 1.0), - orientation=quat_from_euler_xyz(torch.Tensor(0), torch.Tensor(0), torch.rand(1) * torch.pi).numpy(), + orientation=quat_from_euler_xyz(torch.zeros(1), torch.zeros(1), torch.rand(1) * torch.pi)[0].numpy(), ) # Clone the scene envs_prim_paths = [f"/World/envs/env_{i}" for i in range(num_envs)] lab_cloner.usd_replicate(sim.stage, [env_fmt.format(0)], [env_fmt], env_ids, positions=env_origins) - physics_scene_path = sim.get_physics_context().prim_path - lab_cloner.filter_collisions( - sim.stage, physics_scene_path, "/World/collisions", prim_paths=envs_prim_paths, global_paths=["/World/ground"] + # Publish a trivial homogeneous ClonePlan so consumers (e.g. multi-mesh ray-caster's + # target tracker) can drive per-env work via clone_mask. Mirrors InteractiveScene's + # synthesis path for hand-authored scenes that bypass it. + sim.set_clone_plan( + lab_cloner.ClonePlan( + sources=(env_fmt.format(0),), + destinations=(env_fmt,), + clone_mask=torch.ones((1, num_envs), dtype=torch.bool, device=sim.device), + ) + ) + # PhysX-only optimization: filter collisions across env clones. Skip on Newton — + # PhysxSceneAPI isn't applied there and the cloner helper is PhysX-specific. + physics_scene_path = next( + (prim.GetPrimPath().pathString for prim in sim.stage.Traverse() if "PhysxSceneAPI" in prim.GetAppliedSchemas()), + None, ) + if physics_scene_path is not None: + lab_cloner.filter_collisions( + sim.stage, + physics_scene_path, + "/World/collisions", + prim_paths=envs_prim_paths, + global_paths=["/World/ground"], + ) def main(): @@ -127,6 +147,7 @@ def main(): usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", max_init_terrain_level=0, num_envs=1, + env_spacing=10.0, ) _ = TerrainImporter(terrain_importer_cfg) diff --git a/source/isaaclab/test/sensors/check_ray_caster.py b/source/isaaclab/test/sensors/check_ray_caster.py index 821fc247743..87d8a9594ad 100644 --- a/source/isaaclab/test/sensors/check_ray_caster.py +++ b/source/isaaclab/test/sensors/check_ray_caster.py @@ -78,10 +78,21 @@ def design_scene(sim: SimulationContext, num_envs: int = 2048): # Clone the scene envs_prim_paths = [f"/World/envs/env_{i}" for i in range(num_envs)] lab_cloner.usd_replicate(sim.stage, [env_fmt.format(0)], [env_fmt], env_ids, positions=env_origins) - physics_scene_path = sim.get_physics_context().prim_path - lab_cloner.filter_collisions( - sim.stage, physics_scene_path, "/World/collisions", prim_paths=envs_prim_paths, global_paths=["/World/ground"] - ) + # PhysX-only optimization: filter collisions across env clones. Skip on Newton — + # PhysxSceneAPI isn't applied there and the cloner helper is PhysX-specific. + physics_scene_path = None + for prim in sim.stage.Traverse(): + if "PhysxSceneAPI" in prim.GetAppliedSchemas(): + physics_scene_path = prim.GetPrimPath().pathString + break + if physics_scene_path is not None: + lab_cloner.filter_collisions( + sim.stage, + physics_scene_path, + "/World/collisions", + prim_paths=envs_prim_paths, + global_paths=["/World/ground"], + ) def main(): @@ -103,6 +114,7 @@ def main(): usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", max_init_terrain_level=None, num_envs=1, + env_spacing=10.0, ) _ = TerrainImporter(terrain_importer_cfg) diff --git a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py index 0319bf80ef1..28501f1a733 100644 --- a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py +++ b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py @@ -16,6 +16,7 @@ """Rest everything follows.""" import copy +from collections.abc import Callable import numpy as np import pytest @@ -24,18 +25,26 @@ import omni.replicator.core as rep from pxr import Gf +import isaaclab.cloner as lab_cloner import isaaclab.sim as sim_utils +from isaaclab.cloner import ClonePlan from isaaclab.sensors.camera import Camera, CameraCfg from isaaclab.sensors.ray_caster import MultiMeshRayCasterCamera, MultiMeshRayCasterCameraCfg, patterns from isaaclab.sim import PinholeCameraCfg from isaaclab.terrains.trimesh.utils import make_plane from isaaclab.terrains.utils import create_prim_from_mesh +from isaaclab_assets.robots.anymal import ANYMAL_C_CFG +from isaaclab_assets.robots.spot import SPOT_CFG + # sample camera poses (quaternions in xyzw format) POSITION = [2.5, 2.5, 2.5] QUAT_ROS = [0.33985114, 0.82047325, -0.42470819, -0.17591989] QUAT_OPENGL = [0.17591988, 0.42470818, 0.82047324, 0.33985113] QUAT_WORLD = [-0.27984815, -0.1159169, 0.88047623, -0.3647052] +MESH_ID_GROUND = 0 +MESH_ID_OBJECT = 1 +MESH_ID_ROBOT_MIN = 2 def _assert_quat_close(actual, expected, **kwargs): @@ -127,7 +136,7 @@ def test_camera_init_offset(setup_simulation, convention, quat): camera.update(dt) # check that transform is set correctly - np.testing.assert_allclose(camera.data.pos_w[0].cpu().numpy(), cam_cfg_offset.offset.pos) + np.testing.assert_allclose(camera.data.pos_w.torch[0].cpu().numpy(), cam_cfg_offset.offset.pos) del camera @@ -190,7 +199,7 @@ def test_camera_init_intrinsic_matrix(setup_simulation): camera_1 = MultiMeshRayCasterCamera(cfg=camera_cfg) # get intrinsic matrix sim.reset() - intrinsic_matrix = camera_1.data.intrinsic_matrices[0].cpu().flatten().tolist() + intrinsic_matrix = camera_1.data.intrinsic_matrices.torch[0].cpu().flatten().tolist() # initialize from intrinsic matrix intrinsic_camera_cfg = MultiMeshRayCasterCameraCfg( @@ -219,13 +228,13 @@ def test_camera_init_intrinsic_matrix(setup_simulation): # check image data torch.testing.assert_close( - camera_1.data.output["distance_to_image_plane"], - camera_2.data.output["distance_to_image_plane"], + camera_1.data.output["distance_to_image_plane"].torch, + camera_2.data.output["distance_to_image_plane"].torch, ) # check that both intrinsic matrices are the same torch.testing.assert_close( - camera_1.data.intrinsic_matrices[0], - camera_2.data.intrinsic_matrices[0], + camera_1.data.intrinsic_matrices.torch[0], + camera_2.data.intrinsic_matrices.torch[0], ) del camera_1, camera_2 @@ -408,8 +417,8 @@ def test_output_equal_to_usdcamera(setup_simulation, data_types): # check the intrinsic matrices torch.testing.assert_close( - camera_usd.data.intrinsic_matrices, - camera_warp.data.intrinsic_matrices, + camera_usd.data.intrinsic_matrices.torch, + camera_warp.data.intrinsic_matrices.torch, ) # check the apertures @@ -423,8 +432,8 @@ def test_output_equal_to_usdcamera(setup_simulation, data_types): if data_type in camera_usd.data.output and data_type in camera_warp.data.output: if data_type == "distance_to_camera" or data_type == "distance_to_image_plane": torch.testing.assert_close( - camera_usd.data.output[data_type], - camera_warp.data.output[data_type], + camera_usd.data.output[data_type].torch, + camera_warp.data.output[data_type].torch, atol=5e-5, rtol=5e-6, ) @@ -438,13 +447,187 @@ def test_output_equal_to_usdcamera(setup_simulation, data_types): ) else: torch.testing.assert_close( - camera_usd.data.output[data_type], - camera_warp.data.output[data_type], + camera_usd.data.output[data_type].torch, + camera_warp.data.output[data_type].torch, ) del camera_usd, camera_warp +def _create_heterogeneous_clone_scene(sim: sim_utils.SimulationContext, num_envs: int) -> torch.Tensor: + """Create alternating Spot/ANYmal and cube/sphere cloned environments.""" + stage = sim_utils.get_current_stage() + env_fmt = "/World/envs/env_{}" + env_ids = torch.arange(num_envs, dtype=torch.long, device=sim.device) + env_origins, _ = lab_cloner.grid_transforms(num_envs, spacing=4.0, device=sim.device) + + sim_utils.create_prim("/World/envs", "Xform", stage=stage) + for env_id, origin in enumerate(env_origins.cpu().tolist()): + sim_utils.create_prim(env_fmt.format(env_id), "Xform", translation=tuple(origin), stage=stage) + + robot_mask = torch.zeros((2, num_envs), dtype=torch.bool, device=sim.device) + robot_mask[0, 0::2] = True + robot_mask[1, 1::2] = True + object_mask = robot_mask.clone() + + spot_spawn = copy.deepcopy(SPOT_CFG.spawn) + anymal_spawn = copy.deepcopy(ANYMAL_C_CFG.spawn) + spot_spawn.func(env_fmt.format(0) + "/Robot", spot_spawn, translation=SPOT_CFG.init_state.pos) + anymal_spawn.func(env_fmt.format(1) + "/Robot", anymal_spawn, translation=ANYMAL_C_CFG.init_state.pos) + + cube_cfg = sim_utils.CuboidCfg( + size=(0.35, 0.25, 0.25), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.7, 0.2, 0.2)), + ) + sphere_cfg = sim_utils.SphereCfg( + radius=0.18, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.2, 0.2, 0.7)), + ) + cube_spawn = cube_cfg.func + sphere_spawn = sphere_cfg.func + assert isinstance(cube_spawn, Callable) + assert isinstance(sphere_spawn, Callable) + cube_spawn(env_fmt.format(0) + "/Object", cube_cfg, translation=(0.45, 0.0, 0.25)) + sphere_spawn(env_fmt.format(1) + "/Object", sphere_cfg, translation=(0.45, 0.0, 0.25)) + + lab_cloner.usd_replicate( + stage, + [env_fmt.format(i) + f"/{asset_name}" for asset_name in ("Robot", "Object") for i in range(2)], + [env_fmt + "/Robot", env_fmt + "/Robot", env_fmt + "/Object", env_fmt + "/Object"], + env_ids, + mask=torch.cat([robot_mask, object_mask], dim=0), + ) + + sim.set_clone_plan( + ClonePlan( + sources=( + env_fmt.format(0) + "/Robot", + env_fmt.format(1) + "/Robot", + env_fmt.format(0) + "/Object", + env_fmt.format(1) + "/Object", + ), + destinations=( + env_fmt + "/Robot", + env_fmt + "/Robot", + env_fmt + "/Object", + env_fmt + "/Object", + ), + clone_mask=torch.cat([robot_mask, object_mask], dim=0), + ) + ) + sim_utils.update_stage() + return env_origins + + +@pytest.mark.isaacsim_ci +def test_depth_output_equal_to_usd_camera_heterogeneous_scene(setup_simulation): + """Compare ray-caster and USD depth cameras in a heterogeneous cloned scene. + + The scene contains 16 environments with alternating Spot / ANYmal-C robot + prototypes and alternating cube / sphere objects. The ray-caster consumes + the same clone plan used to build the USD scene and should match the batched + USD camera's stable ``distance_to_image_plane`` pixels for every environment. + """ + sim, dt, _ = setup_simulation + num_envs = 16 + env_origins = _create_heterogeneous_clone_scene(sim, num_envs) + + height, width = 96, 128 + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=height, + width=width, + ) + mesh_prim_paths = [ + "/World/defaultGroundPlane", + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + prim_expr="/World/envs/env_.*/Object", + track_mesh_transforms=False, + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + prim_expr="/World/envs/env_.*/Robot/.+", + track_mesh_transforms=True, + ), + ] + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/envs/env_.*/RayCasterCamera", + mesh_prim_paths=mesh_prim_paths, + update_period=0, + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + max_distance=25.0, + data_types=["distance_to_image_plane"], + depth_clipping_behavior="max", + update_mesh_ids=True, + ) + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + camera_cfg_usd = CameraCfg( + height=height, + width=width, + prim_path="/World/envs/env_.*/UsdCamera", + update_period=0, + data_types=["distance_to_image_plane"], + spawn=PinholeCameraCfg( + focal_length=24.0, + focus_distance=400.0, + horizontal_aperture=20.955, + clipping_range=(0.01, 25.0), + ), + ) + camera_usd = Camera(camera_cfg_usd) + + sim.reset() + sim.play() + + eyes = env_origins + torch.tensor((1.8, -2.5, 2.5), dtype=torch.float32, device=sim.device) + targets = env_origins + torch.tensor((0.0, 0.0, 0.0), dtype=torch.float32, device=sim.device) + camera_warp.set_world_poses_from_view(eyes=eyes, targets=targets) + camera_usd.set_world_poses_from_view(eyes=eyes, targets=targets) + + for _ in range(5): + sim.render() + + camera_usd.update(dt) + camera_warp.update(dt) + + ray_depth = camera_warp.data.output["distance_to_image_plane"].torch + usd_depth = camera_usd.data.output["distance_to_image_plane"].torch + assert ray_depth.shape == (num_envs, height, width, 1) + assert usd_depth.shape == ray_depth.shape + depth_diff = (ray_depth - usd_depth).abs() + mesh_ids_proxy = getattr(camera_warp.data, "image_mesh_ids", None) + assert mesh_ids_proxy is not None + mesh_ids = mesh_ids_proxy.torch + assert torch.any(mesh_ids == MESH_ID_OBJECT), "Expected object pixels in the heterogeneous scene" + assert torch.any(mesh_ids >= MESH_ID_ROBOT_MIN), "Expected robot pixels in the heterogeneous scene" + + # The RTX and ray-cast backends can disagree by a pixel along complex robot + # silhouettes. Compare the stable ground pixels after dilating object/robot + # edges and depth discontinuities. + target_mask = mesh_ids[..., 0] != 0 + discontinuity_mask = torch.zeros_like(target_mask) + for depth in (ray_depth, usd_depth): + depth_image = depth[..., 0] + discontinuity_mask[:, 1:, :] |= (depth_image[:, 1:, :] - depth_image[:, :-1, :]).abs() > 0.3 + discontinuity_mask[:, :, 1:] |= (depth_image[:, :, 1:] - depth_image[:, :, :-1]).abs() > 0.3 + edge_mask = target_mask | discontinuity_mask + silhouette_mask = torch.nn.functional.max_pool2d( + edge_mask[:, None, :, :].float(), kernel_size=21, stride=1, padding=10 + ).to(dtype=torch.bool) + stable_mask = ~silhouette_mask[:, 0, :, :, None] + assert stable_mask.float().mean() > 0.7 + stable_ray_depth = ray_depth[stable_mask] + stable_usd_depth = usd_depth[stable_mask] + stable_depth_diff = depth_diff[stable_mask] + stable_close = torch.isclose(stable_ray_depth, stable_usd_depth, atol=5e-5, rtol=5e-6) + assert stable_close.float().mean() > 0.999 + assert torch.quantile(stable_depth_diff, 0.999) < 5.0e-5 + + del camera_usd, camera_warp + + @pytest.mark.isaacsim_ci def test_output_equal_to_usdcamera_offset(setup_simulation): """Test that ray caster camera output equals USD camera output with offset.""" @@ -497,14 +680,14 @@ def test_output_equal_to_usdcamera_offset(setup_simulation): # check image data torch.testing.assert_close( - camera_usd.data.output["distance_to_image_plane"], - camera_warp.data.output["distance_to_image_plane"], + camera_usd.data.output["distance_to_image_plane"].torch, + camera_warp.data.output["distance_to_image_plane"].torch, atol=5e-5, rtol=5e-6, ) torch.testing.assert_close( - camera_usd.data.output["distance_to_camera"], - camera_warp.data.output["distance_to_camera"], + camera_usd.data.output["distance_to_camera"].torch, + camera_warp.data.output["distance_to_camera"].torch, atol=5e-5, rtol=5e-6, ) @@ -593,14 +776,14 @@ def test_output_equal_to_usdcamera_prim_offset(setup_simulation): # check image data torch.testing.assert_close( - camera_usd.data.output["distance_to_image_plane"], - camera_warp.data.output["distance_to_image_plane"], + camera_usd.data.output["distance_to_image_plane"].torch, + camera_warp.data.output["distance_to_image_plane"].torch, atol=5e-5, rtol=5e-6, ) torch.testing.assert_close( - camera_usd.data.output["distance_to_camera"], - camera_warp.data.output["distance_to_camera"], + camera_usd.data.output["distance_to_camera"].torch, + camera_warp.data.output["distance_to_camera"].torch, rtol=4e-6, atol=2e-5, ) @@ -681,15 +864,17 @@ def test_output_equal_to_usd_camera_intrinsics(setup_simulation, height, width): camera_warp.update(dt) # filter nan and inf from output - cam_warp_output = camera_warp.data.output["distance_to_image_plane"].clone() - cam_usd_output = camera_usd.data.output["distance_to_image_plane"].clone() + cam_warp_output = camera_warp.data.output["distance_to_image_plane"].torch.clone() + cam_usd_output = camera_usd.data.output["distance_to_image_plane"].torch.clone() cam_warp_output[torch.isnan(cam_warp_output)] = 0 cam_warp_output[torch.isinf(cam_warp_output)] = 0 cam_usd_output[torch.isnan(cam_usd_output)] = 0 cam_usd_output[torch.isinf(cam_usd_output)] = 0 # check that both have the same intrinsic matrices - torch.testing.assert_close(camera_warp.data.intrinsic_matrices[0], camera_usd.data.intrinsic_matrices[0]) + torch.testing.assert_close( + camera_warp.data.intrinsic_matrices.torch[0], camera_usd.data.intrinsic_matrices.torch[0] + ) # check the apertures torch.testing.assert_close( @@ -782,8 +967,8 @@ def test_output_equal_to_usd_camera_when_intrinsics_set(setup_simulation): # check image data torch.testing.assert_close( - camera_usd.data.output["distance_to_camera"], - camera_warp.data.output["distance_to_camera"], + camera_usd.data.output["distance_to_camera"].torch, + camera_warp.data.output["distance_to_camera"].torch, rtol=5e-3, atol=1e-4, ) @@ -804,8 +989,8 @@ def test_image_mesh_ids_identifies_hit_mesh(setup_simulation): sim.reset() camera.update(dt) - mesh_ids = camera.data.image_mesh_ids # shape (N, H, W, 1), dtype torch.int16 - assert mesh_ids is not None, "image_mesh_ids should not be None when update_mesh_ids=True" + assert camera.data.image_mesh_ids is not None, "image_mesh_ids should not be None when update_mesh_ids=True" + mesh_ids = camera.data.image_mesh_ids.torch # shape (N, H, W, 1), dtype torch.int16 assert mesh_ids.shape[-1] == 1 assert mesh_ids.dtype == torch.int16 @@ -813,11 +998,11 @@ def test_image_mesh_ids_identifies_hit_mesh(setup_simulation): # (the default), which leaves missed rays at the Warp-kernel fill value of inf. # Under "max" clipping, missed rays would be clamped to a finite max_distance, making # the inf comparison incorrect. - hit_mask = camera.data.output["distance_to_camera"][0, :, :, 0] < float("inf") + hit_mask = camera.data.output["distance_to_camera"].torch[0, :, :, 0] < float("inf") assert hit_mask.any(), "Expected at least some rays to hit the ground plane" - # All hits against the single registered mesh must carry mesh_id=0 (first mesh index). + # All hits against the single registered mesh must carry the ground mesh id. hit_mesh_ids = mesh_ids[0, :, :, 0][hit_mask] - assert torch.all(hit_mesh_ids == 0), ( - f"All hits against the single ground mesh must have mesh_id=0, got: {hit_mesh_ids.unique()}" + assert torch.all(hit_mesh_ids == MESH_ID_GROUND), ( + f"All hits against the single ground mesh must have mesh_id={MESH_ID_GROUND}, got: {hit_mesh_ids.unique()}" ) diff --git a/source/isaaclab/test/sensors/test_ray_caster_camera.py b/source/isaaclab/test/sensors/test_ray_caster_camera.py index d8ac47e95a6..db252f9456c 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_camera.py +++ b/source/isaaclab/test/sensors/test_ray_caster_camera.py @@ -211,45 +211,52 @@ def test_depth_clipping(setup_sim): camera_max.update(dt) # none clipping should contain inf values - assert torch.isinf(camera_none.data.output["distance_to_camera"]).any() - assert torch.isnan(camera_none.data.output["distance_to_image_plane"]).any() + assert torch.isinf(camera_none.data.output["distance_to_camera"].torch).any() + assert torch.isnan(camera_none.data.output["distance_to_image_plane"].torch).any() assert ( - camera_none.data.output["distance_to_camera"][~torch.isinf(camera_none.data.output["distance_to_camera"])].max() + camera_none.data.output["distance_to_camera"] + .torch[~torch.isinf(camera_none.data.output["distance_to_camera"].torch)] + .max() > camera_cfg_zero.max_distance ) assert ( - camera_none.data.output["distance_to_image_plane"][ - ~torch.isnan(camera_none.data.output["distance_to_image_plane"]) - ].max() + camera_none.data.output["distance_to_image_plane"] + .torch[~torch.isnan(camera_none.data.output["distance_to_image_plane"].torch)] + .max() > camera_cfg_zero.max_distance ) # zero clipping should result in zero values assert torch.all( - camera_zero.data.output["distance_to_camera"][torch.isinf(camera_none.data.output["distance_to_camera"])] == 0.0 + camera_zero.data.output["distance_to_camera"].torch[ + torch.isinf(camera_none.data.output["distance_to_camera"].torch) + ] + == 0.0 ) assert torch.all( - camera_zero.data.output["distance_to_image_plane"][ - torch.isnan(camera_none.data.output["distance_to_image_plane"]) + camera_zero.data.output["distance_to_image_plane"].torch[ + torch.isnan(camera_none.data.output["distance_to_image_plane"].torch) ] == 0.0 ) - assert camera_zero.data.output["distance_to_camera"].max() <= camera_cfg_zero.max_distance - assert camera_zero.data.output["distance_to_image_plane"].max() <= camera_cfg_zero.max_distance + assert camera_zero.data.output["distance_to_camera"].torch.max() <= camera_cfg_zero.max_distance + assert camera_zero.data.output["distance_to_image_plane"].torch.max() <= camera_cfg_zero.max_distance # max clipping should result in max values assert torch.all( - camera_max.data.output["distance_to_camera"][torch.isinf(camera_none.data.output["distance_to_camera"])] + camera_max.data.output["distance_to_camera"].torch[ + torch.isinf(camera_none.data.output["distance_to_camera"].torch) + ] == camera_cfg_zero.max_distance ) assert torch.all( - camera_max.data.output["distance_to_image_plane"][ - torch.isnan(camera_none.data.output["distance_to_image_plane"]) + camera_max.data.output["distance_to_image_plane"].torch[ + torch.isnan(camera_none.data.output["distance_to_image_plane"].torch) ] == camera_cfg_zero.max_distance ) - assert camera_max.data.output["distance_to_camera"].max() <= camera_cfg_zero.max_distance - assert camera_max.data.output["distance_to_image_plane"].max() <= camera_cfg_zero.max_distance + assert camera_max.data.output["distance_to_camera"].torch.max() <= camera_cfg_zero.max_distance + assert camera_max.data.output["distance_to_image_plane"].torch.max() <= camera_cfg_zero.max_distance @pytest.mark.isaacsim_ci @@ -297,9 +304,9 @@ def test_camera_init_offset(setup_sim): camera_ros.update(dt) # check that all transforms are set correctly - np.testing.assert_allclose(camera_ros.data.pos_w[0].cpu().numpy(), cam_cfg_offset_ros.offset.pos) - np.testing.assert_allclose(camera_opengl.data.pos_w[0].cpu().numpy(), cam_cfg_offset_opengl.offset.pos) - np.testing.assert_allclose(camera_world.data.pos_w[0].cpu().numpy(), cam_cfg_offset_world.offset.pos) + np.testing.assert_allclose(camera_ros.data.pos_w.torch[0].cpu().numpy(), cam_cfg_offset_ros.offset.pos) + np.testing.assert_allclose(camera_opengl.data.pos_w.torch[0].cpu().numpy(), cam_cfg_offset_opengl.offset.pos) + np.testing.assert_allclose(camera_world.data.pos_w.torch[0].cpu().numpy(), cam_cfg_offset_world.offset.pos) # check if transform correctly set in output np.testing.assert_allclose(camera_ros.data.pos_w[0].cpu().numpy(), cam_cfg_offset_ros.offset.pos, rtol=1e-5) @@ -316,7 +323,7 @@ def test_camera_init_intrinsic_matrix(setup_sim): camera_1 = RayCasterCamera(cfg=camera_cfg) # get intrinsic matrix sim.reset() - intrinsic_matrix = camera_1.data.intrinsic_matrices[0].cpu().flatten().tolist() + intrinsic_matrix = camera_1.data.intrinsic_matrices.torch[0].cpu().flatten().tolist() teardown(sim) # reinit the first camera sim, camera_cfg, dt = setup() @@ -350,13 +357,13 @@ def test_camera_init_intrinsic_matrix(setup_sim): # check image data torch.testing.assert_close( - camera_1.data.output["distance_to_image_plane"], - camera_2.data.output["distance_to_image_plane"], + camera_1.data.output["distance_to_image_plane"].torch, + camera_2.data.output["distance_to_image_plane"].torch, ) # check that both intrinsic matrices are the same torch.testing.assert_close( - camera_1.data.intrinsic_matrices[0], - camera_2.data.intrinsic_matrices[0], + camera_1.data.intrinsic_matrices.torch[0], + camera_2.data.intrinsic_matrices.torch[0], ) @@ -520,8 +527,8 @@ def test_output_equal_to_usdcamera(setup_sim): # check the intrinsic matrices torch.testing.assert_close( - camera_usd.data.intrinsic_matrices, - camera_warp.data.intrinsic_matrices, + camera_usd.data.intrinsic_matrices.torch, + camera_warp.data.intrinsic_matrices.torch, ) # check the apertures @@ -540,14 +547,14 @@ def test_output_equal_to_usdcamera(setup_sim): # check image data torch.testing.assert_close( - camera_usd.data.output["distance_to_image_plane"], - camera_warp.data.output["distance_to_image_plane"], + camera_usd.data.output["distance_to_image_plane"].torch, + camera_warp.data.output["distance_to_image_plane"].torch, rtol=1e-5, atol=1e-4, ) torch.testing.assert_close( - camera_usd.data.output["distance_to_camera"], - camera_warp.data.output["distance_to_camera"], + camera_usd.data.output["distance_to_camera"].torch, + camera_warp.data.output["distance_to_camera"].torch, atol=5e-5, rtol=5e-6, ) @@ -616,14 +623,14 @@ def test_output_equal_to_usdcamera_offset(setup_sim): # check image data torch.testing.assert_close( - camera_usd.data.output["distance_to_image_plane"], - camera_warp.data.output["distance_to_image_plane"], + camera_usd.data.output["distance_to_image_plane"].torch, + camera_warp.data.output["distance_to_image_plane"].torch, rtol=1e-3, atol=1e-5, ) torch.testing.assert_close( - camera_usd.data.output["distance_to_camera"], - camera_warp.data.output["distance_to_camera"], + camera_usd.data.output["distance_to_camera"].torch, + camera_warp.data.output["distance_to_camera"].torch, rtol=1e-3, atol=1e-5, ) @@ -710,14 +717,14 @@ def test_output_equal_to_usdcamera_prim_offset(setup_sim): # check image data torch.testing.assert_close( - camera_usd.data.output["distance_to_image_plane"], - camera_warp.data.output["distance_to_image_plane"], + camera_usd.data.output["distance_to_image_plane"].torch, + camera_warp.data.output["distance_to_image_plane"].torch, rtol=1e-3, atol=1e-5, ) torch.testing.assert_close( - camera_usd.data.output["distance_to_camera"], - camera_warp.data.output["distance_to_camera"], + camera_usd.data.output["distance_to_camera"].torch, + camera_warp.data.output["distance_to_camera"].torch, rtol=4e-6, atol=2e-5, ) @@ -800,15 +807,17 @@ def test_output_equal_to_usd_camera_intrinsics(setup_sim, focal_length): camera_warp.update(dt) # filter nan and inf from output - cam_warp_output = camera_warp.data.output["distance_to_image_plane"].clone() - cam_usd_output = camera_usd.data.output["distance_to_image_plane"].clone() + cam_warp_output = camera_warp.data.output["distance_to_image_plane"].torch.clone() + cam_usd_output = camera_usd.data.output["distance_to_image_plane"].torch.clone() cam_warp_output[torch.isnan(cam_warp_output)] = 0 cam_warp_output[torch.isinf(cam_warp_output)] = 0 cam_usd_output[torch.isnan(cam_usd_output)] = 0 cam_usd_output[torch.isinf(cam_usd_output)] = 0 # check that both have the same intrinsic matrices - torch.testing.assert_close(camera_warp.data.intrinsic_matrices[0], camera_usd.data.intrinsic_matrices[0]) + torch.testing.assert_close( + camera_warp.data.intrinsic_matrices.torch[0], camera_usd.data.intrinsic_matrices.torch[0] + ) # check the apertures torch.testing.assert_close( @@ -932,14 +941,16 @@ def test_output_equal_to_usd_camera_when_intrinsics_set(setup_sim, focal_length_ import matplotlib.pyplot as plt fig, axs = plt.subplots(1, 3, figsize=(15, 5)) - usd_plt = axs[0].imshow(camera_usd.data.output["distance_to_camera"][0].cpu().numpy()) + usd_plt = axs[0].imshow(camera_usd.data.output["distance_to_camera"].torch[0].cpu().numpy()) fig.colorbar(usd_plt, ax=axs[0]) axs[0].set_title("USD") - warp_plt = axs[1].imshow(camera_warp.data.output["distance_to_camera"][0].cpu().numpy()) + warp_plt = axs[1].imshow(camera_warp.data.output["distance_to_camera"].torch[0].cpu().numpy()) fig.colorbar(warp_plt, ax=axs[1]) axs[1].set_title("WARP") diff_plt = axs[2].imshow( - torch.abs(camera_usd.data.output["distance_to_camera"] - camera_warp.data.output["distance_to_camera"])[0] + torch.abs( + camera_usd.data.output["distance_to_camera"].torch - camera_warp.data.output["distance_to_camera"].torch + )[0] .cpu() .numpy() ) @@ -956,8 +967,8 @@ def test_output_equal_to_usd_camera_when_intrinsics_set(setup_sim, focal_length_ if focal_length != 0.193: # FIXME: 0.193 is not working on the IsaacSim/ UsdGeom side, add back once fixed torch.testing.assert_close( - camera_usd.data.output["distance_to_camera"], - camera_warp.data.output["distance_to_camera"], + camera_usd.data.output["distance_to_camera"].torch, + camera_warp.data.output["distance_to_camera"].torch, rtol=5e-3, atol=1e-4, ) @@ -1030,10 +1041,10 @@ def test_depth_clipping_d2ip_and_d2c_are_independent(setup_sim): cam_d2ip.update(dt) cam_d2c.update(dt) - d2ip_joint = cam_joint.data.output["distance_to_image_plane"] - d2c_joint = cam_joint.data.output["distance_to_camera"] - d2ip_solo = cam_d2ip.data.output["distance_to_image_plane"] - d2c_solo = cam_d2c.data.output["distance_to_camera"] + d2ip_joint = cam_joint.data.output["distance_to_image_plane"].torch + d2c_joint = cam_joint.data.output["distance_to_camera"].torch + d2ip_solo = cam_d2ip.data.output["distance_to_image_plane"].torch + d2c_solo = cam_d2c.data.output["distance_to_camera"].torch # Joint camera must match solo cameras (clipping one must not affect the other) torch.testing.assert_close(d2ip_joint, d2ip_solo, atol=1e-5, rtol=1e-5) @@ -1091,7 +1102,7 @@ def test_set_intrinsic_matrices_updates_output(setup_sim): for _ in range(3): sim.step() camera.update(dt) - output_before = camera.data.output["distance_to_camera"].clone() + output_before = camera.data.output["distance_to_camera"].torch.clone() # Change to a very different focal length (longer → tighter FOV → depth values differ at edges) new_matrix = torch.tensor( @@ -1103,7 +1114,7 @@ def test_set_intrinsic_matrices_updates_output(setup_sim): for _ in range(3): sim.step() camera.update(dt) - output_after = camera.data.output["distance_to_camera"].clone() + output_after = camera.data.output["distance_to_camera"].torch.clone() # Outputs must differ after intrinsics change (different ray angles → different depths) assert not torch.allclose(output_before, output_after, atol=1e-3), ( diff --git a/source/isaaclab/test/sensors/test_ray_caster_integration.py b/source/isaaclab/test/sensors/test_ray_caster_integration.py index 36b6d8ff469..91fa9c12fed 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_integration.py +++ b/source/isaaclab/test/sensors/test_ray_caster_integration.py @@ -10,10 +10,10 @@ These tests require Isaac Sim (AppLauncher). They cover the integration-level items from ``TODO_ray_caster_kernel_tests.md``: -- ``_get_view_transforms_wp`` ArticulationView and RigidBodyView paths +- ``_get_sensor_transforms_wp`` ArticulationView and RigidBodyView paths - ``MultiMeshRayCaster`` env_mask behavior - ``MultiMeshRayCasterCamera.set_intrinsic_matrices`` propagation -- ``_update_mesh_transforms`` non-identity orientation offset (known bug, xfail) +- ``_update_mesh_transforms`` non-identity orientation offset - Depth clipping ordering for ``MultiMeshRayCasterCamera`` """ @@ -22,6 +22,7 @@ simulation_app = AppLauncher(headless=True, enable_cameras=True).app import copy +from typing import Any, cast import numpy as np import pytest @@ -31,6 +32,7 @@ from pxr import UsdGeom, UsdPhysics import isaaclab.sim as sim_utils +from isaaclab.cloner.clone_plan import ClonePlan from isaaclab.sensors.ray_caster import ( MultiMeshRayCaster, MultiMeshRayCasterCamera, @@ -75,6 +77,14 @@ def _single_downward_ray_cfg(prim_path: str) -> RayCasterCfg: ) +def _spawn_cube_part(part_path: str, translation: tuple[float, float, float]) -> None: + """Create a small mesh-bearing part under an Xform target.""" + stage = sim_utils.get_current_stage() + sim_utils.create_prim(part_path, "Xform", translation=translation, stage=stage) + cube = cast(Any, UsdGeom.Cube.Define(stage, f"{part_path}/Mesh")) + cube.CreateSizeAttr().Set(0.35) + + @pytest.fixture def sim_ground(): sim = _make_sim_and_ground() @@ -84,7 +94,7 @@ def sim_ground(): # --------------------------------------------------------------------------- -# _get_view_transforms_wp: ArticulationView path +# _get_sensor_transforms_wp: ArticulationView path # --------------------------------------------------------------------------- @@ -95,7 +105,7 @@ def test_articulation_view_path(sim_ground): Verifies that sensor pos_w matches the prim's initial position and that the downward ray hits the ground plane. This exercises the ``ArticulationView.get_root_transforms()`` quaternion-convention path in - :meth:`_get_view_transforms_wp`. + :meth:`_get_sensor_transforms_wp`. """ sim = sim_ground expected_pos = (3.0, 4.0, 5.0) @@ -107,11 +117,11 @@ def test_articulation_view_path(sim_ground): UsdPhysics.RigidBodyAPI.Apply(prim) UsdPhysics.ArticulationRootAPI.Apply(prim) # Mass is needed for physics; collision is needed for PhysX to track the body. - mass_api = UsdPhysics.MassAPI.Apply(prim) + mass_api = cast(Any, UsdPhysics.MassAPI.Apply(prim)) mass_api.CreateMassAttr().Set(1.0) # Create a small collision cube so PhysX treats this as a real body. cube_path = f"{prim_path}/CollisionCube" - cube_geom = UsdGeom.Cube.Define(stage, cube_path) + cube_geom = cast(Any, UsdGeom.Cube.Define(stage, cube_path)) cube_geom.CreateSizeAttr().Set(0.1) UsdPhysics.CollisionAPI.Apply(stage.GetPrimAtPath(cube_path)) sim_utils.update_stage() @@ -133,7 +143,7 @@ def test_articulation_view_path(sim_ground): # --------------------------------------------------------------------------- -# _get_view_transforms_wp: RigidBodyView path +# _get_sensor_transforms_wp: RigidBodyView path # --------------------------------------------------------------------------- @@ -142,7 +152,7 @@ def test_rigid_body_view_path(sim_ground): """Mount a ray caster on a prim with RigidBodyAPI (no ArticulationRootAPI). Exercises the ``RigidBodyView.get_transforms()`` path in - :meth:`_get_view_transforms_wp`. + :meth:`_get_sensor_transforms_wp`. """ sim = sim_ground expected_pos = (1.0, 2.0, 6.0) @@ -152,10 +162,10 @@ def test_rigid_body_view_path(sim_ground): stage = sim_utils.get_current_stage() prim = stage.GetPrimAtPath(prim_path) UsdPhysics.RigidBodyAPI.Apply(prim) - mass_api = UsdPhysics.MassAPI.Apply(prim) + mass_api = cast(Any, UsdPhysics.MassAPI.Apply(prim)) mass_api.CreateMassAttr().Set(1.0) cube_path = f"{prim_path}/CollisionCube" - cube_geom = UsdGeom.Cube.Define(stage, cube_path) + cube_geom = cast(Any, UsdGeom.Cube.Define(stage, cube_path)) cube_geom.CreateSizeAttr().Set(0.1) UsdPhysics.CollisionAPI.Apply(stage.GetPrimAtPath(cube_path)) sim_utils.update_stage() @@ -226,7 +236,8 @@ def test_multi_mesh_camera_set_intrinsic_matrices(sim_ground_camera): for _ in range(3): sim.step() camera.update(_DT) - output_before = camera.data.output["distance_to_camera"].clone() + camera_output = cast(dict[str, Any], camera.data.output) + output_before = camera_output["distance_to_camera"].torch.clone() # Change to a very different intrinsic matrix (different FOV) new_matrix = torch.tensor( @@ -238,7 +249,8 @@ def test_multi_mesh_camera_set_intrinsic_matrices(sim_ground_camera): for _ in range(3): sim.step() camera.update(_DT) - output_after = camera.data.output["distance_to_camera"].clone() + camera_output = cast(dict[str, Any], camera.data.output) + output_after = camera_output["distance_to_camera"].torch.clone() assert not torch.allclose(output_before, output_after, atol=1e-3), ( "MultiMeshRayCasterCamera: depth output must change after set_intrinsic_matrices; " @@ -256,9 +268,8 @@ def test_multi_mesh_camera_set_intrinsic_matrices(sim_ground_camera): def test_multi_mesh_camera_d2ip_and_d2c_independent(sim_ground_camera): """Requesting both d2ip and d2c simultaneously must produce correct independent results. - The ``distance_to_image_plane`` computation reads ``_ray_distance`` before - ``distance_to_camera`` clips it in-place. This test verifies the two data - types do not interfere with each other. + The two outputs are clipped independently from the same raw ray-distance buffer. + This test verifies the data types do not interfere with each other. """ sim, base_cfg = sim_ground_camera @@ -292,10 +303,13 @@ def test_multi_mesh_camera_d2ip_and_d2c_independent(sim_ground_camera): cam_d2ip.update(_DT) cam_d2c.update(_DT) - d2ip_joint = cam_joint.data.output["distance_to_image_plane"] - d2c_joint = cam_joint.data.output["distance_to_camera"] - d2ip_solo = cam_d2ip.data.output["distance_to_image_plane"] - d2c_solo = cam_d2c.data.output["distance_to_camera"] + joint_output = cast(dict[str, Any], cam_joint.data.output) + d2ip_output = cast(dict[str, Any], cam_d2ip.data.output) + d2c_output = cast(dict[str, Any], cam_d2c.data.output) + d2ip_joint = joint_output["distance_to_image_plane"].torch + d2c_joint = joint_output["distance_to_camera"].torch + d2ip_solo = d2ip_output["distance_to_image_plane"].torch + d2c_solo = d2c_output["distance_to_camera"].torch # Joint camera must match solo cameras (clipping one must not corrupt the other) torch.testing.assert_close(d2ip_joint, d2ip_solo, atol=1e-5, rtol=1e-5) @@ -307,6 +321,85 @@ def test_multi_mesh_camera_d2ip_and_d2c_independent(sim_ground_camera): # --------------------------------------------------------------------------- +@pytest.mark.isaacsim_ci +def test_multi_mesh_uses_clone_plan_geometry_and_backend_object_pose(sim_ground): + """ClonePlan supplies source geometry while PhysX supplies per-env object poses.""" + sim = sim_ground + num_envs = 3 + stage = sim_utils.get_current_stage() + + def _create_object_body(path: str) -> None: + sim_utils.create_prim(path, "Xform", stage=stage) + body_prim = cast(Any, stage.GetPrimAtPath(path)) + UsdPhysics.RigidBodyAPI.Apply(body_prim) + mass_api = cast(Any, UsdPhysics.MassAPI.Apply(body_prim)) + mass_api.CreateMassAttr().Set(1.0) + body_prim.GetAttribute("physics:kinematicEnabled").Set(True) + collision = cast(Any, UsdGeom.Cube.Define(stage, f"{path}/Collision")) + collision.CreateSizeAttr().Set(0.1) + UsdPhysics.CollisionAPI.Apply(stage.GetPrimAtPath(f"{path}/Collision")) + + sim_utils.create_prim("/World/envs", "Xform", stage=stage) + for env_id in range(num_envs): + sim_utils.create_prim(f"/World/envs/env_{env_id}", "Xform", translation=(3.0 * env_id, 0.0, 0.0), stage=stage) + sim_utils.create_prim(f"/World/envs/env_{env_id}/Sensor", "Xform", translation=(0.0, 0.0, 3.0), stage=stage) + _create_object_body(f"/World/envs/env_{env_id}/Object") + + # Representative source assets live in the first concrete cloned instances, + # matching the scene convention used by the cloner. Env 2 has an object + # body for backend pose tracking, but intentionally has no destination mesh. + _spawn_cube_part("/World/envs/env_0/Object/part_0", (0.0, 0.0, 0.0)) + _spawn_cube_part("/World/envs/env_1/Object/part_0", (0.0, 0.0, 0.0)) + + # This test intentionally does not author /env_2/Object/part_0. ClonePlan + # selects source geometry; the object body view supplies env_2's live pose. + sim.set_clone_plan( + ClonePlan( + sources=("/World/envs/env_0/Object", "/World/envs/env_1/Object"), + destinations=("/World/envs/env_{}/Object", "/World/envs/env_{}/Object"), + clone_mask=torch.tensor([[True, False, True], [False, True, False]], dtype=torch.bool, device=sim.device), + ) + ) + sim_utils.update_stage() + + cfg = MultiMeshRayCasterCfg( + prim_path="/World/envs/env_.*/Sensor", + mesh_prim_paths=[ + MultiMeshRayCasterCfg.RaycastTargetCfg( + prim_expr="/World/envs/env_.*/Object/part_.*", + track_mesh_transforms=True, + ), + ], + update_period=0, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(0.0, 0.0, 0.0, 1.0)), + debug_vis=False, + pattern_cfg=patterns.GridPatternCfg(resolution=0.5, size=(1.0, 0.0), direction=(0.0, 0.0, -1.0)), + ray_alignment="world", + ) + sensor = MultiMeshRayCaster(cfg) + sim.reset() + sensor.update(_DT, force_recompute=True) + + env0_object = stage.GetPrimAtPath("/World/envs/env_0/Object") + env1_object = stage.GetPrimAtPath("/World/envs/env_1/Object") + assert env0_object is not None and env0_object.IsValid() + assert env1_object is not None and env1_object.IsValid() + env2_part = stage.GetPrimAtPath("/World/envs/env_2/Object/part_0") + assert env2_part is None or not env2_part.IsValid() + + # Geometry is selected from ClonePlan rows, but poses come from the batched object view. + mesh_ids = wp.to_torch(sensor._mesh_ids_wp).cpu() + mesh_positions = wp.to_torch(sensor._mesh_positions_w).cpu() + assert sensor._mesh_ids_wp.shape == (num_envs, 1) + assert mesh_ids[2, 0] == mesh_ids[0, 0] + torch.testing.assert_close(mesh_positions[:, 0, 0], torch.tensor([0.0, 3.0, 6.0]), atol=0.15, rtol=0.0) + + hits = sensor.data.ray_hits_w.torch + assert torch.isfinite(hits[0]).any(), "env_0 should hit the env_0 source geometry" + assert torch.isfinite(hits[1]).any(), "env_1 should hit the env_1 source geometry" + assert torch.isfinite(hits[2]).any(), "env_2 should hit env_0 geometry at env_2's backend object pose" + + @pytest.mark.isaacsim_ci def test_multi_mesh_env_mask_preserves_masked_buffers(sim_ground): """Masked environments must retain their pre-update buffer values. @@ -357,17 +450,18 @@ def test_multi_mesh_env_mask_preserves_masked_buffers(sim_ground): @pytest.mark.isaacsim_ci def test_update_mesh_transforms_non_identity_offset(sim_ground): - """Tracked mesh position must account for body orientation when applying offset. + """Tracked mesh geometry must account for body orientation when applying offset. Setup: a kinematic rigid body at (0, 0, 2) rotated 90 deg around Z, with a child mesh offset by (1, 0, 0) in the body's local frame. - Correct world position of mesh = body_pos + rotate(body_ori, local_offset) + Correct world position of the baked geometry = body_pos + rotate(body_ori, local_offset) = (0, 0, 2) + rotate(90degZ, (1, 0, 0)) = (0, 0, 2) + (0, 1, 0) = (0, 1, 2) - Naive subtraction (the old bug) would give: body_pos - offset = (-1, 0, 2). + The mesh pose table stores the owner-body pose; the offset is baked into the + Warp mesh vertices so backend body poses remain the single runtime pose source. """ sim = sim_ground @@ -381,29 +475,29 @@ def test_update_mesh_transforms_non_identity_offset(sim_ground): body_path = "/World/DynamicBody" sim_utils.create_prim(body_path, "Xform", translation=(0.0, 0.0, 2.0), orientation=yaw90_xyzw) stage = sim_utils.get_current_stage() - body_prim = stage.GetPrimAtPath(body_path) + body_prim = cast(Any, stage.GetPrimAtPath(body_path)) UsdPhysics.RigidBodyAPI.Apply(body_prim) - mass_api = UsdPhysics.MassAPI.Apply(body_prim) + mass_api = cast(Any, UsdPhysics.MassAPI.Apply(body_prim)) mass_api.CreateMassAttr().Set(1.0) body_prim.GetAttribute("physics:kinematicEnabled").Set(True) # Create a child Xform offset by (1, 0, 0) in the body's local frame, - # then place mesh geometry under it. The Xform translation is the offset - # that _obtain_trackable_prim_view / resolve_prim_pose will discover. + # then place mesh geometry under it. The tracked target view resolves this + # local offset and updates from the backend body pose. child_mesh_path = f"{body_path}/OffsetMesh" sim_utils.create_prim(child_mesh_path, "Xform", translation=(1.0, 0.0, 0.0)) mesh_data = make_plane(size=(2, 2), height=0.0, center_zero=True) create_prim_from_mesh(f"{child_mesh_path}/Plane", mesh_data) # Add collision so PhysX tracks the body col_path = f"{body_path}/CollisionCube" - cube_geom = UsdGeom.Cube.Define(stage, col_path) + cube_geom = cast(Any, UsdGeom.Cube.Define(stage, col_path)) cube_geom.CreateSizeAttr().Set(0.1) UsdPhysics.CollisionAPI.Apply(stage.GetPrimAtPath(col_path)) sim_utils.update_stage() # Create a sensor prim to mount the MultiMeshRayCaster on sensor_path = "/World/SensorMount" - sim_utils.create_prim(sensor_path, "Xform", translation=(0.0, 0.0, 5.0)) + sim_utils.create_prim(sensor_path, "Xform", translation=(0.0, 1.0, 5.0)) # Configure MultiMeshRayCaster to track the child mesh cfg = MultiMeshRayCasterCfg( @@ -422,18 +516,19 @@ def test_update_mesh_transforms_non_identity_offset(sim_ground): ) sensor = MultiMeshRayCaster(cfg) sim.reset() - sensor.update(_DT) + sensor.update(_DT, force_recompute=True) - # Verify mesh position: body at (0,0,2) rotated 90deg Z, child offset (1,0,0) local - # Expected: (0, 0, 2) + rotate(90degZ, (1,0,0)) = (0, 0, 2) + (0, 1, 0) = (0, 1, 2) - mesh_pos = sensor._mesh_positions_w_torch.clone() + # The runtime pose buffer stores the body pose; the child offset is in the mesh vertices. + mesh_pos = wp.to_torch(sensor._mesh_positions_w).clone() np.testing.assert_allclose( mesh_pos[0, 0].cpu().numpy(), + [0.0, 0.0, 2.0], + atol=0.15, + err_msg="Tracked mesh pose should be the owner body pose.", + ) + np.testing.assert_allclose( + sensor.data.ray_hits_w.torch[0, 0].cpu().numpy(), [0.0, 1.0, 2.0], atol=0.15, - err_msg=( - "Mesh position should be (0, 1, 2) via proper frame decomposition: " - "body_pos + rotate(body_ori, local_offset). " - "If this fails, the offset is not being rotated by the body orientation." - ), + err_msg="Ray hit should include the child mesh offset baked into the Warp geometry.", ) diff --git a/source/isaaclab/test/sensors/test_ray_caster_kernels.py b/source/isaaclab/test/sensors/test_ray_caster_kernels.py index cc57e4f1eec..b0853c2a3bc 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_kernels.py +++ b/source/isaaclab/test/sensors/test_ray_caster_kernels.py @@ -54,9 +54,12 @@ _warp_mod = importlib.util.module_from_spec(_warp_spec) _warp_spec.loader.exec_module(_warp_mod) -compute_distance_to_image_plane_masked_kernel = _sensor_mod.compute_distance_to_image_plane_masked_kernel -apply_depth_clipping_masked_kernel = _sensor_mod.apply_depth_clipping_masked_kernel +compute_distance_to_image_plane_to_image_masked_kernel = ( + _sensor_mod.compute_distance_to_image_plane_to_image_masked_kernel +) apply_z_drift_kernel = _sensor_mod.apply_z_drift_kernel +copy_float2d_to_image1_depth_clipped_masked_kernel = _sensor_mod.copy_float2d_to_image1_depth_clipped_masked_kernel +fill_ray_hits_distance_inf_kernel = _sensor_mod.fill_ray_hits_distance_inf_kernel quat_yaw_only = _sensor_mod.quat_yaw_only raycast_dynamic_meshes_kernel = _warp_mod.raycast_dynamic_meshes_kernel @@ -120,6 +123,66 @@ def _to_numpy(a: wp.array) -> np.ndarray: return a.numpy() +# --------------------------------------------------------------------------- +# Tests: fill_ray_hits_distance_inf_kernel +# --------------------------------------------------------------------------- + + +class TestFillRayHitsDistanceInfKernel: + """Tests for :func:`fill_ray_hits_distance_inf_kernel`.""" + + def test_active_envs_are_filled_and_masked_envs_are_preserved(self): + """Active environments are filled with infinity while masked environments retain prior values.""" + env_mask = wp.array(np.array([False, True], dtype=np.bool_), dtype=wp.bool, device=DEVICE) + ray_hits = wp.array( + np.array( + [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + dtype=np.float32, + ), + dtype=wp.vec3f, + device=DEVICE, + ) + ray_distance = wp.array( + np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), + dtype=wp.float32, + device=DEVICE, + ) + ray_normals = wp.array( + np.array( + [ + [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0]], + [[1.0, 0.0, 0.0], [1.0, 1.0, 0.0]], + ], + dtype=np.float32, + ), + dtype=wp.vec3f, + device=DEVICE, + ) + + wp.launch( + fill_ray_hits_distance_inf_kernel, + dim=(2, 2), + inputs=[env_mask, True], + outputs=[ray_hits, ray_distance, ray_normals], + device=DEVICE, + ) + wp.synchronize_device(DEVICE) + + hits_np = _to_numpy(ray_hits) + distance_np = _to_numpy(ray_distance) + normals_np = _to_numpy(ray_normals) + + np.testing.assert_allclose(hits_np[0], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], atol=ATOL) + np.testing.assert_allclose(distance_np[0], [1.0, 2.0], atol=ATOL) + np.testing.assert_allclose(normals_np[0], [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0]], atol=ATOL) + assert np.isinf(hits_np[1]).all() + assert np.isinf(distance_np[1]).all() + assert np.isinf(normals_np[1]).all() + + # --------------------------------------------------------------------------- # Tests: raycast_dynamic_meshes_kernel # --------------------------------------------------------------------------- @@ -320,168 +383,101 @@ def test_equidistant_meshes(self): assert out["mesh_id"][0, 0] in (0, 1) -# --------------------------------------------------------------------------- -# Tests: compute_distance_to_image_plane_masked_kernel -# --------------------------------------------------------------------------- - - -class TestComputeDistanceToImagePlaneMaskedKernel: - """Tests for :func:`compute_distance_to_image_plane_masked_kernel`.""" - - @staticmethod - def _launch( - quat_xyzw: list[float], - ray_distance: list[list[float]], - ray_dirs: list[list[list[float]]], - env_mask: list[bool] | None = None, - ) -> np.ndarray: - """Launch kernel and return distance_to_image_plane as numpy.""" - num_envs = len(ray_distance) - num_rays = len(ray_distance[0]) - if env_mask is None: - env_mask = [True] * num_envs - - mask_wp = wp.array(np.array(env_mask, dtype=np.bool_), dtype=wp.bool, device=DEVICE) - quat_np = np.array([quat_xyzw] * num_envs, dtype=np.float32) - quat_wp = wp.array(quat_np, dtype=wp.quatf, device=DEVICE) - ray_dist_wp = wp.array(np.array(ray_distance, dtype=np.float32), dtype=wp.float32, device=DEVICE) - dirs_wp = wp.array(np.array(ray_dirs, dtype=np.float32), dtype=wp.vec3f, device=DEVICE) - out_wp = wp.zeros((num_envs, num_rays), dtype=wp.float32, device=DEVICE) +class TestComputeDistanceToImagePlaneToImageMaskedKernel: + """Tests for :func:`compute_distance_to_image_plane_to_image_masked_kernel`.""" + + def test_compute_clip_and_copy_to_image(self): + """Distance-to-image-plane is computed, clipped, reshaped, and masked in one kernel.""" + env_mask = wp.array(np.array([False, True], dtype=np.bool_), dtype=wp.bool, device=DEVICE) + quat_w = wp.array(np.array([[0, 0, 0, 1], [0, 0, 0, 1]], dtype=np.float32), dtype=wp.quatf, device=DEVICE) + ray_distance = wp.array(np.array([[1.0, 2.0], [3.0, 7.0]], dtype=np.float32), dtype=wp.float32, device=DEVICE) + ray_directions_w = wp.array( + np.array( + [ + [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + ], + dtype=np.float32, + ), + dtype=wp.vec3f, + device=DEVICE, + ) + dst = wp.array(np.full((2, 1, 2, 1), -1.0, dtype=np.float32), dtype=wp.float32, device=DEVICE) wp.launch( - compute_distance_to_image_plane_masked_kernel, - dim=(num_envs, num_rays), - inputs=[mask_wp, quat_wp, ray_dist_wp, dirs_wp], - outputs=[out_wp], + compute_distance_to_image_plane_to_image_masked_kernel, + dim=(2, 2), + inputs=[env_mask, quat_w, ray_distance, ray_directions_w, 2, True, 5.0, 0.0], + outputs=[dst], device=DEVICE, ) wp.synchronize_device(DEVICE) - return _to_numpy(out_wp) - def test_known_camera_orientation(self): - """Identity camera, ray along +X at distance 5 -- d2ip equals 5.""" - result = self._launch( - quat_xyzw=[0, 0, 0, 1], - ray_distance=[[5.0]], - ray_dirs=[[[1, 0, 0]]], - ) - assert result[0, 0] == pytest.approx(5.0, abs=ATOL) + dst_np = _to_numpy(dst) + np.testing.assert_allclose(dst_np[0, :, :, 0], [[-1.0, -1.0]], atol=ATOL) + np.testing.assert_allclose(dst_np[1, :, :, 0], [[3.0, 0.0]], atol=ATOL) - def test_off_axis_camera(self): - """Camera pitched 45 deg around Y, ray going world -Z. + def test_off_axis_camera_without_clipping(self): + """Camera pitched 45 deg around Y, ray going world -Z.""" + env_mask = wp.array(np.array([True], dtype=np.bool_), dtype=wp.bool, device=DEVICE) + quat_w = wp.array( + np.array([_euler_to_quat_xyzw(0, math.pi / 4, 0)], dtype=np.float32), dtype=wp.quatf, device=DEVICE + ) + ray_distance = wp.array(np.array([[10.0]], dtype=np.float32), dtype=wp.float32, device=DEVICE) + ray_directions_w = wp.array(np.array([[[0.0, 0.0, -1.0]]], dtype=np.float32), dtype=wp.vec3f, device=DEVICE) + dst = wp.array(np.full((1, 1, 1, 1), -1.0, dtype=np.float32), dtype=wp.float32, device=DEVICE) - Camera forward (+X_cam) in world = (cos45, 0, -sin45). - Displacement = 10 * (0, 0, -1) = (0, 0, -10). - Projection onto camera forward = dot((0,0,-10), (cos45,0,-sin45)) - = 10 * sin(45 deg). - """ - pitch45 = list(_euler_to_quat_xyzw(0, math.pi / 4, 0)) - result = self._launch( - quat_xyzw=pitch45, - ray_distance=[[10.0]], - ray_dirs=[[[0, 0, -1]]], + wp.launch( + compute_distance_to_image_plane_to_image_masked_kernel, + dim=(1, 1), + inputs=[env_mask, quat_w, ray_distance, ray_directions_w, 1, False, 5.0, 0.0], + outputs=[dst], + device=DEVICE, ) + wp.synchronize_device(DEVICE) + expected = 10.0 * math.sin(math.pi / 4) - assert result[0, 0] == pytest.approx(expected, abs=ATOL) + assert _to_numpy(dst)[0, 0, 0, 0] == pytest.approx(expected, abs=ATOL) - def test_inf_distance(self): - """Inf distance produces NaN through the projection (inf * 0 = NaN). + def test_inf_distance_is_clipped(self): + """Inf distance produces NaN through projection and is replaced by fill value.""" + env_mask = wp.array(np.array([True], dtype=np.bool_), dtype=wp.bool, device=DEVICE) + quat_w = wp.array(np.array([[0, 0, 0, 1]], dtype=np.float32), dtype=wp.quatf, device=DEVICE) + ray_distance = wp.array(np.array([[float("inf")]], dtype=np.float32), dtype=wp.float32, device=DEVICE) + ray_directions_w = wp.array(np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32), dtype=wp.vec3f, device=DEVICE) + dst = wp.array(np.full((1, 1, 1, 1), -1.0, dtype=np.float32), dtype=wp.float32, device=DEVICE) - When a ray misses, ray_distance is inf. Multiplying inf by zero-valued - ray-direction components yields NaN (IEEE 754), which propagates through - the quaternion rotation. The downstream - :func:`apply_depth_clipping_masked_kernel` handles NaN correctly via - ``wp.isnan()``, so the overall pipeline is sound. - """ - result = self._launch( - quat_xyzw=[0, 0, 0, 1], - ray_distance=[[float("inf")]], - ray_dirs=[[[1, 0, 0]]], + wp.launch( + compute_distance_to_image_plane_to_image_masked_kernel, + dim=(1, 1), + inputs=[env_mask, quat_w, ray_distance, ray_directions_w, 1, True, 5.0, 0.0], + outputs=[dst], + device=DEVICE, ) - assert np.isnan(result[0, 0]), f"Expected NaN from inf*0 contamination, got {result[0, 0]}" - - -# --------------------------------------------------------------------------- -# Tests: apply_depth_clipping_masked_kernel -# --------------------------------------------------------------------------- + wp.synchronize_device(DEVICE) + assert _to_numpy(dst)[0, 0, 0, 0] == pytest.approx(0.0, abs=ATOL) -class TestApplyDepthClippingMaskedKernel: - """Tests for :func:`apply_depth_clipping_masked_kernel`.""" - @staticmethod - def _launch( - depth_values: list[list[float]], - max_dist: float, - fill_val: float, - env_mask: list[bool] | None = None, - ) -> np.ndarray: - """Launch kernel and return clipped depth as numpy.""" - num_envs = len(depth_values) - num_rays = len(depth_values[0]) - if env_mask is None: - env_mask = [True] * num_envs +class TestCopyFloat2dToImage1DepthClippedMaskedKernel: + """Tests for :func:`copy_float2d_to_image1_depth_clipped_masked_kernel`.""" - mask_wp = wp.array(np.array(env_mask, dtype=np.bool_), dtype=wp.bool, device=DEVICE) - depth_wp = wp.array(np.array(depth_values, dtype=np.float32), dtype=wp.float32, device=DEVICE) + def test_clip_and_copy_to_image(self): + """Flat distances are optionally clipped while copying to image layout.""" + env_mask = wp.array(np.array([True], dtype=np.bool_), dtype=wp.bool, device=DEVICE) + src = wp.array(np.array([[1.0, 7.0, np.nan]], dtype=np.float32), dtype=wp.float32, device=DEVICE) + dst = wp.array(np.full((1, 1, 3, 1), -1.0, dtype=np.float32), dtype=wp.float32, device=DEVICE) wp.launch( - apply_depth_clipping_masked_kernel, - dim=(num_envs, num_rays), - inputs=[mask_wp, max_dist, fill_val], - outputs=[depth_wp], + copy_float2d_to_image1_depth_clipped_masked_kernel, + dim=(1, 3), + inputs=[env_mask, src, 3, True, 5.0, 0.0], + outputs=[dst], device=DEVICE, ) wp.synchronize_device(DEVICE) - return _to_numpy(depth_wp) - - def test_boundary_at_max_dist(self): - """Value at exactly max_dist is preserved (not clipped).""" - result = self._launch([[10.0]], max_dist=10.0, fill_val=0.0) - assert result[0, 0] == pytest.approx(10.0, abs=ATOL) - - def test_above_max_dist(self): - """Value above max_dist is replaced with fill_val.""" - result = self._launch([[10.001]], max_dist=10.0, fill_val=0.0) - assert result[0, 0] == pytest.approx(0.0, abs=ATOL) - - def test_nan_value(self): - """NaN value is replaced with fill_val.""" - result = self._launch([[float("nan")]], max_dist=10.0, fill_val=0.0) - assert result[0, 0] == pytest.approx(0.0, abs=ATOL) - - def test_inf_value(self): - """Inf is clipped (inf > max_dist is true).""" - result = self._launch([[float("inf")]], max_dist=10.0, fill_val=0.0) - assert result[0, 0] == pytest.approx(0.0, abs=ATOL) - - def test_negative_depth(self): - """Negative depth passes through unclipped (valid for distance-to-image-plane).""" - result = self._launch([[-3.5]], max_dist=10.0, fill_val=0.0) - assert result[0, 0] == pytest.approx(-3.5, abs=ATOL) - - def test_env_mask(self): - """Masked env retains original value -- clipping is not applied.""" - result = self._launch( - depth_values=[[15.0], [15.0]], - max_dist=10.0, - fill_val=0.0, - env_mask=[False, True], - ) - # Env 0 (masked): unchanged - assert result[0, 0] == pytest.approx(15.0, abs=ATOL) - # Env 1 (active): clipped - assert result[1, 0] == pytest.approx(0.0, abs=ATOL) - - def test_fill_val_zero_vs_max(self): - """fill_val=0.0 and fill_val=max_dist produce correct replacements.""" - max_dist = 10.0 - - result_zero = self._launch([[15.0]], max_dist=max_dist, fill_val=0.0) - assert result_zero[0, 0] == pytest.approx(0.0, abs=ATOL) - result_max = self._launch([[15.0]], max_dist=max_dist, fill_val=max_dist) - assert result_max[0, 0] == pytest.approx(max_dist, abs=ATOL) + np.testing.assert_allclose(_to_numpy(dst)[0, :, :, 0], [[1.0, 0.0, 0.0]], atol=ATOL) # --------------------------------------------------------------------------- diff --git a/source/isaaclab/test/sensors/test_ray_caster_sensor.py b/source/isaaclab/test/sensors/test_ray_caster_sensor.py index 3326c524f95..a3571fc222d 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_sensor.py +++ b/source/isaaclab/test/sensors/test_ray_caster_sensor.py @@ -7,6 +7,8 @@ """Tests for RayCaster sensor behavior: alignment modes and reset.""" +from typing import Literal + from isaaclab.app import AppLauncher simulation_app = AppLauncher(headless=True).app @@ -40,7 +42,7 @@ def _make_sim_and_ground(): return sim -def _ray_caster_cfg(prim_path: str, alignment: str) -> RayCasterCfg: +def _ray_caster_cfg(prim_path: str, alignment: Literal["base", "yaw", "world"]) -> RayCasterCfg: """Single downward ray, no offset from prim.""" return RayCasterCfg( prim_path=prim_path, @@ -257,6 +259,7 @@ def test_ray_caster_reset_resamples_drift(sim_ground): # reset() resamples drift; values should remain within the configured range # Call reset() multiple times until we get a different sample (probability of same is near zero # for continuous uniform distribution, but we retry to avoid flakiness). + drift_after: torch.Tensor = drift_before.clone() for _ in range(5): sensor.reset() drift_after = sensor.drift.clone() @@ -269,3 +272,49 @@ 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" ) + + +@pytest.mark.isaacsim_ci +def test_ray_caster_tracks_physics_body_parent_motion(sim_ground): + """RayCaster pose must follow its physics-body parent after simulation steps.""" + from pxr import UsdGeom, UsdPhysics # noqa: PLC0415 + + sim = sim_ground + dt = 0.01 + parent_path = "/World/PhysicsParent" + + stage = sim_utils.get_current_stage() + sim_utils.create_prim(parent_path, "Xform", translation=(0.0, 0.0, 5.0), stage=stage) + parent_prim = stage.GetPrimAtPath(parent_path) + UsdPhysics.RigidBodyAPI.Apply(parent_prim) + UsdPhysics.ArticulationRootAPI.Apply(parent_prim) + mass_api = UsdPhysics.MassAPI.Apply(parent_prim) + if mass_api is None: + raise RuntimeError(f"Failed to apply MassAPI to {parent_path}.") + mass_api.CreateMassAttr().Set(1.0) + + cube_path = f"{parent_path}/CollisionCube" + cube = UsdGeom.Cube.Define(stage, cube_path) + if cube is None: + raise RuntimeError(f"Failed to create collision cube at {cube_path}.") + cube.CreateSizeAttr().Set(0.1) + UsdPhysics.CollisionAPI.Apply(stage.GetPrimAtPath(cube_path)) + sim_utils.update_stage() + + sensor = RayCaster(_ray_caster_cfg(parent_path, "world")) + sim.reset() + sensor.update(dt, force_recompute=True) + pos_before = sensor.data.pos_w.torch[0].clone() + + for _ in range(100): + sim.step(render=False) + sensor.update(dt) + + sensor.update(dt, force_recompute=True) + pos_after = sensor.data.pos_w.torch[0] + drift_z = (pos_before[2] - pos_after[2]).item() + + assert drift_z > 0.5, ( + f"RayCaster pose did not follow its physics body parent. " + f"z before={pos_before[2].item():.4f} z after={pos_after[2].item():.4f} drift={drift_z:.4f}m." + ) diff --git a/source/isaaclab_newton/changelog.d/octi-raycaster-backend-split.minor.rst b/source/isaaclab_newton/changelog.d/octi-raycaster-backend-split.minor.rst new file mode 100644 index 00000000000..6dedaf9cbd7 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/octi-raycaster-backend-split.minor.rst @@ -0,0 +1,21 @@ +Added +^^^^^ + +* Added Newton backend for :class:`~isaaclab.sensors.ray_caster.RayCaster` / + :class:`~isaaclab.sensors.ray_caster.RayCasterCamera` / + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCaster` / + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCasterCamera`. Site-based, + matching :class:`~isaaclab_newton.sensors.pva.Pva` and + :class:`~isaaclab_newton.sensors.frame_transformer.FrameTransformer`: + registers body-attached sites via + :meth:`~isaaclab_newton.physics.NewtonManager.cl_register_site` for both + the sensor frame and any tracked target meshes, and reads per-step + transforms off :class:`~newton.sensors.SensorFrameTransform` against a + world-origin reference. Static parents/targets bypass the site + machinery and serve cached per-env ``wp.transformf`` arrays. + +Changed +^^^^^^^ + +* Changed Newton tracked target mesh updates to copy site poses directly into + Warp mesh pose tables instead of staging through torch views. diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/sensors/__init__.pyi index e536b281952..3cc07b029af 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/sensors/__init__.pyi @@ -15,6 +15,10 @@ __all__ = [ "JointWrenchSensorData", "Pva", "PvaData", + "MultiMeshRayCaster", + "MultiMeshRayCasterCamera", + "RayCaster", + "RayCasterCamera", ] from .contact_sensor import ContactSensor, ContactSensorData, ContactSensorCfg @@ -22,3 +26,4 @@ from .frame_transformer import FrameTransformer, FrameTransformerData from .imu import Imu, ImuData from .joint_wrench import JointWrenchSensor, JointWrenchSensorData from .pva import Pva, PvaData +from .ray_caster import MultiMeshRayCaster, MultiMeshRayCasterCamera, RayCaster, RayCasterCamera diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/__init__.py b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/__init__.py new file mode 100644 index 00000000000..52c309ac938 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""Sub-module for Newton ray-caster sensor.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/__init__.pyi new file mode 100644 index 00000000000..7562409a7e9 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/__init__.pyi @@ -0,0 +1,16 @@ +# 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 + +__all__ = [ + "MultiMeshRayCaster", + "MultiMeshRayCasterCamera", + "RayCaster", + "RayCasterCamera", +] + +from .multi_mesh_ray_caster import MultiMeshRayCaster +from .multi_mesh_ray_caster_camera import MultiMeshRayCasterCamera +from .ray_caster import RayCaster +from .ray_caster_camera import RayCasterCamera diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/multi_mesh_ray_caster.py b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/multi_mesh_ray_caster.py new file mode 100644 index 00000000000..74d18c56c43 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/multi_mesh_ray_caster.py @@ -0,0 +1,14 @@ +# 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 + +from isaaclab.sensors.ray_caster.base_multi_mesh_ray_caster import BaseMultiMeshRayCaster + +from .ray_caster import _NewtonRayCasterMixin + + +class MultiMeshRayCaster(_NewtonRayCasterMixin, BaseMultiMeshRayCaster): + """Newton MultiMeshRayCaster implementation.""" diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/multi_mesh_ray_caster_camera.py b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/multi_mesh_ray_caster_camera.py new file mode 100644 index 00000000000..3be4ca93fd6 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/multi_mesh_ray_caster_camera.py @@ -0,0 +1,14 @@ +# 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 + +from isaaclab.sensors.ray_caster.base_multi_mesh_ray_caster_camera import BaseMultiMeshRayCasterCamera + +from .ray_caster import _NewtonRayCasterMixin + + +class MultiMeshRayCasterCamera(_NewtonRayCasterMixin, BaseMultiMeshRayCasterCamera): + """Newton MultiMeshRayCasterCamera implementation.""" diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster.py b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster.py new file mode 100644 index 00000000000..e10076d884c --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster.py @@ -0,0 +1,312 @@ +# 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 + +# pyright: reportInvalidTypeForm=none, reportPrivateUsage=none +import re +from typing import Any + +import numpy as np +import warp as wp + +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +from isaaclab.sensors.ray_caster.base_ray_caster import BaseRayCaster +from isaaclab.sensors.ray_caster.kernels import ( + ALIGNMENT_BASE, + copy_mesh_poses_to_table_kernel, + update_ray_caster_kernel, +) +from isaaclab.utils.warp import ProxyArray + +from isaaclab_newton.physics import NewtonManager + + +@wp.kernel +def _newton_site_world_poses_kernel( + site_indices: wp.array(dtype=wp.int32), + shape_body: wp.array(dtype=wp.int32), + shape_transform: wp.array(dtype=wp.transform), + body_q: wp.array(dtype=wp.transform), + out_pose: wp.array(dtype=wp.transformf), + out_pos: wp.array(dtype=wp.vec3f), + out_quat: wp.array(dtype=wp.quatf), +): + i = wp.tid() + site_idx = site_indices[i] + body_idx = shape_body[site_idx] + site_xform = shape_transform[site_idx] + if body_idx == -1: + world_xform = site_xform + else: + world_xform = wp.transform_multiply(body_q[body_idx], site_xform) + out_pose[i] = world_xform + out_pos[i] = wp.transform_get_translation(world_xform) + out_quat[i] = wp.transform_get_rotation(world_xform) + + +@wp.kernel +def _gather_pose_by_index_kernel( + indices: wp.array(dtype=wp.int32), + pos_src: wp.array(dtype=wp.vec3f), + quat_src: wp.array(dtype=wp.quatf), + pos_dst: wp.array(dtype=wp.vec3f), + quat_dst: wp.array(dtype=wp.quatf), +): + i = wp.tid() + src_idx = indices[i] + pos_dst[i] = pos_src[src_idx] + quat_dst[i] = quat_src[src_idx] + + +def _find_physics_ancestor(prim): + """Return the nearest rigid-body ancestor for a sensor or target prim.""" + ancestor = prim + while ancestor and ancestor.IsValid() and ancestor.GetPath().pathString != "/": + if ancestor.HasAPI(UsdPhysics.RigidBodyAPI): + return ancestor + ancestor = ancestor.GetParent() + return None + + +def _newton_body_pattern(body_path: str) -> str: + """Convert a concrete env index to a regex wildcard for prototype body matching.""" + body_path = body_path.replace("{}", ".*") + return re.sub(r"^(/World/envs/)env_\d+/", r"\1env_.*/", body_path) + + +def _identity_offsets(count: int, device: str) -> tuple[wp.array, wp.array]: + """Create identity sensor offsets for site poses that already include the offset.""" + offset_pos_wp = wp.zeros(count, dtype=wp.vec3f, device=device) + identity_quat = np.zeros((count, 4), dtype=np.float32) + identity_quat[:, 3] = 1.0 + return offset_pos_wp, wp.array(identity_quat, dtype=wp.quatf, device=device) + + +class _NewtonRayCasterMixin: + """Newton site registration and pose tracking for ray-caster sensors. + + Sites must be registered during construction so Newton can inject them into + prototype builders before cloning. Once physics is ready, the mixin resolves + those labels to concrete site indices and updates the sensor-owned buffers + directly from Newton model/state arrays. + """ + + @property + def count(self: Any) -> int: + """Number of resolved Newton sites tracked as sensor frames.""" + return self._view_count + + def __init__(self: Any, cfg): + """Register sensor and dynamic target sites before cloning occurs.""" + super().__init__(cfg) # pyright: ignore[reportCallIssue] + self._sensor_site_labels = self._register_sites_for_expr(self.cfg.prim_path) + self._tracked_site_labels_by_expr: dict[str | tuple[str, ...], list[str]] = {} + for target_cfg in getattr(self, "_raycast_targets_cfg", []): + if target_cfg.track_mesh_transforms: + owner_exprs = self._resolve_target_owner_exprs(target_cfg.prim_expr) + labels = self._register_target_sites_for_exprs(owner_exprs) + self._tracked_site_labels_by_expr[target_cfg.prim_expr] = labels + self._tracked_site_labels_by_expr[tuple(owner_exprs)] = labels + + def _register_sites_for_expr(self, prim_expr: str) -> list[str]: + """Register Newton sites for a prim expression and return site labels.""" + prims = sim_utils.find_matching_prims(prim_expr) + labels: list[str] = [] + if len(prims) == 0: + identity = wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat(0.0, 0.0, 0.0, 1.0)) + return [NewtonManager.cl_register_site(_newton_body_pattern(prim_expr), identity)] + + for prim in prims: + body = _find_physics_ancestor(prim) + if body is None: + pos, quat = sim_utils.resolve_prim_pose(prim) + xform = wp.transform(wp.vec3(*[float(v) for v in pos]), wp.quat(*[float(v) for v in quat])) + labels.append(NewtonManager.cl_register_site(None, xform)) + else: + pos, quat = sim_utils.resolve_prim_pose(prim, body) + xform = wp.transform(wp.vec3(*[float(v) for v in pos]), wp.quat(*[float(v) for v in quat])) + labels.append(NewtonManager.cl_register_site(_newton_body_pattern(str(body.GetPath())), xform)) + # Keep the first copy of each label; cloned envs can report the same prototype site more than once. + return list(dict.fromkeys(labels)) + + def _resolve_target_owner_exprs(self, prim_expr: str) -> list[str]: + """Resolve mesh target expressions to owning rigid-body expressions.""" + prims = sim_utils.find_matching_prims(prim_expr) + if len(prims) == 0: + return [_newton_body_pattern(prim_expr)] + + owner_exprs: list[str] = [] + for prim in prims: + body = _find_physics_ancestor(prim) + if body is None: + raise RuntimeError( + f"Cannot track non-physics ray-cast target '{prim_expr}' with Newton. " + "Set track_mesh_transforms=False for static targets, or apply RigidBodyAPI to dynamic targets." + ) + owner_exprs.append(_newton_body_pattern(str(body.GetPath()))) + return list(dict.fromkeys(owner_exprs)) + + def _register_target_sites_for_exprs(self, owner_exprs: list[str]) -> list[str]: + """Register identity-pose Newton sites on target owner bodies.""" + identity = wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat(0.0, 0.0, 0.0, 1.0)) + labels = [NewtonManager.cl_register_site(owner_expr, identity) for owner_expr in owner_exprs] + return list(dict.fromkeys(labels)) + + def _initialize_pose_tracking(self: Any) -> None: + """Resolve registered site labels and allocate sensor-owned pose buffers.""" + site_indices = self._resolve_site_indices(self._sensor_site_labels, self.cfg.prim_path, self._num_envs) + # The base classes still use ``self._view.count`` in a few generic + # places. Point it at the sensor instead of constructing an adapter. + self._view = self + self._view_count = len(site_indices) + self._sensor_site_indices = wp.array(site_indices, dtype=wp.int32, device=self._device) + self._newton_pose_w = wp.empty(self._view_count, dtype=wp.transformf, device=self._device) + self._newton_pos_w = ProxyArray(wp.empty(self._view_count, dtype=wp.vec3f, device=self._device)) + self._newton_quat_w = ProxyArray(wp.empty(self._view_count, dtype=wp.quatf, device=self._device)) + self._offset_pos_wp, self._offset_quat_wp = _identity_offsets(self._view_count, self._device) + + def _update_ray_infos(self: Any, env_mask: wp.array): + """Update Newton site poses and transform local rays in a single ray-caster kernel.""" + self._update_newton_site_transforms( + self._sensor_site_indices, self._newton_pose_w, self._newton_pos_w.warp, self._newton_quat_w.warp + ) + pos_w = self._data.pos_w.warp + quat_w = self._data.quat_w_world.warp if hasattr(self._data, "quat_w_world") else self._data.quat_w.warp + ray_starts = self.ray_starts.warp if hasattr(self.ray_starts, "warp") else self._ray_starts_local + ray_directions = ( + self.ray_directions.warp if hasattr(self.ray_directions, "warp") else self._ray_directions_local + ) + alignment_mode = int(ALIGNMENT_BASE) if hasattr(self._data, "quat_w_world") else self._alignment_mode + wp.launch( + update_ray_caster_kernel, + dim=(self._num_envs, self.num_rays), + inputs=[ + self._newton_pose_w, + env_mask, + self._offset_pos_wp, + self._offset_quat_wp, + self.drift.warp, + self.ray_cast_drift.warp, + ray_starts, + ray_directions, + alignment_mode, + ], + outputs=[ + pos_w, + quat_w, + self._ray_starts_w, + self._ray_directions_w, + ], + device=self._device, + ) + + def get_world_poses(self: Any, indices=None): + """Return world poses for camera helpers that still use pose tuples.""" + self._update_newton_site_transforms( + self._sensor_site_indices, self._newton_pose_w, self._newton_pos_w.warp, self._newton_quat_w.warp + ) + if indices is None: + return self._newton_pos_w, self._newton_quat_w + if not isinstance(indices, wp.array): + indices = wp.array(indices, dtype=wp.int32, device=self._device) + pos_w = wp.empty(indices.shape[0], dtype=wp.vec3f, device=self._device) + quat_w = wp.empty(indices.shape[0], dtype=wp.quatf, device=self._device) + wp.launch( + _gather_pose_by_index_kernel, + dim=indices.shape[0], + inputs=[indices, self._newton_pos_w.warp, self._newton_quat_w.warp], + outputs=[pos_w, quat_w], + device=self._device, + ) + return ProxyArray(pos_w), ProxyArray(quat_w) + + def _create_tracked_target_view(self: Any, target_prim_path: str | list[str]): + """Resolve dynamic multi-mesh target sites to raw Newton site indices.""" + target_key = tuple(target_prim_path) if isinstance(target_prim_path, list) else target_prim_path + labels = self._tracked_site_labels_by_expr.get(target_key) + if labels is None: + target_exprs = target_prim_path if isinstance(target_prim_path, list) else [target_prim_path] + labels = self._register_target_sites_for_exprs([_newton_body_pattern(expr) for expr in target_exprs]) + self._tracked_site_labels_by_expr[target_key] = labels + site_indices = self._resolve_site_indices(labels, str(target_prim_path), self._num_envs) + return wp.array(site_indices, dtype=wp.int32, device=self._device) + + def _update_mesh_transforms(self: Any) -> None: + """Refresh dynamic multi-mesh targets directly from Newton sites.""" + if not hasattr(self, "_mesh_views"): + return + mesh_idx = 0 + for site_indices, 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 + + site_count = site_indices.shape[0] + pos_buf = wp.empty(site_count, dtype=wp.vec3f, device=self._device) + quat_buf = wp.empty(site_count, dtype=wp.quatf, device=self._device) + pose_buf = wp.empty(site_count, dtype=wp.transformf, device=self._device) + self._update_newton_site_transforms(site_indices, pose_buf, pos_buf, quat_buf) + meshes_per_env = site_count + if site_count != 1: + # Newton sites arrive as a flat list across envs; the mesh table is indexed per env. + meshes_per_env = site_count // self._num_envs + + wp.launch( + copy_mesh_poses_to_table_kernel, + dim=(self._num_envs, meshes_per_env), + inputs=[ + pos_buf, + quat_buf, + int(meshes_per_env), + int(mesh_idx), + bool(site_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_newton_site_transforms( + self: Any, + site_indices: wp.array, + pose_buf: wp.array, + pos_buf: wp.array, + quat_buf: wp.array, + ) -> None: + """Launch the Newton site pose kernel into caller-provided buffers.""" + model = NewtonManager._model + state = NewtonManager._state_0 + if model is None or state is None: + raise RuntimeError("Newton simulation state is not initialized.") + wp.launch( + _newton_site_world_poses_kernel, + dim=site_indices.shape[0], + inputs=[site_indices, model.shape_body, model.shape_transform, state.body_q], + outputs=[pose_buf, pos_buf, quat_buf], + device=self._device, + ) + + @staticmethod + def _resolve_site_indices(labels: list[str], prim_expr: str, num_envs: int) -> list[int]: + """Expand registered site labels into per-environment Newton site indices.""" + site_map = NewtonManager._cl_site_index_map + site_indices: list[int] = [] + for env_idx in range(num_envs): + for label in labels: + error_prefix = f"RayCaster target '{prim_expr}' site label '{label}'" + if label not in site_map: + raise ValueError(f"{error_prefix} was not found in NewtonManager._cl_site_index_map.") + global_idx, per_world = site_map[label] + env_site_indices = [global_idx] if per_world is None else per_world[env_idx] + site_indices.extend(env_site_indices) + return site_indices + + +class RayCaster(_NewtonRayCasterMixin, BaseRayCaster): + """Newton ray-caster implementation.""" diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster_camera.py b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster_camera.py new file mode 100644 index 00000000000..6e8f56fd20c --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster_camera.py @@ -0,0 +1,14 @@ +# 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 + +from isaaclab.sensors.ray_caster.base_ray_caster_camera import BaseRayCasterCamera + +from .ray_caster import _NewtonRayCasterMixin + + +class RayCasterCamera(_NewtonRayCasterMixin, BaseRayCasterCamera): + """Newton RayCasterCamera implementation.""" diff --git a/source/isaaclab_physx/changelog.d/octi-raycaster-backend-split.minor.rst b/source/isaaclab_physx/changelog.d/octi-raycaster-backend-split.minor.rst new file mode 100644 index 00000000000..f3130755127 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/octi-raycaster-backend-split.minor.rst @@ -0,0 +1,26 @@ +Added +^^^^^ + +* Added PhysX backend for :class:`~isaaclab.sensors.ray_caster.RayCaster` / + :class:`~isaaclab.sensors.ray_caster.RayCasterCamera` / + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCaster` / + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCasterCamera`. Sensor + body and tracked target meshes both run off ``RigidObjectView`` — + per-step compose via small warp kernels, no + :class:`~isaaclab_physx.sim.views.FabricFrameView` path. Static + parents/targets serve cached per-env ``wp.transformf`` arrays. + +Fixed +^^^^^ + +* Fixed all four ray-caster sensors (:class:`~isaaclab.sensors.ray_caster.RayCaster`, + :class:`~isaaclab.sensors.ray_caster.RayCasterCamera`, + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCaster`, + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCasterCamera`) returning + their spawn-time pose forever when parented under a rigid body. Previous + path went through :class:`~isaaclab_physx.sim.views.FabricFrameView` + which regressed in #5179; the new backend reads body pose directly from + PhysX. The same fix applies to tracked target meshes + (``track_mesh_transforms=True``) parented under rigid bodies. +* Fixed PhysX tracked target mesh updates to write directly into Warp mesh + pose tables instead of staging through torch views. diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/__init__.pyi b/source/isaaclab_physx/isaaclab_physx/sensors/__init__.pyi index e536b281952..3cc07b029af 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/__init__.pyi +++ b/source/isaaclab_physx/isaaclab_physx/sensors/__init__.pyi @@ -15,6 +15,10 @@ __all__ = [ "JointWrenchSensorData", "Pva", "PvaData", + "MultiMeshRayCaster", + "MultiMeshRayCasterCamera", + "RayCaster", + "RayCasterCamera", ] from .contact_sensor import ContactSensor, ContactSensorData, ContactSensorCfg @@ -22,3 +26,4 @@ from .frame_transformer import FrameTransformer, FrameTransformerData from .imu import Imu, ImuData from .joint_wrench import JointWrenchSensor, JointWrenchSensorData from .pva import Pva, PvaData +from .ray_caster import MultiMeshRayCaster, MultiMeshRayCasterCamera, RayCaster, RayCasterCamera diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/__init__.py b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/__init__.py new file mode 100644 index 00000000000..60e772b29dc --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/__init__.py @@ -0,0 +1,10 @@ +# 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 + +"""Sub-module for PhysX ray-caster sensor.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/__init__.pyi b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/__init__.pyi new file mode 100644 index 00000000000..7562409a7e9 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/__init__.pyi @@ -0,0 +1,16 @@ +# 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 + +__all__ = [ + "MultiMeshRayCaster", + "MultiMeshRayCasterCamera", + "RayCaster", + "RayCasterCamera", +] + +from .multi_mesh_ray_caster import MultiMeshRayCaster +from .multi_mesh_ray_caster_camera import MultiMeshRayCasterCamera +from .ray_caster import RayCaster +from .ray_caster_camera import RayCasterCamera diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/multi_mesh_ray_caster.py b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/multi_mesh_ray_caster.py new file mode 100644 index 00000000000..554bcb41351 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/multi_mesh_ray_caster.py @@ -0,0 +1,14 @@ +# 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 + +from isaaclab.sensors.ray_caster.base_multi_mesh_ray_caster import BaseMultiMeshRayCaster + +from .ray_caster import _PhysXRayCasterMixin + + +class MultiMeshRayCaster(_PhysXRayCasterMixin, BaseMultiMeshRayCaster): + """PhysX MultiMeshRayCaster implementation.""" diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/multi_mesh_ray_caster_camera.py b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/multi_mesh_ray_caster_camera.py new file mode 100644 index 00000000000..afd5436fa0d --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/multi_mesh_ray_caster_camera.py @@ -0,0 +1,14 @@ +# 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 + +from isaaclab.sensors.ray_caster.base_multi_mesh_ray_caster_camera import BaseMultiMeshRayCasterCamera + +from .ray_caster import _PhysXRayCasterMixin + + +class MultiMeshRayCasterCamera(_PhysXRayCasterMixin, BaseMultiMeshRayCasterCamera): + """PhysX MultiMeshRayCasterCamera implementation.""" 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 new file mode 100644 index 00000000000..9597abd785d --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py @@ -0,0 +1,210 @@ +# 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 + +from types import SimpleNamespace +from typing import Any + +import torch +import warp as wp + +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +from isaaclab.sensors.ray_caster.base_ray_caster import BaseRayCaster +from isaaclab.sensors.ray_caster.kernels import copy_mesh_transforms_to_table_kernel + +from isaaclab_physx.physics import PhysxManager + + +def _find_physics_ancestor(prim): + """Return the nearest rigid-body ancestor for a sensor or target prim.""" + ancestor = prim + while ancestor and ancestor.IsValid() and ancestor.GetPath().pathString != "/": + if ancestor.HasAPI(UsdPhysics.RigidBodyAPI): + return ancestor + ancestor = ancestor.GetParent() + return None + + +def _body_expr_from_sensor_expr(sensor_expr: str, first_sensor_prim, first_body_prim) -> str: + """Convert a sensor/target expression to the matching rigid-body expression.""" + sensor_path = first_sensor_prim.GetPath().pathString + body_path = first_body_prim.GetPath().pathString + if sensor_path == body_path: + return sensor_expr + # Example: ``.../Robot/base/sensor`` target -> ``.../Robot/base`` body view. + suffix = sensor_path[len(body_path) :] + if suffix and sensor_expr.endswith(suffix): + return sensor_expr[: -len(suffix)] + return body_path + + +def _physx_body_glob(body_expr: str) -> str: + """Convert internal env regex/template expressions to PhysX glob syntax.""" + return body_expr.replace("{}", "*").replace(".*", "*") + + +class _PhysXRayCasterMixin: + """PhysX pose tracking for ray-caster sensors. + + PhysX can provide live rigid-body transforms after physics is ready. Static + non-physics prims are cached once at initialization; they are intentionally + not polled through USD during sensor updates. + """ + + @property + def count(self: Any) -> int: + """Number of tracked sensor frames.""" + return self._view_count + + def _initialize_pose_tracking(self: Any) -> None: + """Initialize direct PhysX body tracking or a cached static pose table.""" + prims = sim_utils.find_matching_prims(self.cfg.prim_path) + if len(prims) == 0: + raise RuntimeError(f"No sensor prims matched: {self.cfg.prim_path}") + + # The base classes still use ``self._view.count`` in a few generic + # places. Point it at the sensor instead of constructing an adapter. + self._view = self + body = _find_physics_ancestor(prims[0]) + if body is None: + self._initialize_static_pose_tracking(prims) + return + + requested_prim_path = getattr(self, "_requested_prim_path", self.cfg.prim_path) + # When the public prim path pointed at a rigid body, BaseRayCaster + # spawned a child sensor prim and preserved the original body path. + body_expr = ( + requested_prim_path + if self.cfg.prim_path != requested_prim_path + else _body_expr_from_sensor_expr(self.cfg.prim_path, prims[0], body) + ) + physics_sim_view = PhysxManager.get_physics_sim_view() + if physics_sim_view is None: + raise RuntimeError("PhysX simulation view is not initialized.") + self._physx_body_view = physics_sim_view.create_rigid_body_view(body_expr.replace(".*", "*")) + self._view_count = self._physx_body_view.count + + offset_pos = [] + offset_quat = [] + for prim in prims: + body_prim = _find_physics_ancestor(prim) + p, q = sim_utils.resolve_prim_pose(prim, body_prim) + offset_pos.append(p) + offset_quat.append(q) + if len(offset_pos) == 1 and self._view_count > 1: + offset_pos = offset_pos * self._view_count + offset_quat = offset_quat * self._view_count + self._offset_pos_wp = wp.array(offset_pos[: self._view_count], dtype=wp.vec3f, device=self._device) + self._offset_quat_contiguous = torch.tensor( + offset_quat[: self._view_count], dtype=torch.float32, device=self._device + ) + self._offset_quat_wp = wp.from_torch(self._offset_quat_contiguous, dtype=wp.quatf) + + def _initialize_static_pose_tracking(self: Any, prims) -> None: + """Cache authored poses for non-physics sensor frames.""" + poses = [] + for prim in prims: + pos, quat = sim_utils.resolve_prim_pose(prim) + poses.append((*pos, *quat)) + self._static_view_transforms_torch = torch.tensor(poses, dtype=torch.float32, device=self._device).contiguous() + self._static_view_transforms_wp = wp.from_torch(self._static_view_transforms_torch).view(wp.transformf) + self._physx_body_view = None + self._view_count = len(prims) + self._offset_pos_wp = wp.zeros(self._view_count, dtype=wp.vec3f, device=self._device) + identity_quat = torch.zeros(self._view_count, 4, device=self._device) + identity_quat[:, 3] = 1.0 + self._offset_quat_contiguous = identity_quat.contiguous() + self._offset_quat_wp = wp.from_torch(self._offset_quat_contiguous, dtype=wp.quatf) + + def _get_view_transforms_wp(self: Any) -> wp.array: + """Return tracked sensor-frame transforms as ``wp.transformf``.""" + if self._physx_body_view is None: + return self._static_view_transforms_wp + transforms = self._physx_body_view.get_transforms() + if isinstance(transforms, wp.array): + return transforms.view(wp.transformf) + return wp.from_torch(transforms.contiguous()).view(wp.transformf) + + def get_world_poses(self: Any, indices=None): + """Return world poses for camera helpers that still use pose tuples.""" + transforms = self._get_view_transforms_wp() + transforms_t = wp.to_torch(transforms).reshape(-1, 7) + if indices is not None: + idx = wp.to_torch(indices).to(dtype=torch.long) if isinstance(indices, wp.array) else indices + transforms_t = transforms_t[idx] + return SimpleNamespace(torch=transforms_t[:, 0:3]), SimpleNamespace(torch=transforms_t[:, 3:7]) + + def _create_tracked_target_view(self: Any, target_prim_paths: str | list[str]): + """Create a PhysX rigid-body view for dynamic multi-mesh targets.""" + if isinstance(target_prim_paths, str): + target_prim_paths = [target_prim_paths] + body_paths = [] + for target_prim_path in target_prim_paths: + prims = sim_utils.find_matching_prims(target_prim_path) + if len(prims) == 0: + # ClonePlan-backed targets may not have destination mesh prims. + # In that case BaseMultiMeshRayCaster passes the destination owner-body expression. + body_paths.append(target_prim_path) + continue + for prim in prims: + body = _find_physics_ancestor(prim) + if body is None: + raise RuntimeError( + f"Cannot track non-physics ray-cast target '{target_prim_path}' with PhysX. " + "Set track_mesh_transforms=False for static targets, or apply RigidBodyAPI to dynamic targets." + ) + body_paths.append(body.GetPath().pathString) + + if len(body_paths) == 0: + raise RuntimeError(f"No tracked target bodies resolved from: {target_prim_paths}") + physics_sim_view = PhysxManager.get_physics_sim_view() + if physics_sim_view is None: + raise RuntimeError("PhysX simulation view is not initialized.") + return physics_sim_view.create_rigid_body_view([_physx_body_glob(path) for path in body_paths]) + + def _update_mesh_transforms(self: Any) -> None: + """Refresh dynamic multi-mesh targets directly from PhysX views.""" + if not hasattr(self, "_mesh_views"): + return + 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 + + transforms = view.get_transforms() + transforms_wp = ( + transforms.view(wp.transformf) + if isinstance(transforms, wp.array) + else wp.from_torch(transforms.contiguous()).view(wp.transformf) + ) + + view_count = view.count + meshes_per_env = view_count + if view_count != 1: + # PhysX 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_transforms_to_table_kernel, + dim=(self._num_envs, meshes_per_env), + inputs=[ + transforms_wp, + 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] + + +class RayCaster(_PhysXRayCasterMixin, BaseRayCaster): + """PhysX ray-caster implementation.""" diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster_camera.py b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster_camera.py new file mode 100644 index 00000000000..ff84a3cd037 --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster_camera.py @@ -0,0 +1,14 @@ +# 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 + +from isaaclab.sensors.ray_caster.base_ray_caster_camera import BaseRayCasterCamera + +from .ray_caster import _PhysXRayCasterMixin + + +class RayCasterCamera(_PhysXRayCasterMixin, BaseRayCasterCamera): + """PhysX RayCasterCamera implementation.""" diff --git a/source/isaaclab_tasks/changelog.d/octi-raycaster-backend-split.minor.rst b/source/isaaclab_tasks/changelog.d/octi-raycaster-backend-split.minor.rst new file mode 100644 index 00000000000..aecca907c48 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/octi-raycaster-backend-split.minor.rst @@ -0,0 +1,8 @@ +Added +^^^^^ + +* Added raycaster-camera depth presets (``raycaster_depth64``, ``raycaster_depth128``, + ``raycaster_depth256``) for both base and wrist views in the Dexsuite Kuka-Allegro + manipulation task, backed by + :class:`~isaaclab.sensors.ray_caster.MultiMeshRayCasterCamera`. Targets the table, + ground plane, manipulated object, and robot visuals. diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/camera_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/camera_cfg.py index 3ad4be97b0f..83e4564553e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/camera_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/config/kuka_allegro/camera_cfg.py @@ -9,7 +9,7 @@ from isaaclab.managers import ObservationGroupCfg as ObsGroup from isaaclab.managers import ObservationTermCfg as ObsTerm from isaaclab.managers import SceneEntityCfg -from isaaclab.sensors import CameraCfg +from isaaclab.sensors import CameraCfg, MultiMeshRayCasterCameraCfg, patterns from isaaclab.utils.configclass import configclass from isaaclab.utils.noise import UniformNoiseCfg as Unoise @@ -50,6 +50,54 @@ renderer_cfg=MultiBackendRendererCfg(), ) +RAY_PATTERN = patterns.PinholeCameraPatternCfg(focal_length=24.0, horizontal_aperture=20.955) + +RAYCASTER_CAMERA_MESH_PRIM_PATHS = [ + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + prim_expr="/World/envs/env_.*/table", + track_mesh_transforms=False, + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + prim_expr="/World/GroundPlane", + track_mesh_transforms=False, + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + prim_expr="/World/envs/env_.*/Object", + track_mesh_transforms=True, + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + prim_expr="/World/envs/env_.*/Robot/.*/visuals", + track_mesh_transforms=True, + ), +] + +BASE_RAYCASTER_CAMERA_CFG = MultiMeshRayCasterCameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=MultiMeshRayCasterCameraCfg.OffsetCfg( + pos=(0.57, -0.8, 0.5), + rot=(0.6124, 0.3536, 0.3536, 0.6124), + convention="opengl", + ), + mesh_prim_paths=RAYCASTER_CAMERA_MESH_PRIM_PATHS, + max_distance=2.5, + data_types=["distance_to_image_plane"], + pattern_cfg=MISSING, +) + +WRIST_RAYCASTER_CAMERA_CFG = MultiMeshRayCasterCameraCfg( + prim_path="/World/envs/env_.*/Robot/ee_link/palm_link/Camera", + offset=MultiMeshRayCasterCameraCfg.OffsetCfg( + pos=(0.038, -0.38, -0.18), + rot=(0.641, 0.641, -0.299, 0.299), + convention="opengl", + ), + mesh_prim_paths=RAYCASTER_CAMERA_MESH_PRIM_PATHS, + max_distance=2.5, + data_types=["distance_to_image_plane"], + pattern_cfg=MISSING, + debug_vis=False, +) + @configclass class BaseTiledCameraCfg(PresetCfg): @@ -88,6 +136,10 @@ class BaseTiledCameraCfg(PresetCfg): semantic_segmentation64 = BASE_CAMERA_CFG.replace(data_types=["semantic_segmentation"], width=64, height=64) semantic_segmentation128 = BASE_CAMERA_CFG.replace(data_types=["semantic_segmentation"], width=128, height=128) semantic_segmentation256 = BASE_CAMERA_CFG.replace(data_types=["semantic_segmentation"], width=256, height=256) + # raycaster camera presets + raycaster_depth64 = BASE_RAYCASTER_CAMERA_CFG.replace(pattern_cfg=RAY_PATTERN.replace(width=64, height=64)) + raycaster_depth128 = BASE_RAYCASTER_CAMERA_CFG.replace(pattern_cfg=RAY_PATTERN.replace(width=128, height=128)) + raycaster_depth256 = BASE_RAYCASTER_CAMERA_CFG.replace(pattern_cfg=RAY_PATTERN.replace(width=256, height=256)) default = rgb64 @@ -128,6 +180,10 @@ class WristTiledCameraCfg(PresetCfg): semantic_segmentation64 = WRIST_CAMERA_CFG.replace(data_types=["semantic_segmentation"], width=64, height=64) semantic_segmentation128 = WRIST_CAMERA_CFG.replace(data_types=["semantic_segmentation"], width=128, height=128) semantic_segmentation256 = WRIST_CAMERA_CFG.replace(data_types=["semantic_segmentation"], width=256, height=256) + # raycaster camera presets + raycaster_depth64 = WRIST_RAYCASTER_CAMERA_CFG.replace(pattern_cfg=RAY_PATTERN.replace(width=64, height=64)) + raycaster_depth128 = WRIST_RAYCASTER_CAMERA_CFG.replace(pattern_cfg=RAY_PATTERN.replace(width=128, height=128)) + raycaster_depth256 = WRIST_RAYCASTER_CAMERA_CFG.replace(pattern_cfg=RAY_PATTERN.replace(width=256, height=256)) default = rgb64 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py index 06faa0c2598..eeea53b4841 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/dexsuite_env_cfg.py @@ -476,7 +476,7 @@ def validate_config(self): cam = getattr(self.scene, cam_attr, None) if cam is None: continue - renderer_type = getattr(cam.renderer_cfg, "renderer_type", None) + renderer_type = getattr(getattr(cam, "renderer_cfg", None), "renderer_type", None) if renderer_type == "newton_warp": unsupported = set(cam.data_types) - warp_supported if unsupported: