diff --git a/docs/source/api/lab/isaaclab.utils.rst b/docs/source/api/lab/isaaclab.utils.rst index 9ef89739f79..9ae7239951c 100644 --- a/docs/source/api/lab/isaaclab.utils.rst +++ b/docs/source/api/lab/isaaclab.utils.rst @@ -14,6 +14,7 @@ dict interpolation math + mesh modifiers noise string @@ -89,6 +90,14 @@ Math operations :inherited-members: :show-inheritance: +Mesh operations +~~~~~~~~~~~~~~~ + +.. automodule:: isaaclab.utils.mesh + :members: + :imported-members: + :show-inheritance: + Modifier operations ~~~~~~~~~~~~~~~~~~~ diff --git a/scripts/demos/sensors/multi_mesh_raycaster.py b/scripts/demos/sensors/multi_mesh_raycaster.py new file mode 100644 index 00000000000..0769d4598c2 --- /dev/null +++ b/scripts/demos/sensors/multi_mesh_raycaster.py @@ -0,0 +1,296 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Example on using the MultiMesh Raycaster sensor. + +Usage: + `python scripts/demos/sensors/multi_mesh_raycaster.py --num_envs 16 --asset_type ` +""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Example on using the raycaster sensor.") +parser.add_argument("--num_envs", type=int, default=16, help="Number of environments to spawn.") +parser.add_argument( + "--asset_type", + type=str, + default="allegro_hand", + help="Asset type to use.", + choices=["allegro_hand", "anymal_d", "multi"], +) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import random +import torch + +import omni.usd +from pxr import Gf, Sdf + +## +# Pre-defined configs +## +from isaaclab_assets.robots.allegro import ALLEGRO_HAND_CFG +from isaaclab_assets.robots.anymal import ANYMAL_D_CFG + +import isaaclab.sim as sim_utils +from isaaclab.assets import Articulation, AssetBaseCfg, RigidObjectCfg +from isaaclab.markers.config import VisualizationMarkersCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sensors.ray_caster import MultiMeshRayCasterCfg, patterns +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + +RAY_CASTER_MARKER_CFG = VisualizationMarkersCfg( + markers={ + "hit": sim_utils.SphereCfg( + radius=0.01, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)), + ), + }, +) + + +if args_cli.asset_type == "allegro_hand": + asset_cfg = ALLEGRO_HAND_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot", + update_period=1 / 60, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0, -0.1, 0.3)), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/thumb_link_.*/visuals_xform"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/index_link.*/visuals_xform"), + MultiMeshRayCasterCfg.RaycastTargetCfg( + target_prim_expr="{ENV_REGEX_NS}/Robot/middle_link_.*/visuals_xform" + ), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/ring_link_.*/visuals_xform"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/palm_link/visuals_xform"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/allegro_mount/visuals_xform"), + ], + ray_alignment="world", + pattern_cfg=patterns.GridPatternCfg(resolution=0.005, size=(0.4, 0.4), direction=(0, 0, -1)), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) + +elif args_cli.asset_type == "anymal_d": + asset_cfg = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot", + update_period=1 / 60, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0, -0.1, 0.3)), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LF_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RF_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LH_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RH_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_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)), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) + +elif args_cli.asset_type == "multi": + asset_cfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Object", + spawn=sim_utils.MultiAssetSpawnerCfg( + assets_cfg=[ + sim_utils.CuboidCfg( + size=(0.3, 0.3, 0.3), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2), + ), + sim_utils.SphereCfg( + radius=0.3, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0), metallic=0.2), + ), + sim_utils.CylinderCfg( + radius=0.2, + height=0.5, + axis="Y", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + ), + sim_utils.CapsuleCfg( + radius=0.15, + height=0.5, + axis="Z", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 1.0, 0.0), metallic=0.2), + ), + sim_utils.ConeCfg( + radius=0.2, + height=0.5, + axis="Z", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 1.0), metallic=0.2), + ), + ], + random_choice=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 2.0)), + ) + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="{ENV_REGEX_NS}/Object", + update_period=1 / 60, + offset=MultiMeshRayCasterCfg.OffsetCfg(pos=(0, 0.0, 0.6)), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Object"), + ], + ray_alignment="world", + pattern_cfg=patterns.GridPatternCfg(resolution=0.01, size=(0.6, 0.6), direction=(0, 0, -1)), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) +else: + raise ValueError(f"Unknown asset type: {args_cli.asset_type}") + + +@configclass +class RaycasterSensorSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the asset.""" + + # ground plane + ground = AssetBaseCfg( + prim_path="/World/Ground", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", + scale=(1, 1, 1), + ), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # asset + asset = asset_cfg + # ray caster + ray_caster = ray_caster_cfg + + +def randomize_shape_color(prim_path_expr: str): + """Randomize the color of the geometry.""" + + # acquire stage + stage = omni.usd.get_context().get_stage() + # resolve prim paths for spawning and cloning + prim_paths = sim_utils.find_matching_prim_paths(prim_path_expr) + # manually clone prims if the source prim path is a regex expression + + with Sdf.ChangeBlock(): + for prim_path in prim_paths: + print("Applying prim scale to:", prim_path) + # spawn single instance + prim_spec = Sdf.CreatePrimInLayer(stage.GetRootLayer(), prim_path) + + # DO YOUR OWN OTHER KIND OF RANDOMIZATION HERE! + # Note: Just need to acquire the right attribute about the property you want to set + # Here is an example on setting color randomly + color_spec = prim_spec.GetAttributeAtPath(prim_path + "/geometry/material/Shader.inputs:diffuseColor") + color_spec.default = Gf.Vec3f(random.random(), random.random(), random.random()) + + # randomize scale + scale_spec = prim_spec.GetAttributeAtPath(prim_path + ".xformOp:scale") + scale_spec.default = Gf.Vec3f(random.uniform(0.5, 1.5), random.uniform(0.5, 1.5), random.uniform(0.5, 1.5)) + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + """Run the simulator.""" + # Define simulation stepping + sim_dt = sim.get_physics_dt() + sim_time = 0.0 + count = 0 + + # Simulate physics + while simulation_app.is_running(): + + if count % 500 == 0: + # reset counter + count = 0 + # reset the scene entities + # root state + root_state = scene["asset"].data.default_root_state.clone() + root_state[:, :3] += scene.env_origins + scene["asset"].write_root_pose_to_sim(root_state[:, :7]) + scene["asset"].write_root_velocity_to_sim(root_state[:, 7:]) + + if isinstance(scene["asset"], Articulation): + # set joint positions with some noise + joint_pos, joint_vel = ( + scene["asset"].data.default_joint_pos.clone(), + scene["asset"].data.default_joint_vel.clone(), + ) + joint_pos += torch.rand_like(joint_pos) * 0.1 + scene["asset"].write_joint_state_to_sim(joint_pos, joint_vel) + # clear internal buffers + scene.reset() + print("[INFO]: Resetting Asset state...") + + if isinstance(scene["asset"], Articulation): + # -- generate actions/commands + targets = scene["asset"].data.default_joint_pos + 5 * ( + torch.rand_like(scene["asset"].data.default_joint_pos) - 0.5 + ) + # -- apply action to the asset + scene["asset"].set_joint_position_target(targets) + # -- write data to sim + scene.write_data_to_sim() + # perform step + sim.step() + # update sim-time + sim_time += sim_dt + count += 1 + # update buffers + scene.update(sim_dt) + + +def main(): + """Main function.""" + + # Initialize the simulation context + sim_cfg = sim_utils.SimulationCfg(dt=0.005, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + # Set main camera + sim.set_camera_view(eye=[3.5, 3.5, 3.5], target=[0.0, 0.0, 0.0]) + # design scene + scene_cfg = RaycasterSensorSceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0, replicate_physics=False) + scene = InteractiveScene(scene_cfg) + + if args_cli.asset_type == "multi": + randomize_shape_color(scene_cfg.asset.prim_path.format(ENV_REGEX_NS="/World/envs/env_.*")) + + # Play the simulator + sim.reset() + # Now we are ready! + print("[INFO]: Setup complete...") + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/demos/sensors/multi_mesh_raycaster_camera.py b/scripts/demos/sensors/multi_mesh_raycaster_camera.py new file mode 100644 index 00000000000..865e4523b3e --- /dev/null +++ b/scripts/demos/sensors/multi_mesh_raycaster_camera.py @@ -0,0 +1,332 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Example on using the MultiMesh Raycaster Camera sensor. + +Usage: + `python scripts/demos/sensors/multi_mesh_raycaster_camera.py --num_envs 16 --asset_type ` +""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Example on using the raycaster sensor.") +parser.add_argument("--num_envs", type=int, default=16, help="Number of environments to spawn.") +parser.add_argument( + "--asset_type", + type=str, + default="allegro_hand", + help="Asset type to use.", + choices=["allegro_hand", "anymal_d", "multi"], +) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import random +import torch + +import omni.usd +from pxr import Gf, Sdf + +## +# Pre-defined configs +## +from isaaclab_assets.robots.allegro import ALLEGRO_HAND_CFG +from isaaclab_assets.robots.anymal import ANYMAL_D_CFG + +import isaaclab.sim as sim_utils +from isaaclab.assets import Articulation, AssetBaseCfg, RigidObjectCfg +from isaaclab.markers.config import VisualizationMarkersCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sensors.ray_caster import MultiMeshRayCasterCameraCfg, patterns +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + +RAY_CASTER_MARKER_CFG = VisualizationMarkersCfg( + markers={ + "hit": sim_utils.SphereCfg( + radius=0.01, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0)), + ), + }, +) + +if args_cli.asset_type == "allegro_hand": + asset_cfg = ALLEGRO_HAND_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + ray_caster_cfg = MultiMeshRayCasterCameraCfg( + prim_path="{ENV_REGEX_NS}/Robot", + update_period=1 / 60, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg( + pos=(-0.70, -0.7, -0.25), rot=(0.268976, 0.268976, 0.653951, 0.653951) + ), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + target_prim_expr="{ENV_REGEX_NS}/Robot/thumb_link_.*/visuals_xform" + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + target_prim_expr="{ENV_REGEX_NS}/Robot/index_link.*/visuals_xform" + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + target_prim_expr="{ENV_REGEX_NS}/Robot/middle_link_.*/visuals_xform" + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + target_prim_expr="{ENV_REGEX_NS}/Robot/ring_link_.*/visuals_xform" + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + target_prim_expr="{ENV_REGEX_NS}/Robot/palm_link/visuals_xform" + ), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + target_prim_expr="{ENV_REGEX_NS}/Robot/allegro_mount/visuals_xform" + ), + ], + pattern_cfg=patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=120, + width=240, + ), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) + +elif args_cli.asset_type == "anymal_d": + asset_cfg = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + ray_caster_cfg = MultiMeshRayCasterCameraCfg( + prim_path="{ENV_REGEX_NS}/Robot", + update_period=1 / 60, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0, -0.1, 1.5), rot=(0.0, 1.0, 0.0, 0.0)), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCameraCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LF_.*/visuals"), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RF_.*/visuals"), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LH_.*/visuals"), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RH_.*/visuals"), + MultiMeshRayCasterCameraCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/base/visuals"), + ], + pattern_cfg=patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=120, + width=240, + ), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) + +elif args_cli.asset_type == "multi": + asset_cfg = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Object", + spawn=sim_utils.MultiAssetSpawnerCfg( + assets_cfg=[ + sim_utils.CuboidCfg( + size=(0.3, 0.3, 0.3), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), metallic=0.2), + ), + sim_utils.SphereCfg( + radius=0.3, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0), metallic=0.2), + ), + sim_utils.CylinderCfg( + radius=0.2, + height=0.5, + axis="Y", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + ), + sim_utils.CapsuleCfg( + radius=0.15, + height=0.5, + axis="Z", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 1.0, 0.0), metallic=0.2), + ), + sim_utils.ConeCfg( + radius=0.2, + height=0.5, + axis="Z", + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 1.0), metallic=0.2), + ), + ], + random_choice=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + mass_props=sim_utils.MassPropertiesCfg(mass=1.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 2.0)), + ) + ray_caster_cfg = MultiMeshRayCasterCameraCfg( + prim_path="{ENV_REGEX_NS}/Object", + update_period=1 / 60, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0, 0.0, 1.5), rot=(0.0, 1.0, 0.0, 0.0)), + mesh_prim_paths=[ + "/World/Ground", + MultiMeshRayCasterCameraCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Object"), + ], + pattern_cfg=patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=120, + width=240, + ), + debug_vis=not args_cli.headless, + visualizer_cfg=RAY_CASTER_MARKER_CFG.replace(prim_path="/Visuals/RayCaster"), + ) +else: + raise ValueError(f"Unknown asset type: {args_cli.asset_type}") + + +@configclass +class RaycasterSensorSceneCfg(InteractiveSceneCfg): + """Design the scene with sensors on the asset.""" + + # ground plane + ground = AssetBaseCfg( + prim_path="/World/Ground", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", + scale=(1, 1, 1), + ), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # asset + asset = asset_cfg + # ray caster + ray_caster = ray_caster_cfg + + +def randomize_shape_color(prim_path_expr: str): + """Randomize the color of the geometry.""" + + # acquire stage + stage = omni.usd.get_context().get_stage() + # resolve prim paths for spawning and cloning + prim_paths = sim_utils.find_matching_prim_paths(prim_path_expr) + # manually clone prims if the source prim path is a regex expression + + with Sdf.ChangeBlock(): + for prim_path in prim_paths: + print("Applying prim scale to:", prim_path) + # spawn single instance + prim_spec = Sdf.CreatePrimInLayer(stage.GetRootLayer(), prim_path) + + # DO YOUR OWN OTHER KIND OF RANDOMIZATION HERE! + # Note: Just need to acquire the right attribute about the property you want to set + # Here is an example on setting color randomly + color_spec = prim_spec.GetAttributeAtPath(prim_path + "/geometry/material/Shader.inputs:diffuseColor") + color_spec.default = Gf.Vec3f(random.random(), random.random(), random.random()) + + # randomize scale + scale_spec = prim_spec.GetAttributeAtPath(prim_path + ".xformOp:scale") + scale_spec.default = Gf.Vec3f(random.uniform(0.5, 1.5), random.uniform(0.5, 1.5), random.uniform(0.5, 1.5)) + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + """Run the simulator.""" + # Define simulation stepping + sim_dt = sim.get_physics_dt() + sim_time = 0.0 + count = 0 + + triggered = True + countdown = 42 + + # Simulate physics + while simulation_app.is_running(): + + if count % 500 == 0: + # reset counter + count = 0 + # reset the scene entities + # root state + root_state = scene["asset"].data.default_root_state.clone() + root_state[:, :3] += scene.env_origins + scene["asset"].write_root_pose_to_sim(root_state[:, :7]) + scene["asset"].write_root_velocity_to_sim(root_state[:, 7:]) + + if isinstance(scene["asset"], Articulation): + # set joint positions with some noise + joint_pos, joint_vel = ( + scene["asset"].data.default_joint_pos.clone(), + scene["asset"].data.default_joint_vel.clone(), + ) + joint_pos += torch.rand_like(joint_pos) * 0.1 + scene["asset"].write_joint_state_to_sim(joint_pos, joint_vel) + # clear internal buffers + scene.reset() + print("[INFO]: Resetting Asset state...") + + if isinstance(scene["asset"], Articulation): + # -- generate actions/commands + targets = scene["asset"].data.default_joint_pos + 5 * ( + torch.rand_like(scene["asset"].data.default_joint_pos) - 0.5 + ) + # -- apply action to the asset + scene["asset"].set_joint_position_target(targets) + # -- write data to sim + scene.write_data_to_sim() + # perform step + sim.step() + # update sim-time + sim_time += sim_dt + count += 1 + # update buffers + scene.update(sim_dt) + + if not triggered: + if countdown > 0: + countdown -= 1 + continue + + data = scene["ray_caster"].data.ray_hits_w.cpu().numpy() # noqa: F841 + triggered = True + else: + continue + + +def main(): + """Main function.""" + + # Initialize the simulation context + sim_cfg = sim_utils.SimulationCfg(dt=0.005, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + # Set main camera + sim.set_camera_view(eye=[3.5, 3.5, 3.5], target=[0.0, 0.0, 0.0]) + # design scene + scene_cfg = RaycasterSensorSceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0, replicate_physics=False) + scene = InteractiveScene(scene_cfg) + + if args_cli.asset_type == "multi": + randomize_shape_color(scene_cfg.asset.prim_path.format(ENV_REGEX_NS="/World/envs/env_.*")) + + # Play the simulator + sim.reset() + # Now we are ready! + print("[INFO]: Setup complete...") + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py b/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py index a4fe1bce519..1652dc9bd4c 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py @@ -6,6 +6,12 @@ """Sub-module for Warp-based ray-cast sensor.""" from . import patterns +from .multi_mesh_ray_caster import MultiMeshRayCaster +from .multi_mesh_ray_caster_camera import MultiMeshRayCasterCamera +from .multi_mesh_ray_caster_camera_cfg import MultiMeshRayCasterCameraCfg +from .multi_mesh_ray_caster_camera_data import MultiMeshRayCasterCameraData +from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg +from .multi_mesh_ray_caster_data import MultiMeshRayCasterData from .ray_caster import RayCaster from .ray_caster_camera import RayCasterCamera from .ray_caster_camera_cfg import RayCasterCameraCfg 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 new file mode 100644 index 00000000000..922805b4873 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py @@ -0,0 +1,424 @@ +# Copyright (c) 2022-2025, 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 + +"""Multi-mesh ray casting sensor implementation. + +This file adds support for ray casting against multiple (possibly regex-selected) mesh targets. +""" + +import numpy as np +import re +import torch +import trimesh +from collections.abc import Sequence +from typing import TYPE_CHECKING, ClassVar + +import carb +import omni.log +import omni.physics.tensors.impl.api as physx +import warp as wp +from isaacsim.core.prims import XFormPrim + +import isaaclab.sim as sim_utils +from isaaclab.utils.math import matrix_from_quat, quat_mul +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, raycast_dynamic_meshes + +from .multi_mesh_ray_caster_data import MultiMeshRayCasterData +from .prim_utils import obtain_world_pose_from_view +from .ray_caster import RayCaster + +if TYPE_CHECKING: + from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg + + +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:`warp_meshes` list. 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, …) 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. + + 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(target_prim_expr="{ENV_REGEX_NS}/Robot/LF_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RF_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/LH_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="{ENV_REGEX_NS}/Robot/RH_.*/visuals"), + MultiMeshRayCasterCfg.RaycastTargetCfg(target_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_offsets: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + + mesh_views: ClassVar[dict[str, XFormPrim | physx.ArticulationView | physx.RigidBodyView]] = {} + """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. + """ + # Initialize base class + super().__init__(cfg) + + # Create empty variables for storing output data + self._num_meshes_per_env: dict[str, int] = {} + """Keeps track of the number of meshes per env for each ray_cast target. + Since we allow regex indexing (e.g. env_*/object_*) they can differ + """ + + self._raycast_targets_cfg: list[MultiMeshRayCasterCfg.RaycastTargetCfg] = [] + for target in self.cfg.mesh_prim_paths: + # Legacy support for string targets. Treat them as global targets. + if isinstance(target, str): + self._raycast_targets_cfg.append( + cfg.RaycastTargetCfg(target_prim_expr=target, track_mesh_transforms=False) + ) + else: + self._raycast_targets_cfg.append(target) + + # Resolve regex namespace if set + for cfg in self._raycast_targets_cfg: + cfg.target_prim_expr = cfg.target_prim_expr.format(ENV_REGEX_NS="/World/envs/env_.*") + + # overwrite the data class + 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: + # update sensors if needed + self._update_outdated_buffers() + # return the data + 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 to ray cast against + target_prim_path = target_cfg.target_prim_expr + # # check if mesh already casted into warp mesh and skip if so. + if target_prim_path in multi_mesh_ids: + carb.log_warn( + 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 + + # find all matching prim paths to provided expression of the target + 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 + ) # If only one prim is found, treat it as a global prim. Either it's a single global object (e.g. ground) or we are only using one env. + + loaded_vertices: list[np.ndarray | None] = [] + wp_mesh_ids = [] + + for target_prim in target_prims: + # Reuse previously parsed shared mesh instance if possible. + 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())) # + if curr_prim_base_path in MultiMeshRayCaster.meshes: + MultiMeshRayCaster.meshes[str(target_prim.GetPath())] = MultiMeshRayCaster.meshes[ + curr_prim_base_path + ] + # Reuse mesh imported by another ray-cast sensor (global cache). + if str(target_prim.GetPath()) in MultiMeshRayCaster.meshes: + wp_mesh_ids.append(MultiMeshRayCaster.meshes[str(target_prim.GetPath())].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()}'" + carb.log_warn(warn_msg) + continue + + trimesh_meshes = [] + + for mesh_prim in mesh_prims: + # check if valid + 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) + + # add to list of parsed meshes + trimesh_meshes.append(mesh) + + if len(trimesh_meshes) == 1: + trimesh_mesh = trimesh_meshes[0] + elif target_cfg.merge_prim_meshes: + # combine all trimesh meshes into a single mesh + 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." + ) + + # check if the mesh is already registered, if so only reference the mesh + registered_idx = _registered_points_idx(trimesh_mesh.vertices, loaded_vertices) + if registered_idx != -1 and self.cfg.reference_meshes: + omni.log.info("Found a duplicate mesh, only reference the mesh.") + # 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())] = wp_mesh + wp_mesh_ids.append(wp_mesh.id) + + # print info + if registered_idx != -1: + omni.log.info(f"Found duplicate mesh for mesh prims under path '{target_prim.GetPath()}'.") + else: + omni.log.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: + # reference the mesh for each environment to ray cast against + 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: + # split up the meshes for each environment. Little bit ugly, since + # the current order is interleaved (env1_obj1, env1_obj2, env2_obj1, env2_obj2, ...) + 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], MultiMeshRayCaster.mesh_offsets[target_prim_path] = ( + self._get_trackable_prim_view(target_prim_path) + ) + + # throw an error if no meshes are found + if all([target_cfg.target_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 = torch.zeros(self._num_envs, total_n_meshes_per_env, 3, device=self.device) + self._mesh_orientations_w = torch.zeros(self._num_envs, total_n_meshes_per_env, 4, device=self.device) + + # Update the mesh positions and rotations + mesh_idx = 0 + for target_cfg in self._raycast_targets_cfg: + n_meshes = self._num_meshes_per_env[target_cfg.target_prim_expr] + + # update position of the target meshes + pos_w, ori_w = [], [] + for prim in sim_utils.find_matching_prims(target_cfg.target_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[:, mesh_idx : mesh_idx + n_meshes] = pos_w + self._mesh_orientations_w[:, mesh_idx : mesh_idx + n_meshes] = ori_w + mesh_idx += n_meshes + + # flatten the list of meshes that are included in mesh_prim_paths of the specific ray caster + 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.target_prim_expr][env_idx]) + multi_mesh_ids_flattened.append(meshes_in_env) + + self._mesh_views = [ + self.mesh_views[target_cfg.target_prim_expr] if target_cfg.track_mesh_transforms else None + for target_cfg in self._raycast_targets_cfg + ] + + # save a warp array with mesh ids that is passed to the raycast function + 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() + if self.cfg.update_mesh_ids: + self._data.ray_mesh_ids = torch.zeros( + self._num_envs, self.num_rays, 1, device=self.device, dtype=torch.int16 + ) + + def _update_buffers_impl(self, env_ids: Sequence[int]): + """Fills the buffers of the sensor data. + + Args: + env_ids: The environment ids to update. + """ + + self._update_ray_infos(env_ids) + + # Update the mesh positions and rotations + 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.target_prim_expr] + continue + + # update position of the target meshes + pos_w, ori_w = obtain_world_pose_from_view(view, None) + 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 + + if target_cfg.target_prim_expr in MultiMeshRayCaster.mesh_offsets: + pos_offset, ori_offset = MultiMeshRayCaster.mesh_offsets[target_cfg.target_prim_expr] + pos_w -= pos_offset + ori_w = quat_mul(ori_offset.expand(ori_w.shape[0], -1), ori_w) + + count = view.count + if count != 1: # Mesh is not global, i.e. we have different meshes for each env + 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[:, mesh_idx : mesh_idx + count] = pos_w + self._mesh_orientations_w[:, mesh_idx : mesh_idx + count] = ori_w + mesh_idx += count + + self._data.ray_hits_w[env_ids], _, _, _, mesh_ids = raycast_dynamic_meshes( + self._ray_starts_w[env_ids], + self._ray_directions_w[env_ids], + mesh_ids_wp=self._mesh_ids_wp, # list with shape num_envs x num_meshes_per_env + max_dist=self.cfg.max_distance, + mesh_positions_w=self._mesh_positions_w[env_ids], + mesh_orientations_w=self._mesh_orientations_w[env_ids], + return_mesh_id=self.cfg.update_mesh_ids, + ) + + if self.cfg.update_mesh_ids: + self._data.ray_mesh_ids[env_ids] = mesh_ids + + def __del__(self): + super().__del__() + if RayCaster._instance_count == 0: + MultiMeshRayCaster.mesh_offsets.clear() + 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 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 new file mode 100644 index 00000000000..259c64af53a --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py @@ -0,0 +1,220 @@ +# Copyright (c) 2022-2025, 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 torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.utils.warp import raycast_dynamic_meshes + +from .multi_mesh_ray_caster import MultiMeshRayCaster +from .multi_mesh_ray_caster_camera_data import MultiMeshRayCasterCameraData +from .prim_utils import obtain_world_pose_from_view +from .ray_caster_camera import RayCasterCamera + +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): + # 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 stars 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] + # create buffer to store ray hits + self.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self._device) + # set offsets + 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_quat = quat_w.repeat(self._view.count, 1) + self._offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device).repeat(self._view.count, 1) + + self._data.quat_w = torch.zeros(self._view.count, 4, device=self.device) + self._data.pos_w = torch.zeros(self._view.count, 3, device=self.device) + + self._ray_starts_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + self._ray_directions_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + + def _update_ray_infos(self, env_ids: Sequence[int]): + """Updates the ray information buffers.""" + + # compute poses from current view + pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + ) + # update the data + self._data.pos_w[env_ids] = pos_w + self._data.quat_w_world[env_ids] = quat_w + self._data.quat_w_ros[env_ids] = quat_w + + # note: full orientation is considered + ray_starts_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids]) + ray_starts_w += pos_w.unsqueeze(1) + ray_directions_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_directions[env_ids]) + + self._ray_starts_w[env_ids] = ray_starts_w + self._ray_directions_w[env_ids] = ray_directions_w + + def _update_buffers_impl(self, env_ids: Sequence[int] | torch.Tensor | None): + """Fills the buffers of the sensor data.""" + self._update_ray_infos(env_ids) + + # increment frame count + if env_ids is None: + env_ids = torch.arange(self._num_envs, device=self.device) + elif not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, device=self.device) + + self._frame[env_ids] += 1 + + # Update the mesh positions and rotations + 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.target_prim_expr] + continue + + # update position of the target meshes + pos_w, ori_w = obtain_world_pose_from_view(view, None) + 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 + + if target_cfg.target_prim_expr in MultiMeshRayCaster.mesh_offsets: + pos_offset, ori_offset = MultiMeshRayCaster.mesh_offsets[target_cfg.target_prim_expr] + pos_w -= pos_offset + ori_w = math_utils.quat_mul(ori_offset.expand(ori_w.shape[0], -1), ori_w) + + count = view.count + if count != 1: # Mesh is not global, i.e. we have different meshes for each env + 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[:, mesh_idx : mesh_idx + count] = pos_w + self._mesh_orientations_w[:, mesh_idx : mesh_idx + count] = ori_w + mesh_idx += count + + # ray cast and store the hits + self.ray_hits_w, ray_depth, ray_normal, _, ray_mesh_ids = raycast_dynamic_meshes( + self._ray_starts_w[env_ids], + self._ray_directions_w[env_ids], + mesh_ids_wp=self._mesh_ids_wp, # list with shape num_envs x num_meshes_per_env + max_dist=self.cfg.max_distance, + mesh_positions_w=self._mesh_positions_w[env_ids], + mesh_orientations_w=self._mesh_orientations_w[env_ids], + return_distance=any( + [name in self.cfg.data_types for name in ["distance_to_image_plane", "distance_to_camera"]] + ), + return_normal="normals" in self.cfg.data_types, + return_mesh_id=self.cfg.update_mesh_ids, + ) + + # update output buffers + if "distance_to_image_plane" in self.cfg.data_types: + # note: data is in camera frame so we only take the first component (z-axis of camera frame) + distance_to_image_plane = ( + math_utils.quat_apply( + math_utils.quat_inv(self._data.quat_w_world[env_ids]).repeat(1, self.num_rays), + (ray_depth[:, :, None] * self._ray_directions_w[env_ids]), + ) + )[:, :, 0] + # apply the maximum distance after the transformation + if self.cfg.depth_clipping_behavior == "max": + distance_to_image_plane = torch.clip(distance_to_image_plane, max=self.cfg.max_distance) + distance_to_image_plane[torch.isnan(distance_to_image_plane)] = self.cfg.max_distance + elif self.cfg.depth_clipping_behavior == "zero": + distance_to_image_plane[distance_to_image_plane > self.cfg.max_distance] = 0.0 + distance_to_image_plane[torch.isnan(distance_to_image_plane)] = 0.0 + self._data.output["distance_to_image_plane"][env_ids] = distance_to_image_plane.view( + -1, *self.image_shape, 1 + ) + + if "distance_to_camera" in self.cfg.data_types: + if self.cfg.depth_clipping_behavior == "max": + ray_depth = torch.clip(ray_depth, max=self.cfg.max_distance) + elif self.cfg.depth_clipping_behavior == "zero": + ray_depth[ray_depth > self.cfg.max_distance] = 0.0 + self._data.output["distance_to_camera"][env_ids] = ray_depth.view(-1, *self.image_shape, 1) + + if "normals" in self.cfg.data_types: + self._data.output["normals"][env_ids] = ray_normal.view(-1, *self.image_shape, 3) + + if self.cfg.update_mesh_ids: + self._data.image_mesh_ids[env_ids] = ray_mesh_ids.view(-1, *self.image_shape, 1) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_cfg.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_cfg.py new file mode 100644 index 00000000000..85478eaef27 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_cfg.py @@ -0,0 +1,32 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the ray-cast camera sensor.""" + +import carb + +from isaaclab.utils import configclass + +from .multi_mesh_ray_caster_camera import MultiMeshRayCasterCamera +from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg +from .ray_caster_camera_cfg import RayCasterCameraCfg + + +@configclass +class MultiMeshRayCasterCameraCfg(RayCasterCameraCfg, MultiMeshRayCasterCfg): + """Configuration for the multi-mesh ray-cast camera sensor.""" + + class_type: type = MultiMeshRayCasterCamera + + def __post_init__(self): + super().__post_init__() + + # Camera only supports 'base' ray alignment. Ensure this is set correctly. + if self.ray_alignment != "base": + carb.log_warn( + "Ray alignment for MultiMeshRayCasterCameraCfg only supports 'base' alignment. Overriding from" + f"'{self.ray_alignment}' to 'base'." + ) + self.ray_alignment = "base" 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 new file mode 100644 index 00000000000..fe9eef128ae --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera_data.py @@ -0,0 +1,21 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Data container for the multi-mesh ray-cast camera sensor.""" +import torch + +from isaaclab.sensors.camera import CameraData + +from .ray_caster_data import RayCasterData + + +class MultiMeshRayCasterCameraData(CameraData, RayCasterData): + """Data container for the multi-mesh ray-cast sensor.""" + + image_mesh_ids: torch.Tensor = 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, + and 1 is the number of mesh ids per pixel.""" diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_cfg.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_cfg.py new file mode 100644 index 00000000000..26f775c531a --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_cfg.py @@ -0,0 +1,70 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Configuration for the ray-cast sensor.""" + + +from dataclasses import MISSING + +from isaaclab.utils import configclass + +from .multi_mesh_ray_caster import MultiMeshRayCaster +from .ray_caster_cfg import RayCasterCfg + + +@configclass +class MultiMeshRayCasterCfg(RayCasterCfg): + """Configuration for the multi-mesh ray-cast sensor.""" + + @configclass + class RaycastTargetCfg: + """Configuration for different ray-cast targets.""" + + target_prim_expr: str = MISSING + """The regex to specify the target prim to ray cast against.""" + + is_shared: bool = False + """Whether the target prim is assumed to be the same mesh across all environments. Defaults to False. + + If True, only the first mesh is read and then reused for all environments, rather than re-parsed. + This provides a startup performance boost when there are many environments that all use the same asset. + + .. note:: + If :attr:`MultiMeshRayCasterCfg.reference_meshes` is False, this flag has no effect. + """ + + merge_prim_meshes: bool = True + """Whether to merge the parsed meshes for a prim that contains multiple meshes. Defaults to True. + + This will create a new mesh that combines all meshes in the parsed prim. The raycast hits mesh IDs will then refer to the single + merged mesh. + """ + + track_mesh_transforms: bool = True + """Whether the mesh transformations should be tracked. Defaults to True. + + .. note:: + Not tracking the mesh transformations is recommended when the meshes are static to increase performance. + """ + + class_type: type = MultiMeshRayCaster + + mesh_prim_paths: list[str | RaycastTargetCfg] = MISSING + """The list of mesh primitive paths to ray cast against. + + If an entry is a string, it is internally converted to :class:`RaycastTargetCfg` with `~RaycastTargetCfg.track_mesh_transforms` disabled. These settings ensure backwards compatibility with the default raycaster. + """ + + update_mesh_ids: bool = False + """Whether to update the mesh ids of the ray hits in the :attr:`data` container.""" + + reference_meshes: bool = True + """Whether to reference duplicated meshes instead of loading each one separately into memory. + Defaults to True. + + When enabled, the raycaster parses all meshes in all environments, but reuses references + for duplicates instead of storing multiple copies. This reduces memory footprint. + """ 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 new file mode 100644 index 00000000000..6bf5e97cf4a --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_data.py @@ -0,0 +1,20 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Data container for the multi-mesh ray-cast sensor.""" +import torch + +from .ray_caster_data import RayCasterData + + +class MultiMeshRayCasterData(RayCasterData): + """Data container for the multi-mesh ray-cast sensor.""" + + ray_mesh_ids: torch.Tensor = 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 + in the scan pattern per sensor.""" diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/prim_utils.py b/source/isaaclab/isaaclab/sensors/ray_caster/prim_utils.py new file mode 100644 index 00000000000..52eba9e8dd6 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/prim_utils.py @@ -0,0 +1,43 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch + +import omni.physics.tensors.impl.api as physx +from isaacsim.core.prims import XFormPrim + +from isaaclab.utils.math import convert_quat + + +def obtain_world_pose_from_view( + physx_view: XFormPrim | physx.ArticulationView | physx.RigidBodyView, + env_ids: torch.Tensor, + clone: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Get the world poses of the prim referenced by the prim view. + Args: + physx_view: The prim view to get the world poses from. + env_ids: The environment ids of the prims to get the world poses for. + clone: Whether to clone the returned tensors (default: False). + Returns: + A tuple containing the world positions and orientations of the prims. Orientation is in wxyz format. + Raises: + NotImplementedError: If the prim view is not of the correct type. + """ + if isinstance(physx_view, XFormPrim): + pos_w, quat_w = physx_view.get_world_poses(env_ids) + elif isinstance(physx_view, physx.ArticulationView): + pos_w, quat_w = physx_view.get_root_transforms()[env_ids].split([3, 4], dim=-1) + quat_w = convert_quat(quat_w, to="wxyz") + elif isinstance(physx_view, physx.RigidBodyView): + pos_w, quat_w = physx_view.get_transforms()[env_ids].split([3, 4], dim=-1) + quat_w = convert_quat(quat_w, to="wxyz") + else: + raise NotImplementedError(f"Cannot get world poses for prim view of type '{type(physx_view)}'.") + + if clone: + return pos_w.clone(), quat_w.clone() + else: + return pos_w, quat_w diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py index 8e2f6541fe9..53dce3afefc 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster.py @@ -9,10 +9,10 @@ import re import torch from collections.abc import Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar +import isaacsim.core.utils.stage as stage_utils import omni.log -import omni.physics.tensors.impl.api as physx import warp as wp from isaacsim.core.prims import XFormPrim from isaacsim.core.simulation_manager import SimulationManager @@ -22,10 +22,11 @@ import isaaclab.utils.math as math_utils from isaaclab.markers import VisualizationMarkers from isaaclab.terrains.trimesh.utils import make_plane -from isaaclab.utils.math import convert_quat, quat_apply, quat_apply_yaw +from isaaclab.utils.math import quat_apply, quat_apply_yaw from isaaclab.utils.warp import convert_to_warp_mesh, raycast_mesh from ..sensor_base import SensorBase +from .prim_utils import obtain_world_pose_from_view from .ray_caster_data import RayCasterData if TYPE_CHECKING: @@ -51,12 +52,21 @@ class RayCaster(SensorBase): cfg: RayCasterCfg """The configuration parameters.""" + # Class variables to share meshes across instances + meshes: ClassVar[dict[str, wp.Mesh]] = {} + """A dictionary to store warp meshes for raycasting, shared across all instances. + + The keys correspond to the prim path for the meshes, and values are the corresponding warp Mesh objects.""" + _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 # check if sensor path is valid # note: currently we do not handle environment indices if there is a regex pattern in the leaf # For example, if the prim path is "/World/Sensor_[1,2]". @@ -64,15 +74,13 @@ def __init__(self, cfg: RayCasterCfg): sensor_path_is_regex = re.match(r"^[a-zA-Z0-9/_]+$", sensor_path) is None if sensor_path_is_regex: raise RuntimeError( - f"Invalid prim path for the ray-caster sensor: {self.cfg.prim_path}." + f"Invalid prim path for the ray-caster sensor: {cfg.prim_path}." "\n\tHint: Please ensure that the prim path does not contain any regex patterns in the leaf." ) # Initialize base class super().__init__(cfg) # Create empty variables for storing output data self._data = RayCasterData() - # the warp meshes used for raycasting. - self.meshes: dict[str, wp.Mesh] = {} def __str__(self) -> str: """Returns: A string containing information about the instance.""" @@ -80,7 +88,7 @@ def __str__(self) -> str: 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(self.meshes)}\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}" @@ -131,28 +139,16 @@ def reset(self, env_ids: Sequence[int] | None = None): def _initialize_impl(self): super()._initialize_impl() # obtain global simulation view + self._physics_sim_view = SimulationManager.get_physics_sim_view() - # check if the prim at path is an articulated or rigid prim - # we do this since for physics-based view classes we can access their data directly - # otherwise we need to use the xform view class which is slower - found_supported_prim_class = False prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) if prim is None: - raise RuntimeError(f"Failed to find a prim at path expression: {self.cfg.prim_path}") - # create view based on the type of prim - if prim.HasAPI(UsdPhysics.ArticulationRootAPI): - self._view = self._physics_sim_view.create_articulation_view(self.cfg.prim_path.replace(".*", "*")) - found_supported_prim_class = True - elif prim.HasAPI(UsdPhysics.RigidBodyAPI): - self._view = self._physics_sim_view.create_rigid_body_view(self.cfg.prim_path.replace(".*", "*")) - found_supported_prim_class = True - else: - self._view = XFormPrim(self.cfg.prim_path, reset_xform_properties=False) - found_supported_prim_class = True - omni.log.warn(f"The prim at path {prim.GetPath().pathString} is not a physics prim! Using XFormPrim.") - # check if prim view class is found - if not found_supported_prim_class: - raise RuntimeError(f"Failed to find a valid prim view class for the prim paths: {self.cfg.prim_path}") + available_prims = ",".join([str(p.GetPath()) for p in stage_utils.get_current_stage().Traverse()]) + raise RuntimeError( + f"Failed to find a prim at path expression: {self.cfg.prim_path}. Available prims: {available_prims}" + ) + + self._view, self._offset = self._get_trackable_prim_view(self.cfg.prim_path) # load the meshes by parsing the stage self._initialize_warp_meshes() @@ -168,6 +164,10 @@ def _initialize_warp_meshes(self): # read prims to ray-cast for mesh_prim_path in self.cfg.mesh_prim_paths: + # check if mesh already casted into warp mesh + if mesh_prim_path in RayCaster.meshes: + continue + # check if the prim is a plane - handle PhysX plane as a special case # if a plane exists then we need to create an infinite mesh that is a plane mesh_prim = sim_utils.get_first_matching_child_prim( @@ -201,10 +201,10 @@ def _initialize_warp_meshes(self): # print info omni.log.info(f"Created infinite plane mesh prim: {mesh_prim.GetPath()}.") # add the warp mesh to the list - self.meshes[mesh_prim_path] = wp_mesh + RayCaster.meshes[mesh_prim_path] = wp_mesh # throw an error if no meshes are found - if all([mesh_prim_path not in self.meshes for mesh_prim_path in self.cfg.mesh_prim_paths]): + if all([mesh_prim_path 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}" ) @@ -225,26 +225,19 @@ def _initialize_rays_impl(self): self.drift = torch.zeros(self._view.count, 3, device=self.device) self.ray_cast_drift = torch.zeros(self._view.count, 3, device=self.device) # fill the data buffer - self._data.pos_w = torch.zeros(self._view.count, 3, device=self._device) - self._data.quat_w = torch.zeros(self._view.count, 4, device=self._device) - self._data.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self._device) - - def _update_buffers_impl(self, env_ids: Sequence[int]): - """Fills the buffers of the sensor data.""" - # obtain the poses of the sensors - if isinstance(self._view, XFormPrim): - pos_w, quat_w = self._view.get_world_poses(env_ids) - elif isinstance(self._view, physx.ArticulationView): - pos_w, quat_w = self._view.get_root_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = convert_quat(quat_w, to="wxyz") - elif isinstance(self._view, physx.RigidBodyView): - pos_w, quat_w = self._view.get_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = convert_quat(quat_w, to="wxyz") - else: - raise RuntimeError(f"Unsupported view type: {type(self._view)}") - # note: we clone here because we are read-only operations - pos_w = pos_w.clone() - quat_w = quat_w.clone() + self._data.pos_w = torch.zeros(self._view.count, 3, device=self.device) + self._data.quat_w = torch.zeros(self._view.count, 4, device=self.device) + self._data.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + self._ray_starts_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + self._ray_directions_w = torch.zeros(self._view.count, self.num_rays, 3, device=self.device) + + def _update_ray_infos(self, env_ids: Sequence[int]): + """Updates the ray information buffers.""" + + pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, self._offset[0][env_ids], self._offset[1][env_ids] + ) # apply drift to ray starting position in world frame pos_w += self.drift[env_ids] # store the poses @@ -291,13 +284,20 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): else: raise RuntimeError(f"Unsupported ray_alignment type: {self.cfg.ray_alignment}.") + self._ray_starts_w[env_ids] = ray_starts_w + self._ray_directions_w[env_ids] = ray_directions_w + + def _update_buffers_impl(self, env_ids: Sequence[int]): + """Fills the buffers of the sensor data.""" + self._update_ray_infos(env_ids) + # ray cast and store the hits # TODO: Make this work for multiple meshes? self._data.ray_hits_w[env_ids] = raycast_mesh( - ray_starts_w, - ray_directions_w, + self._ray_starts_w[env_ids], + self._ray_directions_w[env_ids], max_dist=self.cfg.max_distance, - mesh=self.meshes[self.cfg.mesh_prim_paths[0]], + mesh=RayCaster.meshes[self.cfg.mesh_prim_paths[0]], )[0] # apply vertical drift to ray starting position in ray caster frame @@ -316,12 +316,72 @@ def _set_debug_vis_impl(self, debug_vis: bool): self.ray_visualizer.set_visibility(False) def _debug_vis_callback(self, event): + if self._data.ray_hits_w is None: + return # remove possible inf values viz_points = self._data.ray_hits_w.reshape(-1, 3) viz_points = viz_points[~torch.any(torch.isinf(viz_points), dim=1)] - # show ray hit positions + self.ray_visualizer.visualize(viz_points) + def _get_trackable_prim_view( + self, target_prim_path: str + ) -> tuple[XFormPrim | any, tuple[torch.Tensor, torch.Tensor]]: + """Get a prim view that can be used to track the pose of the mesh prims. Additionally, it resolves the + relative pose between the mesh and its corresponding physics prim. This is especially useful if the + mesh is not directly parented to the physics prim. + """ + + mesh_prim = sim_utils.find_first_matching_prim(target_prim_path) + current_prim = mesh_prim + current_path_expr = target_prim_path + + prim_view = None + + while prim_view is None: + if current_prim.HasAPI(UsdPhysics.ArticulationRootAPI): + prim_view = self._physics_sim_view.create_articulation_view(current_path_expr.replace(".*", "*")) + omni.log.info(f"Created articulation view for mesh prim at path: {target_prim_path}") + break + + if current_prim.HasAPI(UsdPhysics.RigidBodyAPI): + prim_view = self._physics_sim_view.create_rigid_body_view(current_path_expr.replace(".*", "*")) + omni.log.info(f"Created rigid body view for mesh prim at path: {target_prim_path}") + break + + new_root_prim = current_prim.GetParent() + current_path_expr = current_path_expr.rsplit("/", 1)[0] + if not new_root_prim.IsValid(): + prim_view = XFormPrim(target_prim_path, reset_xform_properties=False) + current_path_expr = target_prim_path + omni.log.warn( + f"The prim at path {target_prim_path} which is used for raycasting is not a physics prim." + " Defaulting to XFormPrim. \n The pose of the mesh will most likely not" + " be updated correctly when running in headless mode and position lookups will be much slower. \n" + " If possible, ensure that the mesh or its parent is a physics prim (rigid body or articulation)." + ) + break + current_prim = new_root_prim + + mesh_prims = sim_utils.find_matching_prims(target_prim_path) + target_prims = sim_utils.find_matching_prims(current_path_expr) + if len(mesh_prims) != len(target_prims): + raise RuntimeError( + f"The number of mesh prims ({len(mesh_prims)}) does not match the number of physics prims" + f" ({len(target_prims)})Please specify the correct mesh and physics prim paths more" + " specifically in your target expressions." + ) + positions = [] + quaternions = [] + for mesh, target in zip(mesh_prims, target_prims): + pos, orientation = sim_utils.resolve_prim_pose(mesh, target) + positions.append(torch.tensor(pos, dtype=torch.float32, device=self.device)) + quaternions.append(torch.tensor(orientation, dtype=torch.float32, device=self.device)) + + positions = torch.stack(positions).to(device=self.device, dtype=torch.float32) + quaternions = torch.stack(quaternions).to(device=self.device, dtype=torch.float32) + return prim_view, (positions, quaternions) + """ Internal simulation callbacks. """ @@ -332,3 +392,8 @@ def _invalidate_initialize_callback(self, event): super()._invalidate_initialize_callback(event) # set all existing views to None to invalidate them self._view = None + + def __del__(self): + RayCaster._instance_count -= 1 + if RayCaster._instance_count == 0: + RayCaster.meshes.clear() 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 3bf8729b78b..eb877cedace 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera.py @@ -10,13 +10,13 @@ from typing import TYPE_CHECKING, ClassVar, Literal import isaacsim.core.utils.stage as stage_utils -import omni.physics.tensors.impl.api as physx -from isaacsim.core.prims import XFormPrim +import omni.log import isaaclab.utils.math as math_utils from isaaclab.sensors.camera import CameraData from isaaclab.utils.warp import raycast_mesh +from .prim_utils import obtain_world_pose_from_view from .ray_caster import RayCaster if TYPE_CHECKING: @@ -86,7 +86,7 @@ def __str__(self) -> str: 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(self.meshes)}\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" @@ -143,11 +143,14 @@ def reset(self, env_ids: Sequence[int] | None = None): # reset the timestamps super().reset(env_ids) # resolve None - if env_ids is None: - env_ids = slice(None) + if env_ids is None or isinstance(env_ids, slice): + env_ids = self._ALL_INDICES # reset the data # note: this recomputation is useful if one performs events such as randomizations on the camera poses. - pos_w, quat_w = self._compute_camera_world_poses(env_ids) + pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, 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 @@ -184,11 +187,11 @@ def set_world_poses( RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first. """ # resolve env_ids - if env_ids is None: + if env_ids is None or isinstance(env_ids, slice): env_ids = self._ALL_INDICES # get current positions - pos_w, quat_w = self._compute_view_world_poses(env_ids) + pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids) if positions is not None: # transform to camera frame pos_offset_world_frame = positions - pos_w @@ -201,7 +204,10 @@ def set_world_poses( self._offset_quat[env_ids] = math_utils.quat_mul(math_utils.quat_inv(quat_w), quat_w_set) # update the data - pos_w, quat_w = self._compute_camera_world_poses(env_ids) + pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, 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 @@ -260,7 +266,10 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): self._frame[env_ids] += 1 # compute poses from current view - pos_w, quat_w = self._compute_camera_world_poses(env_ids) + pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + pos_w, quat_w = math_utils.combine_frame_transforms( + pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids] + ) # update the data self._data.pos_w[env_ids] = pos_w self._data.quat_w_world[env_ids] = quat_w @@ -280,7 +289,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): self.ray_hits_w, ray_depth, ray_normal, _ = raycast_mesh( ray_starts_w, ray_directions_w, - mesh=self.meshes[self.cfg.mesh_prim_paths[0]], + mesh=RayCaster.meshes[self.cfg.mesh_prim_paths[0]], max_dist=1e6, return_distance=any( [name in self.cfg.data_types for name in ["distance_to_image_plane", "distance_to_camera"]] @@ -395,39 +404,46 @@ def _compute_intrinsic_matrices(self): 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.0: + This function will be removed in a future release in favor of implementation :meth:`obtain_world_pose_from_view`. + Returns: A tuple of the position (in meters) and quaternion (w, x, y, z). + + """ - # obtain the poses of the sensors - # note: clone arg doesn't exist for xform prim view so we need to do this manually - if isinstance(self._view, XFormPrim): - if isinstance(env_ids, slice): # catch the case where env_ids is a slice - env_ids = self._ALL_INDICES - pos_w, quat_w = self._view.get_world_poses(env_ids) - elif isinstance(self._view, physx.ArticulationView): - pos_w, quat_w = self._view.get_root_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = math_utils.convert_quat(quat_w, to="wxyz") - elif isinstance(self._view, physx.RigidBodyView): - pos_w, quat_w = self._view.get_transforms()[env_ids].split([3, 4], dim=-1) - quat_w = math_utils.convert_quat(quat_w, to="wxyz") - else: - raise RuntimeError(f"Unsupported view type: {type(self._view)}") - # return the pose - return pos_w.clone(), quat_w.clone() + # deprecation + omni.log.warn( + "The function '_compute_view_world_poses' will be deprecated in favor of the util method" + " 'obtain_world_pose_from_view'. Please use 'obtain_world_pose_from_view' instead...." + ) + + return obtain_world_pose_from_view(self._view, env_ids, clone=True) 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.0: + This function will be removed in a future release. Instead, use the code block below: + + .. code-block:: python + + pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + 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 (w, x, y, z) in "world" convention. """ - # get the pose of the view the camera is attached to - pos_w, quat_w = self._compute_view_world_poses(env_ids) - # apply offsets - # need to apply quat because offset relative to parent frame - pos_w += math_utils.quat_apply(quat_w, self._offset_pos[env_ids]) - quat_w = math_utils.quat_mul(quat_w, self._offset_quat[env_ids]) - return pos_w, quat_w + # deprecation + omni.log.warn( + "The function '_compute_camera_world_poses' will be deprecated in favor of the combination of methods" + " 'obtain_world_pose_from_view' and 'math_utils.combine_frame_transforms'. Please use" + " 'obtain_world_pose_from_view' and 'math_utils.combine_frame_transforms' instead...." + ) + + # get the pose of the view the camera is attached to + pos_w, quat_w = obtain_world_pose_from_view(self._view, env_ids, clone=True) + return math_utils.combine_frame_transforms(pos_w, quat_w, self._offset_pos[env_ids], self._offset_quat[env_ids]) diff --git a/source/isaaclab/isaaclab/utils/__init__.py b/source/isaaclab/isaaclab/utils/__init__.py index a5365948699..036d7c9c0c6 100644 --- a/source/isaaclab/isaaclab/utils/__init__.py +++ b/source/isaaclab/isaaclab/utils/__init__.py @@ -10,6 +10,7 @@ from .configclass import configclass from .dict import * from .interpolation import * +from .mesh import * from .modifiers import * from .string import * from .timer import Timer diff --git a/source/isaaclab/isaaclab/utils/mesh.py b/source/isaaclab/isaaclab/utils/mesh.py new file mode 100644 index 00000000000..cfeeacabbff --- /dev/null +++ b/source/isaaclab/isaaclab/utils/mesh.py @@ -0,0 +1,170 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +"""Utility functions for working with meshes.""" + +import numpy as np +import trimesh +from collections.abc import Callable + +from pxr import Usd, UsdGeom + + +def create_trimesh_from_geom_mesh(mesh_prim: Usd.Prim) -> trimesh.Trimesh: + """Reads the vertices and faces of a mesh prim. + + The function reads the vertices and faces of a mesh prim and returns it. If the underlying mesh is a quad mesh, + it converts it to a triangle mesh. + + Args: + mesh_prim: The mesh prim to read the vertices and faces from. + + Returns: + A trimesh.Trimesh object containing the mesh geometry. + """ + if mesh_prim.GetTypeName() != "Mesh": + raise ValueError(f"Prim at path '{mesh_prim.GetPath()}' is not a mesh.") + # cast into UsdGeomMesh + mesh = UsdGeom.Mesh(mesh_prim) + + # read the vertices and faces + points = np.asarray(mesh.GetPointsAttr().Get()).copy() + + # Load faces and convert to triangle if needed. (Default is quads) + num_vertex_per_face = np.asarray(mesh.GetFaceVertexCountsAttr().Get()) + indices = np.asarray(mesh.GetFaceVertexIndicesAttr().Get()) + return trimesh.Trimesh(points, convert_faces_to_triangles(indices, num_vertex_per_face)) + + +def create_trimesh_from_geom_shape(prim: Usd.Prim) -> trimesh.Trimesh: + """Converts a primitive object to a trimesh. + + Args: + prim: The prim that should be converted to a trimesh. + + Returns: + A trimesh object representing the primitive. + + Raises: + ValueError: If the prim is not a supported primitive. Check PRIMITIVE_MESH_TYPES for supported primitives. + """ + + if prim.GetTypeName() not in PRIMITIVE_MESH_TYPES: + raise ValueError(f"Prim at path '{prim.GetPath()}' is not a primitive mesh. Cannot convert to trimesh.") + + return _MESH_CONVERTERS_CALLBACKS[prim.GetTypeName()](prim) + + +def convert_faces_to_triangles(faces: np.ndarray, point_counts: np.ndarray) -> np.ndarray: + """Converts quad mesh face indices into triangle face indices. + + This function expects an array of faces (indices) and the number of points per face. It then converts potential + quads into triangles and returns the new triangle face indices as a numpy array of shape (n_faces_new, 3). + + Args: + faces: The faces of the quad mesh as a one-dimensional array. Shape is (N,). + point_counts: The number of points per face. Shape is (N,). + + Returns: + The new face ids with triangles. Shape is (n_faces_new, 3). + """ + # check if the mesh is already triangulated + if (point_counts == 3).all(): + return faces.reshape(-1, 3) # already triangulated + all_faces = [] + + vertex_counter = 0 + # Iterates over all faces of the mesh to triangulate them. + # could be very slow for large meshes + for num_points in point_counts: + # Triangulate n-gons (n>4) using fan triangulation + for i in range(num_points - 2): + triangle = np.array([faces[vertex_counter], faces[vertex_counter + 1 + i], faces[vertex_counter + 2 + i]]) + all_faces.append(triangle) + + vertex_counter += num_points + return np.asarray(all_faces) + + +def _create_plane_trimesh(prim: Usd.Prim) -> trimesh.Trimesh: + """Creates a trimesh for a plane primitive.""" + size = (2e6, 2e6) + vertices = np.array([[size[0], size[1], 0], [size[0], 0.0, 0], [0.0, size[1], 0], [0.0, 0.0, 0]]) - np.array( + [size[0] / 2.0, size[1] / 2.0, 0.0] + ) + faces = np.array([[1, 0, 2], [2, 3, 1]]) + return trimesh.Trimesh(vertices=vertices, faces=faces) + + +def _create_cube_trimesh(prim: Usd.Prim) -> trimesh.Trimesh: + """Creates a trimesh for a cube primitive.""" + size = prim.GetAttribute("size").Get() + extends = [size, size, size] + return trimesh.creation.box(extends) + + +def _create_sphere_trimesh(prim: Usd.Prim, subdivisions: int = 2) -> trimesh.Trimesh: + """Creates a trimesh for a sphere primitive.""" + radius = prim.GetAttribute("radius").Get() + mesh = trimesh.creation.icosphere(radius=radius, subdivisions=subdivisions) + return mesh + + +def _create_cylinder_trimesh(prim: Usd.Prim) -> trimesh.Trimesh: + """Creates a trimesh for a cylinder primitive.""" + radius = prim.GetAttribute("radius").Get() + height = prim.GetAttribute("height").Get() + mesh = trimesh.creation.cylinder(radius=radius, height=height) + axis = prim.GetAttribute("axis").Get() + if axis == "X": + # rotate −90° about Y to point the length along +X + R = trimesh.transformations.rotation_matrix(np.radians(-90), [0, 1, 0]) + mesh.apply_transform(R) + elif axis == "Y": + # rotate +90° about X to point the length along +Y + R = trimesh.transformations.rotation_matrix(np.radians(90), [1, 0, 0]) + mesh.apply_transform(R) + return mesh + + +def _create_capsule_trimesh(prim: Usd.Prim) -> trimesh.Trimesh: + """Creates a trimesh for a capsule primitive.""" + radius = prim.GetAttribute("radius").Get() + height = prim.GetAttribute("height").Get() + mesh = trimesh.creation.capsule(radius=radius, height=height) + axis = prim.GetAttribute("axis").Get() + if axis == "X": + # rotate −90° about Y to point the length along +X + R = trimesh.transformations.rotation_matrix(np.radians(-90), [0, 1, 0]) + mesh.apply_transform(R) + elif axis == "Y": + # rotate +90° about X to point the length along +Y + R = trimesh.transformations.rotation_matrix(np.radians(90), [1, 0, 0]) + mesh.apply_transform(R) + return mesh + + +def _create_cone_trimesh(prim: Usd.Prim) -> trimesh.Trimesh: + """Creates a trimesh for a cone primitive.""" + radius = prim.GetAttribute("radius").Get() + height = prim.GetAttribute("height").Get() + mesh = trimesh.creation.cone(radius=radius, height=height) + # shift all vertices down by height/2 for usd / trimesh cone primitive definition discrepancy + mesh.apply_translation((0.0, 0.0, -height / 2.0)) + return mesh + + +_MESH_CONVERTERS_CALLBACKS: dict[str, Callable[[Usd.Prim], trimesh.Trimesh]] = { + "Plane": _create_plane_trimesh, + "Cube": _create_cube_trimesh, + "Sphere": _create_sphere_trimesh, + "Cylinder": _create_cylinder_trimesh, + "Capsule": _create_capsule_trimesh, + "Cone": _create_cone_trimesh, +} + +PRIMITIVE_MESH_TYPES = list(_MESH_CONVERTERS_CALLBACKS.keys()) +"""List of supported primitive mesh types that can be converted to a trimesh.""" diff --git a/source/isaaclab/isaaclab/utils/warp/__init__.py b/source/isaaclab/isaaclab/utils/warp/__init__.py index 14c49f25528..8400fb670a0 100644 --- a/source/isaaclab/isaaclab/utils/warp/__init__.py +++ b/source/isaaclab/isaaclab/utils/warp/__init__.py @@ -5,4 +5,4 @@ """Sub-module containing operations based on warp.""" -from .ops import convert_to_warp_mesh, raycast_mesh +from .ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh, raycast_single_mesh diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py index 2fe544651f5..748c01c7344 100644 --- a/source/isaaclab/isaaclab/utils/warp/kernels.py +++ b/source/isaaclab/isaaclab/utils/warp/kernels.py @@ -75,6 +75,173 @@ def raycast_mesh_kernel( ray_face_id[tid] = f +@wp.kernel(enable_backward=False) +def raycast_static_meshes_kernel( + mesh: wp.array2d(dtype=wp.uint64), + ray_starts: wp.array2d(dtype=wp.vec3), + ray_directions: wp.array2d(dtype=wp.vec3), + ray_hits: wp.array2d(dtype=wp.vec3), + ray_distance: wp.array2d(dtype=wp.float32), + ray_normal: wp.array2d(dtype=wp.vec3), + ray_face_id: wp.array2d(dtype=wp.int32), + ray_mesh_id: wp.array2d(dtype=wp.int16), + max_dist: float = 1e6, + return_normal: int = False, + return_face_id: int = False, + return_mesh_id: int = False, +): + """Performs ray-casting against multiple static meshes. + + This function performs ray-casting against the given meshes using the provided ray start positions + and directions. The resulting ray hit positions are stored in the :obj:`ray_hits` array. + + The function utilizes the ``mesh_query_ray`` method from the ``wp`` module to perform the actual ray-casting + operation. The maximum ray-cast distance is set to ``1e6`` units. + + .. note:: + That the ``ray_starts``, ``ray_directions``, and ``ray_hits`` arrays should have compatible shapes + and data types to ensure proper execution. Additionally, they all must be in the same frame. + + This kernel differs from the :meth:`raycast_dynamic_meshes_kernel` in that it does not take into + account the mesh's position and rotation. This kernel is useful for ray-casting against static meshes + that are not expected to move. + + Args: + mesh: The input mesh. The ray-casting is performed against this mesh on the device specified by the + `mesh`'s `device` attribute. + ray_starts: The input ray start positions. Shape is (B, N, 3). + ray_directions: The input ray directions. Shape is (B, N, 3). + ray_hits: The output ray hit positions. Shape is (B, N, 3). + ray_distance: The output ray hit distances. Shape is (B, N,), if ``return_distance`` is True. Otherwise, + this array is not used. + ray_normal: The output ray hit normals. Shape is (B, N, 3), if ``return_normal`` is True. Otherwise, + this array is not used. + ray_face_id: The output ray hit face ids. Shape is (B, N,), if ``return_face_id`` is True. Otherwise, + this array is not used. + ray_mesh_id: The output ray hit mesh ids. Shape is (B, N,), if ``return_mesh_id`` is True. Otherwise, + this array is not used. + max_dist: The maximum ray-cast distance. Defaults to 1e6. + return_normal: Whether to return the ray hit normals. Defaults to False`. + return_face_id: Whether to return the ray hit face ids. Defaults to False. + return_mesh_id: Whether to return the mesh id. Defaults to False. + """ + # get the thread id + tid_mesh_id, tid_env, tid_ray = wp.tid() + + direction = ray_directions[tid_env, tid_ray] + start_pos = ray_starts[tid_env, tid_ray] + + # ray cast against the mesh and store the hit position + mesh_query_ray_t = wp.mesh_query_ray(mesh[tid_env, tid_mesh_id], start_pos, direction, max_dist) + + # if the ray hit, store the hit data + if mesh_query_ray_t.result: + + wp.atomic_min(ray_distance, tid_env, tid_ray, mesh_query_ray_t.t) + # check if hit distance is less than the current hit distance, only then update the memory + # TODO, in theory we could use the output of atomic_min to avoid the non-thread safe next comparison + # however, warp atomic_min is returning the wrong values on gpu currently. + # FIXME https://github.com/NVIDIA/warp/issues/1058 + if mesh_query_ray_t.t == ray_distance[tid_env, tid_ray]: + # convert back to world space and update the hit data + ray_hits[tid_env, tid_ray] = start_pos + mesh_query_ray_t.t * direction + + # update the normal and face id if requested + if return_normal == 1: + ray_normal[tid_env, tid_ray] = mesh_query_ray_t.normal + if return_face_id == 1: + ray_face_id[tid_env, tid_ray] = mesh_query_ray_t.face + if return_mesh_id == 1: + ray_mesh_id[tid_env, tid_ray] = wp.int16(tid_mesh_id) + + +@wp.kernel(enable_backward=False) +def raycast_dynamic_meshes_kernel( + mesh: wp.array2d(dtype=wp.uint64), + ray_starts: wp.array2d(dtype=wp.vec3), + ray_directions: wp.array2d(dtype=wp.vec3), + ray_hits: wp.array2d(dtype=wp.vec3), + ray_distance: wp.array2d(dtype=wp.float32), + ray_normal: wp.array2d(dtype=wp.vec3), + ray_face_id: wp.array2d(dtype=wp.int32), + ray_mesh_id: wp.array2d(dtype=wp.int16), + mesh_positions: wp.array2d(dtype=wp.vec3), + mesh_rotations: wp.array2d(dtype=wp.quat), + max_dist: float = 1e6, + return_normal: int = False, + return_face_id: int = False, + return_mesh_id: int = False, +): + """Performs ray-casting against multiple meshes. + + This function performs ray-casting against the given meshes using the provided ray start positions + and directions. The resulting ray hit positions are stored in the :obj:`ray_hits` array. + + The function utilizes the ``mesh_query_ray`` method from the ``wp`` module to perform the actual ray-casting + operation. The maximum ray-cast distance is set to ``1e6`` units. + + + Note: + That the ``ray_starts``, ``ray_directions``, and ``ray_hits`` arrays should have compatible shapes + and data types to ensure proper execution. Additionally, they all must be in the same frame. + + All arguments are expected to be batched with the first dimension (B, batch) being the number of envs + and the second dimension (N, num_rays) being the number of rays. For Meshes, W is the number of meshes. + + Args: + mesh: The input mesh. The ray-casting is performed against this mesh on the device specified by the + `mesh`'s `device` attribute. + ray_starts: The input ray start positions. Shape is (B, N, 3). + ray_directions: The input ray directions. Shape is (B, N, 3). + ray_hits: The output ray hit positions. Shape is (B, N, 3). + ray_distance: The output ray hit distances. Shape is (B, N,), if ``return_distance`` is True. Otherwise, + this array is not used. + ray_normal: The output ray hit normals. Shape is (B, N, 3), if ``return_normal`` is True. Otherwise, + this array is not used. + ray_face_id: The output ray hit face ids. Shape is (B, N,), if ``return_face_id`` is True. Otherwise, + this array is not used. + ray_mesh_id: The output ray hit mesh ids. Shape is (B, N,), if ``return_mesh_id`` is True. Otherwise, + this array is not used. + mesh_positions: The input mesh positions in world frame. Shape is (W, 3). + mesh_rotations: The input mesh rotations in world frame. Shape is (W, 4). + max_dist: The maximum ray-cast distance. Defaults to 1e6. + return_normal: Whether to return the ray hit normals. Defaults to False`. + return_face_id: Whether to return the ray hit face ids. Defaults to False. + return_mesh_id: Whether to return the mesh id. Defaults to False. + """ + # get the thread id + tid_mesh_id, tid_env, tid_ray = wp.tid() + + mesh_pose = wp.transform(mesh_positions[tid_env, tid_mesh_id], mesh_rotations[tid_env, tid_mesh_id]) + mesh_pose_inv = wp.transform_inverse(mesh_pose) + direction = wp.transform_vector(mesh_pose_inv, ray_directions[tid_env, tid_ray]) + start_pos = wp.transform_point(mesh_pose_inv, ray_starts[tid_env, tid_ray]) + + # ray cast against the mesh and store the hit position + mesh_query_ray_t = wp.mesh_query_ray(mesh[tid_env, tid_mesh_id], start_pos, direction, max_dist) + # if the ray hit, store the hit data + if mesh_query_ray_t.result: + + wp.atomic_min(ray_distance, tid_env, tid_ray, mesh_query_ray_t.t) + # check if hit distance is less than the current hit distance, only then update the memory + # TODO, in theory we could use the output of atomic_min to avoid the non-thread safe next comparison + # however, warp atomic_min is returning the wrong values on gpu currently. + # FIXME https://github.com/NVIDIA/warp/issues/1058 + if mesh_query_ray_t.t == ray_distance[tid_env, tid_ray]: + # convert back to world space and update the hit data + hit_pos = start_pos + mesh_query_ray_t.t * direction + ray_hits[tid_env, tid_ray] = wp.transform_point(mesh_pose, hit_pos) + + # update the normal and face id if requested + if return_normal == 1: + n = wp.transform_vector(mesh_pose, mesh_query_ray_t.normal) + ray_normal[tid_env, tid_ray] = n + if return_face_id == 1: + ray_face_id[tid_env, tid_ray] = mesh_query_ray_t.face + if return_mesh_id == 1: + ray_mesh_id[tid_env, tid_ray] = wp.int16(tid_mesh_id) + + @wp.kernel(enable_backward=False) def reshape_tiled_image( tiled_image_buffer: Any, diff --git a/source/isaaclab/isaaclab/utils/warp/ops.py b/source/isaaclab/isaaclab/utils/warp/ops.py index a2db46c4b52..cb58e51043f 100644 --- a/source/isaaclab/isaaclab/utils/warp/ops.py +++ b/source/isaaclab/isaaclab/utils/warp/ops.py @@ -18,6 +18,8 @@ # initialize the warp module wp.init() +from isaaclab.utils.math import convert_quat + from . import kernels @@ -127,6 +129,257 @@ def raycast_mesh( return ray_hits.to(device).view(shape), ray_distance, ray_normal, ray_face_id +def raycast_single_mesh( + ray_starts: torch.Tensor, + ray_directions: torch.Tensor, + mesh_id: int, + max_dist: float = 1e6, + return_distance: bool = False, + return_normal: bool = False, + return_face_id: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Performs ray-casting against a mesh. + + Note that the :attr:`ray_starts` and :attr:`ray_directions`, and :attr:`ray_hits` should have compatible shapes + and data types to ensure proper execution. Additionally, they all must be in the same frame. + + Args: + ray_starts: The starting position of the rays. Shape (B, N, 3). + ray_directions: The ray directions for each ray. Shape (B, N, 3). + mesh_id: The warp mesh id to ray-cast against. + max_dist: The maximum distance to ray-cast. Defaults to 1e6. + return_distance: Whether to return the distance of the ray until it hits the mesh. Defaults to False. + return_normal: Whether to return the normal of the mesh face the ray hits. Defaults to False. + return_face_id: Whether to return the face id of the mesh face the ray hits. Defaults to False. + + Returns: + The ray hit position. Shape (B, N, 3). + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit distance. Shape (B, N,). + Will only return if :attr:`return_distance` is True, else returns None. + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit normal. Shape (B, N, 3). + Will only return if :attr:`return_normal` is True else returns None. + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit face id. Shape (B, N,). + Will only return if :attr:`return_face_id` is True else returns None. + The returned tensor contains :obj:`int(-1)` for missed hits. + """ + # cast mesh id into array + mesh_ids = wp.array2d( + [[mesh_id] for _ in range(ray_starts.shape[0])], dtype=wp.uint64, device=wp.device_from_torch(ray_starts.device) + ) + ray_hits, ray_distance, ray_normal, ray_face_id, ray_mesh_id = raycast_dynamic_meshes( + ray_starts=ray_starts, + ray_directions=ray_directions, + mesh_ids_wp=mesh_ids, + max_dist=max_dist, + return_distance=return_distance, + return_normal=return_normal, + return_face_id=return_face_id, + ) + + return ray_hits, ray_distance, ray_normal, ray_face_id + + +def raycast_dynamic_meshes( + ray_starts: torch.Tensor, + ray_directions: torch.Tensor, + mesh_ids_wp: wp.Array, + mesh_positions_w: torch.Tensor | None = None, + mesh_orientations_w: torch.Tensor | None = None, + max_dist: float = 1e6, + return_distance: bool = False, + return_normal: bool = False, + return_face_id: bool = False, + return_mesh_id: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Performs ray-casting against multiple, dynamic meshes. + + Note that the :attr:`ray_starts` and :attr:`ray_directions`, and :attr:`ray_hits` should have compatible shapes + and data types to ensure proper execution. Additionally, they all must be in the same frame. + + If mesh positions and rotations are provided, they need to have to have the same shape as the + number of meshes. + + Args: + ray_starts: The starting position of the rays. Shape (B, N, 3). + ray_directions: The ray directions for each ray. Shape (B, N, 3). + mesh_ids_wp: The warp mesh ids to ray-cast against. Length (B, M). + mesh_positions_w: The world positions of the meshes. Shape (B, M, 3). + mesh_orientations_w: The world orientation as quaternion (wxyz) format. Shape (B, M, 4). + max_dist: The maximum distance to ray-cast. Defaults to 1e6. + return_distance: Whether to return the distance of the ray until it hits the mesh. Defaults to False. + return_normal: Whether to return the normal of the mesh face the ray hits. Defaults to False. + return_face_id: Whether to return the face id of the mesh face the ray hits. Defaults to False. + return_mesh_id: Whether to return the mesh id of the mesh face the ray hits. Defaults to False. + NOTE: the type of the returned tensor is torch.int16, so you can't have more than 32767 meshes. + + Returns: + The ray hit position. Shape (B, N, 3). + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit distance. Shape (B, N,). + Will only return if :attr:`return_distance` is True, else returns None. + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit normal. Shape (B, N, 3). + Will only return if :attr:`return_normal` is True else returns None. + The returned tensor contains :obj:`float('inf')` for missed hits. + The ray hit face id. Shape (B, N,). + Will only return if :attr:`return_face_id` is True else returns None. + The returned tensor contains :obj:`int(-1)` for missed hits. + The ray hit mesh id. Shape (B, N,). + Will only return if :attr:`return_mesh_id` is True else returns None. + The returned tensor contains :obj:`-1` for missed hits. + """ + # extract device and shape information + shape = ray_starts.shape + device = ray_starts.device + + # device of the mesh + torch_device = wp.device_to_torch(mesh_ids_wp.device) + n_meshes = mesh_ids_wp.shape[1] + + n_envs = ray_starts.shape[0] + n_rays_per_env = ray_starts.shape[1] + + # reshape the tensors + ray_starts = ray_starts.to(torch_device).view(n_envs, n_rays_per_env, 3).contiguous() + ray_directions = ray_directions.to(torch_device).view(n_envs, n_rays_per_env, 3).contiguous() + + # create output tensor for the ray hits + ray_hits = torch.full((n_envs, n_rays_per_env, 3), float("inf"), device=torch_device).contiguous() + + # map the memory to warp arrays + ray_starts_wp = wp.from_torch(ray_starts, dtype=wp.vec3) + ray_directions_wp = wp.from_torch(ray_directions, dtype=wp.vec3) + ray_hits_wp = wp.from_torch(ray_hits, dtype=wp.vec3) + # required to check if a closer hit is reported, returned only if return_distance is true + ray_distance = torch.full( + ( + n_envs, + n_rays_per_env, + ), + float("inf"), + device=torch_device, + ).contiguous() + ray_distance_wp = wp.from_torch(ray_distance, dtype=wp.float32) + + if return_normal: + ray_normal = torch.full((n_envs, n_rays_per_env, 3), float("inf"), device=torch_device).contiguous() + ray_normal_wp = wp.from_torch(ray_normal, dtype=wp.vec3) + else: + ray_normal = None + ray_normal_wp = wp.empty((1, 1), dtype=wp.vec3, device=torch_device) + + if return_face_id: + ray_face_id = torch.ones( + ( + n_envs, + n_rays_per_env, + ), + dtype=torch.int32, + device=torch_device, + ).contiguous() * (-1) + ray_face_id_wp = wp.from_torch(ray_face_id, dtype=wp.int32) + else: + ray_face_id = None + ray_face_id_wp = wp.empty((1, 1), dtype=wp.int32, device=torch_device) + + if return_mesh_id: + ray_mesh_id = -torch.ones((n_envs, n_rays_per_env), dtype=torch.int16, device=torch_device).contiguous() + ray_mesh_id_wp = wp.from_torch(ray_mesh_id, dtype=wp.int16) + else: + ray_mesh_id = None + ray_mesh_id_wp = wp.empty((1, 1), dtype=wp.int16, device=torch_device) + + ## + # Call the warp kernels + ### + if mesh_positions_w is None and mesh_orientations_w is None: + # Static mesh case, no need to pass in positions and rotations. + # launch the warp kernel + wp.launch( + kernel=kernels.raycast_static_meshes_kernel, + dim=[n_meshes, n_envs, n_rays_per_env], + inputs=[ + mesh_ids_wp, + ray_starts_wp, + ray_directions_wp, + ray_hits_wp, + ray_distance_wp, + ray_normal_wp, + ray_face_id_wp, + ray_mesh_id_wp, + float(max_dist), + int(return_normal), + int(return_face_id), + int(return_mesh_id), + ], + device=torch_device, + ) + else: + # dynamic mesh case + if mesh_positions_w is None: + mesh_positions_wp_w = wp.zeros((n_envs, n_meshes), dtype=wp.vec3, device=torch_device) + else: + mesh_positions_w = mesh_positions_w.to(torch_device).view(n_envs, n_meshes, 3).contiguous() + mesh_positions_wp_w = wp.from_torch(mesh_positions_w, dtype=wp.vec3) + + if mesh_orientations_w is None: + # Note (zrene): This is a little bit ugly, since it requires to initialize torch memory first + # But I couldn't find a better way to initialize a quaternion identity in warp + # wp.zeros(1, dtype=wp.quat, device=torch_device) gives all zero quaternion + quat_identity = torch.tensor([0, 0, 0, 1], dtype=torch.float32, device=torch_device).repeat( + n_envs, n_meshes, 1 + ) + mesh_quat_wp_w = wp.from_torch(quat_identity, dtype=wp.quat) + else: + mesh_orientations_w = convert_quat( + mesh_orientations_w.to(dtype=torch.float32, device=torch_device), "xyzw" + ).contiguous() + mesh_quat_wp_w = wp.from_torch(mesh_orientations_w, dtype=wp.quat) + + # launch the warp kernel + wp.launch( + kernel=kernels.raycast_dynamic_meshes_kernel, + dim=[n_meshes, n_envs, n_rays_per_env], + inputs=[ + mesh_ids_wp, + ray_starts_wp, + ray_directions_wp, + ray_hits_wp, + ray_distance_wp, + ray_normal_wp, + ray_face_id_wp, + ray_mesh_id_wp, + mesh_positions_wp_w, + mesh_quat_wp_w, + float(max_dist), + int(return_normal), + int(return_face_id), + int(return_mesh_id), + ], + device=torch_device, + ) + ## + # Cleanup and convert back to torch tensors + ## + + # NOTE: Synchronize is not needed anymore, but we keep it for now. Check with @dhoeller. + wp.synchronize() + + if return_distance: + ray_distance = ray_distance.to(device).view(shape[:2]) + if return_normal: + ray_normal = ray_normal.to(device).view(shape) + if return_face_id: + ray_face_id = ray_face_id.to(device).view(shape[:2]) + if return_mesh_id: + ray_mesh_id = ray_mesh_id.to(device).view(shape[:2]) + + return ray_hits.to(device).view(shape), ray_distance, ray_normal, ray_face_id, ray_mesh_id + + def convert_to_warp_mesh(points: np.ndarray, indices: np.ndarray, device: str) -> wp.Mesh: """Create a warp mesh object with a mesh defined from vertices and triangles. diff --git a/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py new file mode 100644 index 00000000000..5f54edadb60 --- /dev/null +++ b/source/isaaclab/test/sensors/check_multi_mesh_ray_caster.py @@ -0,0 +1,214 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +""" +This script shows how to use the ray caster from the Isaac Lab framework. + +.. code-block:: bash + + # Usage + ./isaaclab.sh -p source/extensions/omni.isaac.lab/test/sensors/check_multi_mesh_ray_caster.py --headless +""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Ray Caster Test Script") +parser.add_argument("--num_envs", type=int, default=16, help="Number of environments to clone.") +parser.add_argument("--num_objects", type=int, default=0, help="Number of additional objects to clone.") +parser.add_argument( + "--terrain_type", + type=str, + default="generator", + help="Type of terrain to import. Can be 'generator' or 'usd' or 'plane'.", +) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + + +"""Rest everything follows.""" + +import random +import torch + +import isaacsim.core.utils.prims as prim_utils +from isaacsim.core.api.simulation_context import SimulationContext +from isaacsim.core.cloner import GridCloner +from isaacsim.core.prims import RigidPrim +from isaacsim.core.utils.viewports import set_camera_view + +import isaaclab.sim as sim_utils +import isaaclab.terrains as terrain_gen +from isaaclab.sensors.ray_caster import MultiMeshRayCaster, MultiMeshRayCasterCfg, patterns +from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG +from isaaclab.terrains.terrain_importer import TerrainImporter +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.math import quat_from_euler_xyz +from isaaclab.utils.timer import Timer + + +def design_scene(sim: SimulationContext, num_envs: int = 2048): + """Design the scene.""" + # Create interface to clone the scene + cloner = GridCloner(spacing=10.0) + cloner.define_base_env("/World/envs") + # Everything under the namespace "/World/envs/env_0" will be cloned + prim_utils.define_prim("/World/envs/env_0") + # Define the scene + # -- Light + cfg = sim_utils.DistantLightCfg(intensity=2000) + cfg.func("/World/light", cfg) + # -- Balls + cfg = sim_utils.SphereCfg( + radius=0.25, + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.5), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 0.0, 1.0)), + ) + cfg.func("/World/envs/env_0/ball", cfg, translation=(0.0, 0.0, 5.0)) + + for i in range(args_cli.num_objects): + object = sim_utils.CuboidCfg( + size=(0.5 + random.random() * 0.5, 0.5 + random.random() * 0.5, 0.1 + random.random() * 0.05), + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + mass_props=sim_utils.MassPropertiesCfg(mass=0.5), + collision_props=sim_utils.CollisionPropertiesCfg(), + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=(0.0 + i / args_cli.num_objects, 0.0, 1.0 - i / args_cli.num_objects) + ), + ) + object.func( + 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(), + ) + + # Clone the scene + cloner.define_base_env("/World/envs") + envs_prim_paths = cloner.generate_paths("/World/envs/env", num_paths=num_envs) + cloner.clone(source_prim_path="/World/envs/env_0", prim_paths=envs_prim_paths, replicate_physics=True) + physics_scene_path = sim.get_physics_context().prim_path + cloner.filter_collisions( + physics_scene_path, "/World/collisions", prim_paths=envs_prim_paths, global_paths=["/World/ground"] + ) + + +def main(): + """Main function.""" + + # Load kit helper + sim_params = { + "use_gpu": True, + "use_gpu_pipeline": True, + "use_flatcache": True, # deprecated from Isaac Sim 2023.1 onwards + "use_fabric": True, # used from Isaac Sim 2023.1 onwards + "enable_scene_query_support": True, + } + sim = SimulationContext( + physics_dt=1.0 / 60.0, rendering_dt=1.0 / 60.0, sim_params=sim_params, backend="torch", device="cuda:0" + ) + # Set main camera + set_camera_view([0.0, 30.0, 25.0], [0.0, 0.0, -2.5]) + + # Parameters + num_envs = args_cli.num_envs + # Design the scene + design_scene(sim=sim, num_envs=num_envs) + # Handler for terrains importing + terrain_importer_cfg = terrain_gen.TerrainImporterCfg( + prim_path="/World/ground", + terrain_type=args_cli.terrain_type, + terrain_generator=ROUGH_TERRAINS_CFG, + usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd", + max_init_terrain_level=0, + num_envs=1, + ) + _ = TerrainImporter(terrain_importer_cfg) + + mesh_targets: list[MultiMeshRayCasterCfg.RaycastTargetCfg] = [ + MultiMeshRayCasterCfg.RaycastTargetCfg(target_prim_expr="/World/ground", track_mesh_transforms=False), + ] + if args_cli.num_objects != 0: + mesh_targets.append( + MultiMeshRayCasterCfg.RaycastTargetCfg( + target_prim_expr="/World/envs/env_.*/object_.*", track_mesh_transforms=True + ) + ) + # Create a ray-caster sensor + ray_caster_cfg = MultiMeshRayCasterCfg( + prim_path="/World/envs/env_.*/ball", + mesh_prim_paths=mesh_targets, + pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=(1.6, 1.0)), + attach_yaw_only=True, + debug_vis=not args_cli.headless, + ) + ray_caster = MultiMeshRayCaster(cfg=ray_caster_cfg) + # Create a view over all the balls + ball_view = RigidPrim("/World/envs/env_.*/ball", reset_xform_properties=False) + + # Play simulator + sim.reset() + + # Initialize the views + # -- balls + ball_view.initialize() + # Print the sensor information + print(ray_caster) + + # Get the initial positions of the balls + ball_initial_positions, ball_initial_orientations = ball_view.get_world_poses() + ball_initial_velocities = ball_view.get_velocities() + + # Create a counter for resetting the scene + step_count = 0 + # Simulate physics + while simulation_app.is_running(): + # If simulation is stopped, then exit. + if sim.is_stopped(): + break + # If simulation is paused, then skip. + if not sim.is_playing(): + sim.step(render=False) + continue + # Reset the scene + if step_count % 500 == 0: + # sample random indices to reset + reset_indices = torch.randint(0, num_envs, (num_envs // 2,)) + # reset the balls + ball_view.set_world_poses( + ball_initial_positions[reset_indices], ball_initial_orientations[reset_indices], indices=reset_indices + ) + ball_view.set_velocities(ball_initial_velocities[reset_indices], indices=reset_indices) + # reset the sensor + ray_caster.reset(reset_indices) + # reset the counter + step_count = 0 + # Step simulation + sim.step() + # Update the ray-caster + with Timer(f"Ray-caster update with {num_envs} x {ray_caster.num_rays} rays"): + ray_caster.update(dt=sim.get_physics_dt(), force_recompute=True) + # Update counter + step_count += 1 + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster.py b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster.py new file mode 100644 index 00000000000..5d6434970e0 --- /dev/null +++ b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster.py @@ -0,0 +1,251 @@ +# Copyright (c) 2022-2025, 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.app import AppLauncher + +# launch omniverse app. Used for warp. +app_launcher = AppLauncher(headless=True) + +import numpy as np +import torch +import trimesh + +import pytest +import warp as wp + +from isaaclab.utils.math import matrix_from_quat, quat_from_euler_xyz, random_orientation +from isaaclab.utils.warp.ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_single_mesh + + +@pytest.fixture(scope="module") +def device(): + return "cuda" if torch.cuda.is_available() else "cpu" + + +@pytest.fixture +def rays(device): + ray_starts = torch.tensor([[0, -0.35, -5], [0.25, 0.35, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_directions = torch.tensor([[0, 0, 1], [0, 0, 1]], dtype=torch.float32, device=device).unsqueeze(0) + expected_ray_hits = torch.tensor( + [[0, -0.35, -0.5], [0.25, 0.35, -0.5]], dtype=torch.float32, device=device + ).unsqueeze(0) + return ray_starts, ray_directions, expected_ray_hits + + +@pytest.fixture +def trimesh_box(): + return trimesh.creation.box([2, 2, 1]) + + +@pytest.fixture +def single_mesh(trimesh_box, device): + wp_mesh = convert_to_warp_mesh(trimesh_box.vertices, trimesh_box.faces, device) + return wp_mesh, wp_mesh.id + + +def test_raycast_multi_cubes(device, trimesh_box, rays): + """Test raycasting against two cubes.""" + ray_starts, ray_directions, _ = rays + + trimesh_1 = trimesh_box.copy() + wp_mesh_1 = convert_to_warp_mesh(trimesh_1.vertices, trimesh_1.faces, device) + + translation = np.eye(4) + translation[:3, 3] = [0, 2, 0] + trimesh_2 = trimesh_box.copy().apply_transform(translation) + wp_mesh_2 = convert_to_warp_mesh(trimesh_2.vertices, trimesh_2.faces, device) + + # get mesh id array + mesh_ids_wp = wp.array2d([[wp_mesh_1.id, wp_mesh_2.id]], dtype=wp.uint64, device=device) + + # Static positions (no transforms passed) + ray_start = torch.tensor([[0, 0, -5], [0, 2.5, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_ids = raycast_dynamic_meshes( + ray_start, + ray_directions, + mesh_ids_wp, + return_distance=True, + return_normal=True, + return_face_id=True, + return_mesh_id=True, + ) + + torch.testing.assert_close( + ray_hits, torch.tensor([[[0, 0, -0.5], [0, 2.5, -0.5]]], dtype=torch.float32, device=device) + ) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device)) + assert torch.equal(mesh_ids, torch.tensor([[0, 1]], dtype=torch.int32, device=device)) + + # Dynamic positions/orientations + ray_start = torch.tensor([[0, 0, -5], [0, 4.5, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_ids = raycast_dynamic_meshes( + ray_start, + ray_directions, + mesh_ids_wp, + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_positions_w=torch.tensor([[[0, 0, 0], [0, 2, 0]]], dtype=torch.float32, device=device), + mesh_orientations_w=torch.tensor([[[1, 0, 0, 0], [1, 0, 0, 0]]], dtype=torch.float32, device=device), + return_mesh_id=True, + ) + + torch.testing.assert_close( + ray_hits, torch.tensor([[[0, 0, -0.5], [0, 4.5, -0.5]]], dtype=torch.float32, device=device) + ) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device)) + assert torch.equal(mesh_ids, torch.tensor([[0, 1]], dtype=torch.int32, device=device)) + + +def test_raycast_single_cube(device, single_mesh, rays): + """Test raycasting against a single cube.""" + ray_starts, ray_directions, expected_ray_hits = rays + _, single_mesh_id = single_mesh + + ray_hits, ray_distance, ray_normal, ray_face_id = raycast_single_mesh( + ray_starts, + ray_directions, + single_mesh_id, + return_distance=True, + return_normal=True, + return_face_id=True, + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close( + ray_normal, + torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device), + ) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + # check multiple meshes implementation + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + +@pytest.mark.parametrize("num_samples", [10]) +def test_raycast_moving_cube(device, single_mesh, rays, num_samples): + r"""Test raycasting against a single cube with different distances. + |-------------| + |\ | + | \ | + | \ 8 | + | \ | + | \ x_1 | + | \ | + | \ | + | \ | + | \ | + | \ | + | 3 x_2 \ | + | \ | + | \| + |-------------| + + """ + ray_starts, ray_directions, expected_ray_hits = rays + _, single_mesh_id = single_mesh + + # move the cube along the z axis + for distance in torch.linspace(0, 1, num_samples, device=device): + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_id = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + return_mesh_id=True, + mesh_positions_w=torch.tensor([[0, 0, distance]], dtype=torch.float32, device=device), + ) + torch.testing.assert_close( + ray_hits, + expected_ray_hits + + torch.tensor([[0, 0, distance], [0, 0, distance]], dtype=torch.float32, device=device).unsqueeze(0), + ) + torch.testing.assert_close( + ray_distance, distance + torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device) + ) + torch.testing.assert_close( + ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device) + ) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + +def test_raycast_rotated_cube(device, single_mesh, rays): + """Test raycasting against a single cube with different 90deg. orientations.""" + ray_starts, ray_directions, expected_ray_hits = rays + _, single_mesh_id = single_mesh + + cube_rotation = quat_from_euler_xyz(torch.tensor([0.0]), torch.tensor([0.0]), torch.tensor([np.pi])).to(device) + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_orientations_w=cube_rotation.unsqueeze(0), + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], dtype=torch.float32, device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], dtype=torch.float32, device=device)) + # Make sure the face ids are correct. The cube is rotated by 90deg. so the face ids are different. + torch.testing.assert_close(ray_face_id, torch.tensor([[8, 3]], dtype=torch.int32, device=device)) + + +@pytest.mark.parametrize("num_random", [10]) +def test_raycast_random_cube(device, trimesh_box, single_mesh, rays, num_random): + """Test raycasting against a single cube with random poses.""" + ray_starts, ray_directions, _ = rays + _, single_mesh_id = single_mesh + + for orientation in random_orientation(num_random, device): + pos = torch.tensor([[0, 0, torch.rand(1)]], dtype=torch.float32, device=device) + tf_hom = np.eye(4) + tf_hom[:3, :3] = matrix_from_quat(orientation).cpu().numpy() + tf_hom[:3, 3] = pos.cpu().numpy() + tf_mesh = trimesh_box.copy().apply_transform(tf_hom) + + # get raycast for transformed, static mesh + wp_mesh = convert_to_warp_mesh(tf_mesh.vertices, tf_mesh.faces, device) + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[wp_mesh.id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + ) + # get raycast for modified mesh + ray_hits_m, ray_distance_m, ray_normal_m, ray_face_id_m, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_positions_w=pos, + mesh_orientations_w=orientation.view(1, 1, -1), + ) + torch.testing.assert_close(ray_hits, ray_hits_m) + torch.testing.assert_close(ray_distance, ray_distance_m) + torch.testing.assert_close(ray_normal, ray_normal_m) + torch.testing.assert_close(ray_face_id, ray_face_id_m) 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 new file mode 100644 index 00000000000..a68084caf45 --- /dev/null +++ b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py @@ -0,0 +1,863 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# ignore private usage of variables warning +# pyright: reportPrivateUsage=none + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True, enable_cameras=True).app + +"""Rest everything follows.""" + +import copy +import numpy as np +import os +import torch + +import isaacsim.core.utils.prims as prim_utils +import isaacsim.core.utils.stage as stage_utils +import omni.replicator.core as rep +import pytest +from pxr import Gf + +import isaaclab.sim as sim_utils +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.utils import convert_dict_to_backend +from isaaclab.utils.timer import Timer + +# sample camera poses +POSITION = [2.5, 2.5, 2.5] +QUAT_ROS = [-0.17591989, 0.33985114, 0.82047325, -0.42470819] +QUAT_OPENGL = [0.33985113, 0.17591988, 0.42470818, 0.82047324] +QUAT_WORLD = [-0.3647052, -0.27984815, -0.1159169, 0.88047623] + + +@pytest.fixture(scope="function") +def setup_simulation(): + """Fixture to set up and tear down the simulation environment.""" + # Create a new stage + stage_utils.create_new_stage() + # Simulation time-step + dt = 0.01 + # Load kit helper + sim_cfg = sim_utils.SimulationCfg(dt=dt) + sim: sim_utils.SimulationContext = sim_utils.SimulationContext(sim_cfg) + # Ground-plane + mesh = make_plane(size=(100, 100), height=0.0, center_zero=True) + create_prim_from_mesh("/World/defaultGroundPlane", mesh) + # load stage + stage_utils.update_stage() + + camera_cfg = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), + debug_vis=False, + pattern_cfg=patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=480, + width=640, + ), + data_types=["distance_to_image_plane"], + ) + + # create xform because placement of camera directly under world is not supported + prim_utils.create_prim("/World/Camera", "Xform") + + yield sim, dt, camera_cfg + + # Cleanup + # close all the opened viewport from before. + rep.vp_manager.destroy_hydra_textures("Replicator") + # stop simulation + # note: cannot use self.sim.stop() since it does one render step after stopping!! This doesn't make sense :( + sim._timeline.stop() + # clear the stage + sim.clear_all_callbacks() + sim.clear_instance() + + +@pytest.mark.parametrize( + "convention,quat", + [ + ("ros", QUAT_ROS), + ("opengl", QUAT_OPENGL), + ("world", QUAT_WORLD), + ], +) +@pytest.mark.isaacsim_ci +def test_camera_init_offset(setup_simulation, convention, quat): + """Test camera initialization with offset using different conventions.""" + sim, dt, camera_cfg = setup_simulation + + # Create camera config with specific convention + cam_cfg_offset = copy.deepcopy(camera_cfg) + cam_cfg_offset.offset = MultiMeshRayCasterCameraCfg.OffsetCfg( + pos=POSITION, + rot=quat, + convention=convention, + ) + prim_utils.create_prim(f"/World/CameraOffset{convention.capitalize()}", "Xform") + cam_cfg_offset.prim_path = f"/World/CameraOffset{convention.capitalize()}" + + camera = MultiMeshRayCasterCamera(cam_cfg_offset) + + # play sim + sim.reset() + + # update camera + 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) + + del camera + + +@pytest.mark.isaacsim_ci +def test_camera_init(setup_simulation): + """Test camera initialization.""" + sim, dt, camera_cfg = setup_simulation + + # Create camera + camera = MultiMeshRayCasterCamera(cfg=camera_cfg) + # Play sim + sim.reset() + # Check if camera is initialized + assert camera.is_initialized + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + # Check buffers that exists and have correct shapes + assert camera.data.pos_w.shape == (1, 3) + assert camera.data.quat_w_ros.shape == (1, 4) + assert camera.data.quat_w_world.shape == (1, 4) + assert camera.data.quat_w_opengl.shape == (1, 4) + assert camera.data.intrinsic_matrices.shape == (1, 3, 3) + assert camera.data.image_shape == (camera_cfg.pattern_cfg.height, camera_cfg.pattern_cfg.width) + assert camera.data.info == [{camera_cfg.data_types[0]: None}] + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update camera + camera.update(dt) + # check image data + for im_data in camera.data.output.values(): + assert im_data.shape == (1, camera_cfg.pattern_cfg.height, camera_cfg.pattern_cfg.width, 1) + + del camera + + +@pytest.mark.isaacsim_ci +def test_camera_resolution(setup_simulation): + """Test camera resolution is correctly set.""" + sim, dt, camera_cfg = setup_simulation + + # Create camera + camera = MultiMeshRayCasterCamera(cfg=camera_cfg) + # Play sim + sim.reset() + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + camera.update(dt) + # access image data and compare shapes + for im_data in camera.data.output.values(): + assert im_data.shape == (1, camera_cfg.pattern_cfg.height, camera_cfg.pattern_cfg.width, 1) + + del camera + + +@pytest.mark.isaacsim_ci +def test_camera_init_intrinsic_matrix(setup_simulation): + """Test camera initialization from intrinsic matrix.""" + sim, dt, camera_cfg = setup_simulation + + # get the first camera + camera_1 = MultiMeshRayCasterCamera(cfg=camera_cfg) + # get intrinsic matrix + sim.reset() + intrinsic_matrix = camera_1.data.intrinsic_matrices[0].cpu().flatten().tolist() + + # initialize from intrinsic matrix + intrinsic_camera_cfg = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), + debug_vis=False, + pattern_cfg=patterns.PinholeCameraPatternCfg.from_intrinsic_matrix( + intrinsic_matrix=intrinsic_matrix, + height=camera_cfg.pattern_cfg.height, + width=camera_cfg.pattern_cfg.width, + focal_length=camera_cfg.pattern_cfg.focal_length, + ), + data_types=["distance_to_image_plane"], + ) + camera_2 = MultiMeshRayCasterCamera(cfg=intrinsic_camera_cfg) + + # play sim + sim.reset() + sim.play() + + # update cameras + camera_1.update(dt) + camera_2.update(dt) + + # check image data + torch.testing.assert_close( + camera_1.data.output["distance_to_image_plane"], + camera_2.data.output["distance_to_image_plane"], + ) + # check that both intrinsic matrices are the same + torch.testing.assert_close( + camera_1.data.intrinsic_matrices[0], + camera_2.data.intrinsic_matrices[0], + ) + + del camera_1, camera_2 + + +@pytest.mark.isaacsim_ci +def test_multi_camera_init(setup_simulation): + """Test multi-camera initialization.""" + sim, dt, camera_cfg = setup_simulation + + # -- camera 1 + cam_cfg_1 = copy.deepcopy(camera_cfg) + cam_cfg_1.prim_path = "/World/Camera_0" + prim_utils.create_prim("/World/Camera_0", "Xform") + # Create camera + cam_1 = MultiMeshRayCasterCamera(cam_cfg_1) + + # -- camera 2 + cam_cfg_2 = copy.deepcopy(camera_cfg) + cam_cfg_2.prim_path = "/World/Camera_1" + prim_utils.create_prim("/World/Camera_1", "Xform") + # Create camera + cam_2 = MultiMeshRayCasterCamera(cam_cfg_2) + + # play sim + sim.reset() + + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update camera + cam_1.update(dt) + cam_2.update(dt) + # check image data + for cam in [cam_1, cam_2]: + for im_data in cam.data.output.values(): + assert im_data.shape == (1, camera_cfg.pattern_cfg.height, camera_cfg.pattern_cfg.width, 1) + + del cam_1, cam_2 + + +@pytest.mark.isaacsim_ci +def test_camera_set_world_poses(setup_simulation): + """Test camera function to set specific world pose.""" + sim, dt, camera_cfg = setup_simulation + + camera = MultiMeshRayCasterCamera(camera_cfg) + # play sim + sim.reset() + + # convert to torch tensors + position = torch.tensor([POSITION], dtype=torch.float32, device=camera.device) + orientation = torch.tensor([QUAT_WORLD], dtype=torch.float32, device=camera.device) + # set new pose + camera.set_world_poses(position.clone(), orientation.clone(), convention="world") + + # check if transform correctly set in output + torch.testing.assert_close(camera.data.pos_w, position) + torch.testing.assert_close(camera.data.quat_w_world, orientation) + + del camera + + +@pytest.mark.isaacsim_ci +def test_camera_set_world_poses_from_view(setup_simulation): + """Test camera function to set specific world pose from view.""" + sim, dt, camera_cfg = setup_simulation + + camera = MultiMeshRayCasterCamera(camera_cfg) + # play sim + sim.reset() + + # convert to torch tensors + eyes = torch.tensor([POSITION], dtype=torch.float32, device=camera.device) + targets = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=camera.device) + quat_ros_gt = torch.tensor([QUAT_ROS], dtype=torch.float32, device=camera.device) + # set new pose + camera.set_world_poses_from_view(eyes.clone(), targets.clone()) + + # check if transform correctly set in output + torch.testing.assert_close(camera.data.pos_w, eyes) + torch.testing.assert_close(camera.data.quat_w_ros, quat_ros_gt) + + del camera + + +@pytest.mark.parametrize("height,width", [(240, 320), (480, 640)]) +@pytest.mark.isaacsim_ci +def test_intrinsic_matrix(setup_simulation, height, width): + """Checks that the camera's set and retrieve methods work for intrinsic matrix.""" + sim, dt, camera_cfg = setup_simulation + + camera_cfg_copy = copy.deepcopy(camera_cfg) + camera_cfg_copy.pattern_cfg.height = height + camera_cfg_copy.pattern_cfg.width = width + camera = MultiMeshRayCasterCamera(camera_cfg_copy) + # play sim + sim.reset() + # Desired properties (obtained from realsense camera at 320x240 resolution) + rs_intrinsic_matrix = [229.31640625, 0.0, 164.810546875, 0.0, 229.826171875, 122.1650390625, 0.0, 0.0, 1.0] + rs_intrinsic_matrix = torch.tensor(rs_intrinsic_matrix, device=camera.device).reshape(3, 3).unsqueeze(0) + # Set matrix into simulator + camera.set_intrinsic_matrices(rs_intrinsic_matrix.clone()) + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + # Simulate physics + for _ in range(10): + # perform rendering + sim.step() + # update camera + camera.update(dt) + # Check that matrix is correct + torch.testing.assert_close(rs_intrinsic_matrix, camera.data.intrinsic_matrices) + + del camera + + +@pytest.mark.isaacsim_ci +def test_throughput(setup_simulation): + """Test camera throughput for different image sizes.""" + sim, dt, camera_cfg = setup_simulation + + # Create directory temp dir to dump the results + file_dir = os.path.dirname(os.path.realpath(__file__)) + temp_dir = os.path.join(file_dir, "output", "camera", "throughput") + os.makedirs(temp_dir, exist_ok=True) + # Create replicator writer + rep_writer = rep.BasicWriter(output_dir=temp_dir, frame_padding=3) + # create camera + camera_cfg_copy = copy.deepcopy(camera_cfg) + camera_cfg_copy.pattern_cfg.height = 480 + camera_cfg_copy.pattern_cfg.width = 640 + camera = MultiMeshRayCasterCamera(camera_cfg_copy) + + # Play simulator + sim.reset() + + # Set camera pose + eyes = torch.tensor([[2.5, 2.5, 2.5]], dtype=torch.float32, device=camera.device) + targets = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=camera.device) + camera.set_world_poses_from_view(eyes, targets) + + # Simulate for a few steps + # note: This is a workaround to ensure that the textures are loaded. + # Check "Known Issues" section in the documentation for more details. + for _ in range(5): + sim.step() + # Simulate physics + for _ in range(5): + # perform rendering + sim.step() + # update camera + with Timer(f"Time taken for updating camera with shape {camera.image_shape}"): + camera.update(dt) + # Save images + with Timer(f"Time taken for writing data with shape {camera.image_shape} "): + # Pack data back into replicator format to save them using its writer + rep_output = {"annotators": {}} + camera_data = convert_dict_to_backend(camera.data.output, backend="numpy") + for key, data, info in zip(camera_data.keys(), camera_data.values(), camera.data.info[0].values()): + if info is not None: + rep_output["annotators"][key] = {"render_product": {"data": data, **info}} + else: + rep_output["annotators"][key] = {"render_product": {"data": data}} + # Save images + rep_output["trigger_outputs"] = {"on_time": camera.frame[0]} + rep_writer.write(rep_output) + print("----------------------------------------") + # Check image data + for im_data in camera.data.output.values(): + assert im_data.shape == (1, camera_cfg_copy.pattern_cfg.height, camera_cfg_copy.pattern_cfg.width, 1) + + del camera + + +@pytest.mark.parametrize( + "data_types", + [ + ["distance_to_image_plane", "distance_to_camera", "normals"], + ["distance_to_image_plane"], + ["distance_to_camera"], + ], +) +@pytest.mark.isaacsim_ci +def test_output_equal_to_usdcamera(setup_simulation, data_types): + """Test that ray caster camera output equals USD camera output.""" + sim, dt, camera_cfg = setup_simulation + + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=240, + width=320, + ) + prim_utils.create_prim("/World/Camera_warp", "Xform") + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera_warp", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0)), + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + data_types=data_types, + ) + + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + # create usd camera + camera_cfg_usd = CameraCfg( + height=240, + width=320, + prim_path="/World/Camera_usd", + update_period=0, + data_types=data_types, + spawn=PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(1e-4, 1.0e5) + ), + ) + camera_usd = Camera(camera_cfg_usd) + + # play sim + sim.reset() + sim.play() + + # convert to torch tensors + eyes = torch.tensor([[2.5, 2.5, 4.5]], dtype=torch.float32, device=camera_warp.device) + targets = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=camera_warp.device) + # set views + camera_warp.set_world_poses_from_view(eyes, targets) + camera_usd.set_world_poses_from_view(eyes, targets) + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # check the intrinsic matrices + torch.testing.assert_close( + camera_usd.data.intrinsic_matrices, + camera_warp.data.intrinsic_matrices, + ) + + # check the apertures + torch.testing.assert_close( + camera_usd._sensor_prims[0].GetHorizontalApertureAttr().Get(), + camera_cfg_warp.pattern_cfg.horizontal_aperture, + ) + + # check image data + for data_type in 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], + atol=5e-5, + rtol=5e-6, + ) + elif data_type == "normals": + # NOTE: floating point issues of ~1e-5, so using atol and rtol in this case + torch.testing.assert_close( + camera_usd.data.output[data_type][..., :3], + camera_warp.data.output[data_type], + rtol=1e-5, + atol=1e-4, + ) + else: + torch.testing.assert_close( + camera_usd.data.output[data_type], + camera_warp.data.output[data_type], + ) + + 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.""" + sim, dt, camera_cfg = setup_simulation + offset_rot = (-0.1251, 0.3617, 0.8731, -0.3020) + + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=240, + width=320, + ) + prim_utils.create_prim("/World/Camera_warp", "Xform") + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera_warp", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(2.5, 2.5, 4.0), rot=offset_rot, convention="ros"), + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + data_types=["distance_to_image_plane", "distance_to_camera", "normals"], + ) + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + # create usd camera + camera_cfg_usd = CameraCfg( + height=240, + width=320, + prim_path="/World/Camera_usd", + update_period=0, + data_types=["distance_to_image_plane", "distance_to_camera", "normals"], + spawn=PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(1e-6, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(2.5, 2.5, 4.0), rot=offset_rot, convention="ros"), + ) + camera_usd = Camera(camera_cfg_usd) + + # play sim + sim.reset() + sim.play() + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # check image data + torch.testing.assert_close( + camera_usd.data.output["distance_to_image_plane"], + camera_warp.data.output["distance_to_image_plane"], + atol=5e-5, + rtol=5e-6, + ) + torch.testing.assert_close( + camera_usd.data.output["distance_to_camera"], + camera_warp.data.output["distance_to_camera"], + atol=5e-5, + rtol=5e-6, + ) + + # check normals + # NOTE: floating point issues of ~1e-5, so using atol and rtol in this case + torch.testing.assert_close( + camera_usd.data.output["normals"][..., :3], + camera_warp.data.output["normals"], + rtol=1e-5, + atol=1e-4, + ) + + del camera_usd, camera_warp + + +@pytest.mark.isaacsim_ci +def test_output_equal_to_usdcamera_prim_offset(setup_simulation): + """Test that the output of the ray caster camera is equal to the output of the usd camera when both are placed + under an XForm prim that is translated and rotated from the world origin.""" + sim, dt, camera_cfg = setup_simulation + + offset_rot = [-0.1251, 0.3617, 0.8731, -0.3020] + + # gf quat + gf_quatf = Gf.Quatd() + gf_quatf.SetReal(QUAT_OPENGL[0]) + gf_quatf.SetImaginary(tuple(QUAT_OPENGL[1:])) + + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=240, + width=320, + ) + prim_raycast_cam = prim_utils.create_prim("/World/Camera_warp", "Xform") + prim_raycast_cam.GetAttribute("xformOp:translate").Set(tuple(POSITION)) + prim_raycast_cam.GetAttribute("xformOp:orient").Set(gf_quatf) + + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera_warp", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0, 0, 2.0), rot=offset_rot, convention="ros"), + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + data_types=["distance_to_image_plane", "distance_to_camera", "normals"], + ) + + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + # create usd camera + camera_cfg_usd = CameraCfg( + height=240, + width=320, + prim_path="/World/Camera_usd/camera", + update_period=0, + data_types=["distance_to_image_plane", "distance_to_camera", "normals"], + spawn=PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(1e-6, 1.0e5) + ), + offset=CameraCfg.OffsetCfg(pos=(0, 0, 2.0), rot=offset_rot, convention="ros"), + update_latest_camera_pose=True, + ) + prim_usd = prim_utils.create_prim("/World/Camera_usd", "Xform") + prim_usd.GetAttribute("xformOp:translate").Set(tuple(POSITION)) + prim_usd.GetAttribute("xformOp:orient").Set(gf_quatf) + + camera_usd = Camera(camera_cfg_usd) + + # play sim + sim.reset() + sim.play() + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # check if pos and orientation are correct + torch.testing.assert_close(camera_warp.data.pos_w[0], camera_usd.data.pos_w[0]) + torch.testing.assert_close(camera_warp.data.quat_w_ros[0], camera_usd.data.quat_w_ros[0]) + + # check image data + torch.testing.assert_close( + camera_usd.data.output["distance_to_image_plane"], + camera_warp.data.output["distance_to_image_plane"], + atol=5e-5, + rtol=5e-6, + ) + torch.testing.assert_close( + camera_usd.data.output["distance_to_camera"], + camera_warp.data.output["distance_to_camera"], + rtol=4e-6, + atol=2e-5, + ) + + # check normals + # NOTE: floating point issues of ~1e-5, so using atol and rtol in this case + torch.testing.assert_close( + camera_usd.data.output["normals"][..., :3], + camera_warp.data.output["normals"], + rtol=1e-5, + atol=1e-4, + ) + + del camera_usd, camera_warp + + +@pytest.mark.parametrize("height,width", [(540, 960), (240, 320)]) +@pytest.mark.isaacsim_ci +def test_output_equal_to_usd_camera_intrinsics(setup_simulation, height, width): + """Test that the output of the ray caster camera and usd camera are the same when both are + initialized with the same intrinsic matrix.""" + sim, dt, camera_cfg = setup_simulation + + # create cameras + offset_rot = [-0.1251, 0.3617, 0.8731, -0.3020] + offset_pos = (2.5, 2.5, 4.0) + intrinsics = [380.0831, 0.0, width / 2, 0.0, 380.0831, height / 2, 0.0, 0.0, 1.0] + prim_utils.create_prim("/World/Camera_warp", "Xform") + # get camera cfgs + camera_warp_cfg = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera_warp", + mesh_prim_paths=["/World/defaultGroundPlane"], + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=offset_pos, rot=offset_rot, convention="ros"), + debug_vis=False, + pattern_cfg=patterns.PinholeCameraPatternCfg.from_intrinsic_matrix( + intrinsic_matrix=intrinsics, + height=height, + width=width, + focal_length=38.0, + ), + max_distance=25.0, + data_types=["distance_to_image_plane"], + ) + camera_usd_cfg = CameraCfg( + prim_path="/World/Camera_usd", + offset=CameraCfg.OffsetCfg(pos=offset_pos, rot=offset_rot, convention="ros"), + spawn=PinholeCameraCfg.from_intrinsic_matrix( + intrinsic_matrix=intrinsics, + height=height, + width=width, + clipping_range=(0.01, 25), + focal_length=38.0, + ), + height=height, + width=width, + data_types=["distance_to_image_plane"], + ) + + # set aperture offsets to 0, as currently not supported for usd camera + camera_warp_cfg.pattern_cfg.horizontal_aperture_offset = 0 + camera_warp_cfg.pattern_cfg.vertical_aperture_offset = 0 + camera_usd_cfg.spawn.horizontal_aperture_offset = 0 + camera_usd_cfg.spawn.vertical_aperture_offset = 0 + # init cameras + camera_warp = MultiMeshRayCasterCamera(camera_warp_cfg) + camera_usd = Camera(camera_usd_cfg) + + # play sim + sim.reset() + sim.play() + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + 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[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]) + + # check the apertures + torch.testing.assert_close( + camera_usd._sensor_prims[0].GetHorizontalApertureAttr().Get(), + camera_warp_cfg.pattern_cfg.horizontal_aperture, + ) + torch.testing.assert_close( + camera_usd._sensor_prims[0].GetVerticalApertureAttr().Get(), + camera_warp_cfg.pattern_cfg.vertical_aperture, + ) + + # check image data + torch.testing.assert_close( + cam_warp_output, + cam_usd_output, + atol=5e-5, + rtol=5e-6, + ) + + del camera_usd, camera_warp + + +@pytest.mark.isaacsim_ci +def test_output_equal_to_usd_camera_when_intrinsics_set(setup_simulation): + """Test that the output of the ray caster camera is equal to the output of the usd camera when both are placed + under an XForm prim and an intrinsic matrix is set.""" + sim, dt, camera_cfg = setup_simulation + + camera_pattern_cfg = patterns.PinholeCameraPatternCfg( + focal_length=24.0, + horizontal_aperture=20.955, + height=540, + width=960, + ) + camera_cfg_warp = MultiMeshRayCasterCameraCfg( + prim_path="/World/Camera", + mesh_prim_paths=["/World/defaultGroundPlane"], + update_period=0, + offset=MultiMeshRayCasterCameraCfg.OffsetCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0)), + debug_vis=False, + pattern_cfg=camera_pattern_cfg, + data_types=["distance_to_camera"], + ) + + camera_warp = MultiMeshRayCasterCamera(camera_cfg_warp) + + # create usd camera + camera_cfg_usd = CameraCfg( + height=540, + width=960, + prim_path="/World/Camera_usd", + update_period=0, + data_types=["distance_to_camera"], + spawn=PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(1e-4, 1.0e5) + ), + ) + camera_usd = Camera(camera_cfg_usd) + + # play sim + sim.reset() + sim.play() + + # set intrinsic matrix + # NOTE: extend the test to cover aperture offsets once supported by the usd camera + intrinsic_matrix = torch.tensor( + [[380.0831, 0.0, camera_cfg_usd.width / 2, 0.0, 380.0831, camera_cfg_usd.height / 2, 0.0, 0.0, 1.0]], + device=camera_warp.device, + ).reshape(1, 3, 3) + camera_warp.set_intrinsic_matrices(intrinsic_matrix, focal_length=10) + camera_usd.set_intrinsic_matrices(intrinsic_matrix, focal_length=10) + + # set camera position + camera_warp.set_world_poses_from_view( + eyes=torch.tensor([[0.0, 0.0, 5.0]], device=camera_warp.device), + targets=torch.tensor([[0.0, 0.0, 0.0]], device=camera_warp.device), + ) + camera_usd.set_world_poses_from_view( + eyes=torch.tensor([[0.0, 0.0, 5.0]], device=camera_usd.device), + targets=torch.tensor([[0.0, 0.0, 0.0]], device=camera_usd.device), + ) + + # perform steps + for _ in range(5): + sim.step() + + # update camera + camera_usd.update(dt) + camera_warp.update(dt) + + # check image data + torch.testing.assert_close( + camera_usd.data.output["distance_to_camera"], + camera_warp.data.output["distance_to_camera"], + rtol=5e-3, + atol=1e-4, + ) + + del camera_usd, camera_warp diff --git a/source/isaaclab/test/sensors/test_ray_caster.py b/source/isaaclab/test/sensors/test_ray_caster.py new file mode 100644 index 00000000000..c8daa468e4d --- /dev/null +++ b/source/isaaclab/test/sensors/test_ray_caster.py @@ -0,0 +1,242 @@ +# Copyright (c) 2022-2025, 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 numpy as np +import torch +import trimesh + +import pytest + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True, enable_cameras=True).app + +# Import after app launch +import warp as wp + +from isaaclab.utils.math import matrix_from_quat, quat_from_euler_xyz, random_orientation +from isaaclab.utils.warp.ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh + + +@pytest.fixture(scope="module") +def raycast_setup(): + device = "cuda" if torch.cuda.is_available() else "cpu" + # Base trimesh cube and its Warp conversion + trimesh_mesh = trimesh.creation.box([2, 2, 1]) + single_mesh = [ + convert_to_warp_mesh( + trimesh_mesh.vertices, + trimesh_mesh.faces, + device, + ) + ] + single_mesh_id = single_mesh[0].id + + # Rays + ray_starts = torch.tensor([[0, -0.35, -5], [0.25, 0.35, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_directions = torch.tensor([[0, 0, 1], [0, 0, 1]], dtype=torch.float32, device=device).unsqueeze(0) + expected_ray_hits = torch.tensor( + [[0, -0.35, -0.5], [0.25, 0.35, -0.5]], dtype=torch.float32, device=device + ).unsqueeze(0) + + return { + "device": device, + "trimesh_mesh": trimesh_mesh, + "single_mesh_id": single_mesh_id, + "wp_mesh": single_mesh[0], + "ray_starts": ray_starts, + "ray_directions": ray_directions, + "expected_ray_hits": expected_ray_hits, + } + + +def test_raycast_multi_cubes(raycast_setup): + device = raycast_setup["device"] + base_tm = raycast_setup["trimesh_mesh"] + + tm1 = base_tm.copy() + wp_mesh_1 = convert_to_warp_mesh(tm1.vertices, tm1.faces, device) + + translation = np.eye(4) + translation[:3, 3] = [0, 2, 0] + tm2 = base_tm.copy().apply_transform(translation) + wp_mesh_2 = convert_to_warp_mesh(tm2.vertices, tm2.faces, device) + + mesh_ids_wp = wp.array2d([[wp_mesh_1.id, wp_mesh_2.id]], dtype=wp.uint64, device=device) + + ray_directions = raycast_setup["ray_directions"] + + # Case 1 + ray_start = torch.tensor([[0, 0, -5], [0, 2.5, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_ids = raycast_dynamic_meshes( + ray_start, + ray_directions, + mesh_ids_wp, + return_distance=True, + return_normal=True, + return_face_id=True, + return_mesh_id=True, + ) + + torch.testing.assert_close(ray_hits, torch.tensor([[[0, 0, -0.5], [0, 2.5, -0.5]]], device=device)) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + assert torch.equal(mesh_ids, torch.tensor([[0, 1]], dtype=torch.int32, device=device)) + + # Case 2 (explicit poses/orientations) + ray_start = torch.tensor([[0, 0, -5], [0, 4.5, -5]], dtype=torch.float32, device=device).unsqueeze(0) + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_ids = raycast_dynamic_meshes( + ray_start, + ray_directions, + mesh_ids_wp, + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_positions_w=torch.tensor([[[0, 0, 0], [0, 2, 0]]], dtype=torch.float32, device=device), + mesh_orientations_w=torch.tensor([[[1, 0, 0, 0], [1, 0, 0, 0]]], dtype=torch.float32, device=device), + return_mesh_id=True, + ) + + torch.testing.assert_close(ray_hits, torch.tensor([[[0, 0, -0.5], [0, 4.5, -0.5]]], device=device)) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + assert torch.equal(mesh_ids, torch.tensor([[0, 1]], dtype=torch.int32, device=device)) + + +def test_raycast_single_cube(raycast_setup): + device = raycast_setup["device"] + ray_starts = raycast_setup["ray_starts"] + ray_directions = raycast_setup["ray_directions"] + mesh = raycast_setup["wp_mesh"] + expected_ray_hits = raycast_setup["expected_ray_hits"] + single_mesh_id = raycast_setup["single_mesh_id"] + + # Single-mesh helper + ray_hits, ray_distance, ray_normal, ray_face_id = raycast_mesh( + ray_starts, + ray_directions, + mesh, + return_distance=True, + return_normal=True, + return_face_id=True, + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + # Multi-mesh API with one mesh + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + +def test_raycast_moving_cube(raycast_setup): + device = raycast_setup["device"] + ray_starts = raycast_setup["ray_starts"] + ray_directions = raycast_setup["ray_directions"] + single_mesh_id = raycast_setup["single_mesh_id"] + expected_ray_hits = raycast_setup["expected_ray_hits"] + + for distance in torch.linspace(0, 1, 10, device=device): + ray_hits, ray_distance, ray_normal, ray_face_id, mesh_id = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + return_mesh_id=True, + mesh_positions_w=torch.tensor([[0, 0, distance.item()]], dtype=torch.float32, device=device), + ) + offset = torch.tensor([[0, 0, distance.item()], [0, 0, distance.item()]], dtype=torch.float32, device=device) + torch.testing.assert_close(ray_hits, expected_ray_hits + offset.unsqueeze(0)) + torch.testing.assert_close(ray_distance, distance + torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close( + ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32) + ) + torch.testing.assert_close(ray_face_id, torch.tensor([[3, 8]], dtype=torch.int32, device=device)) + + +def test_raycast_rotated_cube(raycast_setup): + device = raycast_setup["device"] + ray_starts = raycast_setup["ray_starts"] + ray_directions = raycast_setup["ray_directions"] + single_mesh_id = raycast_setup["single_mesh_id"] + expected_ray_hits = raycast_setup["expected_ray_hits"] + + cube_rotation = quat_from_euler_xyz( + torch.tensor([0.0], device=device), torch.tensor([0.0], device=device), torch.tensor([np.pi], device=device) + ) + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_orientations_w=cube_rotation.unsqueeze(0), + ) + torch.testing.assert_close(ray_hits, expected_ray_hits) + torch.testing.assert_close(ray_distance, torch.tensor([[4.5, 4.5]], device=device)) + torch.testing.assert_close(ray_normal, torch.tensor([[[0, 0, -1], [0, 0, -1]]], device=device, dtype=torch.float32)) + # Rotated cube swaps face IDs + torch.testing.assert_close(ray_face_id, torch.tensor([[8, 3]], dtype=torch.int32, device=device)) + + +def test_raycast_random_cube(raycast_setup): + device = raycast_setup["device"] + base_tm = raycast_setup["trimesh_mesh"] + ray_starts = raycast_setup["ray_starts"] + ray_directions = raycast_setup["ray_directions"] + single_mesh_id = raycast_setup["single_mesh_id"] + + for orientation in random_orientation(10, device): + pos = torch.tensor([[0.0, 0.0, torch.rand(1, device=device).item()]], dtype=torch.float32, device=device) + + tf_hom = np.eye(4) + tf_hom[:3, :3] = matrix_from_quat(orientation).cpu().numpy() + tf_hom[:3, 3] = pos.squeeze(0).cpu().numpy() + + tf_mesh = base_tm.copy().apply_transform(tf_hom) + wp_mesh = convert_to_warp_mesh(tf_mesh.vertices, tf_mesh.faces, device) + + # Raycast transformed, static mesh + ray_hits, ray_distance, ray_normal, ray_face_id, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[wp_mesh.id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + ) + # Raycast original mesh with pose provided + ray_hits_m, ray_distance_m, ray_normal_m, ray_face_id_m, _ = raycast_dynamic_meshes( + ray_starts, + ray_directions, + wp.array2d([[single_mesh_id]], dtype=wp.uint64, device=device), + return_distance=True, + return_normal=True, + return_face_id=True, + mesh_positions_w=pos, + mesh_orientations_w=orientation.view(1, 1, -1), + ) + + torch.testing.assert_close(ray_hits, ray_hits_m) + torch.testing.assert_close(ray_distance, ray_distance_m) + torch.testing.assert_close(ray_normal, ray_normal_m) + torch.testing.assert_close(ray_face_id, ray_face_id_m)