diff --git a/.gitattributes b/.gitattributes index e3c0ead689d..aec401902b9 100644 --- a/.gitattributes +++ b/.gitattributes @@ -12,3 +12,4 @@ *.hdf5 filter=lfs diff=lfs merge=lfs -text *.bat text eol=crlf +docs/make.bat text eol=crlf diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 00000000000..199bc8361fc --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/robot_model"] + path = source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/robot_model + url = https://github.com/ntnu-arl/robot_model.git diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index ebfc174f0a7..fd064f7939a 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -67,6 +67,7 @@ Guidelines for modifications: * Felix Yu * Gary Lvov * Giulio Romualdi +* Grzegorz Malczyk * Haoran Zhou * Harsh Patel * HoJin Jeon @@ -100,6 +101,7 @@ Guidelines for modifications: * Michael Noseworthy * Michael Lin * Miguel Alonso Jr +* Mihir Kulkarni * Mingyu Lee * Muhong Guo * Narendra Dahile @@ -138,6 +140,7 @@ Guidelines for modifications: * Virgilio Gómez Lambo * Vladimir Fokow * Wei Yang +* Welf Rehberg * Xavier Nal * Xinjie Yao * Xinpeng Liu 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/actuators/__init__.py b/source/isaaclab/isaaclab/actuators/__init__.py index 5ccf5d7b082..a8ef77e7fe1 100644 --- a/source/isaaclab/isaaclab/actuators/__init__.py +++ b/source/isaaclab/isaaclab/actuators/__init__.py @@ -32,6 +32,8 @@ IdealPDActuatorCfg, ImplicitActuatorCfg, RemotizedPDActuatorCfg, + ThrusterCfg, ) from .actuator_net import ActuatorNetLSTM, ActuatorNetMLP from .actuator_pd import DCMotor, DelayedPDActuator, IdealPDActuator, ImplicitActuator, RemotizedPDActuator +from .thruster import Thruster diff --git a/source/isaaclab/isaaclab/actuators/actuator_cfg.py b/source/isaaclab/isaaclab/actuators/actuator_cfg.py index e5351e4fa63..54a525610b9 100644 --- a/source/isaaclab/isaaclab/actuators/actuator_cfg.py +++ b/source/isaaclab/isaaclab/actuators/actuator_cfg.py @@ -11,6 +11,7 @@ from . import actuator_net, actuator_pd from .actuator_base import ActuatorBase +from .thruster import Thruster @configclass @@ -283,3 +284,50 @@ class RemotizedPDActuatorCfg(DelayedPDActuatorCfg): This tensor describes the relationship between the joint angle (rad), the transmission ratio (in/out), and the output torque (N*m). The table is used to interpolate the output torque based on the joint angle. """ + + +@configclass +class ThrusterCfg: + """Configuration for thruster actuator groups. + + This config defines per-actuator-group parameters used by the low-level + thruster/motor models (time-constants, thrust ranges, integration scheme, + and initial state specifications). Fields left as ``MISSING`` are required + and must be provided by the user configuration. + """ + + class_type: type[Thruster] = Thruster + """Concrete Python class that consumes this config.""" + + dt: float = MISSING + """Simulation/integration timestep used by the thruster update [s].""" + + thrust_range: tuple[float, float] = MISSING + """Per-motor thrust clamp range [N]: values are clipped to this interval.""" + + max_thrust_rate: float = 100000.0 + """Per-motor thrust slew-rate limit applied inside the first-order model [N/s].""" + + thrust_const_range: tuple[float, float] = MISSING + """Range for thrust coefficient :math:`k_f` [N/(rps²)].""" + + tau_inc_range: tuple[float, float] = MISSING + """Range of time constants when commanded output is **increasing** (rise dynamics) [s].""" + + tau_dec_range: tuple[float, float] = MISSING + """Range of time constants when commanded output is **decreasing** (fall dynamics) [s].""" + + torque_to_thrust_ratio: float = MISSING + """Yaw-moment coefficient converting thrust to motor torque about +Z [N·m per N]. + Used as ``tau_z = torque_to_thrust_ratio * thrust_z * direction``. + """ + + use_discrete_approximation: bool = True + """If ``True``, use discrete mixing factor ``1/(dt + tau)``; if ``False``, use continuous ``1/tau``.""" + + integration_scheme: Literal["rk4", "euler"] = "rk4" + """Numerical integrator for the first-order model. Choose ``"euler"`` or ``"rk4"``.""" + + thruster_names_expr: list[str] = MISSING + + """Articulation's joint names that are part of the group.""" diff --git a/source/isaaclab/isaaclab/actuators/thruster.py b/source/isaaclab/isaaclab/actuators/thruster.py new file mode 100644 index 00000000000..fad660eb9f5 --- /dev/null +++ b/source/isaaclab/isaaclab/actuators/thruster.py @@ -0,0 +1,227 @@ +# 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.types import ArticulationThrustActions + +if TYPE_CHECKING: + from .actuator_cfg import ThrusterCfg + + +class Thruster: + """Low-level motor/thruster dynamics with separate rise/fall time constants. + + Integration scheme is Euler or RK4. All internal buffers are shaped (num_envs, num_motors). + Units: thrust [N], rates [N/s], time [s]. + """ + + computed_thrust: torch.Tensor + """The computed thrust for the actuator group. Shape is (num_envs, num_thrusters).""" + + applied_thrust: torch.Tensor + """The applied thrust for the actuator group. Shape is (num_envs, num_thrusters). + + This is the thrust obtained after clipping the :attr:`computed_thrust` based on the + actuator characteristics. + """ + + cfg: ThrusterCfg + + def __init__( + self, + cfg: ThrusterCfg, + thruster_names: list[str], + thruster_ids: slice | torch.Tensor, + num_envs: int, + device: str, + init_thruster_rps: torch.Tensor, + ): + """Construct buffers and sample per-motor parameters. + + Args: + cfg: Thruster configuration. + thruster_names: List of thruster names belonging to this group. + thruster_ids: Slice or tensor of indices into the articulation thruster array. + num_envs: Number of parallel/vectorized environments. + device: PyTorch device string or device identifier. + init_thruster_rps: Initial per-thruster rotations-per-second tensor used when + the configuration uses RPM-based thrust modelling. + """ + self.cfg = cfg + self._num_envs = num_envs + self._device = device + self._thruster_names = thruster_names + self._thruster_indices = thruster_ids + self._init_thruster_rps = init_thruster_rps + + # Range tensors, shaped (num_envs, 2, num_motors); [:,0,:]=min, [:,1,:]=max + self.num_motors = len(thruster_names) + self.thrust_r = torch.tensor(cfg.thrust_range).to(self._device) + self.tau_inc_r = torch.tensor(cfg.tau_inc_range).to(self._device) + self.tau_dec_r = torch.tensor(cfg.tau_dec_range).to(self._device) + + self.max_rate = torch.tensor(cfg.max_thrust_rate).expand(self._num_envs, self.num_motors).to(self._device) + + self.max_thrust = self.cfg.thrust_range[1] + self.min_thrust = self.cfg.thrust_range[0] + + # State & randomized per-motor parameters + self.tau_inc_s = math_utils.sample_uniform(*self.tau_inc_r, (self._num_envs, self.num_motors), self._device) + self.tau_dec_s = math_utils.sample_uniform(*self.tau_dec_r, (self._num_envs, self.num_motors), self._device) + + self.thrust_const_r = torch.tensor(cfg.thrust_const_range).to(device) + self.thrust_const = math_utils.sample_uniform( + *self.thrust_const_r, (self._num_envs, self.num_motors), self._device + ) + + self.curr_thrust = ( + torch.ones(self._num_envs, self.num_motors, device=self._device, dtype=torch.float32) + * self.thrust_const + * self._init_thruster_rps**2 + ) + + # Mixing factor (discrete vs continuous form) + if self.cfg.use_discrete_approximation: + self.mixing_factor_function = self.discrete_mixing_factor + else: + self.mixing_factor_function = self.continuous_mixing_factor + + # Choose stepping kernel once (avoids per-step branching) + if self.cfg.integration_scheme == "euler": + self._step_thrust = self.compute_thrust_with_rpm_time_constant + elif self.cfg.integration_scheme == "rk4": + self._step_thrust = self.compute_thrust_with_rpm_time_constant_rk4 + else: + raise ValueError("integration scheme unknown") + + @property + def num_thrusters(self) -> int: + """Number of actuators in the group.""" + return len(self._thruster_names) + + @property + def thruster_names(self) -> list[str]: + """Articulation's thruster names that are part of the group.""" + return self._thruster_names + + @property + def thruster_indices(self) -> slice | torch.Tensor: + """Articulation's thruster indices that are part of the group. + + Note: + If :obj:`slice(None)` is returned, then the group contains all the thrusters in the articulation. + We do this to avoid unnecessary indexing of the thrusters for performance reasons. + """ + return self._thruster_indices + + def compute(self, control_action: ArticulationThrustActions) -> ArticulationThrustActions: + """Advance the thruster state one step. + + Applies saturation, chooses rise/fall tau per motor, computes mixing factor, + and integrates with the selected kernel. + + Args: + control_action: (num_envs, num_motors) commanded per-motor thrust [N]. + + Returns: + (num_envs, num_motors) updated thrust state [N]. + + """ + des_thrust = control_action.thrusts + des_thrust = torch.clamp(des_thrust, *self.thrust_r) + + thrust_decrease_mask = torch.sign(self.curr_thrust) * torch.sign(des_thrust - self.curr_thrust) + motor_tau = torch.where(thrust_decrease_mask < 0, self.tau_dec_s, self.tau_inc_s) + mixing = self.mixing_factor_function(motor_tau) + + self.curr_thrust[:] = self._step_thrust(des_thrust, self.curr_thrust, mixing) + + self.computed_thrust = self.curr_thrust + self.applied_thrust = torch.clamp(self.computed_thrust, self.min_thrust, self.max_thrust) + + control_action.thrusts = self.applied_thrust + + return control_action + + def reset_idx(self, env_ids=None) -> None: + """Re-sample parameters and reinitialize state. + + Args: + env_ids: Env indices to reset. If ``None``, resets all envs. + """ + if env_ids is None: + env_ids = slice(None) + + if isinstance(env_ids, slice): + num_resets = self._num_envs + else: + num_resets = len(env_ids) + + self.tau_inc_s[env_ids] = math_utils.sample_uniform( + *self.tau_inc_r, + (num_resets, self.num_motors), + self._device, + ) + self.tau_dec_s[env_ids] = math_utils.sample_uniform( + *self.tau_dec_r, + (num_resets, self.num_motors), + self._device, + ) + self.thrust_const[env_ids] = math_utils.sample_uniform( + *self.thrust_const_r, + (num_resets, self.num_motors), + self._device, + ) + self.curr_thrust[env_ids] = self.thrust_const[env_ids] * self._init_thruster_rps[env_ids] ** 2 + + def reset(self, env_ids: Sequence[int]) -> None: + """Reset all envs.""" + self.reset_idx(env_ids) + + def motor_model_rate(self, error: torch.Tensor, mixing_factor: torch.Tensor): + return torch.clamp(mixing_factor * (error), -self.max_rate, self.max_rate) + + def rk4_integration(self, error: torch.Tensor, mixing_factor: torch.Tensor): + k1 = self.motor_model_rate(error, mixing_factor) + k2 = self.motor_model_rate(error + 0.5 * self.cfg.dt * k1, mixing_factor) + k3 = self.motor_model_rate(error + 0.5 * self.cfg.dt * k2, mixing_factor) + k4 = self.motor_model_rate(error + self.cfg.dt * k3, mixing_factor) + return (self.cfg.dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + + def discrete_mixing_factor(self, time_constant: torch.Tensor): + return 1.0 / (self.cfg.dt + time_constant) + + def continuous_mixing_factor(self, time_constant: torch.Tensor): + return 1.0 / time_constant + + def compute_thrust_with_rpm_time_constant( + self, + des_thrust: torch.Tensor, + curr_thrust: torch.Tensor, + mixing_factor: torch.Tensor, + ): + current_rpm = torch.sqrt(curr_thrust / self.thrust_const) + desired_rpm = torch.sqrt(des_thrust / self.thrust_const) + rpm_error = desired_rpm - current_rpm + current_rpm += self.motor_model_rate(rpm_error, mixing_factor) * self.cfg.dt + return self.thrust_const * current_rpm**2 + + def compute_thrust_with_rpm_time_constant_rk4( + self, + des_thrust: torch.Tensor, + curr_thrust: torch.Tensor, + mixing_factor: torch.Tensor, + ) -> torch.Tensor: + current_rpm = torch.sqrt(curr_thrust / self.thrust_const) + desired_rpm = torch.sqrt(des_thrust / self.thrust_const) + rpm_error = desired_rpm - current_rpm + current_rpm += self.rk4_integration(rpm_error, mixing_factor) + return self.thrust_const * current_rpm**2 diff --git a/source/isaaclab/isaaclab/assets/__init__.py b/source/isaaclab/isaaclab/assets/__init__.py index 206e5dd9c5c..a5597a9c0b1 100644 --- a/source/isaaclab/isaaclab/assets/__init__.py +++ b/source/isaaclab/isaaclab/assets/__init__.py @@ -38,7 +38,7 @@ the corresponding actuator torques. """ -from .articulation import Articulation, ArticulationCfg, ArticulationData +from .articulation import Articulation, ArticulationCfg, ArticulationData, Multirotor, MultirotorCfg, MultirotorData from .asset_base import AssetBase from .asset_base_cfg import AssetBaseCfg from .deformable_object import DeformableObject, DeformableObjectCfg, DeformableObjectData diff --git a/source/isaaclab/isaaclab/assets/articulation/__init__.py b/source/isaaclab/isaaclab/assets/articulation/__init__.py index 9429f0fec31..d0caee0693e 100644 --- a/source/isaaclab/isaaclab/assets/articulation/__init__.py +++ b/source/isaaclab/isaaclab/assets/articulation/__init__.py @@ -8,3 +8,6 @@ from .articulation import Articulation from .articulation_cfg import ArticulationCfg from .articulation_data import ArticulationData +from .multirotor import Multirotor +from .multirotor_cfg import MultirotorCfg +from .multirotor_data import MultirotorData diff --git a/source/isaaclab/isaaclab/assets/articulation/multirotor.py b/source/isaaclab/isaaclab/assets/articulation/multirotor.py new file mode 100644 index 00000000000..eaaeaae3a12 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation/multirotor.py @@ -0,0 +1,382 @@ +# 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 + +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import omni.log + +import isaaclab.utils.string as string_utils +from isaaclab.actuators import Thruster +from isaaclab.assets.articulation.multirotor_data import MultirotorData +from isaaclab.utils.types import ArticulationThrustActions + +from .articulation import Articulation + +if TYPE_CHECKING: + from .multirotor_cfg import MultirotorCfg + + +class Multirotor(Articulation): + """A multirotor articulation asset class. + + This class extends the base articulation class to support multirotor vehicles + with thruster actuators that can be applied at specific body locations. + """ + + cfg: MultirotorCfg + """Configuration instance for the multirotor.""" + + actuators: dict[str, Thruster] + """Dictionary of thruster actuator instances for the multirotor. + + The keys are the actuator names and the values are the actuator instances. The actuator instances + are initialized based on the actuator configurations specified in the :attr:`MultirotorCfg.actuators` + attribute. They are used to compute the thruster commands during the :meth:`write_data_to_sim` function. + """ + + def __init__(self, cfg: MultirotorCfg): + """Initialize the multirotor articulation. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + + """ + Properties + """ + + @property + def thruster_names(self) -> list[str]: + """Ordered names of thrusters in the multirotor.""" + if not hasattr(self, "actuators") or not self.actuators: + return [] + + thruster_names = [] + for actuator in self.actuators.values(): + if hasattr(actuator, "thruster_names"): + thruster_names.extend(actuator.thruster_names) + else: + raise ValueError("Non thruster actuator found in multirotor actuators. Not supported at the moment.") + + return thruster_names + + @property + def num_thrusters(self) -> int: + """Number of thrusters in the multirotor.""" + return len(self.thruster_names) + + @property + def allocation_matrix(self) -> torch.Tensor: + """Allocation matrix for control allocation.""" + return torch.tensor(self.cfg.allocation_matrix, device=self.device, dtype=torch.float32) + + """ + Operations + """ + + def set_thrust_target( + self, + target: torch.Tensor, + thruster_ids: Sequence[int] | slice | None = None, + env_ids: Sequence[int] | None = None, + ): + """Set target thrust values for thrusters. + + Args: + target: Target thrust values. Shape is (num_envs, num_thrusters) or (num_envs,). + thruster_ids: Indices of thrusters to set. Defaults to None (all thrusters). + env_ids: Environment indices to set. Defaults to None (all environments). + """ + # resolve indices + if env_ids is None: + env_ids = slice(None) + if thruster_ids is None: + thruster_ids = slice(None) + + # broadcast env_ids if needed to allow double indexing + if env_ids != slice(None) and thruster_ids != slice(None): + env_ids = env_ids[:, None] + + # set targets + self._data.thrust_target[env_ids, thruster_ids] = target + + def reset(self, env_ids: Sequence[int] | None = None): + """Reset the multirotor to default state. + + Args: + env_ids: Environment indices to reset. Defaults to None (all environments). + """ + # call parent reset + super().reset(env_ids) + + # reset multirotor-specific data + if env_ids is None: + env_ids = self._ALL_INDICES + elif not isinstance(env_ids, torch.Tensor): + env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device) + + # reset thruster targets to default values + if self._data.thrust_target is not None and self._data.default_thruster_rps is not None: + self._data.thrust_target[env_ids] = self._data.default_thruster_rps[env_ids] + + def write_data_to_sim(self): + """Write thrust and torque commands to the simulation.""" + self._apply_actuator_model() + # apply thruster forces at individual locations + self._apply_combined_wrench() + + def _initialize_impl(self): + """Initialize the multirotor implementation.""" + # call parent initialization + super()._initialize_impl() + + # Replace data container with MultirotorData + self._data = MultirotorData(self.root_physx_view, self.device) + + # Create thruster buffers with correct size (SINGLE PHASE) + self._create_thruster_buffers() + + # Process thruster configuration + self._process_thruster_cfg() + + # Process configuration + self._process_cfg() + + # Update the robot data + self.update(0.0) + + # Log multirotor information + self._log_multirotor_info() + + def _count_thrusters_from_config(self) -> int: + """Count total number of thrusters from actuator configuration. + + Returns: + Total number of thrusters across all actuator groups. + """ + total_thrusters = 0 + + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + if not hasattr(actuator_cfg, "thruster_names_expr"): + continue + + # Use find_bodies to count thrusters for this actuator + body_indices, thruster_names = self.find_bodies(actuator_cfg.thruster_names_expr) + total_thrusters += len(body_indices) + + if total_thrusters == 0: + raise ValueError( + "No thrusters found in actuator configuration. " + "Please check thruster_names_expr in your MultirotorCfg.actuators." + ) + + return total_thrusters + + def _create_thruster_buffers(self): + """Create thruster buffers with correct size.""" + num_instances = self.num_instances + num_thrusters = self._count_thrusters_from_config() + + # Create thruster data tensors with correct size + self._data.default_thruster_rps = torch.zeros(num_instances, num_thrusters, device=self.device) + # thrust after controller and allocation is applied + self._data.thrust_target = torch.zeros(num_instances, num_thrusters, device=self.device) + self._data.computed_thrust = torch.zeros(num_instances, num_thrusters, device=self.device) + self._data.applied_thrust = torch.zeros(num_instances, num_thrusters, device=self.device) + + # Combined wrench buffers + self._thrust_target_sim = torch.zeros_like(self._data.thrust_target) # thrust after actuator model is applied + # wrench target for combined mode + self._internal_wrench_target_sim = torch.zeros(num_instances, 6, device=self.device) + # internal force/torque targets per body for combined mode + self._internal_force_target_sim = torch.zeros(num_instances, self.num_bodies, 3, device=self.device) + self._internal_torque_target_sim = torch.zeros(num_instances, self.num_bodies, 3, device=self.device) + + # Placeholder thruster names (will be filled during actuator creation) + self._data.thruster_names = [f"thruster_{i}" for i in range(num_thrusters)] + + def _process_actuators_cfg(self): + """Override parent method to do nothing - we handle thrusters separately.""" + # Do nothing - we handle thruster processing in _process_thruster_cfg() otherwise this + # gives issues with joint name expressions + pass + + def _process_cfg(self): + """Post processing of multirotor configuration parameters.""" + # Handle root state (like parent does) + default_root_state = ( + tuple(self.cfg.init_state.pos) + + tuple(self.cfg.init_state.rot) + + tuple(self.cfg.init_state.lin_vel) + + tuple(self.cfg.init_state.ang_vel) + ) + default_root_state = torch.tensor(default_root_state, dtype=torch.float, device=self.device) + self._data.default_root_state = default_root_state.repeat(self.num_instances, 1) + + # Handle thruster-specific initial state + if hasattr(self._data, "default_thruster_rps") and hasattr(self.cfg.init_state, "rps"): + # Match against thruster names + indices_list, _, values_list = string_utils.resolve_matching_names_values( + self.cfg.init_state.rps, self.thruster_names + ) + if indices_list: + rps_values = torch.tensor(values_list, device=self.device) + self._data.default_thruster_rps[:, indices_list] = rps_values + self._data.thrust_target[:, indices_list] = rps_values + + def _process_thruster_cfg(self): + """Process and apply multirotor thruster properties.""" + # create actuators + self.actuators = dict() + self._has_implicit_actuators = False + + # Check for mixed configurations (same as before) + has_thrusters = False + has_joints = False + + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + if hasattr(actuator_cfg, "thruster_names_expr"): + has_thrusters = True + elif hasattr(actuator_cfg, "joint_names_expr"): + has_joints = True + + if has_thrusters and has_joints: + raise ValueError("Mixed configurations with both thrusters and regular joints are not supported.") + + if has_joints: + raise ValueError("Regular joint actuators are not supported in Multirotor class.") + + # Store the body-to-thruster mapping + self._thruster_body_mapping = {} + + # Track thruster names as we create actuators + all_thruster_names = [] + + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + body_indices, thruster_names = self.find_bodies(actuator_cfg.thruster_names_expr) + + # Create 0-based thruster array indices starting from current count + start_idx = len(all_thruster_names) + thruster_array_indices = list(range(start_idx, start_idx + len(body_indices))) + + # Track all thruster names + all_thruster_names.extend(thruster_names) + + # Store the mapping + self._thruster_body_mapping[actuator_name] = { + "body_indices": body_indices, + "array_indices": thruster_array_indices, + "thruster_names": thruster_names, + } + + # Create thruster actuator + actuator: Thruster = actuator_cfg.class_type( + cfg=actuator_cfg, + thruster_names=thruster_names, + thruster_ids=thruster_array_indices, + num_envs=self.num_instances, + device=self.device, + init_thruster_rps=self._data.default_thruster_rps[:, thruster_array_indices], + ) + + # Store actuator + self.actuators[actuator_name] = actuator + + # Log information + omni.log.info( + f"Thruster actuator: {actuator_name} with model '{actuator_cfg.class_type.__name__}'" + f" (thruster names: {thruster_names} [{body_indices}])." + ) + + # Update thruster names in data container + self._data.thruster_names = all_thruster_names + + # Log summary + omni.log.info(f"Initialized {len(self.actuators)} thruster actuator(s) for multirotor.") + + def _apply_actuator_model(self): + """Processes thruster commands for the multirotor by forwarding them to the actuators. + + The actions are first processed using actuator models. The thruster actuator models + compute the thruster level simulation commands and sets them into the PhysX buffers. + """ + + # process thruster actions per group + for actuator in self.actuators.values(): + if not isinstance(actuator, Thruster): + continue + + # prepare input for actuator model based on cached data + control_action = ArticulationThrustActions( + thrusts=self._data.thrust_target[:, actuator.thruster_indices], + thruster_indices=actuator.thruster_indices, + ) + + # compute thruster command from the actuator model + control_action = actuator.compute(control_action) + + # update targets (these are set into the simulation) + if control_action.thrusts is not None: + self._thrust_target_sim[:, actuator.thruster_indices] = control_action.thrusts + + # update state of the actuator model + self._data.computed_thrust[:, actuator.thruster_indices] = actuator.computed_thrust + self._data.applied_thrust[:, actuator.thruster_indices] = actuator.applied_thrust + + def _apply_combined_wrench(self): + """Apply combined wrench to the base link (like articulation_with_thrusters.py).""" + # Combine individual thrusts into a wrench vector + self._combine_thrusts() + + self.root_physx_view.apply_forces_and_torques_at_position( + force_data=self._internal_force_target_sim.view(-1, 3), # Shape: (num_envs * num_bodies, 3) + torque_data=self._internal_torque_target_sim.view(-1, 3), # Shape: (num_envs * num_bodies, 3) + position_data=None, # Apply at center of mass + indices=self._ALL_INDICES, + is_global=False, # Forces are in local frame + ) + + def _combine_thrusts(self): + """Combine individual thrusts into a wrench vector.""" + thrusts = self._thrust_target_sim + self._internal_wrench_target_sim = (self.allocation_matrix @ thrusts.T).T + # Apply forces to base link (body index 0) only + self._internal_force_target_sim[:, 0, :] = self._internal_wrench_target_sim[:, :3] + self._internal_torque_target_sim[:, 0, :] = self._internal_wrench_target_sim[:, 3:] + + def _validate_cfg(self): + """Validate the multirotor configuration after processing. + + Note: + This function should be called only after the configuration has been processed and the buffers have been + created. Otherwise, some settings that are altered during processing may not be validated. + """ + # Only validate if actuators have been created + if hasattr(self, "actuators") and self.actuators: + # Validate thruster-specific configuration + for actuator_name in self.actuators: + if isinstance(self.actuators[actuator_name], Thruster): + initial_thrust = self.actuators[actuator_name].curr_thrust + # check that the initial thrust is within the limits + thrust_limits = self.actuators[actuator_name].cfg.thrust_range + if torch.any(initial_thrust < thrust_limits[0]) or torch.any(initial_thrust > thrust_limits[1]): + raise ValueError( + f"Initial thrust for actuator '{actuator_name}' is out of bounds: " + f"{initial_thrust} not in {thrust_limits}" + ) + + def _log_multirotor_info(self): + """Log multirotor-specific information.""" + omni.log.info(f"Multirotor initialized with {self.num_thrusters} thrusters") + omni.log.info(f"Thruster names: {self.thruster_names}") + omni.log.info(f"Thruster force direction: {self.cfg.thruster_force_direction}") diff --git a/source/isaaclab/isaaclab/assets/articulation/multirotor_cfg.py b/source/isaaclab/isaaclab/assets/articulation/multirotor_cfg.py new file mode 100644 index 00000000000..f0e022187d6 --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation/multirotor_cfg.py @@ -0,0 +1,48 @@ +# 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 dataclasses import MISSING + +from isaaclab.actuators import ThrusterCfg +from isaaclab.utils import configclass + +from .articulation_cfg import ArticulationCfg +from .multirotor import Multirotor + + +@configclass +class MultirotorCfg(ArticulationCfg): + """Configuration parameters for a multirotor articulation. + + This extends the base articulation configuration to support multirotor-specific + settings. + """ + + class_type: type = Multirotor + + @configclass + class InitialStateCfg(ArticulationCfg.InitialStateCfg): + """Initial state of the multirotor articulation.""" + + # multirotor-specific initial state + rps: dict[str, float] = {".*": 100.0} + """RPS of the thrusters. Defaults to 100.0 for all thrusters.""" + + # multirotor-specific configuration + init_state: InitialStateCfg = InitialStateCfg() + """Initial state of the multirotor object.""" + + actuators: dict[str, ThrusterCfg] = MISSING + """Thruster actuators for the multirotor with corresponding thruster names.""" + + # multirotor force application settings + thruster_force_direction: tuple[float, float, float] = (0.0, 0.0, 1.0) + """Default force direction in body-local frame for thrusters. Defaults to Z-axis (upward).""" + + allocation_matrix: list[list[float]] | None = None + """allocation matrix for control allocation""" + + rotor_directions: list[int] | None = None + """List of rotor directions, -1 for clockwise, 1 for counter-clockwise.""" diff --git a/source/isaaclab/isaaclab/assets/articulation/multirotor_data.py b/source/isaaclab/isaaclab/assets/articulation/multirotor_data.py new file mode 100644 index 00000000000..b9443bf4fcb --- /dev/null +++ b/source/isaaclab/isaaclab/assets/articulation/multirotor_data.py @@ -0,0 +1,50 @@ +# 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 + +from isaaclab.assets.articulation.articulation_data import ArticulationData + + +class MultirotorData(ArticulationData): + """Data container for a multirotor articulation. + + This class extends the base articulation data container to include multirotor-specific + data such as thruster states and forces. + """ + + thruster_names: list[str] = None + """List of thruster names in the multirotor.""" + + default_thruster_rps: torch.Tensor = None + """Default thruster RPS state of all thrusters. Shape is (num_instances, num_thrusters). + + This quantity is configured through the :attr:`isaaclab.assets.MultirotorCfg.init_state` parameter. + """ + + thrust_target: torch.Tensor = None + """Thrust targets commanded by the user. Shape is (num_instances, num_thrusters). + + This quantity contains the target thrust values set by the user through the + :meth:`isaaclab.assets.Multirotor.set_thrust_target` method. + """ + + ## + # Thruster commands + ## + + computed_thrust: torch.Tensor = None + """Computed thrust from the actuator model (before clipping). Shape is (num_instances, num_thrusters). + + This quantity contains the thrust values computed by the thruster actuator models + before any clipping or saturation is applied. + """ + + applied_thrust: torch.Tensor = None + """Applied thrust from the actuator model (after clipping). Shape is (num_instances, num_thrusters). + + This quantity contains the final thrust values that are applied to the simulation + after all actuator model processing, including clipping and saturation. + """ diff --git a/source/isaaclab/isaaclab/controllers/lee_acceleration_control.py b/source/isaaclab/isaaclab/controllers/lee_acceleration_control.py new file mode 100644 index 00000000000..ad3fcee86d2 --- /dev/null +++ b/source/isaaclab/isaaclab/controllers/lee_acceleration_control.py @@ -0,0 +1,161 @@ +# 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 + +# Copyright (c) 2023, Autonomous Robots Lab, Norwegian University of Science and Technology. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils + +if TYPE_CHECKING: + from isaaclab.assets import Multirotor + + from .lee_acceleration_control_cfg import LeeAccControllerCfg + + +class LeeAccController: + """Implementation of a Lee-style geometric controller (SE(3)). + + Computes a body-frame wrench command ``[Fx, Fy, Fz, Tx, Ty, Tz]`` from a compact action: + collective thrust scale, roll command, pitch command, and yaw-rate command. + Gains may be randomized per environment if enabled in the configuration. + """ + + cfg: LeeAccControllerCfg + + def __init__(self, cfg: LeeAccControllerCfg, asset: Multirotor, num_envs: int, device: str): + """Initialize controller buffers and pre-compute aggregate inertias. + + Args: + cfg: Controller configuration. + env: Owning environment (must provide ``scene['robot']``, ``num_envs``, and ``device``). + """ + self.cfg = cfg + self.robot: Multirotor = asset + + # Aggregate mass and inertia about the robot COM for all bodies + root_quat_exp = self.robot.data.root_link_quat_w.unsqueeze(1).expand(num_envs, self.robot.num_bodies, 4) + body_link_pos_delta = self.robot.data.body_link_pos_w - self.robot.data.root_pos_w.unsqueeze(1) + self.mass, self.robot_inertia, _ = math_utils.aggregate_inertia_about_robot_com( + self.robot.root_physx_view.get_inertias().to(device), + self.robot.root_physx_view.get_inv_masses().to(device), + self.robot.data.body_com_pos_b, + self.robot.data.body_com_quat_b, + math_utils.quat_apply_inverse(root_quat_exp, body_link_pos_delta), + math_utils.quat_mul(math_utils.quat_inv(root_quat_exp), self.robot.data.body_link_quat_w), + ) + self.gravity = torch.tensor(self.cfg.gravity, device=device).expand(num_envs, -1) + + # Gain ranges (single tensor each): shape (num_envs, 2, 3); [:,0]=min, [:,1]=max + self.K_rot_range = torch.tensor(self.cfg.K_rot_range, device=device).repeat(num_envs, 1, 1) + self.K_angvel_range = torch.tensor(self.cfg.K_angvel_range, device=device).repeat(num_envs, 1, 1) + + # Current (possibly randomized) gains + self.K_rot_current = self.K_rot_range.mean(dim=1) + self.K_angvel_current = self.K_angvel_range.mean(dim=1) + + # Buffers (all shapes use num_envs in the first dimension) + self.accel = torch.zeros((num_envs, 3), device=device) + self.wrench_command_b = torch.zeros((num_envs, 6), device=device) # [fx, fy, fz, tx, ty, tz] + self.desired_body_angvel_w = torch.zeros_like(self.robot.data.root_ang_vel_b) + self.desired_quat = torch.zeros_like(self.robot.data.root_quat_w) + self.euler_angle_rates_w = torch.zeros_like(self.robot.data.root_ang_vel_b) + self.buffer_tensor = torch.zeros((num_envs, 3, 3), device=device) + + def compute(self, command): + + self.wrench_command_b[:] = 0.0 + acc = command[:, 0:3].clone() + forces = (acc - self.gravity) * self.mass.view(-1, 1) + self.wrench_command_b[:, 2] = torch.sum( + forces * math_utils.matrix_from_quat(self.robot.data.root_quat_w)[:, :, 2], dim=1 + ) + + robot_euler_w = torch.stack(math_utils.euler_xyz_from_quat(self.robot.data.root_quat_w), dim=-1) + robot_euler_w = math_utils.wrap_to_pi(robot_euler_w) + + self.desired_quat[:] = calculate_desired_orientation_for_position_velocity_control( + forces, robot_euler_w[:, 2], self.buffer_tensor + ) + self.euler_angle_rates_w[:] = 0.0 + self.desired_body_angvel_w[:] = 0.0 + self.wrench_command_b[:, 3:6] = self.compute_body_torque(self.desired_quat, self.desired_body_angvel_w) + + return self.wrench_command_b + + def reset(self): + """Reset controller state for all environments.""" + self.reset_idx(env_ids=None) + + def reset_idx(self, env_ids: torch.Tensor | None): + """Reset controller state (and optionally randomize gains) for selected environments. + + Args: + env_ids: Tensor of environment indices, or ``None`` for all. + """ + if env_ids is None: + env_ids = slice(None) + self.randomize_params(env_ids) + + def randomize_params(self, env_ids: slice | torch.Tensor): + """Randomize controller gains for the given environments if enabled.""" + if not self.cfg.randomize_params: + return + self.K_rot_current[env_ids] = math_utils.rand_range(self.K_rot_range[env_ids, 0], self.K_rot_range[env_ids, 1]) + self.K_angvel_current[env_ids] = math_utils.rand_range( + self.K_angvel_range[env_ids, 0], self.K_angvel_range[env_ids, 1] + ) + + def compute_body_torque(self, setpoint_orientation, setpoint_angvel): + """PD attitude control in body frame with feedforward Coriolis term. + + Args: + setpoint_orientation: (num_envs, 4) desired orientation quaternion (wxyz) in world frame. + setpoint_angvel: (num_envs, 3) desired angular velocity in world frame [rad/s]. + + Returns: + (num_envs, 3) body torque command [N·m]. + """ + setpoint_angvel[:, 2] = torch.clamp(setpoint_angvel[:, 2], -self.cfg.max_yaw_rate, self.cfg.max_yaw_rate) + RT_Rd_quat = math_utils.quat_mul( + math_utils.quat_inv(self.robot.data.root_quat_w), setpoint_orientation + ) # (N,4) wxyz + R_err = math_utils.matrix_from_quat(RT_Rd_quat) + skew_matrix = R_err.transpose(-1, -2) - R_err + rotation_error = 0.5 * torch.stack([-skew_matrix[:, 1, 2], skew_matrix[:, 0, 2], -skew_matrix[:, 0, 1]], dim=1) + angvel_error = self.robot.data.root_ang_vel_b - math_utils.quat_apply(RT_Rd_quat, setpoint_angvel) + feed_forward_body_rates = torch.cross( + self.robot.data.root_ang_vel_b, + torch.bmm(self.robot_inertia, self.robot.data.root_ang_vel_b.unsqueeze(2)).squeeze(2), + dim=1, + ) + torque = -self.K_rot_current * rotation_error - self.K_angvel_current * angvel_error + feed_forward_body_rates + return torque + + +def calculate_desired_orientation_for_position_velocity_control(forces_command, yaw_setpoint, rotation_matrix_desired): + b3_c = torch.div(forces_command, torch.norm(forces_command, dim=1).unsqueeze(1)) + temp_dir = torch.zeros_like(forces_command) + temp_dir[:, 0] = torch.cos(yaw_setpoint) + temp_dir[:, 1] = torch.sin(yaw_setpoint) + + b2_c = torch.cross(b3_c, temp_dir, dim=1) + b2_c = torch.div(b2_c, torch.norm(b2_c, dim=1).unsqueeze(1)) + b1_c = torch.cross(b2_c, b3_c, dim=1) + + rotation_matrix_desired[:, :, 0] = b1_c + rotation_matrix_desired[:, :, 1] = b2_c + rotation_matrix_desired[:, :, 2] = b3_c + q = math_utils.quat_from_matrix(rotation_matrix_desired) + quat_desired = torch.stack((q[:, 0], q[:, 1], q[:, 2], q[:, 3]), dim=1) + + return quat_desired diff --git a/source/isaaclab/isaaclab/controllers/lee_acceleration_control_cfg.py b/source/isaaclab/isaaclab/controllers/lee_acceleration_control_cfg.py new file mode 100644 index 00000000000..1006a5f62a7 --- /dev/null +++ b/source/isaaclab/isaaclab/controllers/lee_acceleration_control_cfg.py @@ -0,0 +1,47 @@ +# 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 + +# Copyright (c) 2023, Autonomous Robots Lab, Norwegian University of Science and Technology. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from .lee_acceleration_control import LeeAccController + + +@configclass +class LeeAccControllerCfg: + """Configuration for a Lee-style geometric quadrotor controller. + + Unless otherwise noted, vectors are ordered as (x, y, z) in the simulation world/body frames. + When :attr:`randomize_params` is True, gains are sampled uniformly per environment between + their corresponding ``*_min`` and ``*_max`` bounds at reset. + """ + + class_type: type[LeeAccController] = LeeAccController + """Concrete controller class to instantiate.""" + + gravity: tuple[float, float, float] = (0.0, 0.0, -9.81) + """World gravity vector used by the controller [m/s^2].""" + + K_rot_range: tuple[tuple[float, float, float], tuple[float, float, float]] = ((1.85, 1.85, 0.4), (1.6, 1.6, 0.25)) + """Orientation (rotation) error proportional gain range about body axes [unitless].""" + + K_angvel_range: tuple[tuple[float, float, float], tuple[float, float, float]] = ( + (0.5, 0.5, 0.09), + (0.4, 0.4, 0.075), + ) + """Body angular-velocity error proportional gain range (roll, pitch, yaw) [unitless].""" + + max_inclination_angle_rad: float = 1.0471975511965976 + """Maximum allowed roll/pitch magnitude (inclination) in radians.""" + + max_yaw_rate: float = 1.0471975511965976 + """Maximum allowed yaw rate command [rad/s].""" + + randomize_params: bool = True + """If True, sample controller gains uniformly between the provided min/max bounds at resets.""" diff --git a/source/isaaclab/isaaclab/controllers/lee_position_control.py b/source/isaaclab/isaaclab/controllers/lee_position_control.py new file mode 100644 index 00000000000..b5902671f62 --- /dev/null +++ b/source/isaaclab/isaaclab/controllers/lee_position_control.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 + +# Copyright (c) 2023, Autonomous Robots Lab, Norwegian University of Science and Technology. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils + +if TYPE_CHECKING: + from isaaclab.assets import Multirotor + + from .lee_position_control_cfg import LeePosControllerCfg + + +class LeePosController: + """Implementation of a Lee-style geometric controller (SE(3)). + + Computes a body-frame wrench command ``[Fx, Fy, Fz, Tx, Ty, Tz]`` from a compact action: + collective thrust scale, roll command, pitch command, and yaw-rate command. + Gains may be randomized per environment if enabled in the configuration. + """ + + cfg: LeePosControllerCfg + + def __init__(self, cfg: LeePosControllerCfg, asset: Multirotor, num_envs: int, device: str): + """Initialize controller buffers and pre-compute aggregate inertias. + + Args: + cfg: Controller configuration. + env: Owning environment (must provide ``scene['robot']``, ``num_envs``, and ``device``). + """ + self.cfg = cfg + self.robot: Multirotor = asset + + # Aggregate mass and inertia about the robot COM for all bodies + root_quat_exp = self.robot.data.root_link_quat_w.unsqueeze(1).expand(num_envs, self.robot.num_bodies, 4) + body_link_pos_delta = self.robot.data.body_link_pos_w - self.robot.data.root_pos_w.unsqueeze(1) + self.mass, self.robot_inertia, _ = math_utils.aggregate_inertia_about_robot_com( + self.robot.root_physx_view.get_inertias().to(device), + self.robot.root_physx_view.get_inv_masses().to(device), + self.robot.data.body_com_pos_b, + self.robot.data.body_com_quat_b, + math_utils.quat_apply_inverse(root_quat_exp, body_link_pos_delta), + math_utils.quat_mul(math_utils.quat_inv(root_quat_exp), self.robot.data.body_link_quat_w), + ) + self.gravity = torch.tensor(self.cfg.gravity, device=device).expand(num_envs, -1) + + # Gain ranges (single tensor each): shape (num_envs, 2, 3); [:,0]=min, [:,1]=max + self.K_pos_range = torch.tensor(self.cfg.K_pos_range, device=device).repeat(num_envs, 1, 1) + self.K_rot_range = torch.tensor(self.cfg.K_rot_range, device=device).repeat(num_envs, 1, 1) + self.K_angvel_range = torch.tensor(self.cfg.K_angvel_range, device=device).repeat(num_envs, 1, 1) + + # Current (possibly randomized) gains + self.K_pos_current = self.K_pos_range.mean(dim=1) + self.K_rot_current = self.K_rot_range.mean(dim=1) + self.K_angvel_current = self.K_angvel_range.mean(dim=1) + + # Buffers (all shapes use num_envs in the first dimension) + self.accel = torch.zeros((num_envs, 3), device=device) + self.wrench_command_b = torch.zeros((num_envs, 6), device=device) # [fx, fy, fz, tx, ty, tz] + self.desired_body_angvel_w = torch.zeros_like(self.robot.data.root_ang_vel_b) + self.desired_quat = torch.zeros_like(self.robot.data.root_quat_w) + self.euler_angle_rates_w = torch.zeros_like(self.robot.data.root_ang_vel_b) + self.buffer_tensor = torch.zeros((num_envs, 3, 3), device=device) + + def compute(self, command): + + self.wrench_command_b[:] = 0.0 + acc = self.compute_acceleration(setpoint_position=command[:, 0:3]) + forces = (acc - self.gravity) * self.mass.view(-1, 1) + self.wrench_command_b[:, 2] = torch.sum( + forces * math_utils.matrix_from_quat(self.robot.data.root_quat_w)[:, :, 2], dim=1 + ) + + robot_euler_w = torch.stack(math_utils.euler_xyz_from_quat(self.robot.data.root_quat_w), dim=-1) + robot_euler_w = math_utils.wrap_to_pi(robot_euler_w) + + self.desired_quat[:] = calculate_desired_orientation_for_position_velocity_control( + forces, robot_euler_w[:, 2], self.buffer_tensor + ) + self.euler_angle_rates_w[:] = 0.0 + self.desired_body_angvel_w[:] = 0.0 + self.wrench_command_b[:, 3:6] = self.compute_body_torque(self.desired_quat, self.desired_body_angvel_w) + + return self.wrench_command_b + + def reset(self): + """Reset controller state for all environments.""" + self.reset_idx(env_ids=None) + + def reset_idx(self, env_ids: torch.Tensor | None): + """Reset controller state (and optionally randomize gains) for selected environments. + + Args: + env_ids: Tensor of environment indices, or ``None`` for all. + """ + if env_ids is None: + env_ids = slice(None) + self.randomize_params(env_ids) + + def randomize_params(self, env_ids: slice | torch.Tensor): + """Randomize controller gains for the given environments if enabled.""" + if not self.cfg.randomize_params: + return + self.K_pos_current[env_ids] = math_utils.rand_range(self.K_pos_range[env_ids, 0], self.K_pos_range[env_ids, 1]) + self.K_rot_current[env_ids] = math_utils.rand_range(self.K_rot_range[env_ids, 0], self.K_rot_range[env_ids, 1]) + self.K_angvel_current[env_ids] = math_utils.rand_range( + self.K_angvel_range[env_ids, 0], self.K_angvel_range[env_ids, 1] + ) + + def compute_acceleration(self, setpoint_position): + position_error_world_frame = setpoint_position - self.robot.data.root_pos_w + + accel_command = self.K_pos_current * position_error_world_frame + return accel_command + + def compute_body_torque(self, setpoint_orientation, setpoint_angvel): + """PD attitude control in body frame with feedforward Coriolis term. + + Args: + setpoint_orientation: (num_envs, 4) desired orientation quaternion (wxyz) in world frame. + setpoint_angvel: (num_envs, 3) desired angular velocity in world frame [rad/s]. + + Returns: + (num_envs, 3) body torque command [N·m]. + """ + setpoint_angvel[:, 2] = torch.clamp(setpoint_angvel[:, 2], -self.cfg.max_yaw_rate, self.cfg.max_yaw_rate) + RT_Rd_quat = math_utils.quat_mul( + math_utils.quat_inv(self.robot.data.root_quat_w), setpoint_orientation + ) # (N,4) wxyz + R_err = math_utils.matrix_from_quat(RT_Rd_quat) + skew_matrix = R_err.transpose(-1, -2) - R_err + rotation_error = 0.5 * torch.stack([-skew_matrix[:, 1, 2], skew_matrix[:, 0, 2], -skew_matrix[:, 0, 1]], dim=1) + angvel_error = self.robot.data.root_ang_vel_b - math_utils.quat_apply(RT_Rd_quat, setpoint_angvel) + feed_forward_body_rates = torch.cross( + self.robot.data.root_ang_vel_b, + torch.bmm(self.robot_inertia, self.robot.data.root_ang_vel_b.unsqueeze(2)).squeeze(2), + dim=1, + ) + torque = -self.K_rot_current * rotation_error - self.K_angvel_current * angvel_error + feed_forward_body_rates + return torque + + +def calculate_desired_orientation_for_position_velocity_control(forces_command, yaw_setpoint, rotation_matrix_desired): + b3_c = torch.div(forces_command, torch.norm(forces_command, dim=1).unsqueeze(1)) + temp_dir = torch.zeros_like(forces_command) + temp_dir[:, 0] = torch.cos(yaw_setpoint) + temp_dir[:, 1] = torch.sin(yaw_setpoint) + + b2_c = torch.cross(b3_c, temp_dir, dim=1) + b2_c = torch.div(b2_c, torch.norm(b2_c, dim=1).unsqueeze(1)) + b1_c = torch.cross(b2_c, b3_c, dim=1) + + rotation_matrix_desired[:, :, 0] = b1_c + rotation_matrix_desired[:, :, 1] = b2_c + rotation_matrix_desired[:, :, 2] = b3_c + q = math_utils.quat_from_matrix(rotation_matrix_desired) + quat_desired = torch.stack((q[:, 0], q[:, 1], q[:, 2], q[:, 3]), dim=1) + + return quat_desired diff --git a/source/isaaclab/isaaclab/controllers/lee_position_control_cfg.py b/source/isaaclab/isaaclab/controllers/lee_position_control_cfg.py new file mode 100644 index 00000000000..6e6d0f29491 --- /dev/null +++ b/source/isaaclab/isaaclab/controllers/lee_position_control_cfg.py @@ -0,0 +1,50 @@ +# 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 + +# Copyright (c) 2023, Autonomous Robots Lab, Norwegian University of Science and Technology. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from .lee_position_control import LeePosController + + +@configclass +class LeePosControllerCfg: + """Configuration for a Lee-style geometric quadrotor controller. + + Unless otherwise noted, vectors are ordered as (x, y, z) in the simulation world/body frames. + When :attr:`randomize_params` is True, gains are sampled uniformly per environment between + their corresponding ``*_min`` and ``*_max`` bounds at reset. + """ + + class_type: type[LeePosController] = LeePosController + """Concrete controller class to instantiate.""" + + gravity: tuple[float, float, float] = (0.0, 0.0, -9.81) + """World gravity vector used by the controller [m/s^2].""" + + K_pos_range: tuple[tuple[float, float, float], tuple[float, float, float]] = ((4.0, 4.0, 2.0), (3.0, 3.0, 2.5)) + """Position error proportional gain range about body axes [unitless].""" + + K_rot_range: tuple[tuple[float, float, float], tuple[float, float, float]] = ((1.85, 1.85, 0.4), (1.6, 1.6, 0.25)) + """Orientation (rotation) error proportional gain range about body axes [unitless].""" + + K_angvel_range: tuple[tuple[float, float, float], tuple[float, float, float]] = ( + (0.5, 0.5, 0.09), + (0.4, 0.4, 0.075), + ) + """Body angular-velocity error proportional gain range (roll, pitch, yaw) [unitless].""" + + max_inclination_angle_rad: float = 1.0471975511965976 + """Maximum allowed roll/pitch magnitude (inclination) in radians.""" + + max_yaw_rate: float = 1.0471975511965976 + """Maximum allowed yaw rate command [rad/s].""" + + randomize_params: bool = True + """If True, sample controller gains uniformly between the provided min/max bounds at resets.""" diff --git a/source/isaaclab/isaaclab/controllers/lee_velocity_control.py b/source/isaaclab/isaaclab/controllers/lee_velocity_control.py new file mode 100644 index 00000000000..9463529b0e6 --- /dev/null +++ b/source/isaaclab/isaaclab/controllers/lee_velocity_control.py @@ -0,0 +1,244 @@ +# 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 + +# Copyright (c) 2023, Autonomous Robots Lab, Norwegian University of Science and Technology. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils + +if TYPE_CHECKING: + from isaaclab.assets import Multirotor + + from .lee_velocity_control_cfg import LeeVelControllerCfg + + +class LeeVelController: + """Implementation of a Lee-style geometric controller (SE(3)). + + Computes a body-frame wrench command ``[Fx, Fy, Fz, Tx, Ty, Tz]`` from a compact action: + collective thrust scale, roll command, pitch command, and yaw-rate command. + Gains may be randomized per environment if enabled in the configuration. + """ + + cfg: LeeVelControllerCfg + + def __init__(self, cfg: LeeVelControllerCfg, asset: Multirotor, num_envs: int, device: str): + """Initialize controller buffers and pre-compute aggregate inertias. + + Args: + cfg: Controller configuration. + env: Owning environment (must provide ``scene['robot']``, ``num_envs``, and ``device``). + """ + self.cfg = cfg + self.robot: Multirotor = asset + + # Aggregate mass and inertia about the robot COM for all bodies + root_quat_exp = self.robot.data.root_link_quat_w.unsqueeze(1).expand(num_envs, self.robot.num_bodies, 4) + body_link_pos_delta = self.robot.data.body_link_pos_w - self.robot.data.root_pos_w.unsqueeze(1) + self.mass, self.robot_inertia, _ = math_utils.aggregate_inertia_about_robot_com( + self.robot.root_physx_view.get_inertias().to(device), + self.robot.root_physx_view.get_inv_masses().to(device), + self.robot.data.body_com_pos_b, + self.robot.data.body_com_quat_b, + math_utils.quat_apply_inverse(root_quat_exp, body_link_pos_delta), + math_utils.quat_mul(math_utils.quat_inv(root_quat_exp), self.robot.data.body_link_quat_w), + ) + self.gravity = torch.tensor(self.cfg.gravity, device=device).expand(num_envs, -1) + + # Gain ranges (single tensor each): shape (num_envs, 2, 3); [:,0]=min, [:,1]=max + self.K_vel_range = torch.tensor(self.cfg.K_vel_range, device=device).repeat(num_envs, 1, 1) + self.K_rot_range = torch.tensor(self.cfg.K_rot_range, device=device).repeat(num_envs, 1, 1) + self.K_angvel_range = torch.tensor(self.cfg.K_angvel_range, device=device).repeat(num_envs, 1, 1) + + # Current (possibly randomized) gains + self.K_vel_current = self.K_vel_range.mean(dim=1) + self.K_rot_current = self.K_rot_range.mean(dim=1) + self.K_angvel_current = self.K_angvel_range.mean(dim=1) + + # Buffers (all shapes use num_envs in the first dimension) + self.accel = torch.zeros((num_envs, 3), device=device) + self.wrench_command_b = torch.zeros((num_envs, 6), device=device) # [fx, fy, fz, tx, ty, tz] + self.desired_body_angvel_w = torch.zeros_like(self.robot.data.root_ang_vel_b) + self.desired_quat = torch.zeros_like(self.robot.data.root_quat_w) + self.euler_angle_rates_w = torch.zeros_like(self.robot.data.root_ang_vel_b) + self.buffer_tensor = torch.zeros((num_envs, 3, 3), device=device) + + def compute(self, command): + + self.wrench_command_b[:] = 0.0 + acc = self.compute_acceleration(setpoint_velocity=command[:, 0:3]) + forces = (acc - self.gravity) * self.mass.view(-1, 1) + self.wrench_command_b[:, 2] = torch.sum( + forces * math_utils.matrix_from_quat(self.robot.data.root_quat_w)[:, :, 2], dim=1 + ) + + robot_euler_w = torch.stack(math_utils.euler_xyz_from_quat(self.robot.data.root_quat_w), dim=-1) + robot_euler_w = math_utils.wrap_to_pi(robot_euler_w) + + self.desired_quat = calculate_desired_orientation_for_position_velocity_control( + forces, robot_euler_w[:, 2], self.buffer_tensor + ) + self.euler_angle_rates_w[:, :2] = 0.0 + self.euler_angle_rates_w[:, 2] = command[:, 3] + self.desired_body_angvel_w[:] = euler_to_body_rate(robot_euler_w, self.euler_angle_rates_w, self.buffer_tensor) + self.wrench_command_b[:, 3:6] = self.compute_body_torque(self.desired_quat, self.desired_body_angvel_w) + + return self.wrench_command_b + + def reset(self): + """Reset controller state for all environments.""" + self.reset_idx(env_ids=None) + + def reset_idx(self, env_ids: torch.Tensor | None): + """Reset controller state (and optionally randomize gains) for selected environments. + + Args: + env_ids: Tensor of environment indices, or ``None`` for all. + """ + if env_ids is None: + env_ids = slice(None) + self.randomize_params(env_ids) + + def randomize_params(self, env_ids: slice | torch.Tensor): + """Randomize controller gains for the given environments if enabled.""" + if not self.cfg.randomize_params: + return + self.K_vel_current[env_ids] = math_utils.rand_range(self.K_vel_range[env_ids, 0], self.K_vel_range[env_ids, 1]) + self.K_rot_current[env_ids] = math_utils.rand_range(self.K_rot_range[env_ids, 0], self.K_rot_range[env_ids, 1]) + self.K_angvel_current[env_ids] = math_utils.rand_range( + self.K_angvel_range[env_ids, 0], self.K_angvel_range[env_ids, 1] + ) + + def compute_acceleration(self, setpoint_velocity): + robot_vehicle_orientation = vehicle_frame_quat_from_quat(self.robot.data.root_quat_w) + setpoint_velocity_world_frame = math_utils.quat_apply(robot_vehicle_orientation, setpoint_velocity) + velocity_error = setpoint_velocity_world_frame - self.robot.data.root_lin_vel_w + + accel_command = self.K_vel_current * velocity_error + return accel_command + + def compute_body_torque(self, setpoint_orientation, setpoint_angvel): + """PD attitude control in body frame with feedforward Coriolis term. + + Args: + setpoint_orientation: (num_envs, 4) desired orientation quaternion (wxyz) in world frame. + setpoint_angvel: (num_envs, 3) desired angular velocity in world frame [rad/s]. + + Returns: + (num_envs, 3) body torque command [N·m]. + """ + setpoint_angvel[:, 2] = torch.clamp(setpoint_angvel[:, 2], -self.cfg.max_yaw_rate, self.cfg.max_yaw_rate) + RT_Rd_quat = math_utils.quat_mul( + math_utils.quat_inv(self.robot.data.root_quat_w), setpoint_orientation + ) # (N,4) wxyz + R_err = math_utils.matrix_from_quat(RT_Rd_quat) + skew_matrix = R_err.transpose(-1, -2) - R_err + rotation_error = 0.5 * torch.stack([-skew_matrix[:, 1, 2], skew_matrix[:, 0, 2], -skew_matrix[:, 0, 1]], dim=1) + angvel_error = self.robot.data.root_ang_vel_b - math_utils.quat_apply(RT_Rd_quat, setpoint_angvel) + feed_forward_body_rates = torch.cross( + self.robot.data.root_ang_vel_b, + torch.bmm(self.robot_inertia, self.robot.data.root_ang_vel_b.unsqueeze(2)).squeeze(2), + dim=1, + ) + torque = -self.K_rot_current * rotation_error - self.K_angvel_current * angvel_error + feed_forward_body_rates + return torque + + +def calculate_desired_orientation_for_position_velocity_control(forces_command, yaw_setpoint, rotation_matrix_desired): + b3_c = torch.div(forces_command, torch.norm(forces_command, dim=1).unsqueeze(1)) + temp_dir = torch.zeros_like(forces_command) + temp_dir[:, 0] = torch.cos(yaw_setpoint) + temp_dir[:, 1] = torch.sin(yaw_setpoint) + + b2_c = torch.cross(b3_c, temp_dir, dim=1) + b2_c = torch.div(b2_c, torch.norm(b2_c, dim=1).unsqueeze(1)) + b1_c = torch.cross(b2_c, b3_c, dim=1) + + rotation_matrix_desired[:, :, 0] = b1_c + rotation_matrix_desired[:, :, 1] = b2_c + rotation_matrix_desired[:, :, 2] = b3_c + q = math_utils.quat_from_matrix(rotation_matrix_desired) + quat_desired = torch.stack((q[:, 0], q[:, 1], q[:, 2], q[:, 3]), dim=1) + + return quat_desired + + +@torch.jit.script +def euler_to_body_rate(curr_euler_rate, des_euler_rate, mat_euler_to_body_rate): + s_pitch = torch.sin(curr_euler_rate[:, 1]) + c_pitch = torch.cos(curr_euler_rate[:, 1]) + + s_roll = torch.sin(curr_euler_rate[:, 0]) + c_roll = torch.cos(curr_euler_rate[:, 0]) + + mat_euler_to_body_rate[:, 0, 0] = 1.0 + mat_euler_to_body_rate[:, 1, 1] = c_roll + mat_euler_to_body_rate[:, 0, 2] = -s_pitch + mat_euler_to_body_rate[:, 2, 1] = -s_roll + mat_euler_to_body_rate[:, 1, 2] = s_roll * c_pitch + mat_euler_to_body_rate[:, 2, 2] = c_roll * c_pitch + + return torch.bmm(mat_euler_to_body_rate, des_euler_rate.unsqueeze(2)).squeeze(2) + + +@torch.jit.script +def copysign(a, b): + # type: (float, Tensor) -> Tensor + a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0]) + return torch.abs(a) * torch.sign(b) + + +@torch.jit.script +def get_euler_xyz_tensor(q): + qx, qy, qz, qw = 1, 2, 3, 0 + # roll (x-axis rotation) + sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz]) + cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz] + roll = torch.atan2(sinr_cosp, cosr_cosp) + + # pitch (y-axis rotation) + sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx]) + pitch = torch.where(torch.abs(sinp) >= 1, copysign(torch.pi / 2.0, sinp), torch.asin(sinp)) + + # yaw (z-axis rotation) + siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy]) + cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz] + yaw = torch.atan2(siny_cosp, cosy_cosp) + + return torch.stack([roll % (2 * torch.pi), pitch % (2 * torch.pi), yaw % (2 * torch.pi)], dim=-1) + + +@torch.jit.script +def quat_from_euler_xyz_tensor(euler_xyz_tensor: torch.Tensor) -> torch.Tensor: + roll = euler_xyz_tensor[..., 0] + pitch = euler_xyz_tensor[..., 1] + yaw = euler_xyz_tensor[..., 2] + cy = torch.cos(yaw * 0.5) + sy = torch.sin(yaw * 0.5) + cr = torch.cos(roll * 0.5) + sr = torch.sin(roll * 0.5) + cp = torch.cos(pitch * 0.5) + sp = torch.sin(pitch * 0.5) + + qw = cy * cr * cp + sy * sr * sp + qx = cy * sr * cp - sy * cr * sp + qy = cy * cr * sp + sy * sr * cp + qz = sy * cr * cp - cy * sr * sp + + return torch.stack([qw, qx, qy, qz], dim=-1) + + +@torch.jit.script +def vehicle_frame_quat_from_quat(body_quat: torch.Tensor) -> torch.Tensor: + body_euler = get_euler_xyz_tensor(body_quat) * torch.tensor([0.0, 0.0, 1.0], device=body_quat.device) + return quat_from_euler_xyz_tensor(body_euler) diff --git a/source/isaaclab/isaaclab/controllers/lee_velocity_control_cfg.py b/source/isaaclab/isaaclab/controllers/lee_velocity_control_cfg.py new file mode 100644 index 00000000000..aeaaf1fc04a --- /dev/null +++ b/source/isaaclab/isaaclab/controllers/lee_velocity_control_cfg.py @@ -0,0 +1,50 @@ +# 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 + +# Copyright (c) 2023, Autonomous Robots Lab, Norwegian University of Science and Technology. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from .lee_velocity_control import LeeVelController + + +@configclass +class LeeVelControllerCfg: + """Configuration for a Lee-style geometric quadrotor controller. + + Unless otherwise noted, vectors are ordered as (x, y, z) in the simulation world/body frames. + When :attr:`randomize_params` is True, gains are sampled uniformly per environment between + their corresponding ``*_min`` and ``*_max`` bounds at reset. + """ + + class_type: type[LeeVelController] = LeeVelController + """Concrete controller class to instantiate.""" + + gravity: tuple[float, float, float] = (0.0, 0.0, -9.81) + """World gravity vector used by the controller [m/s^2].""" + + K_vel_range: tuple[tuple[float, float, float], tuple[float, float, float]] = ((3.3, 3.3, 1.3), (2.7, 2.7, 1.7)) + """Velocity error proportional gain range about body axes [unitless].""" + + K_rot_range: tuple[tuple[float, float, float], tuple[float, float, float]] = ((1.85, 1.85, 0.4), (1.6, 1.6, 0.25)) + """Orientation (rotation) error proportional gain range about body axes [unitless].""" + + K_angvel_range: tuple[tuple[float, float, float], tuple[float, float, float]] = ( + (0.5, 0.5, 0.09), + (0.4, 0.4, 0.075), + ) + """Body angular-velocity error proportional gain range (roll, pitch, yaw) [unitless].""" + + max_inclination_angle_rad: float = 1.0471975511965976 + """Maximum allowed roll/pitch magnitude (inclination) in radians.""" + + max_yaw_rate: float = 1.0471975511965976 + """Maximum allowed yaw rate command [rad/s].""" + + randomize_params: bool = True + """If True, sample controller gains uniformly between the provided min/max bounds at resets.""" diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py index 19a84846a52..1a180b832a5 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py @@ -6,6 +6,9 @@ from dataclasses import MISSING from isaaclab.controllers import DifferentialIKControllerCfg, OperationalSpaceControllerCfg +from isaaclab.controllers.lee_acceleration_control_cfg import LeeAccControllerCfg +from isaaclab.controllers.lee_position_control_cfg import LeePosControllerCfg +from isaaclab.controllers.lee_velocity_control_cfg import LeeVelControllerCfg from isaaclab.managers.action_manager import ActionTerm, ActionTermCfg from isaaclab.utils import configclass @@ -16,6 +19,7 @@ non_holonomic_actions, surface_gripper_actions, task_space_actions, + thrust_actions, ) ## @@ -374,3 +378,79 @@ class SurfaceGripperBinaryActionCfg(ActionTermCfg): """The command value to close the gripper. Defaults to 1.0.""" class_type: type[ActionTerm] = surface_gripper_actions.SurfaceGripperBinaryAction + + +## +# Drone actions. +## + + +@configclass +class ThrustActionCfg(ActionTermCfg): + """Configuration for the thrust action term. + + See :class:`ThrustAction` for more details. + """ + + class_type: type[ActionTerm] = thrust_actions.ThrustAction + + asset_name: str = MISSING + """Name or regex expression of the asset that the action will be mapped to.""" + + scale: float | dict[str, float] = 1.0 + """Scale factor for the action (float or dict of regex expressions). Defaults to 1.0.""" + + offset: float | dict[str, float] = 0.0 + """Offset factor for the action (float or dict of regex expressions). Defaults to 0.0.""" + + preserve_order: bool = False + """Whether to preserve the order of the asset names in the action output. Defaults to False.""" + + use_default_offset: bool = True + """Whether to use default thrust (e.g. hover thrust) configured in the articulation asset as offset. + Defaults to True. + + If True, this flag results in overwriting the values of :attr:`offset` to the default thrust values + from the articulation asset. + """ + + +@configclass +class NavigationActionCfg(ActionTermCfg): + """Configuration for the navigation action term. + + See :class:`NavigationAction` for more details. + """ + + class_type: type[ActionTerm] = thrust_actions.NavigationAction + + asset_name: str = MISSING + """Name or regex expression of the asset that the action will be mapped to.""" + + scale: float | dict[str, float] = 1.0 + """Scale factor for the action (float or dict of regex expressions). Defaults to 1.0.""" + + offset: float | dict[str, float] = 0.0 + """Offset factor for the action (float or dict of regex expressions). Defaults to 0.0.""" + + preserve_order: bool = False + """Whether to preserve the order of the asset names in the action output. Defaults to False.""" + + use_default_offset: bool = False + """Whether to use default thrust (e.g. hover thrust) configured in the articulation asset as offset. + Defaults to False. + + If True, this flag results in overwriting the values of :attr:`offset` to the default thrust values + from the articulation asset. + """ + + command_type: str = "vel" + """Type of command to apply: "vel" for velocity commands, "pos" for position commands. + "acc" for acceleration commands. Defaults to "vel". + """ + + action_dim: dict[str, int] = {"vel": 3, "pos": 4, "acc": 4} + """Dimension of the action space for each command type.""" + + controller_cfg: LeeVelControllerCfg | LeePosControllerCfg | LeeAccControllerCfg = MISSING + """The configuration for the Lee velocity controller.""" diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/thrust_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/thrust_actions.py new file mode 100644 index 00000000000..f7bce0e777f --- /dev/null +++ b/source/isaaclab/isaaclab/envs/mdp/actions/thrust_actions.py @@ -0,0 +1,243 @@ +# 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 omni.log + +import isaaclab.utils.string as string_utils +from isaaclab.assets.articulation import Multirotor +from isaaclab.managers.action_manager import ActionTerm + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor + + from . import actions_cfg + +from isaaclab.controllers.lee_acceleration_control import LeeAccController +from isaaclab.controllers.lee_position_control import LeePosController +from isaaclab.controllers.lee_velocity_control import LeeVelController + + +class ThrustAction(ActionTerm): + """Thrust action term that applies the processed actions as thrust commands.""" + + cfg: actions_cfg.ThrustActionCfg + """The configuration of the action term.""" + _asset: Multirotor + """The articulation asset on which the action term is applied.""" + _scale: torch.Tensor | float + """The scaling factor applied to the input action.""" + _offset: torch.Tensor | float + """The offset applied to the input action.""" + _clip: torch.Tensor + """The clip applied to the input action.""" + + def __init__(self, cfg: actions_cfg.ThrustActionCfg, env: ManagerBasedEnv) -> None: + # initialize the action term + super().__init__(cfg, env) + + thruster_names_expr = self._asset.actuators["thrusters"].cfg.thruster_names_expr + + # resolve the thrusters over which the action term is applied + self._thruster_ids, self._thruster_names = self._asset.find_bodies( + thruster_names_expr, preserve_order=self.cfg.preserve_order + ) + self._num_thrusters = len(self._thruster_ids) + # log the resolved thruster names for debugging + omni.log.info( + f"Resolved thruster names for the action term {self.__class__.__name__}:" + f" {self._thruster_names} [{self._thruster_ids}]" + ) + + # Avoid indexing across all thrusters for efficiency + if self._num_thrusters == self._asset.num_thrusters and not self.cfg.preserve_order: + self._thruster_ids = slice(None) + + # create tensors for raw and processed actions + self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device) + self._processed_actions = torch.zeros_like(self.raw_actions) + + # parse scale + if isinstance(cfg.scale, (float, int)): + self._scale = float(cfg.scale) + elif isinstance(cfg.scale, dict): + self._scale = torch.ones(self.num_envs, self.action_dim, device=self.device) + # resolve the dictionary config + index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.scale, self._thruster_names) + self._scale[:, index_list] = torch.tensor(value_list, device=self.device) + else: + raise ValueError(f"Unsupported scale type: {type(cfg.scale)}. Supported types are float and dict.") + + # parse offset + if isinstance(cfg.offset, (float, int)): + self._offset = float(cfg.offset) + elif isinstance(cfg.offset, dict): + self._offset = torch.zeros_like(self._raw_actions) + # resolve the dictionary config + index_list, _, value_list = string_utils.resolve_matching_names_values( + self.cfg.offset, self._thruster_names + ) + self._offset[:, index_list] = torch.tensor(value_list, device=self.device) + else: + raise ValueError(f"Unsupported offset type: {type(cfg.offset)}. Supported types are float and dict.") + + # parse clip + if cfg.clip is not None: + if isinstance(cfg.clip, dict): + self._clip = torch.tensor([[-float("inf"), float("inf")]], device=self.device).repeat( + self.num_envs, self.action_dim, 1 + ) + index_list, _, value_list = string_utils.resolve_matching_names_values( + self.cfg.clip, self._thruster_names + ) + self._clip[:, index_list] = torch.tensor(value_list, device=self.device) + else: + raise ValueError(f"Unsupported clip type: {type(cfg.clip)}. Supported types are dict.") + + # Handle use_default_offset + if cfg.use_default_offset: + # Use default thruster RPS as offset + self._offset = self._asset.data.default_thruster_rps[:, self._thruster_ids].clone() + + @property + def action_dim(self) -> int: + return self._num_thrusters + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_actions + + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term.""" + super().IO_descriptor + self._IO_descriptor.shape = (self.action_dim,) + self._IO_descriptor.dtype = str(self.raw_actions.dtype) + self._IO_descriptor.action_type = "ThrustAction" + self._IO_descriptor.thruster_names = self._thruster_names + self._IO_descriptor.scale = self._scale + if isinstance(self._offset, torch.Tensor): + self._IO_descriptor.offset = self._offset[0].detach().cpu().numpy().tolist() + else: + self._IO_descriptor.offset = self._offset + if self.cfg.clip is not None: + if isinstance(self._clip, torch.Tensor): + self._IO_descriptor.clip = self._clip[0].detach().cpu().numpy().tolist() + else: + self._IO_descriptor.clip = self._clip + else: + self._IO_descriptor.clip = None + return self._IO_descriptor + + def process_actions(self, actions: torch.Tensor): + """Process actions by applying scaling, offset, and clipping.""" + # store the raw actions + self._raw_actions[:] = actions + # apply the affine transformations + self._processed_actions = self._raw_actions * self._scale + self._offset + # clip actions + if self.cfg.clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, min=self._clip[:, :, 0], max=self._clip[:, :, 1] + ) + + def apply_actions(self): + """Apply the processed actions as thrust commands.""" + # Set thrust targets using thruster IDs + self._asset.set_thrust_target(self.processed_actions, thruster_ids=self._thruster_ids) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Reset the action term.""" + self._raw_actions[env_ids] = 0.0 + + +class NavigationAction(ThrustAction): + """Navigation action term that applies velocity commands to multirotors.""" + + cfg: actions_cfg.NavigationActionCfg + """The configuration of the action term.""" + + def __init__(self, cfg: actions_cfg.NavigationActionCfg, env: ManagerBasedEnv) -> None: + + # Initialize parent class (this handles all the thruster setup) + super().__init__(cfg, env) + + # Validate command type + if self.cfg.command_type not in ["vel", "pos", "acc"]: + raise ValueError( + f"Unsupported command_type {self.cfg.command_type}. Supported types are 'vel', 'pos', 'acc'." + ) + + # Initialize controller + if self.cfg.command_type == "vel": + self._lc = LeeVelController( + cfg=self.cfg.controller_cfg, asset=self._asset, num_envs=self.num_envs, device=self.device + ) + elif self.cfg.command_type == "pos": + self._lc = LeePosController( + cfg=self.cfg.controller_cfg, asset=self._asset, num_envs=self.num_envs, device=self.device + ) + elif self.cfg.command_type == "acc": + self._lc = LeeAccController( + cfg=self.cfg.controller_cfg, asset=self._asset, num_envs=self.num_envs, device=self.device + ) + + @property + def action_dim(self) -> int: + return self.cfg.action_dim[self.cfg.command_type] + + @property + def IO_descriptor(self) -> GenericActionIODescriptor: + """The IO descriptor of the action term.""" + # Get parent IO descriptor + descriptor = super().IO_descriptor + # Override action type for navigation + descriptor.action_type = "NavigationAction" + return descriptor + + def apply_actions(self): + """Apply the processed actions as velocity commands.""" + # process the actions to be in the correct range + clamped_action = torch.clamp(self.processed_actions, min=-1.0, max=1.0) + processed_actions = torch.zeros(self.num_envs, 4, device=self.device) + max_speed = 2.0 # [m/s] + max_yawrate = torch.pi / 3.0 # [rad/s] + max_inclination_angle = torch.pi / 4.0 # [rad] + + clamped_action[:, 0] += 1.0 # only allow positive thrust commands [0, 2] + + processed_actions[:, 0] = ( + clamped_action[:, 0] * torch.cos(max_inclination_angle * clamped_action[:, 1]) * max_speed / 2.0 + ) + processed_actions[:, 1] = 0.0 # set lateral thrust command to 0 + processed_actions[:, 2] = ( + clamped_action[:, 0] * torch.sin(max_inclination_angle * clamped_action[:, 1]) * max_speed / 2.0 + ) + processed_actions[:, 3] = clamped_action[:, 2] * max_yawrate + + # Compute wrench command using controller + wrench_command = self._lc.compute(processed_actions) + + # Convert wrench to thrust commands using allocation matrix + thrust_commands = (torch.pinverse(self._asset.allocation_matrix) @ wrench_command.T).T + + # Apply thrust commands using thruster IDs + self._asset.set_thrust_target(thrust_commands, thruster_ids=self._thruster_ids) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + # Call parent reset + super().reset(env_ids) + # Reset controller internal states + self._lc.reset_idx(env_ids) 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..743df793c09 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_camera.py @@ -0,0 +1,221 @@ +# 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..579b2075da2 --- /dev/null +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster_cfg.py @@ -0,0 +1,72 @@ +# 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. + """ + + is_global: bool = False + + 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/math.py b/source/isaaclab/isaaclab/utils/math.py index f8ad612a916..c9203bff594 100644 --- a/source/isaaclab/isaaclab/utils/math.py +++ b/source/isaaclab/isaaclab/utils/math.py @@ -1966,3 +1966,83 @@ def generate_random_transformation_matrix(pos_boundary: float = 1, rot_boundary: T[:3, 3] = translation return T + + +def aggregate_inertia_about_robot_com( + body_inertias_local: torch.Tensor, + body_inv_mass_local: torch.Tensor, + body_com_pos_b: torch.Tensor, + body_com_quat_b: torch.Tensor, + body_pos_b: torch.Tensor, + body_quat_b: torch.Tensor, + eps=1e-12, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Aggregate per-link inertias into a single inertia about the robot COM, + expressed in the base (root link) frame. + + Shapes: + num_envs=N, num_bodies=B + + Args: + body_inertias_local (N,B,9|3,3): Link inertias in the mass/COM frame. + body_inv_mass_local (N,B): Inverse link masses (<=0 treated as padding). + body_com_pos_b (N,B,3): Link COM position relative to the link frame + (massLocalPose translation); used as body_pos_b + R_link_base @ body_com_pos_b. + body_com_quat_b (N,B,4 wxyz): Mass→link rotation (massLocalPose rotation). + body_pos_b (N,B,3): Link origins in base frame. + body_quat_b (N,B,4 wxyz): Link→base orientation. + eps (float): Small value to guard division by zero. + + Returns: + total_mass (N,): Sum of link masses. + I_total (N,3,3): Inertia about robot COM in base frame (symmetrized). + com_robot_b (N,3): Robot COM in base frame. + + Method (base frame throughout): + 1) COM of each link: com_link_b = body_pos_b + R_link_base @ body_com_pos_b + 2) Robot COM: mass-weighted average of com_link_b + 3) Rotate each link inertia: I_b = (R_link_base @ R_mass_link) I_local (⋯)^T + 4) Parallel-axis: I_pa = m (‖r‖² I - r rᵀ), r = com_link_b - com_robot_b + 5) Sum over links and symmetrize + """ + # Inertia in mass frame (local to COM) + num_envs, num_bodies, _ = body_inertias_local.shape + I_local = body_inertias_local.view(num_envs, num_bodies, 3, 3) + + # Masses + m = torch.where(body_inv_mass_local > 0, 1.0 / body_inv_mass_local, torch.zeros_like(body_inv_mass_local)) + m_sum = m.sum(dim=1, keepdim=True) + valid = (m > 0).float().unsqueeze(-1) + + # Rotations: link->base (R_link_base) and mass->link (R_mass_link) + R_link_base = matrix_from_quat(body_quat_b) + R_mass_link = body_pos_b + (R_link_base @ body_com_pos_b[..., :, None]).squeeze(-1) + + # Robot COM base frame (mass-weighted) + com_robot_b = (m.unsqueeze(-1) * R_mass_link).sum(dim=1) / (m_sum + eps) + + # Rotate inertia from mass frame to world: R = R_link_base * R_mass + R_mass = matrix_from_quat(body_com_quat_b) + R = R_link_base @ R_mass + I_world = R @ I_local @ R.transpose(-1, -2) + + # Parallel-axis to robot COM + r = R_mass_link - com_robot_b[:, None, :] + rrT = r[..., :, None] @ r[..., None, :] + r2 = (r * r).sum(dim=-1, keepdim=True) + I3 = torch.eye(3, device=body_pos_b.device).reshape(1, 1, 3, 3).expand(num_envs, num_bodies, 3, 3) + I_pa = m[..., None, None] * (r2[..., None] * I3 - rrT) + + # Sum over links (ignore zero-mass pads) + I_total = ((I_world + I_pa) * valid[..., None]).sum(dim=1) + I_total = 0.5 * (I_total + I_total.transpose(-1, -2)) + total_mass = m.sum(dim=1) + + return total_mass, I_total, com_robot_b + + +@torch.jit.script +def rand_range(lower, upper): + # type: (torch.Tensor, torch.Tensor) -> torch.Tensor + return (upper - lower) * torch.rand_like(upper) + lower 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/types.py b/source/isaaclab/isaaclab/utils/types.py index aa6f1fbfd45..8d67059b6a6 100644 --- a/source/isaaclab/isaaclab/utils/types.py +++ b/source/isaaclab/isaaclab/utils/types.py @@ -37,3 +37,24 @@ class ArticulationActions: If the joint indices are a slice, this indicates that the indices are continuous and correspond to all the joints of the articulation. We use a slice to make the indexing more efficient. """ + + +@dataclass +class ArticulationThrustActions: + """Data container to store articulation's thruster actions. + + This class is used to store the actions of the thrusters of an articulation. + It is used to store the thrust values and indices. + + If the actions are not provided, the values are set to None. + """ + + thrusts: torch.Tensor | None = None + """The thrusts of the articulation. Defaults to None.""" + + thruster_indices: torch.Tensor | Sequence[int] | slice | None = None + """The thruster indices of the articulation. Defaults to None. + + If the thruster indices are a slice, this indicates that the indices are continuous and correspond + to all the thrusters of the articulation. We use a slice to make the indexing more efficient. + """ 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) diff --git a/source/isaaclab_assets/isaaclab_assets/robots/arl_robot_1.py b/source/isaaclab_assets/isaaclab_assets/robots/arl_robot_1.py new file mode 100644 index 00000000000..7e4e87810e5 --- /dev/null +++ b/source/isaaclab_assets/isaaclab_assets/robots/arl_robot_1.py @@ -0,0 +1,74 @@ +# 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 ARL robots. + +The following configuration parameters are available: + +* :obj:`ARL_ROBOT_1_CFG`: The ARL_Robot_1 with (TODO add motor propeller combination) +""" + +import isaaclab.sim as sim_utils +from isaaclab import ISAACLAB_EXT_DIR +from isaaclab.actuators import ThrusterCfg +from isaaclab.assets.articulation import MultirotorCfg + +## +# Configuration - Actuators. +## + +ARL_ROBOT_1_THRUSTER = ThrusterCfg( + thrust_range=(0.1, 10.0), + thrust_const_range=(9.26312e-06, 1.826312e-05), + tau_inc_range=(0.05, 0.08), + tau_dec_range=(0.005, 0.005), + torque_to_thrust_ratio=0.07, + thruster_names_expr=["back_left_prop", "back_right_prop", "front_left_prop", "front_right_prop"], +) + +## +# Configuration - Articulation. +## + +ARL_ROBOT_1_CFG = MultirotorCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAACLAB_EXT_DIR}/../isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/robot_model/arl_robot_1/arl_robot_1.usd", + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + retain_accelerations=False, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=1000.0, + max_depenetration_velocity=1.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + ), + init_state=MultirotorCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.0), + lin_vel=(0.0, 0.0, 0.0), + ang_vel=(0.0, 0.0, 0.0), + rot=(1.0, 0.0, 0.0, 0.0), + rps={ + "back_left_prop": 200.0, + "back_right_prop": 200.0, + "front_left_prop": 200.0, + "front_right_prop": 200.0, + }, + ), + actuators={"thrusters": ARL_ROBOT_1_THRUSTER}, + rotor_directions=[1, -1, 1, -1], + allocation_matrix=[ + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [1.0, 1.0, 1.0, 1.0], + [-0.13, -0.13, 0.13, 0.13], + [-0.13, 0.13, 0.13, -0.13], + [-0.07, 0.07, -0.07, 0.07], + ], +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/__init__.py new file mode 100644 index 00000000000..ee75c13ca1f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Drone ARL environments.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.py new file mode 100644 index 00000000000..40965c889bb --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/__init__.py @@ -0,0 +1,13 @@ +# 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 sub-module contains the functions that are specific to the drone ARL environments.""" + +from isaaclab.envs.mdp import * # noqa: F401, F403 + +from .commands import * # noqa: F401, F403 +from .curriculums import * # noqa: F401, F403 +from .observations import * # noqa: F401, F403 +from .rewards import * # noqa: F401, F403 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/__init__.py new file mode 100644 index 00000000000..f47a4f597d1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/__init__.py @@ -0,0 +1,9 @@ +# 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 + +"""Various command terms that can be used in the environment.""" + +from .commands_cfg import DroneUniformPoseCommandCfg +from .drone_pose_command import DroneUniformPoseCommand diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/commands_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/commands_cfg.py new file mode 100644 index 00000000000..5de9fab3eea --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/commands_cfg.py @@ -0,0 +1,16 @@ +# 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 isaaclab.envs.mdp.commands.commands_cfg import UniformPoseCommandCfg +from isaaclab.utils import configclass + +from .drone_pose_command import DroneUniformPoseCommand + + +@configclass +class DroneUniformPoseCommandCfg(UniformPoseCommandCfg): + """Configuration for uniform drone pose command generator.""" + + class_type: type = DroneUniformPoseCommand diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/drone_pose_command.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/drone_pose_command.py new file mode 100644 index 00000000000..3b3c7246d2a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/commands/drone_pose_command.py @@ -0,0 +1,67 @@ +# 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 + +"""Sub-module containing command generators for pose tracking.""" + +from __future__ import annotations + +import torch + +from isaaclab.envs.mdp.commands.pose_command import UniformPoseCommand +from isaaclab.utils.math import combine_frame_transforms, compute_pose_error + + +class DroneUniformPoseCommand(UniformPoseCommand): + """Drone-specific UniformPoseCommand extensions. + + This class customizes the generic :class:`UniformPoseCommand` for drone (multirotor) + use-cases. Main differences and additions: + + - Transforms pose commands from the drone's base frame to the world frame before use. + - Accounts for per-environment origin offsets (``scene.env_origins``) when computing + position errors so tasks running on shifted/sub-terrain environments compute + meaningful errors. + - Computes and exposes simple metrics used by higher-level code: ``position_error`` + and ``orientation_error`` (stored in ``self.metrics``). + - Provides a debug visualization callback that renders the goal pose (with + sub-terrain shift) and current body pose using the existing visualizers. + + The implementation overrides :meth:`_update_metrics` and :meth:`_debug_vis_callback` + from the base class to implement these drone-specific behaviors. + """ + + def _update_metrics(self): + # transform command from base frame to simulation world frame + self.pose_command_w[:, :3], self.pose_command_w[:, 3:] = combine_frame_transforms( + self.robot.data.root_pos_w, + self.robot.data.root_quat_w, + self.pose_command_b[:, :3], + self.pose_command_b[:, 3:], + ) + # compute the error + pos_error, rot_error = compute_pose_error( + # Sub-terrain shift for correct position error calculation + self.pose_command_b[:, :3] + self._env.scene.env_origins, + self.pose_command_w[:, 3:], + self.robot.data.body_pos_w[:, self.body_idx], + self.robot.data.body_quat_w[:, self.body_idx], + ) + self.metrics["position_error"] = torch.norm(pos_error, dim=-1) + self.metrics["orientation_error"] = torch.norm(rot_error, dim=-1) + + def _debug_vis_callback(self, event): + # check if robot is initialized + # note: this is needed in-case the robot is de-initialized. we can't access the data + if not self.robot.is_initialized: + return + # update the markers + # -- goal pose + # Sub-terrain shift for visualization purposes + self.goal_pose_visualizer.visualize( + self.pose_command_b[:, :3] + self._env.scene.env_origins, self.pose_command_b[:, 3:] + ) + # -- current body pose + body_link_pose_w = self.robot.data.body_link_pose_w[:, self.body_idx] + self.current_pose_visualizer.visualize(body_link_pose_w[:, :3], body_link_pose_w[:, 3:7]) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/curriculums.py new file mode 100644 index 00000000000..6d50dc41ef5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/curriculums.py @@ -0,0 +1,106 @@ +# 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 + +"""Common functions that can be used to create curriculum for the drone navigation environment. + +The functions can be passed to the :class:`isaaclab.managers.CurriculumTermCfg` object to enable +the curriculum introduced by the function. +""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation +from isaaclab.managers import SceneEntityCfg +from isaaclab.terrains import TerrainImporter + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def get_obstacle_curriculum_state( + env: ManagerBasedRLEnv, + min_difficulty: int = 2, + max_difficulty: int = 10, +) -> dict: + """Get or lazily initialize the obstacle curriculum state dictionary. + + This helper function manages the obstacle curriculum state by initializing it + on first call and returning the existing state on subsequent calls. The state + is stored as a private attribute on the environment instance. + + The curriculum state tracks per-environment difficulty levels used to control + the number of obstacles spawned in each environment. Difficulty progresses + based on agent performance (successful goal reaching vs. collisions). + + Args: + env: The manager-based RL environment instance. + min_difficulty: Minimum difficulty level for obstacle density. Lower values + mean fewer obstacles. Defaults to 2. + max_difficulty: Maximum difficulty level for obstacle density. Higher values + mean more obstacles. Defaults to 10. + + Returns: + Dictionary containing: + - "difficulty_levels": torch.Tensor of shape (num_envs,) with per-environment + difficulty levels, initialized to min_difficulty + - "min_difficulty": int, the minimum difficulty value + - "max_difficulty": int, the maximum difficulty value + + Note: + This function stores state directly on the environment as `_obstacle_curriculum_state`. + It's designed to be called from both curriculum terms and reward functions to ensure + consistent access to the difficulty state. + """ + + if not hasattr(env, "_obstacle_curriculum_state"): + env._obstacle_curriculum_state = { + "difficulty_levels": torch.ones(env.num_envs, device=env.device) * min_difficulty, + "min_difficulty": min_difficulty, + "max_difficulty": max_difficulty, + } + return env._obstacle_curriculum_state + + +def obstacle_density_curriculum( + env: ManagerBasedRLEnv, + env_ids: Sequence[int], + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + command_name: str = "target_pose", + max_difficulty: int = 10, + min_difficulty: int = 2, +) -> float: + """Curriculum that adjusts obstacle density based on performance. + + The difficulty state is managed centrally via get_obstacle_curriculum_state(). + This ensures consistent access across curriculum, reward, and event terms. + """ + # Get or initialize curriculum state + curriculum_state = get_obstacle_curriculum_state(env, min_difficulty, max_difficulty) + + # Extract robot and command + asset: Articulation = env.scene[asset_cfg.name] + command = env.command_manager.get_command(command_name) + + target_position_w = command[:, :3].clone() + current_position = asset.data.root_pos_w - env.scene.env_origins + position_error = torch.norm(target_position_w[env_ids] - current_position[env_ids], dim=1) + + # Decide difficulty changes + crashed = env.termination_manager.terminated[env_ids] + move_up = position_error < 1.5 # Success + move_down = crashed & ~move_up + + # Update difficulty levels + difficulty_levels = curriculum_state["difficulty_levels"] + difficulty_levels[env_ids] += move_up.long() - move_down.long() + difficulty_levels[env_ids] = torch.clamp( + difficulty_levels[env_ids], min=curriculum_state["min_difficulty"], max=curriculum_state["max_difficulty"] - 1 + ) + + return difficulty_levels.float().mean().item() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/observations.py new file mode 100644 index 00000000000..ae6f8c8e090 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/observations.py @@ -0,0 +1,174 @@ +# 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 + +"""Common functions that can be used to create drone observation terms. + +The functions can be passed to the :class:`isaaclab.managers.ObservationTermCfg` object to enable +the observation introduced by the function. +""" + +from __future__ import annotations + +import torch +import torch.jit +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.assets import Articulation, Multirotor +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import Camera, MultiMeshRayCasterCamera, RayCasterCamera, TiledCamera + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv + +from isaaclab.envs.utils.io_descriptors import generic_io_descriptor, record_shape + +""" +State. +""" + + +def base_roll_pitch(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Roll and pitch of the base in the simulation world frame.""" + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + # extract euler angles (in world frame) + roll, pitch, _ = math_utils.euler_xyz_from_quat(asset.data.root_quat_w) + # normalize angle to [-pi, pi] + roll = math_utils.wrap_to_pi(roll) + pitch = math_utils.wrap_to_pi(pitch) + + return torch.cat((roll.unsqueeze(-1), pitch.unsqueeze(-1)), dim=-1) + + +""" +Sensors +""" + + +class VAEModelManager: + """Manager for the VAE model.""" + + _model = None + + @classmethod + def get_model(cls, device): + """Get or load the VAE model.""" + if cls._model is None: + import os + + model_path = os.path.join(os.path.dirname(__file__), "vae_model.pt") + cls._model = torch.jit.load(model_path, map_location=device) + cls._model.eval() + return cls._model + + +def image_latents( + env: ManagerBasedEnv, + sensor_cfg: SceneEntityCfg = SceneEntityCfg("tiled_camera"), + data_type: str = "rgb", + convert_perspective_to_orthogonal: bool = False, + normalize: bool = True, +) -> torch.Tensor: + """Images of a specific datatype from the camera sensor. + + If the flag :attr:`normalize` is True, post-processing of the images are performed based on their + data-types: + + - "rgb": Scales the image to (0, 1) and subtracts with the mean of the current image batch. + - "depth" or "distance_to_camera" or "distance_to_plane": Replaces infinity values with zero. + + Args: + env: The environment the cameras are placed within. + sensor_cfg: The desired sensor to read from. Defaults to SceneEntityCfg("tiled_camera"). + data_type: The data type to pull from the desired camera. Defaults to "rgb". + convert_perspective_to_orthogonal: Whether to orthogonalize perspective depth images. + This is used only when the data type is "distance_to_camera". Defaults to False. + normalize: Whether to normalize the images. This depends on the selected data type. + Defaults to True. + + Returns: + The images produced at the last time-step + """ + # extract the used quantities (to enable type-hinting) + sensor: TiledCamera | Camera | RayCasterCamera | MultiMeshRayCasterCamera = env.scene.sensors[sensor_cfg.name] + + # obtain the input image + images = sensor.data.output[data_type] + + # depth image conversion + if (data_type == "distance_to_camera") and convert_perspective_to_orthogonal: + images = math_utils.orthogonalize_perspective_depth(images, sensor.data.intrinsic_matrices) + + # rgb/depth/normals image normalization + if normalize: + if data_type == "rgb": + images = images.float() / 255.0 + mean_tensor = torch.mean(images, dim=(1, 2), keepdim=True) + images -= mean_tensor + elif "distance_to" in data_type or "depth" in data_type: + images[images == float("inf")] = 10.0 + images[images == -float("inf")] = 10.0 + images[images > 10.0] = 10.0 + images = images / 10.0 # normalize to 0-1 + images[images < 0.02] = -1.0 # set very close values to -1 + elif "normals" in data_type: + images = (images + 1.0) * 0.5 + + _vae_model = VAEModelManager.get_model(env.device) + + with torch.no_grad(): + latents = _vae_model(images.squeeze(-1).half()) + + return latents + + +""" +Actions. +""" + + +@generic_io_descriptor(dtype=torch.float32, observation_type="Action", on_inspect=[record_shape]) +def last_action_navigation(env: ManagerBasedEnv, action_name: str | None = None) -> torch.Tensor: + """The last input action to the environment. + + The name of the action term for which the action is required. If None, the + entire action tensor is returned. + """ + clamped_action = torch.clamp(env.action_manager.action, min=-1.0, max=1.0) + processed_actions = torch.zeros(env.num_envs, 4, device=env.device) + max_speed = 2.0 # [m/s] + max_yawrate = torch.pi / 3.0 # [rad/s] + max_inclination_angle = torch.pi / 4.0 # [rad] + + clamped_action[:, 0] += 1.0 # only allow positive thrust commands [0, 2] + processed_actions[:, 0] = ( + clamped_action[:, 0] * torch.cos(max_inclination_angle * clamped_action[:, 1]) * max_speed / 2.0 + ) + processed_actions[:, 1] = 0.0 # set lateral thrust command to 0 + processed_actions[:, 2] = ( + clamped_action[:, 0] * torch.sin(max_inclination_angle * clamped_action[:, 1]) * max_speed / 2.0 + ) + processed_actions[:, 3] = clamped_action[:, 2] * max_yawrate + return processed_actions + + +""" +Commands. +""" + + +@generic_io_descriptor(dtype=torch.float32, observation_type="Command", on_inspect=[record_shape]) +def generated_drone_commands( + env: ManagerBasedRLEnv, command_name: str | None = None, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """The generated command from command term in the command manager with the given name.""" + asset: Multirotor = env.scene[asset_cfg.name] + current_position_w = asset.data.root_pos_w - env.scene.env_origins + command = env.command_manager.get_command(command_name) + current_position_b = math_utils.quat_apply_inverse(asset.data.root_link_quat_w, command[:, :3] - current_position_w) + current_position_b_dir = current_position_b / (torch.norm(current_position_b, dim=-1, keepdim=True) + 1e-8) + current_position_b_mag = torch.norm(current_position_b, dim=-1, keepdim=True) + return torch.cat((current_position_b_dir, current_position_b_mag), dim=-1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/rewards.py new file mode 100644 index 00000000000..a2cce36f6b7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/rewards.py @@ -0,0 +1,248 @@ +# 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 typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + +from isaaclab.assets import RigidObject +from isaaclab.managers import SceneEntityCfg + +""" +Drone control rewards. +""" + + +def distance_to_goal_exp( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + std: float = 1.0, + command_name: str = "target_pose", +) -> torch.Tensor: + """Reward the distance to a goal position using an exponential kernel. + + This reward computes an exponential falloff of the squared Euclidean distance + between the commanded target position and the asset (robot) root position. + + Args: + env: The manager-based RL environment instance. + asset_cfg: SceneEntityCfg identifying the asset (defaults to "robot"). + std: Standard deviation used in the exponential kernel; larger values + produce a gentler falloff. + command_name: Name of the command to read the target pose from the + environment's command manager. The function expects the command + tensor to contain positions in its first three columns. + + Returns: + A 1-D tensor of shape (num_envs,) containing the per-environment reward + values in [0, 1], with 1.0 when the position error is zero. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + command = env.command_manager.get_command(command_name) + + target_position_w = command[:, :3].clone() + current_position = asset.data.root_pos_w - env.scene.env_origins + + # compute the error + position_error_square = torch.sum(torch.square(target_position_w - current_position), dim=1) + return torch.exp(-position_error_square / std**2) + + +def distance_to_goal_exp_curriculum( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + std: float = 1.0, + command_name: str = "target_pose", +) -> torch.Tensor: + """Reward the distance to a goal position using an exponential kernel with curriculum-based scaling. + + This reward extends the basic exponential distance reward by applying a scaling factor + that increases with the obstacle difficulty level. As the curriculum progresses and + obstacle density increases, the reward weight grows to compensate for the added difficulty. + + The scaling weight is computed as: 1.0 + (difficulty_level / max_difficulty), meaning + the reward can scale from 1.0x (at minimum difficulty) to 2.0x (at maximum difficulty). + + Args: + env: The manager-based RL environment instance. + asset_cfg: SceneEntityCfg identifying the asset (defaults to "robot"). + std: Standard deviation used in the exponential kernel; larger values + produce a gentler falloff. Defaults to 1.0. + command_name: Name of the command to read the target pose from the + environment's command manager. The function expects the command + tensor to contain positions in its first three columns. + + Returns: + A 1-D tensor of shape (num_envs,) containing the per-environment weighted + reward values. Values are in [0, weight], where weight varies based on the + current curriculum difficulty level. + + Note: + If no curriculum is active (i.e., `env._obstacle_difficulty_levels` doesn't exist), + the function behaves identically to :func:`distance_to_goal_exp` with weight=1.0. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + command = env.command_manager.get_command(command_name) + + target_position_w = command[:, :3].clone() + current_position = asset.data.root_pos_w - env.scene.env_origins + + # compute the error + position_error_square = torch.sum(torch.square(target_position_w - current_position), dim=1) + # weight based on the current curriculum level + if hasattr(env, "_obstacle_difficulty_levels"): + weight = 1.0 + env._obstacle_difficulty_levels.float() / float(env._max_obstacle_difficulty) + else: + weight = 1.0 + return weight * torch.exp(-position_error_square / std**2) + + +def ang_vel_xyz_exp( + env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), std: float = 1.0 +) -> torch.Tensor: + """Penalize angular velocity magnitude with an exponential kernel. + + This reward computes exp(-||omega||^2 / std^2) where omega is the body-frame + angular velocity of the asset. It is useful for encouraging low rotational + rates. + + Args: + env: The manager-based RL environment instance. + asset_cfg: SceneEntityCfg identifying the asset (defaults to "robot"). + std: Standard deviation used in the exponential kernel; controls + sensitivity to angular velocity magnitude. + + Returns: + A 1-D tensor of shape (num_envs,) with values in (0, 1], where 1 indicates + zero angular velocity. + """ + + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + + # compute squared magnitude of angular velocity (all axes) + ang_vel_squared = torch.sum(torch.square(asset.data.root_ang_vel_b), dim=1) + + return torch.exp(-ang_vel_squared / std**2) + + +def velocity_to_goal_reward_curriculum( + env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), command_name: str = "target_pose" +) -> torch.Tensor: + """Reward velocity alignment toward the goal with curriculum-based scaling. + + This reward encourages the agent to move in the direction of the goal by computing + the dot product between the asset's velocity vector and the normalized direction + vector to the goal. A curriculum-based scaling factor is applied that increases + with obstacle difficulty. + + The reward is positive when moving toward the goal, negative when moving away, + and zero when moving perpendicular to the goal direction. The magnitude scales + linearly with speed in the goal direction. + + The scaling weight is computed as: 1.0 + (difficulty_level / max_difficulty), + allowing the reward to scale from 1.0x to 2.0x as difficulty increases. + + Args: + env: The manager-based RL environment instance. + asset_cfg: SceneEntityCfg identifying the asset (defaults to "robot"). + command_name: Name of the command to read the target pose from the + environment's command manager. The function expects the command + tensor to contain positions in its first three columns. + + Returns: + A 1-D tensor of shape (num_envs,) containing the per-environment weighted + reward values. Values can be positive (moving toward goal), negative + (moving away), or zero (perpendicular motion), scaled by the curriculum weight. + + Note: + If no curriculum is active (i.e., `env._obstacle_difficulty_levels` doesn't exist), + the function uses weight=1.0 without curriculum scaling. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + # get the center of the environment + command = env.command_manager.get_command(command_name) + + target_position_w = command[:, :3].clone() + current_position = asset.data.root_pos_w - env.scene.env_origins + direction_to_goal = target_position_w - current_position + direction_to_goal = direction_to_goal / (torch.norm(direction_to_goal, dim=1, keepdim=True) + 1e-8) + # compute the reward as the dot product between the velocity and the direction to the goal + velocity_towards_goal = torch.sum(asset.data.root_lin_vel_w * direction_to_goal, dim=1) + # Use obstacle curriculum if it exists + if hasattr(env, "_obstacle_difficulty_levels"): + weight = 1.0 + env._obstacle_difficulty_levels.float() / float(env._max_obstacle_difficulty) + else: + weight = 1.0 + return weight * velocity_towards_goal + + +def lin_vel_xyz_exp( + env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), std: float = 1.0 +) -> torch.Tensor: + """Penalize linear velocity magnitude with an exponential kernel. + + Computes exp(-||v||^2 / std^2) where v is the asset's linear velocity in + world frame. Useful for encouraging the agent to reduce translational speed. + + Args: + env: The manager-based RL environment instance. + asset_cfg: SceneEntityCfg identifying the asset (defaults to "robot"). + std: Standard deviation used in the exponential kernel. + + Returns: + A 1-D tensor of shape (num_envs,) with values in (0, 1], where 1 indicates + zero linear velocity. + """ + + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + + # compute squared magnitude of linear velocity (all axes) + lin_vel_squared = torch.sum(torch.square(asset.data.root_lin_vel_w), dim=1) + + return torch.exp(-lin_vel_squared / std**2) + + +def yaw_aligned( + env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), std: float = 0.5 +) -> torch.Tensor: + """Reward alignment of the vehicle's yaw to zero using an exponential kernel. + + The function extracts the yaw (rotation about Z) from the world-frame root + quaternion and computes exp(-yaw^2 / std^2). This encourages heading + alignment to a zero-yaw reference. + + Args: + env: The manager-based RL environment instance. + asset_cfg: SceneEntityCfg identifying the asset (defaults to "robot"). + std: Standard deviation used in the exponential kernel; smaller values + make the reward more sensitive to yaw deviations. + + Returns: + A 1-D tensor of shape (num_envs,) with values in (0, 1], where 1 indicates + perfect yaw alignment (yaw == 0). + """ + + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + + # extract yaw from current orientation + _, _, yaw = math_utils.euler_xyz_from_quat(asset.data.root_quat_w) + + # normalize yaw to [-pi, pi] (target is 0) + yaw = math_utils.wrap_to_pi(yaw) + + # return exponential reward (1 when yaw=0, approaching 0 when rotated) + return torch.exp(-(yaw**2) / std**2) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/vae_model.pt b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/vae_model.pt new file mode 100644 index 00000000000..1da2481f35c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/mdp/vae_model.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1da3ff7787b8f2614a3e5bae46d675fdec941ef125ccc7bf61a6822d5c8413c2 +size 4257097 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/__init__.py new file mode 100644 index 00000000000..71145604512 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Drone NTNU navigation environments.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/__init__.py new file mode 100644 index 00000000000..dc608831e53 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Configurations for drone NTNU navigation environments.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/__init__.py new file mode 100644 index 00000000000..be1ec37d582 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/__init__.py @@ -0,0 +1,36 @@ +# 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 gymnasium as gym + +from . import agents + +## +# Register Gym environments. +## + +gym.register( + id="Isaac-Navigation-3DObstacles-ARL-robot-1-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.floating_obstacles_env_cfg:FloatingObstacleEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_rough_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:NavigationEnvPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", + }, +) + +gym.register( + id="Isaac-Navigation-3DObstacles-ARL-robot-1-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.floating_obstacles_env_cfg:FloatingObstacleEnvCfg_PLAY", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_rough_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:NavigationEnvPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_rough_ppo_cfg.yaml", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/__init__.py new file mode 100644 index 00000000000..2e924fbf1b1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/rl_games_rough_ppo_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/rl_games_rough_ppo_cfg.yaml new file mode 100644 index 00000000000..ee56f5851c9 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/rl_games_rough_ppo_cfg.yaml @@ -0,0 +1,87 @@ +# 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 + +params: + seed: 42 + + # environment wrapper clipping + env: + clip_actions: 1.0 + algo: + name: a2c_continuous + + model: + name: continuous_a2c_logstd + + network: + name: actor_critic + separate: False + space: + continuous: + mu_activation: None + sigma_activation: None + mu_init: + name: default + sigma_init: + name: const_initializer + val: 0 + fixed_sigma: True + mlp: + units: [256,128,64] + d2rl: False + activation: elu + initializer: + name: default + scale: 2 + rnn: + name: gru + units: 64 + layers: 1 + # before_mlp: False + # layer_norm: True + config: + name: arl_robot_1_navigation + env_name: rlgpu + device: 'cuda:0' + device_name: 'cuda:0' + env_config: + num_envs: 8192 + + reward_shaper: + # min_val: -1 + scale_value: 0.1 + + normalize_advantage: True + gamma: 0.98 + tau: 0.95 + ppo: True + learning_rate: 1e-4 + lr_schedule: adaptive + kl_threshold: 0.016 + save_best_after: 10 + score_to_win: 100000 + grad_norm: 1.0 + entropy_coef: 0 + truncate_grads: True + e_clip: 0.2 + clip_value: False + num_actors: 1024 + horizon_length: 32 + minibatch_size: 2048 + mini_epochs: 4 + critic_coef: 2 + normalize_input: True + bounds_loss_coef: 0.0001 + max_epochs: 1500 + normalize_value: True + use_diagnostics: True + value_bootstrap: True + #weight_decay: 0.0001 + use_smooth_clamp: False + + player: + render: True + deterministic: True + games_num: 100000 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 00000000000..6a33b1c3032 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,37 @@ +# 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 isaaclab.utils import configclass + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg + + +@configclass +class NavigationEnvPPORunnerCfg(RslRlOnPolicyRunnerCfg): + num_steps_per_env = 24 + max_iterations = 1500 + save_interval = 50 + experiment_name = "arl_robot_1_navigation" + empirical_normalization = False + policy = RslRlPpoActorCriticCfg( + init_noise_std=0.5, + actor_hidden_dims=[256, 128, 64], + critic_hidden_dims=[256, 128, 64], + activation="elu", + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.001, + num_learning_epochs=4, + num_mini_batches=4, + learning_rate=4.0e-4, + schedule="adaptive", + gamma=0.98, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/skrl_rough_ppo_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/skrl_rough_ppo_cfg.yaml new file mode 100644 index 00000000000..6a56ce422d7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/agents/skrl_rough_ppo_cfg.yaml @@ -0,0 +1,95 @@ +# 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 + +seed: 42 + + +# Models are instantiated using skrl's model instantiator utility +# https://skrl.readthedocs.io/en/latest/api/utils/model_instantiators.html +models: + separate: False + policy: + class: GaussianMixin + clip_actions: False + clip_log_std: True + min_log_std: -20.0 + max_log_std: 2.0 + initial_log_std: 0.0 + network: + - name: mlp + input: STATES + layers: [256, 128, 64] + activations: elu + - name: gru + input: mlp + type: GRU + layers: [64] + num_layers: 1 + output: ACTIONS + value: + class: DeterministicMixin + clip_actions: False + network: + - name: mlp + input: STATES + layers: [256, 128, 64] + activations: elu + - name: gru + input: mlp + type: GRU + layers: [64] + num_layers: 1 + output: ONE + + +# Rollout memory +# https://skrl.readthedocs.io/en/latest/api/memories/random.html +memory: + class: RandomMemory + memory_size: -1 # automatically determined (same as agent:rollouts) + + +# PPO agent configuration (field names are from PPO_DEFAULT_CONFIG) +# https://skrl.readthedocs.io/en/latest/api/agents/ppo.html +agent: + class: PPO + rollouts: 24 + learning_epochs: 5 + mini_batches: 4 + discount_factor: 0.99 + lambda: 0.95 + learning_rate: 1.0e-03 + learning_rate_scheduler: KLAdaptiveLR + learning_rate_scheduler_kwargs: + kl_threshold: 0.01 + state_preprocessor: null + state_preprocessor_kwargs: null + value_preprocessor: RunningStandardScaler + value_preprocessor_kwargs: null + random_timesteps: 0 + learning_starts: 0 + grad_norm_clip: 1.0 + ratio_clip: 0.2 + value_clip: 0.2 + clip_predicted_values: True + entropy_loss_scale: 0.005 + value_loss_scale: 1.0 + kl_threshold: 0.0 + rewards_shaper_scale: 0.6 + time_limit_bootstrap: False + # logging and checkpoint + experiment: + directory: "arl_robot_1_navigation" + experiment_name: "" + write_interval: auto + checkpoint_interval: auto + + +# Sequential trainer +# https://skrl.readthedocs.io/en/latest/api/trainers/sequential.html +trainer: + class: SequentialTrainer + timesteps: 36000 + environment_info: log diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/floating_obstacles_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/floating_obstacles_env_cfg.py new file mode 100644 index 00000000000..0d8e42a370d --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/floating_obstacles_env_cfg.py @@ -0,0 +1,41 @@ +# 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 + +## +# Pre-defined configs +## +from isaaclab_assets.robots.arl_robot_1 import ARL_ROBOT_1_CFG + +from isaaclab.utils import configclass + +from isaaclab_tasks.manager_based.drone_arl.navigation.config.arl_robot_1.navigation_env_cfg import ( + NavigationVelocityFloatingObstacleEnvCfg, +) + + +@configclass +class FloatingObstacleEnvCfg(NavigationVelocityFloatingObstacleEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + # switch robot to arl_robot_1 + self.scene.robot = ARL_ROBOT_1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + self.scene.robot.actuators["thrusters"].dt = self.sim.dt + + +@configclass +class FloatingObstacleEnvCfg_PLAY(FloatingObstacleEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + self.curriculum.obstacle_levels.params["max_difficulty"] = 40 + self.curriculum.obstacle_levels.params["min_difficulty"] = 39 + + # disable randomization for play + self.observations.policy.enable_corruption = False + # remove random pushing event + self.events.base_external_force_torque = None + self.events.push_robot = None diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/navigation_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/navigation_env_cfg.py new file mode 100644 index 00000000000..b4bb1792a09 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/navigation_env_cfg.py @@ -0,0 +1,314 @@ +# 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 math +from dataclasses import MISSING + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg, MultirotorCfg +from isaaclab.controllers.lee_velocity_control_cfg import LeeVelControllerCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import CurriculumTermCfg as CurrTerm +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors import ContactSensorCfg +from isaaclab.sensors.ray_caster.multi_mesh_ray_caster_camera_cfg import MultiMeshRayCasterCameraCfg +from isaaclab.sensors.ray_caster.patterns import PinholeCameraPatternCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise + +import isaaclab_tasks.manager_based.drone_arl.mdp as mdp + +## +# Pre-defined configs +## +from .obstacles.obstacle_scene import ( + OBSTACLE_SCENE_CFG, + generate_obstacle_collection, + reset_obstacles_with_individual_ranges, +) + + +## +# Scene definition +## +@configclass +class MySceneCfg(InteractiveSceneCfg): + """Scene configuration for drone navigation with obstacles.""" + + # obstacles + object_collection = generate_obstacle_collection(OBSTACLE_SCENE_CFG) + + # robots + robot: MultirotorCfg = MISSING + + # sensors + depth_camera = MultiMeshRayCasterCameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/base_link", + mesh_prim_paths=[ + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + target_prim_expr=f"{{ENV_REGEX_NS}}/obstacle_{wall_name}", is_global=False, track_mesh_transforms=True + ) + for wall_name, _ in OBSTACLE_SCENE_CFG.wall_cfgs.items() + ] + + [ + MultiMeshRayCasterCameraCfg.RaycastTargetCfg( + target_prim_expr=f"{{ENV_REGEX_NS}}/obstacle_{i}", is_global=False, track_mesh_transforms=True + ) + for i in range(OBSTACLE_SCENE_CFG.max_num_obstacles) + ], + offset=MultiMeshRayCasterCameraCfg.OffsetCfg( + pos=(0.15, 0.0, 0.04), rot=(1.0, 0.0, 0.0, 0.0), convention="world" + ), + update_period=0.1, + pattern_cfg=PinholeCameraPatternCfg( + width=480, height=270, focal_length=0.193, horizontal_aperture=0.36, vertical_aperture=0.21 + ), + data_types=["distance_to_image_plane"], + max_distance=10.0, + depth_clipping_behavior="max", + ) + + contact_forces = ContactSensorCfg( + prim_path="{ENV_REGEX_NS}/Robot/.*", + update_period=0.0, + history_length=10, + debug_vis=False, + ) + # lights + sky_light = AssetBaseCfg( + prim_path="/World/skyLight", + spawn=sim_utils.DomeLightCfg( + intensity=750.0, + texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr", + ), + ) + + +## +# MDP settings +## + + +@configclass +class CommandsCfg: + """Command specifications for the MDP.""" + + target_pose = mdp.DroneUniformPoseCommandCfg( + asset_name="robot", + body_name="base_link", + resampling_time_range=(10.0, 10.0), + debug_vis=True, + ranges=mdp.DroneUniformPoseCommandCfg.Ranges( + pos_x=(4.0, 5.0), + pos_y=(-3.0, 3.0), + pos_z=(1.0, 5.0), + roll=(-0.0, 0.0), + pitch=(-0.0, 0.0), + yaw=(-0.0, 0.0), + ), + ) + + +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + velocity_commands = mdp.NavigationActionCfg( + asset_name="robot", + scale=1.0, + offset=0.0, + preserve_order=False, + use_default_offset=False, + command_type="vel", + controller_cfg=LeeVelControllerCfg(), + ) + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group.""" + + # observation terms (order preserved) + base_link_position = ObsTerm( + func=mdp.generated_drone_commands, + params={"command_name": "target_pose", "asset_cfg": SceneEntityCfg("robot")}, + noise=Unoise(n_min=-0.1, n_max=0.1), + ) + base_roll_pitch = ObsTerm(func=mdp.base_roll_pitch, noise=Unoise(n_min=-0.1, n_max=0.1)) + base_lin_vel = ObsTerm(func=mdp.base_lin_vel, noise=Unoise(n_min=-0.1, n_max=0.1)) + base_ang_vel = ObsTerm(func=mdp.base_ang_vel, noise=Unoise(n_min=-0.1, n_max=0.1)) + last_action = ObsTerm(func=mdp.last_action_navigation, noise=Unoise(n_min=-0.0, n_max=0.0)) + depth_latent = ObsTerm( + func=mdp.image_latents, + params={"sensor_cfg": SceneEntityCfg("depth_camera"), "data_type": "distance_to_image_plane"}, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + # observation groups + policy: PolicyCfg = PolicyCfg() + + +@configclass +class EventCfg: + """Configuration for events.""" + + # reset + + reset_base = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": { + "x": (-5.0, -4.5), + "y": (-3.0, 3.0), + "z": (1.0, 5.0), + "yaw": (-math.pi / 6.0, math.pi / 6.0), + }, + "velocity_range": { + "x": (-0.2, 0.2), + "y": (-0.2, 0.2), + "z": (-0.2, 0.2), + "roll": (-0.2, 0.2), + "pitch": (-0.2, 0.2), + "yaw": (-0.2, 0.2), + }, + }, + ) + + reset_obstacles = EventTerm( + func=reset_obstacles_with_individual_ranges, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("object_collection"), + "obstacle_configs": OBSTACLE_SCENE_CFG.obstacle_cfgs, + "wall_configs": OBSTACLE_SCENE_CFG.wall_cfgs, + "env_size": OBSTACLE_SCENE_CFG.env_size, + "use_curriculum": True, + "min_num_obstacles": OBSTACLE_SCENE_CFG.min_num_obstacles, + "max_num_obstacles": OBSTACLE_SCENE_CFG.max_num_obstacles, + "ground_offset": OBSTACLE_SCENE_CFG.ground_offset, + }, + ) + + +@configclass +class RewardsCfg: + """Reward terms for the MDP.""" + + goal_dist_exp1 = RewTerm( + func=mdp.distance_to_goal_exp_curriculum, + weight=2.0, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "std": 7.0, + "command_name": "target_pose", + }, + ) + goal_dist_exp2 = RewTerm( + func=mdp.distance_to_goal_exp_curriculum, + weight=4.0, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "std": 0.5, + "command_name": "target_pose", + }, + ) + velocity_reward = RewTerm( + func=mdp.velocity_to_goal_reward_curriculum, + weight=0.5, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "command_name": "target_pose", + }, + ) + action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.05) + action_magnitude_l2 = RewTerm(func=mdp.action_l2, weight=-0.05) + + termination_penalty = RewTerm( + func=mdp.is_terminated, + weight=-100.0, + ) + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + collision = DoneTerm( + func=mdp.illegal_contact, + params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*"), "threshold": 1.0}, + time_out=False, + ) + + +@configclass +class CurriculumCfg: + """Curriculum terms for the MDP.""" + + obstacle_levels = CurrTerm( + func=mdp.obstacle_density_curriculum, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "max_difficulty": 10, + "min_difficulty": 0, + }, + ) + + +## +# Environment configuration +## + + +@configclass +class NavigationVelocityFloatingObstacleEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the locomotion velocity-tracking environment.""" + + # Scene settings + scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=20.5) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + # MDP settings + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + curriculum: CurriculumCfg = CurriculumCfg() + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 10 + self.episode_length_s = 10.0 + # simulation settings + self.sim.dt = 0.01 + self.sim.render_interval = self.decimation + self.sim.physics_material = sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + static_friction=1.0, + dynamic_friction=1.0, + ) + self.sim.physx.gpu_max_rigid_patch_count = 2**21 + # update sensor update periods + # we tick all the sensors based on the smallest update period (physics update period) + if self.scene.contact_forces is not None: + self.scene.contact_forces.update_period = self.sim.dt diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/obstacles/obstacle_scene.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/obstacles/obstacle_scene.py new file mode 100644 index 00000000000..dba04ab7d11 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/obstacles/obstacle_scene.py @@ -0,0 +1,279 @@ +# 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 numpy as np +import torch + +import isaaclab.sim as sim_utils +import isaaclab.utils.math as math_utils +from isaaclab.assets import RigidObjectCfg, RigidObjectCollection, RigidObjectCollectionCfg +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers import SceneEntityCfg + +from .obstacle_scene_cfg import ObstaclesSceneCfg + +"""Obstacle scene generation and reset functionality for drone navigation environments. + +This module provides utilities for generating dynamic 3D obstacle courses with walls and +floating obstacles. The obstacle configurations support curriculum learning where difficulty +can be progressively increased by adjusting the number of active obstacles. +""" + +OBSTACLE_SCENE_CFG = ObstaclesSceneCfg( + env_size=(12.0, 8.0, 6.0), + min_num_obstacles=20, + max_num_obstacles=40, + ground_offset=3.0, +) + + +def generate_obstacle_collection(cfg: ObstaclesSceneCfg) -> RigidObjectCollectionCfg: + """Generate a rigid object collection configuration for walls and obstacles. + + Creates a complete scene with boundary walls and a variety of floating obstacles + (panels, cubes, rods, etc.) based on the provided configuration. Each obstacle is + assigned random colors and configured with appropriate physics properties. + + Wall objects are configured with very high mass (10^7 kg) and high damping to remain + stationary during collisions. Obstacle objects have moderate mass (100 kg) to move in the right position if reset + in collision. + + Args: + cfg: Configuration object specifying obstacle types, sizes, quantities, and + positioning constraints. + + Returns: + A RigidObjectCollectionCfg containing all wall and obstacle configurations, + ready to be added to a scene. + + Note: + All obstacles are initially placed at origin [0, 0, 0]. Actual positions are + set during environment reset via :func:`reset_obstacles_with_individual_ranges`. + """ + max_num_obstacles = cfg.max_num_obstacles + + rigid_objects = {} + + for wall_name, wall_cfg in cfg.wall_cfgs.items(): + # Walls get their specific size and default center + default_center = [0.0, 0.0, 0.0] # Will be set properly at reset + color = float(np.random.randint(0, 256, size=1, dtype=np.uint8)) / 255.0 + + rigid_objects[wall_name] = RigidObjectCfg( + prim_path=f"{{ENV_REGEX_NS}}/obstacle_{wall_name}", + spawn=sim_utils.CuboidCfg( + size=wall_cfg.size, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.5, 0.5, color), metallic=0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, + solver_velocity_iteration_count=0, + disable_gravity=True, + kinematic_enabled=False, + linear_damping=9999.0, + angular_damping=9999.0, + max_linear_velocity=0.0, + max_angular_velocity=0.0, + ), + # mass of walls needs to be way larger than weight of obstacles to make them not move during reset + mass_props=sim_utils.MassPropertiesCfg(mass=10000000.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=tuple(default_center)), + collision_group=0, + ) + + obstacle_types = list(cfg.obstacle_cfgs.values()) + for i in range(max_num_obstacles): + obj_name = f"obstacle_{i}" + obs_cfg = obstacle_types[i % len(obstacle_types)] + + default_center = [0.0, 0.0, 0.0] + color = np.random.randint(0, 256, size=3, dtype=np.uint8) + color_normalized = tuple(float(c) / 255.0 for c in color) + + rigid_objects[obj_name] = RigidObjectCfg( + prim_path=f"{{ENV_REGEX_NS}}/{obj_name}", + spawn=sim_utils.CuboidCfg( + size=obs_cfg.size, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=color_normalized, metallic=0.2), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, + solver_velocity_iteration_count=0, + disable_gravity=True, + kinematic_enabled=False, + linear_damping=1.0, + angular_damping=1.0, + max_linear_velocity=0.0, + max_angular_velocity=0.0, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=tuple(default_center)), + collision_group=0, + ) + + return RigidObjectCollectionCfg(rigid_objects=rigid_objects) + + +def reset_obstacles_with_individual_ranges( + env: ManagerBasedRLEnv, + env_ids: torch.Tensor, + asset_cfg: SceneEntityCfg, + obstacle_configs: dict, + wall_configs: dict, + env_size: tuple[float, float, float], + use_curriculum: bool = True, + min_num_obstacles: int = 1, + max_num_obstacles: int = 10, + ground_offset: float = 0.1, +) -> None: + """Reset obstacle and wall positions for specified environments without collision checking. + + This function repositions all walls and a curriculum-determined subset of obstacles + within the specified environment bounds. + + Walls are positioned at fixed locations based on their configuration ratios. Obstacles + are randomly placed within their designated zones, with the number of active obstacles + determined by the curriculum difficulty level. Inactive obstacles are moved far below + the scene (-1000m in Z) to effectively remove them from the environment. + + The curriculum scaling works as: + num_obstacles = min + (difficulty / max_difficulty) * (max - min) + + Args: + env: The manager-based RL environment instance. + env_ids: Tensor of environment indices to reset. + asset_cfg: Scene entity configuration identifying the obstacle collection. + obstacle_configs: Dictionary mapping obstacle type names to their BoxCfg + configurations, specifying size and placement ranges. + wall_configs: Dictionary mapping wall names to their BoxCfg configurations. + env_size: Tuple of (length, width, height) defining the environment bounds in meters. + use_curriculum: If True, number of obstacles scales with curriculum difficulty. + If False, spawns max_num_obstacles in every environment. Defaults to True. + min_num_obstacles: Minimum number of obstacles to spawn per environment. + Defaults to 1. + max_num_obstacles: Maximum number of obstacles to spawn per environment. + Defaults to 10. + ground_offset: Z-axis offset to prevent obstacles from spawning at z=0. + Defaults to 0.1 meters. + + Note: + This function expects the environment to have `_obstacle_difficulty_levels` and + `_max_obstacle_difficulty` attributes when `use_curriculum=True`. These are + typically set by :func:`obstacle_density_curriculum`. + """ + obstacles: RigidObjectCollection = env.scene[asset_cfg.name] + + num_objects = obstacles.num_objects + num_envs = len(env_ids) + object_names = obstacles.object_names + + # Get difficulty levels per environment + if use_curriculum and hasattr(env, "_obstacle_difficulty_levels"): + difficulty_levels = env._obstacle_difficulty_levels[env_ids] + max_difficulty = env._max_obstacle_difficulty + else: + difficulty_levels = torch.ones(num_envs, device=env.device) * max_num_obstacles + max_difficulty = max_num_obstacles + + # Calculate active obstacles per env based on difficulty + obstacles_per_env = ( + min_num_obstacles + (difficulty_levels / max_difficulty) * (max_num_obstacles - min_num_obstacles) + ).long() + + # Prepare tensors + all_poses = torch.zeros(num_envs, num_objects, 7, device=env.device) + all_velocities = torch.zeros(num_envs, num_objects, 6, device=env.device) + + wall_names = list(wall_configs.keys()) + obstacle_types = list(obstacle_configs.values()) + env_size_t = torch.tensor(env_size, device=env.device) + + # place walls + for wall_name, wall_cfg in wall_configs.items(): + if wall_name in object_names: + wall_idx = object_names.index(wall_name) + + min_ratio = torch.tensor(wall_cfg.center_ratio_min, device=env.device) + max_ratio = torch.tensor(wall_cfg.center_ratio_max, device=env.device) + + if torch.allclose(min_ratio, max_ratio): + center_ratios = min_ratio.unsqueeze(0).repeat(num_envs, 1) + else: + ratios = torch.rand(num_envs, 3, device=env.device) + center_ratios = ratios * (max_ratio - min_ratio) + min_ratio + + positions = (center_ratios - 0.5) * env_size_t + positions[:, 2] += ground_offset + positions += env.scene.env_origins[env_ids] + + all_poses[:, wall_idx, 0:3] = positions + all_poses[:, wall_idx, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=env.device).repeat(num_envs, 1) + + # Get obstacle indices + obstacle_indices = [idx for idx, name in enumerate(object_names) if name not in wall_names] + + if len(obstacle_indices) == 0: + obstacles.write_object_pose_to_sim(all_poses, env_ids=env_ids) + obstacles.write_object_velocity_to_sim(all_velocities, env_ids=env_ids) + return + + # Determine which obstacles are active per env + active_masks = torch.zeros(num_envs, len(obstacle_indices), dtype=torch.bool, device=env.device) + for env_idx in range(num_envs): + num_active = obstacles_per_env[env_idx].item() + perm = torch.randperm(len(obstacle_indices), device=env.device)[:num_active] + active_masks[env_idx, perm] = True + + # place obstacles + for obj_list_idx in range(len(obstacle_indices)): + obj_idx = obstacle_indices[obj_list_idx] + + # Which envs need this obstacle? + envs_need_obstacle = active_masks[:, obj_list_idx] + + if not envs_need_obstacle.any(): + # Move all to -1000 + all_poses[:, obj_idx, 0:3] = env.scene.env_origins[env_ids] + torch.tensor( + [0.0, 0.0, -1000.0], device=env.device + ) + all_poses[:, obj_idx, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=env.device) + continue + + # Get obstacle config + config_idx = obj_list_idx % len(obstacle_types) + obs_cfg = obstacle_types[config_idx] + + min_ratio = torch.tensor(obs_cfg.center_ratio_min, device=env.device) + max_ratio = torch.tensor(obs_cfg.center_ratio_max, device=env.device) + + # sample object positions + num_active_envs = envs_need_obstacle.sum().item() + ratios = torch.rand(num_active_envs, 3, device=env.device) + positions = (ratios * (max_ratio - min_ratio) + min_ratio - 0.5) * env_size_t + positions[:, 2] += ground_offset + + # Add env origins + active_env_indices = torch.where(envs_need_obstacle)[0] + positions += env.scene.env_origins[env_ids[active_env_indices]] + + # Generate quaternions + quats = math_utils.random_orientation(num_envs, device=env.device) + + # Write poses + all_poses[envs_need_obstacle, obj_idx, 0:3] = positions + all_poses[envs_need_obstacle, obj_idx, 3:7] = quats[envs_need_obstacle] + + # Move inactive obstacles far away + inactive = ~envs_need_obstacle + all_poses[inactive, obj_idx, 0:3] = env.scene.env_origins[env_ids[inactive]] + torch.tensor( + [0.0, 0.0, -1000.0], device=env.device + ) + all_poses[inactive, obj_idx, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=env.device) + + # Write to sim + obstacles.write_object_pose_to_sim(all_poses, env_ids=env_ids) + obstacles.write_object_velocity_to_sim(all_velocities, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/obstacles/obstacle_scene_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/obstacles/obstacle_scene_cfg.py new file mode 100644 index 00000000000..4bfa841cccd --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/navigation/config/arl_robot_1/obstacles/obstacle_scene_cfg.py @@ -0,0 +1,113 @@ +# 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 dataclasses import MISSING + +from isaaclab.utils import configclass + + +@configclass +class ObstaclesSceneCfg: + """Configuration for a terrain with floating obstacles.""" + + min_num_obstacles: int = 1 + max_num_obstacles: int = 40 + ground_offset: float = 3.0 + + env_size: tuple[float, float, float] = MISSING + + @configclass + class BoxCfg: + """Configuration for a box-shaped obstacle or wall. + + Defines the size and placement constraints for rectangular obstacles within + the environment. The center position is specified as ratios of the environment + size, allowing for flexible scaling. + + Attributes: + size: Tuple of (length, width, height) in meters. + center_ratio_min: Minimum position as ratio of env_size (0.0 to 1.0) for + each axis. Used for random placement bounds. + center_ratio_max: Maximum position as ratio of env_size (0.0 to 1.0) for + each axis. For fixed positions, set equal to center_ratio_min. + """ + + size: tuple[float, float, float] = MISSING + center_ratio_min: tuple[float, float, float] = MISSING + center_ratio_max: tuple[float, float, float] = MISSING + + # Obstacle configurations + panel_obs_cfg = BoxCfg() + panel_obs_cfg.size = (0.1, 1.2, 3.0) + panel_obs_cfg.center_ratio_min = (0.3, 0.05, 0.05) + panel_obs_cfg.center_ratio_max = (0.85, 0.95, 0.95) + + small_wall_obs_cfg = BoxCfg() + small_wall_obs_cfg.size = (0.1, 0.5, 0.5) + small_wall_obs_cfg.center_ratio_min = (0.3, 0.05, 0.05) + small_wall_obs_cfg.center_ratio_max = (0.85, 0.9, 0.9) + + big_wall_obs_cfg = BoxCfg() + big_wall_obs_cfg.size = (0.1, 1.0, 1.0) + big_wall_obs_cfg.center_ratio_min = (0.3, 0.05, 0.05) + big_wall_obs_cfg.center_ratio_max = (0.85, 0.9, 0.9) + + small_cube_obs_cfg = BoxCfg() + small_cube_obs_cfg.size = (0.4, 0.4, 0.4) + small_cube_obs_cfg.center_ratio_min = (0.3, 0.05, 0.05) + small_cube_obs_cfg.center_ratio_max = (0.85, 0.9, 0.9) + + rod_obs_cfg = BoxCfg() + rod_obs_cfg.size = (0.1, 0.1, 2.0) + rod_obs_cfg.center_ratio_min = (0.3, 0.05, 0.05) + rod_obs_cfg.center_ratio_max = (0.85, 0.9, 0.9) + + # Wall configurations + left_wall_cfg = BoxCfg() + left_wall_cfg.size = (12.0, 0.2, 6.0) + left_wall_cfg.center_ratio_min = (0.5, 1.0, 0.5) + left_wall_cfg.center_ratio_max = (0.5, 1.0, 0.5) + + right_wall_cfg = BoxCfg() + right_wall_cfg.size = (12.0, 0.2, 6.0) + right_wall_cfg.center_ratio_min = (0.5, 0.0, 0.5) + right_wall_cfg.center_ratio_max = (0.5, 0.0, 0.5) + + back_wall_cfg = BoxCfg() + back_wall_cfg.size = (0.2, 8.0, 6.0) + back_wall_cfg.center_ratio_min = (0.0, 0.5, 0.5) + back_wall_cfg.center_ratio_max = (0.0, 0.5, 0.5) + + front_wall_cfg = BoxCfg() + front_wall_cfg.size = (0.2, 8.0, 6.0) + front_wall_cfg.center_ratio_min = (1.0, 0.5, 0.5) + front_wall_cfg.center_ratio_max = (1.0, 0.5, 0.5) + + top_wall_cfg = BoxCfg() + top_wall_cfg.size = (12.0, 8.0, 0.2) + top_wall_cfg.center_ratio_min = (0.5, 0.5, 1.0) + top_wall_cfg.center_ratio_max = (0.5, 0.5, 1.0) + + bottom_wall_cfg = BoxCfg() + bottom_wall_cfg.size = (12.0, 8.0, 0.2) + bottom_wall_cfg.center_ratio_min = (0.5, 0.5, 0.0) + bottom_wall_cfg.center_ratio_max = (0.5, 0.5, 0.0) + + wall_cfgs = { + "left_wall": left_wall_cfg, + "right_wall": right_wall_cfg, + "back_wall": back_wall_cfg, + "front_wall": front_wall_cfg, + "bottom_wall": bottom_wall_cfg, + "top_wall": top_wall_cfg, + } + + obstacle_cfgs = { + "panel": panel_obs_cfg, + "small_wall": small_wall_obs_cfg, + "big_wall": big_wall_obs_cfg, + "small_cube": small_cube_obs_cfg, + "rod": rod_obs_cfg, + } diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/robot_model b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/robot_model new file mode 160000 index 00000000000..f53d8e81450 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/robot_model @@ -0,0 +1 @@ +Subproject commit f53d8e814503b44c379ba72ff5e1f691b48a4c28 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/__init__.py new file mode 100644 index 00000000000..9a37793ea66 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Drone ARL state-based control environments.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/__init__.py new file mode 100644 index 00000000000..601fefecac2 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/__init__.py @@ -0,0 +1,6 @@ +# 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 + +"""Configurations for state-based control environments.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/__init__.py new file mode 100644 index 00000000000..1657ec3061f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/__init__.py @@ -0,0 +1,36 @@ +# 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 gymnasium as gym + +from . import agents + +## +# Register Gym environments. +## + +gym.register( + id="Isaac-StateBasedControl-ARL-robot-1-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.empty_env_cfg:EmptyEnvCfg", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:StateBasedControlEnvPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", + }, +) + +gym.register( + id="Isaac-StateBasedControl-ARL-robot-1-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.empty_env_cfg:EmptyEnvCfg_PLAY", + "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:StateBasedControlEnvPPORunnerCfg", + "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", + }, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/__init__.py new file mode 100644 index 00000000000..2e924fbf1b1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/__init__.py @@ -0,0 +1,4 @@ +# 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 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/rl_games_ppo_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/rl_games_ppo_cfg.yaml new file mode 100644 index 00000000000..b23c75e2816 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/rl_games_ppo_cfg.yaml @@ -0,0 +1,87 @@ +# 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 + +params: + seed: 42 + + # environment wrapper clipping + env: + clip_actions: 1.0 + algo: + name: a2c_continuous + + model: + name: continuous_a2c_logstd + + network: + name: actor_critic + separate: False + space: + continuous: + mu_activation: None + sigma_activation: None + mu_init: + name: default + sigma_init: + name: const_initializer + val: 0 + fixed_sigma: True + mlp: + units: [256,128,64] + d2rl: False + activation: elu + initializer: + name: default + scale: 2 + rnn: + name: gru + units: 64 + layers: 1 + # before_mlp: False + # layer_norm: True + config: + env_name: rlgpu + device: 'cuda:0' + device_name: 'cuda:0' + env_config: + num_envs: 8192 + + name: arl_robot_1_state_based_control + reward_shaper: + # min_val: -1 + scale_value: 0.1 + + normalize_advantage: True + gamma: 0.98 + tau: 0.95 + ppo: True + learning_rate: 1e-4 + lr_schedule: adaptive + kl_threshold: 0.016 + save_best_after: 10 + score_to_win: 100000 + grad_norm: 1.0 + entropy_coef: 0 + truncate_grads: True + e_clip: 0.2 + clip_value: False + num_actors: 1024 + horizon_length: 32 + minibatch_size: 2048 + mini_epochs: 4 + critic_coef: 2 + normalize_input: True + bounds_loss_coef: 0.0001 + max_epochs: 1500 + normalize_value: True + use_diagnostics: True + value_bootstrap: True + #weight_decay: 0.0001 + use_smooth_clamp: False + + player: + render: True + deterministic: True + games_num: 100000 diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 00000000000..b4fab7d138a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,37 @@ +# 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 isaaclab.utils import configclass + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg + + +@configclass +class StateBasedControlEnvPPORunnerCfg(RslRlOnPolicyRunnerCfg): + num_steps_per_env = 24 + max_iterations = 1500 + save_interval = 50 + experiment_name = "arl_robot_1_state_based_control" + empirical_normalization = False + policy = RslRlPpoActorCriticCfg( + init_noise_std=0.5, + actor_hidden_dims=[256, 128, 64], + critic_hidden_dims=[256, 128, 64], + activation="elu", + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.001, + num_learning_epochs=4, + num_mini_batches=4, + learning_rate=4.0e-4, + schedule="adaptive", + gamma=0.98, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/skrl_ppo_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/skrl_ppo_cfg.yaml new file mode 100644 index 00000000000..1c4b131d0c5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/agents/skrl_ppo_cfg.yaml @@ -0,0 +1,95 @@ +# 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 + +seed: 42 + + +# Models are instantiated using skrl's model instantiator utility +# https://skrl.readthedocs.io/en/latest/api/utils/model_instantiators.html +models: + separate: False + policy: + class: GaussianMixin + clip_actions: False + clip_log_std: True + min_log_std: -20.0 + max_log_std: 2.0 + initial_log_std: 0.0 + network: + - name: mlp + input: STATES + layers: [256, 128, 64] + activations: elu + - name: gru + input: mlp + type: GRU + layers: [64] + num_layers: 1 + output: ACTIONS + value: + class: DeterministicMixin + clip_actions: False + network: + - name: mlp + input: STATES + layers: [256, 128, 64] + activations: elu + - name: gru + input: mlp + type: GRU + layers: [64] + num_layers: 1 + output: ONE + + +# Rollout memory +# https://skrl.readthedocs.io/en/latest/api/memories/random.html +memory: + class: RandomMemory + memory_size: -1 # automatically determined (same as agent:rollouts) + + +# PPO agent configuration (field names are from PPO_DEFAULT_CONFIG) +# https://skrl.readthedocs.io/en/latest/api/agents/ppo.html +agent: + class: PPO + rollouts: 24 + learning_epochs: 5 + mini_batches: 4 + discount_factor: 0.99 + lambda: 0.95 + learning_rate: 1.0e-03 + learning_rate_scheduler: KLAdaptiveLR + learning_rate_scheduler_kwargs: + kl_threshold: 0.01 + state_preprocessor: null + state_preprocessor_kwargs: null + value_preprocessor: RunningStandardScaler + value_preprocessor_kwargs: null + random_timesteps: 0 + learning_starts: 0 + grad_norm_clip: 1.0 + ratio_clip: 0.2 + value_clip: 0.2 + clip_predicted_values: True + entropy_loss_scale: 0.005 + value_loss_scale: 1.0 + kl_threshold: 0.0 + rewards_shaper_scale: 0.6 + time_limit_bootstrap: False + # logging and checkpoint + experiment: + directory: "arl_robot_1_state_based_control" + experiment_name: "" + write_interval: auto + checkpoint_interval: auto + + +# Sequential trainer +# https://skrl.readthedocs.io/en/latest/api/trainers/sequential.html +trainer: + class: SequentialTrainer + timesteps: 36000 + environment_info: log diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/empty_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/empty_env_cfg.py new file mode 100644 index 00000000000..e9cb2a98218 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/empty_env_cfg.py @@ -0,0 +1,41 @@ +# 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 isaaclab_assets.robots.arl_robot_1 import ARL_ROBOT_1_CFG + +from isaaclab.utils import configclass + +from .state_based_control_env_cfg import StateBasedControlEmptyEnvCfg + +## +# Pre-defined configs +## + + +@configclass +class EmptyEnvCfg(StateBasedControlEmptyEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + # switch robot to arl_robot_1 + self.scene.robot = ARL_ROBOT_1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + self.scene.robot.actuators["thrusters"].dt = self.sim.dt + + +@configclass +class EmptyEnvCfg_PLAY(EmptyEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # make a smaller scene for play + self.scene.num_envs = 50 + self.scene.env_spacing = 2.5 + + # disable randomization for play + self.observations.policy.enable_corruption = False + # remove random pushing event + self.events.base_external_force_torque = None + self.events.push_robot = None diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/state_based_control_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/state_based_control_env_cfg.py new file mode 100644 index 00000000000..cfc6a8d0bde --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/drone_arl/state_based_control/config/arl_robot_1/state_based_control_env_cfg.py @@ -0,0 +1,229 @@ +# 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 math +from dataclasses import MISSING + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg, MultirotorCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg + +# from isaaclab.sim import PinholeCameraCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise + +import isaaclab_tasks.manager_based.drone_arl.mdp as mdp + + +## +# Scene definition +## +@configclass +class MySceneCfg(InteractiveSceneCfg): + """Configuration for the terrain scene with a flying robot.""" + + # robots + robot: MultirotorCfg = MISSING + + # lights + sky_light = AssetBaseCfg( + prim_path="/World/skyLight", + spawn=sim_utils.DomeLightCfg( + intensity=750.0, + texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr", + ), + ) + + +## +# MDP settings +## + + +@configclass +class CommandsCfg: + """Command specifications for the MDP.""" + + target_pose = mdp.DroneUniformPoseCommandCfg( + asset_name="robot", + body_name="base_link", + resampling_time_range=(10.0, 10.0), + debug_vis=True, + ranges=mdp.DroneUniformPoseCommandCfg.Ranges( + pos_x=(-0.0, 0.0), + pos_y=(-0.0, 0.0), + pos_z=(-0.0, 0.0), + roll=(-0.0, 0.0), + pitch=(-0.0, 0.0), + yaw=(-0.0, 0.0), + ), + ) + + +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + thrust_command = mdp.ThrustActionCfg( + asset_name="robot", + scale=3.0, + offset=3.0, + preserve_order=False, + use_default_offset=False, + clip={ + "back_left_prop": (0.0, 6.0), + "back_right_prop": (0.0, 6.0), + "front_left_prop": (0.0, 6.0), + "front_right_prop": (0.0, 6.0), + }, + ) + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group.""" + + # observation terms (order preserved) + base_link_position = ObsTerm(func=mdp.root_pos_w, noise=Unoise(n_min=-0.1, n_max=0.1)) + base_orientation = ObsTerm(func=mdp.root_quat_w, noise=Unoise(n_min=-0.1, n_max=0.1)) + base_lin_vel = ObsTerm(func=mdp.base_lin_vel, noise=Unoise(n_min=-0.1, n_max=0.1)) + base_ang_vel = ObsTerm(func=mdp.base_ang_vel, noise=Unoise(n_min=-0.1, n_max=0.1)) + last_action = ObsTerm(func=mdp.last_action, noise=Unoise(n_min=-0.0, n_max=0.0)) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + # observation groups + policy: PolicyCfg = PolicyCfg() + + +@configclass +class EventCfg: + """Configuration for events.""" + + # reset + + reset_base = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": { + "x": (-1.0, 1.0), + "y": (-1.0, 1.0), + "z": (-1.0, 1.0), + "yaw": (-math.pi / 6.0, math.pi / 6.0), + "roll": (-math.pi / 6.0, math.pi / 6.0), + "pitch": (-math.pi / 6.0, math.pi / 6.0), + }, + "velocity_range": { + "x": (-0.2, 0.2), + "y": (-0.2, 0.2), + "z": (-0.2, 0.2), + "roll": (-0.2, 0.2), + "pitch": (-0.2, 0.2), + "yaw": (-0.2, 0.2), + }, + }, + ) + + +@configclass +class RewardsCfg: + """Reward terms for the MDP.""" + + distance_to_goal_exp = RewTerm( + func=mdp.distance_to_goal_exp, + weight=25.0, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "std": 1.5, + "command_name": "target_pose", + }, + ) + flat_orientation_l2 = RewTerm( + func=mdp.flat_orientation_l2, + weight=1.0, + params={"asset_cfg": SceneEntityCfg("robot")}, + ) + yaw_aligned = RewTerm( + func=mdp.yaw_aligned, + weight=2.0, + params={"asset_cfg": SceneEntityCfg("robot"), "std": 1.0}, + ) + lin_vel_xyz_exp = RewTerm( + func=mdp.lin_vel_xyz_exp, + weight=2.5, + params={"asset_cfg": SceneEntityCfg("robot"), "std": 2.0}, + ) + ang_vel_xyz_exp = RewTerm( + func=mdp.ang_vel_xyz_exp, + weight=10.0, + params={"asset_cfg": SceneEntityCfg("robot"), "std": 10.0}, + ) + action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.05) + action_magnitude_l2 = RewTerm(func=mdp.action_l2, weight=-0.05) + + termination_penalty = RewTerm( + func=mdp.is_terminated, + weight=-5.0, + ) + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + time_out = DoneTerm(func=mdp.time_out, time_out=True) + crash = DoneTerm(func=mdp.root_height_below_minimum, params={"minimum_height": -3.0}) + + +## +# Environment configuration +## + + +@configclass +class StateBasedControlEmptyEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the state-based drone pose-control environment.""" + + # Scene settings + scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=2.5) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + # MDP settings + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 10 + self.episode_length_s = 5.0 + # simulation settings + self.sim.dt = 0.01 + self.sim.render_interval = self.decimation + self.sim.physics_material = sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + static_friction=1.0, + dynamic_friction=1.0, + ) + self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15