From 9516cbc7cb94499a84568bd8397da671d822a3c9 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Mon, 23 Feb 2026 11:29:14 -0800 Subject: [PATCH 1/4] add cloner --- source/isaaclab/docs/CHANGELOG.rst | 1 - .../isaaclab/isaaclab/sim/schemas/schemas.py | 24 +- source/isaaclab/isaaclab/sim/utils/prims.py | 14 +- .../isaaclab/terrains/terrain_importer.py | 1 + source/isaaclab_physx/test/test_cloner.py | 217 ++++++++++++++++++ 5 files changed, 239 insertions(+), 18 deletions(-) create mode 100644 source/isaaclab_physx/test/test_cloner.py diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index f067bc8c8a4..8411cba3343 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/source/isaaclab/docs/CHANGELOG.rst @@ -25,7 +25,6 @@ Added (different prototypes across envs), replication switches to per-object instead of whole-env cloning. This reduces PhysX cloning time in heterogeneous scenes. - 4.0.0 (2026-02-22) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab/isaaclab/sim/schemas/schemas.py b/source/isaaclab/isaaclab/sim/schemas/schemas.py index b76c896e6ef..04791e5b8ec 100644 --- a/source/isaaclab/isaaclab/sim/schemas/schemas.py +++ b/source/isaaclab/isaaclab/sim/schemas/schemas.py @@ -156,7 +156,7 @@ def modify_articulation_root_properties( # check if prim has articulation applied on it if not UsdPhysics.ArticulationRootAPI(articulation_prim): return False - # ensure PhysX articulation API is applied + # ensure PhysX articulation API is applied (string-based, no PhysxSchema import) applied_schemas = articulation_prim.GetAppliedSchemas() if "PhysxArticulationAPI" not in applied_schemas: articulation_prim.AddAppliedSchema("PhysxArticulationAPI") @@ -318,7 +318,7 @@ def modify_rigid_body_properties( return False # retrieve the USD rigid-body api usd_rigid_body_api = UsdPhysics.RigidBodyAPI(rigid_body_prim) - # ensure PhysX rigid body API is applied + # ensure PhysX rigid body API is applied (string-based, no PhysxSchema import) applied_schemas = rigid_body_prim.GetAppliedSchemas() if "PhysxRigidBodyAPI" not in applied_schemas: rigid_body_prim.AddAppliedSchema("PhysxRigidBodyAPI") @@ -415,7 +415,7 @@ def modify_collision_properties( return False # retrieve the USD collision api usd_collision_api = UsdPhysics.CollisionAPI(collider_prim) - # ensure PhysX collision API is applied + # ensure PhysX collision API is applied (string-based, no PhysxSchema import) applied_schemas = collider_prim.GetAppliedSchemas() if "PhysxCollisionAPI" not in applied_schemas: collider_prim.AddAppliedSchema("PhysxCollisionAPI") @@ -569,7 +569,7 @@ def activate_contact_sensors(prim_path: str, threshold: float = 0.0, stage: Usd. # if a prim has a rigid body API, it is a rigid body and we don't need to # check its children if child_prim.HasAPI(UsdPhysics.RigidBodyAPI): - # set sleep threshold to zero + # set sleep threshold to zero (string-based, no PhysxSchema import) child_applied = child_prim.GetAppliedSchemas() if "PhysxRigidBodyAPI" not in child_applied: child_prim.AddAppliedSchema("PhysxRigidBodyAPI") @@ -651,9 +651,9 @@ def modify_joint_drive_properties( drive_api_name = "linear" else: return False - # check that prim is not a tendon child prim - applied_schemas_str = str(prim.GetAppliedSchemas()) - if "PhysxTendonAxisAPI" in applied_schemas_str and "PhysxTendonAxisRootAPI" not in applied_schemas_str: + # check that prim is not a tendon child prim (string-based, no PhysxSchema import) + applied_schemas = prim.GetAppliedSchemas() + if "PhysxTendonAxisAPI" in applied_schemas and "PhysxTendonAxisRootAPI" not in applied_schemas: return False # check if prim has joint drive applied on it @@ -661,7 +661,7 @@ def modify_joint_drive_properties( if not usd_drive_api: usd_drive_api = UsdPhysics.DriveAPI.Apply(prim, drive_api_name) # ensure PhysX joint API is applied - if "PhysxJointAPI" not in applied_schemas_str: + if "PhysxJointAPI" not in applied_schemas: prim.AddAppliedSchema("PhysxJointAPI") # mapping from configuration name to USD attribute name @@ -744,7 +744,7 @@ def modify_fixed_tendon_properties( # get USD prim tendon_prim = stage.GetPrimAtPath(prim_path) - # check if prim has fixed tendon applied on it + # check if prim has fixed tendon applied on it (string-based, no PhysxSchema import) applied_schemas = tendon_prim.GetAppliedSchemas() if not any("PhysxTendonAxisRootAPI" in s for s in applied_schemas): return False @@ -809,7 +809,7 @@ def modify_spatial_tendon_properties( stage = get_current_stage() # get USD prim tendon_prim = stage.GetPrimAtPath(prim_path) - # check if prim has spatial tendon applied on it + # check if prim has spatial tendon applied on it (string-based, no PhysxSchema import) applied_schemas = tendon_prim.GetAppliedSchemas() has_spatial = any( "PhysxTendonAttachmentRootAPI" in s or "PhysxTendonAttachmentLeafAPI" in s for s in applied_schemas @@ -884,7 +884,7 @@ def define_deformable_body_properties( # get deformable-body USD prim mesh_prim = matching_prims[0] - # ensure PhysX deformable body API is applied + # ensure PhysX deformable body API is applied (string-based, no PhysxSchema import) mesh_applied = mesh_prim.GetAppliedSchemas() if "PhysxDeformableBodyAPI" not in mesh_applied: mesh_prim.AddAppliedSchema("PhysxDeformableBodyAPI") @@ -943,7 +943,7 @@ def modify_deformable_body_properties( # get deformable-body USD prim deformable_body_prim = stage.GetPrimAtPath(prim_path) - # check if the prim is valid and has the deformable-body API + # check if the prim is valid and has the deformable-body API (string-based, no PhysxSchema import) if not deformable_body_prim.IsValid(): return False if "PhysxDeformableBodyAPI" not in deformable_body_prim.GetAppliedSchemas(): diff --git a/source/isaaclab/isaaclab/sim/utils/prims.py b/source/isaaclab/isaaclab/sim/utils/prims.py index 34cfc63e52c..174484bb532 100644 --- a/source/isaaclab/isaaclab/sim/utils/prims.py +++ b/source/isaaclab/isaaclab/sim/utils/prims.py @@ -687,12 +687,16 @@ def wrapper(prim_path: str | Sdf.Path, cfg: SpawnerCfg, *args, **kwargs): # deal with spaces by replacing them with underscores semantic_type_sanitized = semantic_type.replace(" ", "_") semantic_value_sanitized = semantic_value.replace(" ", "_") - # add labels to the prim - add_labels( - prim, labels=[semantic_value_sanitized], instance_name=semantic_type_sanitized, overwrite=False - ) + # set the semantic API for the instance + instance_name = f"{semantic_type_sanitized}_{semantic_value_sanitized}" + sem = Semantics.SemanticsAPI.Apply(prim, instance_name) + # create semantic type and data attributes + sem.CreateSemanticTypeAttr() + sem.CreateSemanticDataAttr() + sem.GetSemanticTypeAttr().Set(semantic_type) + sem.GetSemanticDataAttr().Set(semantic_value) # activate rigid body contact sensors (lazy import to avoid circular import with schemas) - if hasattr(cfg, "activate_contact_sensors") and cfg.activate_contact_sensors: # type: ignore + if hasattr(cfg, "activate_contact_sensors") and cfg.activate_contact_sensors: from ..schemas import schemas as _schemas _schemas.activate_contact_sensors(prim_spawn_path) diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer.py b/source/isaaclab/isaaclab/terrains/terrain_importer.py index c78f9ae3344..761da73e014 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer.py @@ -13,6 +13,7 @@ import trimesh import isaaclab.sim as sim_utils +from isaaclab.cloner import grid_transforms from isaaclab.markers import VisualizationMarkers from isaaclab.markers.config import FRAME_MARKER_CFG diff --git a/source/isaaclab_physx/test/test_cloner.py b/source/isaaclab_physx/test/test_cloner.py new file mode 100644 index 00000000000..723435a67a5 --- /dev/null +++ b/source/isaaclab_physx/test/test_cloner.py @@ -0,0 +1,217 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for PhysX-specific cloner utilities.""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import pytest +import torch +import warp as wp +from isaaclab_physx.cloner import physx_replicate + +import isaaclab.sim as sim_utils +from isaaclab.cloner import TemplateCloneCfg, clone_from_template, sequential +from isaaclab.sim import build_simulation_context + + +@pytest.fixture(params=["cpu", "cuda"]) +def sim(request): + """Provide a fresh simulation context for each test on CPU and CUDA.""" + with build_simulation_context(device=request.param, dt=0.01, add_lighting=False) as sim: + yield sim + + +def test_physx_replicate_no_error(sim): + """PhysX replicator call runs without raising exceptions for simple mapping.""" + sim_utils.create_prim("/World/envs", "Xform") + sim_utils.create_prim("/World/template", "Xform") + sim_utils.create_prim("/World/template/A", "Xform") + + num_envs = 2 + env_ids = torch.arange(num_envs, dtype=torch.long) + for i in range(num_envs): + sim_utils.create_prim(f"/World/envs/env_{i}", "Xform") + + mapping = torch.ones((1, num_envs), dtype=torch.bool) + + physx_replicate( + sim_utils.get_current_stage(), + sources=["/World/template/A"], + destinations=["/World/envs/env_{}/A"], + env_ids=env_ids, + mapping=mapping, + ) + + +def test_clone_from_template(sim): + """Clone prototypes via TemplateCloneCfg and clone_from_template and exercise both USD and PhysX. + + Steps: + - Create /World/template and /World/envs/env_0..env_31 + - Spawn three prototypes under /World/template/Object/proto_asset_.* + - Clone using TemplateCloneCfg with random_heterogeneous_cloning=False (modulo mapping) + - Verify modulo placement exists; then call sim.reset(), and create PhysX view + """ + num_clones = 32 + clone_cfg = TemplateCloneCfg(device=sim.cfg.device, clone_strategy=sequential) + sim_utils.create_prim(clone_cfg.template_root, "Xform") + sim_utils.create_prim(f"{clone_cfg.template_root}/Object", "Xform") + sim_utils.create_prim("/World/envs", "Xform") + for i in range(num_clones): + sim_utils.create_prim(f"/World/envs/env_{i}", "Xform", translation=(0, 0, 0)) + + cfg = sim_utils.MultiAssetSpawnerCfg( + assets_cfg=[ + sim_utils.ConeCfg( + radius=0.3, + height=0.6, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + ), + 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), + ), + ], + 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(), + ) + prim = cfg.func(f"{clone_cfg.template_root}/Object/{clone_cfg.template_prototype_identifier}_.*", cfg) + assert prim.IsValid() + + stage = sim_utils.get_current_stage() + clone_from_template(stage, num_clones=num_clones, template_clone_cfg=clone_cfg) + + primitive_prims = sim_utils.get_all_matching_child_prims( + "/World/envs", predicate=lambda prim: prim.GetTypeName() in ["Cone", "Cube", "Sphere"] + ) + + for i, primitive_prim in enumerate(primitive_prims): + modulus = i % 3 + if modulus == 0: + assert primitive_prim.GetTypeName() == "Cone" + elif modulus == 1: + assert primitive_prim.GetTypeName() == "Cube" + else: + assert primitive_prim.GetTypeName() == "Sphere" + + sim.reset() + object_view_regex = f"{clone_cfg.clone_regex}/Object".replace(".*", "*") + physics_sim_view = sim.physics_manager.get_physics_sim_view() + physx_view = physics_sim_view.create_rigid_body_view(object_view_regex) + assert physx_view is not None + + +def _run_colocation_collision_filter(sim, asset_cfg, expected_types, assert_count=False): + """Shared harness for colocated collision filter checks across devices.""" + num_clones = 32 + clone_cfg = TemplateCloneCfg(device=sim.cfg.device, clone_strategy=sequential) + sim_utils.create_prim(clone_cfg.template_root, "Xform") + sim_utils.create_prim(f"{clone_cfg.template_root}/Object", "Xform") + sim_utils.create_prim("/World/envs", "Xform") + for i in range(num_clones): + sim_utils.create_prim(f"/World/envs/env_{i}", "Xform", translation=(0, 0, 0)) + + prim = asset_cfg.func(f"{clone_cfg.template_root}/Object/{clone_cfg.template_prototype_identifier}_.*", asset_cfg) + assert prim.IsValid() + + stage = sim_utils.get_current_stage() + clone_from_template(stage, num_clones=num_clones, template_clone_cfg=clone_cfg) + + primitive_prims = sim_utils.get_all_matching_child_prims( + "/World/envs", predicate=lambda prim: prim.GetTypeName() in expected_types + ) + + if assert_count: + assert len(primitive_prims) == num_clones + + for i, primitive_prim in enumerate(primitive_prims): + assert primitive_prim.GetTypeName() == expected_types[i % len(expected_types)] + + sim.reset() + object_view_regex = f"{clone_cfg.clone_regex}/Object".replace(".*", "*") + physics_sim_view = sim.physics_manager.get_physics_sim_view() + physx_view = physics_sim_view.create_rigid_body_view(object_view_regex) + for _ in range(100): + sim.step() + transforms = wp.to_torch(physx_view.get_transforms()) + distance_from_origin = torch.linalg.norm(transforms[:, :2], dim=-1) + assert torch.all(distance_from_origin < 0.1) + + +def test_colocation_collision_filter_homogeneous(sim): + """Verify colocated clones of a single prototype stay stable after PhysX cloning. + + All clones are spawned at exactly the same pose; if the collision filter is wrong the pile + explodes on reset. This asserts the filter keeps the colocated objects stable while stepping + across CPU and CUDA backends. + """ + _run_colocation_collision_filter( + sim, + sim_utils.ConeCfg( + radius=0.3, + height=0.6, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=4, solver_velocity_iteration_count=0 + ), + collision_props=sim_utils.CollisionPropertiesCfg(), + ), + expected_types=["Cone"], + assert_count=True, + ) + + +@pytest.mark.xfail(reason="Heterogeneous cloning with collision filtering not yet fully supported") +def test_colocation_collision_filter_heterogeneous(sim): + """Verify colocated clones of multiple prototypes retain modulo ordering and remain stable. + + The cone, cube, and sphere are all spawned in the identical pose for every clone; an incorrect + collision filter would blow up the simulation on reset. This guards both modulo ordering and + that the colocated set stays stable through PhysX steps across CPU and CUDA. + """ + _run_colocation_collision_filter( + sim, + sim_utils.MultiAssetSpawnerCfg( + assets_cfg=[ + sim_utils.ConeCfg( + radius=0.3, + height=0.6, + visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.0, 1.0, 0.0), metallic=0.2), + mass_props=sim_utils.MassPropertiesCfg(mass=100.0), + ), + 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), + ), + ], + 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(), + ), + expected_types=["Cone", "Cube", "Sphere"], + ) From 4f11d5200a17419236bc3011877804d19c8d52b3 Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Mon, 23 Feb 2026 11:58:16 -0800 Subject: [PATCH 2/4] Squashed commit of the following: commit 46fe3cfb659df669e9512c4dbe4310afcffc9b92 Author: Antoine Richard Date: Mon Feb 23 19:01:02 2026 +0100 WIP commit f34eec3ad33e4e5c6b47dff8fee85a4bac70c094 Author: Antoine Richard Date: Mon Feb 23 16:16:19 2026 +0100 pre-commits commit 683c22ad9530d41740be974aa6f1bc17fb3b6470 Author: Antoine Richard Date: Mon Feb 23 16:15:55 2026 +0100 Looking good commit 728e9cfd5b1a688477087a949167c3409b10e597 Author: Antoine Richard Date: Mon Feb 23 15:30:49 2026 +0100 Getting there? commit 0547fc53d035b14029817ab783b30f563fa9e344 Author: Antoine Richard Date: Mon Feb 23 15:08:29 2026 +0100 One more clean-up for the base abstractions. commit d35c280afb74391936bfb4278d65ca03f24859d5 Author: Antoine Richard Date: Mon Feb 23 13:53:54 2026 +0100 pre-commits commit 31e15895311c478c80ecc61fa645be696992c911 Author: ooctipus Date: Mon Feb 23 13:43:31 2026 +0100 working checkpoint commit 66e3453b26f995041897de805ff88ead23774d00 Author: Antoine Richard Date: Fri Feb 20 18:31:57 2026 +0100 pre-commits commit 8dd323a0431fbcd288f444ece0acf5ece7965e99 Author: Antoine Richard Date: Fri Feb 20 18:30:57 2026 +0100 migrated the tests to use the new APIs. Removes almost all the warnings. Part of the general clean-up on the APIs. commit 27c00859ee2fd03153e901fd04f49d1267903b9b Author: Antoine Richard Date: Fri Feb 20 18:29:48 2026 +0100 Not tested but the code is here? commit 3a3586d1fb7f1317d1b74994e1a094b8aa92374b Author: Antoine Richard Date: Fri Feb 20 17:27:53 2026 +0100 reducing the number of warnings. commit d37f443bce918f8076f83ac2f4d27cabcbf40ab0 Author: Antoine Richard Date: Fri Feb 20 17:27:01 2026 +0100 WIP commit 7e6e5f884ec053d6fd0e97fca54238947c9f4a56 Author: Antoine Richard Date: Fri Feb 20 16:08:49 2026 +0100 moving bits around commit 4e0b07cb642eaf88c1b54075d047ed0c4b43d9c2 Author: Antoine Richard Date: Fri Feb 20 14:52:05 2026 +0100 Now passing all articulation tests. commit 9bd5169e9b7045ae2d0acc2d3368dbe6f7a1e163 Author: Antoine Richard Date: Fri Feb 20 14:12:50 2026 +0100 Working on unifying the new asset API across all backends. commit 27f9ee2bc6e16dbb8a08253f73a8de6a46f3416c Author: Antoine Richard Date: Fri Feb 20 13:30:28 2026 +0100 WIP commit 737e86d9f630e6efd1f5614535254e4bec931038 Author: Antoine Richard Date: Fri Feb 20 13:30:14 2026 +0100 Reworking methods docstrings to be more precise commit 35e31e6ac3b2175e8f1ee50235e2755b72d3c955 Author: Antoine Richard Date: Fri Feb 20 11:39:29 2026 +0100 working on unifying the docstrings, and fixing typos. Trying to be more precise on shapes and dtypes between warp and torch. commit 9963879f495bcc9583338eee616f054a17ea46ac Author: Antoine Richard Date: Wed Feb 18 09:27:08 2026 +0100 WIP commit fadb5cdc21333f87b16495a4281511c1928988c9 Author: Antoine Richard Date: Tue Feb 17 14:48:18 2026 +0100 WIP --- .../migration/migrating_to_isaaclab_3-0.rst | 4 +- .../assets/articulation/base_articulation.py | 649 ++- .../articulation/base_articulation_data.py | 365 +- source/isaaclab/isaaclab/assets/asset_base.py | 98 +- .../assets/rigid_object/base_rigid_object.py | 156 +- .../rigid_object/base_rigid_object_data.py | 185 +- .../base_rigid_object_collection.py | 204 +- .../base_rigid_object_collection_data.py | 154 +- .../isaaclab_newton/assets/__init__.py | 4 + .../assets/articulation/__init__.py | 9 + .../assets/articulation/articulation.py | 3858 +++++++++++++++++ .../assets/articulation/articulation_data.py | 1863 ++++++++ .../assets/articulation/kernels.py | 559 +++ .../isaaclab_newton/assets/kernels.py | 1438 ++++++ .../assets/rigid_object/__init__.py | 9 + .../assets/rigid_object/rigid_object.py | 1183 +++++ .../assets/rigid_object/rigid_object_data.py | 1230 ++++++ .../assets/articulation/articulation.py | 381 +- .../assets/articulation/articulation_data.py | 570 ++- .../deformable_object/deformable_object.py | 10 +- .../assets/rigid_object/rigid_object.py | 193 +- .../assets/rigid_object/rigid_object_data.py | 559 +-- .../rigid_object_collection.py | 295 +- .../rigid_object_collection_data.py | 385 +- .../assets/surface_gripper/surface_gripper.py | 4 +- .../test/assets/test_articulation.py | 287 +- .../test/assets/test_deformable_object.py | 8 +- .../test/assets/test_rigid_object.py | 264 +- .../assets/test_rigid_object_collection.py | 305 +- .../test/assets/test_surface_gripper.py | 6 +- 30 files changed, 13328 insertions(+), 1907 deletions(-) create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/articulation/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/kernels.py create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/rigid_object/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py create mode 100644 source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py diff --git a/docs/source/migration/migrating_to_isaaclab_3-0.rst b/docs/source/migration/migrating_to_isaaclab_3-0.rst index 8f7682953f1..cc6737a2e87 100644 --- a/docs/source/migration/migrating_to_isaaclab_3-0.rst +++ b/docs/source/migration/migrating_to_isaaclab_3-0.rst @@ -683,10 +683,10 @@ The previous ``write_*_to_sim(data, env_ids)`` methods have been removed. robot.write_root_pose_to_sim(pose_data, env_ids) # After (Isaac Lab 3.x) — indexed variant (partial data) - robot.write_root_pose_to_sim_index(pose_data, env_ids) + robot.write_root_pose_to_sim_index(root_pose=pose_data, env_ids=env_ids) # After (Isaac Lab 3.x) — mask variant (full data, boolean mask) - robot.write_root_pose_to_sim_mask(pose_data, env_mask) + robot.write_root_pose_to_sim_mask(root_pose=pose_data, env_mask=env_mask) .. list-table:: Affected write methods (RigidObject / Articulation) :header-rows: 1 diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index e2b87d7d32c..fbe240db505 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -176,13 +176,24 @@ def root_view(self): @property @abstractmethod def instantaneous_wrench_composer(self) -> WrenchComposer: - """Instantaneous wrench composer for the articulation.""" + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ raise NotImplementedError() @property @abstractmethod def permanent_wrench_composer(self) -> WrenchComposer: - """Permanent wrench composer for the articulation.""" + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ raise NotImplementedError() """ @@ -200,7 +211,7 @@ def reset( Args: env_ids: Environment indices. If None, then all indices are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -316,6 +327,7 @@ def find_spatial_tendons( @abstractmethod def write_root_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -324,14 +336,15 @@ def write_root_pose_to_sim_index( The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: - This method expect partial data. + This method expects partial data. .. tip:: For maximum performance we recommend looking at the actual implementation of the method in the backend. Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -339,29 +352,32 @@ def write_root_pose_to_sim_index( @abstractmethod def write_root_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root pose over selected environment indices into the simulation. + """Set the root pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend looking at the actual implementation of the method in the backend. Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_link_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -370,14 +386,15 @@ def write_root_link_pose_to_sim_index( The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: - This method expect partial data. + This method expects partial data. .. tip:: For maximum performance we recommend looking at the actual implementation of the method in the backend. Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -385,36 +402,39 @@ def write_root_link_pose_to_sim_index( @abstractmethod def write_root_link_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root link pose over selected environment indices into the simulation. + """Set the root link pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: - This method expect full data. + This method expects full data. .. tip:: For maximum performance we recommend looking at the actual implementation of the method in the backend. Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_com_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root center of mass pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). - The orientation is the orientation of the principle axes of inertia. + The orientation is the orientation of the principal axes of inertia. .. note:: This method expects partial data. @@ -424,7 +444,8 @@ def write_root_com_pose_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7). + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -432,13 +453,14 @@ def write_root_com_pose_to_sim_index( @abstractmethod def write_root_com_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass pose over selected environment indices into the simulation. + """Set the root center of mass pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). - The orientation is the orientation of the principle axes of inertia. + The orientation is the orientation of the principal axes of inertia. .. note:: This method expects full data. @@ -448,14 +470,16 @@ def write_root_com_pose_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -464,7 +488,7 @@ def write_root_velocity_to_sim_index( The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's center of mass rather than the roots frame. + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects partial data. @@ -474,7 +498,8 @@ def write_root_velocity_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -482,15 +507,16 @@ def write_root_velocity_to_sim_index( @abstractmethod def write_root_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass velocity over selected environment indices into the simulation. + """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's center of mass rather than the roots frame. + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects full data. @@ -500,14 +526,16 @@ def write_root_velocity_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_com_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -516,7 +544,7 @@ def write_root_com_velocity_to_sim_index( The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's center of mass rather than the roots frame. + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects partial data. @@ -526,7 +554,8 @@ def write_root_com_velocity_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -534,15 +563,16 @@ def write_root_com_velocity_to_sim_index( @abstractmethod def write_root_com_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass velocity over selected environment indices into the simulation. + """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's center of mass rather than the roots frame. + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects full data. @@ -552,14 +582,16 @@ def write_root_com_velocity_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_link_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -568,7 +600,7 @@ def write_root_link_velocity_to_sim_index( The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's frame rather than the roots center of mass. + This sets the velocity of the root's frame rather than the root's center of mass. .. note:: This method expects partial data. @@ -578,7 +610,8 @@ def write_root_link_velocity_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6). + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -586,15 +619,16 @@ def write_root_link_velocity_to_sim_index( @abstractmethod def write_root_link_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root link velocity over selected environment indices into the simulation. + """Set the root link velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's frame rather than the roots center of mass. + This sets the velocity of the root's frame rather than the root's center of mass. .. note:: This method expects full data. @@ -604,17 +638,71 @@ def write_root_link_velocity_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_state_to_sim_index( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Write joint positions and velocities over selected environment indices into the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. + + Args: + position: Joint positions. Shape is (len(env_ids), len(joint_ids)). + velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the state for. Defaults to None (all joints). + env_ids: The environment indices to set the state for. Defaults to None (all instances). + """ + raise NotImplementedError() + + @abstractmethod + def write_joint_state_to_sim_mask( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint positions and velocities over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. + + Args: + position: Joint positions. Shape is (num_instances, num_joints). + velocity: Joint velocities. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_joint_position_to_sim_index( self, + *, position: torch.Tensor | wp.array, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write joint positions to the simulation. @@ -628,16 +716,17 @@ def write_joint_position_to_sim_index( Args: position: Joint positions. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). + env_ids: The environment indices to set the targets for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_joint_position_to_sim_mask( self, + *, position: torch.Tensor | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write joint positions to the simulation. @@ -650,17 +739,18 @@ def write_joint_position_to_sim_mask( Args: position: Joint positions. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_joint_velocity_to_sim_index( self, + *, velocity: torch.Tensor | wp.array, - joint_ids: Sequence[int] | slice | None = None, - env_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write joint velocities to the simulation. @@ -674,16 +764,17 @@ def write_joint_velocity_to_sim_index( Args: velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). + env_ids: The environment indices to set the targets for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_joint_velocity_to_sim_mask( self, + *, velocity: torch.Tensor | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write joint velocities to the simulation. @@ -696,8 +787,8 @@ def write_joint_velocity_to_sim_mask( Args: velocity: Joint velocities. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -708,8 +799,9 @@ def write_joint_velocity_to_sim_mask( @abstractmethod def write_joint_stiffness_to_sim_index( self, + *, stiffness: torch.Tensor | float | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write joint stiffness into the simulation. @@ -724,16 +816,17 @@ def write_joint_stiffness_to_sim_index( Args: stiffness: Joint stiffness. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the stiffness for. Defaults to None (all joints). - env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + env_ids: The environment indices to set the stiffness for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_joint_stiffness_to_sim_mask( self, + *, stiffness: torch.Tensor | float | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write joint stiffness into the simulation. @@ -746,16 +839,17 @@ def write_joint_stiffness_to_sim_mask( Args: stiffness: Joint stiffness. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_joint_damping_to_sim_index( self, + *, damping: torch.Tensor | float | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write joint damping into the simulation. @@ -770,16 +864,17 @@ def write_joint_damping_to_sim_index( Args: damping: Joint damping. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the damping for. Defaults to None (all joints). - env_ids: The environment indices to set the damping for. Defaults to None (all environments). + env_ids: The environment indices to set the damping for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_joint_damping_to_sim_mask( self, + *, damping: torch.Tensor | float | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write joint damping into the simulation. @@ -792,16 +887,17 @@ def write_joint_damping_to_sim_mask( Args: damping: Joint damping. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_joint_position_limit_to_sim_index( self, + *, limits: torch.Tensor | float | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, warn_limit_violation: bool = True, ) -> None: @@ -815,9 +911,10 @@ def write_joint_position_limit_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - limits: Joint limits. Shape is (len(env_ids), len(joint_ids), 2). + limits: Joint limits. Shape is (len(env_ids), len(joint_ids), 2) or (len(env_ids), len(joint_ids)) with + dtype wp.vec2f. joint_ids: The joint indices to set the limits for. Defaults to None (all joints). - env_ids: The environment indices to set the limits for. Defaults to None (all environments). + env_ids: The environment indices to set the limits for. Defaults to None (all instances). warn_limit_violation: Whether to use warning or info level logging when default joint positions exceed the new limits. Defaults to True. """ @@ -826,10 +923,11 @@ def write_joint_position_limit_to_sim_index( @abstractmethod def write_joint_position_limit_to_sim_mask( self, + *, limits: torch.Tensor | float | wp.array, - warn_limit_violation: bool = True, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + warn_limit_violation: bool = True, ) -> None: """Write joint position limits into the simulation. @@ -841,19 +939,21 @@ def write_joint_position_limit_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - limits: Joint limits. Shape is (num_instances, num_joints, 2). + limits: Joint limits. Shape is (num_instances, num_joints, 2) or (num_instances, num_joints) with dtype + wp.vec2f. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). warn_limit_violation: Whether to use warning or info level logging when default joint positions exceed the new limits. Defaults to True. - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. """ raise NotImplementedError() @abstractmethod def write_joint_velocity_limit_to_sim_index( self, + *, limits: torch.Tensor | float | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write joint max velocity to the simulation. @@ -872,16 +972,17 @@ def write_joint_velocity_limit_to_sim_index( Args: limits: Joint max velocity. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the max velocity for. Defaults to None (all joints). - env_ids: The environment indices to set the max velocity for. Defaults to None (all environments). + env_ids: The environment indices to set the max velocity for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_joint_velocity_limit_to_sim_mask( self, + *, limits: torch.Tensor | float | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write joint max velocity to the simulation. @@ -898,16 +999,17 @@ def write_joint_velocity_limit_to_sim_mask( Args: limits: Joint max velocity. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_joint_effort_limit_to_sim_index( self, + *, limits: torch.Tensor | float | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write joint effort limits into the simulation. @@ -925,16 +1027,17 @@ def write_joint_effort_limit_to_sim_index( Args: limits: Joint torque limits. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the joint torque limits for. Defaults to None (all joints). - env_ids: The environment indices to set the joint torque limits for. Defaults to None (all environments). + env_ids: The environment indices to set the joint torque limits for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_joint_effort_limit_to_sim_mask( self, + *, limits: torch.Tensor | float | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write joint effort limits into the simulation. @@ -950,16 +1053,17 @@ def write_joint_effort_limit_to_sim_mask( Args: limits: Joint torque limits. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_joint_armature_to_sim_index( self, + *, armature: torch.Tensor | float | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write joint armature into the simulation. @@ -977,16 +1081,17 @@ def write_joint_armature_to_sim_index( Args: armature: Joint armature. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the joint torque limits for. Defaults to None (all joints). - env_ids: The environment indices to set the joint torque limits for. Defaults to None (all environments). + env_ids: The environment indices to set the joint torque limits for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_joint_armature_to_sim_mask( self, + *, armature: torch.Tensor | float | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write joint armature into the simulation. @@ -1002,16 +1107,17 @@ def write_joint_armature_to_sim_mask( Args: armature: Joint armature. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_joint_friction_coefficient_to_sim_index( self, + *, joint_friction_coeff: torch.Tensor | float | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: r"""Write joint static friction coefficients into the simulation. @@ -1030,16 +1136,17 @@ def write_joint_friction_coefficient_to_sim_index( Args: joint_friction_coeff: Joint static friction coefficient. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the joint torque limits for. Defaults to None (all joints). - env_ids: The environment indices to set the joint torque limits for. Defaults to None (all environments). + env_ids: The environment indices to set the joint torque limits for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_joint_friction_coefficient_to_sim_mask( self, + *, joint_friction_coeff: torch.Tensor | float | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: r"""Write joint static friction coefficients into the simulation. @@ -1056,8 +1163,8 @@ def write_joint_friction_coefficient_to_sim_mask( Args: joint_friction_coeff: Joint static friction coefficient. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -1068,8 +1175,9 @@ def write_joint_friction_coefficient_to_sim_mask( @abstractmethod def set_masses_index( self, + *, masses: torch.Tensor | wp.array, - body_ids: Sequence[int] | slice | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set masses of all bodies in the simulation world frame. @@ -1084,16 +1192,17 @@ def set_masses_index( Args: masses: Masses of all bodies. Shape is (len(env_ids), len(body_ids)). body_ids: The body indices to set the masses for. Defaults to None (all bodies). - env_ids: The environment indices to set the masses for. Defaults to None (all environments). + env_ids: The environment indices to set the masses for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_masses_mask( self, + *, masses: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set masses of all bodies in the simulation world frame. @@ -1106,16 +1215,17 @@ def set_masses_mask( Args: masses: Masses of all bodies. Shape is (num_instances, num_bodies). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_mask: Body mask. If None, then all the bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_coms_index( self, + *, coms: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set center of mass positions of all bodies in the simulation world frame. @@ -1128,19 +1238,21 @@ def set_coms_index( Some backends may provide optimized implementations for masks / indices. Args: - coms: Center of mass positions of all bodies. Shape is (len(env_ids), len(body_ids), 3). + coms: Center of mass positions of all bodies. Shape is (len(env_ids), len(body_ids), 3) + or (len(env_ids), len(body_ids)) with dtype wp.vec3f. body_ids: The body indices to set the center of mass positions for. Defaults to None (all bodies). env_ids: The environment indices to set the center of mass positions for. Defaults to None - (all environments). + (all instances). """ raise NotImplementedError() @abstractmethod def set_coms_mask( self, + *, coms: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set center of mass positions of all bodies in the simulation world frame. @@ -1152,17 +1264,19 @@ def set_coms_mask( Some backends may provide optimized implementations for masks / indices. Args: - coms: Center of mass positions of all bodies. Shape is (num_instances, num_bodies, 3). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + coms: Center of mass positions of all bodies. Shape is (num_instances, num_bodies, 3) + or (num_instances, num_bodies) with dtype wp.vec3f. + body_mask: Body mask. If None, then all the bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_inertias_index( self, + *, inertias: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set inertias of all bodies in the simulation world frame. @@ -1177,16 +1291,17 @@ def set_inertias_index( Args: inertias: Inertias of all bodies. Shape is (len(env_ids), len(body_ids), 9). body_ids: The body indices to set the inertias for. Defaults to None (all bodies). - env_ids: The environment indices to set the inertias for. Defaults to None (all environments). + env_ids: The environment indices to set the inertias for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_inertias_mask( self, + *, inertias: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set inertias of all bodies in the simulation world frame. @@ -1199,16 +1314,17 @@ def set_inertias_mask( Args: inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_mask: Body mask. If None, then all the bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_joint_position_target_index( self, + *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set joint position targets into internal buffers. @@ -1226,16 +1342,17 @@ def set_joint_position_target_index( Args: target: Joint position targets. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). + env_ids: The environment indices to set the targets for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_joint_position_target_mask( self, + *, target: torch.Tensor | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set joint position targets into internal buffers. @@ -1251,16 +1368,17 @@ def set_joint_position_target_mask( Args: target: Joint position targets. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_joint_velocity_target_index( self, + *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set joint velocity targets into internal buffers. @@ -1278,16 +1396,17 @@ def set_joint_velocity_target_index( Args: target: Joint velocity targets. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). + env_ids: The environment indices to set the targets for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_joint_velocity_target_mask( self, + *, target: torch.Tensor | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set joint velocity targets into internal buffers. @@ -1303,16 +1422,17 @@ def set_joint_velocity_target_mask( Args: target: Joint velocity targets. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_joint_effort_target_index( self, + *, target: torch.Tensor | wp.array, - joint_ids: Sequence[int] | slice | None = None, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set joint efforts into internal buffers. @@ -1330,16 +1450,17 @@ def set_joint_effort_target_index( Args: target: Joint effort targets. Shape is (len(env_ids), len(joint_ids)). joint_ids: The joint indices to set the targets for. Defaults to None (all joints). - env_ids: The environment indices to set the targets for. Defaults to None (all environments). + env_ids: The environment indices to set the targets for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_joint_effort_target_mask( self, + *, target: torch.Tensor | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set joint efforts into internal buffers. @@ -1355,8 +1476,8 @@ def set_joint_effort_target_mask( Args: target: Joint effort targets. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. - joint_mask: Joint mask. If None, then all joints are used. + joint_mask: Joint mask. If None, then all the joints are updated. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -1367,8 +1488,9 @@ def set_joint_effort_target_mask( @abstractmethod def set_fixed_tendon_stiffness_index( self, + *, stiffness: torch.Tensor | wp.array, - fixed_tendon_ids: Sequence[int] | slice | None = None, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set fixed tendon stiffness into internal buffers. @@ -1387,16 +1509,17 @@ def set_fixed_tendon_stiffness_index( Args: stiffness: Fixed tendon stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). fixed_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + env_ids: The environment indices to set the stiffness for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_stiffness_mask( self, + *, stiffness: torch.Tensor | wp.array, - env_mask: wp.array | None = None, fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set fixed tendon stiffness into internal buffers. @@ -1413,16 +1536,18 @@ def set_fixed_tendon_stiffness_mask( Args: stiffness: Fixed tendon stiffness. Shape is (num_instances, num_fixed_tendons). - env_mask: Environment mask. If None, then all indices are used. - fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_damping_index( self, + *, damping: torch.Tensor | wp.array, - fixed_tendon_ids: Sequence[int] | slice | None = None, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set fixed tendon damping into internal buffers. @@ -1434,19 +1559,24 @@ def set_fixed_tendon_damping_index( .. note:: This method expects partial data. + .. tip:: + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. + Args: damping: Fixed tendon damping. Shape is (len(env_ids), len(fixed_tendon_ids)). fixed_tendon_ids: The tendon indices to set the damping for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the damping for. Defaults to None (all environments). + env_ids: The environment indices to set the damping for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_damping_mask( self, + *, damping: torch.Tensor | wp.array, - env_mask: wp.array | None = None, fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set fixed tendon damping into internal buffers. @@ -1457,18 +1587,24 @@ def set_fixed_tendon_damping_mask( .. note:: This method expects full data. + .. tip:: + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. + Args: damping: Fixed tendon damping. Shape is (num_instances, num_fixed_tendons). - env_mask: Environment mask. If None, then all indices are used. - fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_limit_stiffness_index( self, + *, limit_stiffness: torch.Tensor | wp.array, - fixed_tendon_ids: Sequence[int] | slice | None = None, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set fixed tendon limit stiffness into internal buffers. @@ -1487,16 +1623,17 @@ def set_fixed_tendon_limit_stiffness_index( Args: limit_stiffness: Fixed tendon limit stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). fixed_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). + env_ids: The environment indices to set the limit stiffness for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_limit_stiffness_mask( self, + *, limit_stiffness: torch.Tensor | wp.array, - env_mask: wp.array | None = None, fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set fixed tendon limit stiffness into internal buffers. @@ -1513,16 +1650,18 @@ def set_fixed_tendon_limit_stiffness_mask( Args: limit_stiffness: Fixed tendon limit stiffness. Shape is (num_instances, num_fixed_tendons). - env_mask: Environment mask. If None, then all indices are used. - fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_position_limit_index( self, + *, limit: torch.Tensor | wp.array, - fixed_tendon_ids: Sequence[int] | slice | None = None, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set fixed tendon position limits into internal buffers. @@ -1541,16 +1680,17 @@ def set_fixed_tendon_position_limit_index( Args: limit: Fixed tendon limit. Shape is (len(env_ids), len(fixed_tendon_ids)). fixed_tendon_ids: The tendon indices to set the limit for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the limit for. Defaults to None (all environments). + env_ids: The environment indices to set the limit for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_position_limit_mask( self, + *, limit: torch.Tensor | wp.array, - env_mask: wp.array | None = None, fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set fixed tendon position limits into internal buffers. @@ -1567,16 +1707,18 @@ def set_fixed_tendon_position_limit_mask( Args: limit: Fixed tendon limit. Shape is (num_instances, num_fixed_tendons). - env_mask: Environment mask. If None, then all indices are used. - fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_rest_length_index( self, + *, rest_length: torch.Tensor | wp.array, - fixed_tendon_ids: Sequence[int] | slice | None = None, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set fixed tendon rest length into internal buffers. @@ -1595,16 +1737,17 @@ def set_fixed_tendon_rest_length_index( Args: rest_length: Fixed tendon rest length. Shape is (len(env_ids), len(fixed_tendon_ids)). fixed_tendon_ids: The tendon indices to set the rest length for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the rest length for. Defaults to None (all environments). + env_ids: The environment indices to set the rest length for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_rest_length_mask( self, + *, rest_length: torch.Tensor | wp.array, - env_mask: wp.array | None = None, fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set fixed tendon rest length into internal buffers. @@ -1621,16 +1764,18 @@ def set_fixed_tendon_rest_length_mask( Args: rest_length: Fixed tendon rest length. Shape is (num_instances, num_fixed_tendons). - env_mask: Environment mask. If None, then all indices are used. - fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_offset_index( self, + *, offset: torch.Tensor | wp.array, - fixed_tendon_ids: Sequence[int] | slice | None = None, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set fixed tendon offset into internal buffers. @@ -1649,16 +1794,17 @@ def set_fixed_tendon_offset_index( Args: offset: Fixed tendon offset. Shape is (len(env_ids), len(fixed_tendon_ids)). fixed_tendon_ids: The tendon indices to set the offset for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the offset for. Defaults to None (all environments). + env_ids: The environment indices to set the offset for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_fixed_tendon_offset_mask( self, + *, offset: torch.Tensor | wp.array, - env_mask: wp.array | None = None, fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set fixed tendon offset into internal buffers. @@ -1675,15 +1821,17 @@ def set_fixed_tendon_offset_mask( Args: offset: Fixed tendon offset. Shape is (num_instances, num_fixed_tendons). - env_mask: Environment mask. If None, then all indices are used. - fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_fixed_tendon_properties_to_sim_index( self, - fixed_tendon_ids: Sequence[int] | slice | None = None, + *, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write fixed tendon properties into the simulation. @@ -1697,15 +1845,16 @@ def write_fixed_tendon_properties_to_sim_index( Args: fixed_tendon_ids: The fixed tendon indices to set the limits for. Defaults to None (all fixed tendons). - env_ids: The environment indices to set the limits for. Defaults to None (all environments). + env_ids: The environment indices to set the limits for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_fixed_tendon_properties_to_sim_mask( self, - env_mask: wp.array | None = None, + *, fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write fixed tendon properties into the simulation. @@ -1717,16 +1866,18 @@ def write_fixed_tendon_properties_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - env_mask: Environment mask. If None, then all indices are used. - fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + fixed_tendon_mask: Fixed tendon mask. If None, then all the fixed tendons are updated. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_spatial_tendon_stiffness_index( self, + *, stiffness: torch.Tensor | wp.array, - spatial_tendon_ids: Sequence[int] | slice | None = None, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set spatial tendon stiffness into internal buffers. @@ -1738,19 +1889,24 @@ def set_spatial_tendon_stiffness_index( .. note:: This method expects partial data. + .. tip:: + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. + Args: stiffness: Spatial tendon stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). spatial_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the stiffness for. Defaults to None (all environments). + env_ids: The environment indices to set the stiffness for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_spatial_tendon_stiffness_mask( self, + *, stiffness: torch.Tensor | wp.array, - env_mask: wp.array | None = None, spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set spatial tendon stiffness into internal buffers. @@ -1761,18 +1917,24 @@ def set_spatial_tendon_stiffness_mask( .. note:: This method expects full data. + .. tip:: + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. + Args: stiffness: Spatial tendon stiffness. Shape is (num_instances, num_spatial_tendons). - env_mask: Environment mask. If None, then all indices are used. - spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + spatial_tendon_mask: Spatial tendon mask. If None, then all the spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_spatial_tendon_damping_index( self, + *, damping: torch.Tensor | wp.array, - spatial_tendon_ids: Sequence[int] | slice | None = None, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set spatial tendon damping into internal buffers. @@ -1785,21 +1947,23 @@ def set_spatial_tendon_damping_index( This method expects partial data. .. tip:: - For maximum performance we recommend providing the environment mask instead of the environment indices. + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. Args: damping: Spatial tendon damping. Shape is (len(env_ids), len(spatial_tendon_ids)). spatial_tendon_ids: The tendon indices to set the damping for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the damping for. Defaults to None (all environments). + env_ids: The environment indices to set the damping for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_spatial_tendon_damping_mask( self, + *, damping: torch.Tensor | wp.array, - env_mask: wp.array | None = None, spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set spatial tendon damping into internal buffers. @@ -1811,20 +1975,23 @@ def set_spatial_tendon_damping_mask( This method expects full data. .. tip:: - For maximum performance we recommend providing the environment mask instead of the environment indices. + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. Args: damping: Spatial tendon damping. Shape is (num_instances, num_spatial_tendons). - env_mask: Environment mask. If None, then all indices are used. - spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + spatial_tendon_mask: Spatial tendon mask. If None, then all the spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_spatial_tendon_limit_stiffness_index( self, + *, limit_stiffness: torch.Tensor | wp.array, - spatial_tendon_ids: Sequence[int] | slice | None = None, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set spatial tendon limit stiffness into internal buffers. @@ -1837,22 +2004,24 @@ def set_spatial_tendon_limit_stiffness_index( This method expects partial data. .. tip:: - For maximum performance we recommend providing the environment mask instead of the environment indices. + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. Args: limit_stiffness: Spatial tendon limit stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). spatial_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the limit stiffness for. Defaults to None (all environments). + env_ids: The environment indices to set the limit stiffness for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_spatial_tendon_limit_stiffness_mask( self, + *, limit_stiffness: torch.Tensor | wp.array, - env_mask: wp.array | None = None, spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set spatial tendon limit stiffness into internal buffers. @@ -1864,20 +2033,23 @@ def set_spatial_tendon_limit_stiffness_mask( This method expects full data. .. tip:: - For maximum performance we recommend providing the environment mask instead of the environment indices. + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. Args: limit_stiffness: Spatial tendon limit stiffness. Shape is (num_instances, num_spatial_tendons). - env_mask: Environment mask. If None, then all indices are used. - spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + spatial_tendon_mask: Spatial tendon mask. If None, then all the spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_spatial_tendon_offset_index( self, - offset: torch.Tensor, - spatial_tendon_ids: Sequence[int] | slice | None = None, + *, + offset: torch.Tensor | wp.array, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set spatial tendon offset into internal buffers. @@ -1889,19 +2061,24 @@ def set_spatial_tendon_offset_index( .. note:: This method expects partial data. + .. tip:: + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. + Args: offset: Spatial tendon offset. Shape is (len(env_ids), len(spatial_tendon_ids)). spatial_tendon_ids: The tendon indices to set the offset for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the offset for. Defaults to None (all environments). + env_ids: The environment indices to set the offset for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def set_spatial_tendon_offset_mask( self, - offset: torch.Tensor, - env_mask: wp.array | None = None, + *, + offset: torch.Tensor | wp.array, spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set spatial tendon offset into internal buffers. @@ -1912,17 +2089,23 @@ def set_spatial_tendon_offset_mask( .. note:: This method expects full data. + .. tip:: + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. + Args: offset: Spatial tendon offset. Shape is (num_instances, num_spatial_tendons). - env_mask: Environment mask. If None, then all indices are used. - spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + spatial_tendon_mask: Spatial tendon mask. If None, then all the spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_spatial_tendon_properties_to_sim_index( self, - spatial_tendon_ids: Sequence[int] | slice | None = None, + *, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write spatial tendon properties into the simulation. @@ -1931,20 +2114,22 @@ def write_spatial_tendon_properties_to_sim_index( This method expects partial data. .. tip:: - For maximum performance we recommend providing the environment mask instead of the environment indices. + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. Args: spatial_tendon_ids: The spatial tendon indices to set the properties for. Defaults to None (all spatial tendons). - env_ids: The environment indices to set the properties for. Defaults to None (all environments). + env_ids: The environment indices to set the properties for. Defaults to None (all instances). """ raise NotImplementedError() @abstractmethod def write_spatial_tendon_properties_to_sim_mask( self, - env_mask: wp.array | None = None, + *, spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Write spatial tendon properties into the simulation. @@ -1952,11 +2137,13 @@ def write_spatial_tendon_properties_to_sim_mask( This method expects full data. .. tip:: - For maximum performance we recommend providing the environment mask instead of the environment indices. + For maximum performance we recommend looking at the actual implementation of the method in the backend. + Some backends may provide optimized implementations for masks / indices. Args: - env_mask: Environment mask. If None, then all indices are used. - spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + spatial_tendon_mask: Spatial tendon mask. If None, then all the spatial tendons are updated. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -2042,8 +2229,6 @@ def write_joint_friction_to_sim( joint_friction: torch.Tensor | float | wp.array, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - env_mask: wp.array | None = None, - joint_mask: wp.array | None = None, ) -> None: """Write joint friction coefficients into the simulation. @@ -2056,9 +2241,7 @@ def write_joint_friction_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_friction_coefficient_to_sim( - joint_friction, joint_ids=joint_ids, env_ids=env_ids, env_mask=env_mask, joint_mask=joint_mask - ) + self.write_joint_friction_coefficient_to_sim(joint_friction, joint_ids=joint_ids, env_ids=env_ids) def write_joint_limits_to_sim( self, @@ -2066,8 +2249,6 @@ def write_joint_limits_to_sim( joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, warn_limit_violation: bool = True, - env_mask: wp.array | None = None, - joint_mask: wp.array | None = None, ) -> None: """Write joint limits into the simulation. @@ -2085,8 +2266,6 @@ def write_joint_limits_to_sim( joint_ids=joint_ids, env_ids=env_ids, warn_limit_violation=warn_limit_violation, - env_mask=env_mask, - joint_mask=joint_mask, ) def set_fixed_tendon_limit( @@ -2094,8 +2273,6 @@ def set_fixed_tendon_limit( limit: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - env_mask: wp.array | None = None, - fixed_tendon_mask: wp.array | None = None, ) -> None: """Set fixed tendon position limits into internal buffers. @@ -2112,8 +2289,6 @@ def set_fixed_tendon_limit( limit, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, - env_mask=env_mask, - fixed_tendon_mask=fixed_tendon_mask, ) @abstractmethod @@ -2156,7 +2331,7 @@ def write_root_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_pose_to_sim_index(root_pose, env_ids=env_ids) + self.write_root_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_link_pose_to_sim( self, @@ -2170,7 +2345,7 @@ def write_root_link_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_link_pose_to_sim_index(root_pose, env_ids=env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_com_pose_to_sim( self, @@ -2184,7 +2359,7 @@ def write_root_com_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_com_pose_to_sim_index(root_pose, env_ids=env_ids) + self.write_root_com_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_velocity_to_sim( self, @@ -2198,7 +2373,7 @@ def write_root_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_velocity_to_sim_index(root_velocity, env_ids=env_ids) + self.write_root_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def write_root_com_velocity_to_sim( self, @@ -2212,7 +2387,7 @@ def write_root_com_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_com_velocity_to_sim_index(root_velocity, env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def write_root_link_velocity_to_sim( self, @@ -2226,7 +2401,7 @@ def write_root_link_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_link_velocity_to_sim_index(root_velocity, env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) @abstractmethod def write_joint_state_to_sim( @@ -2253,7 +2428,7 @@ def write_joint_position_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_position_to_sim_index(position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) def write_joint_velocity_to_sim( self, @@ -2268,7 +2443,7 @@ def write_joint_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_velocity_to_sim_index(velocity, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) def write_joint_stiffness_to_sim( self, @@ -2283,7 +2458,7 @@ def write_joint_stiffness_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_stiffness_to_sim_index(stiffness, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_stiffness_to_sim_index(stiffness=stiffness, joint_ids=joint_ids, env_ids=env_ids) def write_joint_damping_to_sim( self, @@ -2298,7 +2473,7 @@ def write_joint_damping_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_damping_to_sim_index(damping, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_damping_to_sim_index(damping=damping, joint_ids=joint_ids, env_ids=env_ids) def write_joint_position_limit_to_sim( self, @@ -2315,7 +2490,7 @@ def write_joint_position_limit_to_sim( stacklevel=2, ) self.write_joint_position_limit_to_sim_index( - limits, joint_ids=joint_ids, env_ids=env_ids, warn_limit_violation=warn_limit_violation + limits=limits, joint_ids=joint_ids, env_ids=env_ids, warn_limit_violation=warn_limit_violation ) def write_joint_velocity_limit_to_sim( @@ -2331,7 +2506,7 @@ def write_joint_velocity_limit_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_velocity_limit_to_sim_index(limits, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_limit_to_sim_index(limits=limits, joint_ids=joint_ids, env_ids=env_ids) def write_joint_effort_limit_to_sim( self, @@ -2346,7 +2521,7 @@ def write_joint_effort_limit_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_effort_limit_to_sim_index(limits, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_effort_limit_to_sim_index(limits=limits, joint_ids=joint_ids, env_ids=env_ids) def write_joint_armature_to_sim( self, @@ -2361,7 +2536,7 @@ def write_joint_armature_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_armature_to_sim_index(armature, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_armature_to_sim_index(armature=armature, joint_ids=joint_ids, env_ids=env_ids) def write_joint_friction_coefficient_to_sim( self, @@ -2376,7 +2551,9 @@ def write_joint_friction_coefficient_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_joint_friction_coefficient_to_sim_index(joint_friction_coeff, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=joint_friction_coeff, joint_ids=joint_ids, env_ids=env_ids + ) def set_masses( self, @@ -2390,7 +2567,7 @@ def set_masses( DeprecationWarning, stacklevel=2, ) - self.set_masses_index(masses, body_ids=body_ids, env_ids=env_ids) + self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids) def set_coms( self, @@ -2404,7 +2581,7 @@ def set_coms( DeprecationWarning, stacklevel=2, ) - self.set_coms_index(coms, body_ids=body_ids, env_ids=env_ids) + self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids) def set_inertias( self, @@ -2419,7 +2596,7 @@ def set_inertias( DeprecationWarning, stacklevel=2, ) - self.set_inertias_index(inertias, body_ids=body_ids, env_ids=env_ids) + self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids) def set_external_force_and_torque( self, @@ -2454,7 +2631,7 @@ def set_joint_position_target( DeprecationWarning, stacklevel=2, ) - self.set_joint_position_target_index(target, joint_ids=joint_ids, env_ids=env_ids) + self.set_joint_position_target_index(target=target, joint_ids=joint_ids, env_ids=env_ids) def set_joint_velocity_target( self, @@ -2469,7 +2646,7 @@ def set_joint_velocity_target( DeprecationWarning, stacklevel=2, ) - self.set_joint_velocity_target_index(target, joint_ids=joint_ids, env_ids=env_ids) + self.set_joint_velocity_target_index(target=target, joint_ids=joint_ids, env_ids=env_ids) def set_joint_effort_target( self, @@ -2484,7 +2661,7 @@ def set_joint_effort_target( DeprecationWarning, stacklevel=2, ) - self.set_joint_effort_target_index(target, joint_ids=joint_ids, env_ids=env_ids) + self.set_joint_effort_target_index(target=target, joint_ids=joint_ids, env_ids=env_ids) def set_fixed_tendon_stiffness( self, @@ -2499,7 +2676,7 @@ def set_fixed_tendon_stiffness( DeprecationWarning, stacklevel=2, ) - self.set_fixed_tendon_stiffness_index(stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) + self.set_fixed_tendon_stiffness_index(stiffness=stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) def set_fixed_tendon_damping( self, @@ -2514,7 +2691,7 @@ def set_fixed_tendon_damping( DeprecationWarning, stacklevel=2, ) - self.set_fixed_tendon_damping_index(damping, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) + self.set_fixed_tendon_damping_index(damping=damping, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) def set_fixed_tendon_limit_stiffness( self, @@ -2529,7 +2706,9 @@ def set_fixed_tendon_limit_stiffness( DeprecationWarning, stacklevel=2, ) - self.set_fixed_tendon_limit_stiffness_index(limit_stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) + self.set_fixed_tendon_limit_stiffness_index( + limit_stiffness=limit_stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids + ) def set_fixed_tendon_position_limit( self, @@ -2544,7 +2723,7 @@ def set_fixed_tendon_position_limit( DeprecationWarning, stacklevel=2, ) - self.set_fixed_tendon_position_limit_index(limit, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) + self.set_fixed_tendon_position_limit_index(limit=limit, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) def set_fixed_tendon_rest_length( self, @@ -2559,7 +2738,9 @@ def set_fixed_tendon_rest_length( DeprecationWarning, stacklevel=2, ) - self.set_fixed_tendon_rest_length_index(rest_length, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) + self.set_fixed_tendon_rest_length_index( + rest_length=rest_length, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids + ) def set_fixed_tendon_offset( self, @@ -2574,7 +2755,7 @@ def set_fixed_tendon_offset( DeprecationWarning, stacklevel=2, ) - self.set_fixed_tendon_offset_index(offset, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) + self.set_fixed_tendon_offset_index(offset=offset, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids) def write_fixed_tendon_properties_to_sim( self, @@ -2604,7 +2785,9 @@ def set_spatial_tendon_stiffness( DeprecationWarning, stacklevel=2, ) - self.set_spatial_tendon_stiffness_index(stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids) + self.set_spatial_tendon_stiffness_index( + stiffness=stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids + ) def set_spatial_tendon_damping( self, @@ -2619,7 +2802,7 @@ def set_spatial_tendon_damping( DeprecationWarning, stacklevel=2, ) - self.set_spatial_tendon_damping_index(damping, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids) + self.set_spatial_tendon_damping_index(damping=damping, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids) def set_spatial_tendon_limit_stiffness( self, @@ -2635,7 +2818,7 @@ def set_spatial_tendon_limit_stiffness( stacklevel=2, ) self.set_spatial_tendon_limit_stiffness_index( - limit_stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids + limit_stiffness=limit_stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids ) def set_spatial_tendon_offset( @@ -2651,7 +2834,7 @@ def set_spatial_tendon_offset( DeprecationWarning, stacklevel=2, ) - self.set_spatial_tendon_offset_index(offset, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids) + self.set_spatial_tendon_offset_index(offset=offset, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids) def write_spatial_tendon_properties_to_sim( self, diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py index e133f46d3ff..61ee8f28e7a 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation_data.py @@ -66,7 +66,8 @@ def update(self, dt: float) -> None: def default_root_pose(self) -> wp.array: """Default root pose ``[pos, quat]`` in the local environment frame. - The position and quaternion are of the articulation root's actor frame. Shape is (num_instances, 7). + The position and quaternion are of the articulation root's actor frame. Shape is (num_instances), + dtype = wp.transformf. In torch this resolves to (num_instances, 7). """ raise NotImplementedError @@ -76,7 +77,7 @@ def default_root_vel(self) -> wp.array: """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. The linear and angular velocities are of the articulation root's center of mass frame. - Shape is (num_instances, 6). + Shape is (num_instances), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). """ raise NotImplementedError @@ -89,7 +90,9 @@ def default_root_state(self) -> wp.array: @property @abstractmethod def default_joint_pos(self) -> wp.array: - """Default joint positions of all joints. Shape is (num_instances, num_joints). + """Default joint positions of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ @@ -98,7 +101,9 @@ def default_joint_pos(self) -> wp.array: @property @abstractmethod def default_joint_vel(self) -> wp.array: - """Default joint velocities of all joints. Shape is (num_instances, num_joints). + """Default joint velocities of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ @@ -111,7 +116,9 @@ def default_joint_vel(self) -> wp.array: @property @abstractmethod def joint_pos_target(self) -> wp.array: - """Joint position targets commanded by the user. Shape is (num_instances, num_joints). + """Joint position targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). For an implicit actuator model, the targets are directly set into the simulation. For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), @@ -122,7 +129,9 @@ def joint_pos_target(self) -> wp.array: @property @abstractmethod def joint_vel_target(self) -> wp.array: - """Joint velocity targets commanded by the user. Shape is (num_instances, num_joints). + """Joint velocity targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). For an implicit actuator model, the targets are directly set into the simulation. For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), @@ -133,7 +142,9 @@ def joint_vel_target(self) -> wp.array: @property @abstractmethod def joint_effort_target(self) -> wp.array: - """Joint effort targets commanded by the user. Shape is (num_instances, num_joints). + """Joint effort targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). For an implicit actuator model, the targets are directly set into the simulation. For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), @@ -148,7 +159,9 @@ def joint_effort_target(self) -> wp.array: @property @abstractmethod def computed_torque(self) -> wp.array: - """Joint torques computed from the actuator model (before clipping). Shape is (num_instances, num_joints). + """Joint torques computed from the actuator model (before clipping). + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is the raw torque output from the actuator mode, before any clipping is applied. It is exposed for users who want to inspect the computations inside the actuator model. @@ -159,7 +172,9 @@ def computed_torque(self) -> wp.array: @property @abstractmethod def applied_torque(self) -> wp.array: - """Joint torques applied from the actuator model (after clipping). Shape is (num_instances, num_joints). + """Joint torques applied from the actuator model (after clipping). + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the actuator model. @@ -173,7 +188,9 @@ def applied_torque(self) -> wp.array: @property @abstractmethod def joint_stiffness(self) -> wp.array: - """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints). + """Joint stiffness provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). In the case of explicit actuators, the value for the corresponding joints is zero. """ @@ -182,7 +199,9 @@ def joint_stiffness(self) -> wp.array: @property @abstractmethod def joint_damping(self) -> wp.array: - """Joint damping provided to the simulation. Shape is (num_instances, num_joints) + """Joint damping provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). In the case of explicit actuators, the value for the corresponding joints is zero. """ @@ -191,31 +210,28 @@ def joint_damping(self) -> wp.array: @property @abstractmethod def joint_armature(self) -> wp.array: - """Joint armature provided to the simulation. Shape is (num_instances, num_joints).""" - raise NotImplementedError + """Joint armature provided to the simulation. - @property - @abstractmethod - def joint_friction_coeff(self) -> wp.array: - """Joint static friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ raise NotImplementedError @property @abstractmethod - def joint_dynamic_friction_coeff(self) -> wp.array: - """Joint dynamic friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" - raise NotImplementedError + def joint_friction_coeff(self) -> wp.array: + """Joint static friction coefficient provided to the simulation. - @property - @abstractmethod - def joint_viscous_friction_coeff(self) -> wp.array: - """Joint viscous friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ raise NotImplementedError @property @abstractmethod def joint_pos_limits(self) -> wp.array: - """Joint position limits provided to the simulation. Shape is (num_instances, num_joints, 2). + """Joint position limits provided to the simulation. + + Shape is (num_instances, num_joints, 2), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_joints, 2). The limits are in the order :math:`[lower, upper]`. """ @@ -224,13 +240,19 @@ def joint_pos_limits(self) -> wp.array: @property @abstractmethod def joint_vel_limits(self) -> wp.array: - """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints).""" + """Joint maximum velocity provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ raise NotImplementedError @property @abstractmethod def joint_effort_limits(self) -> wp.array: - """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints).""" + """Joint maximum effort provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ raise NotImplementedError ## @@ -240,7 +262,10 @@ def joint_effort_limits(self) -> wp.array: @property @abstractmethod def soft_joint_pos_limits(self) -> wp.array: - r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints, 2). + r"""Soft joint positions limits for all joints. + + Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_joints, 2). The limits are in the order :math:`[lower, upper]`.The soft joint position limits are computed as a sub-region of the :attr:`joint_pos_limits` based on the @@ -262,7 +287,9 @@ def soft_joint_pos_limits(self) -> wp.array: @property @abstractmethod def soft_joint_vel_limits(self) -> wp.array: - """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints). + """Soft joint velocity limits for all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model has a variable velocity limit model. For instance, in a variable gear ratio actuator model. @@ -272,7 +299,10 @@ def soft_joint_vel_limits(self) -> wp.array: @property @abstractmethod def gear_ratio(self) -> wp.array: - """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints).""" + """Gear ratio for relating motor torques to applied Joint torques. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ raise NotImplementedError ## @@ -282,37 +312,61 @@ def gear_ratio(self) -> wp.array: @property @abstractmethod def fixed_tendon_stiffness(self) -> wp.array: - """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon stiffness provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ raise NotImplementedError @property @abstractmethod def fixed_tendon_damping(self) -> wp.array: - """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon damping provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ raise NotImplementedError @property @abstractmethod def fixed_tendon_limit_stiffness(self) -> wp.array: - """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon limit stiffness provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ raise NotImplementedError @property @abstractmethod def fixed_tendon_rest_length(self) -> wp.array: - """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon rest length provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ raise NotImplementedError @property @abstractmethod def fixed_tendon_offset(self) -> wp.array: - """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon offset provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ raise NotImplementedError @property @abstractmethod def fixed_tendon_pos_limits(self) -> wp.array: - """Fixed tendon position limits provided to the simulation. Shape is (num_instances, num_fixed_tendons, 2).""" + """Fixed tendon position limits provided to the simulation. + + Shape is (num_instances, num_fixed_tendons, 2), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_fixed_tendons, 2). + """ raise NotImplementedError ## @@ -322,25 +376,41 @@ def fixed_tendon_pos_limits(self) -> wp.array: @property @abstractmethod def spatial_tendon_stiffness(self) -> wp.array: - """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + """Spatial tendon stiffness provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ raise NotImplementedError @property @abstractmethod def spatial_tendon_damping(self) -> wp.array: - """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + """Spatial tendon damping provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ raise NotImplementedError @property @abstractmethod def spatial_tendon_limit_stiffness(self) -> wp.array: - """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + """Spatial tendon limit stiffness provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ raise NotImplementedError @property @abstractmethod def spatial_tendon_offset(self) -> wp.array: - """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + """Spatial tendon offset provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ raise NotImplementedError ## @@ -350,7 +420,9 @@ def spatial_tendon_offset(self) -> wp.array: @property @abstractmethod def root_link_pose_w(self) -> wp.array: - """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + """Root link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the articulation root's actor frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -360,7 +432,9 @@ def root_link_pose_w(self) -> wp.array: @property @abstractmethod def root_link_vel_w(self) -> wp.array: - """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the articulation root's actor frame relative to the world. @@ -370,7 +444,9 @@ def root_link_vel_w(self) -> wp.array: @property @abstractmethod def root_com_pose_w(self) -> wp.array: - """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + """Root center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the articulation root's center of mass frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -380,7 +456,9 @@ def root_com_pose_w(self) -> wp.array: @property @abstractmethod def root_com_vel_w(self) -> wp.array: - """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the articulation root's center of mass frame relative to the world. @@ -412,20 +490,29 @@ def root_com_state_w(self) -> wp.array: @property @abstractmethod def body_mass(self) -> wp.array: - """Body mass ``wp.float32`` in the world frame. Shape is (num_instances, num_bodies).""" + """Body mass ``wp.float32`` in the world frame. + + Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). + """ raise NotImplementedError @property @abstractmethod def body_inertia(self) -> wp.array: - """Body inertia ``wp.mat33`` in the world frame. Shape is (num_instances, num_bodies, 3, 3).""" + """Flattened body inertia in the world frame. + + Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to + (num_instances, num_bodies, 9). + """ raise NotImplementedError @property @abstractmethod def body_link_pose_w(self) -> wp.array: """Body link pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies, 7). + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the articulation links' actor frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -436,7 +523,9 @@ def body_link_pose_w(self) -> wp.array: @abstractmethod def body_link_vel_w(self) -> wp.array: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies, 6). + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the articulation links' actor frame relative to the world. @@ -447,7 +536,9 @@ def body_link_vel_w(self) -> wp.array: @abstractmethod def body_com_pose_w(self) -> wp.array: """Body center of mass pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies, 7). + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the articulation links relative to the world. The orientation is provided in (x, y, z, w) format. @@ -458,7 +549,9 @@ def body_com_pose_w(self) -> wp.array: @abstractmethod def body_com_vel_w(self) -> wp.array: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies, 6). + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the articulation links' center of mass frame relative to the world. @@ -487,7 +580,9 @@ def body_com_state_w(self) -> wp.array: @abstractmethod def body_com_acc_w(self) -> wp.array: """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]``. - Shape is (num_instances, num_bodies, 6). + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). All values are relative to the world. """ @@ -497,7 +592,9 @@ def body_com_acc_w(self) -> wp.array: @abstractmethod def body_com_pose_b(self) -> wp.array: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, 1, 7). + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. The orientation is provided in (x, y, z, w) format. @@ -509,7 +606,8 @@ def body_com_pose_b(self) -> wp.array: def body_incoming_joint_wrench_b(self) -> wp.array: """Joint reaction wrench applied from body parent to child body in parent body frame. - Shape is (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the world of an articulation. For more information on joint wrenches, please check the `PhysX documentation`_ and the @@ -527,19 +625,31 @@ def body_incoming_joint_wrench_b(self) -> wp.array: @property @abstractmethod def joint_pos(self) -> wp.array: - """Joint positions of all joints. Shape is (num_instances, num_joints).""" + """Joint positions of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to + (num_instances, num_joints). + """ raise NotImplementedError @property @abstractmethod def joint_vel(self) -> wp.array: - """Joint velocities of all joints. Shape is (num_instances, num_joints).""" + """Joint velocities of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to + (num_instances, num_joints). + """ raise NotImplementedError @property @abstractmethod def joint_acc(self) -> wp.array: - """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" + """Joint acceleration of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to + (num_instances, num_joints). + """ raise NotImplementedError ## @@ -549,13 +659,18 @@ def joint_acc(self) -> wp.array: @property @abstractmethod def projected_gravity_b(self) -> wp.array: - """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + """Projection of the gravity direction on base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ raise NotImplementedError @property @abstractmethod def heading_w(self) -> wp.array: - """Yaw heading of the base frame (in radians). Shape is (num_instances,). + """Yaw heading of the base frame (in radians). + + Shape is (num_instances), dtype = wp.float32. In torch this resolves to (num_instances,). .. note:: This quantity is computed by assuming that the forward-direction of the base @@ -566,9 +681,11 @@ def heading_w(self) -> wp.array: @property @abstractmethod def root_link_lin_vel_b(self) -> wp.array: - """Root link linear velocity in base frame. Shape is (num_instances, 3). + """Root link linear velocity in base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the linear velocity of the articulation root's actor frame with respect to the + This quantity is the linear velocity of the articulation root's actor frame with respect to its actor frame. """ raise NotImplementedError @@ -576,9 +693,11 @@ def root_link_lin_vel_b(self) -> wp.array: @property @abstractmethod def root_link_ang_vel_b(self) -> wp.array: - """Root link angular velocity in base world frame. Shape is (num_instances, 3). + """Root link angular velocity in base frame. - This quantity is the angular velocity of the articulation root's actor frame with respect to the + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the angular velocity of the articulation root's actor frame with respect to its actor frame. """ raise NotImplementedError @@ -586,9 +705,11 @@ def root_link_ang_vel_b(self) -> wp.array: @property @abstractmethod def root_com_lin_vel_b(self) -> wp.array: - """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the linear velocity of the articulation root's center of mass frame with respect to the + This quantity is the linear velocity of the articulation root's center of mass frame with respect to its actor frame. """ raise NotImplementedError @@ -596,9 +717,11 @@ def root_com_lin_vel_b(self) -> wp.array: @property @abstractmethod def root_com_ang_vel_b(self) -> wp.array: - """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the angular velocity of the articulation root's center of mass frame with respect to the + This quantity is the angular velocity of the articulation root's center of mass frame with respect to its actor frame. """ raise NotImplementedError @@ -610,7 +733,9 @@ def root_com_ang_vel_b(self) -> wp.array: @property @abstractmethod def root_link_pos_w(self) -> wp.array: - """Root link position in simulation world frame. Shape is (num_instances, 3). + """Root link position in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ @@ -619,7 +744,9 @@ def root_link_pos_w(self) -> wp.array: @property @abstractmethod def root_link_quat_w(self) -> wp.array: - """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances, 4). + """Root link orientation (x, y, z, w) in simulation world frame. + + Shape is (num_instances), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ @@ -628,7 +755,9 @@ def root_link_quat_w(self) -> wp.array: @property @abstractmethod def root_link_lin_vel_w(self) -> wp.array: - """Root linear velocity in simulation world frame. Shape is (num_instances, 3). + """Root linear velocity in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ @@ -637,7 +766,9 @@ def root_link_lin_vel_w(self) -> wp.array: @property @abstractmethod def root_link_ang_vel_w(self) -> wp.array: - """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + """Root link angular velocity in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ @@ -646,25 +777,31 @@ def root_link_ang_vel_w(self) -> wp.array: @property @abstractmethod def root_com_pos_w(self) -> wp.array: - """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + """Root center of mass position in simulation world frame. - This quantity is the position of the actor frame of the root rigid body relative to the world. + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ raise NotImplementedError @property @abstractmethod def root_com_quat_w(self) -> wp.array: - """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances, 4). + """Root center of mass orientation (x, y, z, w) in simulation world frame. - This quantity is the orientation of the actor frame of the root rigid body relative to the world. + Shape is (num_instances), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + + This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ raise NotImplementedError @property @abstractmethod def root_com_lin_vel_w(self) -> wp.array: - """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ @@ -673,7 +810,9 @@ def root_com_lin_vel_w(self) -> wp.array: @property @abstractmethod def root_com_ang_vel_w(self) -> wp.array: - """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ @@ -682,7 +821,10 @@ def root_com_ang_vel_w(self) -> wp.array: @property @abstractmethod def body_link_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the position of the articulation bodies' actor frame relative to the world. """ @@ -691,7 +833,10 @@ def body_link_pos_w(self) -> wp.array: @property @abstractmethod def body_link_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). This quantity is the orientation of the articulation bodies' actor frame relative to the world. """ @@ -700,44 +845,58 @@ def body_link_quat_w(self) -> wp.array: @property @abstractmethod def body_link_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). - This quantity is the linear velocity of the articulation bodies' center of mass frame relative to the world. + This quantity is the linear velocity of the articulation bodies' actor frame relative to the world. """ raise NotImplementedError @property @abstractmethod def body_link_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular velocity of all bodies in simulation world frame. - This quantity is the angular velocity of the articulation bodies' center of mass frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the angular velocity of the articulation bodies' actor frame relative to the world. """ raise NotImplementedError @property @abstractmethod def body_com_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). - This quantity is the position of the articulation bodies' actor frame. + This quantity is the position of the articulation bodies' center of mass frame. """ raise NotImplementedError @property @abstractmethod def body_com_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all bodies in simulation world frame. - Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). - This quantity is the orientation of the articulation bodies' actor frame. + This quantity is the orientation of the principal axes of inertia of the articulation bodies. """ raise NotImplementedError @property @abstractmethod def body_com_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear velocity of the articulation bodies' center of mass frame. """ @@ -746,7 +905,10 @@ def body_com_lin_vel_w(self) -> wp.array: @property @abstractmethod def body_com_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular velocity of the articulation bodies' center of mass frame. """ @@ -755,7 +917,10 @@ def body_com_ang_vel_w(self) -> wp.array: @property @abstractmethod def body_com_lin_acc_w(self) -> wp.array: - """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear acceleration of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear acceleration of the articulation bodies' center of mass frame. """ @@ -764,7 +929,10 @@ def body_com_lin_acc_w(self) -> wp.array: @property @abstractmethod def body_com_ang_acc_w(self) -> wp.array: - """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular acceleration of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular acceleration of the articulation bodies' center of mass frame. """ @@ -774,19 +942,24 @@ def body_com_ang_acc_w(self) -> wp.array: @abstractmethod def body_com_pos_b(self) -> wp.array: """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies, 3). - This quantity is the center of mass location relative to its body'slink frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the center of mass location relative to its body's link frame. """ raise NotImplementedError @property @abstractmethod def body_com_quat_b(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all of the bodies in their - respective link frames. Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link + frames. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). - This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ raise NotImplementedError diff --git a/source/isaaclab/isaaclab/assets/asset_base.py b/source/isaaclab/isaaclab/assets/asset_base.py index bdc96e79c56..e6bd12220b3 100644 --- a/source/isaaclab/isaaclab/assets/asset_base.py +++ b/source/isaaclab/isaaclab/assets/asset_base.py @@ -13,12 +13,10 @@ from typing import TYPE_CHECKING, Any import torch -from isaaclab_physx.physics import IsaacEvents, PhysxManager - -import omni.kit.app import isaaclab.sim as sim_utils from isaaclab.physics import PhysicsEvent, PhysicsManager +from isaaclab.sim.simulation_context import SimulationContext from isaaclab.sim.utils.stage import get_current_stage if TYPE_CHECKING: @@ -43,10 +41,10 @@ class AssetBase(ABC): at the configured path. For more information on the spawn configuration, see the :mod:`isaaclab.sim.spawners` module. - Unlike Isaac Sim interface, where one usually needs to call the - :meth:`isaacsim.core.prims.XFormPrim.initialize` method to initialize the PhysX handles, the asset - class automatically initializes and invalidates the PhysX handles when the stage is played/stopped. This - is done by registering callbacks for the stage play/stop events. + Unlike backend-specific interfaces (e.g. Isaac Sim PhysX) where one usually needs to call + initialize explicitly, the asset class automatically initializes and invalidates physics + handles when the simulation is ready or stopped. This is done by registering callbacks + for the physics lifecycle events (:attr:`PhysicsEvent.PHYSICS_READY`, :attr:`PhysicsEvent.STOP`). Additionally, the class registers a callback for debug visualization of the asset if a debug visualization is implemented in the asset class. This can be enabled by setting the :attr:`AssetBaseCfg.debug_vis` attribute @@ -193,16 +191,18 @@ def set_debug_vis(self, debug_vis: bool) -> bool: return False # toggle debug visualization objects self._set_debug_vis_impl(debug_vis) - # toggle debug visualization handles + # toggle debug visualization handles (Kit/omni only for PhysX backend) if debug_vis: - # create a subscriber for the post update event if it doesn't exist if self._debug_vis_handle is None: - app_interface = omni.kit.app.get_app_interface() - self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( - lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) - ) + sim_ctx = SimulationContext.instance() + if sim_ctx and "Physx" in sim_ctx.physics_manager.__name__: + import omni.kit.app + + app_interface = omni.kit.app.get_app_interface() + self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( + lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) + ) else: - # remove the subscriber if it exists if self._debug_vis_handle is not None: self._debug_vis_handle.unsubscribe() self._debug_vis_handle = None @@ -241,7 +241,7 @@ def update(self, dt: float): @abstractmethod def _initialize_impl(self): - """Initializes the PhysX handles and internal buffers.""" + """Initializes the physics handles and internal buffers for the current backend.""" raise NotImplementedError def _set_debug_vis_impl(self, debug_vis: bool): @@ -265,9 +265,9 @@ def _debug_vis_callback(self, event): """ def _register_callbacks(self): - """Registers physics lifecycle and prim deletion callbacks.""" + """Registers physics lifecycle callbacks via the current backend's physics manager.""" + physics_mgr_cls = SimulationContext.instance().physics_manager - # register simulator callbacks (with weakref safety to avoid crashes on deletion) def safe_callback(callback_name, event, obj_ref): """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted.""" try: @@ -280,42 +280,50 @@ def safe_callback(callback_name, event, obj_ref): # note: use weakref on callbacks to ensure that this object can be deleted when its destructor is called. obj_ref = weakref.proxy(self) - # Register PHYSICS_READY callback for initialization (order=10 for lower priority) - self._initialize_handle = PhysxManager.register_callback( + # Backend-agnostic: PHYSICS_READY (init) and STOP (invalidate) + self._initialize_handle = physics_mgr_cls.register_callback( lambda payload, obj_ref=obj_ref: safe_callback("_initialize_callback", payload, obj_ref), PhysicsEvent.PHYSICS_READY, order=10, ) - # Register TIMELINE_STOP callback for invalidation (PhysX-specific) - self._invalidate_initialize_handle = PhysxManager.register_callback( - lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref), - IsaacEvents.TIMELINE_STOP, + self._invalidate_initialize_handle = physics_mgr_cls.register_callback( + lambda payload, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", payload, obj_ref), + PhysicsEvent.STOP, order=10, ) - # Register PRIM_DELETION callback (PhysX-specific) - self._prim_deletion_handle = PhysxManager.register_callback( - lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref), - IsaacEvents.PRIM_DELETION, - ) + # Optional: prim deletion (only supported by PhysX backend) + self._prim_deletion_handle = None + physics_backend = physics_mgr_cls.__name__ + if "Physx" in physics_backend: + from isaaclab_physx.physics import IsaacEvents + + self._prim_deletion_handle = physics_mgr_cls.register_callback( + lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref), + IsaacEvents.PRIM_DELETION, + ) def _initialize_callback(self, event): """Initializes the scene elements. .. note:: - PhysX handles are only enabled once the simulator starts playing. Hence, this function needs to be - called whenever the simulator "plays" from a "stop" state. + Physics handles are only valid once the simulation is ready. This callback runs when + :attr:`PhysicsEvent.PHYSICS_READY` is dispatched by the current backend. """ if not self._is_initialized: - # obtain simulation related information self._backend = PhysicsManager.get_backend() self._device = PhysicsManager.get_device() - # initialize the asset try: self._initialize_impl() except Exception as e: - # Store exception to be raised after callback completes - PhysxManager.store_callback_exception(e) - # set flag + store_fn = getattr( + SimulationContext.instance().physics_manager, + "store_callback_exception", + None, + ) + if callable(store_fn): + store_fn(e) + else: + raise self._is_initialized = True def _invalidate_initialize_callback(self, event): @@ -326,15 +334,12 @@ def _invalidate_initialize_callback(self, event): self._debug_vis_handle = None def _on_prim_deletion(self, event) -> None: - """Invalidates and deletes the callbacks when the prim is deleted. + """Invalidates and clears callbacks when the prim is deleted. - Args: - event: The prim deletion event containing the prim path in payload. - - Note: - This function is called when the prim is deleted. + Only used when the backend supports prim deletion events (e.g. PhysX). """ - prim_path = event.payload["prim_path"] + payload = getattr(event, "payload", event) if not isinstance(event, dict) else event + prim_path = payload.get("prim_path", "") if isinstance(payload, dict) else "" if prim_path == "/": self._clear_callbacks() return @@ -344,12 +349,8 @@ def _on_prim_deletion(self, event) -> None: if result: self._clear_callbacks() - def _clear_callbacks(self, event: Any = None) -> None: - """Clears the callbacks. - - Args: - event: Optional event that triggered the callback (unused but required for event handlers). - """ + def _clear_callbacks(self) -> None: + """Clears all registered callbacks.""" if self._initialize_handle is not None: self._initialize_handle.deregister() self._initialize_handle = None @@ -359,7 +360,6 @@ def _clear_callbacks(self, event: Any = None) -> None: if self._prim_deletion_handle is not None: self._prim_deletion_handle.deregister() self._prim_deletion_handle = None - # Clear debug visualization if self._debug_vis_handle is not None: self._debug_vis_handle.unsubscribe() self._debug_vis_handle = None diff --git a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py index c8988407eb6..a2e57aed541 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py +++ b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object.py @@ -99,13 +99,24 @@ def root_view(self): @property @abstractmethod def instantaneous_wrench_composer(self) -> WrenchComposer: - """Instantaneous wrench composer for the rigid object.""" + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ raise NotImplementedError() @property @abstractmethod def permanent_wrench_composer(self) -> WrenchComposer: - """Permanent wrench composer for the rigid object.""" + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ raise NotImplementedError() """ @@ -123,7 +134,7 @@ def reset( Args: env_ids: Environment indices. If None, then all indices are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -173,6 +184,7 @@ def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = Fal @abstractmethod def write_root_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -188,7 +200,8 @@ def write_root_pose_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) or + (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -196,10 +209,11 @@ def write_root_pose_to_sim_index( @abstractmethod def write_root_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root pose over selected environment indices into the simulation. + """Set the root pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). @@ -211,14 +225,16 @@ def write_root_pose_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) or + (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_link_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -234,7 +250,8 @@ def write_root_link_pose_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7). + root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) or + (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -242,10 +259,11 @@ def write_root_link_pose_to_sim_index( @abstractmethod def write_root_link_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root link pose over selected environment indices into the simulation. + """Set the root link pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). @@ -257,21 +275,23 @@ def write_root_link_pose_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root link poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root link poses in simulation frame. Shape is (num_instances, 7) or + (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_com_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root center of mass pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). - The orientation is the orientation of the principle axes of inertia. + The orientation is the orientation of the principal axes of inertia. .. note:: This method expects partial data. @@ -281,7 +301,8 @@ def write_root_com_pose_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7). + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or + (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -289,13 +310,14 @@ def write_root_com_pose_to_sim_index( @abstractmethod def write_root_com_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass pose over selected environment indices into the simulation. + """Set the root center of mass pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). - The orientation is the orientation of the principle axes of inertia. + The orientation is the orientation of the principal axes of inertia. .. note:: This method expects full data. @@ -305,14 +327,16 @@ def write_root_com_pose_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) or + (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -321,7 +345,7 @@ def write_root_velocity_to_sim_index( The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's center of mass rather than the roots frame. + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects partial data. @@ -331,7 +355,8 @@ def write_root_velocity_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -339,15 +364,16 @@ def write_root_velocity_to_sim_index( @abstractmethod def write_root_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass velocity over selected environment indices into the simulation. + """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's center of mass rather than the roots frame. + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects full data. @@ -357,14 +383,16 @@ def write_root_velocity_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_com_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -373,7 +401,7 @@ def write_root_com_velocity_to_sim_index( The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's center of mass rather than the roots frame. + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects partial data. @@ -383,7 +411,8 @@ def write_root_com_velocity_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -391,15 +420,16 @@ def write_root_com_velocity_to_sim_index( @abstractmethod def write_root_com_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root center of mass velocity over selected environment indices into the simulation. + """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's center of mass rather than the roots frame. + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects full data. @@ -409,14 +439,16 @@ def write_root_com_velocity_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_root_link_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -425,7 +457,7 @@ def write_root_link_velocity_to_sim_index( The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's frame rather than the roots center of mass. + This sets the velocity of the root's frame rather than the root's center of mass. .. note:: This method expects partial data. @@ -435,7 +467,8 @@ def write_root_link_velocity_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6). + root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @@ -443,15 +476,16 @@ def write_root_link_velocity_to_sim_index( @abstractmethod def write_root_link_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: - """Set the root link velocity over selected environment indices into the simulation. + """Set the root link velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. .. note:: - This sets the velocity of the root's frame rather than the roots center of mass. + This sets the velocity of the root's frame rather than the root's center of mass. .. note:: This method expects full data. @@ -461,8 +495,9 @@ def write_root_link_velocity_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -473,8 +508,9 @@ def write_root_link_velocity_to_sim_mask( @abstractmethod def set_masses_index( self, + *, masses: torch.Tensor | wp.array, - body_ids: Sequence[int] | slice | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set masses of all bodies. @@ -496,9 +532,10 @@ def set_masses_index( @abstractmethod def set_masses_mask( self, + *, masses: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set masses of all bodies. @@ -511,16 +548,17 @@ def set_masses_mask( Args: masses: Masses of all bodies. Shape is (num_instances, num_bodies). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_coms_index( self, + *, coms: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set center of mass positions of all bodies. @@ -543,9 +581,10 @@ def set_coms_index( @abstractmethod def set_coms_mask( self, + *, coms: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set center of mass positions of all bodies. @@ -557,17 +596,19 @@ def set_coms_mask( Some backends may provide optimized implementations for masks / indices. Args: - coms: Center of mass positions of all bodies. Shape is (num_instances, num_bodies, 3). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + coms: Center of mass positions of all bodies. Shape is (num_instances, num_bodies, 3) + or (num_instances, num_bodies) with dtype wp.vec3f. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_inertias_index( self, + *, inertias: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set inertias of all bodies. @@ -589,9 +630,10 @@ def set_inertias_index( @abstractmethod def set_inertias_mask( self, + *, inertias: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set inertias of all bodies. @@ -604,8 +646,8 @@ def set_inertias_mask( Args: inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -679,7 +721,7 @@ def write_root_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_pose_to_sim_index(root_pose, env_ids=env_ids) + self.write_root_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_link_pose_to_sim( self, @@ -693,7 +735,7 @@ def write_root_link_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_link_pose_to_sim_index(root_pose, env_ids=env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_com_pose_to_sim( self, @@ -707,7 +749,7 @@ def write_root_com_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_com_pose_to_sim_index(root_pose, env_ids=env_ids) + self.write_root_com_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_velocity_to_sim( self, @@ -721,7 +763,7 @@ def write_root_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_velocity_to_sim_index(root_velocity, env_ids=env_ids) + self.write_root_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def write_root_com_velocity_to_sim( self, @@ -735,7 +777,7 @@ def write_root_com_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_com_velocity_to_sim_index(root_velocity, env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def write_root_link_velocity_to_sim( self, @@ -749,7 +791,7 @@ def write_root_link_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_link_velocity_to_sim_index(root_velocity, env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def set_masses( self, @@ -763,7 +805,7 @@ def set_masses( DeprecationWarning, stacklevel=2, ) - self.set_masses_index(masses, body_ids=body_ids, env_ids=env_ids) + self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids) def set_coms( self, @@ -777,7 +819,7 @@ def set_coms( DeprecationWarning, stacklevel=2, ) - self.set_coms_index(coms, body_ids=body_ids, env_ids=env_ids) + self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids) def set_inertias( self, @@ -792,7 +834,7 @@ def set_inertias( DeprecationWarning, stacklevel=2, ) - self.set_inertias_index(inertias, body_ids=body_ids, env_ids=env_ids) + self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids) def set_external_force_and_torque( self, diff --git a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py index 366916218d1..1f92ee9b5d5 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py +++ b/source/isaaclab/isaaclab/assets/rigid_object/base_rigid_object_data.py @@ -64,7 +64,8 @@ def update(self, dt: float) -> None: def default_root_pose(self) -> wp.array: """Default root pose ``[pos, quat]`` in local environment frame. - The position and quaternion are of the rigid body's actor frame. Shape is (num_instances, 7). + The position and quaternion are of the rigid body's actor frame. + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). """ raise NotImplementedError() @@ -73,7 +74,8 @@ def default_root_pose(self) -> wp.array: def default_root_vel(self) -> wp.array: """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame. - The linear and angular velocities are of the rigid body's center of mass frame. Shape is (num_instances, 6). + The linear and angular velocities are of the rigid body's center of mass frame. + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). """ raise NotImplementedError() @@ -90,7 +92,9 @@ def default_root_state(self) -> wp.array: @property @abstractmethod def root_link_pose_w(self) -> wp.array: - """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + """Root link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the actor frame of the root rigid body relative to the world. The orientation is provided in (x, y, z, w) format. @@ -100,7 +104,9 @@ def root_link_pose_w(self) -> wp.array: @property @abstractmethod def root_link_vel_w(self) -> wp.array: - """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. @@ -110,7 +116,9 @@ def root_link_vel_w(self) -> wp.array: @property @abstractmethod def root_com_pose_w(self) -> wp.array: - """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + """Root center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the center of mass frame of the root rigid body relative to the world. The orientation is provided in (x, y, z, w) format. @@ -120,7 +128,9 @@ def root_com_pose_w(self) -> wp.array: @property @abstractmethod def root_com_vel_w(self) -> wp.array: - """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. @@ -152,7 +162,10 @@ def root_com_state_w(self) -> wp.array: @property @abstractmethod def body_link_pose_w(self) -> wp.array: - """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1, 7). + """Body link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.transformf. + In torch this resolves to (num_instances, 1, 7). This quantity is the pose of the actor frame of the rigid body relative to the world. The orientation is provided in (x, y, z, w) format. @@ -162,7 +175,10 @@ def body_link_pose_w(self) -> wp.array: @property @abstractmethod def body_link_vel_w(self) -> wp.array: - """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1, 6). + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 1, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. @@ -172,7 +188,10 @@ def body_link_vel_w(self) -> wp.array: @property @abstractmethod def body_com_pose_w(self) -> wp.array: - """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1, 7). + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.transformf. + In torch this resolves to (num_instances, 1, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the world. The orientation is provided in (x, y, z, w) format. @@ -183,7 +202,9 @@ def body_com_pose_w(self) -> wp.array: @abstractmethod def body_com_vel_w(self) -> wp.array: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, 1, 6). + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 1, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. @@ -212,7 +233,9 @@ def body_com_state_w(self) -> wp.array: @abstractmethod def body_com_acc_w(self) -> wp.array: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. - Shape is (num_instances, 1, 6). + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. + In torch this resolves to (num_instances, 1, 6). This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. """ @@ -222,7 +245,9 @@ def body_com_acc_w(self) -> wp.array: @abstractmethod def body_com_pose_b(self) -> wp.array: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, 1, 7). + + Shape is (num_instances, 1), dtype = wp.transformf. + In torch this resolves to (num_instances, 1, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. The orientation is provided in (x, y, z, w) format. @@ -232,13 +257,21 @@ def body_com_pose_b(self) -> wp.array: @property @abstractmethod def body_mass(self) -> wp.array: - """Mass of all bodies in the simulation world frame. Shape is (num_instances, 1, 1).""" + """Mass of all bodies in the simulation world frame. + + Shape is (num_instances, 1, 1), dtype = wp.float32. + In torch this resolves to (num_instances, 1, 1). + """ raise NotImplementedError() @property @abstractmethod def body_inertia(self) -> wp.array: - """Inertia of all bodies in the simulation world frame. Shape is (num_instances, 1, 9).""" + """Inertia of all bodies in the simulation world frame. + + Shape is (num_instances, 1, 9), dtype = wp.float32. + In torch this resolves to (num_instances, 1, 9). + """ raise NotImplementedError() ## @@ -248,13 +281,18 @@ def body_inertia(self) -> wp.array: @property @abstractmethod def projected_gravity_b(self) -> wp.array: - """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + """Projection of the gravity direction on base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ raise NotImplementedError() @property @abstractmethod def heading_w(self) -> wp.array: - """Yaw heading of the base frame (in radians). Shape is (num_instances,). + """Yaw heading of the base frame (in radians). + + Shape is (num_instances,), dtype = wp.float32. In torch this resolves to (num_instances,). .. note:: This quantity is computed by assuming that the forward-direction of the base @@ -265,7 +303,9 @@ def heading_w(self) -> wp.array: @property @abstractmethod def root_link_lin_vel_b(self) -> wp.array: - """Root link linear velocity in base frame. Shape is (num_instances, 3). + """Root link linear velocity in base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. @@ -275,7 +315,9 @@ def root_link_lin_vel_b(self) -> wp.array: @property @abstractmethod def root_link_ang_vel_b(self) -> wp.array: - """Root link angular velocity in base world frame. Shape is (num_instances, 3). + """Root link angular velocity in base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. @@ -285,7 +327,9 @@ def root_link_ang_vel_b(self) -> wp.array: @property @abstractmethod def root_com_lin_vel_b(self) -> wp.array: - """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. @@ -295,7 +339,9 @@ def root_com_lin_vel_b(self) -> wp.array: @property @abstractmethod def root_com_ang_vel_b(self) -> wp.array: - """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. @@ -309,7 +355,9 @@ def root_com_ang_vel_b(self) -> wp.array: @property @abstractmethod def root_link_pos_w(self) -> wp.array: - """Root link position in simulation world frame. Shape is (num_instances, 3). + """Root link position in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ @@ -318,7 +366,9 @@ def root_link_pos_w(self) -> wp.array: @property @abstractmethod def root_link_quat_w(self) -> wp.array: - """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances, 4). + """Root link orientation (x, y, z, w) in simulation world frame. + + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ @@ -327,7 +377,9 @@ def root_link_quat_w(self) -> wp.array: @property @abstractmethod def root_link_lin_vel_w(self) -> wp.array: - """Root linear velocity in simulation world frame. Shape is (num_instances, 3). + """Root linear velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ @@ -336,7 +388,9 @@ def root_link_lin_vel_w(self) -> wp.array: @property @abstractmethod def root_link_ang_vel_w(self) -> wp.array: - """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + """Root link angular velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ @@ -345,25 +399,31 @@ def root_link_ang_vel_w(self) -> wp.array: @property @abstractmethod def root_com_pos_w(self) -> wp.array: - """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + """Root center of mass position in simulation world frame. - This quantity is the position of the actor frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ raise NotImplementedError() @property @abstractmethod def root_com_quat_w(self) -> wp.array: - """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances, 4). + """Root center of mass orientation (x, y, z, w) in simulation world frame. - This quantity is the orientation of the actor frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + + This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ raise NotImplementedError() @property @abstractmethod def root_com_lin_vel_w(self) -> wp.array: - """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ @@ -372,7 +432,9 @@ def root_com_lin_vel_w(self) -> wp.array: @property @abstractmethod def root_com_ang_vel_w(self) -> wp.array: - """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ @@ -381,7 +443,9 @@ def root_com_ang_vel_w(self) -> wp.array: @property @abstractmethod def body_link_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ @@ -390,52 +454,64 @@ def body_link_pos_w(self) -> wp.array: @property @abstractmethod def body_link_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, 1, 4). + """Orientation (x, y, z, w) of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). - This quantity is the orientation of the rigid bodies' actor frame relative to the world. + This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ raise NotImplementedError() @property @abstractmethod def body_link_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear velocity of all bodies in simulation world frame. - This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + + This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. """ raise NotImplementedError() @property @abstractmethod def body_link_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). - This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. """ raise NotImplementedError() @property @abstractmethod def body_com_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Positions of all bodies' center of mass in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). - This quantity is the position of the rigid bodies' actor frame. + This quantity is the position of the rigid bodies' center of mass frame. """ raise NotImplementedError() @property @abstractmethod def body_com_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all bodies in simulation world frame. + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. - Shape is (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame. + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + + This quantity is the orientation of the rigid bodies' principal axes of inertia. """ raise NotImplementedError() @property @abstractmethod def body_com_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ @@ -444,7 +520,9 @@ def body_com_lin_vel_w(self) -> wp.array: @property @abstractmethod def body_com_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ @@ -453,7 +531,9 @@ def body_com_ang_vel_w(self) -> wp.array: @property @abstractmethod def body_com_lin_acc_w(self) -> wp.array: - """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear acceleration of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ @@ -462,7 +542,9 @@ def body_com_lin_acc_w(self) -> wp.array: @property @abstractmethod def body_com_ang_acc_w(self) -> wp.array: - """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular acceleration of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ @@ -472,19 +554,22 @@ def body_com_ang_acc_w(self) -> wp.array: @abstractmethod def body_com_pos_b(self) -> wp.array: """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, 1, 3). - This quantity is the center of mass location relative to its body'slink frame. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + + This quantity is the center of mass location relative to its body's link frame. """ raise NotImplementedError() @property @abstractmethod def body_com_quat_b(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all of the bodies in their - respective link frames. Shape is (num_instances, 1, 4). + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their + respective link frames. + + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). - This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ raise NotImplementedError() diff --git a/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection.py b/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection.py index ab9d9a52c77..0384bc67da5 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection.py +++ b/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection.py @@ -97,13 +97,24 @@ def root_view(self): @property @abstractmethod def instantaneous_wrench_composer(self) -> WrenchComposer: - """Instantaneous wrench composer for the rigid object collection.""" + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ raise NotImplementedError() @property @abstractmethod def permanent_wrench_composer(self) -> WrenchComposer: - """Permanent wrench composer for the rigid object collection.""" + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ raise NotImplementedError() """ @@ -125,7 +136,7 @@ def reset( Args: env_ids: Environment indices. If None, then all indices are used. object_ids: Object indices. If None, then all indices are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -177,9 +188,10 @@ def find_bodies( @abstractmethod def write_body_pose_to_sim_index( self, + *, body_poses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - body_ids: slice | torch.Tensor | None = None, ) -> None: """Set the body poses over selected environment and body indices into the simulation. @@ -193,20 +205,22 @@ def write_body_pose_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - body_poses: Body poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7). - env_ids: Environment indices. If None, then all indices are used. + body_poses: Body poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7) + or (len(env_ids), len(body_ids)) with dtype wp.transformf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @abstractmethod def write_body_pose_to_sim_mask( self, + *, body_poses: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: - """Set the body poses over selected environment and body indices into the simulation. + """Set the body poses over selected environment and body mask into the simulation. The body pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). @@ -218,18 +232,20 @@ def write_body_pose_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - body_poses: Body poses in simulation frame. Shape is (num_instances, num_bodies, 7). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_poses: Body poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_body_link_pose_to_sim_index( self, + *, body_poses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - body_ids: slice | torch.Tensor | None = None, ) -> None: """Set the body link pose over selected environment and body indices into the simulation. @@ -243,20 +259,22 @@ def write_body_link_pose_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - body_poses: Body link poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7). - env_ids: Environment indices. If None, then all indices are used. + body_poses: Body link poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7) + or (len(env_ids), len(body_ids)) with dtype wp.transformf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @abstractmethod def write_body_link_pose_to_sim_mask( self, + *, body_poses: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: - """Set the body link pose over selected environment and body indices into the simulation. + """Set the body link pose over selected environment and body mask into the simulation. The body link pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). @@ -268,23 +286,25 @@ def write_body_link_pose_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - body_poses: Body link poses in simulation frame. Shape is (num_instances, num_bodies, 7). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_poses: Body link poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_body_com_pose_to_sim_index( self, + *, body_poses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - body_ids: slice | torch.Tensor | None = None, ) -> None: """Set the body center of mass pose over selected environment and body indices into the simulation. The body center of mass pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). - The orientation is the orientation of the principle axes of inertia. + The orientation is the orientation of the principal axes of inertia. .. note:: This method expects partial data. @@ -294,23 +314,25 @@ def write_body_com_pose_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - body_poses: Body center of mass poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7). - env_ids: Environment indices. If None, then all indices are used. + body_poses: Body center of mass poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7) + or (len(env_ids), len(body_ids)) with dtype wp.transformf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @abstractmethod def write_body_com_pose_to_sim_mask( self, + *, body_poses: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: - """Set the body center of mass pose over selected environment and body indices into the simulation. + """Set the body center of mass pose over selected environment and body mask into the simulation. The body center of mass pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). - The orientation is the orientation of the principle axes of inertia. + The orientation is the orientation of the principal axes of inertia. .. note:: This method expects full data. @@ -320,18 +342,20 @@ def write_body_com_pose_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - body_poses: Body center of mass poses in simulation frame. Shape is (num_instances, num_bodies, 7). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_poses: Body center of mass poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_body_velocity_to_sim_index( self, + *, body_velocities: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - body_ids: slice | torch.Tensor | None = None, ) -> None: """Set the body velocity over selected environment and body indices into the simulation. @@ -348,20 +372,22 @@ def write_body_velocity_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - body_velocities: Body velocities in simulation frame. Shape is (len(env_ids), len(body_ids), 6). - env_ids: Environment indices. If None, then all indices are used. + body_velocities: Body velocities in simulation frame. Shape is (len(env_ids), len(body_ids), 6) + or (len(env_ids), len(body_ids)) with dtype wp.spatial_vectorf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @abstractmethod def write_body_velocity_to_sim_mask( self, + *, body_velocities: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: - """Set the body velocity over selected environment and body indices into the simulation. + """Set the body velocity over selected environment and body mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. @@ -376,18 +402,20 @@ def write_body_velocity_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - body_velocities: Body velocities in simulation frame. Shape is (num_instances, num_bodies, 6). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_velocities: Body velocities in simulation frame. Shape is (num_instances, num_bodies, 6) + or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_body_com_velocity_to_sim_index( self, + *, body_velocities: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - body_ids: slice | torch.Tensor | None = None, ) -> None: """Set the body center of mass velocity over selected environment and body indices into the simulation. @@ -405,20 +433,21 @@ def write_body_com_velocity_to_sim_index( Args: body_velocities: Body center of mass velocities in simulation frame. Shape is - (len(env_ids), len(body_ids), 6). - env_ids: Environment indices. If None, then all indices are used. + (len(env_ids), len(body_ids), 6) or (len(env_ids), len(body_ids)) with dtype wp.spatial_vectorf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @abstractmethod def write_body_com_velocity_to_sim_mask( self, + *, body_velocities: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: - """Set the body center of mass velocity over selected environment and body indices into the simulation. + """Set the body center of mass velocity over selected environment and body mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. @@ -434,18 +463,19 @@ def write_body_com_velocity_to_sim_mask( Args: body_velocities: Body center of mass velocities in simulation frame. Shape is - (num_instances, num_bodies, 6). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + (num_instances, num_bodies, 6) or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def write_body_link_velocity_to_sim_index( self, + *, body_velocities: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, - body_ids: slice | torch.Tensor | None = None, ) -> None: """Set the body link velocity over selected environment and body indices into the simulation. @@ -462,20 +492,22 @@ def write_body_link_velocity_to_sim_index( Some backends may provide optimized implementations for masks / indices. Args: - body_velocities: Body link velocities in simulation frame. Shape is (len(env_ids), len(body_ids), 6). - env_ids: Environment indices. If None, then all indices are used. + body_velocities: Body link velocities in simulation frame. Shape is (len(env_ids), len(body_ids), 6) + or (len(env_ids), len(body_ids)) with dtype wp.spatial_vectorf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. """ raise NotImplementedError() @abstractmethod def write_body_link_velocity_to_sim_mask( self, + *, body_velocities: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: - """Set the body link velocity over selected environment and body indices into the simulation. + """Set the body link velocity over selected environment and body mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. @@ -490,9 +522,10 @@ def write_body_link_velocity_to_sim_mask( Some backends may provide optimized implementations for masks / indices. Args: - body_velocities: Body link velocities in simulation frame. Shape is (num_instances, num_bodies, 6). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_velocities: Body link velocities in simulation frame. Shape is (num_instances, num_bodies, 6) + or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -503,8 +536,9 @@ def write_body_link_velocity_to_sim_mask( @abstractmethod def set_masses_index( self, + *, masses: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set masses of all bodies. @@ -526,9 +560,10 @@ def set_masses_index( @abstractmethod def set_masses_mask( self, + *, masses: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set masses of all bodies. @@ -541,16 +576,17 @@ def set_masses_mask( Args: masses: Masses of all bodies. Shape is (num_instances, num_bodies). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_coms_index( self, + *, coms: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set center of mass positions of all bodies. @@ -573,9 +609,10 @@ def set_coms_index( @abstractmethod def set_coms_mask( self, + *, coms: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set center of mass positions of all bodies. @@ -587,17 +624,19 @@ def set_coms_mask( Some backends may provide optimized implementations for masks / indices. Args: - coms: Center of mass positions of all bodies. Shape is (num_instances, num_bodies, 3). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + coms: Center of mass positions of all bodies. Shape is (num_instances, num_bodies, 3) + or (num_instances, num_bodies) with dtype wp.vec3f. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @abstractmethod def set_inertias_index( self, + *, inertias: torch.Tensor | wp.array, - body_ids: Sequence[int] | None = None, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set inertias of all bodies. @@ -619,9 +658,10 @@ def set_inertias_index( @abstractmethod def set_inertias_mask( self, + *, inertias: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set inertias of all bodies. @@ -634,8 +674,8 @@ def set_inertias_mask( Args: inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ raise NotImplementedError() @@ -735,7 +775,7 @@ def write_body_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_pose_to_sim_index(body_poses, env_ids=env_ids, body_ids=body_ids) + self.write_body_pose_to_sim_index(body_poses=body_poses, env_ids=env_ids, body_ids=body_ids) def write_body_link_pose_to_sim( self, @@ -750,7 +790,7 @@ def write_body_link_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_link_pose_to_sim_index(body_poses, env_ids=env_ids, body_ids=body_ids) + self.write_body_link_pose_to_sim_index(body_poses=body_poses, env_ids=env_ids, body_ids=body_ids) def write_body_com_pose_to_sim( self, @@ -765,7 +805,7 @@ def write_body_com_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_com_pose_to_sim_index(body_poses, env_ids=env_ids, body_ids=body_ids) + self.write_body_com_pose_to_sim_index(body_poses=body_poses, env_ids=env_ids, body_ids=body_ids) def write_body_velocity_to_sim( self, @@ -780,7 +820,7 @@ def write_body_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_velocity_to_sim_index(body_velocities, env_ids=env_ids, body_ids=body_ids) + self.write_body_velocity_to_sim_index(body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids) def write_body_com_velocity_to_sim( self, @@ -795,7 +835,7 @@ def write_body_com_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_com_velocity_to_sim_index(body_velocities, env_ids=env_ids, body_ids=body_ids) + self.write_body_com_velocity_to_sim_index(body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids) def write_body_link_velocity_to_sim( self, @@ -810,7 +850,7 @@ def write_body_link_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_link_velocity_to_sim_index(body_velocities, env_ids=env_ids, body_ids=body_ids) + self.write_body_link_velocity_to_sim_index(body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids) def set_masses( self, @@ -824,7 +864,7 @@ def set_masses( DeprecationWarning, stacklevel=2, ) - self.set_masses_index(masses, body_ids=body_ids, env_ids=env_ids) + self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids) def set_coms( self, @@ -838,7 +878,7 @@ def set_coms( DeprecationWarning, stacklevel=2, ) - self.set_coms_index(coms, body_ids=body_ids, env_ids=env_ids) + self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids) def set_inertias( self, @@ -853,7 +893,7 @@ def set_inertias( DeprecationWarning, stacklevel=2, ) - self.set_inertias_index(inertias, body_ids=body_ids, env_ids=env_ids) + self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids) def set_external_force_and_torque( self, @@ -941,7 +981,7 @@ def write_object_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_pose_to_sim_index(object_pose, env_ids=env_ids, body_ids=object_ids) + self.write_body_pose_to_sim_index(body_poses=object_pose, env_ids=env_ids, body_ids=object_ids) def write_object_link_pose_to_sim( self, @@ -956,7 +996,7 @@ def write_object_link_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_link_pose_to_sim_index(object_pose, env_ids=env_ids, body_ids=object_ids) + self.write_body_link_pose_to_sim_index(body_poses=object_pose, env_ids=env_ids, body_ids=object_ids) def write_object_com_pose_to_sim( self, @@ -971,7 +1011,7 @@ def write_object_com_pose_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_com_pose_to_sim_index(object_pose, env_ids=env_ids, body_ids=object_ids) + self.write_body_com_pose_to_sim_index(body_poses=object_pose, env_ids=env_ids, body_ids=object_ids) def write_object_velocity_to_sim( self, @@ -986,7 +1026,7 @@ def write_object_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_com_velocity_to_sim_index(object_velocity, env_ids=env_ids, body_ids=object_ids) + self.write_body_com_velocity_to_sim_index(body_velocities=object_velocity, env_ids=env_ids, body_ids=object_ids) def write_object_com_velocity_to_sim( self, @@ -1001,7 +1041,7 @@ def write_object_com_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_com_velocity_to_sim_index(object_velocity, env_ids=env_ids, body_ids=object_ids) + self.write_body_com_velocity_to_sim_index(body_velocities=object_velocity, env_ids=env_ids, body_ids=object_ids) def write_object_link_velocity_to_sim( self, @@ -1016,7 +1056,9 @@ def write_object_link_velocity_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_link_velocity_to_sim_index(object_velocity, env_ids=env_ids, body_ids=object_ids) + self.write_body_link_velocity_to_sim_index( + body_velocities=object_velocity, env_ids=env_ids, body_ids=object_ids + ) def find_objects( self, name_keys: str | Sequence[str], preserve_order: bool = False diff --git a/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection_data.py b/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection_data.py index 2783a1143d3..1913bb4bef9 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection_data.py +++ b/source/isaaclab/isaaclab/assets/rigid_object_collection/base_rigid_object_collection_data.py @@ -66,17 +66,19 @@ def update(self, dt: float) -> None: def default_body_pose(self) -> wp.array: """Default body pose ``[pos, quat]`` in local environment frame. - The position and quaternion are of the rigid body's actor frame. Shape is (num_instances, num_bodies, 7). + The position and quaternion are of the rigid body's actor frame. + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). """ raise NotImplementedError() @property @abstractmethod def default_body_vel(self) -> wp.array: - """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. Shape is + """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. - - The linear and angular velocities are of the rigid body's center of mass frame. Shape is + The linear and angular velocities are of the rigid body's center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). """ raise NotImplementedError() @@ -94,7 +96,10 @@ def default_body_state(self) -> wp.array: @property @abstractmethod def body_link_pose_w(self) -> wp.array: - """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies, 7). + """Body link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the actor frame of the rigid body relative to the world. The orientation is provided in (x, y, z, w) format. @@ -104,7 +109,10 @@ def body_link_pose_w(self) -> wp.array: @property @abstractmethod def body_link_vel_w(self) -> wp.array: - """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies, 6). + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. @@ -114,7 +122,10 @@ def body_link_vel_w(self) -> wp.array: @property @abstractmethod def body_com_pose_w(self) -> wp.array: - """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies, 7). + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the world. The orientation is provided in (x, y, z, w) format. @@ -125,7 +136,9 @@ def body_com_pose_w(self) -> wp.array: @abstractmethod def body_com_vel_w(self) -> wp.array: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies, 6). + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. @@ -154,7 +167,9 @@ def body_com_state_w(self) -> wp.array: @abstractmethod def body_com_acc_w(self) -> wp.array: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. - Shape is (num_instances, num_bodies, 6). + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. """ @@ -164,7 +179,9 @@ def body_com_acc_w(self) -> wp.array: @abstractmethod def body_com_pose_b(self) -> wp.array: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, num_bodies, 7). + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. The orientation is provided in (x, y, z, w) format. @@ -174,13 +191,20 @@ def body_com_pose_b(self) -> wp.array: @property @abstractmethod def body_mass(self) -> wp.array: - """Mass of all bodies in the simulation world frame. Shape is (num_instances, num_bodies).""" + """Mass of all bodies in the simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). + """ raise NotImplementedError() @property @abstractmethod def body_inertia(self) -> wp.array: - """Inertia of all bodies in the simulation world frame. Shape is (num_instances, num_bodies, 9).""" + """Inertia of all bodies in the simulation world frame. + + Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to + (num_instances, num_bodies, 9). + """ raise NotImplementedError() ## @@ -190,13 +214,20 @@ def body_inertia(self) -> wp.array: @property @abstractmethod def projected_gravity_b(self) -> wp.array: - """Projection of the gravity direction on base frame. Shape is (num_instances, num_bodies, 3).""" + """Projection of the gravity direction on base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + """ raise NotImplementedError() @property @abstractmethod def heading_w(self) -> wp.array: - """Yaw heading of the base frame (in radians). Shape is (num_instances, num_bodies). + """Yaw heading of the base frame (in radians). + + Shape is (num_instances, num_bodies), dtype = wp.float32. + In torch this resolves to (num_instances, num_bodies). .. note:: This quantity is computed by assuming that the forward-direction of the base @@ -207,7 +238,10 @@ def heading_w(self) -> wp.array: @property @abstractmethod def body_link_lin_vel_b(self) -> wp.array: - """Root link linear velocity in base frame. Shape is (num_instances, num_bodies, 3). + """Root link linear velocity in base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. @@ -217,7 +251,10 @@ def body_link_lin_vel_b(self) -> wp.array: @property @abstractmethod def body_link_ang_vel_b(self) -> wp.array: - """Root link angular velocity in base world frame. Shape is (num_instances, num_bodies, 3). + """Root link angular velocity in base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. @@ -227,7 +264,10 @@ def body_link_ang_vel_b(self) -> wp.array: @property @abstractmethod def body_com_lin_vel_b(self) -> wp.array: - """Root center of mass linear velocity in base frame. Shape is (num_instances, num_bodies, 3). + """Root center of mass linear velocity in base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. @@ -237,7 +277,10 @@ def body_com_lin_vel_b(self) -> wp.array: @property @abstractmethod def body_com_ang_vel_b(self) -> wp.array: - """Root center of mass angular velocity in base world frame. Shape is (num_instances, num_bodies, 3). + """Root center of mass angular velocity in base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. @@ -251,7 +294,10 @@ def body_com_ang_vel_b(self) -> wp.array: @property @abstractmethod def body_link_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ @@ -260,52 +306,70 @@ def body_link_pos_w(self) -> wp.array: @property @abstractmethod def body_link_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). - This quantity is the orientation of the rigid bodies' actor frame relative to the world. + This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ raise NotImplementedError() @property @abstractmethod def body_link_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear velocity of all bodies in simulation world frame. - This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. """ raise NotImplementedError() @property @abstractmethod def body_link_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). - This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. """ raise NotImplementedError() @property @abstractmethod def body_com_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Positions of all bodies' center of mass in simulation world frame. - This quantity is the position of the rigid bodies' actor frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the position of the rigid bodies' center of mass frame. """ raise NotImplementedError() @property @abstractmethod def body_com_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all bodies in simulation world frame. + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). - Shape is (num_instances, num_bodies, 4). This quantity is the orientation of the rigid bodies' actor frame. + This quantity is the orientation of the rigid bodies' principal axes of inertia. """ raise NotImplementedError() @property @abstractmethod def body_com_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ @@ -314,7 +378,10 @@ def body_com_lin_vel_w(self) -> wp.array: @property @abstractmethod def body_com_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ @@ -323,7 +390,10 @@ def body_com_ang_vel_w(self) -> wp.array: @property @abstractmethod def body_com_lin_acc_w(self) -> wp.array: - """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear acceleration of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ @@ -332,7 +402,10 @@ def body_com_lin_acc_w(self) -> wp.array: @property @abstractmethod def body_com_ang_acc_w(self) -> wp.array: - """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular acceleration of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ @@ -342,19 +415,24 @@ def body_com_ang_acc_w(self) -> wp.array: @abstractmethod def body_com_pos_b(self) -> wp.array: """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies, 3). - This quantity is the center of mass location relative to its body'slink frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the center of mass location relative to its body's link frame. """ raise NotImplementedError() @property @abstractmethod def body_com_quat_b(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all of the bodies in their - respective link frames. Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their + respective link frames. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). - This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ raise NotImplementedError() diff --git a/source/isaaclab_newton/isaaclab_newton/assets/__init__.py b/source/isaaclab_newton/isaaclab_newton/assets/__init__.py new file mode 100644 index 00000000000..460a3056908 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/__init__.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/__init__.py new file mode 100644 index 00000000000..791c25607cd --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .articulation import Articulation +from .articulation_data import ArticulationData + +__all__ = ["Articulation", "ArticulationData"] diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py new file mode 100644 index 00000000000..c4dcdd6cace --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -0,0 +1,3858 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Flag for pyright to ignore type errors in this file. +# pyright: reportPrivateUsage=false + +from __future__ import annotations + +import logging +import warnings +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import numpy as np +import torch +import warp as wp +from newton import JointType +from newton.selection import ArticulationView +from newton.solvers import SolverNotifyFlags +from prettytable import PrettyTable + +from pxr import UsdPhysics + +from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator +from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims +from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values +from isaaclab.utils.types import ArticulationActions +from isaaclab.utils.version import get_isaac_sim_version, has_kit +from isaaclab.utils.wrench_composer import WrenchComposer + +from isaaclab_newton.assets import kernels as shared_kernels +from isaaclab_newton.assets.articulation import kernels as articulation_kernels +from isaaclab_newton.physics import NewtonManager as SimulationManager + +from .articulation_data import ArticulationData + +if TYPE_CHECKING: + from isaaclab.assets.articulation.articulation_cfg import ArticulationCfg + +# import logger +logger = logging.getLogger(__name__) + + +class Articulation(BaseArticulation): + """An articulation asset class. + + An articulation is a collection of rigid bodies connected by joints. The joints can be either + fixed or actuated. The joints can be of different types, such as revolute, prismatic, D-6, etc. + However, the articulation class has currently been tested with revolute and prismatic joints. + The class supports both floating-base and fixed-base articulations. The type of articulation + is determined based on the root joint of the articulation. If the root joint is fixed, then + the articulation is considered a fixed-base system. Otherwise, it is considered a floating-base + system. This can be checked using the :attr:`Articulation.is_fixed_base` attribute. + + For an asset to be considered an articulation, the root prim of the asset must have the + `USD ArticulationRootAPI`_. This API is used to define the sub-tree of the articulation using + the reduced coordinate formulation. On playing the simulation, the physics engine parses the + articulation root prim and creates the corresponding articulation in the physics engine. The + articulation root prim can be specified using the :attr:`AssetBaseCfg.prim_path` attribute. + + The articulation class also provides the functionality to augment the simulation of an articulated + system with custom actuator models. These models can either be explicit or implicit, as detailed in + the :mod:`isaaclab.actuators` module. The actuator models are specified using the + :attr:`ArticulationCfg.actuators` attribute. These are then parsed and used to initialize the + corresponding actuator models, when the simulation is played. + + During the simulation step, the articulation class first applies the actuator models to compute + the joint commands based on the user-specified targets. These joint commands are then applied + into the simulation. The joint commands can be either position, velocity, or effort commands. + As an example, the following snippet shows how this can be used for position commands: + + .. code-block:: python + + # an example instance of the articulation class + my_articulation = Articulation(cfg) + + # set joint position targets + my_articulation.set_joint_position_target(position) + # propagate the actuator models and apply the computed commands into the simulation + my_articulation.write_data_to_sim() + + # step the simulation using the simulation context + sim_context.step() + + # update the articulation state, where dt is the simulation time step + my_articulation.update(dt) + + .. _`USD ArticulationRootAPI`: https://openusd.org/dev/api/class_usd_physics_articulation_root_a_p_i.html + + """ + + cfg: ArticulationCfg + """Configuration instance for the articulations.""" + + __backend_name__: str = "newton" + """The name of the backend for the articulation.""" + + actuators: dict + """Dictionary of actuator instances for the articulation. + + 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:`ArticulationCfg.actuators` + attribute. They are used to compute the joint commands during the :meth:`write_data_to_sim` function. + """ + + def __init__(self, cfg: ArticulationCfg): + """Initialize the articulation. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + + """ + Properties + """ + + @property + def data(self) -> ArticulationData: + return self._data + + @property + def num_instances(self) -> int: + return self.root_view.count + + @property + def is_fixed_base(self) -> bool: + """Whether the articulation is a fixed-base or floating-base system.""" + return self.root_view.is_fixed_base + + @property + def num_joints(self) -> int: + """Number of joints in articulation.""" + return self.root_view.joint_dof_count + + @property + def num_fixed_tendons(self) -> int: + """Number of fixed tendons in articulation.""" + return 0 + + @property + def num_spatial_tendons(self) -> int: + """Number of spatial tendons in articulation.""" + return 0 + + @property + def num_bodies(self) -> int: + """Number of bodies in articulation.""" + return self.root_view.link_count + + @property + def num_shapes_per_body(self) -> list[int]: + """Number of collision shapes per body in the articulation. + + This property returns a list where each element represents the number of collision + shapes for the corresponding body in the articulation. This is cached for efficient + access during material property randomization and other operations. + + Returns: + List of integers representing the number of shapes per body. + """ + if not hasattr(self, "_num_shapes_per_body"): + self._num_shapes_per_body = [] + for shapes in self._root_view.body_shapes: + self._num_shapes_per_body.append(len(shapes)) + return self._num_shapes_per_body + + @property + def joint_names(self) -> list[str]: + """Ordered names of joints in articulation.""" + return self.root_view.joint_dof_names + + @property + def fixed_tendon_names(self) -> list[str]: + """Ordered names of fixed tendons in articulation.""" + return [] + + @property + def spatial_tendon_names(self) -> list[str]: + """Ordered names of spatial tendons in articulation.""" + return [] + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies in articulation.""" + return self.root_view.link_names + + @property + def root_view(self) -> ArticulationView: + """Root view for the asset. + + .. note:: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._root_view + + @property + def instantaneous_wrench_composer(self) -> WrenchComposer: + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ + return self._instantaneous_wrench_composer + + @property + def permanent_wrench_composer(self) -> WrenchComposer: + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ + return self._permanent_wrench_composer + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: + """Reset the articulation. + + .. caution:: + If both `env_ids` and `env_mask` are provided, then `env_mask` takes precedence over `env_ids`. + + Args: + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # use ellipses object to skip initial indices. + if (env_ids is None) or (env_ids == slice(None)): + env_ids = slice(None) + # reset actuators + for actuator in self.actuators.values(): + actuator.reset(env_ids) + # reset external wrenches. + self._instantaneous_wrench_composer.reset(env_ids, env_mask) + self._permanent_wrench_composer.reset(env_ids, env_mask) + + def write_data_to_sim(self): + """Write external wrenches and joint commands to the simulation. + + If any explicit actuators are present, then the actuator models are used to compute the + joint commands. Otherwise, the joint commands are directly set into the simulation. + + .. note:: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + # write external wrench + if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: + if self._instantaneous_wrench_composer.active: + # Compose instantaneous wrench with permanent wrench + self._instantaneous_wrench_composer.add_forces_and_torques_index( + forces=self._permanent_wrench_composer.composed_force, + torques=self._permanent_wrench_composer.composed_torque, + body_ids=self._ALL_BODY_INDICES, + env_ids=self._ALL_INDICES, + ) + # Apply both instantaneous and permanent wrench to the simulation + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + self._instantaneous_wrench_composer.composed_force, + self._instantaneous_wrench_composer.composed_torque, + self._data._sim_bind_body_external_wrench, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) + else: + # Apply permanent wrench to the simulation + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + self._permanent_wrench_composer.composed_force, + self._permanent_wrench_composer.composed_torque, + self._data._sim_bind_body_external_wrench, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) + self._instantaneous_wrench_composer.reset() + + # apply actuator models + self._apply_actuator_model() + # write actions into simulation via Newton bindings + wp.copy(self.data._sim_bind_joint_effort, self._joint_effort_target_sim) + # position and velocity targets only for implicit actuators + if self._has_implicit_actuators: + wp.copy(self.data._sim_bind_joint_position_target, self._joint_pos_target_sim) + wp.copy(self.data._sim_bind_joint_velocity_target, self._joint_vel_target_sim) + + def update(self, dt: float): + """Updates the simulation data. + + Args: + dt: The time step size in seconds. + """ + self.data.update(dt) + + """ + Operations - Finders. + """ + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find bodies in the articulation based on the name keys. + + Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body indices and names. + """ + return resolve_matching_names(name_keys, self.body_names, preserve_order) + + def find_joints( + self, name_keys: str | Sequence[str], joint_subset: list[str] | None = None, preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + """Find joints in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the joint names. + joint_subset: A subset of joints to search for. Defaults to None, which means all joints + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the joint indices and names. + """ + if joint_subset is None: + joint_subset = self.joint_names + # find joints + return resolve_matching_names(name_keys, joint_subset, preserve_order) + + def find_fixed_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + """Find fixed tendons in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the joint + names with fixed tendons. + tendon_subsets: A subset of joints with fixed tendons to search for. Defaults to None, which means + all joints in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the tendon indices and names. + """ + if tendon_subsets is None: + # tendons follow the joint names they are attached to + tendon_subsets = self.fixed_tendon_names + # find tendons + return resolve_matching_names(name_keys, tendon_subsets, preserve_order) + + def find_spatial_tendons( + self, name_keys: str | Sequence[str], tendon_subsets: list[str] | None = None, preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + """Find spatial tendons in the articulation based on the name keys. + + Please see the :func:`isaaclab.utils.string.resolve_matching_names` function for more information + on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the tendon names. + tendon_subsets: A subset of tendons to search for. Defaults to None, which means all tendons + in the articulation are searched. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the tendon indices and names. + """ + if tendon_subsets is None: + tendon_subsets = self.spatial_tendon_names + # find tendons + return resolve_matching_names(name_keys, tendon_subsets, preserve_order) + + """ + Operations - State Writers. + """ + + def write_root_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expect partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) + + def write_root_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root pose over selected environment mask into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + self.write_root_link_pose_to_sim_mask(root_pose, env_mask=env_mask) + + def write_root_link_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expect partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.set_root_link_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_pose, + env_ids, + ], + outputs=[ + self.data.root_link_pose_w, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + ], + device=self.device, + ) + # Need to invalidate the buffer to trigger the update with the new state. + # Only invalidate if the buffer has been accessed (not None). + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + self.data._body_link_pose_w_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. + if self.data._body_com_pose_w is not None: + self.data._body_com_pose_w.timestamp = -1.0 + if self.data._body_state_w is not None: + self.data._body_state_w.timestamp = -1.0 + if self.data._body_link_state_w is not None: + self.data._body_link_state_w.timestamp = -1.0 + if self.data._body_com_state_w is not None: + self.data._body_com_state_w.timestamp = -1.0 + + def write_root_link_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root link pose over selected environment mask into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + + wp.launch( + shared_kernels.set_root_link_pose_to_sim_mask, + dim=root_pose.shape[0], + inputs=[ + root_pose, + env_mask, + ], + outputs=[ + self.data.root_link_pose_w, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + ], + device=self.device, + ) + # Need to invalidate the buffer to trigger the update with the new state. + # Only invalidate if the buffer has been accessed (not None). + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + self.data._body_link_pose_w_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. + if self.data._body_com_pose_w is not None: + self.data._body_com_pose_w.timestamp = -1.0 + if self.data._body_state_w is not None: + self.data._body_state_w.timestamp = -1.0 + if self.data._body_link_state_w is not None: + self.data._body_link_state_w.timestamp = -1.0 + if self.data._body_com_state_w is not None: + self.data._body_com_state_w.timestamp = -1.0 + + def write_root_com_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expect partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + # Note: we are doing a single launch for faster performance. Prior versions would call + # write_root_link_pose_to_sim after this. + wp.launch( + shared_kernels.set_root_com_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_pose, + self.data.body_com_pose_b, + env_ids, + ], + outputs=[ + self.data._root_com_pose_w.data, + self.data.root_link_pose_w, + None, # self.data._root_com_state_w.data, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._root_com_pose_w.timestamp = self.data._sim_timestamp + # Need to invalidate the buffer to trigger the update with the new state. + # Only invalidate if the buffer has been accessed (not None). + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + self.data._body_link_pose_w_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. + if self.data._body_com_pose_w is not None: + self.data._body_com_pose_w.timestamp = -1.0 + if self.data._body_state_w is not None: + self.data._body_state_w.timestamp = -1.0 + if self.data._body_link_state_w is not None: + self.data._body_link_state_w.timestamp = -1.0 + if self.data._body_com_state_w is not None: + self.data._body_com_state_w.timestamp = -1.0 + + def write_root_com_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass pose over selected environment mask into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + wp.launch( + shared_kernels.set_root_com_pose_to_sim_mask, + dim=root_pose.shape[0], + inputs=[ + root_pose, + env_mask, + ], + outputs=[ + self.data._root_com_pose_w.data, + self.data.root_link_pose_w, + None, # self.data._root_com_state_w.data, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._root_com_pose_w.timestamp = self.data._sim_timestamp + # Need to invalidate the buffer to trigger the update with the new state. + # Only invalidate if the buffer has been accessed (not None). + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + self.data._body_link_pose_w_timestamp = -1.0 # Forces a kinematic update to get the latest body link poses. + if self.data._body_com_pose_w is not None: + self.data._body_com_pose_w.timestamp = -1.0 + if self.data._body_state_w is not None: + self.data._body_state_w.timestamp = -1.0 + if self.data._body_link_state_w is not None: + self.data._body_link_state_w.timestamp = -1.0 + if self.data._body_com_state_w is not None: + self.data._body_com_state_w.timestamp = -1.0 + + def write_root_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expect partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. + Shape is (len(env_ids), 6) or (len(env_ids),) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) + + def write_root_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. + Shape is (num_instances, 6) or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + self.write_root_com_velocity_to_sim_mask(root_velocity=root_velocity, env_mask=env_mask) + + def write_root_com_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expect partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. + Shape is (len(env_ids), 6) or (len(env_ids),) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_velocity, + env_ids, + self.data._num_bodies, + ], + outputs=[ + self.data.root_com_vel_w, + self.data._body_com_acc_w.data, + None, # self.data._root_state_w.data, + None, # self.data._root_com_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._body_com_acc_w.timestamp = self.data._sim_timestamp + # Only invalidate if the buffer has been accessed (not None). + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + + def write_root_com_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. + Shape is (num_instances, 6) or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_mask, + dim=root_velocity.shape[0], + inputs=[ + root_velocity, + env_mask, + self.data._num_bodies, + ], + outputs=[ + self.data.root_com_vel_w, + self.data._body_com_acc_w.data, + None, # self.data._root_state_w.data, + None, # self.data._root_com_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._body_com_acc_w.timestamp = self.data._sim_timestamp + # Only invalidate if the buffer has been accessed (not None). + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + + def write_root_link_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. + + .. note:: + This method expect partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root frame velocities in simulation world frame. + Shape is (len(env_ids), 6) or (len(env_ids),) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + # Note: we are doing a single launch for faster performance. Prior versions would do multiple launches. + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_ids, + self.data._num_bodies, + ], + outputs=[ + self.data._root_link_vel_w.data, + self.data.root_com_vel_w, + self.data._body_com_acc_w.data, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + None, # self.data._root_com_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._root_link_vel_w.timestamp = self.data._sim_timestamp + self.data._body_com_acc_w.timestamp = self.data._sim_timestamp + # Only invalidate if the buffer has been accessed (not None). + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + + def write_root_link_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root link velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root frame velocities in simulation world frame. + Shape is (num_instances, 6) or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_mask, + dim=root_velocity.shape[0], + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_mask, + self.data._num_bodies, + ], + outputs=[ + self.data._root_link_vel_w.data, + self.data.root_com_vel_w, + self.data._body_com_acc_w.data, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + None, # self.data._root_com_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._root_link_vel_w.timestamp = self.data._sim_timestamp + self.data._body_com_acc_w.timestamp = self.data._sim_timestamp + # Only invalidate if the buffer has been accessed (not None). + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + + def write_joint_state_to_sim( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + """Deprecated, same as :meth:`write_joint_position_to_sim_index` and + :meth:`write_joint_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_joint_state_to_sim' will be deprecated in a future release. Please" + " use 'write_joint_position_to_sim_index' and 'write_joint_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + # set into simulation + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) + + def write_joint_state_to_sim_mask( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint positions and velocities over selected environment mask into the simulation. + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + position: Joint positions. Shape is (num_instances, num_joints). + velocity: Joint velocities. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # set into simulation + self.write_joint_position_to_sim_mask(position, env_mask=env_mask, joint_mask=joint_mask) + self.write_joint_velocity_to_sim_mask(velocity, env_mask=env_mask, joint_mask=joint_mask) + + def write_joint_position_to_sim_index( + self, + *, + position: torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + """Write joint positions over selected environment indices into the simulation. + + .. note:: + This method expect partial data or full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + position: Joint positions. Shape is (len(env_ids), len(joint_ids)). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + position, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_pos, + ], + device=self.device, + ) + # Need to invalidate the buffer to trigger the update with the new root pose. + # Only invalidate if the buffer has been accessed (not None). + if self.data._body_link_vel_w is not None: + self.data._body_link_vel_w.timestamp = -1.0 + if self.data._body_com_pose_b is not None: + self.data._body_com_pose_b.timestamp = -1.0 + if self.data._body_com_pose_w is not None: + self.data._body_com_pose_w.timestamp = -1.0 + if self.data._body_state_w is not None: + self.data._body_state_w.timestamp = -1.0 + if self.data._body_link_state_w is not None: + self.data._body_link_state_w.timestamp = -1.0 + if self.data._body_com_state_w is not None: + self.data._body_com_state_w.timestamp = -1.0 + + def write_joint_position_to_sim_mask( + self, + *, + position: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint positions over selected environment mask into the simulation. + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + position: Joint positions. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + position, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_pos, + ], + device=self.device, + ) + # Need to invalidate the buffer to trigger the update with the new root pose. + # Only invalidate if the buffer has been accessed (not None). + if self.data._body_link_vel_w is not None: + self.data._body_link_vel_w.timestamp = -1.0 + if self.data._body_com_pose_b is not None: + self.data._body_com_pose_b.timestamp = -1.0 + if self.data._body_com_pose_w is not None: + self.data._body_com_pose_w.timestamp = -1.0 + if self.data._body_state_w is not None: + self.data._body_state_w.timestamp = -1.0 + if self.data._body_link_state_w is not None: + self.data._body_link_state_w.timestamp = -1.0 + if self.data._body_com_state_w is not None: + self.data._body_com_state_w.timestamp = -1.0 + + def write_joint_velocity_to_sim_index( + self, + *, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + """Write joint velocities to the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + velocity: Joint velocities. Shape is (len(env_ids), len(joint_ids)) or (num_instances, num_joints). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + articulation_kernels.write_joint_vel_data_index, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + velocity, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_vel, + self.data._previous_joint_vel, + self.data._joint_acc.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._joint_acc.timestamp = self.data._sim_timestamp + + def write_joint_velocity_to_sim_mask( + self, + *, + velocity: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint velocities over selected environment mask into the simulation. + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + velocity: Joint velocities. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + wp.launch( + articulation_kernels.write_joint_vel_data_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + velocity, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_vel, + self.data._previous_joint_vel, + self.data._joint_acc.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._joint_acc.timestamp = self.data._sim_timestamp + + """ + Operations - Simulation Parameters Writers. + """ + + def write_joint_stiffness_to_sim_index( + self, + *, + stiffness: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + """Write joint stiffness over selected environment indices into the simulation. + + .. note:: + This method expect partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + stiffness: Joint stiffness. Shape is (len(env_ids), len(joint_ids)). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + if isinstance(stiffness, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + stiffness, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_stiffness, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + stiffness, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_stiffness, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_stiffness_to_sim_mask( + self, + *, + stiffness: torch.Tensor | wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint stiffness over selected environment mask into the simulation. + + .. note:: + This method expect full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + stiffness: Joint stiffness. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + if isinstance(stiffness, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + stiffness, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_stiffness, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + stiffness, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_stiffness, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_damping_to_sim_index( + self, + *, + damping: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + """Write joint damping over selected environment indices into the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + damping: Joint damping. Shape is (len(env_ids), len(joint_ids)). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # Note This function isn't setting the values for actuator models. (#128) + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + if isinstance(damping, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + damping, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_damping, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + damping, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_damping, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_damping_to_sim_mask( + self, + *, + damping: torch.Tensor | wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint damping over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + damping: Joint damping. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + if isinstance(damping, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + damping, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_damping, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + damping, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_damping, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_position_limit_to_sim_index( + self, + *, + limits: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + warn_limit_violation: bool = True, + ): + """Write joint position limits over selected environment indices into the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limits: Joint limits. Shape is (len(env_ids), len(joint_ids), 2). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + warn_limit_violation: Whether to use warning or info level logging when default joint positions + exceed the new limits. Defaults to True. + """ + # Note This function isn't setting the values for actuator models. (#128) + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + + clamped_defaults = wp.zeros(1, dtype=wp.int32, device=self.device) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + # Note: we are doing a single launch for faster performance. Prior versions would do this in multiple launches. + if isinstance(limits, float): + raise ValueError("Joint position limits must be a tensor or array, not a float.") + wp.launch( + articulation_kernels.write_joint_limit_data_to_buffer_index, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + limits, + self.cfg.soft_joint_pos_limit_factor, + env_ids, + joint_ids, + ], + outputs=[ + self.data._sim_bind_joint_pos_limits_lower, + self.data._sim_bind_joint_pos_limits_upper, + self.data._soft_joint_pos_limits, + self.data._default_joint_pos, + clamped_defaults, + ], + device=self.device, + ) + # Log a warning if the default joint positions are outside of the new limits. + if clamped_defaults.numpy()[0] > 0: + violation_message = ( + "Some default joint positions are outside of the range of the new joint limits. Default joint positions" + " will be clamped to be within the new joint limits." + ) + if warn_limit_violation: + logger.warning(violation_message) + else: + logger.info(violation_message) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_position_limit_to_sim_mask( + self, + *, + limits: torch.Tensor | wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + warn_limit_violation: bool = True, + ): + """Write joint position limits over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limits: Joint limits. Shape is (num_instances, num_joints, 2). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + warn_limit_violation: Whether to use warning or info level logging when default joint positions + exceed the new limits. Defaults to True. + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + clamped_defaults = wp.zeros(1, dtype=wp.int32, device=self.device) + if isinstance(limits, float): + raise ValueError("Joint position limits must be a tensor or array, not a float.") + wp.launch( + articulation_kernels.write_joint_limit_data_to_buffer_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + limits, + self.cfg.soft_joint_pos_limit_factor, + env_mask, + joint_mask, + ], + outputs=[ + self.data._sim_bind_joint_pos_limits_lower, + self.data._sim_bind_joint_pos_limits_upper, + self.data._soft_joint_pos_limits, + self.data._default_joint_pos, + clamped_defaults, + ], + device=self.device, + ) + # Log a warning if the default joint positions are outside of the new limits. + if clamped_defaults.numpy()[0] > 0: + violation_message = ( + "Some default joint positions are outside of the range of the new joint limits. Default joint positions" + " will be clamped to be within the new joint limits." + ) + if warn_limit_violation: + logger.warning(violation_message) + else: + logger.info(violation_message) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_velocity_limit_to_sim_index( + self, + *, + limits: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + """Write joint max velocity over selected environment indices into the simulation. + + The velocity limit is used to constrain the joint velocities in the physics engine. The joint will only + be able to reach this velocity if the joint's effort limit is sufficiently large. If the joint is moving + faster than this velocity, the physics engine will actually try to brake the joint to reach this velocity. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limits: Joint max velocity. Shape is (len(env_ids), len(joint_ids)). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + if isinstance(limits, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + limits, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_vel_limits, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + limits, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_vel_limits, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_velocity_limit_to_sim_mask( + self, + *, + limits: torch.Tensor | wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write joint max velocity over selected environment mask into the simulation. + + The velocity limit is used to constrain the joint velocities in the physics engine. The joint will only + be able to reach this velocity if the joint's effort limit is sufficiently large. If the joint is moving + faster than this velocity, the physics engine will actually try to brake the joint to reach this velocity. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limits: Joint max velocity. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + if isinstance(limits, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + limits, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_vel_limits, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + limits, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_vel_limits, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_effort_limit_to_sim_index( + self, + *, + limits: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + """Write joint effort limits over selected environment indices into the simulation. + + The effort limit is used to constrain the computed joint efforts in the physics engine. If the + computed effort exceeds this limit, the physics engine will clip the effort to this value. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limits: Joint torque limits. Shape is (len(env_ids), len(joint_ids)). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # Note This function isn't setting the values for actuator models. (#128) + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + if isinstance(limits, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + limits, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_effort_limits, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + limits, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_effort_limits, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_effort_limit_to_sim_mask( + self, + *, + limits: torch.Tensor | wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint effort limits over selected environment mask into the simulation. + + The effort limit is used to constrain the computed joint efforts in the physics engine. If the + computed effort exceeds this limit, the physics engine will clip the effort to this value. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limits: Joint torque limits. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + if isinstance(limits, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + limits, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_effort_limits, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + limits, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_effort_limits, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_armature_to_sim_index( + self, + *, + armature: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + """Write joint armature over selected environment indices into the simulation. + + The armature is directly added to the corresponding joint-space inertia. It helps improve the + simulation stability by reducing the joint velocities. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + armature: Joint armature. Shape is (len(env_ids), len(joint_ids)). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + if isinstance(armature, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + armature, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_armature, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + armature, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_armature, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_armature_to_sim_mask( + self, + *, + armature: torch.Tensor | wp.array | float, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + """Write joint armature over selected environment mask into the simulation. + + The armature is directly added to the corresponding joint-space inertia. It helps improve the + simulation stability by reducing the joint velocities. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + armature: Joint armature. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + if isinstance(armature, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + armature, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_armature, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + armature, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_armature, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_friction_coefficient_to_sim_index( + self, + *, + joint_friction_coeff: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ): + r"""Write joint friction coefficients over selected environment indices into the simulation. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + joint_friction_coeff: Static friction coefficient :math:`\mu_s`. + Shape is (len(env_ids), len(joint_ids)). + joint_ids: Joint indices. If None, then all joints are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + if isinstance(joint_friction_coeff, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + joint_friction_coeff, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_friction_coeff, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + joint_friction_coeff, + env_ids, + joint_ids, + ], + outputs=[ + self.data.joint_friction_coeff, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + def write_joint_friction_coefficient_to_sim_mask( + self, + *, + joint_friction_coeff: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ): + r"""Write joint friction coefficients over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + joint_friction_coeff: Static friction coefficient :math:`\mu_s`. + Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + if isinstance(joint_friction_coeff, float): + wp.launch( + articulation_kernels.float_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + joint_friction_coeff, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_friction_coeff, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + joint_friction_coeff, + env_mask, + joint_mask, + ], + outputs=[ + self.data.joint_friction_coeff, + ], + device=self.device, + ) + # tell the physics engine that some of the joint properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.JOINT_DOF_PROPERTIES) + + """ + Operations - Setters. + """ + + def set_masses_index( + self, + *, + masses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set masses of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + masses: Masses of all bodies. Shape is (len(env_ids), len(body_ids)). + body_ids: Body indices. If None, then all bodies are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + masses, + env_ids, + body_ids, + ], + outputs=[ + self.data.body_mass, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_masses_mask( + self, + *, + masses: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set masses of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + masses: Masses of all bodies. Shape is (num_instances, num_bodies). + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + masses, + env_mask, + body_mask, + ], + outputs=[ + self.data.body_mass, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_coms_index( + self, + *, + coms: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set center of mass position of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + .. caution:: + Unlike the PhysX version of this method, this method does not set the center of mass orientation. + Only the position is set. This is because Newton considers the center of mass orientation to always be + aligned with the body frame. + + Args: + coms: Center of mass position of all bodies. Shape is (len(env_ids), len(body_ids), 3). + body_ids: Body indices. If None, then all bodies are used. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_body_com_position_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + coms, + env_ids, + body_ids, + ], + outputs=[ + self.data.body_com_pos_b, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_coms_mask( + self, + *, + coms: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set center of mass position of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + .. caution:: + Unlike the PhysX version of this method, this method does not set the center of mass orientation. + Only the position is set. This is because Newton considers the center of mass orientation to always be + aligned with the body frame. + + Args: + coms: Center of mass position of all bodies. Shape is (num_instances, num_bodies, 3). + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + wp.launch( + shared_kernels.write_body_com_position_to_buffer_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + coms, + env_mask, + body_mask, + ], + outputs=[ + self.data.body_com_pos_b, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_inertias_index( + self, + *, + inertias: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set inertias of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + inertias: Inertias of all bodies. Shape is (len(env_ids), len(body_ids), 9). + body_ids: The body indices to set the inertias for. Defaults to None (all bodies). + env_ids: The environment indices to set the inertias for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_body_inertia_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + inertias, + env_ids, + body_ids, + ], + outputs=[ + self.data.body_inertia, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_inertias_mask( + self, + *, + inertias: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set inertias of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). + body_mask: Body mask. If None, then all bodies are used. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + wp.launch( + shared_kernels.write_body_inertia_to_buffer_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + inertias, + env_mask, + body_mask, + ], + outputs=[ + self.data.body_inertia, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_joint_position_target_index( + self, + *, + target: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set joint position targets into internal buffers using indices. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + target: Joint position targets. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + target, + env_ids, + joint_ids, + ], + outputs=[ + self.data._joint_pos_target, + ], + device=self.device, + ) + # Only updates internal buffers, does not apply the targets to the simulation. + + def set_joint_position_target_mask( + self, + *, + target: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint position targets into internal buffers using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + target: Joint position targets. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + target, + env_mask, + joint_mask, + ], + outputs=[ + self.data._joint_pos_target, + ], + device=self.device, + ) + # Only updates internal buffers, does not apply the targets to the simulation. + + def set_joint_velocity_target_index( + self, + *, + target: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set joint velocity targets into internal buffers using indices. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + target: Joint velocity targets. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + target, + env_ids, + joint_ids, + ], + outputs=[ + self.data._joint_vel_target, + ], + device=self.device, + ) + # Only updates internal buffers, does not apply the targets to the simulation. + + def set_joint_velocity_target_mask( + self, + *, + target: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint velocity targets into internal buffers using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + target: Joint velocity targets. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # Resolve masks. + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + target, + env_mask, + joint_mask, + ], + outputs=[ + self.data._joint_vel_target, + ], + device=self.device, + ) + # Only updates internal buffers, does not apply the targets to the simulation. + + def set_joint_effort_target_index( + self, + *, + target: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set joint efforts into internal buffers using indices. + + This function does not apply the joint targets to the simulation. It only fills the buffers with + the desired values. To apply the joint targets, call the :meth:`write_data_to_sim` function. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + target: Joint effort targets. Shape is (len(env_ids), len(joint_ids)). + joint_ids: The joint indices to set the targets for. Defaults to None (all joints). + env_ids: The environment indices to set the targets for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + joint_ids = self._resolve_joint_ids(joint_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], joint_ids.shape[0]), + inputs=[ + target, + env_ids, + joint_ids, + ], + outputs=[ + self.data._joint_effort_target, + ], + device=self.device, + ) + # Only updates internal buffers, does not apply the targets to the simulation. + + def set_joint_effort_target_mask( + self, + *, + target: torch.Tensor | wp.array, + joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set joint efforts into internal buffers using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + target: Joint effort targets. Shape is (num_instances, num_joints). + joint_mask: Joint mask. If None, then all joints are used. Shape is (num_joints,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if joint_mask is None: + joint_mask = self._ALL_JOINT_MASK + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], joint_mask.shape[0]), + inputs=[ + target, + env_mask, + joint_mask, + ], + outputs=[ + self.data._joint_effort_target, + ], + device=self.device, + ) + # Only updates internal buffers, does not apply the targets to the simulation. + + """ + Operations - Tendons. + """ + + def set_fixed_tendon_stiffness_index( + self, + *, + stiffness: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set fixed tendon stiffness into internal buffers using indices. + + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the + :meth:`write_fixed_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + stiffness: Fixed tendon stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all fixed tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_fixed_tendon_stiffness_mask( + self, + *, + stiffness: torch.Tensor | wp.array, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon stiffness into internal buffers using masks. + + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the + :meth:`write_fixed_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + stiffness: Fixed tendon stiffness. Shape is (num_instances, num_fixed_tendons). + fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_fixed_tendon_damping_index( + self, + *, + damping: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set fixed tendon damping into internal buffers using indices. + + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the :meth:`write_fixed_tendon_properties_to_sim_index` + function. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + damping: Fixed tendon damping. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the damping for. Defaults to None (all fixed tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_fixed_tendon_damping_mask( + self, + *, + damping: torch.Tensor | wp.array, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon damping into internal buffers using masks. + + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the + :meth:`write_fixed_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + damping: Fixed tendon damping. Shape is (num_instances, num_fixed_tendons). + fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_fixed_tendon_limit_stiffness_index( + self, + *, + limit_stiffness: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set fixed tendon limit stiffness into internal buffers using indices. + + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the + :meth:`write_fixed_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limit_stiffness: Fixed tendon limit stiffness. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None (all fixed tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_fixed_tendon_limit_stiffness_mask( + self, + *, + limit_stiffness: torch.Tensor | wp.array, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon limit stiffness into internal buffers using masks. + + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the + :meth:`write_fixed_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limit_stiffness: Fixed tendon limit stiffness. Shape is (num_instances, num_fixed_tendons). + fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_fixed_tendon_position_limit_index( + self, + *, + limit: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set fixed tendon position limit into internal buffers using indices. + + This function does not apply the tendon position limit to the simulation. It only fills the buffers with + the desired values. To apply the tendon position limit, call the + :meth:`write_fixed_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limit: Fixed tendon position limit. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the position limit for. Defaults to None (all fixed tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_fixed_tendon_position_limit_mask( + self, + *, + limit: torch.Tensor | wp.array, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon position limit into internal buffers using masks. + + This function does not apply the tendon position limit to the simulation. It only fills the buffers with + the desired values. To apply the tendon position limit, call the + :meth:`write_fixed_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limit: Fixed tendon position limit. Shape is (num_instances, num_fixed_tendons). + fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_fixed_tendon_rest_length_index( + self, + *, + rest_length: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set fixed tendon rest length into internal buffers using indices. + + This function does not apply the tendon rest length to the simulation. It only fills the buffers with + the desired values. To apply the tendon rest length, call the + :meth:`write_fixed_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + rest_length: Fixed tendon rest length. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the rest length for. Defaults to None (all fixed tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_fixed_tendon_rest_length_mask( + self, + *, + rest_length: torch.Tensor | wp.array, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon rest length into internal buffers using masks. + + This function does not apply the tendon rest length to the simulation. It only fills the buffers with + the desired values. To apply the tendon rest length, call the + :meth:`write_fixed_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + rest_length: Fixed tendon rest length. Shape is (num_instances, num_fixed_tendons). + fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_fixed_tendon_offset_index( + self, + *, + offset: torch.Tensor | wp.array, + fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set fixed tendon offset into internal buffers using indices. + + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the + :meth:`write_fixed_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + offset: Fixed tendon offset. Shape is (len(env_ids), len(fixed_tendon_ids)). + fixed_tendon_ids: The tendon indices to set the offset for. Defaults to None (all fixed tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_fixed_tendon_offset_mask( + self, + *, + offset: torch.Tensor | wp.array, + fixed_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set fixed tendon offset into internal buffers using masks. + + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the + :meth:`write_fixed_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + offset: Fixed tendon offset. Shape is (num_instances, num_fixed_tendons). + fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. + Shape is (num_fixed_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def write_fixed_tendon_properties_to_sim_index( + self, + *, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Write fixed tendon properties into the simulation using indices. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + fixed_tendon_ids: The fixed tendon indices to write the properties for. Defaults to None + (all fixed tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def write_fixed_tendon_properties_to_sim_mask( + self, + *, + env_mask: wp.array | None = None, + ) -> None: + """Write fixed tendon properties into the simulation using masks. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_spatial_tendon_stiffness_index( + self, + *, + stiffness: torch.Tensor | wp.array, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set spatial tendon stiffness into internal buffers using indices. + + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the + :meth:`write_spatial_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + stiffness: Spatial tendon stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). + spatial_tendon_ids: The tendon indices to set the stiffness for. Defaults to None (all spatial tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_spatial_tendon_stiffness_mask( + self, + *, + stiffness: torch.Tensor | wp.array, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set spatial tendon stiffness into internal buffers using masks. + + This function does not apply the tendon stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon stiffness, call the + :meth:`write_spatial_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + stiffness: Spatial tendon stiffness. Shape is (num_instances, num_spatial_tendons). + spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_spatial_tendon_damping_index( + self, + *, + damping: torch.Tensor | wp.array, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set spatial tendon damping into internal buffers using indices. + + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the + :meth:`write_spatial_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + damping: Spatial tendon damping. Shape is (len(env_ids), len(spatial_tendon_ids)). + spatial_tendon_ids: The tendon indices to set the damping for. Defaults to None (all spatial tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_spatial_tendon_damping_mask( + self, + *, + damping: torch.Tensor | wp.array, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set spatial tendon damping into internal buffers using masks. + + This function does not apply the tendon damping to the simulation. It only fills the buffers with + the desired values. To apply the tendon damping, call the + :meth:`write_spatial_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + damping: Spatial tendon damping. Shape is (num_instances, num_spatial_tendons). + spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_spatial_tendon_limit_stiffness_index( + self, + *, + limit_stiffness: torch.Tensor | wp.array, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set spatial tendon limit stiffness into internal buffers using indices. + + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the + :meth:`write_spatial_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limit_stiffness: Spatial tendon limit stiffness. Shape is (len(env_ids), len(spatial_tendon_ids)). + spatial_tendon_ids: The tendon indices to set the limit stiffness for. Defaults to None + (all spatial tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_spatial_tendon_limit_stiffness_mask( + self, + *, + limit_stiffness: torch.Tensor | wp.array, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set spatial tendon limit stiffness into internal buffers using masks. + + This function does not apply the tendon limit stiffness to the simulation. It only fills the buffers with + the desired values. To apply the tendon limit stiffness, call the + :meth:`write_spatial_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + limit_stiffness: Spatial tendon limit stiffness. Shape is (num_instances, num_spatial_tendons). + spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def set_spatial_tendon_offset_index( + self, + *, + offset: torch.Tensor | wp.array, + spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set spatial tendon offset into internal buffers using indices. + + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the + :meth:`write_spatial_tendon_properties_to_sim_index` method. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + offset: Spatial tendon offset. Shape is (len(env_ids), len(spatial_tendon_ids)). + spatial_tendon_ids: The tendon indices to set the offset for. Defaults to None (all spatial tendons). + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def set_spatial_tendon_offset_mask( + self, + *, + offset: torch.Tensor | wp.array, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set spatial tendon offset into internal buffers using masks. + + This function does not apply the tendon offset to the simulation. It only fills the buffers with + the desired values. To apply the tendon offset, call the + :meth:`write_spatial_tendon_properties_to_sim_mask` method. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + offset: Spatial tendon offset. Shape is (num_instances, num_spatial_tendons). + spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + Shape is (num_spatial_tendons,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + def write_spatial_tendon_properties_to_sim_index( + self, + *, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Write spatial tendon properties into the simulation using indices. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + env_ids: Environment indices. If None, then all indices are used. + """ + raise NotImplementedError() + + def write_spatial_tendon_properties_to_sim_mask( + self, + *, + spatial_tendon_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Write spatial tendon properties into the simulation using masks. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + raise NotImplementedError() + + """ + Internal helper. + """ + + def _initialize_impl(self): + # obtain global simulation view + self._physics_sim_view = SimulationManager.get_physics_sim_view() + + if self.cfg.articulation_root_prim_path is not None: + # The articulation root prim path is specified explicitly, so we can just use this. + root_prim_path_expr = self.cfg.prim_path + self.cfg.articulation_root_prim_path + else: + # No articulation root prim path was specified, so we need to search + # for it. We search for this in the first environment and then + # create a regex that matches all environments. + first_env_matching_prim = find_first_matching_prim(self.cfg.prim_path) + if first_env_matching_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{self.cfg.prim_path}'.") + first_env_matching_prim_path = first_env_matching_prim.GetPath().pathString + + # Find all articulation root prims in the first environment. + first_env_root_prims = get_all_matching_child_prims( + first_env_matching_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), + traverse_instance_prims=False, + ) + if len(first_env_root_prims) == 0: + raise RuntimeError( + f"Failed to find an articulation when resolving '{first_env_matching_prim_path}'." + " Please ensure that the prim has 'USD ArticulationRootAPI' applied." + ) + if len(first_env_root_prims) > 1: + raise RuntimeError( + f"Failed to find a single articulation when resolving '{first_env_matching_prim_path}'." + f" Found multiple '{first_env_root_prims}' under '{first_env_matching_prim_path}'." + " Please ensure that there is only one articulation in the prim path tree." + ) + + # Now we convert the found articulation root from the first + # environment back into a regex that matches all environments. + first_env_root_prim_path = first_env_root_prims[0].GetPath().pathString + root_prim_path_relative_to_prim_path = first_env_root_prim_path[len(first_env_matching_prim_path) :] + root_prim_path_expr = self.cfg.prim_path + root_prim_path_relative_to_prim_path + + # -- articulation + self._root_view = ArticulationView( + SimulationManager.get_model(), + root_prim_path_expr.replace(".*", "*"), + verbose=False, + exclude_joint_types=[JointType.FREE, JointType.FIXED], + ) + + # log information about the articulation + logger.info(f"Articulation initialized at: {self.cfg.prim_path} with root '{root_prim_path_expr}'.") + logger.info(f"Is fixed root: {self.is_fixed_base}") + logger.info(f"Number of bodies: {self.num_bodies}") + logger.info(f"Body names: {self.body_names}") + logger.info(f"Number of joints: {self.num_joints}") + logger.info(f"Joint names: {self.joint_names}") + logger.info(f"Number of fixed tendons: {self.num_fixed_tendons}") + + # container for data access + self._data = ArticulationData(self.root_view, self.device) + + # create buffers + self._create_buffers() + # process configuration + self._process_cfg() + self._process_actuators_cfg() + self._process_tendons() + # validate configuration + self._validate_cfg() + # update the robot data + self.update(0.0) + # log joint information + self._log_articulation_info() + # Let the articulation data know that it is fully instantiated and ready to use. + self.data.is_primed = True + + def _create_buffers(self): + self._ALL_INDICES = wp.array(np.arange(self.num_instances, dtype=np.int32), device=self.device) + self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device) + self._ALL_JOINT_INDICES = wp.array(np.arange(self.num_joints, dtype=np.int32), device=self.device) + self._ALL_JOINT_MASK = wp.ones((self.num_joints,), dtype=wp.bool, device=self.device) + self._ALL_BODY_INDICES = wp.array(np.arange(self.num_bodies, dtype=np.int32), device=self.device) + self._ALL_BODY_MASK = wp.ones((self.num_bodies,), dtype=wp.bool, device=self.device) + self._ALL_FIXED_TENDON_INDICES = wp.array(np.arange(self.num_fixed_tendons, dtype=np.int32), device=self.device) + self._ALL_FIXED_TENDON_MASK = wp.ones((self.num_fixed_tendons,), dtype=wp.bool, device=self.device) + self._ALL_SPATIAL_TENDON_INDICES = wp.array( + np.arange(self.num_spatial_tendons, dtype=np.int32), device=self.device + ) + self._ALL_SPATIAL_TENDON_MASK = wp.ones((self.num_spatial_tendons,), dtype=wp.bool, device=self.device) + + # external wrench composer + self._instantaneous_wrench_composer = WrenchComposer(self) + self._permanent_wrench_composer = WrenchComposer(self) + + # asset named data + self.data.joint_names = self.joint_names + self.data.body_names = self.body_names + # tendon names are set in _process_tendons function + + # -- joint commands (sent to the simulation after actuator processing) + self._joint_pos_target_sim = wp.zeros_like(self.data.joint_pos_target, device=self.device) + self._joint_vel_target_sim = wp.zeros_like(self.data.joint_pos_target, device=self.device) + self._joint_effort_target_sim = wp.zeros_like(self.data.joint_pos_target, device=self.device) + + # soft joint position limits (recommended not to be too close to limits). + wp.launch( + articulation_kernels.update_soft_joint_pos_limits, + dim=(self.num_instances, self.num_joints), + inputs=[ + self.data.joint_pos_limits, + self.cfg.soft_joint_pos_limit_factor, + ], + outputs=[ + self.data.soft_joint_pos_limits, + ], + device=self.device, + ) + + def _process_cfg(self): + """Post processing of configuration parameters.""" + # default state + # -- root state + # Note we cast to tuple to avoid torch/numpy type mismatch. + default_root_pose = tuple(self.cfg.init_state.pos) + tuple(self.cfg.init_state.rot) + default_root_vel = tuple(self.cfg.init_state.lin_vel) + tuple(self.cfg.init_state.ang_vel) + default_root_pose = np.tile(np.array(default_root_pose, dtype=np.float32), (self.num_instances, 1)) + default_root_vel = np.tile(np.array(default_root_vel, dtype=np.float32), (self.num_instances, 1)) + self.data.default_root_pose = wp.array(default_root_pose, dtype=wp.transformf, device=self.device) + self.data.default_root_vel = wp.array(default_root_vel, dtype=wp.spatial_vectorf, device=self.device) + + # -- joint state + pos_idx_list, _, pos_val_list = resolve_matching_names_values(self.cfg.init_state.joint_pos, self.joint_names) + vel_idx_list, _, vel_val_list = resolve_matching_names_values(self.cfg.init_state.joint_vel, self.joint_names) + wp.launch( + articulation_kernels.update_default_joint_values, + dim=(self.num_instances, len(pos_idx_list)), + inputs=[ + wp.array(pos_val_list, dtype=wp.float32, device=self.device), + wp.array(pos_idx_list, dtype=wp.int32, device=self.device), + ], + outputs=[ + self.data.default_joint_pos, + ], + device=self.device, + ) + wp.launch( + articulation_kernels.update_default_joint_values, + dim=(self.num_instances, len(vel_idx_list)), + inputs=[ + wp.array(vel_val_list, dtype=wp.float32, device=self.device), + wp.array(vel_idx_list, dtype=wp.int32, device=self.device), + ], + outputs=[ + self.data.default_joint_vel, + ], + device=self.device, + ) + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + self._root_view = None + + """ + Internal helpers -- Actuators. + """ + + def _process_actuators_cfg(self): + """Process and apply articulation joint properties.""" + # create actuators + self.actuators = dict() + # flag for implicit actuators + # if this is false, we by-pass certain checks when doing actuator-related operations + self._has_implicit_actuators = False + + # iterate over all actuator configurations + for actuator_name, actuator_cfg in self.cfg.actuators.items(): + # type annotation for type checkers + actuator_cfg: ActuatorBaseCfg + # create actuator group + joint_ids, joint_names = self.find_joints(actuator_cfg.joint_names_expr) + # check if any joints are found + if len(joint_names) == 0: + raise ValueError( + f"No joints found for actuator group: {actuator_name} with joint name expression:" + f" {actuator_cfg.joint_names_expr}." + ) + # resolve joint indices + # we pass a slice if all joints are selected to avoid indexing overhead + if len(joint_names) == self.num_joints: + joint_ids = slice(None) + else: + joint_ids = torch.tensor(joint_ids, device=self.device, dtype=torch.int32) + # create actuator collection + # note: for efficiency avoid indexing when over all indices + actuator: ActuatorBase = actuator_cfg.class_type( + cfg=actuator_cfg, + joint_names=joint_names, + joint_ids=joint_ids, + num_envs=self.num_instances, + device=self.device, + stiffness=wp.to_torch(self._data.joint_stiffness)[:, joint_ids], + damping=wp.to_torch(self._data.joint_damping)[:, joint_ids], + armature=wp.to_torch(self._data.joint_armature)[:, joint_ids], + friction=wp.to_torch(self._data.joint_friction_coeff)[:, joint_ids], + effort_limit=wp.to_torch(self._data.joint_effort_limits)[:, joint_ids].clone(), + velocity_limit=wp.to_torch(self._data.joint_vel_limits)[:, joint_ids], + ) + # log information on actuator groups + model_type = "implicit" if actuator.is_implicit_model else "explicit" + logger.info( + f"Actuator collection: {actuator_name} with model '{actuator_cfg.class_type.__name__}'" + f" (type: {model_type}) and joint names: {joint_names} [{joint_ids}]." + ) + # store actuator group + self.actuators[actuator_name] = actuator + # set the passed gains and limits into the simulation + if isinstance(actuator, ImplicitActuator): + self._has_implicit_actuators = True + # the gains and limits are set into the simulation since actuator model is implicit + self.write_joint_stiffness_to_sim_index(stiffness=actuator.stiffness, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=actuator.damping, joint_ids=actuator.joint_indices) + else: + # the gains and limits are processed by the actuator model + # we set gains to zero, and torque limit to a high value in simulation to avoid any interference + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) + + # Set common properties into the simulation + self.write_joint_effort_limit_to_sim_index( + limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices + ) + self.write_joint_velocity_limit_to_sim_index( + limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices + ) + self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices + ) + + # Store the configured values from the actuator model + # note: this is the value configured in the actuator model (for implicit and explicit actuators) + joint_ids = actuator.joint_indices + if joint_ids == slice(None): + joint_ids = self._ALL_JOINT_INDICES + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, joint_ids.shape[0]), + inputs=[ + actuator.stiffness, + self._ALL_INDICES, + joint_ids, + ], + outputs=[ + self.data._sim_bind_joint_stiffness_sim, + ], + device=self.device, + ) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, joint_ids.shape[0]), + inputs=[ + actuator.damping, + self._ALL_INDICES, + joint_ids, + ], + outputs=[ + self.data._sim_bind_joint_damping_sim, + ], + device=self.device, + ) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, joint_ids.shape[0]), + inputs=[ + actuator.armature, + self._ALL_INDICES, + joint_ids, + ], + outputs=[ + self.data._sim_bind_joint_armature, + ], + device=self.device, + ) + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(self.num_instances, joint_ids.shape[0]), + inputs=[ + actuator.friction, + self._ALL_INDICES, + joint_ids, + ], + outputs=[ + self.data._sim_bind_joint_friction_coeff, + ], + device=self.device, + ) + + # perform some sanity checks to ensure actuators are prepared correctly + total_act_joints = sum(actuator.num_joints for actuator in self.actuators.values()) + if total_act_joints != (self.num_joints - self.num_fixed_tendons): + logger.warning( + "Not all actuators are configured! Total number of actuated joints not equal to number of" + f" joints available: {total_act_joints} != {self.num_joints - self.num_fixed_tendons}." + ) + + if self.cfg.actuator_value_resolution_debug_print: + t = PrettyTable(["Group", "Property", "Name", "ID", "USD Value", "ActutatorCfg Value", "Applied"]) + for actuator_group, actuator in self.actuators.items(): + group_count = 0 + for property, resolution_details in actuator.joint_property_resolution_table.items(): + for prop_idx, resolution_detail in enumerate(resolution_details): + actuator_group_str = actuator_group if group_count == 0 else "" + property_str = property if prop_idx == 0 else "" + fmt = [f"{v:.2e}" if isinstance(v, float) else str(v) for v in resolution_detail] + t.add_row([actuator_group_str, property_str, *fmt]) + group_count += 1 + logger.warning(f"\nActuatorCfg-USD Value Discrepancy Resolution (matching values are skipped): \n{t}") + + def _process_tendons(self): + """Process fixed and spatial tendons.""" + # create a list to store the fixed tendon names + self._fixed_tendon_names = list() + self._spatial_tendon_names = list() + # parse fixed tendons properties if they exist + if self.num_fixed_tendons > 0 or self.num_spatial_tendons > 0: + raise NotImplementedError("Fixed and spatial tendons are not supported yet.") + + def _apply_actuator_model(self): + """Processes joint commands for the articulation by forwarding them to the actuators. + + The actions are first processed using actuator models. Depending on the robot configuration, + the actuator models compute the joint level simulation commands and sets them into the PhysX buffers. + """ + # process actions per group + for actuator in self.actuators.values(): + # prepare input for actuator model based on cached data + # TODO : A tensor dict would be nice to do the indexing of all tensors together + control_action = ArticulationActions( + joint_positions=wp.to_torch(self._data.joint_pos_target)[:, actuator.joint_indices], + joint_velocities=wp.to_torch(self._data.joint_vel_target)[:, actuator.joint_indices], + joint_efforts=wp.to_torch(self._data.joint_effort_target)[:, actuator.joint_indices], + joint_indices=actuator.joint_indices, + ) + # compute joint command from the actuator model + control_action = actuator.compute( + control_action, + joint_pos=wp.to_torch(self._data.joint_pos)[:, actuator.joint_indices], + joint_vel=wp.to_torch(self._data.joint_vel)[:, actuator.joint_indices], + ) + # update targets (these are set into the simulation) + joint_indices = actuator.joint_indices + if actuator.joint_indices == slice(None) or actuator.joint_indices is None: + joint_indices = self._ALL_JOINT_INDICES + if hasattr(actuator, "gear_ratio"): + gear_ratio = actuator.gear_ratio + else: + gear_ratio = None + wp.launch( + articulation_kernels.update_targets, + dim=(self.num_instances, joint_indices.shape[0]), + inputs=[ + control_action.joint_positions, + control_action.joint_velocities, + control_action.joint_efforts, + joint_indices, + ], + outputs=[ + self._joint_pos_target_sim, + self._joint_vel_target_sim, + self._joint_effort_target_sim, + ], + device=self.device, + ) + # update state of the actuator model + wp.launch( + articulation_kernels.update_actuator_state_model, + dim=(self.num_instances, joint_indices.shape[0]), + inputs=[ + actuator.computed_effort, + actuator.applied_effort, + gear_ratio, + actuator.velocity_limit, + joint_indices, + ], + outputs=[ + self._data.computed_torque, + self._data.applied_torque, + self._data.gear_ratio, + self._data.soft_joint_vel_limits, + ], + device=self.device, + ) + + """ + Internal helpers -- Debugging. + """ + + def _validate_cfg(self): + """Validate the 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. + For instance, the actuator models may change the joint max velocity limits. + """ + # Skip validation if there are no joints (e.g., fixed-base articulation with 0 DOF) + if self.num_joints == 0: + return + + # check that the default values are within the limits + joint_pos_limits_lower = wp.to_torch(self._data.joint_pos_limits_lower)[0] + joint_pos_limits_upper = wp.to_torch(self._data.joint_pos_limits_upper)[0] + default_joint_pos = wp.to_torch(self._data.default_joint_pos)[0] + out_of_range = default_joint_pos < joint_pos_limits_lower + out_of_range |= default_joint_pos > joint_pos_limits_upper + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + # throw error if any of the default joint positions are out of the limits + if len(violated_indices) > 0: + # prepare message for violated joints + msg = "The following joints have default positions out of the limits: \n" + for idx in violated_indices: + joint_name = self.data.joint_names[idx] + joint_limit = [joint_pos_limits_lower[idx], joint_pos_limits_upper[idx]] + joint_pos = default_joint_pos[idx] + # add to message + msg += f"\t- '{joint_name}': {joint_pos:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" + raise ValueError(msg) + + # check that the default joint velocities are within the limits + joint_max_vel = wp.to_torch(self._data.joint_vel_limits)[0] + default_joint_vel = wp.to_torch(self._data.default_joint_vel)[0] + out_of_range = torch.abs(default_joint_vel) > joint_max_vel + violated_indices = torch.nonzero(out_of_range, as_tuple=False).squeeze(-1) + if len(violated_indices) > 0: + # prepare message for violated joints + msg = "The following joints have default velocities out of the limits: \n" + for idx in violated_indices: + joint_name = self.data.joint_names[idx] + joint_limit = [-joint_max_vel[idx], joint_max_vel[idx]] + joint_vel = default_joint_vel[idx] + # add to message + msg += f"\t- '{joint_name}': {joint_vel:.3f} not in [{joint_limit[0]:.3f}, {joint_limit[1]:.3f}]\n" + raise ValueError(msg) + + def _log_articulation_info(self): + """Log information about the articulation. + + .. note:: We purposefully read the values from the simulator to ensure that the values are configured as + expected. + """ + + # define custom formatters for large numbers and limit ranges + def format_large_number(_, v: float) -> str: + """Format large numbers using scientific notation.""" + if abs(v) >= 1e3: + return f"{v:.1e}" + else: + return f"{v:.3f}" + + def format_limits(_, v: tuple[float, float]) -> str: + """Format limit ranges using scientific notation.""" + if abs(v[0]) >= 1e3 or abs(v[1]) >= 1e3: + return f"[{v[0]:.1e}, {v[1]:.1e}]" + else: + return f"[{v[0]:.3f}, {v[1]:.3f}]" + + # read out all joint parameters from simulation + # -- gains + # Use data properties which have already been cloned and stored during initialization + # This avoids issues with indexedarray or empty arrays from root_view + stiffnesses = wp.to_torch(self.data.joint_stiffness)[0].cpu().tolist() + dampings = wp.to_torch(self.data.joint_damping)[0].cpu().tolist() + # -- properties + armatures = wp.to_torch(self.data.joint_armature)[0].cpu().tolist() + # For friction, use the individual components from data + friction_coeff = wp.to_torch(self.data.joint_friction_coeff)[0].cpu() + static_frictions = friction_coeff.tolist() + # -- limits + # joint_pos_limits is vec2f array, convert to torch and extract [lower, upper] pairs + position_limits_torch = wp.to_torch(self.data.joint_pos_limits)[0].cpu() # shape: (num_joints, 2) + position_limits = [tuple(pos_limit.tolist()) for pos_limit in position_limits_torch] + velocity_limits = wp.to_torch(self.data.joint_vel_limits)[0].cpu().tolist() + effort_limits = wp.to_torch(self.data.joint_effort_limits)[0].cpu().tolist() + # create table for term information + joint_table = PrettyTable() + joint_table.title = f"Simulation Joint Information (Prim path: {self.cfg.prim_path})" + # build field names based on Isaac Sim version + field_names = ["Index", "Name", "Stiffness", "Damping", "Armature"] + field_names.extend(["Static Friction"]) + field_names.extend(["Position Limits", "Velocity Limits", "Effort Limits"]) + joint_table.field_names = field_names + + # apply custom formatters to numeric columns + joint_table.custom_format["Stiffness"] = format_large_number + joint_table.custom_format["Damping"] = format_large_number + joint_table.custom_format["Armature"] = format_large_number + joint_table.custom_format["Static Friction"] = format_large_number + joint_table.custom_format["Position Limits"] = format_limits + joint_table.custom_format["Velocity Limits"] = format_large_number + joint_table.custom_format["Effort Limits"] = format_large_number + + # set alignment of table columns + joint_table.align["Name"] = "l" + # add info on each term + for index, name in enumerate(self.joint_names): + # build row data based on Isaac Sim version + row_data = [index, name, stiffnesses[index], dampings[index], armatures[index]] + if has_kit() and get_isaac_sim_version().major < 5: + row_data.append(static_frictions[index]) + else: + row_data.extend([static_frictions[index]]) + row_data.extend([position_limits[index], velocity_limits[index], effort_limits[index]]) + # add row to table + joint_table.add_row(row_data) + # convert table to string + logger.info(f"Simulation parameters for joints in {self.cfg.prim_path}:\n" + joint_table.get_string()) + + # read out all fixed tendon parameters from simulation + if self.num_fixed_tendons > 0: + raise NotImplementedError("Fixed tendons are not supported yet.") + + if self.num_spatial_tendons > 0: + raise NotImplementedError("Spatial tendons are not supported yet.") + + def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array: + """Resolve environment indices to a warp array. + + .. note:: + We need to convert torch tensors to warp arrays since the TensorAPI views only support warp arrays. + + Args: + env_ids: Environment indices. If None, then all indices are used. + + Returns: + A warp array of environment indices. + """ + if (env_ids is None) or (env_ids == slice(None)): + return self._ALL_INDICES + if isinstance(env_ids, torch.Tensor): + # Convert int64 to int32 if needed, as warp expects int32 + if env_ids.dtype == torch.int64: + env_ids = env_ids.to(torch.int32) + return wp.from_torch(env_ids, dtype=wp.int32) + if isinstance(env_ids, list): + return wp.array(env_ids, dtype=wp.int32, device=self.device) + return env_ids + + def _resolve_joint_ids(self, joint_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + """Resolve joint indices to a warp array or tensor. + + .. note:: + We do not need to convert torch tensors to warp arrays since they never get passed to the TensorAPI views. + + Args: + joint_ids: Joint indices. If None, then all indices are used. + + Returns: + A warp array of joint indices or a tensor of joint indices. + """ + if isinstance(joint_ids, list): + return wp.array(joint_ids, dtype=wp.int32, device=self.device) + if (joint_ids is None) or (joint_ids == slice(None)): + return self._ALL_JOINT_INDICES + return joint_ids + + def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + """Resolve body indices to a warp array or tensor. + + Args: + body_ids: Body indices. If None, then all indices are used. + + Returns: + A warp array of body indices or a tensor of body indices. + """ + if isinstance(body_ids, list): + return wp.array(body_ids, dtype=wp.int32, device=self.device) + if (body_ids is None) or (body_ids == slice(None)): + return self._ALL_BODY_INDICES + return body_ids + + def _resolve_fixed_tendon_ids( + self, tendon_ids: Sequence[int] | torch.Tensor | wp.array | None + ) -> wp.array | torch.Tensor: + """Resolve tendon indices to a warp array or tensor. + + Args: + tendon_ids: Tendon indices. If None, then all indices are used. + + Returns: + A warp array of tendon indices or a tensor of tendon indices. + """ + if isinstance(tendon_ids, list): + return wp.array(tendon_ids, dtype=wp.int32, device=self.device) + if (tendon_ids is None) or (tendon_ids == slice(None)): + return self._ALL_FIXED_TENDON_INDICES + return tendon_ids + + def _resolve_spatial_tendon_ids( + self, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None + ) -> wp.array | torch.Tensor: + """Resolve spatial tendon indices to a warp array or tensor. + + Args: + spatial_tendon_ids: Spatial tendon indices. If None, then all indices are used. + + Returns: + A warp array of spatial tendon indices or a tensor of spatial tendon indices. + """ + if isinstance(spatial_tendon_ids, list): + return wp.array(spatial_tendon_ids, dtype=wp.int32, device=self.device) + if (spatial_tendon_ids is None) or (spatial_tendon_ids == slice(None)): + return self._ALL_SPATIAL_TENDON_INDICES + return spatial_tendon_ids + + """ + Deprecated methods. + """ + + def write_joint_friction_coefficient_to_sim( + self, + joint_friction_coeff: torch.Tensor | wp.array | float, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ): + """Deprecated, same as :meth:`write_joint_friction_coefficient_to_sim_index`.""" + warnings.warn( + "The function 'write_joint_friction_coefficient_to_sim' will be deprecated in a future release. Please" + " use 'write_joint_friction_coefficient_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + full_data=full_data, + ) + + def write_root_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + + def write_root_com_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_root_com_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_com_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_com_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_com_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + + def write_root_link_state_to_sim( + self, + root_state: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Deprecated, same as :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_link_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_link_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_link_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py new file mode 100644 index 00000000000..38b6b3396e8 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py @@ -0,0 +1,1863 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + +import logging +import warnings +import weakref +from typing import TYPE_CHECKING + +import torch +import warp as wp + +from isaaclab.assets.articulation.base_articulation_data import BaseArticulationData +from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.math import normalize + +from isaaclab_newton.assets import kernels as shared_kernels +from isaaclab_newton.assets.articulation import kernels as articulation_kernels +from isaaclab_newton.physics import NewtonManager as SimulationManager + +if TYPE_CHECKING: + from newton.selection import ArticulationView + +# import logger +logger = logging.getLogger(__name__) + + +class ArticulationData(BaseArticulationData): + """Data container for an articulation. + + This class contains the data for an articulation in the simulation. The data includes the state of + the root rigid body, the state of all the bodies in the articulation, and the joint state. The data is + stored in the simulation world frame unless otherwise specified. + + An articulation is comprised of multiple rigid bodies or links. For a rigid body, there are two frames + of reference that are used: + + - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim + with the rigid body schema. + - Center of mass frame: The frame of reference of the center of mass of the rigid body. + + Depending on the settings, the two frames may not coincide with each other. In the robotics sense, the actor frame + can be interpreted as the link frame. + """ + + __backend_name__: str = "newton" + """The name of the backend for the articulation data.""" + + def __init__(self, root_view: ArticulationView, device: str): + """Initializes the articulation data. + + Args: + root_view: The root articulation view. + device: The device used for processing. + """ + super().__init__(root_view, device) + # Set the root articulation view + # note: this is stored as a weak reference to avoid circular references between the asset class + # and the data container. This is important to avoid memory leaks. + self._root_view: ArticulationView = weakref.proxy(root_view) + + # Set initial time stamp + self._sim_timestamp = 0.0 + self._is_primed = False + + # Convert to direction vector + gravity = wp.to_torch(SimulationManager.get_model().gravity)[0] + gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) + gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) + gravity_dir = gravity_dir.repeat(self._root_view.count, 1) + forward_vec = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self._root_view.count, 1) + + # Initialize constants + self.GRAVITY_VEC_W = wp.from_torch(gravity_dir, dtype=wp.vec3f) + self.FORWARD_VEC_B = wp.from_torch(forward_vec, dtype=wp.vec3f) + + self._create_simulation_bindings() + self._create_buffers() + + @property + def is_primed(self) -> bool: + """Whether the articulation data is fully instantiated and ready to use.""" + return self._is_primed + + @is_primed.setter + def is_primed(self, value: bool) -> None: + """Set whether the articulation data is fully instantiated and ready to use. + + .. note:: + Once this quantity is set to True, it cannot be changed. + + Args: + value: The primed state. + + Raises: + ValueError: If the articulation data is already primed. + """ + if self._is_primed: + raise ValueError("The articulation data is already primed.") + self._is_primed = True + + def update(self, dt: float) -> None: + """Updates the data for the articulation. + + Args: + dt: The time step for the update. This must be a positive value. + """ + # update the simulation timestamp + self._sim_timestamp += dt + # Trigger an update of the joint and body com acceleration buffers at a higher frequency + # since we do finite differencing. + self.joint_acc + self.body_com_acc_w + + """ + Names. + """ + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + joint_names: list[str] = None + """Joint names in the order parsed by the simulation view.""" + + fixed_tendon_names: list[str] = None + """Fixed tendon names in the order parsed by the simulation view.""" + + spatial_tendon_names: list[str] = None + """Spatial tendon names in the order parsed by the simulation view.""" + + """ + Defaults - Initial state. + """ + + @property + def default_root_pose(self) -> wp.array: + """Default root pose ``[pos, quat]`` in the local environment frame. + + The position and quaternion are of the articulation root's actor frame. Shape is (num_instances), + dtype = wp.transformf. In torch this resolves to (num_instances, 7). + """ + return self._default_root_pose + + @default_root_pose.setter + def default_root_pose(self, value: wp.array) -> None: + """Set the default root pose. + + Args: + value: The default root pose. Shape is (num_instances, 7). + + Raises: + ValueError: If the articulation data is already primed. + """ + if self.is_primed: + raise ValueError("The articulation data is already primed.") + self._default_root_pose.assign(value) + + @property + def default_root_vel(self) -> wp.array: + """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. + + The linear and angular velocities are of the articulation root's center of mass frame. + Shape is (num_instances), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + """ + return self._default_root_vel + + @default_root_vel.setter + def default_root_vel(self, value: wp.array) -> None: + """Set the default root velocity. + + Args: + value: The default root velocity. Shape is (num_instances, 6). + + Raises: + ValueError: If the articulation data is already primed. + """ + if self.is_primed: + raise ValueError("The articulation data is already primed.") + self._default_root_vel.assign(value) + + @property + def default_joint_pos(self) -> wp.array: + """Default joint positions of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self._default_joint_pos + + @default_joint_pos.setter + def default_joint_pos(self, value: wp.array) -> None: + """Set the default joint positions. + + Args: + value: The default joint positions. Shape is (num_instances, num_joints). + + Raises: + ValueError: If the articulation data is already primed. + """ + if self.is_primed: + raise ValueError("The articulation data is already primed.") + self._default_joint_pos.assign(value) + + @property + def default_joint_vel(self) -> wp.array: + """Default joint velocities of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + return self._default_joint_vel + + @default_joint_vel.setter + def default_joint_vel(self, value: wp.array) -> None: + """Set the default joint velocities. + + Args: + value: The default joint velocities. Shape is (num_instances, num_joints). + + Raises: + ValueError: If the articulation data is already primed. + """ + if self.is_primed: + raise ValueError("The articulation data is already primed.") + self._default_joint_vel.assign(value) + + """ + Joint commands -- Set into simulation. + """ + + @property + def joint_pos_target(self) -> wp.array: + """Joint position targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + return self._joint_pos_target + + @property + def joint_vel_target(self) -> wp.array: + """Joint velocity targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + return self._joint_vel_target + + @property + def joint_effort_target(self) -> wp.array: + """Joint effort targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + For an implicit actuator model, the targets are directly set into the simulation. + For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), + which are then set into the simulation. + """ + return self._joint_effort_target + + """ + Joint commands -- Explicit actuators. + """ + + @property + def computed_torque(self) -> wp.array: + """Joint torques computed from the actuator model (before clipping). + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + This quantity is the raw torque output from the actuator mode, before any clipping is applied. + It is exposed for users who want to inspect the computations inside the actuator model. + For instance, to penalize the learning agent for a difference between the computed and applied torques. + """ + return self._computed_torque + + @property + def applied_torque(self) -> wp.array: + """Joint torques applied from the actuator model (after clipping). + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the + actuator model. + """ + return self._applied_torque + + """ + Joint properties + """ + + @property + def joint_stiffness(self) -> wp.array: + """Joint stiffness provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + In the case of explicit actuators, the value for the corresponding joints is zero. + """ + return self._sim_bind_joint_stiffness_sim + + @property + def joint_damping(self) -> wp.array: + """Joint damping provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + In the case of explicit actuators, the value for the corresponding joints is zero. + """ + return self._sim_bind_joint_damping_sim + + @property + def joint_armature(self) -> wp.array: + """Joint armature provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ + return self._sim_bind_joint_armature + + @property + def joint_friction_coeff(self) -> wp.array: + """Joint static friction coefficient provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ + return self._sim_bind_joint_friction_coeff + + @property + def joint_pos_limits_lower(self) -> wp.array: + """Joint position limits lower provided to the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_pos_limits_lower + + @property + def joint_pos_limits_upper(self) -> wp.array: + """Joint position limits upper provided to the simulation. Shape is (num_instances, num_joints).""" + return self._sim_bind_joint_pos_limits_upper + + @property + def joint_pos_limits(self) -> wp.array: + """Joint position limits provided to the simulation. + + Shape is (num_instances, num_joints, 2), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`. + """ + if self._joint_pos_limits is None: + self._joint_pos_limits = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.vec2f, device=self.device + ) + wp.launch( + articulation_kernels.concat_joint_pos_limits_lower_and_upper, + dim=(self._num_instances, self._num_joints), + inputs=[ + self._sim_bind_joint_pos_limits_lower, + self._sim_bind_joint_pos_limits_upper, + ], + outputs=[ + self._joint_pos_limits, + ], + device=self.device, + ) + return self._joint_pos_limits + + @property + def joint_vel_limits(self) -> wp.array: + """Joint maximum velocity provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ + return self._sim_bind_joint_vel_limits_sim + + @property + def joint_effort_limits(self) -> wp.array: + """Joint maximum effort provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ + return self._sim_bind_joint_effort_limits_sim + + """ + Joint properties - Custom. + """ + + @property + def soft_joint_pos_limits(self) -> wp.array: + r"""Soft joint positions limits for all joints. + + Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_joints, 2). + + The limits are in the order :math:`[lower, upper]`.The soft joint position limits are computed as + a sub-region of the :attr:`joint_pos_limits` based on the + :attr:`~isaaclab.assets.ArticulationCfg.soft_joint_pos_limit_factor` parameter. + + Consider the joint position limits :math:`[lower, upper]` and the soft joint position limits + :math:`[soft_lower, soft_upper]`. The soft joint position limits are computed as: + + .. math:: + + soft\_lower = (lower + upper) / 2 - factor * (upper - lower) / 2 + soft\_upper = (lower + upper) / 2 + factor * (upper - lower) / 2 + + The soft joint position limits help specify a safety region around the joint limits. It isn't used by the + simulation, but is useful for learning agents to prevent the joint positions from violating the limits. + """ + return self._soft_joint_pos_limits + + @property + def soft_joint_vel_limits(self) -> wp.array: + """Soft joint velocity limits for all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + + These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model + has a variable velocity limit model. For instance, in a variable gear ratio actuator model. + """ + return self._soft_joint_vel_limits + + @property + def gear_ratio(self) -> wp.array: + """Gear ratio for relating motor torques to applied Joint torques. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ + return self._gear_ratio + + """ + Fixed tendon properties. + """ + + @property + def fixed_tendon_stiffness(self) -> wp.array: + """Fixed tendon stiffness provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ + raise NotImplementedError + + @property + def fixed_tendon_damping(self) -> wp.array: + """Fixed tendon damping provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ + raise NotImplementedError + + @property + def fixed_tendon_limit_stiffness(self) -> wp.array: + """Fixed tendon limit stiffness provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ + raise NotImplementedError + + @property + def fixed_tendon_rest_length(self) -> wp.array: + """Fixed tendon rest length provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ + raise NotImplementedError + + @property + def fixed_tendon_offset(self) -> wp.array: + """Fixed tendon offset provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ + raise NotImplementedError + + @property + def fixed_tendon_pos_limits(self) -> wp.array: + """Fixed tendon position limits provided to the simulation. + + Shape is (num_instances, num_fixed_tendons, 2), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_fixed_tendons, 2). + """ + raise NotImplementedError + + """ + Spatial tendon properties. + """ + + @property + def spatial_tendon_stiffness(self) -> wp.array: + """Spatial tendon stiffness provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ + raise NotImplementedError + + @property + def spatial_tendon_damping(self) -> wp.array: + """Spatial tendon damping provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ + raise NotImplementedError + + @property + def spatial_tendon_limit_stiffness(self) -> wp.array: + """Spatial tendon limit stiffness provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ + raise NotImplementedError + + @property + def spatial_tendon_offset(self) -> wp.array: + """Spatial tendon offset provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ + raise NotImplementedError + + """ + Root state properties. + """ + + @property + def root_link_pose_w(self) -> wp.array: + """Root link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + + This quantity is the pose of the articulation root's actor frame relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self._sim_bind_root_link_pose_w + + @property + def root_link_vel_w(self) -> wp.array: + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + + This quantity contains the linear and angular velocities of the articulation root's actor frame + relative to the world. + """ + if self._root_link_vel_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.get_root_link_vel_from_root_com_vel, + dim=self._num_instances, + inputs=[ + self.root_com_vel_w, + self.root_link_quat_w, + self.body_com_pos_b, + ], + outputs=[ + self._root_link_vel_w.data, + ], + device=self.device, + ) + self._root_link_vel_w.timestamp = self._sim_timestamp + + return self._root_link_vel_w.data + + @property + def root_com_pose_w(self) -> wp.array: + """Root center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + + This quantity is the pose of the articulation root's center of mass frame relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + if self._root_com_pose_w.timestamp < self._sim_timestamp: + # apply local transform to center of mass frame + wp.launch( + shared_kernels.get_root_com_pose_from_root_link_pose, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.body_com_pos_b, + ], + outputs=[ + self._root_com_pose_w.data, + ], + device=self.device, + ) + self._root_com_pose_w.timestamp = self._sim_timestamp + + return self._root_com_pose_w.data + + @property + def root_com_vel_w(self) -> wp.array: + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + + This quantity contains the linear and angular velocities of the articulation root's center of mass frame + relative to the world. + """ + return self._sim_bind_root_com_vel_w + + """ + Body state properties. + """ + + @property + def body_mass(self) -> wp.array: + """Body mass ``wp.float32`` in the world frame. + + Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). + """ + return self._sim_bind_body_mass + + @property + def body_inertia(self) -> wp.array: + """Flattened body inertia in the world frame. + + Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to + (num_instances, num_bodies, 9). + """ + return self._sim_bind_body_inertia + + @property + def body_link_pose_w(self) -> wp.array: + """Body link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). + + This quantity is the pose of the articulation links' actor frame relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self._sim_bind_body_link_pose_w + + @property + def body_link_vel_w(self) -> wp.array: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). + + This quantity contains the linear and angular velocities of the articulation links' actor frame + relative to the world. + """ + if self._body_link_vel_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.get_body_link_vel_from_body_com_vel, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_com_vel_w, + self.body_link_pose_w, + self.body_com_pos_b, + ], + outputs=[ + self._body_link_vel_w.data, + ], + device=self.device, + ) + self._body_link_vel_w.timestamp = self._sim_timestamp + + return self._body_link_vel_w.data + + @property + def body_com_pose_w(self) -> wp.array: + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). + + This quantity is the pose of the center of mass frame of the articulation links relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + if self._body_com_pose_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.get_body_com_pose_from_body_link_pose, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_com_pos_b, + ], + outputs=[ + self._body_com_pose_w.data, + ], + device=self.device, + ) + self._body_com_pose_w.timestamp = self._sim_timestamp + + return self._body_com_pose_w.data + + @property + def body_com_vel_w(self) -> wp.array: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). + + This quantity contains the linear and angular velocities of the articulation links' center of mass frame + relative to the world. + """ + return self._sim_bind_body_com_vel_w + + @property + def body_com_acc_w(self): + """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]``. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). + + All values are relative to the world. + """ + if self._body_com_acc_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.derive_body_acceleration_from_body_com_velocities, + dim=(self._num_instances, self._num_bodies), + device=self.device, + inputs=[ + self._sim_bind_body_com_vel_w, + SimulationManager.get_dt(), + self._previous_body_com_vel, + ], + outputs=[ + self._body_com_acc_w.data, + ], + ) + # set the buffer data and timestamp + self._body_com_acc_w.timestamp = self._sim_timestamp + # update the previous velocity + return self._body_com_acc_w.data + + @property + def body_com_pos_b(self) -> wp.array: + """Center of mass position of all of the bodies in their respective link frames. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the center of mass location relative to its body's link frame. + """ + return self._sim_bind_body_com_pos_b + + @property + def body_com_pose_b(self) -> wp.array: + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. + + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). + + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + The orientation is provided in (x, y, z, w) format. + """ + warnings.warn( + "In Newton, body com pose always has unit quaternion. Consider using body_com_pos_b instead." + "Querying this property requires appending a unit quaternion to the position which is expensive.", + category=UserWarning, + stacklevel=2, + ) + if self._body_com_pose_b.timestamp < self._sim_timestamp: + # set the buffer data and timestamp + wp.launch( + shared_kernels.make_dummy_body_com_pose_b, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_com_pos_b, + ], + outputs=[ + self._body_com_pose_b.data, + ], + device=self.device, + ) + self._body_com_pose_b.timestamp = self._sim_timestamp + return self._body_com_pose_b.data + + @property + def body_incoming_joint_wrench_b(self) -> wp.array: + """Joint reaction wrench applied from body parent to child body in parent body frame. + + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the + world of an articulation. + + For more information on joint wrenches, please check the `PhysX documentation`_ and the + underlying `PhysX Tensor API`_. + + .. _PhysX documentation: https://nvidia-omniverse.github.io/PhysX/physx/5.5.1/docs/Articulations.html#link-incoming-joint-force + .. _PhysX Tensor API: https://docs.omniverse.nvidia.com/kit/docs/omni_physics/latest/extensions/runtime/source/omni.physics.tensors/docs/api/python.html#omni.physics.tensors.impl.api.ArticulationView.get_link_incoming_joint_force + """ + raise NotImplementedError + + """ + Joint state properties. + """ + + @property + def joint_pos(self) -> wp.array: + """Joint positions of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to + (num_instances, num_joints). + """ + return self._sim_bind_joint_pos + + @property + def joint_vel(self) -> wp.array: + """Joint velocities of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to + (num_instances, num_joints). + """ + return self._sim_bind_joint_vel + + @property + def joint_acc(self) -> wp.array: + """Joint acceleration of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to + (num_instances, num_joints). + """ + if self._joint_acc.timestamp < self._sim_timestamp: + # note: we use finite differencing to compute acceleration + time_elapsed = self._sim_timestamp - self._joint_acc.timestamp + wp.launch( + articulation_kernels.get_joint_acc_from_joint_vel, + dim=(self._num_instances, self._num_joints), + inputs=[ + self.joint_vel, + self._previous_joint_vel, + time_elapsed, + ], + outputs=[ + self._joint_acc.data, + ], + device=self.device, + ) + self._joint_acc.timestamp = self._sim_timestamp + return self._joint_acc.data + + """ + Derived Properties. + """ + + @property + def projected_gravity_b(self): + """Projection of the gravity direction on base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + if self._projected_gravity_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.GRAVITY_VEC_W, self.root_link_quat_w], + outputs=[self._projected_gravity_b.data], + device=self.device, + ) + self._projected_gravity_b.timestamp = self._sim_timestamp + return self._projected_gravity_b.data + + @property + def heading_w(self): + """Yaw heading of the base frame (in radians). + + Shape is (num_instances), dtype = wp.float32. In torch this resolves to (num_instances,). + + .. note:: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + if self._heading_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.root_heading_w, + dim=self._num_instances, + inputs=[self.FORWARD_VEC_B, self.root_link_quat_w], + outputs=[self._heading_w.data], + device=self.device, + ) + self._heading_w.timestamp = self._sim_timestamp + return self._heading_w.data + + @property + def root_link_lin_vel_b(self) -> wp.array: + """Root link linear velocity in base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the linear velocity of the articulation root's actor frame with respect to + its actor frame. + """ + if self._root_link_lin_vel_b is None: + self._root_link_lin_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.vec3f, device=self.device + ) + if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_link_lin_vel_w, self.root_link_quat_w], + outputs=[self._root_link_lin_vel_b.data], + device=self.device, + ) + self._root_link_lin_vel_b.timestamp = self._sim_timestamp + return self._root_link_lin_vel_b.data + + @property + def root_link_ang_vel_b(self) -> wp.array: + """Root link angular velocity in base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the angular velocity of the articulation root's actor frame with respect to + its actor frame. + """ + if self._root_link_ang_vel_b is None: + self._root_link_ang_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.vec3f, device=self.device + ) + if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_link_ang_vel_w, self.root_link_quat_w], + outputs=[self._root_link_ang_vel_b.data], + device=self.device, + ) + self._root_link_ang_vel_b.timestamp = self._sim_timestamp + return self._root_link_ang_vel_b.data + + @property + def root_com_lin_vel_b(self) -> wp.array: + """Root center of mass linear velocity in base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the linear velocity of the articulation root's center of mass frame with respect to + its actor frame. + """ + if self._root_com_lin_vel_b is None: + self._root_com_lin_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.vec3f, device=self.device + ) + if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_com_lin_vel_w, self.root_link_quat_w], + outputs=[self._root_com_lin_vel_b.data], + device=self.device, + ) + self._root_com_lin_vel_b.timestamp = self._sim_timestamp + return self._root_com_lin_vel_b.data + + @property + def root_com_ang_vel_b(self) -> wp.array: + """Root center of mass angular velocity in base frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the angular velocity of the articulation root's center of mass frame with respect to + its actor frame. + """ + if self._root_com_ang_vel_b is None: + self._root_com_ang_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.vec3f, device=self.device + ) + if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_com_ang_vel_w, self.root_link_quat_w], + outputs=[self._root_com_ang_vel_b.data], + device=self.device, + ) + self._root_com_ang_vel_b.timestamp = self._sim_timestamp + return self._root_com_ang_vel_b.data + + """ + Sliced properties. + """ + + @property + def root_link_pos_w(self) -> wp.array: + """Root link position in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + return self._get_pos_from_transform(self._root_link_pos_w, self.root_link_pose_w) + + @property + def root_link_quat_w(self) -> wp.array: + """Root link orientation (x, y, z, w) in simulation world frame. + + Shape is (num_instances), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + + This quantity is the orientation of the actor frame of the root rigid body. + """ + return self._get_quat_from_transform(self._root_link_quat_w, self.root_link_pose_w) + + @property + def root_link_lin_vel_w(self) -> wp.array: + """Root linear velocity in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + return self._get_top_from_spatial_vector(self._root_link_lin_vel_w, self.root_link_vel_w) + + @property + def root_link_ang_vel_w(self) -> wp.array: + """Root link angular velocity in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + return self._get_bottom_from_spatial_vector(self._root_link_ang_vel_w, self.root_link_vel_w) + + @property + def root_com_pos_w(self) -> wp.array: + """Root center of mass position in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the position of the center of mass frame of the root rigid body relative to the world. + """ + return self._get_pos_from_transform(self._root_com_pos_w, self.root_com_pose_w) + + @property + def root_com_quat_w(self) -> wp.array: + """Root center of mass orientation (x, y, z, w) in simulation world frame. + + Shape is (num_instances), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + + This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. + """ + return self._get_quat_from_transform(self._root_com_quat_w, self.root_com_pose_w) + + @property + def root_com_lin_vel_w(self) -> wp.array: + """Root center of mass linear velocity in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + return self._get_top_from_spatial_vector(self._root_com_lin_vel_w, self.root_com_vel_w) + + @property + def root_com_ang_vel_w(self) -> wp.array: + """Root center of mass angular velocity in simulation world frame. + + Shape is (num_instances), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + return self._get_bottom_from_spatial_vector(self._root_com_ang_vel_w, self.root_com_vel_w) + + @property + def body_link_pos_w(self) -> wp.array: + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the position of the articulation bodies' actor frame relative to the world. + """ + return self._get_pos_from_transform(self._body_link_pos_w, self.body_link_pose_w) + + @property + def body_link_quat_w(self) -> wp.array: + """Orientation (x, y, z, w) of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + + This quantity is the orientation of the articulation bodies' actor frame relative to the world. + """ + return self._get_quat_from_transform(self._body_link_quat_w, self.body_link_pose_w) + + @property + def body_link_lin_vel_w(self) -> wp.array: + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the linear velocity of the articulation bodies' actor frame relative to the world. + """ + return self._get_top_from_spatial_vector(self._body_link_lin_vel_w, self.body_link_vel_w) + + @property + def body_link_ang_vel_w(self) -> wp.array: + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the angular velocity of the articulation bodies' actor frame relative to the world. + """ + return self._get_bottom_from_spatial_vector(self._body_link_ang_vel_w, self.body_link_vel_w) + + @property + def body_com_pos_w(self) -> wp.array: + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the position of the articulation bodies' center of mass frame. + """ + return self._get_pos_from_transform(self._body_com_pos_w, self.body_com_pose_w) + + @property + def body_com_quat_w(self) -> wp.array: + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + + This quantity is the orientation of the principal axes of inertia of the articulation bodies. + """ + return self._get_quat_from_transform(self._body_com_quat_w, self.body_com_pose_w) + + @property + def body_com_lin_vel_w(self) -> wp.array: + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the linear velocity of the articulation bodies' center of mass frame. + """ + return self._get_top_from_spatial_vector(self._body_com_lin_vel_w, self.body_com_vel_w) + + @property + def body_com_ang_vel_w(self) -> wp.array: + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the angular velocity of the articulation bodies' center of mass frame. + """ + return self._get_bottom_from_spatial_vector(self._body_com_ang_vel_w, self.body_com_vel_w) + + @property + def body_com_lin_acc_w(self) -> wp.array: + """Linear acceleration of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the linear acceleration of the articulation bodies' center of mass frame. + """ + return self._get_top_from_spatial_vector(self._body_com_lin_acc_w, self.body_com_acc_w) + + @property + def body_com_ang_acc_w(self) -> wp.array: + """Angular acceleration of all bodies in simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + + This quantity is the angular acceleration of the articulation bodies' center of mass frame. + """ + return self._get_bottom_from_spatial_vector(self._body_com_ang_acc_w, self.body_com_acc_w) + + @property + def body_com_quat_b(self) -> wp.array: + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their respective link + frames. + + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. + """ + return self._get_quat_from_transform(self._body_com_quat_b, self.body_com_pose_b) + + def _create_simulation_bindings(self) -> None: + """Create simulation bindings for the root data. + + Direct simulation bindings are pointers to the simulation data, their data is not copied, and should + only be updated using warp kernels. Any modifications made to the bindings will be reflected in the simulation. + Hence we encourage users to carefully think about the data they modify and in which order it should be updated. + + .. caution:: This is possible if and only if the properties that we access are strided from newton and not + indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. + """ + # Short-hand for the number of instances, number of links, and number of joints. + n_view = self._root_view.count + n_dof = self._root_view.joint_dof_count + + # -- root properties + if self._root_view.is_fixed_base: + self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[ + :, 0, 0 + ] + else: + self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[:, 0] + self._sim_bind_root_com_vel_w = self._root_view.get_root_velocities(SimulationManager.get_state_0()) + if self._sim_bind_root_com_vel_w is not None: + if self._root_view.is_fixed_base: + self._sim_bind_root_com_vel_w = self._sim_bind_root_com_vel_w[:, 0, 0] + else: + self._sim_bind_root_com_vel_w = self._sim_bind_root_com_vel_w[:, 0] + # -- body properties + self._sim_bind_body_com_pos_b = self._root_view.get_attribute("body_com", SimulationManager.get_model())[:, 0] + self._sim_bind_body_link_pose_w = self._root_view.get_link_transforms(SimulationManager.get_state_0())[:, 0] + self._sim_bind_body_com_vel_w = self._root_view.get_link_velocities(SimulationManager.get_state_0()) + if self._sim_bind_body_com_vel_w is not None: + self._sim_bind_body_com_vel_w = self._sim_bind_body_com_vel_w[:, 0] + self._sim_bind_body_mass = self._root_view.get_attribute("body_mass", SimulationManager.get_model())[:, 0] + self._sim_bind_body_inertia = self._root_view.get_attribute("body_inertia", SimulationManager.get_model())[:, 0] + self._sim_bind_body_external_wrench = self._root_view.get_attribute("body_f", SimulationManager.get_state_0())[ + :, 0 + ] + # -- joint properties + if n_dof > 0: + self._sim_bind_joint_pos_limits_lower = self._root_view.get_attribute( + "joint_limit_lower", SimulationManager.get_model() + )[:, 0] + self._sim_bind_joint_pos_limits_upper = self._root_view.get_attribute( + "joint_limit_upper", SimulationManager.get_model() + )[:, 0] + self._sim_bind_joint_stiffness_sim = self._root_view.get_attribute( + "joint_target_ke", SimulationManager.get_model() + )[:, 0] + self._sim_bind_joint_damping_sim = self._root_view.get_attribute( + "joint_target_kd", SimulationManager.get_model() + )[:, 0] + self._sim_bind_joint_armature = self._root_view.get_attribute( + "joint_armature", SimulationManager.get_model() + )[:, 0] + self._sim_bind_joint_friction_coeff = self._root_view.get_attribute( + "joint_friction", SimulationManager.get_model() + )[:, 0] + self._sim_bind_joint_vel_limits_sim = self._root_view.get_attribute( + "joint_velocity_limit", SimulationManager.get_model() + )[:, 0] + self._sim_bind_joint_effort_limits_sim = self._root_view.get_attribute( + "joint_effort_limit", SimulationManager.get_model() + )[:, 0] + # -- joint states + self._sim_bind_joint_pos = self._root_view.get_dof_positions(SimulationManager.get_state_0())[:, 0] + self._sim_bind_joint_vel = self._root_view.get_dof_velocities(SimulationManager.get_state_0())[:, 0] + # -- joint commands (sent to the simulation) + self._sim_bind_joint_effort = self._root_view.get_attribute("joint_f", SimulationManager.get_control())[ + :, 0 + ] + self._sim_bind_joint_position_target = self._root_view.get_attribute( + "joint_target_pos", SimulationManager.get_control() + )[:, 0] + self._sim_bind_joint_velocity_target = self._root_view.get_attribute( + "joint_target_vel", SimulationManager.get_control() + )[:, 0] + else: + # No joints (e.g., free-floating rigid body) - set bindings to empty arrays + self._sim_bind_joint_pos_limits_lower = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_pos_limits_upper = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_stiffness_sim = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_damping_sim = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_armature = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_friction_coeff = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_vel_limits_sim = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_effort_limits_sim = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_pos = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_vel = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_effort = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_position_target = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + self._sim_bind_joint_velocity_target = wp.zeros((n_view, 0), dtype=wp.float32, device=self.device) + + def _create_buffers(self) -> None: + super()._create_buffers() + """Create buffers for the root data.""" + + # Short-hand for the number of instances, number of links, and number of joints. + self._num_instances = self._root_view.count + self._num_joints = self._root_view.joint_dof_count + self._num_bodies = self._root_view.link_count + self._num_fixed_tendons = 0 # self._root_view.max_fixed_tendons + self._num_spatial_tendons = 0 # self._root_view.max_spatial_tendons + + # Initialize history for finite differencing. If the articulation is fixed, the root com velocity is not + # available, so we use zeros. + if self._root_view.get_root_velocities(SimulationManager.get_state_0()) is None: + logger.warning( + "Failed to get root com velocity. If the articulation is fixed, this is expected." + "Setting root com velocity to zeros." + ) + self._sim_bind_root_com_vel_w = wp.zeros( + (self._num_instances), dtype=wp.spatial_vectorf, device=self.device + ) + self._sim_bind_body_com_vel_w = wp.zeros( + (self._num_instances, self._num_bodies), dtype=wp.spatial_vectorf, device=self.device + ) + # -- default root pose and velocity + self._default_root_pose = wp.zeros((self._num_instances,), dtype=wp.transformf, device=self.device) + self._default_root_vel = wp.zeros((self._num_instances,), dtype=wp.spatial_vectorf, device=self.device) + # -- default joint positions and velocities + self._default_joint_pos = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._default_joint_vel = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + # -- joint commands (sent to the actuator from the user) + self._joint_pos_target = wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + self._joint_vel_target = wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + self._joint_effort_target = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + # -- computed joint efforts from the actuator models + self._computed_torque = wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + self._applied_torque = wp.zeros((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + # -- joint properties for the actuator models + if self._num_joints > 0: + self._actuator_stiffness = wp.clone(self._sim_bind_joint_stiffness_sim) + self._actuator_damping = wp.clone(self._sim_bind_joint_damping_sim) + else: + self._actuator_stiffness = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + self._actuator_damping = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + # -- other data that are filled based on explicit actuator models + self._joint_dynamic_friction = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._joint_viscous_friction = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._soft_joint_vel_limits = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + self._gear_ratio = wp.ones((self._num_instances, self._num_joints), dtype=wp.float32, device=self.device) + # -- update the soft joint position limits + self._soft_joint_pos_limits = wp.zeros( + (self._num_instances, self._num_joints), dtype=wp.vec2f, device=self.device + ) + + # Initialize history for finite differencing + if self._num_joints > 0: + self._previous_joint_vel = wp.clone( + self._root_view.get_dof_velocities(SimulationManager.get_state_0())[:, 0] + ) + else: + self._previous_joint_vel = wp.zeros((self._num_instances, 0), dtype=wp.float32, device=self.device) + self._previous_body_com_vel = wp.clone(self._sim_bind_body_com_vel_w) + + # Initialize the lazy buffers. + # -- link frame w.r.t. world frame + self._root_link_vel_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._root_link_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._body_link_vel_w = TimestampedBuffer( + shape=(self._num_instances, self._num_bodies), dtype=wp.spatial_vectorf, device=self.device + ) + # -- com frame w.r.t. link frame + self._body_com_pose_b = TimestampedBuffer( + shape=(self._num_instances, self._num_bodies), dtype=wp.transformf, device=self.device + ) + # -- com frame w.r.t. world frame + self._root_com_pose_w = TimestampedBuffer(shape=(self._num_instances,), dtype=wp.transformf, device=self.device) + self._root_com_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._root_com_acc_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._body_com_pose_w = TimestampedBuffer( + shape=(self._num_instances, self._num_bodies), dtype=wp.transformf, device=self.device + ) + self._body_com_acc_w = TimestampedBuffer( + shape=(self._num_instances, self._num_bodies), dtype=wp.spatial_vectorf, device=self.device + ) + # -- derived properties (these are cached to avoid repeated memory allocations) + self._projected_gravity_b = TimestampedBuffer(shape=(self._num_instances,), dtype=wp.vec3f, device=self.device) + self._heading_w = TimestampedBuffer(shape=(self._num_instances,), dtype=wp.float32, device=self.device) + # -- joint state + self._joint_acc = TimestampedBuffer( + shape=(self._num_instances, self._num_joints), dtype=wp.float32, device=self.device + ) + # self._body_incoming_joint_wrench_b = TimestampedWarpBuffer( + # shape=(self._num_instances, self._num_joints), dtype=wp.spatial_vectorf, device=self.device + # ) + # Empty memory pre-allocations + self._root_link_lin_vel_b = None + self._root_link_ang_vel_b = None + self._root_com_lin_vel_b = None + self._root_com_ang_vel_b = None + self._joint_pos_limits = None + self._root_state_w = None + self._root_link_state_w = None + self._root_com_state_w = None + self._body_com_quat_b = None + self._root_link_pos_w = None + self._root_link_quat_w = None + self._root_link_lin_vel_w = None + self._root_link_ang_vel_w = None + self._root_com_pos_w = None + self._root_com_quat_w = None + self._root_com_lin_vel_w = None + self._root_com_ang_vel_w = None + self._body_state_w = None + self._body_link_state_w = None + self._body_com_state_w = None + self._body_link_pos_w = None + self._body_link_quat_w = None + self._body_link_lin_vel_w = None + self._body_link_ang_vel_w = None + self._body_com_pos_w = None + self._body_com_quat_w = None + self._body_com_lin_vel_w = None + self._body_com_ang_vel_w = None + self._body_com_lin_acc_w = None + self._body_com_ang_acc_w = None + self._default_root_state = None + + """ + Internal helpers. + """ + + def _get_pos_from_transform(self, source: wp.array | None, transform: wp.array) -> wp.array: + """Generates a position array from a transform array. + + Args: + transform: The transform array. Shape is (N) dtype=wp.transformf. + + Returns: + The position array. Shape is (N) dtype=wp.vec3f. + """ + # Check if we already created the lazy buffer. + if source is None: + if transform.is_contiguous: + # Check if the array is contiguous. If so, we can just return a strided array. + # Then this update becomes a no-op. + return wp.array( + ptr=transform.ptr, + shape=transform.shape, + dtype=wp.vec3f, + strides=transform.strides, + device=self.device, + ) + else: + # If the array is not contiguous, we need to create a new array to write to. + # Shape matches transform.shape since each element is vec3f (already contains 3 floats) + source = wp.zeros(transform.shape, dtype=wp.vec3f, device=self.device) + + # If the array is not contiguous, we need to launch the kernel to get the position part of the transform. + if not transform.is_contiguous: + # Launch the right kernel based on the shape of the transform array. + if len(transform.shape) > 1: + wp.launch( + shared_kernels.split_transform_to_pos_2d, + dim=transform.shape, + inputs=[transform], + outputs=[source], + device=self.device, + ) + else: + wp.launch( + shared_kernels.split_transform_to_pos_1d, + dim=transform.shape, + inputs=[transform], + outputs=[source], + device=self.device, + ) + return source + + def _get_quat_from_transform(self, source: wp.array | None, transform: wp.array) -> wp.array: + """Generates a quaternion array from a transform array. + + Args: + transform: The transform array. Shape is (N) dtype=wp.transformf. + + Returns: + The quaternion array. Shape is (N) dtype=wp.quatf. + """ + # Check if we already created the lazy buffer. + if source is None: + if transform.is_contiguous: + # Check if the array is contiguous. If so, we can just return a strided array. + # Then this update becomes a no-op. + return wp.array( + ptr=transform.ptr + 3 * 4, + shape=transform.shape, + dtype=wp.quatf, + strides=transform.strides, + device=self.device, + ) + else: + # If the array is not contiguous, we need to create a new array to write to. + # Shape matches transform.shape since each element is quatf (already contains 4 floats) + source = wp.zeros(transform.shape, dtype=wp.quatf, device=self.device) + + # If the array is not contiguous, we need to launch the kernel to get the quaternion part of the transform. + if not transform.is_contiguous: + # Launch the right kernel based on the shape of the transform array. + if len(transform.shape) > 1: + wp.launch( + shared_kernels.split_transform_to_quat_2d, + dim=transform.shape, + inputs=[transform], + outputs=[source], + device=self.device, + ) + else: + wp.launch( + shared_kernels.split_transform_to_quat_1d, + dim=transform.shape, + inputs=[transform], + outputs=[source], + device=self.device, + ) + # Return the source array. (no-op if the array is contiguous.) + return source + + def _get_top_from_spatial_vector(self, source: wp.array | None, spatial_vector: wp.array) -> wp.array: + """Gets the top part of a spatial vector array. + + For instance the linear velocity is the top part of a velocity vector. + + Args: + spatial_vector: The spatial vector array. Shape is (N) dtype=wp.spatial_vectorf. + + Returns: + The top part of the spatial vector array. Shape is (N) dtype=wp.vec3f. + """ + # Check if we already created the lazy buffer. + if source is None: + if spatial_vector.is_contiguous: + # Check if the array is contiguous. If so, we can just return a strided array. + # Then this update becomes a no-op. + return wp.array( + ptr=spatial_vector.ptr, + shape=spatial_vector.shape, + dtype=wp.vec3f, + strides=spatial_vector.strides, + device=self.device, + ) + else: + # If the array is not contiguous, we need to create a new array to write to. + # Shape matches spatial_vector.shape since each element is vec3f (already contains 3 floats) + source = wp.zeros(spatial_vector.shape, dtype=wp.vec3f, device=self.device) + + # If the array is not contiguous, we need to launch the kernel to get the top part of the spatial vector. + if not spatial_vector.is_contiguous: + # Launch the right kernel based on the shape of the spatial_vector array. + if len(spatial_vector.shape) > 1: + wp.launch( + shared_kernels.split_spatial_vector_to_top_2d, + dim=spatial_vector.shape, + inputs=[spatial_vector], + outputs=[source], + device=self.device, + ) + else: + wp.launch( + shared_kernels.split_spatial_vector_to_top_1d, + dim=spatial_vector.shape, + inputs=[spatial_vector], + outputs=[source], + device=self.device, + ) + # Return the source array. (no-op if the array is contiguous.) + return source + + def _get_bottom_from_spatial_vector(self, source: wp.array | None, spatial_vector: wp.array) -> wp.array: + """Gets the bottom part of a spatial vector array. + + For instance the angular velocity is the bottom part of a velocity vector. + + Args: + spatial_vector: The spatial vector array. Shape is (N) dtype=wp.spatial_vectorf. + + Returns: + The bottom part of the spatial vector array. Shape is (N) dtype=wp.vec3f. + """ + # Check if we already created the lazy buffer. + if source is None: + if spatial_vector.is_contiguous: + # Check if the array is contiguous. If so, we can just return a strided array. + # Then this update becomes a no-op. + return wp.array( + ptr=spatial_vector.ptr + 3 * 4, + shape=spatial_vector.shape, + dtype=wp.vec3f, + strides=spatial_vector.strides, + device=self.device, + ) + else: + # If the array is not contiguous, we need to create a new array to write to. + # Shape matches spatial_vector.shape since each element is vec3f (already contains 3 floats) + source = wp.zeros(spatial_vector.shape, dtype=wp.vec3f, device=self.device) + + # If the array is not contiguous, we need to launch the kernel to get the bottom part of the spatial vector. + if not spatial_vector.is_contiguous: + # Launch the right kernel based on the shape of the spatial_vector array. + if len(spatial_vector.shape) > 1: + wp.launch( + shared_kernels.split_spatial_vector_to_bottom_2d, + dim=spatial_vector.shape, + inputs=[spatial_vector], + outputs=[source], + device=self.device, + ) + else: + wp.launch( + shared_kernels.split_spatial_vector_to_bottom_1d, + dim=spatial_vector.shape, + inputs=[spatial_vector], + outputs=[source], + device=self.device, + ) + # Return the source array. (no-op if the array is contiguous.) + return source + + """ + Deprecated properties. + """ + + @property + def root_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" + warnings.warn( + "The `root_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_state_w is None: + self._root_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) + if self._root_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=(self._num_instances), + inputs=[ + self.root_link_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_state_w.data, + ], + device=self.device, + ) + self._root_state_w.timestamp = self._sim_timestamp + + return self._root_state_w.data + + @property + def root_link_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" + warnings.warn( + "The `root_link_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_link_state_w is None: + self._root_link_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) + if self._root_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_link_vel_w, + ], + outputs=[ + self._root_link_state_w.data, + ], + device=self.device, + ) + self._root_link_state_w.timestamp = self._sim_timestamp + + return self._root_link_state_w.data + + @property + def root_com_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" + warnings.warn( + "The `root_com_state_w` property will be deprecated in a IsaacLab 4.0. Please use `root_com_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_com_state_w is None: + self._root_com_state_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=shared_kernels.vec13f, device=self.device + ) + if self._root_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_com_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_com_state_w.data, + ], + device=self.device, + ) + self._root_com_state_w.timestamp = self._sim_timestamp + + return self._root_com_state_w.data + + @property + def default_root_state(self) -> wp.array: + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in the local environment frame. + + + The position and quaternion are of the articulation root's actor frame. Meanwhile, the linear and angular + velocities are of its center of mass frame. Shape is (num_instances, 13). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + warnings.warn( + "Reading the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " + "Please use the default_root_pose and default_root_vel properties instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._default_root_state is None: + self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self._default_root_pose, + self._default_root_vel, + ], + outputs=[ + self._default_root_state, + ], + device=self.device, + ) + return self._default_root_state + + @property + def body_state_w(self): + """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position and quaternion are of all the articulation links' actor frame. Meanwhile, the linear and angular + velocities are of the articulation links's center of mass frame. + """ + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_state_w is None: + self._body_state_w = TimestampedBuffer( + (self._num_instances, self._num_bodies), self.device, shared_kernels.vec13f + ) + if self._body_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_com_vel_w, + ], + outputs=[ + self._body_state_w.data, + ], + device=self.device, + ) + self._body_state_w.timestamp = self._sim_timestamp + + return self._body_state_w.data + + @property + def body_link_state_w(self): + """State of all bodies' link frame`[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. + """ + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_link_state_w is None: + self._body_link_state_w = TimestampedBuffer( + (self._num_instances, self._num_bodies), self.device, shared_kernels.vec13f + ) + if self._body_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_link_vel_w, + ], + outputs=[ + self._body_link_state_w.data, + ], + device=self.device, + ) + self._body_link_state_w.timestamp = self._sim_timestamp + + return self._body_link_state_w.data + + @property + def body_com_state_w(self): + """State of all bodies center of mass `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the + world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the + principal inertia. + """ + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_com_state_w is None: + self._body_com_state_w = TimestampedBuffer( + (self._num_instances, self._num_bodies), self.device, shared_kernels.vec13f + ) + if self._body_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_com_pose_w, + self.body_com_vel_w, + ], + outputs=[ + self._body_com_state_w.data, + ], + device=self.device, + ) + self._body_com_state_w.timestamp = self._sim_timestamp + + return self._body_com_state_w.data diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py new file mode 100644 index 00000000000..6dc714d4e9c --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/kernels.py @@ -0,0 +1,559 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + +""" +Articulation-specific warp functions. +""" + + +@wp.func +def compute_soft_joint_pos_limits_func( + joint_pos_limits: wp.vec2f, + soft_limit_factor: wp.float32, +): + """Compute the soft joint position limits. + + Args: + joint_pos_limits: The joint position limits. + soft_limit_factor: The soft limit factor. + + Returns: + The soft joint position limits. + """ + joint_pos_mean = (joint_pos_limits[0] + joint_pos_limits[1]) / 2.0 + joint_pos_range = joint_pos_limits[1] - joint_pos_limits[0] + return wp.vec2f( + joint_pos_mean - 0.5 * joint_pos_range * soft_limit_factor, + joint_pos_mean + 0.5 * joint_pos_range * soft_limit_factor, + ) + + +""" +Articulation-specific warp kernels. +""" + + +@wp.kernel +def get_joint_acc_from_joint_vel( + joint_vel: wp.array2d(dtype=wp.float32), + prev_joint_vel: wp.array2d(dtype=wp.float32), + dt: wp.float32, + joint_acc: wp.array2d(dtype=wp.float32), +): + """Compute the joint acceleration from the joint velocity using finite differencing. + + This kernel computes the joint acceleration by taking the difference between the current + and previous joint velocities, divided by the time step. It also updates the previous + joint velocity buffer with the current values. + + Args: + joint_vel: Input array of current joint velocities. Shape is (num_envs, num_joints). + prev_joint_vel: Input/output array of previous joint velocities. Shape is (num_envs, num_joints). + This buffer is updated with the current joint velocities after computing acceleration. + dt: Input time step (scalar) used for finite differencing. + joint_acc: Output array where joint accelerations are written. Shape is (num_envs, num_joints). + """ + i, j = wp.tid() + joint_acc[i, j] = (joint_vel[i, j] - prev_joint_vel[i, j]) / dt + prev_joint_vel[i, j] = joint_vel[i, j] + + +@wp.kernel +def write_joint_vel_data_index( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + joint_vel: wp.array2d(dtype=wp.float32), + prev_joint_vel: wp.array2d(dtype=wp.float32), + joint_acc: wp.array2d(dtype=wp.float32), +): + """Write joint velocity data to the output buffers. + + This kernel writes joint velocity data from the input array to the output buffers. + It also updates the previous joint velocity buffer and resets the joint acceleration to 0.0. + + Args: + in_data: Input array containing joint velocity data. Shape is (num_selected_envs, num_selected_joints). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. Shape is (num_selected_joints,). + joint_vel: Output array where joint velocities are written. Shape is (num_envs, num_joints). + prev_joint_vel: Output array where previous joint velocities are written. Shape is + (num_envs, num_joints). + joint_acc: Output array where joint accelerations are reset to 0.0. Shape is + (num_envs, num_joints). + """ + i, j = wp.tid() + joint_vel[env_ids[i], joint_ids[j]] = in_data[i, j] + prev_joint_vel[env_ids[i], joint_ids[j]] = in_data[i, j] + joint_acc[env_ids[i], joint_ids[j]] = 0.0 + + +@wp.kernel +def write_joint_vel_data_mask( + in_data: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + joint_vel: wp.array2d(dtype=wp.float32), + prev_joint_vel: wp.array2d(dtype=wp.float32), + joint_acc: wp.array2d(dtype=wp.float32), +): + """Write joint velocity data to the output buffers. + + This kernel writes joint velocity data from the input array to the output buffers. + It also updates the previous joint velocity buffer and resets the joint acceleration to 0.0. + + Args: + in_data: Input array containing joint velocity data. Shape is (num_envs, num_joints). + env_mask: Input array of environment mask. Shape is (num_envs,). + joint_mask: Input array of joint mask. Shape is (num_joints,). + joint_vel: Output array where joint velocities are written. Shape is (num_envs, num_joints). + prev_joint_vel: Output array where previous joint velocities are written. Shape is + (num_envs, num_joints). + joint_acc: Output array where joint accelerations are reset to 0.0. Shape is + (num_envs, num_joints). + """ + i, j = wp.tid() + if env_mask[i] and joint_mask[j]: + joint_vel[i, j] = in_data[i, j] + prev_joint_vel[i, j] = in_data[i, j] + joint_acc[i, j] = 0.0 + + +@wp.kernel +def write_joint_limit_data_to_buffer_index( + in_data: wp.array2d(dtype=wp.vec2f), + soft_limit_factor: wp.float32, + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + joint_pos_limits_lower: wp.array2d(dtype=wp.float32), + joint_pos_limits_upper: wp.array2d(dtype=wp.float32), + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), + default_joint_pos: wp.array2d(dtype=wp.float32), + clamped_defaults: wp.array(dtype=wp.int32), +): + """Write joint limit data to the output buffers and compute soft limits. + + This kernel writes joint position limits from the input array to the output buffer, + computes soft joint position limits, and clamps default joint positions if they + fall outside the limits. + + Args: + in_data: Input array containing joint position limits as vec2f (lower, upper). + Shape is (num_selected_envs, num_selected_joints). + soft_limit_factor: Input scalar factor for computing soft limits (typically 0.0-1.0). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. Shape is (num_selected_joints,). + joint_pos_limits_lower: Output array where joint position limits lower are written. Shape is + (num_envs, num_joints). + joint_pos_limits_upper: Output array where joint position limits upper are written. Shape is + (num_envs, num_joints). + soft_joint_pos_limits: Output array where soft joint position limits are written. + Shape is (num_envs, num_joints). + default_joint_pos: Input/output array of default joint positions. If any values fall + outside the limits, they are clamped. Shape is (num_envs, num_joints). + clamped_defaults: Output 1-element array flag indicating whether any default joint + positions were clamped. Non-zero if any clamping occurred. Shape is (1,). + """ + i, j = wp.tid() + joint_pos_limits_lower[env_ids[i], joint_ids[j]] = in_data[i, j][0] + joint_pos_limits_upper[env_ids[i], joint_ids[j]] = in_data[i, j][1] + if ( + default_joint_pos[env_ids[i], joint_ids[j]] < joint_pos_limits_lower[env_ids[i], joint_ids[j]] + ) or default_joint_pos[env_ids[i], joint_ids[j]] > joint_pos_limits_upper[env_ids[i], joint_ids[j]]: + wp.atomic_add(clamped_defaults, 0, 1) + default_joint_pos[env_ids[i], joint_ids[j]] = wp.clamp( + default_joint_pos[env_ids[i], joint_ids[j]], + joint_pos_limits_lower[env_ids[i], joint_ids[j]], + joint_pos_limits_upper[env_ids[i], joint_ids[j]], + ) + soft_joint_pos_limits[env_ids[i], joint_ids[j]] = compute_soft_joint_pos_limits_func( + wp.vec2f(joint_pos_limits_lower[env_ids[i], joint_ids[j]], joint_pos_limits_upper[env_ids[i], joint_ids[j]]), + soft_limit_factor, + ) + + +@wp.kernel +def write_joint_limit_data_to_buffer_mask( + in_data: wp.array2d(dtype=wp.vec2f), + soft_limit_factor: wp.float32, + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + joint_pos_limits_lower: wp.array2d(dtype=wp.float32), + joint_pos_limits_upper: wp.array2d(dtype=wp.float32), + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), + default_joint_pos: wp.array2d(dtype=wp.float32), + clamped_defaults: wp.array(dtype=wp.int32), +): + """Write joint limit data to the output buffers and compute soft limits. + + This kernel writes joint position limits from the input array to the output buffer, + computes soft joint position limits, and clamps default joint positions if they + fall outside the limits. + + Args: + in_data: Input array containing joint position limits as vec2f (lower, upper). + Shape is (num_envs, num_joints). + soft_limit_factor: Input scalar factor for computing soft limits (typically 0.0-1.0). + env_mask: Input array of environment mask. Shape is (num_envs,). + joint_mask: Input array of joint mask. Shape is (num_joints,). + joint_pos_limits_lower: Output array where joint position limits lower are written. Shape is + (num_envs, num_joints). + joint_pos_limits_upper: Output array where joint position limits upper are written. Shape is + (num_envs, num_joints). + soft_joint_pos_limits: Output array where soft joint position limits are written. + Shape is (num_envs, num_joints). + default_joint_pos: Input/output array of default joint positions. If any values fall + outside the limits, they are clamped. Shape is (num_envs, num_joints). + clamped_defaults: Output 1-element array flag indicating whether any default joint + positions were clamped. Non-zero if any clamping occurred. Shape is (1,). + """ + i, j = wp.tid() + if env_mask[i] and joint_mask[j]: + joint_pos_limits_lower[i, j] = in_data[i, j][0] + joint_pos_limits_upper[i, j] = in_data[i, j][1] + if (default_joint_pos[i, j] < joint_pos_limits_lower[i, j]) or default_joint_pos[i, j] > joint_pos_limits_upper[ + i, j + ]: + wp.atomic_add(clamped_defaults, 0, 1) + default_joint_pos[i, j] = wp.clamp( + default_joint_pos[i, j], + joint_pos_limits_lower[i, j], + joint_pos_limits_upper[i, j], + ) + soft_joint_pos_limits[i, j] = compute_soft_joint_pos_limits_func( + wp.vec2f(joint_pos_limits_lower[i, j], joint_pos_limits_upper[i, j]), soft_limit_factor + ) + + +@wp.kernel +def write_joint_friction_data_to_buffer( + in_friction: wp.array2d(dtype=wp.float32), + in_dynamic_friction: wp.array2d(dtype=wp.float32), + in_viscous_friction: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + from_mask: bool, + out_friction: wp.array2d(dtype=wp.float32), + out_dynamic_friction: wp.array2d(dtype=wp.float32), + out_viscous_friction: wp.array2d(dtype=wp.float32), + friction_props: wp.array3d(dtype=wp.float32), +): + """Write joint friction data to the output buffers. + + This kernel writes joint friction coefficients from input arrays to output buffers + and updates the friction properties array used by the physics simulation. + + Args: + in_friction: Input array containing joint friction coefficients. Shape is + (num_envs, num_joints) or (num_selected_envs, num_selected_joints) depending + on from_mask. Can be None if not provided. + in_dynamic_friction: Input array containing joint dynamic friction coefficients. + Shape is (num_envs, num_joints) or (num_selected_envs, num_selected_joints). + Can be None if not provided. + in_viscous_friction: Input array containing joint viscous friction coefficients. + Shape is (num_envs, num_joints) or (num_selected_envs, num_selected_joints). + Can be None if not provided. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. Shape is (num_selected_joints,). + from_mask: Input flag indicating whether to use masked indexing. If True, indices from + env_ids and joint_ids are used to index into input arrays. If False, input arrays + are indexed directly using the thread indices. + out_friction: Output array where joint friction coefficients are written. Shape is + (num_envs, num_joints). + out_dynamic_friction: Output array where joint dynamic friction coefficients are written. + Shape is (num_envs, num_joints). + out_viscous_friction: Output array where joint viscous friction coefficients are written. + Shape is (num_envs, num_joints). + friction_props: Output array where friction properties are written for the physics + simulation. Shape is (num_envs, num_joints, 3) where the last dimension contains + [friction, dynamic_friction, viscous_friction]. + """ + i, j = wp.tid() + # First update the output buffers + if from_mask: + out_friction[env_ids[i], joint_ids[j]] = in_friction[env_ids[i], joint_ids[j]] + if in_dynamic_friction: + out_dynamic_friction[env_ids[i], joint_ids[j]] = in_dynamic_friction[env_ids[i], joint_ids[j]] + if in_viscous_friction: + out_viscous_friction[env_ids[i], joint_ids[j]] = in_viscous_friction[env_ids[i], joint_ids[j]] + else: + out_friction[env_ids[i], joint_ids[j]] = in_friction[i, j] + if in_dynamic_friction: + out_dynamic_friction[env_ids[i], joint_ids[j]] = in_dynamic_friction[i, j] + if in_viscous_friction: + out_viscous_friction[env_ids[i], joint_ids[j]] = in_viscous_friction[i, j] + # Then update the friction properties + friction_props[env_ids[i], joint_ids[j], 0] = out_friction[env_ids[i], joint_ids[j]] + if in_dynamic_friction: + friction_props[env_ids[i], joint_ids[j], 1] = out_dynamic_friction[env_ids[i], joint_ids[j]] + if in_viscous_friction: + friction_props[env_ids[i], joint_ids[j], 2] = out_viscous_friction[env_ids[i], joint_ids[j]] + + +@wp.kernel +def write_joint_friction_param_to_buffer( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + buffer_index: wp.int32, + from_mask: bool, + out_data: wp.array2d(dtype=wp.float32), + out_buffer: wp.array3d(dtype=wp.float32), +): + """Write a joint friction parameter to the output buffers. + + This kernel writes a single joint friction parameter (e.g., dynamic or viscous friction) + from the input array to both a 2D output array and a specific slice of a 3D buffer array. + + Args: + in_data: Input array containing joint friction parameter values. Shape is + (num_envs, num_joints) or (num_selected_envs, num_selected_joints) depending + on from_mask. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. Shape is (num_selected_joints,). + buffer_index: Input scalar index specifying which slice of the 3D buffer to write to. + Typically 0 for friction, 1 for dynamic friction, or 2 for viscous friction. + from_mask: Input flag indicating whether to use masked indexing. If True, indices from + env_ids and joint_ids are used to index into in_data. If False, in_data is indexed + directly using the thread indices. + out_data: Output array where friction parameter values are written. Shape is + (num_envs, num_joints). + out_buffer: Output 3D array where friction parameter values are written to the specified + slice. Shape is (num_envs, num_joints, num_friction_params). + """ + i, j = wp.tid() + if from_mask: + out_data[env_ids[i], joint_ids[j]] = in_data[env_ids[i], joint_ids[j]] + out_buffer[env_ids[i], joint_ids[j], buffer_index] = in_data[env_ids[i], joint_ids[j]] + else: + out_data[env_ids[i], joint_ids[j]] = in_data[i, j] + out_buffer[env_ids[i], joint_ids[j], buffer_index] = in_data[i, j] + + +@wp.kernel +def float_data_to_buffer_with_indices( + in_data: wp.float32, + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.float32), +): + """Write a scalar float value to a 2D buffer at specified indices. + + This kernel broadcasts a single scalar float value to all specified (env_id, joint_id) + locations in the output buffer. + + Args: + in_data: Input scalar float value to broadcast. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + joint_ids: Input array of joint indices to write to. Shape is (num_selected_joints,). + out_data: Output array where the scalar value is written. Shape is (num_envs, num_joints). + """ + i, j = wp.tid() + out_data[env_ids[i], joint_ids[j]] = in_data + + +@wp.kernel +def float_data_to_buffer_with_mask( + in_data: wp.float32, + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + out_data: wp.array2d(dtype=wp.float32), +): + """Write a scalar float value to a 2D buffer at specified mask. + + This kernel broadcasts a single scalar float value to all the positions that are marked as True in the environment + and joint masks. + + Args: + in_data: Input scalar float value to broadcast. + env_mask: Input array of environment mask. Shape is (num_envs,). + joint_mask: Input array of joint mask. Shape is (num_joints,). + out_data: Output array where the scalar value is written. Shape is (num_envs, num_joints). + """ + i, j = wp.tid() + if env_mask[i] and joint_mask[j]: + out_data[i, j] = in_data + + +@wp.kernel +def update_soft_joint_pos_limits( + joint_pos_limits: wp.array2d(dtype=wp.vec2f), + soft_limit_factor: wp.float32, + soft_joint_pos_limits: wp.array2d(dtype=wp.vec2f), +): + """Update soft joint position limits based on hard limits and a soft limit factor. + + This kernel computes soft joint position limits from hard joint position limits using + a soft limit factor. Soft limits are typically used to provide a safety margin before + reaching the hard limits. + + Args: + joint_pos_limits: Input array of hard joint position limits as vec2f (lower, upper). + Shape is (num_envs, num_joints). + soft_limit_factor: Input scalar factor for computing soft limits (typically 0.0-1.0). + A value of 1.0 means soft limits equal hard limits, while smaller values create + a tighter range. + soft_joint_pos_limits: Output array where soft joint position limits are written. + Shape is (num_envs, num_joints). + """ + i, j = wp.tid() + soft_joint_pos_limits[i, j] = compute_soft_joint_pos_limits_func(joint_pos_limits[i, j], soft_limit_factor) + + +@wp.kernel +def update_default_joint_values( + source: wp.array(dtype=wp.float32), + ids: wp.array(dtype=wp.int32), + target: wp.array2d(dtype=wp.float32), +): + """Update default joint values from a source array using joint indices. + + This kernel writes values from a 1D source array to specific joint indices in a 2D + target array for all environments. + + Args: + source: Input array containing joint values to write. Shape is (num_joints,). + ids: Input array of joint indices specifying which joints to update. Shape is + (num_selected_joints,). + target: Output array where joint values are written. Shape is (num_envs, num_joints). + Values are written to target[i, ids[j]] for all environments i. + """ + i, j = wp.tid() + target[i, ids[j]] = source[j] + + +@wp.kernel +def update_targets( + source_joint_positions: wp.array2d(dtype=wp.float32), + source_joint_velocities: wp.array2d(dtype=wp.float32), + source_joint_efforts: wp.array2d(dtype=wp.float32), + joint_indices: wp.array(dtype=wp.int32), + target_joint_positions: wp.array2d(dtype=wp.float32), + target_joint_velocities: wp.array2d(dtype=wp.float32), + target_joint_efforts: wp.array2d(dtype=wp.float32), +): + """Update joint target values from source arrays using joint indices. + + This kernel copies joint positions, velocities, and efforts from source arrays to + target arrays, remapping joint indices using the provided joint_indices array. + Only non-None source arrays are processed. + + Args: + source_joint_positions: Input array of source joint positions. Shape is + (num_envs, num_selected_joints). Can be None if not provided. + source_joint_velocities: Input array of source joint velocities. Shape is + (num_envs, num_selected_joints). Can be None if not provided. + source_joint_efforts: Input array of source joint efforts. Shape is + (num_envs, num_selected_joints). Can be None if not provided. + joint_indices: Input array of joint indices for remapping. Shape is + (num_selected_joints,). Specifies which joints in the target arrays to update. + target_joint_positions: Output array where joint positions are written. Shape is + (num_envs, num_joints). + target_joint_velocities: Output array where joint velocities are written. Shape is + (num_envs, num_joints). + target_joint_efforts: Output array where joint efforts are written. Shape is + (num_envs, num_joints). + """ + i, j = wp.tid() + if source_joint_positions: + target_joint_positions[i, joint_indices[j]] = source_joint_positions[i, j] + if source_joint_velocities: + target_joint_velocities[i, joint_indices[j]] = source_joint_velocities[i, j] + if source_joint_efforts: + target_joint_efforts[i, joint_indices[j]] = source_joint_efforts[i, j] + + +@wp.kernel +def update_actuator_state_model( + source_computed_effort: wp.array2d(dtype=wp.float32), + source_applied_effort: wp.array2d(dtype=wp.float32), + source_gear_ratio: wp.array2d(dtype=wp.float32), + source_vel_limits: wp.array2d(dtype=wp.float32), + joint_indices: wp.array(dtype=wp.int32), + target_computed_effort: wp.array2d(dtype=wp.float32), + target_applied_effort: wp.array2d(dtype=wp.float32), + target_gear_ratio: wp.array2d(dtype=wp.float32), + target_soft_joint_vel_limits: wp.array2d(dtype=wp.float32), +): + """Update actuator state model parameters from source arrays using joint indices. + + This kernel copies actuator state model parameters (computed effort, applied effort, + gear ratio, and velocity limits) from source arrays to target arrays, remapping + joint indices using the provided joint_indices array. + + Args: + source_computed_effort: Input array of source computed effort values. Shape is + (num_envs, num_selected_joints). + source_applied_effort: Input array of source applied effort values. Shape is + (num_envs, num_selected_joints). + source_gear_ratio: Input array of source gear ratio values. Shape is + (num_envs, num_selected_joints). Can be None if not provided. + source_vel_limits: Input array of source velocity limit values. Shape is + (num_envs, num_selected_joints). + joint_indices: Input array of joint indices for remapping. Shape is + (num_selected_joints,). Specifies which joints in the target arrays to update. + target_computed_effort: Output array where computed effort values are written. + Shape is (num_envs, num_joints). + target_applied_effort: Output array where applied effort values are written. + Shape is (num_envs, num_joints). + target_gear_ratio: Output array where gear ratio values are written. Shape is + (num_envs, num_joints). + target_soft_joint_vel_limits: Output array where soft joint velocity limits are + written. Shape is (num_envs, num_joints). + """ + i, j = wp.tid() + target_computed_effort[i, joint_indices[j]] = source_computed_effort[i, j] + target_applied_effort[i, joint_indices[j]] = source_applied_effort[i, j] + target_soft_joint_vel_limits[i, joint_indices[j]] = source_vel_limits[i, j] + if source_gear_ratio: + target_gear_ratio[i, joint_indices[j]] = source_gear_ratio[i, j] + + +@wp.kernel +def extract_friction_properties( + friction_props: wp.array3d(dtype=wp.float32), + out_friction: wp.array2d(dtype=wp.float32), + out_dynamic_friction: wp.array2d(dtype=wp.float32), + out_viscous_friction: wp.array2d(dtype=wp.float32), +): + """Extract friction properties from a 3D array into separate 2D arrays. + + This kernel extracts the three friction components (static friction, dynamic friction, + and viscous friction) from a 3D friction properties array into three separate 2D arrays. + + Args: + friction_props: Input 3D array containing friction properties. Shape is + (num_envs, num_joints, 3) where the last dimension contains + [friction, dynamic_friction, viscous_friction]. + out_friction: Output array where static friction coefficients are written. + Shape is (num_envs, num_joints). + out_dynamic_friction: Output array where dynamic friction coefficients are written. + Shape is (num_envs, num_joints). + out_viscous_friction: Output array where viscous friction coefficients are written. + Shape is (num_envs, num_joints). + """ + i, j = wp.tid() + out_friction[i, j] = friction_props[i, j, 0] + out_dynamic_friction[i, j] = friction_props[i, j, 1] + out_viscous_friction[i, j] = friction_props[i, j, 2] + + +@wp.kernel +def concat_joint_pos_limits_lower_and_upper( + joint_pos_limits_lower: wp.array2d(dtype=wp.float32), + joint_pos_limits_upper: wp.array2d(dtype=wp.float32), + joint_pos_limits: wp.array2d(dtype=wp.vec2f), +): + """Concatenate joint position limits lower and upper in a single array. + + Args: + joint_pos_limits_lower: Input array of joint position limits lower. Shape is (num_envs, num_joints). + joint_pos_limits_upper: Input array of joint position limits upper. Shape is (num_envs, num_joints). + joint_pos_limits: Output array where joint position limits are written. Shape is (num_envs, num_joints, 2). + """ + i, j = wp.tid() + joint_pos_limits[i, j] = wp.vec2f(joint_pos_limits_lower[i, j], joint_pos_limits_upper[i, j]) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py new file mode 100644 index 00000000000..436bd4cf1f6 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/kernels.py @@ -0,0 +1,1438 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + +vec13f = wp.types.vector(length=13, dtype=wp.float32) + +""" +Shared @wp.func helpers. +""" + + +@wp.func +def update_wrench_with_force_and_torque( + force: wp.vec3f, + torque: wp.vec3f, +) -> wp.spatial_vectorf: + return wp.spatial_vector(force, torque, wp.float32) + + +@wp.func +def get_link_vel_from_root_com_vel_func( + com_vel: wp.spatial_vectorf, + link_pose: wp.transformf, + body_com_pos_b: wp.vec3f, +): + """Compute link velocity from center-of-mass velocity. + + Transforms a COM spatial velocity into a link-frame velocity by projecting + the angular velocity contribution from the COM offset relative to the link frame. + + Args: + com_vel: COM spatial velocity (angular, linear). + link_pose: Link pose in world frame. + body_com_pos_b: COM position in body (link) frame. + + Returns: + Link spatial velocity (angular, linear). + """ + projected_vel = wp.cross( + wp.spatial_bottom(com_vel), + wp.quat_rotate(wp.transform_get_rotation(link_pose), -body_com_pos_b), + ) + return wp.spatial_vector(wp.spatial_top(com_vel) + projected_vel, wp.spatial_bottom(com_vel)) + + +@wp.func +def get_com_pose_from_link_pose_func( + link_pose: wp.transformf, + body_com_pos_b: wp.vec3f, +): + """Compute COM pose in world frame from link pose and body-frame COM offset. + + Args: + link_pose: Link pose in world frame. + body_com_pos_b: COM position in body (link) frame. + + Returns: + COM pose in world frame. + """ + return link_pose * wp.transformf(body_com_pos_b, wp.quatf(0.0, 0.0, 0.0, 1.0)) + + +@wp.func +def concat_pose_and_vel_to_state_func( + pose: wp.transformf, + vel: wp.spatial_vectorf, +) -> vec13f: + """Concatenate a pose and velocity into a 13-element state vector. + + The state vector layout is [pos(3), quat(4), ang_vel(3), lin_vel(3)]. + + Args: + pose: Pose as a transform (position + quaternion). + vel: Spatial velocity (angular, linear). + + Returns: + 13-element state vector. + """ + return vec13f( + pose[0], pose[1], pose[2], pose[3], pose[4], pose[5], pose[6], vel[0], vel[1], vel[2], vel[3], vel[4], vel[5] + ) + + +@wp.func +def compute_heading_w_func( + forward_vec: wp.vec3f, + quat: wp.quatf, +): + """Compute heading angle (yaw) in world frame from a forward vector and orientation. + + Rotates the forward vector by the quaternion and computes atan2(y, x). + + Args: + forward_vec: Forward direction vector in body frame. + quat: Orientation quaternion. + + Returns: + Heading angle in radians. + """ + forward_w = wp.quat_rotate(quat, forward_vec) + return wp.atan2(forward_w[1], forward_w[0]) + + +@wp.func +def set_state_transforms_func( + state: vec13f, + transform: wp.transformf, +) -> vec13f: + """Set the pose portion (first 7 elements) of a 13-element state vector. + + Overwrites elements [0..6] (position + quaternion) with the given transform, + leaving the velocity portion [7..12] unchanged. + + Args: + state: 13-element state vector to modify. + transform: New pose (position + quaternion). + + Returns: + Updated 13-element state vector. + """ + state[0] = transform[0] + state[1] = transform[1] + state[2] = transform[2] + state[3] = transform[3] + state[4] = transform[4] + state[5] = transform[5] + state[6] = transform[6] + return state + + +@wp.func +def set_state_velocities_func( + state: vec13f, + velocity: wp.spatial_vectorf, +) -> vec13f: + """Set the velocity portion (last 6 elements) of a 13-element state vector. + + Overwrites elements [7..12] (angular + linear velocity) with the given spatial velocity, + leaving the pose portion [0..6] unchanged. + + Args: + state: 13-element state vector to modify. + velocity: New spatial velocity (angular, linear). + + Returns: + Updated 13-element state vector. + """ + state[7] = velocity[0] + state[8] = velocity[1] + state[9] = velocity[2] + state[10] = velocity[3] + state[11] = velocity[4] + state[12] = velocity[5] + return state + + +@wp.func +def get_link_velocity_in_com_frame_func( + link_velocity_w: wp.spatial_vectorf, + link_pose_w: wp.transformf, + body_com_pose_b: wp.transformf, +): + """Compute COM velocity from link velocity by accounting for the COM offset. + + Transforms a link-frame spatial velocity into a COM-frame velocity by adding + the cross-product contribution of the COM offset rotated into the world frame. + + Args: + link_velocity_w: Link spatial velocity in world frame (angular, linear). + link_pose_w: Link pose in world frame. + body_com_pose_b: COM pose in body (link) frame. + + Returns: + COM spatial velocity in world frame (angular, linear). + """ + return wp.spatial_vector( + wp.spatial_top(link_velocity_w) + + wp.cross( + wp.spatial_bottom(link_velocity_w), + wp.quat_rotate(wp.transform_get_rotation(link_pose_w), wp.transform_get_translation(body_com_pose_b)), + ), + wp.spatial_bottom(link_velocity_w), + ) + + +@wp.func +def get_com_pose_in_link_frame_func( + com_pose_w: wp.transformf, + com_pose_b: wp.transformf, +): + """Compute link pose in world frame from COM pose by inverting the body-frame COM offset. + + This is the inverse of ``get_com_pose_from_link_pose_func``. Given the COM pose in + world frame and the COM offset in body frame, it recovers the link pose in world frame. + + Args: + com_pose_w: COM pose in world frame. + com_pose_b: COM pose in body (link) frame. + + Returns: + Link pose in world frame. + """ + T2 = wp.transform( + wp.quat_rotate( + wp.quat_inverse(wp.transform_get_rotation(com_pose_b)), -wp.transform_get_translation(com_pose_b) + ), + wp.quat_inverse(wp.transform_get_rotation(com_pose_b)), + ) + link_pose_w = com_pose_w * T2 + return link_pose_w + + +""" +Root-level @wp.kernel (1D — used by RigidObject + Articulation). +""" + + +@wp.kernel +def get_root_link_vel_from_root_com_vel( + com_vel: wp.array(dtype=wp.spatial_vectorf), + link_pose: wp.array(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + link_vel: wp.array(dtype=wp.spatial_vectorf), +): + """Compute root link velocity from root center-of-mass velocity. + + This kernel transforms the root COM velocity into link-frame velocity by projecting + the angular velocity contribution from the COM offset. + + Args: + com_vel: Input array of root COM spatial velocities. Shape is (num_envs,). + link_pose: Input array of root link poses in world frame. Shape is (num_envs,). + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + Only the first body (index 0) is used for the root. + link_vel: Output array where root link velocities are written. Shape is (num_envs,). + """ + i = wp.tid() + link_vel[i] = get_link_vel_from_root_com_vel_func(com_vel[i], link_pose[i], body_com_pos_b[i, 0]) + + +@wp.kernel +def get_root_com_pose_from_root_link_pose( + link_pose: wp.array(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + com_pose_w: wp.array(dtype=wp.transformf), +): + """Compute root COM pose from root link pose. + + This kernel transforms the root link pose to the root COM pose using the body COM offset. + + Args: + link_pose: Input array of root link poses in world frame. Shape is (num_envs,). + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + Only the first body (index 0) is used for the root. + com_pose_w: Output array where root COM poses are written. Shape is (num_envs,). + """ + i = wp.tid() + com_pose_w[i] = get_com_pose_from_link_pose_func(link_pose[i], body_com_pos_b[i, 0]) + + +@wp.kernel +def concat_root_pose_and_vel_to_state( + pose: wp.array(dtype=wp.transformf), + vel: wp.array(dtype=wp.spatial_vectorf), + state: wp.array(dtype=vec13f), +): + """Concatenate root pose and velocity into a 13-element state vector. + + This kernel combines a 7-element pose (pos + quat) and a 6-element velocity + (angular + linear) into a single 13-element state vector. + + Args: + pose: Input array of root poses in world frame. Shape is (num_envs,). + vel: Input array of root spatial velocities. Shape is (num_envs,). + state: Output array where concatenated state vectors are written. Shape is (num_envs,). + """ + i = wp.tid() + state[i] = concat_pose_and_vel_to_state_func(pose[i], vel[i]) + + +@wp.kernel +def split_state_to_root_pose_and_vel( + state: wp.array2d(dtype=wp.float32), + pose: wp.array(dtype=wp.transformf), + vel: wp.array(dtype=wp.spatial_vectorf), +): + """Split a 13-element state vector into root pose and velocity. + + This kernel extracts a 7-element pose (pos + quat) and a 6-element velocity + (angular + linear) from a 13-element state vector. + + Args: + state: Input array of root states. Shape is (num_envs, 13). + pose: Output array where root poses are written. Shape is (num_envs,). + vel: Output array where root spatial velocities are written. Shape is (num_envs,). + """ + i = wp.tid() + # Extract pose: [pos(3), quat(4)] = state[0:7] + pose[i] = wp.transform( + wp.vec3f(state[i, 0], state[i, 1], state[i, 2]), wp.quatf(state[i, 3], state[i, 4], state[i, 5], state[i, 6]) + ) + # Extract velocity: [ang_vel(3), lin_vel(3)] = state[7:13] + vel[i] = wp.spatial_vector( + wp.vec3f(state[i, 7], state[i, 8], state[i, 9]), # angular velocity + wp.vec3f(state[i, 10], state[i, 11], state[i, 12]), # linear velocity + ) + + +""" +Body-level @wp.kernel (2D — used by Articulation + RigidObjectCollection). +""" + + +@wp.kernel +def get_body_link_vel_from_body_com_vel( + body_com_vel: wp.array2d(dtype=wp.spatial_vectorf), + body_link_pose: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + body_link_vel: wp.array2d(dtype=wp.spatial_vectorf), +): + """Compute body link velocities from body COM velocities for all bodies. + + This kernel transforms COM velocities into link-frame velocities by projecting + the angular velocity contribution from the COM offset, for each body in each environment. + + Args: + body_com_vel: Input array of body COM spatial velocities. Shape is (num_envs, num_bodies). + body_link_pose: Input array of body link poses in world frame. Shape is (num_envs, num_bodies). + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + body_link_vel: Output array where body link velocities are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + body_link_vel[i, j] = get_link_vel_from_root_com_vel_func( + body_com_vel[i, j], body_link_pose[i, j], body_com_pos_b[i, j] + ) + + +@wp.kernel +def get_body_com_pose_from_body_link_pose( + body_link_pose: wp.array2d(dtype=wp.transformf), + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + body_com_pose_w: wp.array2d(dtype=wp.transformf), +): + """Compute body COM poses from body link poses for all bodies. + + This kernel transforms link poses to COM poses using the body COM offset in the body frame. + + Args: + body_link_pose: Input array of body link poses in world frame. Shape is (num_envs, num_bodies). + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + body_com_pose_w: Output array where body COM poses in world frame are written. + Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + body_com_pose_w[i, j] = get_com_pose_from_link_pose_func(body_link_pose[i, j], body_com_pos_b[i, j]) + + +@wp.kernel +def concat_body_pose_and_vel_to_state( + pose: wp.array2d(dtype=wp.transformf), + vel: wp.array2d(dtype=wp.spatial_vectorf), + state: wp.array2d(dtype=vec13f), +): + """Concatenate body pose and velocity into 13-element state vectors for all bodies. + + This kernel combines a 7-element pose (pos + quat) and a 6-element velocity + (angular + linear) into a single 13-element state vector, for each body in each environment. + + Args: + pose: Input array of body poses in world frame. Shape is (num_envs, num_bodies). + vel: Input array of body spatial velocities. Shape is (num_envs, num_bodies). + state: Output array where concatenated state vectors are written. + Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + state[i, j] = concat_pose_and_vel_to_state_func(pose[i, j], vel[i, j]) + + +""" +Derived property kernels. +""" + + +@wp.kernel +def quat_apply_inverse_1D_kernel( + gravity: wp.array(dtype=wp.vec3f), + quat: wp.array(dtype=wp.quatf), + projected_gravity: wp.array(dtype=wp.vec3f), +): + """Apply inverse quaternion rotation to gravity vectors (1D). + + This kernel rotates gravity vectors into the local frame of each environment + using the inverse of the provided quaternion. + + Args: + gravity: Input array of gravity vectors in world frame. Shape is (num_envs,). + quat: Input array of quaternions representing orientations. Shape is (num_envs,). + projected_gravity: Output array where projected gravity vectors are written. + Shape is (num_envs,). + """ + i = wp.tid() + projected_gravity[i] = wp.quat_rotate_inv(quat[i], gravity[i]) + + +@wp.kernel +def root_heading_w( + forward_vec: wp.array(dtype=wp.vec3f), + quat: wp.array(dtype=wp.quatf), + heading_w: wp.array(dtype=wp.float32), +): + """Compute root heading angle in the world frame. + + This kernel computes the heading angle (yaw) by rotating the forward vector + by the root quaternion and computing atan2 of the resulting x and y components. + + Args: + forward_vec: Input array of forward direction vectors. Shape is (num_envs,). + quat: Input array of root quaternions. Shape is (num_envs,). + heading_w: Output array where heading angles (radians) are written. Shape is (num_envs,). + """ + i = wp.tid() + heading_w[i] = compute_heading_w_func(forward_vec[i], quat[i]) + + +@wp.kernel +def quat_apply_inverse_2D_kernel( + vec: wp.array2d(dtype=wp.vec3f), + quat: wp.array2d(dtype=wp.quatf), + result: wp.array2d(dtype=wp.vec3f), +): + """Apply inverse quaternion rotation to vectors (2D). + + This kernel rotates vectors into the local frame of each body in each environment + using the inverse of the provided quaternion. + + Args: + vec: Input array of vectors in world frame. Shape is (num_envs, num_bodies). + quat: Input array of quaternions representing orientations. Shape is (num_envs, num_bodies). + result: Output array where rotated vectors are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + result[i, j] = wp.quat_rotate_inv(quat[i, j], vec[i, j]) + + +@wp.kernel +def body_heading_w( + forward_vec: wp.array2d(dtype=wp.vec3f), + quat: wp.array2d(dtype=wp.quatf), + heading_w: wp.array2d(dtype=wp.float32), +): + """Compute body heading angles in the world frame for all bodies. + + This kernel computes heading angles (yaw) by rotating forward vectors + by body quaternions and computing atan2 of the resulting x and y components. + + Args: + forward_vec: Input array of forward direction vectors. Shape is (num_envs, num_bodies). + quat: Input array of body quaternions. Shape is (num_envs, num_bodies). + heading_w: Output array where heading angles (radians) are written. + Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + heading_w[i, j] = compute_heading_w_func(forward_vec[i, j], quat[i, j]) + + +""" +Root-level write kernels (1D — used by RigidObject + Articulation). +""" + + +@wp.kernel +def set_root_link_pose_to_sim_index( + data: wp.array(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + root_link_pose_w: wp.array(dtype=wp.transformf), + root_link_state_w: wp.array(dtype=vec13f), + root_state_w: wp.array(dtype=vec13f), +): + """Write root link pose data to simulation buffers. + + This kernel writes root link poses from the input array to the output buffers + and optionally updates the corresponding state vectors. + + Args: + data: Input array of root link poses. Shape is (num_selected_envs,). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + root_link_pose_w: Output array where root link poses are written. Shape is (num_envs,). + root_link_state_w: Output array where root link states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + root_state_w: Output array where root states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + """ + i = wp.tid() + root_link_pose_w[env_ids[i]] = data[i] + if root_link_state_w: + root_link_state_w[env_ids[i]] = set_state_transforms_func(root_link_state_w[env_ids[i]], data[i]) + if root_state_w: + root_state_w[env_ids[i]] = set_state_transforms_func(root_state_w[env_ids[i]], data[i]) + + +@wp.kernel +def set_root_link_pose_to_sim_mask( + data: wp.array(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + root_link_pose_w: wp.array(dtype=wp.transformf), + root_link_state_w: wp.array(dtype=vec13f), + root_state_w: wp.array(dtype=vec13f), +): + """Write root link pose data to simulation buffers. + + This kernel writes root link poses from the input array to the output buffers + and optionally updates the corresponding state vectors. + + Args: + data: Input array of root link poses. Shape is (num_instances,). + env_mask: Input array of environment mask. Shape is (num_instances,). + root_link_pose_w: Output array where root link poses are written. Shape is (num_envs,). + root_link_state_w: Output array where root link states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + root_state_w: Output array where root states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + """ + i = wp.tid() + if env_mask[i]: + root_link_pose_w[i] = data[i] + if root_link_state_w: + root_link_state_w[i] = set_state_transforms_func(root_link_state_w[i], data[i]) + if root_state_w: + root_state_w[i] = set_state_transforms_func(root_state_w[i], data[i]) + + +@wp.kernel +def set_root_com_pose_to_sim_index( + data: wp.array(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + root_com_pose_w: wp.array(dtype=wp.transformf), + root_link_pose_w: wp.array(dtype=wp.transformf), + root_com_state_w: wp.array(dtype=vec13f), + root_link_state_w: wp.array(dtype=vec13f), + root_state_w: wp.array(dtype=vec13f), +): + """Write root COM pose data to simulation buffers. + + This kernel writes root COM poses from the input array to the output buffers, + computes the corresponding link pose from the COM pose, and optionally updates + the corresponding state vectors. + + Args: + data: Input array of root COM poses. Shape is (num_selected_envs,). + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). Only the first body (index 0) is used for the root. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + root_com_pose_w: Output array where root COM poses are written. Shape is (num_envs,). + root_link_pose_w: Output array where root link poses (derived from COM) are written. + Shape is (num_envs,). + root_com_state_w: Output array where root COM states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + root_link_state_w: Output array where root link states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + root_state_w: Output array where root states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + """ + i = wp.tid() + root_com_pose_w[env_ids[i]] = data[i] + if root_com_state_w: + root_com_state_w[env_ids[i]] = set_state_transforms_func(root_com_state_w[env_ids[i]], data[i]) + # Get the com pose in the link frame + root_link_pose_w[env_ids[i]] = get_com_pose_in_link_frame_func( + root_com_pose_w[env_ids[i]], body_com_pose_b[env_ids[i], 0] + ) + if root_link_state_w: + root_link_state_w[env_ids[i]] = set_state_transforms_func( + root_link_state_w[env_ids[i]], root_link_pose_w[env_ids[i]] + ) + if root_state_w: + root_state_w[env_ids[i]] = set_state_transforms_func(root_state_w[env_ids[i]], root_link_pose_w[env_ids[i]]) + + +@wp.kernel +def set_root_com_pose_to_sim_mask( + data: wp.array(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + root_com_pose_w: wp.array(dtype=wp.transformf), + root_link_pose_w: wp.array(dtype=wp.transformf), + root_com_state_w: wp.array(dtype=vec13f), + root_link_state_w: wp.array(dtype=vec13f), + root_state_w: wp.array(dtype=vec13f), +): + """Write root COM pose data to simulation buffers. + + This kernel writes root COM poses from the input array to the output buffers, + computes the corresponding link pose from the COM pose, and optionally updates + the corresponding state vectors. + + Args: + data: Input array of root COM poses. Shape is (num_instances,). + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). Only the first body (index 0) is used for the root. + env_mask: Input array of environment mask. Shape is (num_instances,). + root_com_pose_w: Output array where root COM poses are written. Shape is (num_envs,). + root_link_pose_w: Output array where root link poses (derived from COM) are written. + Shape is (num_envs,). + root_com_state_w: Output array where root COM states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + root_link_state_w: Output array where root link states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + root_state_w: Output array where root states are updated (pose portion). + Shape is (num_envs,). Can be None if not needed. + """ + i = wp.tid() + if env_mask[i]: + root_com_pose_w[i] = data[i] + if root_com_state_w: + root_com_state_w[i] = set_state_transforms_func(root_com_state_w[i], data[i]) + # Get the com pose in the link frame + root_link_pose_w[i] = get_com_pose_in_link_frame_func(root_com_pose_w[i], body_com_pose_b[i, 0]) + if root_link_state_w: + root_link_state_w[i] = set_state_transforms_func(root_link_state_w[i], root_link_pose_w[i]) + if root_state_w: + root_state_w[i] = set_state_transforms_func(root_state_w[i], root_link_pose_w[i]) + + +@wp.kernel +def set_root_com_velocity_to_sim_index( + data: wp.array(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.int32), + num_bodies: wp.int32, + root_com_velocity_w: wp.array(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), + root_state_w: wp.array(dtype=vec13f), + root_com_state_w: wp.array(dtype=vec13f), +): + """Write root COM velocity data to simulation buffers. + + This kernel writes root COM velocities from the input array to the output buffers, + optionally updates the corresponding state vectors, and zeros out the body + acceleration buffer to prevent reporting stale values. + + Args: + data: Input array of root COM spatial velocities. Shape is (num_selected_envs,). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + num_bodies: Input scalar number of bodies per environment. + root_com_velocity_w: Output array where root COM velocities are written. Shape is (num_envs,). + body_acc_w: Output array where body accelerations are zeroed. Shape is + (num_envs, num_bodies). + root_state_w: Output array where root states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + root_com_state_w: Output array where root COM states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + """ + i = wp.tid() + root_com_velocity_w[env_ids[i]] = data[i] + if root_state_w: + root_state_w[env_ids[i]] = set_state_velocities_func(root_state_w[env_ids[i]], data[i]) + if root_com_state_w: + root_com_state_w[env_ids[i]] = set_state_velocities_func(root_com_state_w[env_ids[i]], data[i]) + # Make the acceleration zero to prevent reporting old values + for j in range(num_bodies): + body_acc_w[env_ids[i], j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def set_root_com_velocity_to_sim_mask( + data: wp.array(dtype=wp.spatial_vectorf), + env_mask: wp.array(dtype=wp.bool), + num_bodies: wp.int32, + root_com_velocity_w: wp.array(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), + root_state_w: wp.array(dtype=vec13f), + root_com_state_w: wp.array(dtype=vec13f), +): + """Write root COM velocity data to simulation buffers. + + This kernel writes root COM velocities from the input array to the output buffers, + optionally updates the corresponding state vectors, and zeros out the body + acceleration buffer to prevent reporting stale values. + + Args: + data: Input array of root COM spatial velocities. Shape is (num_instances,). + env_mask: Input array of environment mask. Shape is (num_instances,). + num_bodies: Input scalar number of bodies per environment. + root_com_velocity_w: Output array where root COM velocities are written. Shape is (num_envs,). + body_acc_w: Output array where body accelerations are zeroed. Shape is + (num_envs, num_bodies). + root_state_w: Output array where root states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + root_com_state_w: Output array where root COM states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + """ + i = wp.tid() + if env_mask[i]: + root_com_velocity_w[i] = data[i] + if root_state_w: + root_state_w[i] = set_state_velocities_func(root_state_w[i], data[i]) + if root_com_state_w: + root_com_state_w[i] = set_state_velocities_func(root_com_state_w[i], data[i]) + # Make the acceleration zero to prevent reporting old values + for j in range(num_bodies): + body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def set_root_link_velocity_to_sim_index( + data: wp.array(dtype=wp.spatial_vectorf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + link_pose_w: wp.array(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + num_bodies: wp.int32, + root_link_velocity_w: wp.array(dtype=wp.spatial_vectorf), + root_com_velocity_w: wp.array(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), + root_link_state_w: wp.array(dtype=vec13f), + root_state_w: wp.array(dtype=vec13f), + root_com_state_w: wp.array(dtype=vec13f), +): + """Write root link velocity data to simulation buffers. + + This kernel writes root link velocities from the input array to the output buffers, + computes the corresponding COM velocity from the link velocity, optionally updates + the corresponding state vectors, and zeros out the body acceleration buffer. + + Args: + data: Input array of root link spatial velocities. Shape is (num_selected_envs,). + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). Only the first body (index 0) is used for the root. + link_pose_w: Input array of root link poses in world frame. Shape is (num_envs,). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + num_bodies: Input scalar number of bodies per environment. + root_link_velocity_w: Output array where root link velocities are written. + Shape is (num_envs,). + root_com_velocity_w: Output array where root COM velocities (derived from link) + are written. Shape is (num_envs,). + body_acc_w: Output array where body accelerations are zeroed. + Shape is (num_envs, num_bodies). + root_link_state_w: Output array where root link states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + root_state_w: Output array where root states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + root_com_state_w: Output array where root COM states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + """ + i = wp.tid() + root_link_velocity_w[env_ids[i]] = data[i] + if root_link_state_w: + root_link_state_w[env_ids[i]] = set_state_velocities_func(root_link_state_w[env_ids[i]], data[i]) + # Get the link velocity in the com frame + root_com_velocity_w[env_ids[i]] = get_link_velocity_in_com_frame_func( + root_link_velocity_w[env_ids[i]], link_pose_w[env_ids[i]], body_com_pose_b[env_ids[i], 0] + ) + if root_com_state_w: + root_com_state_w[env_ids[i]] = set_state_velocities_func( + root_com_state_w[env_ids[i]], root_com_velocity_w[env_ids[i]] + ) + if root_state_w: + root_state_w[env_ids[i]] = set_state_velocities_func(root_state_w[env_ids[i]], root_com_velocity_w[env_ids[i]]) + # Make the acceleration zero to prevent reporting old values + for j in range(num_bodies): + body_acc_w[env_ids[i], j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def set_root_link_velocity_to_sim_mask( + data: wp.array(dtype=wp.spatial_vectorf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + link_pose_w: wp.array(dtype=wp.transformf), + env_mask: wp.array(dtype=wp.bool), + num_bodies: wp.int32, + root_link_velocity_w: wp.array(dtype=wp.spatial_vectorf), + root_com_velocity_w: wp.array(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), + root_link_state_w: wp.array(dtype=vec13f), + root_state_w: wp.array(dtype=vec13f), + root_com_state_w: wp.array(dtype=vec13f), +): + """Write root link velocity data to simulation buffers. + + This kernel writes root link velocities from the input array to the output buffers, + computes the corresponding COM velocity from the link velocity, optionally updates + the corresponding state vectors, and zeros out the body acceleration buffer. + + Args: + data: Input array of root link spatial velocities. Shape is (num_instances,). + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). Only the first body (index 0) is used for the root. + link_pose_w: Input array of root link poses in world frame. Shape is (num_envs,). + env_mask: Input array of environment mask. Shape is (num_instances,). + num_bodies: Input scalar number of bodies per environment. + root_link_velocity_w: Output array where root link velocities are written. + Shape is (num_envs,). + root_com_velocity_w: Output array where root COM velocities (derived from link) + are written. Shape is (num_envs,). + body_acc_w: Output array where body accelerations are zeroed. + Shape is (num_envs, num_bodies). + root_link_state_w: Output array where root link states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + root_state_w: Output array where root states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + root_com_state_w: Output array where root COM states are updated (velocity portion). + Shape is (num_envs,). Can be None if not needed. + """ + i = wp.tid() + if env_mask[i]: + root_link_velocity_w[i] = data[i] + if root_link_state_w: + root_link_state_w[i] = set_state_velocities_func(root_link_state_w[i], data[i]) + # Get the link velocity in the com frame + root_com_velocity_w[i] = get_link_velocity_in_com_frame_func( + root_link_velocity_w[i], link_pose_w[i], body_com_pose_b[i, 0] + ) + if root_com_state_w: + root_com_state_w[i] = set_state_velocities_func(root_com_state_w[i], root_com_velocity_w[i]) + if root_state_w: + root_state_w[i] = set_state_velocities_func(root_state_w[i], root_com_velocity_w[i]) + # Make the acceleration zero to prevent reporting old values + for j in range(num_bodies): + body_acc_w[i, j] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +""" +Body-level write kernels (2D — used by RigidObjectCollection). +""" + + +@wp.kernel +def set_body_link_pose_to_sim( + data: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + from_mask: bool, + body_link_pose_w: wp.array2d(dtype=wp.transformf), + body_link_state_w: wp.array2d(dtype=vec13f), + body_state_w: wp.array2d(dtype=vec13f), +): + """Write body link pose data to simulation buffers. + + This kernel writes body link poses from the input array to the output buffers + and optionally updates the corresponding state vectors, for each body in each environment. + + Args: + data: Input array of body link poses. Shape is (num_envs, num_bodies) or + (num_selected_envs, num_selected_bodies) depending on from_mask. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + from_mask: Input flag indicating whether to use masked indexing. + body_link_pose_w: Output array where body link poses are written. + Shape is (num_envs, num_bodies). + body_link_state_w: Output array where body link states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_state_w: Output array where body states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + """ + i, j = wp.tid() + if from_mask: + body_link_pose_w[env_ids[i], body_ids[j]] = data[env_ids[i], body_ids[j]] + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_link_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + else: + body_link_pose_w[env_ids[i], body_ids[j]] = data[i, j] + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_link_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + + +@wp.kernel +def set_body_com_pose_to_sim( + data: wp.array2d(dtype=wp.transformf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + from_mask: bool, + body_com_pose_w: wp.array2d(dtype=wp.transformf), + body_link_pose_w: wp.array2d(dtype=wp.transformf), + body_com_state_w: wp.array2d(dtype=vec13f), + body_link_state_w: wp.array2d(dtype=vec13f), + body_state_w: wp.array2d(dtype=vec13f), +): + """Write body COM pose data to simulation buffers. + + This kernel writes body COM poses from the input array to the output buffers, + computes the corresponding link poses from the COM poses, and optionally updates + the corresponding state vectors, for each body in each environment. + + Args: + data: Input array of body COM poses. Shape is (num_envs, num_bodies) or + (num_selected_envs, num_selected_bodies) depending on from_mask. + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + from_mask: Input flag indicating whether to use masked indexing. + body_com_pose_w: Output array where body COM poses are written. + Shape is (num_envs, num_bodies). + body_link_pose_w: Output array where body link poses (derived from COM) are written. + Shape is (num_envs, num_bodies). + body_com_state_w: Output array where body COM states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_link_state_w: Output array where body link states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_state_w: Output array where body states are updated (pose portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + """ + i, j = wp.tid() + if from_mask: + body_com_pose_w[env_ids[i], body_ids[j]] = data[env_ids[i], body_ids[j]] + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_com_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + else: + body_com_pose_w[env_ids[i], body_ids[j]] = data[i, j] + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_com_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + # Get the link pose from com pose + body_link_pose_w[env_ids[i], body_ids[j]] = get_com_pose_in_link_frame_func( + body_com_pose_w[env_ids[i], body_ids[j]], body_com_pose_b[env_ids[i], body_ids[j]] + ) + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_link_state_w[env_ids[i], body_ids[j]], body_link_pose_w[env_ids[i], body_ids[j]] + ) + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_transforms_func( + body_state_w[env_ids[i], body_ids[j]], body_link_pose_w[env_ids[i], body_ids[j]] + ) + + +@wp.kernel +def set_body_com_velocity_to_sim( + data: wp.array2d(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + from_mask: bool, + body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), + body_state_w: wp.array2d(dtype=vec13f), + body_com_state_w: wp.array2d(dtype=vec13f), +): + """Write body COM velocity data to simulation buffers. + + This kernel writes body COM velocities from the input array to the output buffers, + optionally updates the corresponding state vectors, and zeros out the body + acceleration buffer, for each body in each environment. + + Args: + data: Input array of body COM spatial velocities. Shape is (num_envs, num_bodies) or + (num_selected_envs, num_selected_bodies) depending on from_mask. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + from_mask: Input flag indicating whether to use masked indexing. + body_com_velocity_w: Output array where body COM velocities are written. + Shape is (num_envs, num_bodies). + body_acc_w: Output array where body accelerations are zeroed. + Shape is (num_envs, num_bodies). + body_state_w: Output array where body states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_com_state_w: Output array where body COM states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + """ + i, j = wp.tid() + if from_mask: + body_com_velocity_w[env_ids[i], body_ids[j]] = data[env_ids[i], body_ids[j]] + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_com_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + else: + body_com_velocity_w[env_ids[i], body_ids[j]] = data[i, j] + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_com_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + # Make the acceleration zero to prevent reporting old values + body_acc_w[env_ids[i], body_ids[j]] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +@wp.kernel +def set_body_link_velocity_to_sim( + data: wp.array2d(dtype=wp.spatial_vectorf), + body_com_pose_b: wp.array2d(dtype=wp.transformf), + body_link_pose_w: wp.array2d(dtype=wp.transformf), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + from_mask: bool, + body_link_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_com_velocity_w: wp.array2d(dtype=wp.spatial_vectorf), + body_acc_w: wp.array2d(dtype=wp.spatial_vectorf), + body_link_state_w: wp.array2d(dtype=vec13f), + body_state_w: wp.array2d(dtype=vec13f), + body_com_state_w: wp.array2d(dtype=vec13f), +): + """Write body link velocity data to simulation buffers. + + This kernel writes body link velocities from the input array to the output buffers, + computes the corresponding COM velocities from the link velocities, optionally updates + the corresponding state vectors, and zeros out the body acceleration buffer. + + Args: + data: Input array of body link spatial velocities. Shape is (num_envs, num_bodies) + or (num_selected_envs, num_selected_bodies) depending on from_mask. + body_com_pose_b: Input array of body COM poses in body frame. Shape is + (num_envs, num_bodies). + body_link_pose_w: Input array of body link poses in world frame. Shape is + (num_envs, num_bodies). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + from_mask: Input flag indicating whether to use masked indexing. + body_link_velocity_w: Output array where body link velocities are written. + Shape is (num_envs, num_bodies). + body_com_velocity_w: Output array where body COM velocities (derived from link) + are written. Shape is (num_envs, num_bodies). + body_acc_w: Output array where body accelerations are zeroed. + Shape is (num_envs, num_bodies). + body_link_state_w: Output array where body link states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_state_w: Output array where body states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + body_com_state_w: Output array where body COM states are updated (velocity portion). + Shape is (num_envs, num_bodies). Can be None if not needed. + """ + i, j = wp.tid() + if from_mask: + body_link_velocity_w[env_ids[i], body_ids[j]] = data[env_ids[i], body_ids[j]] + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_link_state_w[env_ids[i], body_ids[j]], data[env_ids[i], body_ids[j]] + ) + else: + body_link_velocity_w[env_ids[i], body_ids[j]] = data[i, j] + if body_link_state_w: + body_link_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_link_state_w[env_ids[i], body_ids[j]], data[i, j] + ) + # Get the link velocity in the com frame + body_com_velocity_w[env_ids[i], body_ids[j]] = get_link_velocity_in_com_frame_func( + body_link_velocity_w[env_ids[i], body_ids[j]], + body_link_pose_w[env_ids[i], body_ids[j]], + body_com_pose_b[env_ids[i], body_ids[j]], + ) + if body_com_state_w: + body_com_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_com_state_w[env_ids[i], body_ids[j]], body_com_velocity_w[env_ids[i], body_ids[j]] + ) + if body_state_w: + body_state_w[env_ids[i], body_ids[j]] = set_state_velocities_func( + body_state_w[env_ids[i], body_ids[j]], body_com_velocity_w[env_ids[i], body_ids[j]] + ) + # Make the acceleration zero to prevent reporting old values + body_acc_w[env_ids[i], body_ids[j]] = wp.spatial_vectorf(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + + +""" +Generic buffer-writing kernels (used by Articulation + RigidObject + RigidObjectCollection). +""" + + +@wp.kernel +def write_2d_data_to_buffer_with_indices( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + joint_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.float32), +): + """Write 2D float data to a buffer at specified indices. + + This kernel copies float data from an input array to an output buffer at the specified + environment and joint/body indices. + + Args: + in_data: Input array containing float data. Shape is (num_selected_envs, num_selected_joints). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + joint_ids: Input array of joint/body indices to write to. Shape is (num_selected_joints,). + out_data: Output array where data is written. Shape is (num_envs, num_joints). + """ + i, j = wp.tid() + out_data[env_ids[i], joint_ids[j]] = in_data[i, j] + + +@wp.kernel +def write_2d_data_to_buffer_with_mask( + in_data: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), + out_data: wp.array2d(dtype=wp.float32), +): + """Write 2D float data to a buffer at specified indices. + + This kernel copies float data from an input array to an output buffer at the specified + environment and joint/body indices. + + Args: + in_data: Input array containing float data. Shape is (num_instances, num_joints). + env_mask: Input array of environment mask. Shape is (num_instances,). + joint_mask: Input array of joint/body mask. Shape is (num_instances, num_joints). + out_data: Output array where data is written. Shape is (num_instances, num_joints). + """ + i, j = wp.tid() + if env_mask[i] and joint_mask[j]: + out_data[i, j] = in_data[i, j] + + +@wp.kernel +def write_body_inertia_to_buffer_index( + in_data: wp.array3d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + out_data: wp.array3d(dtype=wp.float32), +): + """Write body inertia data to a buffer at specified indices. + + This kernel copies 3x3 inertia tensor data (stored as 9 floats) from an input array + to an output buffer at the specified environment and body indices. + + Args: + in_data: Input array containing inertia data. Shape is (num_selected_envs, num_selected_bodies, 9). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + out_data: Output array where inertia data is written. Shape is (num_envs, num_bodies, 9). + """ + i, j = wp.tid() + for k in range(9): + out_data[env_ids[i], body_ids[j], k] = in_data[i, j, k] + + +@wp.kernel +def write_body_inertia_to_buffer_mask( + in_data: wp.array3d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + out_data: wp.array3d(dtype=wp.float32), +): + """Write body inertia data to a buffer at specified indices. + + This kernel copies 3x3 inertia tensor data (stored as 9 floats) from an input array + to an output buffer at the specified environment and body indices. + + Args: + in_data: Input array containing inertia data. Shape is (num_selected_envs, num_selected_bodies, 9). + env_mask: Input array of environment mask. Shape is (num_selected_envs,). + body_mask: Input array of body mask. Shape is (num_selected_bodies,). + out_data: Output array where inertia data is written. Shape is (num_envs, num_bodies, 9). + """ + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + for k in range(9): + out_data[i, j, k] = in_data[i, j, k] + + +@wp.kernel +def write_single_body_inertia_to_buffer( + in_data: wp.array2d(dtype=wp.float32), + env_ids: wp.array(dtype=wp.int32), + from_mask: bool, + out_data: wp.array2d(dtype=wp.float32), +): + """Write body inertia data to a buffer at specified indices. + + This kernel copies 3x3 inertia tensor data (stored as 9 floats) from an input array + to an output buffer at the specified environment and body indices. + + Args: + in_data: Input array containing inertia data. Shape is (num_envs, 9) or + (num_selected_envs, 9) depending on from_mask. + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + from_mask: Input flag indicating whether to use masked indexing. + out_data: Output array where inertia data is written. Shape is (num_envs, 9). + """ + i = wp.tid() + if from_mask: + for k in range(9): + out_data[env_ids[i], k] = in_data[env_ids[i], k] + else: + for k in range(9): + out_data[env_ids[i], k] = in_data[i, k] + + +@wp.kernel +def write_body_com_position_to_buffer_index( + in_data: wp.array2d(dtype=wp.vec3f), + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + out_data: wp.array2d(dtype=wp.vec3f), +): + """Write body COM position data to a buffer at specified indices. + + This kernel copies body COM position data from an input array to an output buffer at the + specified environment and body indices. + + Args: + in_data: Input array containing body COM positions. Shape is (num_selected_envs, num_selected_bodies). + env_ids: Input array of environment indices to write to. Shape is (num_selected_envs,). + body_ids: Input array of body indices to write to. Shape is (num_selected_bodies,). + out_data: Output array where body COM positions are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + out_data[env_ids[i], body_ids[j]] = in_data[i, j] + + +@wp.kernel +def write_body_com_position_to_buffer_mask( + in_data: wp.array2d(dtype=wp.vec3f), + env_mask: wp.array(dtype=wp.bool), + body_mask: wp.array(dtype=wp.bool), + out_data: wp.array2d(dtype=wp.vec3f), +): + """Write body COM position data to a buffer at specified masks. + + This kernel copies body COM position data from an input array to an output buffer at the + specified environment and body masks. + + Args: + in_data: Input array containing body COM positions. Shape is (num_instances, num_bodies). + env_mask: Input array of environment mask. Shape is (num_instances,). + body_mask: Input array of body mask. Shape is (num_bodies). + out_data: Output array where body COM positions are written. Shape is (num_instances, num_bodies). + """ + i, j = wp.tid() + if env_mask[i] and body_mask[j]: + out_data[i, j] = in_data[i, j] + + +@wp.kernel +def split_transform_to_pos_1d( + transform: wp.array(dtype=wp.transformf), + pos: wp.array(dtype=wp.vec3f), +): + """Split a 1D transform array into a position array. + + This kernel splits a 1D transform array into a position array. + + Args: + transform: Input array of transforms. Shape is (num_envs, 7). + pos: Output array where positions are written. Shape is (num_envs, 3). + """ + i = wp.tid() + pos[i] = wp.transform_get_translation(transform[i]) + + +@wp.kernel +def split_transform_to_quat_1d( + transform: wp.array(dtype=wp.transformf), + quat: wp.array(dtype=wp.quatf), +): + """Split a 1D transform array into a quaternion array. + + This kernel splits a 1D transform array into a quaternion array. + + Args: + transform: Input array of transforms. Shape is (num_envs, 7). + quat: Output array where quaternions are written. Shape is (num_envs, 4). + """ + i = wp.tid() + quat[i] = wp.transform_get_rotation(transform[i]) + + +@wp.kernel +def split_transform_to_pos_2d( + transform: wp.array2d(dtype=wp.transformf), + pos: wp.array2d(dtype=wp.vec3f), +): + """Split a 2D transform array into a position array. + + This kernel splits a 2D transform array into a position array. + + Args: + transform: Input array of transforms. Shape is (num_envs, num_bodies, 7). + pos: Output array where positions are written. Shape is (num_envs, num_bodies, 3). + """ + i, j = wp.tid() + pos[i, j] = wp.transform_get_translation(transform[i, j]) + + +@wp.kernel +def split_transform_to_quat_2d( + transform: wp.array2d(dtype=wp.transformf), + quat: wp.array2d(dtype=wp.quatf), +): + """Split a 2D transform array into a quaternion array. + + This kernel splits a 2D transform array into a quaternion array. + + Args: + transform: Input array of transforms. Shape is (num_envs, num_bodies, 7). + quat: Output array where quaternions are written. Shape is (num_envs, num_bodies, 4). + """ + i, j = wp.tid() + quat[i, j] = wp.transform_get_rotation(transform[i, j]) + + +@wp.kernel +def split_spatial_vector_to_top_1d( + spatial_vector: wp.array(dtype=wp.spatial_vectorf), + top_part: wp.array(dtype=wp.vec3f), +): + """Split a 1D spatial vector array into a top part array. + + This kernel splits a 1D spatial vector array into a top part array. + + Args: + spatial_vector: Input array of spatial vectors. Shape is (num_envs, 6). + top_part: Output array where top parts are written. Shape is (num_envs, 3). + """ + i = wp.tid() + top_part[i] = wp.spatial_top(spatial_vector[i]) + + +@wp.kernel +def split_spatial_vector_to_bottom_1d( + spatial_vector: wp.array(dtype=wp.spatial_vectorf), + bottom_part: wp.array(dtype=wp.vec3f), +): + """Split a 1D spatial vector array into a bottom part array. + + This kernel splits a 1D spatial vector array into a bottom part array. + + Args: + spatial_vector: Input array of spatial vectors. Shape is (num_envs, 6). + bottom_part: Output array where bottom parts are written. Shape is (num_envs, 3). + """ + i = wp.tid() + bottom_part[i] = wp.spatial_bottom(spatial_vector[i]) + + +@wp.kernel +def split_spatial_vector_to_top_2d( + spatial_vector: wp.array2d(dtype=wp.spatial_vectorf), + top_part: wp.array2d(dtype=wp.vec3f), +): + """Split a 2D spatial vector array into a top part array. + + This kernel splits a 2D spatial vector array into a top part array. + + Args: + spatial_vector: Input array of spatial vectors. Shape is (num_envs, num_bodies, 6). + top_part: Output array where top parts are written. Shape is (num_envs, num_bodies, 3). + """ + i, j = wp.tid() + top_part[i, j] = wp.spatial_top(spatial_vector[i, j]) + + +@wp.kernel +def split_spatial_vector_to_bottom_2d( + spatial_vector: wp.array2d(dtype=wp.spatial_vectorf), + bottom_part: wp.array2d(dtype=wp.vec3f), +): + """Split a 2D spatial vector array into a bottom part array. + + This kernel splits a 2D spatial vector array into a bottom part array. + + Args: + spatial_vector: Input array of spatial vectors. Shape is (num_envs, num_bodies, 6). + bottom_part: Output array where bottom parts are written. Shape is (num_envs, num_bodies, 3). + """ + i, j = wp.tid() + bottom_part[i, j] = wp.spatial_bottom(spatial_vector[i, j]) + + +@wp.kernel +def make_dummy_body_com_pose_b( + body_com_pos_b: wp.array2d(dtype=wp.vec3f), + body_com_pose_b: wp.array2d(dtype=wp.transformf), +): + """Make a dummy body COM pose in body frame. + + This kernel makes a dummy body COM pose in body frame. + + Args: + body_com_pos_b: Input array of body COM positions in body frame. Shape is (num_envs, num_bodies). + body_com_pose_b: Output array where body COM poses are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + # Concatenate the position and a unit quaternion + body_com_pose_b[i, j] = wp.transformf(body_com_pos_b[i, j], wp.quatf(0.0, 0.0, 0.0, 1.0)) + + +@wp.kernel +def derive_body_acceleration_from_body_com_velocities( + body_com_vel: wp.array2d(dtype=wp.spatial_vectorf), + dt: wp.float32, + prev_body_com_vel: wp.array2d(dtype=wp.spatial_vectorf), + body_acc: wp.array2d(dtype=wp.spatial_vectorf), +): + """Derive body acceleration from body COM velocities. + + This kernel derives body acceleration from body COM velocities using finite differencing. + + Args: + body_com_vel: Input array of body COM velocities. Shape is (num_envs, num_bodies). + dt: Input time step (scalar) used for finite differencing. + prev_body_com_vel: Input/output array of previous body COM velocities. Shape is (num_envs, num_bodies). + body_acc: Output array where body accelerations are written. Shape is (num_envs, num_bodies). + """ + i, j = wp.tid() + # Compute the acceleration + body_acc[i, j] = (body_com_vel[i, j] - prev_body_com_vel[i, j]) / dt + # Update the previous body COM velocity + prev_body_com_vel[i, j] = body_com_vel[i, j] + + +@wp.kernel +def update_wrench_array_with_force_and_torque( + forces: wp.array2d(dtype=wp.vec3f), + torques: wp.array2d(dtype=wp.vec3f), + wrench: wp.array2d(dtype=wp.spatial_vectorf), + env_ids: wp.array(dtype=wp.bool), + body_ids: wp.array(dtype=wp.bool), +): + env_index, body_index = wp.tid() + if env_ids[env_index] and body_ids[body_index]: + wrench[env_index, body_index] = update_wrench_with_force_and_torque( + forces[env_index, body_index], + torques[env_index, body_index], + ) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/__init__.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/__init__.py new file mode 100644 index 00000000000..fb04f3b5982 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .rigid_object import RigidObject +from .rigid_object_data import RigidObjectData + +__all__ = ["RigidObject", "RigidObjectData"] diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py new file mode 100644 index 00000000000..a107d6d0dbc --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -0,0 +1,1183 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import logging +import warnings +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import numpy as np +import torch +import warp as wp +from newton.selection import ArticulationView +from newton.solvers import SolverNotifyFlags + +from pxr import UsdPhysics + +import isaaclab.sim as sim_utils +import isaaclab.utils.string as string_utils +from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject +from isaaclab.utils.wrench_composer import WrenchComposer + +from isaaclab_newton.assets import kernels as shared_kernels +from isaaclab_newton.physics import NewtonManager as SimulationManager + +from .rigid_object_data import RigidObjectData + +if TYPE_CHECKING: + from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg + +# import logger +logger = logging.getLogger(__name__) + + +class RigidObject(BaseRigidObject): + """A rigid object asset class. + + Rigid objects are assets comprising of rigid bodies. They can be used to represent dynamic objects + such as boxes, spheres, etc. A rigid body is described by its pose, velocity and mass distribution. + + For an asset to be considered a rigid object, the root prim of the asset must have the `USD RigidBodyAPI`_ + applied to it. This API is used to define the simulation properties of the rigid body. On playing the + simulation, the physics engine will automatically register the rigid body and create a corresponding + rigid body handle. This handle can be accessed using the :attr:`root_view` attribute. + + .. note:: + + For users familiar with Isaac Sim, the PhysX view class API is not the exactly same as Isaac Sim view + class API. Similar to Isaac Lab, Isaac Sim wraps around the PhysX view API. However, as of now (2023.1 release), + we see a large difference in initializing the view classes in Isaac Sim. This is because the view classes + in Isaac Sim perform additional USD-related operations which are slow and also not required. + + .. _`USD RigidBodyAPI`: https://openusd.org/dev/api/class_usd_physics_rigid_body_a_p_i.html + """ + + cfg: RigidObjectCfg + """Configuration instance for the rigid object.""" + + __backend_name__: str = "newton" + """The name of the backend for the rigid object.""" + + def __init__(self, cfg: RigidObjectCfg): + """Initialize the rigid object. + + Args: + cfg: A configuration instance. + """ + super().__init__(cfg) + + """ + Properties + """ + + @property + def data(self) -> RigidObjectData: + return self._data + + @property + def num_instances(self) -> int: + return self.root_view.count + + @property + def num_bodies(self) -> int: + """Number of bodies in the asset. + + This is always 1 since each object is a single rigid body. + """ + return 1 + + @property + def body_names(self) -> list[str]: + """Ordered names of bodies in the rigid object.""" + return self.root_view.link_names + + @property + def root_view(self) -> ArticulationView: + """Root view for the asset. + + .. note:: + Use this view with caution. It requires handling of tensors in a specific way. + """ + return self._root_view + + @property + def instantaneous_wrench_composer(self) -> WrenchComposer: + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ + return self._instantaneous_wrench_composer + + @property + def permanent_wrench_composer(self) -> WrenchComposer: + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ + return self._permanent_wrench_composer + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: + """Reset the rigid object. + + Args: + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve all indices + if (env_ids is None) or (env_ids == slice(None)): + env_ids = slice(None) + # reset external wrench + self._instantaneous_wrench_composer.reset(env_ids) + self._permanent_wrench_composer.reset(env_ids) + + def write_data_to_sim(self) -> None: + """Write external wrench to the simulation. + + .. note:: + We write external wrench to the simulation here since this function is called before the simulation step. + This ensures that the external wrench is applied at every simulation step. + """ + # write external wrench + if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: + if self._instantaneous_wrench_composer.active: + # Compose instantaneous wrench with permanent wrench + self._instantaneous_wrench_composer.add_forces_and_torques( + forces=self._permanent_wrench_composer.composed_force, + torques=self._permanent_wrench_composer.composed_torque, + body_ids=self._ALL_BODY_INDICES, + env_ids=self._ALL_INDICES, + ) + # Apply both instantaneous and permanent wrench to the simulation + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + self._instantaneous_wrench_composer.composed_force, + self._instantaneous_wrench_composer.composed_torque, + self._data._sim_bind_body_external_wrench, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) + else: + # Apply permanent wrench to the simulation + wp.launch( + shared_kernels.update_wrench_array_with_force_and_torque, + dim=(self.num_instances, self.num_bodies), + device=self.device, + inputs=[ + self._permanent_wrench_composer.composed_force, + self._permanent_wrench_composer.composed_torque, + self._data._sim_bind_body_external_wrench, + self._ALL_ENV_MASK, + self._ALL_BODY_MASK, + ], + ) + self._instantaneous_wrench_composer.reset() + + def update(self, dt: float) -> None: + """Updates the simulation data. + + Args: + dt: The time step size in seconds. + """ + self.data.update(dt) + + """ + Operations - Finders. + """ + + def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = False) -> tuple[list[int], list[str]]: + """Find bodies in the rigid body based on the name keys. + + Please check the :meth:`isaaclab.utils.string_utils.resolve_matching_names` function for more + information on the name matching. + + Args: + name_keys: A regular expression or a list of regular expressions to match the body names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple of lists containing the body indices and names. + """ + return string_utils.resolve_matching_names(name_keys, self.body_names, preserve_order) + + """ + Operations - Write to simulation. + """ + + def write_root_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) + + def write_root_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root pose over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + self.write_root_link_pose_to_sim_mask(root_pose=root_pose, env_mask=env_mask) + + def write_root_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) + + def write_root_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment mask into the simulation. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + self.write_root_com_velocity_to_sim_mask(root_velocity=root_velocity, env_mask=env_mask) + + def write_root_link_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7), + or (len(env_ids),) / (num_instances,) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.set_root_link_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_pose, + env_ids, + ], + outputs=[ + self.data.root_link_pose_w, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + ], + device=self.device, + ) + # Need to invalidate the buffer to trigger the update with the new state. + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + + def write_root_link_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root link pose over selected environment mask into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + + wp.launch( + shared_kernels.set_root_link_pose_to_sim_mask, + dim=root_pose.shape[0], + inputs=[ + root_pose, + env_mask, + ], + outputs=[ + self.data.root_link_pose_w, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + ], + device=self.device, + ) + # Need to invalidate the buffer to trigger the update with the new state. + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + + def write_root_com_pose_to_sim_index( + self, + *, + root_pose: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass pose over selected environment indices into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7), + or (len(env_ids),) / (num_instances,) with dtype wp.transformf. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + # Note: we are doing a single launch for faster performance. Prior versions would call + # write_root_link_pose_to_sim after this. + wp.launch( + shared_kernels.set_root_com_pose_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_pose, + self.data.body_com_pose_b, + env_ids, + ], + outputs=[ + self.data._root_com_pose_w.data, + self.data.root_link_pose_w, + None, # self.data._root_com_state_w.data, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._root_com_pose_w.timestamp = self.data._sim_timestamp + # Need to invalidate the buffer to trigger the update with the new state. + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + + def write_root_com_pose_to_sim_mask( + self, + *, + root_pose: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass pose over selected environment mask into the simulation. + + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + wp.launch( + shared_kernels.set_root_com_pose_to_sim_mask, + dim=root_pose.shape[0], + inputs=[ + root_pose, + env_mask, + ], + outputs=[ + self.data._root_com_pose_w.data, + self.data.root_link_pose_w, + None, # self.data._root_com_state_w.data, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._root_com_pose_w.timestamp = self.data._sim_timestamp + # Need to invalidate the buffer to trigger the update with the new state. + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + if self.data._root_link_state_w is not None: + self.data._root_link_state_w.timestamp = -1.0 + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + + def write_root_com_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. + Shape is (len(env_ids), 6) or (num_instances, 6), + or (len(env_ids),) / (num_instances,) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_velocity, + env_ids, + 1, + ], + outputs=[ + self.data.root_com_vel_w, + self.data._body_com_acc_w.data, + None, # self.data._root_state_w.data, + None, # self.data._root_com_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._body_com_acc_w.timestamp = self.data._sim_timestamp + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + + def write_root_com_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root center of mass velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + wp.launch( + shared_kernels.set_root_com_velocity_to_sim_mask, + dim=root_velocity.shape[0], + inputs=[ + root_velocity, + env_mask, + 1, + ], + outputs=[ + self.data.root_com_vel_w, + self.data._body_com_acc_w.data, + None, # self.data._root_state_w.data, + None, # self.data._root_com_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._body_com_acc_w.timestamp = self.data._sim_timestamp + if self.data._root_state_w is not None: + self.data._root_state_w.timestamp = -1.0 + if self.data._root_com_state_w is not None: + self.data._root_com_state_w.timestamp = -1.0 + + def write_root_link_velocity_to_sim_index( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set the root link velocity over selected environment indices into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. + + .. note:: + This method expects partial data or full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root frame velocities in simulation world frame. + Shape is (len(env_ids), 6) or (num_instances, 6), + or (len(env_ids),) / (num_instances,) with dtype wp.spatial_vectorf. + env_ids: Environment indices. If None, then all indices are used. + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + # Note: we are doing a single launch for faster performance. Prior versions would do multiple launches. + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_index, + dim=env_ids.shape[0], + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_ids, + 1, + ], + outputs=[ + self.data._root_link_vel_w.data, + self.data.root_com_vel_w, + self.data._body_com_acc_w.data, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + None, # self.data._root_com_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._root_link_vel_w.timestamp = self.data._sim_timestamp + self.data._body_com_acc_w.timestamp = self.data._sim_timestamp + self.data._root_link_state_w.timestamp = -1.0 + self.data._root_state_w.timestamp = -1.0 + self.data._root_com_state_w.timestamp = -1.0 + + def write_root_link_velocity_to_sim_mask( + self, + *, + root_velocity: torch.Tensor | wp.array, + env_mask: wp.array | None = None, + ) -> None: + """Set the root link velocity over selected environment mask into the simulation. + + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + if env_mask is None: + env_mask = self._ALL_ENV_MASK + wp.launch( + shared_kernels.set_root_link_velocity_to_sim_mask, + dim=root_velocity.shape[0], + inputs=[ + root_velocity, + self.data.body_com_pose_b, + self.data.root_link_pose_w, + env_mask, + 1, + ], + outputs=[ + self.data._root_link_vel_w.data, + self.data.root_com_vel_w, + self.data._body_com_acc_w.data, + None, # self.data._root_link_state_w.data, + None, # self.data._root_state_w.data, + None, # self.data._root_com_state_w.data, + ], + device=self.device, + ) + # Update the timestamps + self.data._root_link_vel_w.timestamp = self.data._sim_timestamp + self.data._body_com_acc_w.timestamp = self.data._sim_timestamp + self.data._root_link_state_w.timestamp = -1.0 + self.data._root_state_w.timestamp = -1.0 + self.data._root_com_state_w.timestamp = -1.0 + + """ + Operations - Setters. + """ + + def set_masses_index( + self, + *, + masses: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set masses of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + masses: Masses of all bodies. Shape is (len(env_ids), len(body_ids)). + body_ids: The body indices to set the masses for. Defaults to None (all bodies). + env_ids: The environment indices to set the masses for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_indices, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + masses, + env_ids, + body_ids, + ], + outputs=[ + self.data.body_mass, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_masses_mask( + self, + *, + masses: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set masses of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + masses: Masses of all bodies. Shape is (num_instances, num_bodies). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + wp.launch( + shared_kernels.write_2d_data_to_buffer_with_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + masses, + env_mask, + body_mask, + ], + outputs=[ + self.data.body_mass, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_coms_index( + self, + *, + coms: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set center of mass position of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + .. caution:: + Unlike the PhysX version of this method, this method does not set the center of mass orientation. + Only the position is set. This is because Newton considers the center of mass orientation to always be + aligned with the body frame. + + Args: + coms: Center of mass position of all bodies. Shape is (len(env_ids), len(body_ids), 3). + body_ids: The body indices to set the center of mass pose for. Defaults to None (all bodies). + env_ids: The environment indices to set the center of mass pose for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_body_com_position_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + coms, + env_ids, + body_ids, + ], + outputs=[ + self.data.body_com_pos_b, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_coms_mask( + self, + *, + coms: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set center of mass position of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + .. caution:: + Unlike the PhysX version of this method, this method does not set the center of mass orientation. + Only the position is set. This is because Newton considers the center of mass orientation to always be + aligned with the body frame. + + Args: + coms: Center of mass position of all bodies. Shape is (num_instances, num_bodies, 3). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + wp.launch( + shared_kernels.write_body_com_position_to_buffer_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + coms, + env_mask, + body_mask, + ], + outputs=[ + self.data.body_com_pos_b, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_inertias_index( + self, + *, + inertias: torch.Tensor | wp.array, + body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Set inertias of all bodies using indices. + + .. note:: + This method expects partial data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + inertias: Inertias of all bodies. Shape is (len(env_ids), len(body_ids), 9). + body_ids: The body indices to set the inertias for. Defaults to None (all bodies). + env_ids: The environment indices to set the inertias for. Defaults to None (all environments). + """ + # resolve all indices + env_ids = self._resolve_env_ids(env_ids) + body_ids = self._resolve_body_ids(body_ids) + # Warp kernels can ingest torch tensors directly, so we don't need to convert to warp arrays here. + wp.launch( + shared_kernels.write_body_inertia_to_buffer_index, + dim=(env_ids.shape[0], body_ids.shape[0]), + inputs=[ + inertias, + env_ids, + body_ids, + ], + outputs=[ + self.data.body_inertia, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + def set_inertias_mask( + self, + *, + inertias: torch.Tensor | wp.array, + body_mask: wp.array | None = None, + env_mask: wp.array | None = None, + ) -> None: + """Set inertias of all bodies using masks. + + .. note:: + This method expects full data. + + .. tip:: + Both the index and mask methods have dedicated optimized implementations. Performance is similar for both. + However, to allow graphed pipelines, the mask method must be used. + + Args: + inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). + body_mask: Body mask. If None, then all bodies are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # resolve masks + if env_mask is None: + env_mask = self._ALL_ENV_MASK + if body_mask is None: + body_mask = self._ALL_BODY_MASK + wp.launch( + shared_kernels.write_body_inertia_to_buffer_mask, + dim=(env_mask.shape[0], body_mask.shape[0]), + inputs=[ + inertias, + env_mask, + body_mask, + ], + outputs=[ + self.data.body_inertia, + ], + device=self.device, + ) + # tell the physics engine that some of the body properties have been updated + SimulationManager.add_model_change(SolverNotifyFlags.BODY_INERTIAL_PROPERTIES) + + """ + Internal helper. + """ + + def _initialize_impl(self): + # obtain global simulation view + self._physics_sim_view = SimulationManager.get_physics_sim_view() + # obtain the first prim in the regex expression (all others are assumed to be a copy of this) + template_prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) + if template_prim is None: + raise RuntimeError(f"Failed to find prim for expression: '{self.cfg.prim_path}'.") + template_prim_path = template_prim.GetPath().pathString + + # find rigid root prims + root_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI), + traverse_instance_prims=False, + ) + if len(root_prims) == 0: + raise RuntimeError( + f"Failed to find a rigid body when resolving '{self.cfg.prim_path}'." + " Please ensure that the prim has 'USD RigidBodyAPI' applied." + ) + if len(root_prims) > 1: + raise RuntimeError( + f"Failed to find a single rigid body when resolving '{self.cfg.prim_path}'." + f" Found multiple '{root_prims}' under '{template_prim_path}'." + " Please ensure that there is only one rigid body in the prim path tree." + ) + + articulation_prims = sim_utils.get_all_matching_child_prims( + template_prim_path, + predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI), + traverse_instance_prims=False, + ) + if len(articulation_prims) != 0: + if articulation_prims[0].GetAttribute("physxArticulation:articulationEnabled").Get(): + raise RuntimeError( + f"Found an articulation root when resolving '{self.cfg.prim_path}' for rigid objects. These are" + f" located at: '{articulation_prims}' under '{template_prim_path}'. Please disable the articulation" + " root in the USD or from code by setting the parameter" + " 'ArticulationRootPropertiesCfg.articulation_enabled' to False in the spawn configuration." + ) + + # resolve root prim back into regex expression + root_prim_path = root_prims[0].GetPath().pathString + root_prim_path_expr = self.cfg.prim_path + root_prim_path[len(template_prim_path) :] + # -- object view + self._root_view = self._physics_sim_view.create_rigid_body_view(root_prim_path_expr.replace(".*", "*")) + + # check if the rigid body was created + if self.root_view._backend is None: + raise RuntimeError(f"Failed to create rigid body at: {self.cfg.prim_path}. Please check PhysX logs.") + + # log information about the rigid body + logger.info(f"Rigid body initialized at: {self.cfg.prim_path} with root '{root_prim_path_expr}'.") + logger.info(f"Number of instances: {self.num_instances}") + logger.info(f"Number of bodies: {self.num_bodies}") + logger.info(f"Body names: {self.body_names}") + + # container for data access + self._data = RigidObjectData(self.root_view, self.device) + + # create buffers + self._create_buffers() + # process configuration + self._process_cfg() + # update the rigid body data + self.update(0.0) + # Let the rigid object data know that it is fully instantiated and ready to use. + self.data.is_primed = True + + def _create_buffers(self): + """Create buffers for storing data.""" + # constants + self._ALL_INDICES = wp.array(np.arange(self.num_instances, dtype=np.int32), device=self.device) + self._ALL_ENV_MASK = wp.ones((self.num_instances,), dtype=wp.bool, device=self.device) + self._ALL_BODY_INDICES = wp.array(np.arange(self.num_bodies, dtype=np.int32), device=self.device) + self._ALL_BODY_MASK = wp.ones((self.num_bodies,), dtype=wp.bool, device=self.device) + + # external wrench composer + self._instantaneous_wrench_composer = WrenchComposer(self) + self._permanent_wrench_composer = WrenchComposer(self) + + # set information about rigid body into data + self._data.body_names = self.body_names + + def _process_cfg(self) -> None: + """Post processing of configuration parameters.""" + # default state + # -- root state + # note: we cast to tuple to avoid torch/numpy type mismatch. + default_root_pose = tuple(self.cfg.init_state.pos) + tuple(self.cfg.init_state.rot) + default_root_vel = tuple(self.cfg.init_state.lin_vel) + tuple(self.cfg.init_state.ang_vel) + default_root_pose = np.tile(np.array(default_root_pose, dtype=np.float32), (self.num_instances, 1)) + default_root_vel = np.tile(np.array(default_root_vel, dtype=np.float32), (self.num_instances, 1)) + self._data.default_root_pose = wp.array(default_root_pose, dtype=wp.transformf, device=self.device) + self._data.default_root_vel = wp.array(default_root_vel, dtype=wp.spatial_vectorf, device=self.device) + + def _resolve_env_ids(self, env_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + """Resolve environment indices to a warp array or tensor. + + .. note:: + We need to convert torch tensors to warp arrays since the TensorAPI views only support warp arrays. + + Args: + env_ids: Environment indices. If None, then all indices are used. + + Returns: + A warp array of environment indices or a tensor of environment indices. + """ + if (env_ids is None) or (env_ids == slice(None)): + return self._ALL_INDICES + elif isinstance(env_ids, list): + return wp.array(env_ids, dtype=wp.int32, device=self.device) + if isinstance(env_ids, torch.Tensor): + return wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + return env_ids + + def _resolve_body_ids(self, body_ids: Sequence[int] | torch.Tensor | wp.array | None) -> wp.array | torch.Tensor: + """Resolve body indices to a warp array or tensor. + + .. note:: + We do not need to convert torch tensors to warp arrays since they never get passed to the TensorAPI views. + + Args: + body_ids: Body indices. If None, then all indices are used. + + Returns: + A warp array of body indices or a tensor of body indices. + """ + if (body_ids is None) or (body_ids == slice(None)): + return self._ALL_BODY_INDICES + elif isinstance(body_ids, list): + return wp.array(body_ids, dtype=wp.int32, device=self.device) + return body_ids + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + # set all existing views to None to invalidate them + self._root_view = None + + def write_root_state_to_sim(self, root_state: torch.Tensor, env_ids: Sequence[int] | torch.Tensor | None = None): + """Deprecated, same as :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + + def write_root_com_state_to_sim( + self, root_state: torch.Tensor, env_ids: Sequence[int] | torch.Tensor | None = None + ): + """Deprecated, same as :meth:`write_root_com_pose_to_sim_index` and + :meth:`write_root_com_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_com_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_com_pose_to_sim_index' and 'write_root_com_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_com_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + + def write_root_link_state_to_sim( + self, root_state: torch.Tensor, env_ids: Sequence[int] | torch.Tensor | None = None + ): + """Deprecated, same as :meth:`write_root_link_pose_to_sim_index` and + :meth:`write_root_link_velocity_to_sim_index`.""" + warnings.warn( + "The function 'write_root_link_state_to_sim' will be deprecated in a future release. Please" + " use 'write_root_link_pose_to_sim_index' and 'write_root_link_velocity_to_sim_index' instead.", + DeprecationWarning, + stacklevel=2, + ) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py new file mode 100644 index 00000000000..623af53af8a --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object_data.py @@ -0,0 +1,1230 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + +import logging +import warnings +import weakref +from typing import TYPE_CHECKING + +import torch +import warp as wp + +from isaaclab.assets.rigid_object.base_rigid_object_data import BaseRigidObjectData +from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.math import normalize + +from isaaclab_newton.assets import kernels as shared_kernels +from isaaclab_newton.physics import NewtonManager as SimulationManager + +if TYPE_CHECKING: + from newton.selection import ArticulationView + + +# import logger +logger = logging.getLogger(__name__) + + +class RigidObjectData(BaseRigidObjectData): + """Data container for a rigid object. + + This class contains the data for a rigid object in the simulation. The data includes the state of + the root rigid body and the state of all the bodies in the object. The data is stored in the simulation + world frame unless otherwise specified. + + For a rigid body, there are two frames of reference that are used: + + - Actor frame: The frame of reference of the rigid body prim. This typically corresponds to the Xform prim + with the rigid body schema. + - Center of mass frame: The frame of reference of the center of mass of the rigid body. + + Depending on the settings of the simulation, the actor frame and the center of mass frame may be the same. + This needs to be taken into account when interpreting the data. + + The data is lazily updated, meaning that the data is only updated when it is accessed. This is useful + when the data is expensive to compute or retrieve. The data is updated when the timestamp of the buffer + is older than the current simulation timestamp. The timestamp is updated whenever the data is updated. + """ + + __backend_name__: str = "newton" + """The name of the backend for the rigid object data.""" + + def __init__(self, root_view: ArticulationView, device: str): + """Initializes the rigid object data. + + Args: + root_view: The root rigid body view. + device: The device used for processing. + """ + super().__init__(root_view, device) + # Set the root rigid body view + # note: this is stored as a weak reference to avoid circular references between the asset class + # and the data container. This is important to avoid memory leaks. + self._root_view: ArticulationView = weakref.proxy(root_view) + + # Set initial time stamp + self._sim_timestamp = 0.0 + self._is_primed = False + + # Convert to direction vector + gravity = wp.to_torch(SimulationManager.get_model().gravity)[0] + gravity_dir = torch.tensor((gravity[0], gravity[1], gravity[2]), device=self.device) + gravity_dir = normalize(gravity_dir.unsqueeze(0)).squeeze(0) + gravity_dir = gravity_dir.repeat(self._root_view.count, 1) + forward_vec = torch.tensor((1.0, 0.0, 0.0), device=self.device).repeat(self._root_view.count, 1) + + # Initialize constants + self.GRAVITY_VEC_W = wp.from_torch(gravity_dir, dtype=wp.vec3f) + self.FORWARD_VEC_B = wp.from_torch(forward_vec, dtype=wp.vec3f) + + self._create_buffers() + + @property + def is_primed(self) -> bool: + """Whether the rigid object data is fully instantiated and ready to use.""" + return self._is_primed + + @is_primed.setter + def is_primed(self, value: bool) -> None: + """Set whether the rigid object data is fully instantiated and ready to use. + + .. note:: + Once this quantity is set to True, it cannot be changed. + + Args: + value: The primed state. + + Raises: + ValueError: If the rigid object data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object data is already primed.") + self._is_primed = value + + def update(self, dt: float) -> None: + """Updates the data for the rigid object. + + Args: + dt: The time step for the update. This must be a positive value. + """ + # update the simulation timestamp + self._sim_timestamp += dt + # Trigger an update of the body com acceleration buffer at a higher frequency + # since we do finite differencing. + self.body_com_acc_w + + """ + Names. + """ + + body_names: list[str] = None + """Body names in the order parsed by the simulation view.""" + + """ + Defaults. + """ + + @property + def default_root_pose(self) -> wp.array: + """Default root pose ``[pos, quat]`` in local environment frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + The position and quaternion are of the rigid body's actor frame. + """ + return self._default_root_pose + + @default_root_pose.setter + def default_root_pose(self, value: wp.array) -> None: + """Set the default root pose. + + Args: + value: The default root pose. Shape is (num_instances, 7). + + Raises: + ValueError: If the rigid object data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object data is already primed.") + self._default_root_pose.assign(value) + + @property + def default_root_vel(self) -> wp.array: + """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + The linear and angular velocities are of the rigid body's center of mass frame. + """ + return self._default_root_vel + + @default_root_vel.setter + def default_root_vel(self, value: wp.array) -> None: + """Set the default root velocity. + + Args: + value: The default root velocity. Shape is (num_instances, 6). + + Raises: + ValueError: If the rigid object data is already primed. + """ + if self._is_primed: + raise ValueError("The rigid object data is already primed.") + self._default_root_vel.assign(value) + + """ + Root state properties. + """ + + @property + def root_link_pose_w(self) -> wp.array: + """Root link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + This quantity is the pose of the actor frame of the root rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self._sim_bind_root_link_pose_w + + @property + def root_link_vel_w(self) -> wp.array: + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + This quantity contains the linear and angular velocities of the actor frame of the root + rigid body relative to the world. + """ + if self._root_link_vel_w.timestamp < self._sim_timestamp: + # read the CoM velocity and compute link velocity + wp.launch( + shared_kernels.get_root_link_vel_from_root_com_vel, + dim=self._num_instances, + inputs=[ + self.root_com_vel_w, + self.root_link_pose_w, + self.body_com_pose_b, + ], + outputs=[ + self._root_link_vel_w.data, + ], + device=self.device, + ) + self._root_link_vel_w.timestamp = self._sim_timestamp + + return self._root_link_vel_w.data + + @property + def root_com_pose_w(self) -> wp.array: + """Root center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). + This quantity is the pose of the center of mass frame of the root rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + if self._root_com_pose_w.timestamp < self._sim_timestamp: + # apply local transform to center of mass frame + wp.launch( + shared_kernels.get_root_com_pose_from_root_link_pose, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.body_com_pose_b, + ], + outputs=[ + self._root_com_pose_w.data, + ], + device=self.device, + ) + self._root_com_pose_w.timestamp = self._sim_timestamp + + return self._root_com_pose_w.data + + @property + def root_com_vel_w(self) -> wp.array: + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). + This quantity contains the linear and angular velocities of the root rigid body's center of mass frame + relative to the world. + """ + return self._sim_bind_root_com_vel_w + + """ + Body state properties. + """ + + @property + def body_mass(self) -> wp.array: + """Mass of all bodies in the simulation world frame. + + Shape is (num_instances, 1, 1), dtype = wp.float32. + In torch this resolves to (num_instances, 1, 1). + """ + return self._body_mass + + @property + def body_inertia(self) -> wp.array: + """Inertia of all bodies in the simulation world frame. + + Shape is (num_instances, 1, 9), dtype = wp.float32. + In torch this resolves to (num_instances, 1, 9). + """ + return self._body_inertia + + @property + def body_link_pose_w(self) -> wp.array: + """Body link pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). + This quantity is the pose of the actor frame of the rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self._sim_bind_body_link_pose_w + + @property + def body_link_vel_w(self) -> wp.array: + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). + This quantity contains the linear and angular velocities of the actor frame of the root + rigid body relative to the world. + """ + return self.root_link_vel_w.reshape((self._num_instances, 1)) + + @property + def body_com_pose_w(self) -> wp.array: + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). + This quantity is the pose of the center of mass frame of the rigid body relative to the world. + The orientation is provided in (x, y, z, w) format. + """ + return self.root_com_pose_w.reshape((self._num_instances, 1)) + + @property + def body_com_vel_w(self) -> wp.array: + """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). + This quantity contains the linear and angular velocities of the root rigid body's center of mass frame + relative to the world. + """ + return self._sim_bind_body_com_vel_w + + @property + def body_com_acc_w(self) -> wp.array: + """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. + + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). + This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. + """ + if self._body_com_acc_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.derive_body_acceleration_from_body_com_velocities, + dim=(self._num_instances, 1), + device=self.device, + inputs=[ + self._sim_bind_body_com_vel_w, + SimulationManager.get_dt(), + self._previous_body_com_vel, + ], + outputs=[ + self._body_com_acc_w.data, + ], + ) + # set the buffer data and timestamp + self._body_com_acc_w.timestamp = self._sim_timestamp + # update the previous velocity + return self._body_com_acc_w.data + + @property + def body_com_pos_b(self) -> wp.array: + """Center of mass position of all of the bodies in their respective link frames. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the center of mass location relative to its body's link frame. + """ + return self._sim_bind_body_com_pos_b + + @property + def body_com_pose_b(self) -> wp.array: + """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. + + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). + This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. + The orientation is provided in (x, y, z, w) format. + """ + warnings.warn( + "In Newton, body com pose always has unit quaternion. Consider using body_com_pos_b instead." + "Querying this property requires appending a unit quaternion to the position which is expensive.", + category=UserWarning, + stacklevel=2, + ) + if self._body_com_pose_b.timestamp < self._sim_timestamp: + # set the buffer data and timestamp + wp.launch( + shared_kernels.make_dummy_body_com_pose_b, + dim=(self._num_instances, 1), + inputs=[ + self.body_com_pos_b, + ], + outputs=[ + self._body_com_pose_b.data, + ], + device=self.device, + ) + self._body_com_pose_b.timestamp = self._sim_timestamp + return self._body_com_pose_b.data + + """ + Derived Properties. + """ + + @property + def projected_gravity_b(self) -> wp.array: + """Projection of the gravity direction on base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ + if self._projected_gravity_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.GRAVITY_VEC_W, self.root_link_quat_w], + outputs=[self._projected_gravity_b.data], + device=self.device, + ) + self._projected_gravity_b.timestamp = self._sim_timestamp + return self._projected_gravity_b.data + + @property + def heading_w(self) -> wp.array: + """Yaw heading of the base frame (in radians). + + Shape is (num_instances,), dtype = wp.float32. In torch this resolves to (num_instances,). + + .. note:: + This quantity is computed by assuming that the forward-direction of the base + frame is along x-direction, i.e. :math:`(1, 0, 0)`. + """ + if self._heading_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.root_heading_w, + dim=self._num_instances, + inputs=[self.FORWARD_VEC_B, self.root_link_quat_w], + outputs=[self._heading_w.data], + device=self.device, + ) + self._heading_w.timestamp = self._sim_timestamp + return self._heading_w.data + + @property + def root_link_lin_vel_b(self) -> wp.array: + """Root link linear velocity in base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_link_lin_vel_w, self.root_link_quat_w], + outputs=[self._root_link_lin_vel_b.data], + device=self.device, + ) + self._root_link_lin_vel_b.timestamp = self._sim_timestamp + return self._root_link_lin_vel_b.data + + @property + def root_link_ang_vel_b(self) -> wp.array: + """Root link angular velocity in base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the + rigid body's actor frame. + """ + if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_link_ang_vel_w, self.root_link_quat_w], + outputs=[self._root_link_ang_vel_b.data], + device=self.device, + ) + self._root_link_ang_vel_b.timestamp = self._sim_timestamp + return self._root_link_ang_vel_b.data + + @property + def root_com_lin_vel_b(self) -> wp.array: + """Root center of mass linear velocity in base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_com_lin_vel_w, self.root_link_quat_w], + outputs=[self._root_com_lin_vel_b.data], + device=self.device, + ) + self._root_com_lin_vel_b.timestamp = self._sim_timestamp + return self._root_com_lin_vel_b.data + + @property + def root_com_ang_vel_b(self) -> wp.array: + """Root center of mass angular velocity in base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the + rigid body's actor frame. + """ + if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.quat_apply_inverse_1D_kernel, + dim=self._num_instances, + inputs=[self.root_com_ang_vel_w, self.root_link_quat_w], + outputs=[self._root_com_ang_vel_b.data], + device=self.device, + ) + self._root_com_ang_vel_b.timestamp = self._sim_timestamp + return self._root_com_ang_vel_b.data + + """ + Sliced properties. + """ + + @property + def root_link_pos_w(self) -> wp.array: + """Root link position in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the position of the actor frame of the root rigid body relative to the world. + """ + return self._get_pos_from_transform(self._root_link_pos_w, self.root_link_pose_w) + + @property + def root_link_quat_w(self) -> wp.array: + """Root link orientation (x, y, z, w) in simulation world frame. + + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + This quantity is the orientation of the actor frame of the root rigid body. + """ + return self._get_quat_from_transform(self._root_link_quat_w, self.root_link_pose_w) + + @property + def root_link_lin_vel_w(self) -> wp.array: + """Root linear velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the linear velocity of the root rigid body's actor frame relative to the world. + """ + return self._get_top_from_spatial_vector(self._root_link_lin_vel_w, self.root_link_vel_w) + + @property + def root_link_ang_vel_w(self) -> wp.array: + """Root link angular velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. + """ + return self._get_bottom_from_spatial_vector(self._root_link_ang_vel_w, self.root_link_vel_w) + + @property + def root_com_pos_w(self) -> wp.array: + """Root center of mass position in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the position of the center of mass frame of the root rigid body relative to the world. + """ + return self._get_pos_from_transform(self._root_com_pos_w, self.root_com_pose_w) + + @property + def root_com_quat_w(self) -> wp.array: + """Root center of mass orientation (x, y, z, w) in simulation world frame. + + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. + """ + return self._get_quat_from_transform(self._root_com_quat_w, self.root_com_pose_w) + + @property + def root_com_lin_vel_w(self) -> wp.array: + """Root center of mass linear velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. + """ + return self._get_top_from_spatial_vector(self._root_com_lin_vel_w, self.root_com_vel_w) + + @property + def root_com_ang_vel_w(self) -> wp.array: + """Root center of mass angular velocity in simulation world frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. + """ + return self._get_bottom_from_spatial_vector(self._root_com_ang_vel_w, self.root_com_vel_w) + + @property + def body_link_pos_w(self) -> wp.array: + """Positions of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the position of the rigid bodies' actor frame relative to the world. + """ + return self._get_pos_from_transform(self._body_link_pos_w, self.body_link_pose_w) + + @property + def body_link_quat_w(self) -> wp.array: + """Orientation (x, y, z, w) of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the rigid bodies' actor frame relative to the world. + """ + return self._get_quat_from_transform(self._body_link_quat_w, self.body_link_pose_w) + + @property + def body_link_lin_vel_w(self) -> wp.array: + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. + """ + return self._get_top_from_spatial_vector(self._body_link_lin_vel_w, self.body_link_vel_w) + + @property + def body_link_ang_vel_w(self) -> wp.array: + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. + """ + return self._get_bottom_from_spatial_vector(self._body_link_ang_vel_w, self.body_link_vel_w) + + @property + def body_com_pos_w(self) -> wp.array: + """Positions of all bodies' center of mass in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the position of the rigid bodies' center of mass frame. + """ + return self._get_pos_from_transform(self._body_com_pos_w, self.body_com_pose_w) + + @property + def body_com_quat_w(self) -> wp.array: + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the principal axes of inertia of the rigid bodies. + """ + return self._get_quat_from_transform(self._body_com_quat_w, self.body_com_pose_w) + + @property + def body_com_lin_vel_w(self) -> wp.array: + """Linear velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the linear velocity of the rigid bodies' center of mass frame. + """ + return self._get_top_from_spatial_vector(self._body_com_lin_vel_w, self.body_com_vel_w) + + @property + def body_com_ang_vel_w(self) -> wp.array: + """Angular velocity of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the angular velocity of the rigid bodies' center of mass frame. + """ + return self._get_bottom_from_spatial_vector(self._body_com_ang_vel_w, self.body_com_vel_w) + + @property + def body_com_lin_acc_w(self) -> wp.array: + """Linear acceleration of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the linear acceleration of the rigid bodies' center of mass frame. + """ + return self._get_top_from_spatial_vector(self._body_com_lin_acc_w, self.body_com_acc_w) + + @property + def body_com_ang_acc_w(self) -> wp.array: + """Angular acceleration of all bodies in simulation world frame. + + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the angular acceleration of the rigid bodies' center of mass frame. + """ + return self._get_bottom_from_spatial_vector(self._body_com_ang_acc_w, self.body_com_acc_w) + + @property + def body_com_quat_b(self) -> wp.array: + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their + respective link frames. + + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. + """ + return self._get_quat_from_transform(self._body_com_quat_b, self.body_com_pose_b) + + def _create_buffers(self) -> None: + super()._create_buffers() + # Initialize the lazy buffers. + # -- link frame w.r.t. world frame + self._root_link_pose_w = TimestampedBuffer((self._num_instances), self.device, wp.transformf) + self._root_link_vel_w = TimestampedBuffer((self._num_instances), self.device, wp.spatial_vectorf) + # -- com frame w.r.t. link frame + self._body_com_pose_b = TimestampedBuffer((self._num_instances, 1), self.device, wp.transformf) + # -- com frame w.r.t. world frame + self._root_com_pose_w = TimestampedBuffer((self._num_instances), self.device, wp.transformf) + self._root_com_vel_w = TimestampedBuffer((self._num_instances), self.device, wp.spatial_vectorf) + self._body_com_acc_w = TimestampedBuffer((self._num_instances, 1), self.device, wp.spatial_vectorf) + # -- combined state (these are cached as they concatenate) + self._root_state_w = TimestampedBuffer((self._num_instances), self.device, shared_kernels.vec13f) + self._root_link_state_w = TimestampedBuffer((self._num_instances), self.device, shared_kernels.vec13f) + self._root_com_state_w = TimestampedBuffer((self._num_instances), self.device, shared_kernels.vec13f) + # -- derived properties (these are cached to avoid repeated memory allocations) + self._projected_gravity_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + self._heading_w = TimestampedBuffer((self._num_instances), self.device, wp.float32) + self._root_link_lin_vel_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + self._root_link_ang_vel_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + self._root_com_lin_vel_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + self._root_com_ang_vel_b = TimestampedBuffer((self._num_instances), self.device, wp.vec3f) + + # -- Default state + self._default_root_pose = wp.zeros((self._num_instances), dtype=wp.transformf, device=self.device) + self._default_root_vel = wp.zeros((self._num_instances), dtype=wp.spatial_vectorf, device=self.device) + self._default_root_state = None + + # -- Body properties + self._body_mass = wp.clone(self._root_view.get_masses(), device=self.device) + self._body_inertia = wp.clone(self._root_view.get_inertias(), device=self.device) + + def _create_simulation_bindings(self) -> None: + """Create simulation bindings for the root data. + + Direct simulation bindings are pointers to the simulation data, their data is not copied, and should + only be updated using warp kernels. Any modifications made to the bindings will be reflected in the simulation. + Hence we encourage users to carefully think about the data they modify and in which order it should be updated. + + .. caution:: This is possible if and only if the properties that we access are strided from newton and not + indexed. Newton willing this is the case all the time, but we should pay attention to this if things look off. + """ + # -- root properties + if self._root_view.is_fixed_base: + self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[ + :, 0, 0 + ] + else: + self._sim_bind_root_link_pose_w = self._root_view.get_root_transforms(SimulationManager.get_state_0())[:, 0] + self._sim_bind_root_com_vel_w = self._root_view.get_root_velocities(SimulationManager.get_state_0()) + if self._sim_bind_root_com_vel_w is not None: + if self._root_view.is_fixed_base: + self._sim_bind_root_com_vel_w = self._sim_bind_root_com_vel_w[:, 0, 0] + else: + self._sim_bind_root_com_vel_w = self._sim_bind_root_com_vel_w[:, 0] + # -- body properties + self._sim_bind_body_com_pos_b = self._root_view.get_attribute("body_com", SimulationManager.get_model())[:, 0] + self._sim_bind_body_link_pose_w = self._root_view.get_link_transforms(SimulationManager.get_state_0())[:, 0] + self._sim_bind_body_com_vel_w = self._root_view.get_link_velocities(SimulationManager.get_state_0())[:, 0] + self._sim_bind_body_mass = self._root_view.get_attribute("body_mass", SimulationManager.get_model())[:, 0] + self._sim_bind_body_inertia = self._root_view.get_attribute("body_inertia", SimulationManager.get_model())[:, 0] + self._sim_bind_body_external_wrench = self._root_view.get_attribute("body_f", SimulationManager.get_state_0())[ + :, 0 + ] + + def _create_buffers(self) -> None: + """Create buffers for the root data.""" + self._num_instances = self._root_view.count + # Initialize history for finite differencing. If the rigid object is fixed, the root com velocity is not + # available, so we use zeros. + if self._root_view.get_root_velocities(SimulationManager.get_state_0()) is None: + logger.warning( + "Failed to get root com velocity. If the rigid object is fixed, this is expected." + "Setting root com velocity to zeros." + ) + self._sim_bind_root_com_vel_w = wp.zeros( + (self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._sim_bind_body_com_vel_w = wp.zeros( + (self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + # -- default root pose and velocity + self._default_root_pose = wp.zeros((self._num_instances,), dtype=wp.transformf, device=self.device) + self._default_root_vel = wp.zeros((self._num_instances,), dtype=wp.spatial_vectorf, device=self.device) + + # Initialize history for finite differencing + self._previous_body_com_vel = wp.clone(self._root_view.get_link_velocities(SimulationManager.get_state_0()))[ + :, 0 + ] + + # Initialize the lazy buffers. + # -- link frame w.r.t. world frame + self._root_link_vel_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._root_link_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._projected_gravity_b = TimestampedBuffer(shape=(self._num_instances,), dtype=wp.vec3f, device=self.device) + self._heading_w = TimestampedBuffer(shape=(self._num_instances,), dtype=wp.float32, device=self.device) + self._body_link_vel_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + # -- com frame w.r.t. world frame + self._root_com_pose_w = TimestampedBuffer(shape=(self._num_instances,), dtype=wp.transformf, device=self.device) + self._root_com_vel_b = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._root_com_acc_w = TimestampedBuffer( + shape=(self._num_instances,), dtype=wp.spatial_vectorf, device=self.device + ) + self._body_com_acc_w = TimestampedBuffer( + shape=(self._num_instances, 1), dtype=wp.spatial_vectorf, device=self.device + ) + # Empty memory pre-allocations + self._root_state_w = None + self._root_link_state_w = None + self._root_com_state_w = None + self._body_com_quat_b = None + self._root_link_lin_vel_b = None + self._root_link_ang_vel_b = None + self._root_com_lin_vel_b = None + self._root_com_ang_vel_b = None + self._root_link_pos_w = None + self._root_link_quat_w = None + self._root_link_lin_vel_w = None + self._root_link_ang_vel_w = None + self._root_com_pos_w = None + self._root_com_quat_w = None + self._root_com_lin_vel_w = None + self._root_com_ang_vel_w = None + self._body_link_pos_w = None + self._body_link_quat_w = None + self._body_link_lin_vel_w = None + self._body_link_ang_vel_w = None + self._body_com_pos_w = None + self._body_com_quat_w = None + self._body_com_lin_vel_w = None + self._body_com_ang_vel_w = None + self._body_com_lin_acc_w = None + self._body_com_ang_acc_w = None + self._body_com_pose_b = None + + """ + Internal helpers. + """ + + def _get_pos_from_transform(self, source: wp.array | None, transform: wp.array) -> wp.array: + """Generates a position array from a transform array. + + Args: + transform: The transform array. Shape is (N) dtype=wp.transformf. + + Returns: + The position array. Shape is (N) dtype=wp.vec3f. + """ + # Check if we already created the lazy buffer. + if source is None: + if transform.is_contiguous(): + # Check if the array is contiguous. If so, we can just return a strided array. + # Then this update becomes a no-op. + return wp.array( + ptr=transform.ptr, + shape=transform.shape, + dtype=wp.vec3f, + strides=transform.strides, + device=self.device, + ) + else: + # If the array is no contiguous, we need to create a new array to write to. + source = wp.zeros((transform.shape[0], 3), dtype=wp.vec3f, device=self.device) + + # If the array is not contiguous, we need to launch the kernel to get the position part of the transform. + if not transform.is_contiguous(): + # Launch the right kernel based on the shape of the source array. + if len(source.shape) > 1: + wp.launch( + shared_kernels.split_transform_to_pos_2d, + dim=source.shape, + inputs=[ + source, + ], + outputs=[ + source, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.split_transform_to_pos_1d, + dim=source.shape, + inputs=[ + source, + ], + outputs=[ + source, + ], + device=self.device, + ) + return source + + def _get_quat_from_transform(self, source: wp.array | None, transform: wp.array) -> wp.array: + """Generates a quaternion array from a transform array. + + Args: + transform: The transform array. Shape is (N) dtype=wp.transformf. + + Returns: + The quaternion array. Shape is (N) dtype=wp.quatf. + """ + # Check if we already created the lazy buffer. + if source is None: + if transform.is_contiguous(): + # Check if the array is contiguous. If so, we can just return a strided array. + # Then this update becomes a no-op. + return wp.array( + ptr=transform.ptr + 3 * 4, + shape=transform.shape, + dtype=wp.quatf, + strides=transform.strides, + device=self.device, + ) + else: + # If the array is no contiguous, we need to create a new array to write to. + source = wp.zeros((transform.shape[0], 4), dtype=wp.quatf, device=self.device) + + # If the array is not contiguous, we need to launch the kernel to get the quaternion part of the transform. + if not transform.is_contiguous(): + # Launch the right kernel based on the shape of the source array. + if len(source.shape) > 1: + wp.launch( + shared_kernels.split_transform_to_quat_2d, + dim=source.shape, + inputs=[ + source, + ], + outputs=[ + source, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.split_transform_to_quat_1d, + dim=source.shape, + inputs=[ + source, + ], + outputs=[ + source, + ], + device=self.device, + ) + # Return the source array. (no-op if the array is contiguous.) + return source + + def _get_top_from_spatial_vector(self, source: wp.array | None, spatial_vector: wp.array) -> wp.array: + """Gets the top part of a spatial vector array. + + For instance the linear velocity is the top part of a velocity vector. + + Args: + spatial_vector: The spatial vector array. Shape is (N) dtype=wp.spatial_vectorf. + + Returns: + The top part of the spatial vector array. Shape is (N) dtype=wp.vec3f. + """ + # Check if we already created the lazy buffer. + if source is None: + if spatial_vector.is_contiguous(): + # Check if the array is contiguous. If so, we can just return a strided array. + # Then this update becomes a no-op. + return wp.array( + ptr=spatial_vector.ptr, + shape=spatial_vector.shape, + dtype=wp.vec3f, + strides=spatial_vector.strides, + device=self.device, + ) + else: + # If the array is no contiguous, we need to create a new array to write to. + source = wp.zeros((spatial_vector.shape[0], 3), dtype=wp.vec3f, device=self.device) + + # If the array is not contiguous, we need to launch the kernel to get the top part of the spatial vector. + if not spatial_vector.is_contiguous(): + # Launch the right kernel based on the shape of the source array. + if len(source.shape) > 1: + wp.launch( + shared_kernels.split_spatial_vector_to_top_2d, + dim=source.shape, + inputs=[ + source, + ], + outputs=[ + source, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.split_spatial_vector_to_top_1d, + dim=source.shape, + inputs=[ + source, + ], + outputs=[ + source, + ], + device=self.device, + ) + # Return the source array. (no-op if the array is contiguous.) + return source + + def _get_bottom_from_spatial_vector(self, source: wp.array | None, spatial_vector: wp.array) -> wp.array: + """Gets the bottom part of a spatial vector array. + + For instance the angular velocity is the bottom part of a velocity vector. + + Args: + spatial_vector: The spatial vector array. Shape is (N) dtype=wp.spatial_vectorf. + + Returns: + The bottom part of the spatial vector array. Shape is (N) dtype=wp.vec3f. + """ + # Check if we already created the lazy buffer. + if source is None: + if spatial_vector.is_contiguous(): + # Check if the array is contiguous. If so, we can just return a strided array. + # Then this update becomes a no-op. + return wp.array( + ptr=spatial_vector.ptr + 3 * 4, + shape=spatial_vector.shape, + dtype=wp.vec3f, + strides=spatial_vector.strides, + device=self.device, + ) + else: + # If the array is no contiguous, we need to create a new array to write to. + source = wp.zeros((spatial_vector.shape[0], 3), dtype=wp.vec3f, device=self.device) + + # If the array is not contiguous, we need to launch the kernel to get the bottom part of the spatial vector. + if not spatial_vector.is_contiguous(): + # Launch the right kernel based on the shape of the source array. + if len(source.shape) > 1: + wp.launch( + shared_kernels.split_spatial_vector_to_bottom_2d, + dim=source.shape, + inputs=[ + source, + ], + outputs=[ + source, + ], + device=self.device, + ) + else: + wp.launch( + shared_kernels.split_spatial_vector_to_bottom_1d, + dim=source.shape, + inputs=[ + source, + ], + outputs=[ + source, + ], + device=self.device, + ) + # Return the source array. (no-op if the array is contiguous.) + return source + + """ + Deprecated properties. + """ + + @property + def root_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" + warnings.warn( + "The `root_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_state_w.data, + ], + device=self.device, + ) + self._root_state_w.timestamp = self._sim_timestamp + + return self._root_state_w.data + + @property + def root_link_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" + warnings.warn( + "The `root_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_link_vel_w, + ], + outputs=[ + self._root_link_state_w.data, + ], + device=self.device, + ) + self._root_link_state_w.timestamp = self._sim_timestamp + + return self._root_link_state_w.data + + @property + def root_com_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" + warnings.warn( + "The `root_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_com_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_com_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_com_state_w.data, + ], + device=self.device, + ) + self._root_com_state_w.timestamp = self._sim_timestamp + + return self._root_com_state_w.data + + @property + def default_root_state(self) -> wp.array: + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. + + The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities + are of the center of mass frame. Shape is (num_instances, 13). + """ + warnings.warn( + "Reading the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " + "Please use the default_root_pose and default_root_vel properties instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._default_root_state is None: + self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self._default_root_pose, + self._default_root_vel, + ], + outputs=[ + self._default_root_state, + ], + device=self.device, + ) + return self._default_root_state + + @property + def body_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + # Access internal buffer directly to avoid cascading deprecation warnings from root_state_w + if self._root_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_state_w.data, + ], + device=self.device, + ) + self._root_state_w.timestamp = self._sim_timestamp + return self._root_state_w.data.reshape((self._num_instances, 1)) + + @property + def body_link_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + # Access internal buffer directly to avoid cascading deprecation warnings from root_link_state_w + if self._root_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_link_vel_w, + ], + outputs=[ + self._root_link_state_w.data, + ], + device=self.device, + ) + self._root_link_state_w.timestamp = self._sim_timestamp + return self._root_link_state_w.data.reshape((self._num_instances, 1)) + + @property + def body_com_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_com_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_com_state_w.data, + ], + device=self.device, + ) + self._root_com_state_w.timestamp = self._sim_timestamp + return self._root_com_state_w.data.reshape((self._num_instances, 1)) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 002460c9a74..892e20d2937 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -25,7 +25,7 @@ from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values from isaaclab.utils.types import ArticulationActions -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -181,12 +181,23 @@ def root_view(self) -> physx.ArticulationView: @property def instantaneous_wrench_composer(self) -> WrenchComposer: - """Instantaneous wrench composer for the articulation.""" + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ return self._instantaneous_wrench_composer @property def permanent_wrench_composer(self) -> WrenchComposer: - """Permanent wrench composer for the articulation.""" + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ return self._permanent_wrench_composer """ @@ -201,7 +212,7 @@ def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None Args: env_ids: Environment indices. If None, then all indices are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # use ellipses object to skip initial indices. if (env_ids is None) or (env_ids == slice(None)): @@ -227,7 +238,7 @@ def write_data_to_sim(self): if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques( + self._instantaneous_wrench_composer.add_forces_and_torques_index( forces=self._permanent_wrench_composer.composed_force, torques=self._permanent_wrench_composer.composed_torque, body_ids=self._ALL_BODY_INDICES, @@ -362,6 +373,7 @@ def find_spatial_tendons( def write_root_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: @@ -377,13 +389,15 @@ def write_root_pose_to_sim_index( is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ - self.write_root_link_pose_to_sim_index(root_pose, env_ids=env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: @@ -399,13 +413,15 @@ def write_root_pose_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.write_root_link_pose_to_sim_mask(root_pose, env_mask=env_mask) + self.write_root_link_pose_to_sim_mask(root_pose=root_pose, env_mask=env_mask) def write_root_link_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, @@ -422,7 +438,8 @@ def write_root_link_pose_to_sim_index( is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7), + or (len(env_ids),) / (num_instances,) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ @@ -459,6 +476,7 @@ def write_root_link_pose_to_sim_index( def write_root_link_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: @@ -474,8 +492,9 @@ def write_root_link_pose_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -483,10 +502,11 @@ def write_root_link_pose_to_sim_mask( else: env_ids = self._ALL_INDICES # Set full data to True to ensure the the right code path is taken inside the kernel. - self.write_root_link_pose_to_sim_index(root_pose, env_ids=env_ids, full_data=True) + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids, full_data=True) def write_root_com_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, @@ -494,7 +514,7 @@ def write_root_com_pose_to_sim_index( """Set the root center of mass pose over selected environment indices into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). - The orientation is the orientation of the principle axes of inertia. + The orientation is the orientation of the principal axes of inertia. .. note:: This method expect partial data or full data. @@ -504,7 +524,8 @@ def write_root_com_pose_to_sim_index( is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7), + or (len(env_ids),) / (num_instances,) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ @@ -548,13 +569,14 @@ def write_root_com_pose_to_sim_index( def write_root_com_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: """Set the root center of mass pose over selected environment mask into the simulation. The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). - The orientation is the orientation of the principle axes of inertia. + The orientation is the orientation of the principal axes of inertia. .. note:: This method expect full data. @@ -564,8 +586,9 @@ def write_root_com_pose_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -573,17 +596,19 @@ def write_root_com_pose_to_sim_mask( else: env_ids = self._ALL_INDICES # Set full data to True to ensure the the right code path is taken inside the kernel. - self.write_root_com_pose_to_sim_index(root_pose, env_ids=env_ids, full_data=True) + self.write_root_com_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids, full_data=True) def write_root_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root center of mass velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - .. note:: This sets the velocity of the root's center of mass rather than the roots frame. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expect partial data. @@ -593,20 +618,23 @@ def write_root_velocity_to_sim_index( is only supporting indexing, hence masks need to be converted to indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). + root_velocity: Root center of mass velocities in simulation world frame. + Shape is (len(env_ids), 6) or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def write_root_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - .. note:: This sets the velocity of the root's center of mass rather than the roots frame. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expect full data. @@ -616,13 +644,15 @@ def write_root_velocity_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root center of mass velocities in simulation world frame. + Shape is (num_instances, 6) or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ self.write_root_com_velocity_to_sim_mask(root_velocity=root_velocity, env_mask=env_mask) def write_root_com_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, @@ -630,7 +660,8 @@ def write_root_com_velocity_to_sim_index( """Set the root center of mass velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - .. note:: This sets the velocity of the root's center of mass rather than the roots frame. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expect partial data or full data. @@ -641,7 +672,7 @@ def write_root_com_velocity_to_sim_index( Args: root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) or - (num_instances, 6). + (num_instances, 6), or (len(env_ids),) / (num_instances,) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ @@ -675,13 +706,15 @@ def write_root_com_velocity_to_sim_index( def write_root_com_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - .. note:: This sets the velocity of the root's center of mass rather than the roots frame. + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expect full data. @@ -691,8 +724,9 @@ def write_root_com_velocity_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root center of mass velocities in simulation world frame. + Shape is (num_instances, 6) or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -701,10 +735,11 @@ def write_root_com_velocity_to_sim_mask( env_ids = self._ALL_INDICES # Set full data to True to ensure the the right code path is taken inside the kernel. - self.write_root_com_velocity_to_sim_index(root_velocity, env_ids=env_ids, full_data=True) + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids, full_data=True) def write_root_link_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, @@ -712,7 +747,8 @@ def write_root_link_velocity_to_sim_index( """Set the root link velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - .. note:: This sets the velocity of the root's frame rather than the roots center of mass. + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. .. note:: This method expect partial data or full data. @@ -723,7 +759,7 @@ def write_root_link_velocity_to_sim_index( Args: root_velocity: Root frame velocities in simulation world frame. Shape is (len(env_ids), 6) or - (num_instances, 6). + (num_instances, 6), or (len(env_ids),) / (num_instances,) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ @@ -764,13 +800,15 @@ def write_root_link_velocity_to_sim_index( def write_root_link_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: """Set the root link velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - .. note:: This sets the velocity of the root's frame rather than the roots center of mass. + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. .. note:: This method expect full data. @@ -780,8 +818,9 @@ def write_root_link_velocity_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root frame velocities in simulation world frame. + Shape is (num_instances, 6) or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -789,10 +828,11 @@ def write_root_link_velocity_to_sim_mask( else: env_ids = self._ALL_INDICES # Set full data to True to ensure the the right code path is taken inside the kernel. - self.write_root_link_velocity_to_sim_index(root_velocity, env_ids=env_ids, full_data=True) + self.write_root_link_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids, full_data=True) def write_joint_state_to_sim( self, + *, position: torch.Tensor | wp.array, velocity: torch.Tensor | wp.array, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -807,15 +847,16 @@ def write_joint_state_to_sim( stacklevel=2, ) # set into simulation - self.write_joint_position_to_sim_index(position, joint_ids=joint_ids, env_ids=env_ids) - self.write_joint_velocity_to_sim_index(velocity, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) def write_joint_state_to_sim_mask( self, + *, position: torch.Tensor | wp.array, velocity: torch.Tensor | wp.array, - env_mask: wp.array | None = None, joint_mask: wp.array | None = None, + env_mask: wp.array | None = None, ): """Write joint positions and velocities over selected environment mask into the simulation. @@ -829,15 +870,16 @@ def write_joint_state_to_sim_mask( Args: position: Joint positions. Shape is (num_instances, num_joints). velocity: Joint velocities. Shape is (num_instances, num_joints). - env_mask: Environment mask. If None, then all indices are used. joint_mask: Joint mask. If None, then all joints are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # set into simulation - self.write_joint_position_to_sim_mask(position, env_mask=env_mask, joint_mask=joint_mask) - self.write_joint_velocity_to_sim_mask(velocity, env_mask=env_mask, joint_mask=joint_mask) + self.write_joint_position_to_sim_mask(position=position, env_mask=env_mask, joint_mask=joint_mask) + self.write_joint_velocity_to_sim_mask(velocity=velocity, env_mask=env_mask, joint_mask=joint_mask) def write_joint_position_to_sim_index( self, + *, position: torch.Tensor, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -893,6 +935,7 @@ def write_joint_position_to_sim_index( def write_joint_position_to_sim_mask( self, + *, position: torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -909,7 +952,7 @@ def write_joint_position_to_sim_mask( Args: position: Joint positions. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -921,10 +964,11 @@ def write_joint_position_to_sim_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the the right code path is taken inside the kernel. - self.write_joint_position_to_sim_index(position, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids, full_data=True) def write_joint_velocity_to_sim_index( self, + *, velocity: torch.Tensor | wp.array, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -973,6 +1017,7 @@ def write_joint_velocity_to_sim_index( def write_joint_velocity_to_sim_mask( self, + *, velocity: torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -989,7 +1034,7 @@ def write_joint_velocity_to_sim_mask( Args: velocity: Joint velocities. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -1001,7 +1046,7 @@ def write_joint_velocity_to_sim_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the the right code path is taken inside the kernel. - self.write_joint_velocity_to_sim_index(velocity, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids, full_data=True) """ Operations - Simulation Parameters Writers. @@ -1009,6 +1054,7 @@ def write_joint_velocity_to_sim_mask( def write_joint_stiffness_to_sim_index( self, + *, stiffness: torch.Tensor | wp.array | float, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1068,6 +1114,7 @@ def write_joint_stiffness_to_sim_index( def write_joint_stiffness_to_sim_mask( self, + *, stiffness: torch.Tensor | wp.array | float, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1084,7 +1131,7 @@ def write_joint_stiffness_to_sim_mask( Args: stiffness: Joint stiffness. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -1096,10 +1143,13 @@ def write_joint_stiffness_to_sim_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the the right code path is taken inside the kernel. - self.write_joint_stiffness_to_sim_index(stiffness, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.write_joint_stiffness_to_sim_index( + stiffness=stiffness, joint_ids=joint_ids, env_ids=env_ids, full_data=True + ) def write_joint_damping_to_sim_index( self, + *, damping: torch.Tensor | wp.array | float, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1160,6 +1210,7 @@ def write_joint_damping_to_sim_index( def write_joint_damping_to_sim_mask( self, + *, damping: torch.Tensor | wp.array | float, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1176,7 +1227,7 @@ def write_joint_damping_to_sim_mask( Args: damping: Joint damping. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -1188,10 +1239,11 @@ def write_joint_damping_to_sim_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.write_joint_damping_to_sim_index(damping, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.write_joint_damping_to_sim_index(damping=damping, joint_ids=joint_ids, env_ids=env_ids, full_data=True) def write_joint_position_limit_to_sim_index( self, + *, limits: torch.Tensor | wp.array | float, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1259,6 +1311,7 @@ def write_joint_position_limit_to_sim_index( def write_joint_position_limit_to_sim_mask( self, + *, limits: torch.Tensor | wp.array | float, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1276,7 +1329,7 @@ def write_joint_position_limit_to_sim_mask( Args: limits: Joint limits. Shape is (num_instances, num_joints, 2). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). warn_limit_violation: Whether to use warning or info level logging when default joint positions exceed the new limits. Defaults to True. """ @@ -1291,11 +1344,16 @@ def write_joint_position_limit_to_sim_mask( joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_position_limit_to_sim_index( - limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True, warn_limit_violation=warn_limit_violation + limits=limits, + joint_ids=joint_ids, + env_ids=env_ids, + full_data=True, + warn_limit_violation=warn_limit_violation, ) def write_joint_velocity_limit_to_sim_index( self, + *, limits: torch.Tensor | wp.array | float, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1359,6 +1417,7 @@ def write_joint_velocity_limit_to_sim_index( def write_joint_velocity_limit_to_sim_mask( self, + *, limits: torch.Tensor | wp.array | float, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1379,7 +1438,7 @@ def write_joint_velocity_limit_to_sim_mask( Args: limits: Joint max velocity. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -1391,10 +1450,13 @@ def write_joint_velocity_limit_to_sim_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.write_joint_velocity_limit_to_sim_index(limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.write_joint_velocity_limit_to_sim_index( + limits=limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True + ) def write_joint_effort_limit_to_sim_index( self, + *, limits: torch.Tensor | wp.array | float, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1458,6 +1520,7 @@ def write_joint_effort_limit_to_sim_index( def write_joint_effort_limit_to_sim_mask( self, + *, limits: torch.Tensor | wp.array | float, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1477,7 +1540,7 @@ def write_joint_effort_limit_to_sim_mask( Args: limits: Joint torque limits. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -1489,10 +1552,11 @@ def write_joint_effort_limit_to_sim_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.write_joint_effort_limit_to_sim_index(limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.write_joint_effort_limit_to_sim_index(limits=limits, joint_ids=joint_ids, env_ids=env_ids, full_data=True) def write_joint_armature_to_sim_index( self, + *, armature: torch.Tensor | wp.array | float, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1555,6 +1619,7 @@ def write_joint_armature_to_sim_index( def write_joint_armature_to_sim_mask( self, + *, armature: torch.Tensor | wp.array | float, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1574,7 +1639,7 @@ def write_joint_armature_to_sim_mask( Args: armature: Joint armature. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -1586,10 +1651,11 @@ def write_joint_armature_to_sim_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.write_joint_armature_to_sim_index(armature, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.write_joint_armature_to_sim_index(armature=armature, joint_ids=joint_ids, env_ids=env_ids, full_data=True) def write_joint_friction_coefficient_to_sim_index( self, + *, joint_friction_coeff: torch.Tensor | wp.array | float, joint_dynamic_friction_coeff: torch.Tensor | wp.array | float | None = None, joint_viscous_friction_coeff: torch.Tensor | wp.array | float | None = None, @@ -1660,6 +1726,7 @@ def write_joint_friction_coefficient_to_sim_index( def write_joint_friction_coefficient_to_sim_mask( self, + *, joint_friction_coeff: torch.Tensor | wp.array, joint_dynamic_friction_coeff: torch.Tensor | wp.array | None = None, joint_viscous_friction_coeff: torch.Tensor | wp.array | None = None, @@ -1694,7 +1761,7 @@ def write_joint_friction_coefficient_to_sim_mask( joint_viscous_friction_coeff: Viscous friction coefficient :math:`c_v`. Same shape as above. If None, the viscous coefficient is not updated. joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -1707,7 +1774,7 @@ def write_joint_friction_coefficient_to_sim_mask( joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff, + joint_friction_coeff=joint_friction_coeff, joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, joint_viscous_friction_coeff=joint_viscous_friction_coeff, joint_ids=joint_ids, @@ -1717,6 +1784,7 @@ def write_joint_friction_coefficient_to_sim_mask( def write_joint_dynamic_friction_coefficient_to_sim_index( self, + *, joint_dynamic_friction_coeff: torch.Tensor | wp.array, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1767,6 +1835,7 @@ def write_joint_dynamic_friction_coefficient_to_sim_index( def write_joint_dynamic_friction_coefficient_to_sim_mask( self, + *, joint_dynamic_friction_coeff: torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1783,7 +1852,7 @@ def write_joint_dynamic_friction_coefficient_to_sim_mask( Args: joint_dynamic_friction_coeff: Joint dynamic friction coefficient. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -1796,11 +1865,15 @@ def write_joint_dynamic_friction_coefficient_to_sim_mask( joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_dynamic_friction_coefficient_to_sim_index( - joint_dynamic_friction_coeff, joint_ids=joint_ids, env_ids=env_ids, full_data=True + joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + full_data=True, ) def write_joint_viscous_friction_coefficient_to_sim_index( self, + *, joint_viscous_friction_coeff: torch.Tensor | wp.array, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1822,7 +1895,7 @@ def write_joint_viscous_friction_coefficient_to_sim_index( env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: logger.warning("Setting joint viscous friction coefficients are not supported in Isaac Sim < 5.0") return # resolve all indices @@ -1854,6 +1927,7 @@ def write_joint_viscous_friction_coefficient_to_sim_index( def write_joint_viscous_friction_coefficient_to_sim_mask( self, + *, joint_viscous_friction_coeff: torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1870,7 +1944,7 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( Args: joint_viscous_friction_coeff: Joint viscous friction coefficient. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve masks if env_mask is not None: @@ -1883,7 +1957,10 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.write_joint_viscous_friction_coefficient_to_sim_index( - joint_viscous_friction_coeff, joint_ids=joint_ids, env_ids=env_ids, full_data=True + joint_viscous_friction_coeff=joint_viscous_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + full_data=True, ) """ @@ -1892,6 +1969,7 @@ def write_joint_viscous_friction_coefficient_to_sim_mask( def set_masses_index( self, + *, masses: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1938,6 +2016,7 @@ def set_masses_index( def set_masses_mask( self, + *, masses: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -1954,7 +2033,7 @@ def set_masses_mask( Args: masses: Masses of all bodies. Shape is (num_instances, num_bodies). body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -1966,10 +2045,11 @@ def set_masses_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_masses_index(masses, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) def set_coms_index( self, + *, coms: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -1986,7 +2066,8 @@ def set_coms_index( Args: coms: Center of mass pose of all bodies. Shape is (len(env_ids), len(body_ids), 7) or - (num_instances, num_bodies, 7) if full_data. + (num_instances, num_bodies, 7) if full_data, or (len(env_ids), len(body_ids)) / + (num_instances, num_bodies) with dtype wp.transformf. body_ids: The body indices to set the center of mass pose for. Defaults to None (all bodies). env_ids: The environment indices to set the center of mass pose for. Defaults to None (all environments). full_data: Whether to expect full data. Defaults to False. @@ -2021,6 +2102,7 @@ def set_coms_index( def set_coms_mask( self, + *, coms: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2035,9 +2117,10 @@ def set_coms_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - coms: Center of mass pose of all bodies. Shape is (num_instances, num_bodies, 7). + coms: Center of mass pose of all bodies. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2049,10 +2132,11 @@ def set_coms_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_coms_index(coms, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) def set_inertias_index( self, + *, inertias: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2098,6 +2182,7 @@ def set_inertias_index( def set_inertias_mask( self, + *, inertias: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2114,7 +2199,7 @@ def set_inertias_mask( Args: inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2126,10 +2211,11 @@ def set_inertias_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_inertias_index(inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) def set_joint_position_target_index( self, + *, target: torch.Tensor | wp.array, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2176,6 +2262,7 @@ def set_joint_position_target_index( def set_joint_position_target_mask( self, + *, target: torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2192,7 +2279,7 @@ def set_joint_position_target_mask( Args: target: Joint position targets. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2204,10 +2291,11 @@ def set_joint_position_target_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_joint_position_target_index(target, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.set_joint_position_target_index(target=target, joint_ids=joint_ids, env_ids=env_ids, full_data=True) def set_joint_velocity_target_index( self, + *, target: torch.Tensor | wp.array, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2254,6 +2342,7 @@ def set_joint_velocity_target_index( def set_joint_velocity_target_mask( self, + *, target: torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2270,7 +2359,7 @@ def set_joint_velocity_target_mask( Args: target: Joint velocity targets. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2282,10 +2371,11 @@ def set_joint_velocity_target_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_joint_velocity_target_index(target, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.set_joint_velocity_target_index(target=target, joint_ids=joint_ids, env_ids=env_ids, full_data=True) def set_joint_effort_target_index( self, + *, target: torch.Tensor | wp.array, joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2332,6 +2422,7 @@ def set_joint_effort_target_index( def set_joint_effort_target_mask( self, + *, target: torch.Tensor | wp.array, joint_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2348,7 +2439,7 @@ def set_joint_effort_target_mask( Args: target: Joint effort targets. Shape is (num_instances, num_joints). joint_mask: Joint mask. If None, then all joints are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2360,7 +2451,7 @@ def set_joint_effort_target_mask( else: joint_ids = self._ALL_JOINT_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_joint_effort_target_index(target, joint_ids=joint_ids, env_ids=env_ids, full_data=True) + self.set_joint_effort_target_index(target=target, joint_ids=joint_ids, env_ids=env_ids, full_data=True) """ Operations - Tendons. @@ -2368,6 +2459,7 @@ def set_joint_effort_target_mask( def set_fixed_tendon_stiffness_index( self, + *, stiffness: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2415,6 +2507,7 @@ def set_fixed_tendon_stiffness_index( def set_fixed_tendon_stiffness_mask( self, + *, stiffness: torch.Tensor | wp.array, fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2435,7 +2528,7 @@ def set_fixed_tendon_stiffness_mask( Args: stiffness: Fixed tendon stiffness. Shape is (num_instances, num_fixed_tendons). fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2448,11 +2541,12 @@ def set_fixed_tendon_stiffness_mask( fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_stiffness_index( - stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True + stiffness=stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True ) def set_fixed_tendon_damping_index( self, + *, damping: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2500,6 +2594,7 @@ def set_fixed_tendon_damping_index( def set_fixed_tendon_damping_mask( self, + *, damping: torch.Tensor | wp.array, fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2520,7 +2615,7 @@ def set_fixed_tendon_damping_mask( Args: damping: Fixed tendon damping. Shape is (num_instances, num_fixed_tendons). fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2532,10 +2627,13 @@ def set_fixed_tendon_damping_mask( else: fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_fixed_tendon_damping_index(damping, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True) + self.set_fixed_tendon_damping_index( + damping=damping, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True + ) def set_fixed_tendon_limit_stiffness_index( self, + *, limit_stiffness: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2583,6 +2681,7 @@ def set_fixed_tendon_limit_stiffness_index( def set_fixed_tendon_limit_stiffness_mask( self, + *, limit_stiffness: torch.Tensor | wp.array, fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2603,7 +2702,7 @@ def set_fixed_tendon_limit_stiffness_mask( Args: limit_stiffness: Fixed tendon limit stiffness. Shape is (num_instances, num_fixed_tendons). fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2616,11 +2715,12 @@ def set_fixed_tendon_limit_stiffness_mask( fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_limit_stiffness_index( - limit_stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True + limit_stiffness=limit_stiffness, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True ) def set_fixed_tendon_position_limit_index( self, + *, limit: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2668,6 +2768,7 @@ def set_fixed_tendon_position_limit_index( def set_fixed_tendon_position_limit_mask( self, + *, limit: torch.Tensor | wp.array, fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2688,7 +2789,7 @@ def set_fixed_tendon_position_limit_mask( Args: limit: Fixed tendon position limit. Shape is (num_instances, num_fixed_tendons). fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2701,11 +2802,12 @@ def set_fixed_tendon_position_limit_mask( fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_position_limit_index( - limit, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True + limit=limit, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True ) def set_fixed_tendon_rest_length_index( self, + *, rest_length: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2753,6 +2855,7 @@ def set_fixed_tendon_rest_length_index( def set_fixed_tendon_rest_length_mask( self, + *, rest_length: torch.Tensor | wp.array, fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2773,7 +2876,7 @@ def set_fixed_tendon_rest_length_mask( Args: rest_length: Fixed tendon rest length. Shape is (num_instances, num_fixed_tendons). fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2786,11 +2889,12 @@ def set_fixed_tendon_rest_length_mask( fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.set_fixed_tendon_rest_length_index( - rest_length, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True + rest_length=rest_length, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True ) def set_fixed_tendon_offset_index( self, + *, offset: torch.Tensor | wp.array, fixed_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2838,6 +2942,7 @@ def set_fixed_tendon_offset_index( def set_fixed_tendon_offset_mask( self, + *, offset: torch.Tensor | wp.array, fixed_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2858,7 +2963,7 @@ def set_fixed_tendon_offset_mask( Args: offset: Fixed tendon offset. Shape is (num_instances, num_fixed_tendons). fixed_tendon_mask: Fixed tendon mask. If None, then all fixed tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2870,10 +2975,13 @@ def set_fixed_tendon_offset_mask( else: fixed_tendon_ids = self._ALL_FIXED_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_fixed_tendon_offset_index(offset, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True) + self.set_fixed_tendon_offset_index( + offset=offset, fixed_tendon_ids=fixed_tendon_ids, env_ids=env_ids, full_data=True + ) def write_fixed_tendon_properties_to_sim_index( self, + *, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write fixed tendon properties into the simulation using indices. @@ -2902,6 +3010,7 @@ def write_fixed_tendon_properties_to_sim_index( def write_fixed_tendon_properties_to_sim_mask( self, + *, env_mask: wp.array | None = None, ) -> None: """Write fixed tendon properties into the simulation using masks. @@ -2911,7 +3020,7 @@ def write_fixed_tendon_properties_to_sim_mask( is only supporting indexing, hence masks need to be converted to indices. Args: - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -2922,6 +3031,7 @@ def write_fixed_tendon_properties_to_sim_mask( def set_spatial_tendon_stiffness_index( self, + *, stiffness: torch.Tensor | wp.array, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -2969,6 +3079,7 @@ def set_spatial_tendon_stiffness_index( def set_spatial_tendon_stiffness_mask( self, + *, stiffness: torch.Tensor | wp.array, spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -2989,7 +3100,7 @@ def set_spatial_tendon_stiffness_mask( Args: stiffness: Spatial tendon stiffness. Shape is (num_instances, num_spatial_tendons). spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -3002,11 +3113,12 @@ def set_spatial_tendon_stiffness_mask( spatial_tendon_ids = self._ALL_SPATIAL_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_stiffness_index( - stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True + stiffness=stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True ) def set_spatial_tendon_damping_index( self, + *, damping: torch.Tensor | wp.array, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -3054,6 +3166,7 @@ def set_spatial_tendon_damping_index( def set_spatial_tendon_damping_mask( self, + *, damping: torch.Tensor | wp.array, spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -3074,7 +3187,7 @@ def set_spatial_tendon_damping_mask( Args: damping: Spatial tendon damping. Shape is (num_instances, num_spatial_tendons). spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -3087,11 +3200,12 @@ def set_spatial_tendon_damping_mask( spatial_tendon_ids = self._ALL_SPATIAL_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_damping_index( - damping, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True + damping=damping, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True ) def set_spatial_tendon_limit_stiffness_index( self, + *, limit_stiffness: torch.Tensor | wp.array, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -3140,6 +3254,7 @@ def set_spatial_tendon_limit_stiffness_index( def set_spatial_tendon_limit_stiffness_mask( self, + *, limit_stiffness: torch.Tensor | wp.array, spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -3160,7 +3275,7 @@ def set_spatial_tendon_limit_stiffness_mask( Args: limit_stiffness: Spatial tendon limit stiffness. Shape is (num_instances, num_spatial_tendons). spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -3173,11 +3288,12 @@ def set_spatial_tendon_limit_stiffness_mask( spatial_tendon_ids = self._ALL_SPATIAL_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_limit_stiffness_index( - limit_stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True + limit_stiffness=limit_stiffness, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True ) def set_spatial_tendon_offset_index( self, + *, offset: torch.Tensor | wp.array, spatial_tendon_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -3225,6 +3341,7 @@ def set_spatial_tendon_offset_index( def set_spatial_tendon_offset_mask( self, + *, offset: torch.Tensor | wp.array, spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -3245,7 +3362,7 @@ def set_spatial_tendon_offset_mask( Args: offset: Spatial tendon offset. Shape is (num_instances, num_spatial_tendons). spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -3258,11 +3375,12 @@ def set_spatial_tendon_offset_mask( spatial_tendon_ids = self._ALL_SPATIAL_TENDON_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. self.set_spatial_tendon_offset_index( - offset, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True + offset=offset, spatial_tendon_ids=spatial_tendon_ids, env_ids=env_ids, full_data=True ) def write_spatial_tendon_properties_to_sim_index( self, + *, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Write spatial tendon properties into the simulation using indices. @@ -3290,6 +3408,7 @@ def write_spatial_tendon_properties_to_sim_index( def write_spatial_tendon_properties_to_sim_mask( self, + *, spatial_tendon_mask: wp.array | None = None, env_mask: wp.array | None = None, ) -> None: @@ -3301,7 +3420,7 @@ def write_spatial_tendon_properties_to_sim_mask( Args: spatial_tendon_mask: Spatial tendon mask. If None, then all spatial tendons are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -3624,24 +3743,30 @@ def _process_actuators_cfg(self): if isinstance(actuator, ImplicitActuator): self._has_implicit_actuators = True # the gains and limits are set into the simulation since actuator model is implicit - self.write_joint_stiffness_to_sim_index(actuator.stiffness, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim_index(actuator.damping, joint_ids=actuator.joint_indices) + self.write_joint_stiffness_to_sim_index(stiffness=actuator.stiffness, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=actuator.damping, joint_ids=actuator.joint_indices) else: # the gains and limits are processed by the actuator model # we set gains to zero, and torque limit to a high value in simulation to avoid any interference - self.write_joint_stiffness_to_sim_index(0.0, joint_ids=actuator.joint_indices) - self.write_joint_damping_to_sim_index(0.0, joint_ids=actuator.joint_indices) + self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=actuator.joint_indices) + self.write_joint_damping_to_sim_index(damping=0.0, joint_ids=actuator.joint_indices) # Set common properties into the simulation - self.write_joint_effort_limit_to_sim_index(actuator.effort_limit_sim, joint_ids=actuator.joint_indices) - self.write_joint_velocity_limit_to_sim_index(actuator.velocity_limit_sim, joint_ids=actuator.joint_indices) - self.write_joint_armature_to_sim_index(actuator.armature, joint_ids=actuator.joint_indices) - self.write_joint_friction_coefficient_to_sim_index(actuator.friction, joint_ids=actuator.joint_indices) + self.write_joint_effort_limit_to_sim_index( + limits=actuator.effort_limit_sim, joint_ids=actuator.joint_indices + ) + self.write_joint_velocity_limit_to_sim_index( + limits=actuator.velocity_limit_sim, joint_ids=actuator.joint_indices + ) + self.write_joint_armature_to_sim_index(armature=actuator.armature, joint_ids=actuator.joint_indices) + self.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=actuator.friction, joint_ids=actuator.joint_indices + ) self.write_joint_dynamic_friction_coefficient_to_sim_index( - actuator.dynamic_friction, joint_ids=actuator.joint_indices + joint_dynamic_friction_coeff=actuator.dynamic_friction, joint_ids=actuator.joint_indices ) self.write_joint_viscous_friction_coefficient_to_sim_index( - actuator.viscous_friction, joint_ids=actuator.joint_indices + joint_viscous_friction_coeff=actuator.viscous_friction, joint_ids=actuator.joint_indices ) # perform some sanity checks to ensure actuators are prepared correctly @@ -3871,7 +3996,7 @@ def format_limits(_, v: tuple[float, float]) -> str: for index, name in enumerate(self.joint_names): # build row data based on Isaac Sim version row_data = [index, name, stiffnesses[index], dampings[index], armatures[index]] - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: row_data.append(static_frictions[index]) else: row_data.extend([static_frictions[index], dynamic_frictions[index], viscous_frictions[index]]) @@ -4090,7 +4215,7 @@ def write_joint_friction_coefficient_to_sim( stacklevel=2, ) self.write_joint_friction_coefficient_to_sim_index( - joint_friction_coeff, + joint_friction_coeff=joint_friction_coeff, joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, joint_viscous_friction_coeff=joint_viscous_friction_coeff, joint_ids=joint_ids, @@ -4113,7 +4238,10 @@ def write_joint_viscous_friction_coefficient_to_sim( stacklevel=2, ) self.write_joint_viscous_friction_coefficient_to_sim_index( - joint_viscous_friction_coeff, joint_ids=joint_ids, env_ids=env_ids, full_data=full_data + joint_viscous_friction_coeff=joint_viscous_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + full_data=full_data, ) def write_joint_dynamic_friction_coefficient_to_sim( @@ -4131,7 +4259,10 @@ def write_joint_dynamic_friction_coefficient_to_sim( stacklevel=2, ) self.write_joint_dynamic_friction_coefficient_to_sim_index( - joint_dynamic_friction_coeff, joint_ids=joint_ids, env_ids=env_ids, full_data=full_data + joint_dynamic_friction_coeff=joint_dynamic_friction_coeff, + joint_ids=joint_ids, + env_ids=env_ids, + full_data=full_data, ) def write_root_state_to_sim( @@ -4147,8 +4278,10 @@ def write_root_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) def write_root_com_state_to_sim( self, @@ -4163,8 +4296,10 @@ def write_root_com_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_com_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_com_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) def write_root_link_state_to_sim( self, @@ -4179,5 +4314,7 @@ def write_root_link_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_link_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + if isinstance(root_state, wp.array): + raise ValueError("The root state must be a torch tensor, not a warp array.") + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py index 5d945cb4a2a..4972015fb35 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation_data.py @@ -21,7 +21,7 @@ from isaaclab_physx.physics import PhysxManager as SimulationManager if TYPE_CHECKING: - from isaaclab.assets.articulation.articulation_view import ArticulationView + import omni.physics.tensors.impl.api as physx # import logger logger = logging.getLogger(__name__) @@ -48,7 +48,7 @@ class ArticulationData(BaseArticulationData): __backend_name__: str = "physx" """The name of the backend for the articulation data.""" - def __init__(self, root_view: ArticulationView, device: str): + def __init__(self, root_view: physx.ArticulationView, device: str): """Initializes the articulation data. Args: @@ -59,7 +59,7 @@ def __init__(self, root_view: ArticulationView, device: str): # Set the root articulation view # note: this is stored as a weak reference to avoid circular references between the asset class # and the data container. This is important to avoid memory leaks. - self._root_view: ArticulationView = weakref.proxy(root_view) + self._root_view: physx.ArticulationView = weakref.proxy(root_view) # Set initial time stamp self._sim_timestamp = 0.0 @@ -89,7 +89,8 @@ def is_primed(self) -> bool: def is_primed(self, value: bool) -> None: """Set whether the articulation data is fully instantiated and ready to use. - .. note:: Once this quantity is set to True, it cannot be changed. + .. note:: + Once this quantity is set to True, it cannot be changed. Args: value: The primed state. @@ -137,7 +138,8 @@ def update(self, dt: float) -> None: def default_root_pose(self) -> wp.array: """Default root pose ``[pos, quat]`` in the local environment frame. - The position and quaternion are of the articulation root's actor frame. Shape is (num_instances, 7). + The position and quaternion are of the articulation root's actor frame. + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). """ return self._default_root_pose @@ -160,7 +162,7 @@ def default_root_vel(self) -> wp.array: """Default root velocity ``[lin_vel, ang_vel]`` in the local environment frame. The linear and angular velocities are of the articulation root's center of mass frame. - Shape is (num_instances, 6). + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). """ return self._default_root_vel @@ -178,41 +180,11 @@ def default_root_vel(self, value: wp.array) -> None: raise ValueError("The articulation data is already primed.") self._default_root_vel.assign(value) - @property - def default_root_state(self) -> wp.array: - """Default root state ``[pos, quat, lin_vel, ang_vel]`` in the local environment frame. - - - The position and quaternion are of the articulation root's actor frame. Meanwhile, the linear and angular - velocities are of its center of mass frame. Shape is (num_instances, 13). - - This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. - """ - warnings.warn( - "Reading the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " - "Please use the default_root_pose and default_root_vel properties instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._default_root_state is None: - self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) - wp.launch( - shared_kernels.concat_root_pose_and_vel_to_state, - dim=self._num_instances, - inputs=[ - self._default_root_pose, - self._default_root_vel, - ], - outputs=[ - self._default_root_state, - ], - device=self.device, - ) - return self._default_root_state - @property def default_joint_pos(self) -> wp.array: - """Default joint positions of all joints. Shape is (num_instances, num_joints). + """Default joint positions of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ @@ -234,7 +206,9 @@ def default_joint_pos(self, value: wp.array) -> None: @property def default_joint_vel(self) -> wp.array: - """Default joint velocities of all joints. Shape is (num_instances, num_joints). + """Default joint velocities of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. """ @@ -260,7 +234,9 @@ def default_joint_vel(self, value: wp.array) -> None: @property def joint_pos_target(self) -> wp.array: - """Joint position targets commanded by the user. Shape is (num_instances, num_joints). + """Joint position targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). For an implicit actuator model, the targets are directly set into the simulation. For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), @@ -270,7 +246,9 @@ def joint_pos_target(self) -> wp.array: @property def joint_vel_target(self) -> wp.array: - """Joint velocity targets commanded by the user. Shape is (num_instances, num_joints). + """Joint velocity targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). For an implicit actuator model, the targets are directly set into the simulation. For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), @@ -280,7 +258,9 @@ def joint_vel_target(self) -> wp.array: @property def joint_effort_target(self) -> wp.array: - """Joint effort targets commanded by the user. Shape is (num_instances, num_joints). + """Joint effort targets commanded by the user. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). For an implicit actuator model, the targets are directly set into the simulation. For an explicit actuator model, the targets are used to compute the joint torques (see :attr:`applied_torque`), @@ -294,7 +274,9 @@ def joint_effort_target(self) -> wp.array: @property def computed_torque(self) -> wp.array: - """Joint torques computed from the actuator model (before clipping). Shape is (num_instances, num_joints). + """Joint torques computed from the actuator model (before clipping). + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). This quantity is the raw torque output from the actuator mode, before any clipping is applied. It is exposed for users who want to inspect the computations inside the actuator model. @@ -304,7 +286,9 @@ def computed_torque(self) -> wp.array: @property def applied_torque(self) -> wp.array: - """Joint torques applied from the actuator model (after clipping). Shape is (num_instances, num_joints). + """Joint torques applied from the actuator model (after clipping). + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). These torques are set into the simulation, after clipping the :attr:`computed_torque` based on the actuator model. @@ -317,7 +301,9 @@ def applied_torque(self) -> wp.array: @property def joint_stiffness(self) -> wp.array: - """Joint stiffness provided to the simulation. Shape is (num_instances, num_joints). + """Joint stiffness provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). In the case of explicit actuators, the value for the corresponding joints is zero. """ @@ -325,7 +311,9 @@ def joint_stiffness(self) -> wp.array: @property def joint_damping(self) -> wp.array: - """Joint damping provided to the simulation. Shape is (num_instances, num_joints) + """Joint damping provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). In the case of explicit actuators, the value for the corresponding joints is zero. """ @@ -333,27 +321,42 @@ def joint_damping(self) -> wp.array: @property def joint_armature(self) -> wp.array: - """Joint armature provided to the simulation. Shape is (num_instances, num_joints).""" + """Joint armature provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ return self._joint_armature @property def joint_friction_coeff(self) -> wp.array: - """Joint static friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + """Joint static friction coefficient provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ return self._joint_friction_coeff @property def joint_dynamic_friction_coeff(self) -> wp.array: - """Joint dynamic friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + """Joint dynamic friction coefficient provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ return self._joint_dynamic_friction_coeff @property def joint_viscous_friction_coeff(self) -> wp.array: - """Joint viscous friction coefficient provided to the simulation. Shape is (num_instances, num_joints).""" + """Joint viscous friction coefficient provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ return self._joint_viscous_friction_coeff @property def joint_pos_limits(self) -> wp.array: - """Joint position limits provided to the simulation. Shape is (num_instances, num_joints, 2). + """Joint position limits provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_joints, 2). The limits are in the order :math:`[lower, upper]`. """ @@ -361,12 +364,18 @@ def joint_pos_limits(self) -> wp.array: @property def joint_vel_limits(self) -> wp.array: - """Joint maximum velocity provided to the simulation. Shape is (num_instances, num_joints).""" + """Joint maximum velocity provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ return self._joint_vel_limits @property def joint_effort_limits(self) -> wp.array: - """Joint maximum effort provided to the simulation. Shape is (num_instances, num_joints).""" + """Joint maximum effort provided to the simulation. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ return self._joint_effort_limits """ @@ -375,7 +384,10 @@ def joint_effort_limits(self) -> wp.array: @property def soft_joint_pos_limits(self) -> wp.array: - r"""Soft joint positions limits for all joints. Shape is (num_instances, num_joints, 2). + r"""Soft joint positions limits for all joints. + + Shape is (num_instances, num_joints), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_joints, 2). The limits are in the order :math:`[lower, upper]`.The soft joint position limits are computed as a sub-region of the :attr:`joint_pos_limits` based on the @@ -396,7 +408,9 @@ def soft_joint_pos_limits(self) -> wp.array: @property def soft_joint_vel_limits(self) -> wp.array: - """Soft joint velocity limits for all joints. Shape is (num_instances, num_joints). + """Soft joint velocity limits for all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). These are obtained from the actuator model. It may differ from :attr:`joint_vel_limits` if the actuator model has a variable velocity limit model. For instance, in a variable gear ratio actuator model. @@ -405,7 +419,10 @@ def soft_joint_vel_limits(self) -> wp.array: @property def gear_ratio(self) -> wp.array: - """Gear ratio for relating motor torques to applied Joint torques. Shape is (num_instances, num_joints).""" + """Gear ratio for relating motor torques to applied Joint torques. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ return self._gear_ratio """ @@ -414,32 +431,56 @@ def gear_ratio(self) -> wp.array: @property def fixed_tendon_stiffness(self) -> wp.array: - """Fixed tendon stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon stiffness provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ return self._fixed_tendon_stiffness @property def fixed_tendon_damping(self) -> wp.array: - """Fixed tendon damping provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon damping provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ return self._fixed_tendon_damping @property def fixed_tendon_limit_stiffness(self) -> wp.array: - """Fixed tendon limit stiffness provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon limit stiffness provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ return self._fixed_tendon_limit_stiffness @property def fixed_tendon_rest_length(self) -> wp.array: - """Fixed tendon rest length provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon rest length provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ return self._fixed_tendon_rest_length @property def fixed_tendon_offset(self) -> wp.array: - """Fixed tendon offset provided to the simulation. Shape is (num_instances, num_fixed_tendons).""" + """Fixed tendon offset provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_fixed_tendons). + """ return self._fixed_tendon_offset @property def fixed_tendon_pos_limits(self) -> wp.array: - """Fixed tendon position limits provided to the simulation. Shape is (num_instances, num_fixed_tendons, 2).""" + """Fixed tendon position limits provided to the simulation. + + Shape is (num_instances, num_fixed_tendons), dtype = wp.vec2f. In torch this resolves to + (num_instances, num_fixed_tendons, 2). + """ return self._fixed_tendon_pos_limits """ @@ -448,22 +489,38 @@ def fixed_tendon_pos_limits(self) -> wp.array: @property def spatial_tendon_stiffness(self) -> wp.array: - """Spatial tendon stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + """Spatial tendon stiffness provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ return self._spatial_tendon_stiffness @property def spatial_tendon_damping(self) -> wp.array: - """Spatial tendon damping provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + """Spatial tendon damping provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ return self._spatial_tendon_damping @property def spatial_tendon_limit_stiffness(self) -> wp.array: - """Spatial tendon limit stiffness provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + """Spatial tendon limit stiffness provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ return self._spatial_tendon_limit_stiffness @property def spatial_tendon_offset(self) -> wp.array: - """Spatial tendon offset provided to the simulation. Shape is (num_instances, num_spatial_tendons).""" + """Spatial tendon offset provided to the simulation. + + Shape is (num_instances, num_spatial_tendons), dtype = wp.float32. In torch this resolves to + (num_instances, num_spatial_tendons). + """ return self._spatial_tendon_offset """ @@ -472,7 +529,8 @@ def spatial_tendon_offset(self) -> wp.array: @property def root_link_pose_w(self) -> wp.array: - """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + """Root link pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the articulation root's actor frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -486,7 +544,8 @@ def root_link_pose_w(self) -> wp.array: @property def root_link_vel_w(self) -> wp.array: - """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the articulation root's actor frame relative to the world. @@ -511,7 +570,8 @@ def root_link_vel_w(self) -> wp.array: @property def root_com_pose_w(self) -> wp.array: - """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + """Root center of mass pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the articulation root's center of mass frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -536,7 +596,8 @@ def root_com_pose_w(self) -> wp.array: @property def root_com_vel_w(self) -> wp.array: - """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the articulation root's center of mass frame relative to the world. @@ -553,20 +614,28 @@ def root_com_vel_w(self) -> wp.array: @property def body_mass(self) -> wp.array: - """Body mass ``wp.float32`` in the world frame. Shape is (num_instances, num_bodies).""" + """Body mass in the world frame. + + Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). + """ self._body_mass.assign(self._root_view.get_masses()) return self._body_mass @property def body_inertia(self) -> wp.array: - """Body inertia ``wp.mat33`` in the world frame. Shape is (num_instances, num_bodies, 3, 3).""" + """Flattened body inertia in the world frame. + + Shape is (num_instances, num_bodies, 9), dtype = wp.float32. In torch this resolves to + (num_instances, num_bodies, 9). + """ self._body_inertia.assign(self._root_view.get_inertias()) return self._body_inertia @property def body_link_pose_w(self) -> wp.array: """Body link pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies, 7). + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the articulation links' actor frame relative to the world. The orientation is provided in (x, y, z, w) format. @@ -583,7 +652,8 @@ def body_link_pose_w(self) -> wp.array: @property def body_link_vel_w(self) -> wp.array: """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies, 6). + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the articulation links' actor frame relative to the world. @@ -609,7 +679,8 @@ def body_link_vel_w(self) -> wp.array: @property def body_com_pose_w(self) -> wp.array: """Body center of mass pose ``[pos, quat]`` in simulation world frame. - Shape is (num_instances, num_bodies, 7). + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the articulation links relative to the world. The orientation is provided in (x, y, z, w) format. @@ -634,7 +705,8 @@ def body_com_pose_w(self) -> wp.array: @property def body_com_vel_w(self) -> wp.array: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies, 6). + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the articulation links' center of mass frame relative to the world. @@ -645,85 +717,11 @@ def body_com_vel_w(self) -> wp.array: return self._body_com_vel_w.data - @property - def body_state_w(self): - """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. - Shape is (num_instances, num_bodies, 13). - - The position and quaternion are of all the articulation links' actor frame. Meanwhile, the linear and angular - velocities are of the articulation links's center of mass frame. - """ - if self._body_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_body_pose_and_vel_to_state, - dim=(self._num_instances, self._num_bodies), - inputs=[ - self.body_link_pose_w, - self.body_com_vel_w, - ], - outputs=[ - self._body_state_w.data, - ], - device=self.device, - ) - self._body_state_w.timestamp = self._sim_timestamp - - return self._body_state_w.data - - @property - def body_link_state_w(self): - """State of all bodies' link frame`[pos, quat, lin_vel, ang_vel]` in simulation world frame. - Shape is (num_instances, num_bodies, 13). - - The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. - """ - if self._body_link_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_body_pose_and_vel_to_state, - dim=(self._num_instances, self._num_bodies), - inputs=[ - self.body_link_pose_w, - self.body_link_vel_w, - ], - outputs=[ - self._body_link_state_w.data, - ], - device=self.device, - ) - self._body_link_state_w.timestamp = self._sim_timestamp - - return self._body_link_state_w.data - - @property - def body_com_state_w(self): - """State of all bodies center of mass `[pos, quat, lin_vel, ang_vel]` in simulation world frame. - Shape is (num_instances, num_bodies, 13). - - The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the - world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the - principle inertia. - """ - if self._body_com_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_body_pose_and_vel_to_state, - dim=(self._num_instances, self._num_bodies), - inputs=[ - self.body_com_pose_w, - self.body_com_vel_w, - ], - outputs=[ - self._body_com_state_w.data, - ], - device=self.device, - ) - self._body_com_state_w.timestamp = self._sim_timestamp - - return self._body_com_state_w.data - @property def body_com_acc_w(self): """Acceleration of all bodies center of mass ``[lin_acc, ang_acc]``. - Shape is (num_instances, num_bodies, 6). + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). All values are relative to the world. """ @@ -737,7 +735,8 @@ def body_com_acc_w(self): @property def body_com_pose_b(self) -> wp.array: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, 1, 7). + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. The orientation is provided in (x, y, z, w) format. @@ -756,7 +755,7 @@ def body_incoming_joint_wrench_b(self) -> wp.array: Shape is (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the world of an articulation. - For more information on joint wrenches, please check the`PhysX documentation`_ and the underlying + For more information on joint wrenches, please check the `PhysX documentation`_ and the underlying `PhysX Tensor API`_. .. _`PhysX documentation`: https://nvidia-omniverse.github.io/PhysX/physx/5.5.1/docs/Articulations.html#link-incoming-joint-force @@ -776,7 +775,10 @@ def body_incoming_joint_wrench_b(self) -> wp.array: @property def joint_pos(self) -> wp.array: - """Joint positions of all joints. Shape is (num_instances, num_joints).""" + """Joint positions of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ if self._joint_pos.timestamp < self._sim_timestamp: # read data from simulation and set the buffer data and timestamp self._joint_pos.data = self._root_view.get_dof_positions() @@ -785,7 +787,10 @@ def joint_pos(self) -> wp.array: @property def joint_vel(self) -> wp.array: - """Joint velocities of all joints. Shape is (num_instances, num_joints).""" + """Joint velocities of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ if self._joint_vel.timestamp < self._sim_timestamp: # read data from simulation and set the buffer data and timestamp self._joint_vel.data = self._root_view.get_dof_velocities() @@ -794,7 +799,10 @@ def joint_vel(self) -> wp.array: @property def joint_acc(self) -> wp.array: - """Joint acceleration of all joints. Shape is (num_instances, num_joints).""" + """Joint acceleration of all joints. + + Shape is (num_instances, num_joints), dtype = wp.float32. In torch this resolves to (num_instances, num_joints). + """ if self._joint_acc.timestamp < self._sim_timestamp: # note: we use finite differencing to compute acceleration time_elapsed = self._sim_timestamp - self._joint_acc.timestamp @@ -820,7 +828,8 @@ def joint_acc(self) -> wp.array: @property def projected_gravity_b(self): - """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + """Projection of the gravity direction on base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3).""" if self._projected_gravity_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, @@ -834,7 +843,7 @@ def projected_gravity_b(self): @property def heading_w(self): - """Yaw heading of the base frame (in radians). Shape is (num_instances,). + """Yaw heading of the base frame (in radians). Shape is (num_instances,), dtype = wp.float32. .. note:: This quantity is computed by assuming that the forward-direction of the base @@ -853,10 +862,10 @@ def heading_w(self): @property def root_link_lin_vel_b(self) -> wp.array: - """Root link linear velocity in base frame. Shape is (num_instances, 3). + """Root link linear velocity in base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the linear velocity of the articulation root's actor frame with respect to the - its actor frame. + This quantity is the linear velocity of the articulation root's actor frame with respect to its actor frame. """ if self._root_link_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( @@ -871,10 +880,10 @@ def root_link_lin_vel_b(self) -> wp.array: @property def root_link_ang_vel_b(self) -> wp.array: - """Root link angular velocity in base world frame. Shape is (num_instances, 3). + """Root link angular velocity in base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the angular velocity of the articulation root's actor frame with respect to the - its actor frame. + This quantity is the angular velocity of the articulation root's actor frame with respect to its actor frame. """ if self._root_link_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( @@ -889,10 +898,11 @@ def root_link_ang_vel_b(self) -> wp.array: @property def root_com_lin_vel_b(self) -> wp.array: - """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the linear velocity of the articulation root's center of mass frame with respect to the - its actor frame. + This quantity is the linear velocity of the articulation root's center of mass frame + with respect to its actor frame. """ if self._root_com_lin_vel_b.timestamp < self._sim_timestamp: wp.launch( @@ -907,10 +917,11 @@ def root_com_lin_vel_b(self) -> wp.array: @property def root_com_ang_vel_b(self) -> wp.array: - """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the angular velocity of the articulation root's center of mass frame with respect to the - its actor frame. + This quantity is the angular velocity of the articulation root's center of mass frame + with respect to its actor frame. """ if self._root_com_ang_vel_b.timestamp < self._sim_timestamp: wp.launch( @@ -929,7 +940,8 @@ def root_com_ang_vel_b(self) -> wp.array: @property def root_link_pos_w(self) -> wp.array: - """Root link position in simulation world frame. Shape is (num_instances, 3). + """Root link position in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ @@ -937,7 +949,8 @@ def root_link_pos_w(self) -> wp.array: @property def root_link_quat_w(self) -> wp.array: - """Root link orientation (x, y, z, w) in simulation world frame. Shape is (num_instances, 4). + """Root link orientation (x, y, z, w) in simulation world frame. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ @@ -945,7 +958,8 @@ def root_link_quat_w(self) -> wp.array: @property def root_link_lin_vel_w(self) -> wp.array: - """Root linear velocity in simulation world frame. Shape is (num_instances, 3). + """Root linear velocity in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ @@ -953,7 +967,8 @@ def root_link_lin_vel_w(self) -> wp.array: @property def root_link_ang_vel_w(self) -> wp.array: - """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + """Root link angular velocity in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ @@ -961,23 +976,26 @@ def root_link_ang_vel_w(self) -> wp.array: @property def root_com_pos_w(self) -> wp.array: - """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + """Root center of mass position in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). - This quantity is the position of the actor frame of the root rigid body relative to the world. + This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ return self._get_pos_from_transform(self.root_com_pose_w) @property def root_com_quat_w(self) -> wp.array: - """Root center of mass orientation (x, y, z, w) in simulation world frame. Shape is (num_instances, 4). + """Root center of mass orientation (x, y, z, w) in simulation world frame. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). - This quantity is the orientation of the actor frame of the root rigid body relative to the world. + This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ return self._get_quat_from_transform(self.root_com_pose_w) @property def root_com_lin_vel_w(self) -> wp.array: - """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ @@ -985,7 +1003,8 @@ def root_com_lin_vel_w(self) -> wp.array: @property def root_com_ang_vel_w(self) -> wp.array: - """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ @@ -993,7 +1012,9 @@ def root_com_ang_vel_w(self) -> wp.array: @property def body_link_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Positions of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the position of the articulation bodies' actor frame relative to the world. """ @@ -1001,7 +1022,9 @@ def body_link_pos_w(self) -> wp.array: @property def body_link_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). This quantity is the orientation of the articulation bodies' actor frame relative to the world. """ @@ -1009,40 +1032,49 @@ def body_link_quat_w(self) -> wp.array: @property def body_link_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear velocity of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). - This quantity is the linear velocity of the articulation bodies' center of mass frame relative to the world. + This quantity is the linear velocity of the articulation bodies' actor frame relative to the world. """ return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) @property def body_link_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular velocity of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). - This quantity is the angular velocity of the articulation bodies' center of mass frame relative to the world. + This quantity is the angular velocity of the articulation bodies' actor frame relative to the world. """ return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) @property def body_com_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Positions of all bodies' center of mass in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). - This quantity is the position of the articulation bodies' actor frame. + This quantity is the position of the articulation bodies' center of mass frame. """ return self._get_pos_from_transform(self.body_com_pose_w) @property def body_com_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all bodies in simulation world frame. - Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). - This quantity is the orientation of the articulation bodies' actor frame. + This quantity is the orientation of the articulation bodies' principal axes of inertia. """ return self._get_quat_from_transform(self.body_com_pose_w) @property def body_com_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear velocity of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear velocity of the articulation bodies' center of mass frame. """ @@ -1050,7 +1082,9 @@ def body_com_lin_vel_w(self) -> wp.array: @property def body_com_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular velocity of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular velocity of the articulation bodies' center of mass frame. """ @@ -1058,7 +1092,9 @@ def body_com_ang_vel_w(self) -> wp.array: @property def body_com_lin_acc_w(self) -> wp.array: - """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear acceleration of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear acceleration of the articulation bodies' center of mass frame. """ @@ -1066,7 +1102,9 @@ def body_com_lin_acc_w(self) -> wp.array: @property def body_com_ang_acc_w(self) -> wp.array: - """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular acceleration of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular acceleration of the articulation bodies' center of mass frame. """ @@ -1075,18 +1113,21 @@ def body_com_ang_acc_w(self) -> wp.array: @property def body_com_pos_b(self) -> wp.array: """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies, 3). + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). - This quantity is the center of mass location relative to its body'slink frame. + This quantity is the center of mass location relative to its body's link frame. """ return self._get_pos_from_transform(self.body_com_pose_b) @property def body_com_quat_b(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all of the bodies in their - respective link frames. Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their + respective link frames. + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). - This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ return self._get_quat_from_transform(self.body_com_pose_b) @@ -1323,6 +1364,38 @@ def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array Deprecated properties. """ + @property + def default_root_state(self) -> wp.array: + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in the local environment frame. + + + The position and quaternion are of the articulation root's actor frame. Meanwhile, the linear and angular + velocities are of its center of mass frame. Shape is (num_instances, 13). + + This quantity is configured through the :attr:`isaaclab.assets.ArticulationCfg.init_state` parameter. + """ + warnings.warn( + "Reading the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " + "Please use the default_root_pose and default_root_vel properties instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._default_root_state is None: + self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self._default_root_pose, + self._default_root_vel, + ], + outputs=[ + self._default_root_state, + ], + device=self.device, + ) + return self._default_root_state + @property def root_state_w(self) -> wp.array: """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" @@ -1400,3 +1473,96 @@ def root_com_state_w(self) -> wp.array: self._root_com_state_w.timestamp = self._sim_timestamp return self._root_com_state_w.data + + @property + def body_state_w(self): + """State of all bodies `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position and quaternion are of all the articulation links' actor frame. Meanwhile, the linear and angular + velocities are of the articulation links' center of mass frame. + """ + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_com_vel_w, + ], + outputs=[ + self._body_state_w.data, + ], + device=self.device, + ) + self._body_state_w.timestamp = self._sim_timestamp + + return self._body_state_w.data + + @property + def body_link_state_w(self): + """State of all bodies' link frame`[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's link frame relative to the world. + """ + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_link_vel_w, + ], + outputs=[ + self._body_link_state_w.data, + ], + device=self.device, + ) + self._body_link_state_w.timestamp = self._sim_timestamp + + return self._body_link_state_w.data + + @property + def body_com_state_w(self): + """State of all bodies center of mass `[pos, quat, lin_vel, ang_vel]` in simulation world frame. + Shape is (num_instances, num_bodies, 13). + + The position, quaternion, and linear/angular velocity are of the body's center of mass frame relative to the + world. Center of mass frame is assumed to be the same orientation as the link rather than the orientation of the + principal inertia. + """ + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self._num_instances, self._num_bodies), + inputs=[ + self.body_com_pose_w, + self.body_com_vel_w, + ], + outputs=[ + self._body_com_state_w.data, + ], + device=self.device, + ) + self._body_com_state_w.timestamp = self._sim_timestamp + + return self._body_com_state_w.data diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py index efe15c3f48d..1e18790c854 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object.py @@ -152,8 +152,14 @@ def max_collision_vertices_per_body(self) -> int: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None): - # Think: Should we reset the kinematic targets when resetting the object? + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: + """Reset the deformable object. + + Args: + env_ids: Environment indices. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). + """ + # TODO: Should we reset the kinematic targets when resetting the object? # This is not done in the current implementation. We assume users will reset the kinematic targets. pass diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py index d61a033526d..825128c8b30 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object.py @@ -14,7 +14,6 @@ import torch import warp as wp -import omni.physics.tensors.impl.api as physx from pxr import UsdPhysics import isaaclab.sim as sim_utils @@ -28,6 +27,8 @@ from .rigid_object_data import RigidObjectData if TYPE_CHECKING: + import omni.physics.tensors.impl.api as physx + from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg # import logger @@ -129,12 +130,12 @@ def permanent_wrench_composer(self) -> WrenchComposer: Operations. """ - def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: """Reset the rigid object. Args: env_ids: Environment indices. If None, then all indices are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # resolve all indices if (env_ids is None) or (env_ids == slice(None)): @@ -154,7 +155,7 @@ def write_data_to_sim(self) -> None: if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques( + self._instantaneous_wrench_composer.add_forces_and_torques_index( forces=self._permanent_wrench_composer.composed_force, torques=self._permanent_wrench_composer.composed_torque, body_ids=self._ALL_BODY_INDICES, @@ -212,28 +213,31 @@ def find_bodies(self, name_keys: str | Sequence[str], preserve_order: bool = Fal def write_root_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root pose over selected environment indices into the simulation. - The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: This method expects partial data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7). + root_pose: Root poses in simulation frame. Shape is (len(env_ids), 7) + or (len(env_ids),) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. """ - self.write_root_link_pose_to_sim_index(root_pose, env_ids=env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids) def write_root_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: @@ -243,40 +247,46 @@ def write_root_pose_to_sim_mask( This method expects full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ - self.write_root_link_pose_to_sim_mask(root_pose, env_mask=env_mask) + self.write_root_link_pose_to_sim_mask(root_pose=root_pose, env_mask=env_mask) def write_root_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the root center of mass velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects partial data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6). + root_velocity: Root center of mass velocities in simulation world frame. Shape is (len(env_ids), 6) + or (len(env_ids),) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. """ self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids) def write_root_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: @@ -286,34 +296,37 @@ def write_root_velocity_to_sim_mask( This method expects full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ self.write_root_com_velocity_to_sim_mask(root_velocity=root_velocity, env_mask=env_mask) def write_root_link_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, ) -> None: """Set the root link pose over selected environment indices into the simulation. - The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: This method expects partial data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + root_pose: Root link poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7), + or (len(env_ids),) / (num_instances,) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ @@ -346,50 +359,54 @@ def write_root_link_pose_to_sim_index( def write_root_link_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: """Set the root link pose over selected environment mask into the simulation. - The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). .. note:: This method expects full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ if env_mask is not None: env_ids = wp.nonzero(env_mask) else: env_ids = self._ALL_INDICES - self.write_root_link_pose_to_sim_index(root_pose, env_ids=env_ids, full_data=True) + self.write_root_link_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids, full_data=True) def write_root_com_pose_to_sim_index( self, + *, root_pose: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, ) -> None: """Set the root center of mass pose over selected environment indices into the simulation. - The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). - The orientation is the orientation of the principle axes of inertia. + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. .. note:: This method expects partial data or full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7). + root_pose: Root center of mass poses in simulation frame. Shape is (len(env_ids), 7) or (num_instances, 7), + or (len(env_ids),) / (num_instances,) with dtype wp.transformf. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ @@ -425,33 +442,36 @@ def write_root_com_pose_to_sim_index( def write_root_com_pose_to_sim_mask( self, + *, root_pose: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: """Set the root center of mass pose over selected environment mask into the simulation. - The root pose comprises of the cartesian position and quaternion orientation in (w, x, y, z). - The orientation is the orientation of the principle axes of inertia. + The root pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. .. note:: This method expects full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7). - env_mask: Environment mask. If None, then all indices are used. + root_pose: Root center of mass poses in simulation frame. Shape is (num_instances, 7) + or (num_instances,) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ if env_mask is not None: env_ids = wp.nonzero(env_mask) else: env_ids = self._ALL_INDICES - self.write_root_com_pose_to_sim_index(root_pose, env_ids=env_ids, full_data=True) + self.write_root_com_pose_to_sim_index(root_pose=root_pose, env_ids=env_ids, full_data=True) def write_root_com_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, @@ -459,18 +479,21 @@ def write_root_com_velocity_to_sim_index( """Set the root center of mass velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects partial data or full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: root_velocity: Root center of mass velocities in simulation world frame. - Shape is (len(env_ids), 6) or (num_instances, 6). + Shape is (len(env_ids), 6) or (num_instances, 6), + or (len(env_ids),) / (num_instances,) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ @@ -506,33 +529,38 @@ def write_root_com_velocity_to_sim_index( def write_root_com_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: """Set the root center of mass velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - NOTE: This sets the velocity of the root's center of mass rather than the roots frame. + + .. note:: + This sets the velocity of the root's center of mass rather than the root's frame. .. note:: This method expects full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root center of mass velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ if env_mask is not None: env_ids = wp.nonzero(env_mask) else: env_ids = self._ALL_INDICES - self.write_root_com_velocity_to_sim_index(root_velocity, env_ids=env_ids, full_data=True) + self.write_root_com_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids, full_data=True) def write_root_link_velocity_to_sim_index( self, + *, root_velocity: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, @@ -540,18 +568,21 @@ def write_root_link_velocity_to_sim_index( """Set the root link velocity over selected environment indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - NOTE: This sets the velocity of the root's frame rather than the roots center of mass. + + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. .. note:: This method expects partial data or full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: root_velocity: Root frame velocities in simulation world frame. - Shape is (len(env_ids), 6) or (num_instances, 6). + Shape is (len(env_ids), 6) or (num_instances, 6), + or (len(env_ids),) / (num_instances,) with dtype wp.spatial_vectorf. env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ @@ -592,30 +623,34 @@ def write_root_link_velocity_to_sim_index( def write_root_link_velocity_to_sim_mask( self, + *, root_velocity: torch.Tensor | wp.array, env_mask: wp.array | None = None, ) -> None: """Set the root link velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. - NOTE: This sets the velocity of the root's frame rather than the roots center of mass. + + .. note:: + This sets the velocity of the root's frame rather than the root's center of mass. .. note:: This method expects full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6). - env_mask: Environment mask. If None, then all indices are used. + root_velocity: Root frame velocities in simulation world frame. Shape is (num_instances, 6) + or (num_instances,) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ if env_mask is not None: env_ids = wp.nonzero(env_mask) else: env_ids = self._ALL_INDICES - self.write_root_link_velocity_to_sim_index(root_velocity, env_ids=env_ids, full_data=True) + self.write_root_link_velocity_to_sim_index(root_velocity=root_velocity, env_ids=env_ids, full_data=True) """ Operations - Setters. @@ -623,6 +658,7 @@ def write_root_link_velocity_to_sim_mask( def set_masses_index( self, + *, masses: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -672,6 +708,7 @@ def set_masses_index( def set_masses_mask( self, + *, masses: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -688,7 +725,7 @@ def set_masses_mask( Args: masses: Masses of all bodies. Shape is (num_instances, num_bodies). body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -700,10 +737,11 @@ def set_masses_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_masses_index(masses, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) def set_coms_index( self, + *, coms: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -752,6 +790,7 @@ def set_coms_index( def set_coms_mask( self, + *, coms: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -768,7 +807,7 @@ def set_coms_mask( Args: coms: Center of mass pose of all bodies. Shape is (num_instances, num_bodies, 7). body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -780,10 +819,11 @@ def set_coms_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_coms_index(coms, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) def set_inertias_index( self, + *, inertias: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -831,6 +871,7 @@ def set_inertias_index( def set_inertias_mask( self, + *, inertias: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -847,7 +888,7 @@ def set_inertias_mask( Args: inertias: Inertias of all bodies. Shape is (num_instances, num_bodies, 9). body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -859,7 +900,7 @@ def set_inertias_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_inertias_index(inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) """ Internal helper. @@ -1029,8 +1070,8 @@ def write_root_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) def write_root_com_state_to_sim( self, root_state: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None @@ -1043,8 +1084,8 @@ def write_root_com_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_com_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_com_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + self.write_root_com_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_com_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) def write_root_link_state_to_sim( self, root_state: torch.Tensor | wp.array, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None @@ -1057,5 +1098,5 @@ def write_root_link_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_root_link_pose_to_sim_index(root_state[:, :7], env_ids=env_ids) - self.write_root_link_velocity_to_sim_index(root_state[:, 7:], env_ids=env_ids) + self.write_root_link_pose_to_sim_index(root_pose=root_state[:, :7], env_ids=env_ids) + self.write_root_link_velocity_to_sim_index(root_velocity=root_state[:, 7:], env_ids=env_ids) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object_data.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object_data.py index 89d435d76c9..6cf82afd592 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object_data.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object/rigid_object_data.py @@ -20,7 +20,7 @@ from isaaclab_physx.physics import PhysxManager as SimulationManager if TYPE_CHECKING: - from isaaclab.assets.rigid_object.rigid_object_view import RigidObjectView + import omni.physics.tensors.impl.api as physx # import logger logger = logging.getLogger(__name__) @@ -50,7 +50,7 @@ class RigidObjectData(BaseRigidObjectData): __backend_name__: str = "physx" """The name of the backend for the rigid object data.""" - def __init__(self, root_view: RigidObjectView, device: str): + def __init__(self, root_view: physx.RigidBodyView, device: str): """Initializes the rigid object data. Args: @@ -61,7 +61,7 @@ def __init__(self, root_view: RigidObjectView, device: str): # Set the root rigid body view # note: this is stored as a weak reference to avoid circular references between the asset class # and the data container. This is important to avoid memory leaks. - self._root_view: RigidObjectView = weakref.proxy(root_view) + self._root_view: physx.RigidBodyView = weakref.proxy(root_view) # Set initial time stamp self._sim_timestamp = 0.0 @@ -92,7 +92,8 @@ def is_primed(self) -> bool: def is_primed(self, value: bool) -> None: """Set whether the rigid object data is fully instantiated and ready to use. - .. note:: Once this quantity is set to True, it cannot be changed. + .. note:: + Once this quantity is set to True, it cannot be changed. Args: value: The primed state. @@ -126,8 +127,9 @@ def update(self, dt: float) -> None: @property def default_root_pose(self) -> wp.array: - """Default root pose ``[pos, quat]`` in local environment frame. Shape is (num_instances, 7). + """Default root pose ``[pos, quat]`` in local environment frame. + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). The position and quaternion are of the rigid body's actor frame. """ return self._default_root_pose @@ -148,8 +150,9 @@ def default_root_pose(self, value: wp.array) -> None: @property def default_root_vel(self) -> wp.array: - """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame. Shape is (num_instances, 6). + """Default root velocity ``[lin_vel, ang_vel]`` in local environment frame. + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). The linear and angular velocities are of the rigid body's center of mass frame. """ return self._default_root_vel @@ -168,66 +171,17 @@ def default_root_vel(self, value: wp.array) -> None: raise ValueError("The rigid object data is already primed.") self._default_root_vel.assign(value) - @property - def default_root_state(self) -> wp.array: - """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. - - The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities - are of the center of mass frame. Shape is (num_instances, 13). - """ - warnings.warn( - "Reading the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " - "Please use the default_root_pose and default_root_vel properties instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._default_root_state is None: - self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) - wp.launch( - shared_kernels.concat_root_pose_and_vel_to_state, - dim=self._num_instances, - inputs=[ - self._default_root_pose, - self._default_root_vel, - ], - outputs=[ - self._default_root_state, - ], - device=self.device, - ) - return self._default_root_state - - @default_root_state.setter - def default_root_state(self, value: torch.Tensor) -> None: - """Set the default root state. - - Args: - value: The default root state. Shape is (num_instances, 13). - - Raises: - ValueError: If the rigid object data is already primed. - """ - warnings.warn( - "Setting the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " - "Please use the default_root_pose and default_root_vel properties instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._is_primed: - raise ValueError("The rigid object data is already primed.") - self._default_root_pose.assign(wp.from_torch(value[:, :7]).view(wp.transformf)) - self._default_root_vel.assign(wp.from_torch(value[:, 7:]).view(wp.spatial_vectorf)) - """ Root state properties. """ @property def root_link_pose_w(self) -> wp.array: - """Root link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + """Root link pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the actor frame of the root rigid body relative to the world. - The orientation is provided in (w, x, y, z) format. + The orientation is provided in (x, y, z, w) format. """ if self._root_link_pose_w.timestamp < self._sim_timestamp: # read data from simulation @@ -238,8 +192,9 @@ def root_link_pose_w(self) -> wp.array: @property def root_link_vel_w(self) -> wp.array: - """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + """Root link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. """ @@ -264,10 +219,11 @@ def root_link_vel_w(self) -> wp.array: @property def root_com_pose_w(self) -> wp.array: - """Root center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 7). + """Root center of mass pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances,), dtype = wp.transformf. In torch this resolves to (num_instances, 7). This quantity is the pose of the center of mass frame of the root rigid body relative to the world. - The orientation is provided in (w, x, y, z) format. + The orientation is provided in (x, y, z, w) format. """ if self._root_com_pose_w.timestamp < self._sim_timestamp: # apply local transform to center of mass frame @@ -289,8 +245,9 @@ def root_com_pose_w(self) -> wp.array: @property def root_com_vel_w(self) -> wp.array: - """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 6). + """Root center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances,), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. """ @@ -300,111 +257,43 @@ def root_com_vel_w(self) -> wp.array: return self._root_com_vel_w.data - @property - def root_state_w(self) -> wp.array: - """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" - warnings.warn( - "The `root_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " - "`root_com_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._root_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_root_pose_and_vel_to_state, - dim=self._num_instances, - inputs=[ - self.root_link_pose_w, - self.root_com_vel_w, - ], - outputs=[ - self._root_state_w.data, - ], - device=self.device, - ) - self._root_state_w.timestamp = self._sim_timestamp - - return self._root_state_w.data - - @property - def root_link_state_w(self) -> wp.array: - """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" - warnings.warn( - "The `root_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " - "`root_link_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._root_link_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_root_pose_and_vel_to_state, - dim=self._num_instances, - inputs=[ - self.root_link_pose_w, - self.root_link_vel_w, - ], - outputs=[ - self._root_link_state_w.data, - ], - device=self.device, - ) - self._root_link_state_w.timestamp = self._sim_timestamp - - return self._root_link_state_w.data - - @property - def root_com_state_w(self) -> wp.array: - """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" - warnings.warn( - "The `root_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_com_pose_w` and " - "`root_com_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._root_com_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_root_pose_and_vel_to_state, - dim=self._num_instances, - inputs=[ - self.root_com_pose_w, - self.root_com_vel_w, - ], - outputs=[ - self._root_com_state_w.data, - ], - device=self.device, - ) - self._root_com_state_w.timestamp = self._sim_timestamp - - return self._root_com_state_w.data - """ Body state properties. """ @property def body_mass(self) -> wp.array: - """Mass of all bodies in the simulation world frame. Shape is (num_instances, 1, 1).""" + """Mass of all bodies in the simulation world frame. + + Shape is (num_instances, 1, 1), dtype = wp.float32. + In torch this resolves to (num_instances, 1, 1). + """ return self._body_mass @property def body_inertia(self) -> wp.array: - """Inertia of all bodies in the simulation world frame. Shape is (num_instances, 1, 3, 3).""" + """Inertia of all bodies in the simulation world frame. + + Shape is (num_instances, 1, 9), dtype = wp.float32. + In torch this resolves to (num_instances, 1, 9). + """ return self._body_inertia @property def body_link_pose_w(self) -> wp.array: - """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1, 7). + """Body link pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). This quantity is the pose of the actor frame of the rigid body relative to the world. - The orientation is provided in (w, x, y, z) format. + The orientation is provided in (x, y, z, w) format. """ return self.root_link_pose_w.reshape((self._num_instances, 1)) @property def body_link_vel_w(self) -> wp.array: - """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, 1, 6). + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. """ @@ -412,106 +301,29 @@ def body_link_vel_w(self) -> wp.array: @property def body_com_pose_w(self) -> wp.array: - """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, 1, 7). + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the world. - The orientation is provided in (w, x, y, z) format. + The orientation is provided in (x, y, z, w) format. """ return self.root_com_pose_w.reshape((self._num_instances, 1)) @property def body_com_vel_w(self) -> wp.array: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, 1, 6). + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. """ return self.root_com_vel_w.reshape((self._num_instances, 1)) - @property - def body_state_w(self) -> wp.array: - """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" - warnings.warn( - "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " - "`body_com_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - # Access internal buffer directly to avoid cascading deprecation warnings from root_state_w - if self._root_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_root_pose_and_vel_to_state, - dim=self._num_instances, - inputs=[ - self.root_link_pose_w, - self.root_com_vel_w, - ], - outputs=[ - self._root_state_w.data, - ], - device=self.device, - ) - self._root_state_w.timestamp = self._sim_timestamp - return self._root_state_w.data.reshape((self._num_instances, 1)) - - @property - def body_link_state_w(self) -> wp.array: - """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" - warnings.warn( - "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " - "`body_link_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - # Access internal buffer directly to avoid cascading deprecation warnings from root_link_state_w - if self._root_link_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_root_pose_and_vel_to_state, - dim=self._num_instances, - inputs=[ - self.root_link_pose_w, - self.root_link_vel_w, - ], - outputs=[ - self._root_link_state_w.data, - ], - device=self.device, - ) - self._root_link_state_w.timestamp = self._sim_timestamp - return self._root_link_state_w.data.reshape((self._num_instances, 1)) - - @property - def body_com_state_w(self) -> wp.array: - """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" - warnings.warn( - "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " - "`body_com_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - # Access internal buffer directly to avoid cascading deprecation warnings from root_com_state_w - if self._root_com_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_root_pose_and_vel_to_state, - dim=self._num_instances, - inputs=[ - self.root_com_pose_w, - self.root_com_vel_w, - ], - outputs=[ - self._root_com_state_w.data, - ], - device=self.device, - ) - self._root_com_state_w.timestamp = self._sim_timestamp - return self._root_com_state_w.data.reshape((self._num_instances, 1)) - @property def body_com_acc_w(self) -> wp.array: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. - Shape is (num_instances, 1, 6). + Shape is (num_instances, 1), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, 1, 6). This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. """ if self._body_com_acc_w.timestamp < self._sim_timestamp: @@ -525,10 +337,10 @@ def body_com_acc_w(self) -> wp.array: @property def body_com_pose_b(self) -> wp.array: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, 1, 7). + Shape is (num_instances, 1), dtype = wp.transformf. In torch this resolves to (num_instances, 1, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. - The orientation is provided in (w, x, y, z) format. + The orientation is provided in (x, y, z, w) format. """ if self._body_com_pose_b.timestamp < self._sim_timestamp: # read data from simulation @@ -545,7 +357,10 @@ def body_com_pose_b(self) -> wp.array: @property def projected_gravity_b(self) -> wp.array: - """Projection of the gravity direction on base frame. Shape is (num_instances, 3).""" + """Projection of the gravity direction on base frame. + + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + """ if self._projected_gravity_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_1D_kernel, @@ -559,7 +374,9 @@ def projected_gravity_b(self) -> wp.array: @property def heading_w(self) -> wp.array: - """Yaw heading of the base frame (in radians). Shape is (num_instances,). + """Yaw heading of the base frame (in radians). + + Shape is (num_instances,), dtype = wp.float32. In torch this resolves to (num_instances,). .. note:: This quantity is computed by assuming that the forward-direction of the base @@ -578,8 +395,9 @@ def heading_w(self) -> wp.array: @property def root_link_lin_vel_b(self) -> wp.array: - """Root link linear velocity in base frame. Shape is (num_instances, 3). + """Root link linear velocity in base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ @@ -596,8 +414,9 @@ def root_link_lin_vel_b(self) -> wp.array: @property def root_link_ang_vel_b(self) -> wp.array: - """Root link angular velocity in base world frame. Shape is (num_instances, 3). + """Root link angular velocity in base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ @@ -614,8 +433,9 @@ def root_link_ang_vel_b(self) -> wp.array: @property def root_com_lin_vel_b(self) -> wp.array: - """Root center of mass linear velocity in base frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ @@ -632,8 +452,9 @@ def root_com_lin_vel_b(self) -> wp.array: @property def root_com_ang_vel_b(self) -> wp.array: - """Root center of mass angular velocity in base world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in base frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ @@ -654,144 +475,162 @@ def root_com_ang_vel_b(self) -> wp.array: @property def root_link_pos_w(self) -> wp.array: - """Root link position in simulation world frame. Shape is (num_instances, 3). + """Root link position in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the position of the actor frame of the root rigid body relative to the world. """ return self._get_pos_from_transform(self.root_link_pose_w) @property def root_link_quat_w(self) -> wp.array: - """Root link orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + """Root link orientation (x, y, z, w) in simulation world frame. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). This quantity is the orientation of the actor frame of the root rigid body. """ return self._get_quat_from_transform(self.root_link_pose_w) @property def root_link_lin_vel_w(self) -> wp.array: - """Root linear velocity in simulation world frame. Shape is (num_instances, 3). + """Root linear velocity in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's actor frame relative to the world. """ return self._get_lin_vel_from_spatial_vector(self.root_link_vel_w) @property def root_link_ang_vel_w(self) -> wp.array: - """Root link angular velocity in simulation world frame. Shape is (num_instances, 3). + """Root link angular velocity in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the actor frame of the root rigid body relative to the world. """ return self._get_ang_vel_from_spatial_vector(self.root_link_vel_w) @property def root_com_pos_w(self) -> wp.array: - """Root center of mass position in simulation world frame. Shape is (num_instances, 3). + """Root center of mass position in simulation world frame. - This quantity is the position of the actor frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). + This quantity is the position of the center of mass frame of the root rigid body relative to the world. """ return self._get_pos_from_transform(self.root_com_pose_w) @property def root_com_quat_w(self) -> wp.array: - """Root center of mass orientation (w, x, y, z) in simulation world frame. Shape is (num_instances, 4). + """Root center of mass orientation (x, y, z, w) in simulation world frame. - This quantity is the orientation of the actor frame of the root rigid body relative to the world. + Shape is (num_instances,), dtype = wp.quatf. In torch this resolves to (num_instances, 4). + This quantity is the orientation of the principal axes of inertia of the root rigid body relative to the world. """ return self._get_quat_from_transform(self.root_com_pose_w) @property def root_com_lin_vel_w(self) -> wp.array: - """Root center of mass linear velocity in simulation world frame. Shape is (num_instances, 3). + """Root center of mass linear velocity in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the linear velocity of the root rigid body's center of mass frame relative to the world. """ return self._get_lin_vel_from_spatial_vector(self.root_com_vel_w) @property def root_com_ang_vel_w(self) -> wp.array: - """Root center of mass angular velocity in simulation world frame. Shape is (num_instances, 3). + """Root center of mass angular velocity in simulation world frame. + Shape is (num_instances,), dtype = wp.vec3f. In torch this resolves to (num_instances, 3). This quantity is the angular velocity of the root rigid body's center of mass frame relative to the world. """ return self._get_ang_vel_from_spatial_vector(self.root_com_vel_w) @property def body_link_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Positions of all bodies in simulation world frame. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ return self._get_pos_from_transform(self.body_link_pose_w) @property def body_link_quat_w(self) -> wp.array: - """Orientation (w, x, y, z) of all bodies in simulation world frame. Shape is (num_instances, 1, 4). + """Orientation (x, y, z, w) of all bodies in simulation world frame. - This quantity is the orientation of the rigid bodies' actor frame relative to the world. + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ return self._get_quat_from_transform(self.body_link_pose_w) @property def body_link_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear velocity of all bodies in simulation world frame. - This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. """ return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) @property def body_link_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular velocity of all bodies in simulation world frame. - This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. """ return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) @property def body_com_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Positions of all bodies' center of mass in simulation world frame. - This quantity is the position of the rigid bodies' actor frame. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the position of the rigid bodies' center of mass frame. """ return self._get_pos_from_transform(self.body_com_pose_w) @property def body_com_quat_w(self) -> wp.array: - """Orientation (w, x, y, z) of the principle axis of inertia of all bodies in simulation world frame. + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. - Shape is (num_instances, 1, 4). This quantity is the orientation of the rigid bodies' actor frame. + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the principal axes of inertia of the rigid bodies. """ return self._get_quat_from_transform(self.body_com_pose_w) @property def body_com_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear velocity of all bodies in simulation world frame. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ return self._get_lin_vel_from_spatial_vector(self.body_com_vel_w) @property def body_com_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular velocity of all bodies in simulation world frame. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ return self._get_ang_vel_from_spatial_vector(self.body_com_vel_w) @property def body_com_lin_acc_w(self) -> wp.array: - """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Linear acceleration of all bodies in simulation world frame. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ return self._get_lin_vel_from_spatial_vector(self.body_com_acc_w) @property def body_com_ang_acc_w(self) -> wp.array: - """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, 1, 3). + """Angular acceleration of all bodies in simulation world frame. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ return self._get_ang_vel_from_spatial_vector(self.body_com_acc_w) @@ -799,18 +638,19 @@ def body_com_ang_acc_w(self) -> wp.array: @property def body_com_pos_b(self) -> wp.array: """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, 1, 3). - This quantity is the center of mass location relative to its body'slink frame. + Shape is (num_instances, 1), dtype = wp.vec3f. In torch this resolves to (num_instances, 1, 3). + This quantity is the center of mass location relative to its body's link frame. """ return self._get_pos_from_transform(self.body_com_pose_b) @property def body_com_quat_b(self) -> wp.array: - """Orientation (w, x, y, z) of the principle axis of inertia of all of the bodies in their - respective link frames. Shape is (num_instances, 1, 4). + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their + respective link frames. - This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + Shape is (num_instances, 1), dtype = wp.quatf. In torch this resolves to (num_instances, 1, 4). + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ return self._get_quat_from_transform(self.body_com_pose_b) @@ -890,3 +730,192 @@ def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array strides=spatial_vector.strides, device=self.device, ) + + """ + Deprecated properties. + """ + + @property + def default_root_state(self) -> wp.array: + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. + + The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities + are of the center of mass frame. Shape is (num_instances, 13). + """ + warnings.warn( + "Reading the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " + "Please use the default_root_pose and default_root_vel properties instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._default_root_state is None: + self._default_root_state = wp.zeros((self._num_instances), dtype=shared_kernels.vec13f, device=self.device) + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self._default_root_pose, + self._default_root_vel, + ], + outputs=[ + self._default_root_state, + ], + device=self.device, + ) + return self._default_root_state + + @property + def root_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_com_vel_w`.""" + warnings.warn( + "The `root_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_state_w.data, + ], + device=self.device, + ) + self._root_state_w.timestamp = self._sim_timestamp + + return self._root_state_w.data + + @property + def root_link_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_link_pose_w` and :attr:`root_link_vel_w`.""" + warnings.warn( + "The `root_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_link_pose_w` and " + "`root_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_link_vel_w, + ], + outputs=[ + self._root_link_state_w.data, + ], + device=self.device, + ) + self._root_link_state_w.timestamp = self._sim_timestamp + + return self._root_link_state_w.data + + @property + def root_com_state_w(self) -> wp.array: + """Deprecated, same as :attr:`root_com_pose_w` and :attr:`root_com_vel_w`.""" + warnings.warn( + "The `root_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `root_com_pose_w` and " + "`root_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._root_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_com_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_com_state_w.data, + ], + device=self.device, + ) + self._root_com_state_w.timestamp = self._sim_timestamp + + return self._root_com_state_w.data + + @property + def body_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + # Access internal buffer directly to avoid cascading deprecation warnings from root_state_w + if self._root_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_state_w.data, + ], + device=self.device, + ) + self._root_state_w.timestamp = self._sim_timestamp + return self._root_state_w.data.reshape((self._num_instances, 1)) + + @property + def body_link_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + # Access internal buffer directly to avoid cascading deprecation warnings from root_link_state_w + if self._root_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_link_pose_w, + self.root_link_vel_w, + ], + outputs=[ + self._root_link_state_w.data, + ], + device=self.device, + ) + self._root_link_state_w.timestamp = self._sim_timestamp + return self._root_link_state_w.data.reshape((self._num_instances, 1)) + + @property + def body_com_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + # Access internal buffer directly to avoid cascading deprecation warnings from root_com_state_w + if self._root_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_root_pose_and_vel_to_state, + dim=self._num_instances, + inputs=[ + self.root_com_pose_w, + self.root_com_vel_w, + ], + outputs=[ + self._root_com_state_w.data, + ], + device=self.device, + ) + self._root_com_state_w.timestamp = self._sim_timestamp + return self._root_com_state_w.data.reshape((self._num_instances, 1)) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py index c2163203dc0..0c2e37d94c7 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection.py @@ -131,19 +131,36 @@ def root_view(self): @property def instantaneous_wrench_composer(self) -> WrenchComposer: - """Instantaneous wrench composer for the rigid object collection.""" + """Instantaneous wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set + to this object are discarded. This is useful to apply forces that change all the time, things like drag forces + for instance. + """ return self._instantaneous_wrench_composer @property def permanent_wrench_composer(self) -> WrenchComposer: - """Permanent wrench composer for the rigid object collection.""" + """Permanent wrench composer. + + Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench + composer are persistent and are applied to the simulation at every step. This is useful to apply forces that + are constant over a period of time, things like the thrust of a motor for instance. + """ return self._permanent_wrench_composer """ Operations. """ - def reset(self, env_ids: torch.Tensor | None = None, object_ids: slice | torch.Tensor | None = None) -> None: + def reset( + self, + env_ids: torch.Tensor | None = None, + object_ids: slice | torch.Tensor | None = None, + env_mask: wp.array | None = None, + object_mask: wp.array | None = None, + ) -> None: """Resets all internal buffers of selected environments and objects. Args: @@ -170,7 +187,7 @@ def write_data_to_sim(self) -> None: if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active: if self._instantaneous_wrench_composer.active: # Compose instantaneous wrench with permanent wrench - self._instantaneous_wrench_composer.add_forces_and_torques( + self._instantaneous_wrench_composer.add_forces_and_torques_index( forces=self._permanent_wrench_composer.composed_force, torques=self._permanent_wrench_composer.composed_torque, body_ids=self._ALL_BODY_INDICES, @@ -234,7 +251,7 @@ def find_bodies( Returns: A tuple of lists containing the body indices and names. """ - obj_ids, obj_names = string_utils.resolve_matching_names(name_keys, self.object_names, preserve_order) + obj_ids, obj_names = string_utils.resolve_matching_names(name_keys, self.body_names, preserve_order) return torch.tensor(obj_ids, device=self.device, dtype=torch.int32), obj_names """ @@ -243,9 +260,10 @@ def find_bodies( def write_body_pose_to_sim_index( self, + *, body_poses: torch.Tensor | wp.array, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the body pose over selected environment and body indices into the simulation. @@ -255,22 +273,23 @@ def write_body_pose_to_sim_index( This method expects partial data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - body_poses: Body poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7). - env_ids: Environment indices. If None, then all indices are used. + body_poses: Body poses in simulation frame. Shape is (len(env_ids), len(body_ids), 7) + or (len(env_ids), len(body_ids)) with dtype wp.transformf. body_ids: Body indices. If None, then all indices are used. - full_data: Whether to expect full data. Defaults to False. + env_ids: Environment indices. If None, then all indices are used. """ - self.write_body_link_pose_to_sim_index(body_poses, env_ids=env_ids, body_ids=body_ids) + self.write_body_link_pose_to_sim_index(body_poses=body_poses, env_ids=env_ids, body_ids=body_ids) def write_body_pose_to_sim_mask( self, + *, body_poses: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set the body pose over selected environment mask into the simulation. @@ -280,12 +299,14 @@ def write_body_pose_to_sim_mask( This method expects full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: - body_poses: Body poses in simulation frame. Shape is (num_instances, num_bodies, 7). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + body_poses: Body poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ if env_mask is not None: env_ids = wp.nonzero(env_mask) @@ -295,55 +316,67 @@ def write_body_pose_to_sim_mask( body_ids = wp.nonzero(body_mask) else: body_ids = self._ALL_BODY_INDICES - self.write_body_link_pose_to_sim_index(body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True) + self.write_body_link_pose_to_sim_index( + body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True + ) def write_body_velocity_to_sim_index( self, + *, body_velocities: torch.Tensor | wp.array, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, ) -> None: """Set the body velocity over selected environment and body indices into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + .. note:: This method expects partial data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: body_velocities: Body velocities in simulation frame. - Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6). - env_ids: Environment indices. If None, then all indices are used. + Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.spatial_vectorf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. """ - self.write_body_com_velocity_to_sim_index(body_velocities, env_ids=env_ids, body_ids=body_ids) + self.write_body_com_velocity_to_sim_index(body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids) def write_body_velocity_to_sim_mask( self, + *, body_velocities: torch.Tensor | wp.array, - env_mask: wp.array | None = None, body_mask: wp.array | None = None, + env_mask: wp.array | None = None, ) -> None: """Set the body velocity over selected environment mask into the simulation. The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + .. note:: This method expects full data. .. tip:: - For maximum performance we recommend looking at the actual implementation of the method in the backend. - Some backends may provide optimized implementations for masks / indices. + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. Args: body_velocities: Body velocities in simulation frame. - Shape is (num_instances, num_bodies, 6). - env_mask: Environment mask. If None, then all indices are used. - body_mask: Body mask. If None, then all bodies are used. + Shape is (num_instances, num_bodies, 6) + or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ if env_mask is not None: env_ids = wp.nonzero(env_mask) @@ -353,22 +386,35 @@ def write_body_velocity_to_sim_mask( body_ids = wp.nonzero(body_mask) else: body_ids = self._ALL_BODY_INDICES - self.write_body_com_velocity_to_sim_index(body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True) + self.write_body_com_velocity_to_sim_index( + body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True + ) def write_body_link_pose_to_sim_index( self, + *, body_poses: torch.Tensor | wp.array, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, ) -> None: """Set the body link pose over selected environment and body indices into the simulation. + The body link pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects partial data. + + .. tip:: + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. + Args: body_poses: Body link poses in simulation frame. - Shape is (len(env_ids), len(body_ids), 7) or (num_instances, num_bodies, 7). - env_ids: Environment indices. If None, then all indices are used. + Shape is (len(env_ids), len(body_ids), 7) or (num_instances, num_bodies, 7), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.transformf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ env_ids = self._resolve_env_ids(env_ids) @@ -405,37 +451,62 @@ def write_body_link_pose_to_sim_index( def write_body_link_pose_to_sim_mask( self, + *, body_poses: torch.Tensor | wp.array, env_mask: wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, ) -> None: """Set the body link pose over selected environment mask into the simulation. + The body link pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + + .. note:: + This method expects full data. + + .. tip:: + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. + Args: - body_poses: Body link poses in simulation frame. Shape is (num_instances, num_bodies, 7). - env_mask: Environment mask. If None, then all indices are used. + body_poses: Body link poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). body_ids: Body indices. If None, then all indices are used. """ if env_mask is not None: env_ids = wp.nonzero(env_mask) else: env_ids = self._ALL_ENV_INDICES - self.write_body_link_pose_to_sim_index(body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True) + self.write_body_link_pose_to_sim_index( + body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True + ) def write_body_com_pose_to_sim_index( self, + *, body_poses: torch.Tensor | wp.array, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, ) -> None: """Set the body center of mass pose over selected environment and body indices into the simulation. + The body center of mass pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects partial data. + + .. tip:: + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. + Args: body_poses: Body center of mass poses in simulation frame. - Shape is (len(env_ids), len(body_ids), 7) or (num_instances, num_bodies, 7). - env_ids: Environment indices. If None, then all indices are used. + Shape is (len(env_ids), len(body_ids), 7) or (num_instances, num_bodies, 7), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.transformf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ env_ids = self._resolve_env_ids(env_ids) @@ -475,37 +546,63 @@ def write_body_com_pose_to_sim_index( def write_body_com_pose_to_sim_mask( self, + *, body_poses: torch.Tensor | wp.array, env_mask: wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, ) -> None: """Set the body center of mass pose over selected environment mask into the simulation. + The body center of mass pose comprises of the cartesian position and quaternion orientation in (x, y, z, w). + The orientation is the orientation of the principal axes of inertia. + + .. note:: + This method expects full data. + + .. tip:: + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. + Args: - body_poses: Body center of mass poses in simulation frame. Shape is (num_instances, num_bodies, 7). - env_mask: Environment mask. If None, then all indices are used. + body_poses: Body center of mass poses in simulation frame. Shape is (num_instances, num_bodies, 7) + or (num_instances, num_bodies) with dtype wp.transformf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). body_ids: Body indices. If None, then all indices are used. """ if env_mask is not None: env_ids = wp.nonzero(env_mask) else: env_ids = self._ALL_ENV_INDICES - self.write_body_com_pose_to_sim_index(body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True) + self.write_body_com_pose_to_sim_index(body_poses=body_poses, env_ids=env_ids, body_ids=body_ids, full_data=True) def write_body_com_velocity_to_sim_index( self, + *, body_velocities: torch.Tensor | wp.array, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, ) -> None: """Set the body center of mass velocity over selected environment and body indices into the simulation. + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects partial data. + + .. tip:: + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. + Args: body_velocities: Body center of mass velocities in simulation frame. - Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6). - env_ids: Environment indices. If None, then all indices are used. + Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.spatial_vectorf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ env_ids = self._resolve_env_ids(env_ids) @@ -544,38 +641,68 @@ def write_body_com_velocity_to_sim_index( def write_body_com_velocity_to_sim_mask( self, + *, body_velocities: torch.Tensor | wp.array, env_mask: wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, ) -> None: """Set the body center of mass velocity over selected environment mask into the simulation. + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's center of mass rather than the body's frame. + + .. note:: + This method expects full data. + + .. tip:: + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. + Args: body_velocities: Body center of mass velocities in simulation frame. - Shape is (num_instances, num_bodies, 6). - env_mask: Environment mask. If None, then all indices are used. + Shape is (num_instances, num_bodies, 6) + or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). body_ids: Body indices. If None, then all indices are used. """ if env_mask is not None: env_ids = wp.nonzero(env_mask) else: env_ids = self._ALL_ENV_INDICES - self.write_body_com_velocity_to_sim_index(body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True) + self.write_body_com_velocity_to_sim_index( + body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True + ) def write_body_link_velocity_to_sim_index( self, + *, body_velocities: torch.Tensor | wp.array, - env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, full_data: bool = False, ) -> None: """Set the body link velocity over selected environment and body indices into the simulation. + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's frame rather than the body's center of mass. + + .. note:: + This method expects partial data. + + .. tip:: + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. + Args: body_velocities: Body link velocities in simulation frame. - Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6). - env_ids: Environment indices. If None, then all indices are used. + Shape is (len(env_ids), len(body_ids), 6) or (num_instances, num_bodies, 6), + or (len(env_ids), len(body_ids)) / (num_instances, num_bodies) with dtype wp.spatial_vectorf. body_ids: Body indices. If None, then all indices are used. + env_ids: Environment indices. If None, then all indices are used. full_data: Whether to expect full data. Defaults to False. """ env_ids = self._resolve_env_ids(env_ids) @@ -619,22 +746,38 @@ def write_body_link_velocity_to_sim_index( def write_body_link_velocity_to_sim_mask( self, + *, body_velocities: torch.Tensor | wp.array, env_mask: wp.array | None = None, body_ids: Sequence[int] | torch.Tensor | wp.array | slice | None = None, ) -> None: """Set the body link velocity over selected environment mask into the simulation. + The velocity comprises linear velocity (x, y, z) and angular velocity (x, y, z) in that order. + + .. note:: + This sets the velocity of the body's frame rather than the body's center of mass. + + .. note:: + This method expects full data. + + .. tip:: + For maximum performance we recommend using the index method. This is because in PhysX, the tensor API + is only supporting indexing, hence masks need to be converted to indices. + Args: - body_velocities: Body link velocities in simulation frame. Shape is (num_instances, num_bodies, 6). - env_mask: Environment mask. If None, then all indices are used. + body_velocities: Body link velocities in simulation frame. Shape is (num_instances, num_bodies, 6) + or (num_instances, num_bodies) with dtype wp.spatial_vectorf. + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). body_ids: Body indices. If None, then all indices are used. """ if env_mask is not None: env_ids = wp.nonzero(env_mask) else: env_ids = self._ALL_ENV_INDICES - self.write_body_link_velocity_to_sim_index(body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True) + self.write_body_link_velocity_to_sim_index( + body_velocities=body_velocities, env_ids=env_ids, body_ids=body_ids, full_data=True + ) """ Operations - Setters. @@ -642,6 +785,7 @@ def write_body_link_velocity_to_sim_mask( def set_masses_index( self, + *, masses: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -689,6 +833,7 @@ def set_masses_index( def set_masses_mask( self, + *, masses: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -704,8 +849,8 @@ def set_masses_mask( Args: masses: Masses of all bodies. Shape is ``(num_instances, num_bodies)``. - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -717,10 +862,11 @@ def set_masses_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_masses_index(masses, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_masses_index(masses=masses, body_ids=body_ids, env_ids=env_ids, full_data=True) def set_coms_index( self, + *, coms: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -771,6 +917,7 @@ def set_coms_index( def set_coms_mask( self, + *, coms: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -786,8 +933,8 @@ def set_coms_mask( Args: coms: Center of mass pose of all bodies. Shape is ``(num_instances, num_bodies, 7)``. - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -799,10 +946,11 @@ def set_coms_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_coms_index(coms, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_coms_index(coms=coms, body_ids=body_ids, env_ids=env_ids, full_data=True) def set_inertias_index( self, + *, inertias: torch.Tensor | wp.array, body_ids: Sequence[int] | torch.Tensor | wp.array | None = None, env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, @@ -850,6 +998,7 @@ def set_inertias_index( def set_inertias_mask( self, + *, inertias: torch.Tensor | wp.array, body_mask: wp.array | None = None, env_mask: wp.array | None = None, @@ -865,8 +1014,8 @@ def set_inertias_mask( Args: inertias: Inertias of all bodies. Shape is ``(num_instances, num_bodies, 9)``. - body_mask: Body mask. If None, then all bodies are used. - env_mask: Environment mask. If None, then all indices are used. + body_mask: Body mask. If None, then all bodies are updated. Shape is (num_bodies,). + env_mask: Environment mask. If None, then all the instances are updated. Shape is (num_instances,). """ # Resolve masks. if env_mask is not None: @@ -878,7 +1027,7 @@ def set_inertias_mask( else: body_ids = self._ALL_BODY_INDICES # Set full data to True to ensure the right code path is taken inside the kernel. - self.set_inertias_index(inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) + self.set_inertias_index(inertias=inertias, body_ids=body_ids, env_ids=env_ids, full_data=True) """ Helper functions. @@ -1229,8 +1378,10 @@ def write_body_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_link_pose_to_sim_index(body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) - self.write_body_com_velocity_to_sim_index(body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids) + self.write_body_link_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_com_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) def write_body_com_state_to_sim( self, @@ -1246,8 +1397,10 @@ def write_body_com_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_com_pose_to_sim_index(body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) - self.write_body_com_velocity_to_sim_index(body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids) + self.write_body_com_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_com_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) def write_body_link_state_to_sim( self, @@ -1263,5 +1416,7 @@ def write_body_link_state_to_sim( DeprecationWarning, stacklevel=2, ) - self.write_body_link_pose_to_sim_index(body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) - self.write_body_link_velocity_to_sim_index(body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids) + self.write_body_link_pose_to_sim_index(body_poses=body_states[:, :, :7], env_ids=env_ids, body_ids=body_ids) + self.write_body_link_velocity_to_sim_index( + body_velocities=body_states[:, :, 7:], env_ids=env_ids, body_ids=body_ids + ) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection_data.py b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection_data.py index 443dbec3518..5d6a6835931 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection_data.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/rigid_object_collection/rigid_object_collection_data.py @@ -20,7 +20,7 @@ from isaaclab_physx.physics import PhysxManager as SimulationManager if TYPE_CHECKING: - from isaaclab.assets.rigid_object_collection.rigid_object_collection_view import RigidObjectCollectionView + import omni.physics.tensors.impl.api as physx # import logger logger = logging.getLogger(__name__) @@ -50,7 +50,7 @@ class RigidObjectCollectionData(BaseRigidObjectCollectionData): __backend_name__: str = "physx" """The name of the backend for the rigid object collection data.""" - def __init__(self, root_view: RigidObjectCollectionView, num_bodies: int, device: str): + def __init__(self, root_view: physx.RigidBodyView, num_bodies: int, device: str): """Initializes the rigid object data. Args: @@ -63,7 +63,7 @@ def __init__(self, root_view: RigidObjectCollectionView, num_bodies: int, device # Set the root rigid body view # note: this is stored as a weak reference to avoid circular references between the asset class # and the data container. This is important to avoid memory leaks. - self._root_view: RigidObjectCollectionView = weakref.proxy(root_view) + self._root_view: physx.RigidBodyView = weakref.proxy(root_view) self.num_instances = self._root_view.count // self.num_bodies # Set initial time stamp @@ -95,7 +95,8 @@ def is_primed(self) -> bool: def is_primed(self, value: bool) -> None: """Set whether the rigid object collection data is fully instantiated and ready to use. - .. note:: Once this quantity is set to True, it cannot be changed. + .. note:: + Once this quantity is set to True, it cannot be changed. Args: value: The primed state. @@ -131,7 +132,9 @@ def update(self, dt: float) -> None: def default_body_pose(self) -> wp.array: """Default body pose ``[pos, quat]`` in local environment frame. - The position and quaternion are of the rigid body's actor frame. Shape is (num_instances, num_bodies, 7). + The position and quaternion are of the rigid body's actor frame. + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). """ return self._default_body_pose @@ -151,9 +154,10 @@ def default_body_pose(self, value: wp.array) -> None: @property def default_body_vel(self) -> wp.array: - """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. Shape is + """Default body velocity ``[lin_vel, ang_vel]`` in local environment frame. - The linear and angular velocities are of the rigid body's center of mass frame. Shape is + The linear and angular velocities are of the rigid body's center of mass frame. + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to (num_instances, num_bodies, 6). """ return self._default_body_vel @@ -172,66 +176,16 @@ def default_body_vel(self, value: wp.array) -> None: raise ValueError("The rigid object collection data is already primed.") self._default_body_vel.assign(value) - @property - def default_body_state(self) -> wp.array: - """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. - - The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities - are of the center of mass frame. Shape is (num_instances, num_bodies, 13). - """ - warnings.warn( - "Reading the body state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " - "Please use the default_body_pose and default_body_vel properties instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._default_body_state is None: - self._default_body_state = wp.zeros( - (self.num_instances, self.num_bodies), dtype=shared_kernels.vec13f, device=self.device - ) - wp.launch( - shared_kernels.concat_body_pose_and_vel_to_state, - dim=(self.num_instances, self.num_bodies), - inputs=[ - self._default_body_pose, - self._default_body_vel, - ], - outputs=[ - self._default_body_state, - ], - device=self.device, - ) - return self._default_body_state - - @default_body_state.setter - def default_body_state(self, value: torch.Tensor) -> None: - """Set the default body state. - - Args: - value: The default body state. Shape is (num_instances, num_bodies, 13). - - Raises: - ValueError: If the rigid object collection data is already primed. - """ - warnings.warn( - "Setting the root state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " - "Please use the default_root_pose and default_root_vel properties instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._is_primed: - raise ValueError("The rigid object collection data is already primed.") - self._default_body_pose.assign(wp.from_torch(value[:, :, :7]).view(wp.transformf)) - self._default_body_vel.assign(wp.from_torch(value[:, :, 7:]).view(wp.spatial_vectorf)) - """ Body state properties. """ @property def body_link_pose_w(self) -> wp.array: - """Body link pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies, 7). + """Body link pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the actor frame of the rigid body relative to the world. The orientation is provided in (x, y, z, w) format. """ @@ -246,8 +200,10 @@ def body_link_pose_w(self) -> wp.array: @property def body_link_vel_w(self) -> wp.array: - """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. Shape is (num_instances, num_bodies, 6). + """Body link velocity ``[lin_vel, ang_vel]`` in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the actor frame of the root rigid body relative to the world. """ @@ -271,8 +227,10 @@ def body_link_vel_w(self) -> wp.array: @property def body_com_pose_w(self) -> wp.array: - """Body center of mass pose ``[pos, quat]`` in simulation world frame. Shape is (num_instances, num_bodies, 7). + """Body center of mass pose ``[pos, quat]`` in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the world. The orientation is provided in (x, y, z, w) format. """ @@ -296,8 +254,9 @@ def body_com_pose_w(self) -> wp.array: @property def body_com_vel_w(self) -> wp.array: """Body center of mass velocity ``[lin_vel, ang_vel]`` in simulation world frame. - Shape is (num_instances, num_bodies, 6). + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity contains the linear and angular velocities of the root rigid body's center of mass frame relative to the world. """ @@ -308,89 +267,12 @@ def body_com_vel_w(self) -> wp.array: return self._body_com_vel_w.data - @property - def body_state_w(self) -> wp.array: - """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" - warnings.warn( - "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " - "`body_com_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._body_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_body_pose_and_vel_to_state, - dim=(self.num_instances, self.num_bodies), - inputs=[ - self.body_link_pose_w, - self.body_com_vel_w, - ], - outputs=[ - self._body_state_w.data, - ], - device=self.device, - ) - self._body_state_w.timestamp = self._sim_timestamp - - return self._body_state_w.data - - @property - def body_link_state_w(self) -> wp.array: - """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" - warnings.warn( - "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " - "`body_link_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._body_link_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_body_pose_and_vel_to_state, - dim=(self.num_instances, self.num_bodies), - inputs=[ - self.body_link_pose_w, - self.body_link_vel_w, - ], - outputs=[ - self._body_link_state_w.data, - ], - device=self.device, - ) - self._body_link_state_w.timestamp = self._sim_timestamp - - return self._body_link_state_w.data - - @property - def body_com_state_w(self) -> wp.array: - """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" - warnings.warn( - "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " - "`body_com_vel_w` instead.", - DeprecationWarning, - stacklevel=2, - ) - if self._body_com_state_w.timestamp < self._sim_timestamp: - wp.launch( - shared_kernels.concat_body_pose_and_vel_to_state, - dim=(self.num_instances, self.num_bodies), - inputs=[ - self.body_com_pose_w, - self.body_com_vel_w, - ], - outputs=[ - self._body_com_state_w.data, - ], - device=self.device, - ) - self._body_com_state_w.timestamp = self._sim_timestamp - - return self._body_com_state_w.data - @property def body_com_acc_w(self) -> wp.array: """Acceleration of all bodies ``[lin_acc, ang_acc]`` in the simulation world frame. - Shape is (num_instances, num_bodies, 6). + Shape is (num_instances, num_bodies), dtype = wp.spatial_vectorf. In torch this resolves to + (num_instances, num_bodies, 6). This quantity is the acceleration of the rigid bodies' center of mass frame relative to the world. """ if self._body_com_acc_w.timestamp < self._sim_timestamp: @@ -402,8 +284,9 @@ def body_com_acc_w(self) -> wp.array: @property def body_com_pose_b(self) -> wp.array: """Center of mass pose ``[pos, quat]`` of all bodies in their respective body's link frames. - Shape is (num_instances, num_bodies, 7). + Shape is (num_instances, num_bodies), dtype = wp.transformf. In torch this resolves to + (num_instances, num_bodies, 7). This quantity is the pose of the center of mass frame of the rigid body relative to the body's link frame. The orientation is provided in (x, y, z, w) format. """ @@ -418,12 +301,20 @@ def body_com_pose_b(self) -> wp.array: @property def body_mass(self) -> wp.array: - """Mass of all bodies in the simulation world frame. Shape is (num_instances, num_bodies).""" + """Mass of all bodies in the simulation world frame. + + Shape is (num_instances, num_bodies), dtype = wp.float32. + In torch this resolves to (num_instances, num_bodies). + """ return self._body_mass @property def body_inertia(self) -> wp.array: - """Inertia of all bodies in the simulation world frame. Shape is (num_instances, num_bodies, 9).""" + """Inertia of all bodies in the simulation world frame. + + Shape is (num_instances, num_bodies, 9), dtype = wp.float32. + In torch this resolves to (num_instances, num_bodies, 9). + """ return self._body_inertia """ @@ -432,7 +323,11 @@ def body_inertia(self) -> wp.array: @property def projected_gravity_b(self) -> wp.array: - """Projection of the gravity direction on base frame. Shape is (num_instances, num_bodies, 3).""" + """Projection of the gravity direction on base frame. + + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + """ if self._projected_gravity_b.timestamp < self._sim_timestamp: wp.launch( shared_kernels.quat_apply_inverse_2D_kernel, @@ -446,7 +341,9 @@ def projected_gravity_b(self) -> wp.array: @property def heading_w(self) -> wp.array: - """Yaw heading of the base frame (in radians). Shape is (num_instances, num_bodies). + """Yaw heading of the base frame (in radians). + + Shape is (num_instances, num_bodies), dtype = wp.float32. In torch this resolves to (num_instances, num_bodies). .. note:: This quantity is computed by assuming that the forward-direction of the base @@ -465,8 +362,10 @@ def heading_w(self) -> wp.array: @property def body_link_lin_vel_b(self) -> wp.array: - """Root link linear velocity in base frame. Shape is (num_instances, num_bodies, 3). + """Root link linear velocity in base frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ @@ -483,8 +382,10 @@ def body_link_lin_vel_b(self) -> wp.array: @property def body_link_ang_vel_b(self) -> wp.array: - """Root link angular velocity in base world frame. Shape is (num_instances, num_bodies, 3). + """Root link angular velocity in base frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular velocity of the actor frame of the root rigid body frame with respect to the rigid body's actor frame. """ @@ -501,8 +402,10 @@ def body_link_ang_vel_b(self) -> wp.array: @property def body_com_lin_vel_b(self) -> wp.array: - """Root center of mass linear velocity in base frame. Shape is (num_instances, num_bodies, 3). + """Root center of mass linear velocity in base frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ @@ -519,8 +422,10 @@ def body_com_lin_vel_b(self) -> wp.array: @property def body_com_ang_vel_b(self) -> wp.array: - """Root center of mass angular velocity in base world frame. Shape is (num_instances, num_bodies, 3). + """Root center of mass angular velocity in base frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular velocity of the root rigid body's center of mass frame with respect to the rigid body's actor frame. """ @@ -541,80 +446,100 @@ def body_com_ang_vel_b(self) -> wp.array: @property def body_link_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Positions of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the position of the rigid bodies' actor frame relative to the world. """ return self._get_pos_from_transform(self.body_link_pose_w) @property def body_link_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of all bodies in simulation world frame. - This quantity is the orientation of the rigid bodies' actor frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + This quantity is the orientation of the rigid bodies' actor frame relative to the world. """ return self._get_quat_from_transform(self.body_link_pose_w) @property def body_link_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear velocity of all bodies in simulation world frame. - This quantity is the linear velocity of the rigid bodies' center of mass frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the linear velocity of the rigid bodies' actor frame relative to the world. """ return self._get_lin_vel_from_spatial_vector(self.body_link_vel_w) @property def body_link_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular velocity of all bodies in simulation world frame. - This quantity is the angular velocity of the rigid bodies' center of mass frame relative to the world. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the angular velocity of the rigid bodies' actor frame relative to the world. """ return self._get_ang_vel_from_spatial_vector(self.body_link_vel_w) @property def body_com_pos_w(self) -> wp.array: - """Positions of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Positions of all bodies' center of mass in simulation world frame. - This quantity is the position of the rigid bodies' actor frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the position of the rigid bodies' center of mass frame. """ return self._get_pos_from_transform(self.body_com_pose_w) @property def body_com_quat_w(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all bodies in simulation world frame. + """Orientation (x, y, z, w) of the principal axes of inertia of all bodies in simulation world frame. - Shape is (num_instances, num_bodies, 4). This quantity is the orientation of the rigid bodies' actor frame. + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + This quantity is the orientation of the principal axes of inertia of the rigid bodies. """ return self._get_quat_from_transform(self.body_com_pose_w) @property def body_com_lin_vel_w(self) -> wp.array: - """Linear velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear velocity of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear velocity of the rigid bodies' center of mass frame. """ return self._get_lin_vel_from_spatial_vector(self.body_com_vel_w) @property def body_com_ang_vel_w(self) -> wp.array: - """Angular velocity of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular velocity of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular velocity of the rigid bodies' center of mass frame. """ return self._get_ang_vel_from_spatial_vector(self.body_com_vel_w) @property def body_com_lin_acc_w(self) -> wp.array: - """Linear acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Linear acceleration of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the linear acceleration of the rigid bodies' center of mass frame. """ return self._get_lin_vel_from_spatial_vector(self.body_com_acc_w) @property def body_com_ang_acc_w(self) -> wp.array: - """Angular acceleration of all bodies in simulation world frame. Shape is (num_instances, num_bodies, 3). + """Angular acceleration of all bodies in simulation world frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). This quantity is the angular acceleration of the rigid bodies' center of mass frame. """ return self._get_ang_vel_from_spatial_vector(self.body_com_acc_w) @@ -622,18 +547,21 @@ def body_com_ang_acc_w(self) -> wp.array: @property def body_com_pos_b(self) -> wp.array: """Center of mass position of all of the bodies in their respective link frames. - Shape is (num_instances, num_bodies, 3). - This quantity is the center of mass location relative to its body'slink frame. + Shape is (num_instances, num_bodies), dtype = wp.vec3f. In torch this resolves to + (num_instances, num_bodies, 3). + This quantity is the center of mass location relative to its body's link frame. """ return self._get_pos_from_transform(self.body_com_pose_b) @property def body_com_quat_b(self) -> wp.array: - """Orientation (x, y, z, w) of the principle axis of inertia of all of the bodies in their - respective link frames. Shape is (num_instances, num_bodies, 4). + """Orientation (x, y, z, w) of the principal axes of inertia of all of the bodies in their + respective link frames. - This quantity is the orientation of the principles axes of inertia relative to its body's link frame. + Shape is (num_instances, num_bodies), dtype = wp.quatf. In torch this resolves to + (num_instances, num_bodies, 4). + This quantity is the orientation of the principal axes of inertia relative to its body's link frame. """ return self._get_quat_from_transform(self.body_com_pose_b) @@ -651,7 +579,7 @@ def _create_buffers(self) -> None: self._body_com_pose_w = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.transformf) self._body_com_vel_w = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.spatial_vectorf) self._body_com_acc_w = TimestampedBuffer((self.num_instances, self.num_bodies), self.device, wp.spatial_vectorf) - # -- combined state(these are cached as they concatenate) + # -- combined state (these are cached as they concatenate) self._body_state_w = TimestampedBuffer( (self.num_instances, self.num_bodies), self.device, shared_kernels.vec13f ) @@ -777,3 +705,116 @@ def _get_ang_vel_from_spatial_vector(self, spatial_vector: wp.array) -> wp.array strides=spatial_vector.strides, device=self.device, ) + + """ + Deprecated properties. + """ + + @property + def default_body_state(self) -> wp.array: + """Default root state ``[pos, quat, lin_vel, ang_vel]`` in local environment frame. + + The position and quaternion are of the rigid body's actor frame. Meanwhile, the linear and angular velocities + are of the center of mass frame. Shape is (num_instances, num_bodies, 13). + """ + warnings.warn( + "Reading the body state directly is deprecated since IsaacLab 3.0 and will be removed in a future version. " + "Please use the default_body_pose and default_body_vel properties instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._default_body_state is None: + self._default_body_state = wp.zeros( + (self.num_instances, self.num_bodies), dtype=shared_kernels.vec13f, device=self.device + ) + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self._default_body_pose, + self._default_body_vel, + ], + outputs=[ + self._default_body_state, + ], + device=self.device, + ) + return self._default_body_state + + @property + def body_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_com_vel_w, + ], + outputs=[ + self._body_state_w.data, + ], + device=self.device, + ) + self._body_state_w.timestamp = self._sim_timestamp + + return self._body_state_w.data + + @property + def body_link_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_link_pose_w` and :attr:`body_link_vel_w`.""" + warnings.warn( + "The `body_link_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_link_pose_w` and " + "`body_link_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_link_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_link_pose_w, + self.body_link_vel_w, + ], + outputs=[ + self._body_link_state_w.data, + ], + device=self.device, + ) + self._body_link_state_w.timestamp = self._sim_timestamp + + return self._body_link_state_w.data + + @property + def body_com_state_w(self) -> wp.array: + """Deprecated, same as :attr:`body_com_pose_w` and :attr:`body_com_vel_w`.""" + warnings.warn( + "The `body_com_state_w` property will be deprecated in IsaacLab 4.0. Please use `body_com_pose_w` and " + "`body_com_vel_w` instead.", + DeprecationWarning, + stacklevel=2, + ) + if self._body_com_state_w.timestamp < self._sim_timestamp: + wp.launch( + shared_kernels.concat_body_pose_and_vel_to_state, + dim=(self.num_instances, self.num_bodies), + inputs=[ + self.body_com_pose_w, + self.body_com_vel_w, + ], + outputs=[ + self._body_com_state_w.data, + ], + device=self.device, + ) + self._body_com_state_w.timestamp = self._sim_timestamp + + return self._body_com_state_w.data diff --git a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py index 4ba26e2bb12..28da4654454 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper.py @@ -17,7 +17,7 @@ import isaaclab.sim as sim_utils from isaaclab.assets import AssetBase -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit if TYPE_CHECKING: from isaacsim.robot.surface_gripper import GripperView @@ -88,7 +88,7 @@ def __init__(self, cfg: SurfaceGripperCfg): self._cfg = cfg.copy() # checks for Isaac Sim v5.0 to ensure that the surface gripper is supported - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: raise NotImplementedError( "SurfaceGrippers are only supported by IsaacSim 5.0 and newer. Current version is" f" '{get_isaac_sim_version()}'. Please update to IsaacSim 5.0 or newer to use this feature." diff --git a/source/isaaclab_physx/test/assets/test_articulation.py b/source/isaaclab_physx/test/assets/test_articulation.py index 7be1ad935b8..679e8cf5095 100644 --- a/source/isaaclab_physx/test/assets/test_articulation.py +++ b/source/isaaclab_physx/test/assets/test_articulation.py @@ -33,7 +33,7 @@ from isaaclab.managers import SceneEntityCfg from isaaclab.sim import build_simulation_context from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit ## # Pre-defined configs @@ -373,10 +373,12 @@ def test_initialization_fixed_base(sim, num_articulations, device): articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_state = wp.to_torch(articulation.data.default_root_state).clone() - default_root_state[:, :3] = default_root_state[:, :3] + translations + default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_state_w), default_root_state) + torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) + torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -437,10 +439,12 @@ def test_initialization_fixed_base_single_joint(sim, num_articulations, device, articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_state = wp.to_torch(articulation.data.default_root_state).clone() - default_root_state[:, :3] = default_root_state[:, :3] + translations + default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_state_w), default_root_state) + torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) + torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -550,10 +554,12 @@ def test_initialization_floating_base_made_fixed_base(sim, num_articulations, de articulation.update(sim.cfg.dt) # check that the root is at the correct state - its default state as it is fixed base - default_root_state = wp.to_torch(articulation.data.default_root_state).clone() - default_root_state[:, :3] = default_root_state[:, :3] + translations + default_root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + default_root_vel = wp.to_torch(articulation.data.default_root_vel).clone() + default_root_pose[:, :3] = default_root_pose[:, :3] + translations - torch.testing.assert_close(wp.to_torch(articulation.data.root_state_w), default_root_state) + torch.testing.assert_close(wp.to_torch(articulation.data.root_link_pose_w), default_root_pose) + torch.testing.assert_close(wp.to_torch(articulation.data.root_com_vel_w), default_root_vel) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -698,7 +704,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 - articulation.write_joint_position_limit_to_sim(limits) + articulation.write_joint_position_limit_to_sim_index(limits=limits) # Check new limits are in place torch.testing.assert_close(wp.to_torch(articulation._data.joint_pos_limits), limits) @@ -710,7 +716,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) limits[..., 0] = (torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0) * -1.0 limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) + 5.0 - articulation.write_joint_position_limit_to_sim(limits, env_ids=env_ids, joint_ids=joint_ids) + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) # Check new limits are in place torch.testing.assert_close(wp.to_torch(articulation._data.joint_pos_limits)[env_ids][:, joint_ids], limits) @@ -720,7 +726,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) limits[..., 0] = torch.rand(num_articulations, articulation.num_joints, device=device) * -0.1 limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) * 0.1 - articulation.write_joint_position_limit_to_sim(limits) + articulation.write_joint_position_limit_to_sim_index(limits=limits) # Check if all values are within the bounds default_joint_pos_torch = wp.to_torch(articulation._data.default_joint_pos) @@ -731,7 +737,7 @@ def test_joint_pos_limits(sim, num_articulations, device, add_ground_plane): limits = torch.zeros(env_ids.shape[0], joint_ids.shape[0], 2, device=device) limits[..., 0] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * -0.1 limits[..., 1] = torch.rand(env_ids.shape[0], joint_ids.shape[0], device=device) * 0.1 - articulation.write_joint_position_limit_to_sim(limits, env_ids=env_ids, joint_ids=joint_ids) + articulation.write_joint_position_limit_to_sim_index(limits=limits, env_ids=env_ids, joint_ids=joint_ids) # Check if all values are within the bounds default_joint_pos_torch = wp.to_torch(articulation._data.default_joint_pos) @@ -799,15 +805,16 @@ def test_external_force_buffer(sim, num_articulations, device): body_ids, _ = articulation.find_bodies("base") # reset root state - root_state = wp.to_torch(articulation.data.default_root_state).clone() - articulation.write_root_state_to_sim(root_state) + articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) + articulation.write_root_velocity_to_sim_index(root_velocity=wp.to_torch(articulation.data.default_root_vel).clone()) # reset dof state joint_pos, joint_vel = ( wp.to_torch(articulation.data.default_joint_pos), wp.to_torch(articulation.data.default_joint_vel), ) - articulation.write_joint_state_to_sim(joint_pos, joint_vel) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) # reset articulation articulation.reset() @@ -829,10 +836,9 @@ def test_external_force_buffer(sim, num_articulations, device): external_wrench_b[:, :, 3] = force # apply force - # TODO: Replace with wrench composer once the deprecation is complete - articulation.set_external_force_and_torque( - external_wrench_b[..., :3], - external_wrench_b[..., 3:], + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], body_ids=body_ids, ) @@ -849,7 +855,7 @@ def test_external_force_buffer(sim, num_articulations, device): ) # apply action to the articulation - articulation.set_joint_position_target(wp.to_torch(articulation.data.default_joint_pos).clone()) + articulation.set_joint_position_target_index(target=wp.to_torch(articulation.data.default_joint_pos).clone()) articulation.write_data_to_sim() # perform step @@ -888,27 +894,29 @@ def test_external_force_on_single_body(sim, num_articulations, device): # Now we are ready! for _ in range(5): # reset root state - root_state = wp.to_torch(articulation.data.default_root_state).clone() - - articulation.write_root_pose_to_sim(root_state[:, :7]) - articulation.write_root_velocity_to_sim(root_state[:, 7:]) + articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) + articulation.write_root_velocity_to_sim_index( + root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() + ) # reset dof state joint_pos, joint_vel = ( wp.to_torch(articulation.data.default_joint_pos), wp.to_torch(articulation.data.default_joint_vel), ) - articulation.write_joint_state_to_sim(joint_pos, joint_vel) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) # reset articulation articulation.reset() # apply force - # TODO: Replace with wrench composer once the deprecation is complete - articulation.set_external_force_and_torque( - external_wrench_b[..., :3], external_wrench_b[..., 3:], body_ids=body_ids + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids ) # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target(wp.to_torch(articulation.data.default_joint_pos).clone()) + articulation.set_joint_position_target_index( + target=wp.to_torch(articulation.data.default_joint_pos).clone() + ) articulation.write_data_to_sim() # perform step sim.step() @@ -957,17 +965,20 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic # Now we are ready! for i in range(5): # reset root state - root_state = wp.to_torch(articulation.data.default_root_state).clone() - root_state[0, 0] = 2.5 # space them apart by 2.5m + root_pose = wp.to_torch(articulation.data.default_root_pose).clone() + root_pose[0, 0] = 2.5 # space them apart by 2.5m - articulation.write_root_pose_to_sim(root_state[:, :7]) - articulation.write_root_velocity_to_sim(root_state[:, 7:]) + articulation.write_root_pose_to_sim_index(root_pose=root_pose) + articulation.write_root_velocity_to_sim_index( + root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() + ) # reset dof state joint_pos, joint_vel = ( wp.to_torch(articulation.data.default_joint_pos), wp.to_torch(articulation.data.default_joint_vel), ) - articulation.write_joint_state_to_sim(joint_pos, joint_vel) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) # reset articulation articulation.reset() # apply force @@ -985,14 +996,14 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic external_wrench_positions_b[..., 1] = 1.0 external_wrench_positions_b[..., 2] = 0.0 - articulation.permanent_wrench_composer.set_forces_and_torques( + articulation.permanent_wrench_composer.set_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=external_wrench_positions_b, body_ids=body_ids, is_global=is_global, ) - articulation.permanent_wrench_composer.add_forces_and_torques( + articulation.permanent_wrench_composer.add_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=external_wrench_positions_b, @@ -1002,7 +1013,9 @@ def test_external_force_on_single_body_at_position(sim, num_articulations, devic # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target(wp.to_torch(articulation.data.default_joint_pos).clone()) + articulation.set_joint_position_target_index( + target=wp.to_torch(articulation.data.default_joint_pos).clone() + ) articulation.write_data_to_sim() # perform step sim.step() @@ -1043,25 +1056,29 @@ def test_external_force_on_multiple_bodies(sim, num_articulations, device): # Now we are ready! for _ in range(5): # reset root state - articulation.write_root_pose_to_sim(wp.to_torch(articulation.data.default_root_state).clone()[:, :7]) - articulation.write_root_velocity_to_sim(wp.to_torch(articulation.data.default_root_state).clone()[:, 7:]) + articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) + articulation.write_root_velocity_to_sim_index( + root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() + ) # reset dof state joint_pos, joint_vel = ( wp.to_torch(articulation.data.default_joint_pos), wp.to_torch(articulation.data.default_joint_vel), ) - articulation.write_joint_state_to_sim(joint_pos, joint_vel) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) # reset articulation articulation.reset() # apply force - # TODO: Replace with wrench composer once the deprecation is complete - articulation.set_external_force_and_torque( - external_wrench_b[..., :3], external_wrench_b[..., 3:], body_ids=body_ids + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids ) # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target(wp.to_torch(articulation.data.default_joint_pos).clone()) + articulation.set_joint_position_target_index( + target=wp.to_torch(articulation.data.default_joint_pos).clone() + ) articulation.write_data_to_sim() # perform step sim.step() @@ -1112,14 +1129,17 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d # Now we are ready! for i in range(5): # reset root state - articulation.write_root_pose_to_sim(wp.to_torch(articulation.data.default_root_state).clone()[:, :7]) - articulation.write_root_velocity_to_sim(wp.to_torch(articulation.data.default_root_state).clone()[:, 7:]) + articulation.write_root_pose_to_sim_index(root_pose=wp.to_torch(articulation.data.default_root_pose).clone()) + articulation.write_root_velocity_to_sim_index( + root_velocity=wp.to_torch(articulation.data.default_root_vel).clone() + ) # reset dof state joint_pos, joint_vel = ( wp.to_torch(articulation.data.default_joint_pos), wp.to_torch(articulation.data.default_joint_vel), ) - articulation.write_joint_state_to_sim(joint_pos, joint_vel) + articulation.write_joint_position_to_sim_index(position=joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=joint_vel) # reset articulation articulation.reset() @@ -1137,14 +1157,14 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d external_wrench_positions_b[..., 2] = 0.0 # apply force - articulation.permanent_wrench_composer.set_forces_and_torques( + articulation.permanent_wrench_composer.set_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=external_wrench_positions_b, body_ids=body_ids, is_global=is_global, ) - articulation.permanent_wrench_composer.add_forces_and_torques( + articulation.permanent_wrench_composer.add_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=external_wrench_positions_b, @@ -1154,7 +1174,9 @@ def test_external_force_on_multiple_bodies_at_position(sim, num_articulations, d # perform simulation for _ in range(100): # apply action to the articulation - articulation.set_joint_position_target(wp.to_torch(articulation.data.default_joint_pos).clone()) + articulation.set_joint_position_target_index( + target=wp.to_torch(articulation.data.default_joint_pos).clone() + ) articulation.write_data_to_sim() # perform step sim.step() @@ -1560,12 +1582,11 @@ def test_reset(sim, num_articulations, device): if num_articulations > 1: num_bodies = articulation.num_bodies - # TODO: Replace with wrench composer once the deprecation is complete - articulation.set_external_force_and_torque( + articulation.permanent_wrench_composer.set_forces_and_torques_index( forces=torch.ones((num_articulations, num_bodies, 3), device=device), torques=torch.ones((num_articulations, num_bodies, 3), device=device), ) - articulation.instantaneous_wrench_composer.add_forces_and_torques( + articulation.instantaneous_wrench_composer.add_forces_and_torques_index( forces=torch.ones((num_articulations, num_bodies, 3), device=device), torques=torch.ones((num_articulations, num_bodies, 3), device=device), ) @@ -1613,7 +1634,7 @@ def test_apply_joint_command(sim, num_articulations, device, add_ground_plane): joint_pos[:, 3] = 0.0 # apply action to the articulation - articulation.set_joint_position_target(joint_pos) + articulation.set_joint_position_target_index(target=joint_pos) articulation.write_data_to_sim() for _ in range(100): @@ -1685,12 +1706,14 @@ def test_body_root_state(sim, num_articulations, device, with_offset): articulation.update(sim.cfg.dt) # get state properties - root_state_w = wp.to_torch(articulation.data.root_state_w) - root_link_state_w = wp.to_torch(articulation.data.root_link_state_w) - root_com_state_w = wp.to_torch(articulation.data.root_com_state_w) - body_state_w = wp.to_torch(articulation.data.body_state_w) - body_link_state_w = wp.to_torch(articulation.data.body_link_state_w) - body_com_state_w = wp.to_torch(articulation.data.body_com_state_w) + root_link_pose_w = wp.to_torch(articulation.data.root_link_pose_w) + root_link_vel_w = wp.to_torch(articulation.data.root_link_vel_w) + root_com_pose_w = wp.to_torch(articulation.data.root_com_pose_w) + root_com_vel_w = wp.to_torch(articulation.data.root_com_vel_w) + body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) + body_link_vel_w = wp.to_torch(articulation.data.body_link_vel_w) + body_com_pose_w = wp.to_torch(articulation.data.body_com_pose_w) + body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w) if with_offset: # get joint state @@ -1699,8 +1722,8 @@ def test_body_root_state(sim, num_articulations, device, with_offset): # LINK state # pose - torch.testing.assert_close(root_state_w[..., :7], root_link_state_w[..., :7]) - torch.testing.assert_close(body_state_w[..., :7], body_link_state_w[..., :7]) + torch.testing.assert_close(root_link_pose_w, root_link_pose_w) + torch.testing.assert_close(body_link_pose_w, body_link_pose_w) # lin_vel arm lin_vel_gt = torch.zeros(num_articulations, num_bodies, 3, device=device) @@ -1710,13 +1733,13 @@ def test_body_root_state(sim, num_articulations, device, with_offset): lin_vel_gt[:, 1, :] = torch.cat([vx, vy, vz], dim=-1).squeeze(-2) # linear velocity of root link should be zero - torch.testing.assert_close(lin_vel_gt[:, 0, :], root_link_state_w[..., 7:10], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(lin_vel_gt[:, 0, :], root_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) # linear velocity of pendulum link should be - torch.testing.assert_close(lin_vel_gt, body_link_state_w[..., 7:10], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(lin_vel_gt, body_link_vel_w[..., :3], atol=1e-3, rtol=1e-1) # ang_vel - torch.testing.assert_close(root_state_w[..., 10:], root_link_state_w[..., 10:]) - torch.testing.assert_close(body_state_w[..., 10:], body_link_state_w[..., 10:]) + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) # COM state # position and orientation shouldn't match for the _state_com_w but everything else will @@ -1726,24 +1749,24 @@ def test_body_root_state(sim, num_articulations, device, with_offset): pz = (link_offset[0] + offset[0]) * torch.sin(joint_pos) pos_gt[:, 1, :] = torch.cat([px, py, pz], dim=-1).squeeze(-2) pos_gt += env_pos.unsqueeze(-2).repeat(1, num_bodies, 1) - torch.testing.assert_close(pos_gt[:, 0, :], root_com_state_w[..., :3], atol=1e-3, rtol=1e-1) - torch.testing.assert_close(pos_gt, body_com_state_w[..., :3], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(pos_gt[:, 0, :], root_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) + torch.testing.assert_close(pos_gt, body_com_pose_w[..., :3], atol=1e-3, rtol=1e-1) # orientation com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) - com_quat_w = math_utils.quat_mul(body_link_state_w[..., 3:7], com_quat_b) - torch.testing.assert_close(com_quat_w, body_com_state_w[..., 3:7]) - torch.testing.assert_close(com_quat_w[:, 0, :], root_com_state_w[..., 3:7]) + com_quat_w = math_utils.quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w[:, 0, :], root_com_pose_w[..., 3:]) # linear vel, and angular vel - torch.testing.assert_close(root_state_w[..., 7:], root_com_state_w[..., 7:]) - torch.testing.assert_close(body_state_w[..., 7:], body_com_state_w[..., 7:]) + torch.testing.assert_close(root_com_vel_w, root_com_vel_w) + torch.testing.assert_close(body_com_vel_w, body_com_vel_w) else: # single joint center of masses are at link frames so they will be the same - torch.testing.assert_close(root_state_w, root_link_state_w) - torch.testing.assert_close(root_state_w, root_com_state_w) - torch.testing.assert_close(body_state_w, body_link_state_w) - torch.testing.assert_close(body_state_w, body_com_state_w) + torch.testing.assert_close(root_link_pose_w, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_com_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -1793,8 +1816,8 @@ def test_write_root_state(sim, num_articulations, device, with_offset, state_loc # check they are set torch.testing.assert_close(wp.to_torch(articulation.root_view.get_coms()), com) - rand_state = torch.zeros_like(wp.to_torch(articulation.data.root_state_w)) - rand_state[..., :7] = wp.to_torch(articulation.data.default_root_state)[..., :7] + rand_state = torch.zeros(num_articulations, 13, device=device) + rand_state[..., :7] = wp.to_torch(articulation.data.default_root_pose) rand_state[..., :3] += env_pos # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -1808,19 +1831,25 @@ def test_write_root_state(sim, num_articulations, device, with_offset, state_loc if state_location == "com": if i % 2 == 0: - articulation.write_root_com_state_to_sim(rand_state) + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) else: - articulation.write_root_com_state_to_sim(rand_state, env_ids=env_idx) + articulation.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + articulation.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) elif state_location == "link": if i % 2 == 0: - articulation.write_root_link_state_to_sim(rand_state) + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) else: - articulation.write_root_link_state_to_sim(rand_state, env_ids=env_idx) + articulation.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + articulation.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) if state_location == "com": - torch.testing.assert_close(rand_state, wp.to_torch(articulation.data.root_com_state_w)) + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(articulation.data.root_com_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(articulation.data.root_com_vel_w)) elif state_location == "link": - torch.testing.assert_close(rand_state, wp.to_torch(articulation.data.root_link_state_w)) + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(articulation.data.root_link_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(articulation.data.root_link_vel_w)) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -1854,15 +1883,18 @@ def test_body_incoming_joint_wrench_b_single_joint(sim, num_articulations, devic # apply action to the articulation joint_pos = torch.ones_like(wp.to_torch(articulation.data.joint_pos)) * 1.5708 / 2.0 - articulation.write_joint_state_to_sim( - torch.ones_like(wp.to_torch(articulation.data.joint_pos)), - torch.zeros_like(wp.to_torch(articulation.data.joint_vel)), + articulation.write_joint_position_to_sim_index( + position=torch.ones_like(wp.to_torch(articulation.data.joint_pos)), + ) + articulation.write_joint_velocity_to_sim_index( + velocity=torch.zeros_like(wp.to_torch(articulation.data.joint_vel)), ) - articulation.set_joint_position_target(joint_pos) + articulation.set_joint_position_target_index(target=joint_pos) articulation.write_data_to_sim() for _ in range(50): - # TODO: Replace with wrench composer once the deprecation is complete - articulation.set_external_force_and_torque(forces=external_force_vector_b, torques=external_torque_vector_b) + articulation.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_force_vector_b, torques=external_torque_vector_b + ) articulation.write_data_to_sim() # perform step sim.step() @@ -1974,7 +2006,7 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav limits = torch.zeros(num_articulations, articulation.num_joints, 2, device=device) limits[..., 0] = (torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0) * -1.0 limits[..., 1] = torch.rand(num_articulations, articulation.num_joints, device=device) + 5.0 - articulation.write_joint_position_limit_to_sim(limits) + articulation.write_joint_position_limit_to_sim_index(limits=limits) from torch.distributions import Uniform @@ -1983,30 +2015,35 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav pos_dist = Uniform(joint_pos_limits[..., 0], joint_pos_limits[..., 1]) vel_dist = Uniform(-joint_vel_limits, joint_vel_limits) - original_body_states = wp.to_torch(articulation.data.body_state_w).clone() + original_body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w).clone() + original_body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w).clone() rand_joint_pos = pos_dist.sample() rand_joint_vel = vel_dist.sample() - articulation.write_joint_state_to_sim(rand_joint_pos, rand_joint_vel) + articulation.write_joint_position_to_sim_index(position=rand_joint_pos) + articulation.write_joint_velocity_to_sim_index(velocity=rand_joint_vel) # make sure valued updated - body_state_w = wp.to_torch(articulation.data.body_state_w) + body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) + body_com_vel_w = wp.to_torch(articulation.data.body_com_vel_w) + original_body_states = torch.cat([original_body_link_pose_w, original_body_com_vel_w], dim=-1) + body_state_w = torch.cat([body_link_pose_w, body_com_vel_w], dim=-1) assert torch.count_nonzero(original_body_states[:, 1:] != body_state_w[:, 1:]) > ( len(original_body_states[:, 1:]) / 2 ) # validate body - link consistency - body_link_state_w = wp.to_torch(articulation.data.body_link_state_w) - torch.testing.assert_close(body_state_w[..., :7], body_link_state_w[..., :7]) - # skip 7:10 because they differs from link frame, this should be fine because we are only checking + body_link_vel_w = wp.to_torch(articulation.data.body_link_vel_w) + torch.testing.assert_close(body_link_pose_w, wp.to_torch(articulation.data.body_link_pose_w)) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity - torch.testing.assert_close(body_state_w[..., 10:], body_link_state_w[..., 10:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) # validate link - com conistency body_com_pos_b = wp.to_torch(articulation.data.body_com_pos_b) body_com_quat_b = wp.to_torch(articulation.data.body_com_quat_b) expected_com_pos, expected_com_quat = math_utils.combine_frame_transforms( - body_link_state_w[..., :3].view(-1, 3), - body_link_state_w[..., 3:7].view(-1, 4), + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:].view(-1, 4), body_com_pos_b.view(-1, 3), body_com_quat_b.view(-1, 4), ) @@ -2018,8 +2055,8 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav # validate body - com consistency body_com_lin_vel_w = wp.to_torch(articulation.data.body_com_lin_vel_w) body_com_ang_vel_w = wp.to_torch(articulation.data.body_com_ang_vel_w) - torch.testing.assert_close(body_state_w[..., 7:10], body_com_lin_vel_w) - torch.testing.assert_close(body_state_w[..., 10:], body_com_ang_vel_w) + torch.testing.assert_close(body_com_vel_w[..., :3], body_com_lin_vel_w) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_ang_vel_w) # validate pos_w, quat_w, pos_b, quat_b is consistent with pose_w and pose_b expected_com_pose_w = torch.cat((body_com_pos_w, body_com_quat_w), dim=2) @@ -2033,20 +2070,20 @@ def test_write_joint_state_data_consistency(sim, num_articulations, device, grav body_com_pose_w = wp.to_torch(articulation.data.body_com_pose_w) body_com_pose_b = wp.to_torch(articulation.data.body_com_pose_b) body_pose_w = wp.to_torch(articulation.data.body_pose_w) - body_link_pose_w = wp.to_torch(articulation.data.body_link_pose_w) + body_link_pose_w_fresh = wp.to_torch(articulation.data.body_link_pose_w) torch.testing.assert_close(body_com_pose_w, expected_com_pose_w) torch.testing.assert_close(body_com_pose_b, expected_com_pose_b) torch.testing.assert_close(body_pose_w, expected_body_pose_w) - torch.testing.assert_close(body_link_pose_w, expected_body_link_pose_w) + torch.testing.assert_close(body_link_pose_w_fresh, expected_body_link_pose_w) - # validate pose_w is consistent state[..., :7] + # validate pose_w is consistent with individual properties body_vel_w = wp.to_torch(articulation.data.body_vel_w) - body_com_state_w = wp.to_torch(articulation.data.body_com_state_w) - torch.testing.assert_close(body_pose_w, body_state_w[..., :7]) - torch.testing.assert_close(body_vel_w, body_state_w[..., 7:]) - torch.testing.assert_close(body_link_pose_w, body_link_state_w[..., :7]) - torch.testing.assert_close(body_com_pose_w, body_com_state_w[..., :7]) - torch.testing.assert_close(body_vel_w, body_state_w[..., 7:]) + body_com_vel_w_fresh = wp.to_torch(articulation.data.body_com_vel_w) + torch.testing.assert_close(body_pose_w, body_link_pose_w) + torch.testing.assert_close(body_vel_w, body_com_vel_w) + torch.testing.assert_close(body_link_pose_w_fresh, body_link_pose_w) + torch.testing.assert_close(body_com_pose_w, wp.to_torch(articulation.data.body_com_pose_w)) + torch.testing.assert_close(body_vel_w, body_com_vel_w_fresh) @pytest.mark.parametrize("num_articulations", [1, 2]) @@ -2064,7 +2101,7 @@ def test_spatial_tendons(sim, num_articulations, device): device: The device to run the simulation on """ # skip test if Isaac Sim version is less than 5.0 - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: pytest.skip("Spatial tendons are not supported in Isaac Sim < 5.0. Please update to Isaac Sim 5.0 or later.") return articulation_cfg = generate_articulation_cfg(articulation_type="spatial_tendon_test_asset") @@ -2087,10 +2124,10 @@ def test_spatial_tendons(sim, num_articulations, device): assert wp.to_torch(articulation.data.body_inertia).shape == (num_articulations, articulation.num_bodies, 9) assert articulation.num_spatial_tendons == 1 - articulation.set_spatial_tendon_stiffness(torch.tensor([10.0], device=device)) - articulation.set_spatial_tendon_limit_stiffness(torch.tensor([10.0], device=device)) - articulation.set_spatial_tendon_damping(torch.tensor([10.0], device=device)) - articulation.set_spatial_tendon_offset(torch.tensor([10.0], device=device)) + articulation.set_spatial_tendon_stiffness_index(stiffness=torch.tensor([10.0], device=device)) + articulation.set_spatial_tendon_limit_stiffness_index(limit_stiffness=torch.tensor([10.0], device=device)) + articulation.set_spatial_tendon_damping_index(damping=torch.tensor([10.0], device=device)) + articulation.set_spatial_tendon_offset_index(offset=torch.tensor([10.0], device=device)) # Simulate physics for _ in range(10): @@ -2129,9 +2166,11 @@ def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground # The static friction must be set first to be sure the dynamic friction is not greater than static # when both are set. - articulation.write_joint_friction_coefficient_to_sim(friction) - articulation.write_joint_dynamic_friction_coefficient_to_sim(dynamic_friction) - articulation.write_joint_viscous_friction_coefficient_to_sim(viscous_friction) + articulation.write_joint_friction_coefficient_to_sim_index( + joint_friction_coeff=friction, + joint_dynamic_friction_coeff=dynamic_friction, + joint_viscous_friction_coeff=viscous_friction, + ) articulation.write_data_to_sim() for _ in range(100): @@ -2150,7 +2189,7 @@ def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground # For Isaac Sim >= 5.0: also test the combined API that can set dynamic and viscous via # write_joint_friction_coefficient_to_sim; reset the sim to isolate this path. - if get_isaac_sim_version().major >= 5: + if has_kit() and get_isaac_sim_version().major >= 5: # Reset simulator to ensure a clean state for the alternative API path sim.reset() @@ -2168,7 +2207,7 @@ def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground dynamic_friction_2 = torch.min(dynamic_friction_2, friction_2) # Use the combined setter to write all three at once - articulation.write_joint_friction_coefficient_to_sim( + articulation.write_joint_friction_coefficient_to_sim_index( joint_friction_coeff=friction_2, joint_dynamic_friction_coeff=dynamic_friction_2, joint_viscous_friction_coeff=viscous_friction_2, diff --git a/source/isaaclab_physx/test/assets/test_deformable_object.py b/source/isaaclab_physx/test/assets/test_deformable_object.py index da17666282e..31992cad5de 100644 --- a/source/isaaclab_physx/test/assets/test_deformable_object.py +++ b/source/isaaclab_physx/test/assets/test_deformable_object.py @@ -246,7 +246,7 @@ def test_set_nodal_state(sim, num_cubes): ], dim=-1, ) - cube_object.write_nodal_state_to_sim(nodal_state) + cube_object.write_nodal_state_to_sim_index(nodal_state) torch.testing.assert_close( wp.to_torch(cube_object.data.nodal_state_w), nodal_state, rtol=1e-5, atol=1e-5 @@ -295,7 +295,7 @@ def test_set_nodal_state_with_applied_transform(num_cubes, randomize_pos, random else: torch.testing.assert_close(mean_nodal_pos_init, mean_nodal_pos_default + pos_w, rtol=1e-5, atol=1e-5) - cube_object.write_nodal_state_to_sim(nodal_state) + cube_object.write_nodal_state_to_sim_index(nodal_state) cube_object.reset() for _ in range(50): @@ -318,7 +318,7 @@ def test_set_kinematic_targets(sim, num_cubes): nodal_kinematic_targets = wp.to_torch(cube_object.root_view.get_sim_kinematic_targets()).clone() for _ in range(5): - cube_object.write_nodal_state_to_sim(wp.to_torch(cube_object.data.default_nodal_state_w)) + cube_object.write_nodal_state_to_sim_index(wp.to_torch(cube_object.data.default_nodal_state_w)) default_root_pos = wp.to_torch(cube_object.data.default_nodal_state_w).mean(dim=1) @@ -327,7 +327,7 @@ def test_set_kinematic_targets(sim, num_cubes): nodal_kinematic_targets[1:, :, 3] = 1.0 nodal_kinematic_targets[0, :, 3] = 0.0 nodal_kinematic_targets[0, :, :3] = wp.to_torch(cube_object.data.default_nodal_state_w)[0, :, :3] - cube_object.write_nodal_kinematic_target_to_sim( + cube_object.write_nodal_kinematic_target_to_sim_index( nodal_kinematic_targets[0], env_ids=torch.tensor([0], device=sim.device) ) diff --git a/source/isaaclab_physx/test/assets/test_rigid_object.py b/source/isaaclab_physx/test/assets/test_rigid_object.py index 7a75b8365ef..fa467523c39 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object.py @@ -162,9 +162,11 @@ def test_initialization_with_kinematic_enabled(num_cubes, device): # update object cube_object.update(sim.cfg.dt) # check that the object is kinematic - default_root_state = wp.to_torch(cube_object.data.default_root_state).clone() - default_root_state[:, :3] += origins - torch.testing.assert_close(wp.to_torch(cube_object.data.root_state_w), default_root_state) + default_root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + default_root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + default_root_pose[:, :3] += origins + torch.testing.assert_close(wp.to_torch(cube_object.data.root_link_pose_w), default_root_pose) + torch.testing.assert_close(wp.to_torch(cube_object.data.root_com_vel_w), default_root_vel) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -243,7 +245,7 @@ def test_external_force_buffer(device): external_wrench_b[:, :, 3] = force # apply force - cube_object.permanent_wrench_composer.set_forces_and_torques( + cube_object.permanent_wrench_composer.set_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids, @@ -255,7 +257,7 @@ def test_external_force_buffer(device): assert wp.to_torch(cube_object._permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force # Check if the instantaneous wrench is correctly added to the permanent wrench - cube_object.permanent_wrench_composer.add_forces_and_torques( + cube_object.permanent_wrench_composer.add_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=body_ids, @@ -302,12 +304,13 @@ def test_external_force_on_single_body(num_cubes, device): # Now we are ready! for i in range(5): # reset root state - root_state = wp.to_torch(cube_object.data.default_root_state).clone() + root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() # need to shift the position of the cubes otherwise they will be on top of each other - root_state[:, :3] = origins - cube_object.write_root_pose_to_sim(root_state[:, :7]) - cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + root_pose[:, :3] = origins + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) # reset object cube_object.reset() @@ -320,7 +323,7 @@ def test_external_force_on_single_body(num_cubes, device): positions = None # apply force - cube_object.permanent_wrench_composer.set_forces_and_torques( + cube_object.permanent_wrench_composer.set_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=positions, @@ -383,12 +386,13 @@ def test_external_force_on_single_body_at_position(num_cubes, device): # Now we are ready! for i in range(5): # reset root state - root_state = wp.to_torch(cube_object.data.default_root_state).clone() + root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() # need to shift the position of the cubes otherwise they will be on top of each other - root_state[:, :3] = origins - cube_object.write_root_pose_to_sim(root_state[:, :7]) - cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + root_pose[:, :3] = origins + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) # reset object cube_object.reset() @@ -407,14 +411,14 @@ def test_external_force_on_single_body_at_position(num_cubes, device): external_wrench_positions_b[..., 2] = 0.0 # apply force - cube_object.permanent_wrench_composer.set_forces_and_torques( + cube_object.permanent_wrench_composer.set_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=external_wrench_positions_b, body_ids=body_ids, is_global=is_global, ) - cube_object.permanent_wrench_composer.add_forces_and_torques( + cube_object.permanent_wrench_composer.add_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=external_wrench_positions_b, @@ -494,18 +498,17 @@ def test_set_rigid_object_state(num_cubes, device): # perform simulation for _ in range(5): - root_state = torch.cat( - [ - state_dict["root_pos_w"], - state_dict["root_quat_w"], - state_dict["root_lin_vel_w"], - state_dict["root_ang_vel_w"], - ], + root_pose = torch.cat( + [state_dict["root_pos_w"], state_dict["root_quat_w"]], + dim=-1, + ) + root_vel = torch.cat( + [state_dict["root_lin_vel_w"], state_dict["root_ang_vel_w"]], dim=-1, ) # reset root state - cube_object.write_root_pose_to_sim(root_state[:, :7]) - cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) sim.step() @@ -538,13 +541,14 @@ def test_reset_rigid_object(num_cubes, device): cube_object.update(sim.cfg.dt) # Move the object to a random position - root_state = wp.to_torch(cube_object.data.default_root_state).clone() - root_state[:, :3] = torch.randn(num_cubes, 3, device=sim.device) + root_pose = wp.to_torch(cube_object.data.default_root_pose).clone() + root_pose[:, :3] = torch.randn(num_cubes, 3, device=sim.device) # Random orientation - root_state[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) - cube_object.write_root_pose_to_sim(root_state[:, :7]) - cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + root_pose[:, 3:7] = random_orientation(num=num_cubes, device=sim.device) + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + root_vel = wp.to_torch(cube_object.data.default_root_vel).clone() + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) if i % 2 == 0: # reset object @@ -640,7 +644,7 @@ def test_rigid_body_no_friction(num_cubes, device): initial_velocity = torch.zeros((num_cubes, 6), device=sim.cfg.device) initial_velocity[:, 0] = 0.1 - cube_object.write_root_velocity_to_sim(initial_velocity) + cube_object.write_root_velocity_to_sim_index(root_velocity=initial_velocity) # Simulate physics for _ in range(5): @@ -707,7 +711,7 @@ def test_rigid_body_with_static_friction(num_cubes, device): for _ in range(100): sim.step() cube_object.update(sim.cfg.dt) - cube_object.write_root_velocity_to_sim(torch.zeros((num_cubes, 6), device=sim.device)) + cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=sim.device)) cube_mass = wp.to_torch(cube_object.root_view.get_masses()) gravity_magnitude = abs(sim.cfg.gravity[2]) # 2 cases: force applied is below and above mu @@ -715,7 +719,7 @@ def test_rigid_body_with_static_friction(num_cubes, device): # above mu: block should move as the force applied is > mu for force in "below_mu", "above_mu": # set initial velocity to zero - cube_object.write_root_velocity_to_sim(torch.zeros((num_cubes, 6), device=sim.device)) + cube_object.write_root_velocity_to_sim_index(root_velocity=torch.zeros((num_cubes, 6), device=sim.device)) external_wrench_b = torch.zeros((num_cubes, 1, 6), device=sim.device) if force == "below_mu": @@ -723,10 +727,9 @@ def test_rigid_body_with_static_friction(num_cubes, device): else: external_wrench_b[..., 0] = static_friction_coefficient * cube_mass * gravity_magnitude * 1.01 - # TODO: Replace with wrench composer once the deprecation is complete - cube_object.set_external_force_and_torque( - external_wrench_b[..., :3], - external_wrench_b[..., 3:], + cube_object.permanent_wrench_composer.set_forces_and_torques_index( + forces=external_wrench_b[..., :3], + torques=external_wrench_b[..., 3:], ) # Get root state @@ -744,7 +747,7 @@ def test_rigid_body_with_static_friction(num_cubes, device): wp.to_torch(cube_object.data.root_pos_w), initial_root_pos, rtol=2e-3, atol=2e-3 ) if force == "above_mu": - assert (wp.to_torch(cube_object.data.root_state_w)[..., 0] - initial_root_pos[..., 0] > 0.02).all() + assert (wp.to_torch(cube_object.data.root_pos_w)[..., 0] - initial_root_pos[..., 0] > 0.02).all() @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -782,15 +785,16 @@ def test_rigid_body_with_restitution(num_cubes, device): # Play sim sim.reset() - root_state = torch.zeros(num_cubes, 13, device=sim.device) - root_state[:, 3] = 1.0 # To make orientation a quaternion + root_pose = torch.zeros(num_cubes, 7, device=sim.device) + root_pose[:, 3] = 1.0 # To make orientation a quaternion for i in range(num_cubes): - root_state[i, 1] = 1.0 * i - root_state[:, 2] = 1.0 # Set an initial drop height - root_state[:, 9] = -1.0 # Set an initial downward velocity + root_pose[i, 1] = 1.0 * i + root_pose[:, 2] = 1.0 # Set an initial drop height + root_vel = torch.zeros(num_cubes, 6, device=sim.device) + root_vel[:, 2] = -1.0 # Set an initial downward velocity - cube_object.write_root_pose_to_sim(root_state[:, :7]) - cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + cube_object.write_root_pose_to_sim_index(root_pose=root_pose) + cube_object.write_root_velocity_to_sim_index(root_velocity=root_vel) static_friction = torch.zeros(num_cubes, 1) dynamic_friction = torch.zeros(num_cubes, 1) @@ -952,70 +956,76 @@ def test_body_root_state_properties(num_cubes, device, with_offset): # Simulate physics for _ in range(100): # spin the object around Z axis (com) - cube_object.write_root_velocity_to_sim(spin_twist.repeat(num_cubes, 1)) + cube_object.write_root_velocity_to_sim_index(root_velocity=spin_twist.repeat(num_cubes, 1)) # perform rendering sim.step() # update object cube_object.update(sim.cfg.dt) # get state properties - root_state_w = wp.to_torch(cube_object.data.root_state_w) - root_link_state_w = wp.to_torch(cube_object.data.root_link_state_w) - root_com_state_w = wp.to_torch(cube_object.data.root_com_state_w) - body_state_w = wp.to_torch(cube_object.data.body_state_w) - body_link_state_w = wp.to_torch(cube_object.data.body_link_state_w) - body_com_state_w = wp.to_torch(cube_object.data.body_com_state_w) + root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) + root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) + root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) + body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + body_link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + body_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + body_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) # if offset is [0,0,0] all root_state_%_w will match and all body_%_w will match if not with_offset: - torch.testing.assert_close(root_state_w, root_com_state_w) - torch.testing.assert_close(root_state_w, root_link_state_w) - torch.testing.assert_close(body_state_w, body_com_state_w) - torch.testing.assert_close(body_state_w, body_link_state_w) + torch.testing.assert_close(root_link_pose_w, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(root_link_pose_w, root_link_pose_w) + torch.testing.assert_close(root_com_vel_w, root_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_com_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) + torch.testing.assert_close(body_link_pose_w, body_link_pose_w) + torch.testing.assert_close(body_com_vel_w, body_link_vel_w) else: # cubes are spinning around center of mass # position will not match # center of mass position will be constant (i.e. spinning around com) - torch.testing.assert_close(env_pos + offset, root_com_state_w[..., :3]) - torch.testing.assert_close(env_pos + offset, body_com_state_w[..., :3].squeeze(-2)) + torch.testing.assert_close(env_pos + offset, root_com_pose_w[..., :3]) + torch.testing.assert_close(env_pos + offset, body_com_pose_w[..., :3].squeeze(-2)) # link position will be moving but should stay constant away from center of mass root_link_state_pos_rel_com = quat_apply_inverse( - root_link_state_w[..., 3:7], - root_link_state_w[..., :3] - root_com_state_w[..., :3], + root_link_pose_w[..., 3:], + root_link_pose_w[..., :3] - root_com_pose_w[..., :3], ) torch.testing.assert_close(-offset, root_link_state_pos_rel_com) body_link_state_pos_rel_com = quat_apply_inverse( - body_link_state_w[..., 3:7], - body_link_state_w[..., :3] - body_com_state_w[..., :3], + body_link_pose_w[..., 3:], + body_link_pose_w[..., :3] - body_com_pose_w[..., :3], ) torch.testing.assert_close(-offset, body_link_state_pos_rel_com.squeeze(-2)) # orientation of com will be a constant rotation from link orientation com_quat_b = wp.to_torch(cube_object.data.body_com_quat_b) - com_quat_w = quat_mul(body_link_state_w[..., 3:7], com_quat_b) - torch.testing.assert_close(com_quat_w, body_com_state_w[..., 3:7]) - torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_state_w[..., 3:7]) + com_quat_w = quat_mul(body_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, body_com_pose_w[..., 3:]) + torch.testing.assert_close(com_quat_w.squeeze(-2), root_com_pose_w[..., 3:]) # orientation of link will match root state will always match - torch.testing.assert_close(root_state_w[..., 3:7], root_link_state_w[..., 3:7]) - torch.testing.assert_close(body_state_w[..., 3:7], body_link_state_w[..., 3:7]) + torch.testing.assert_close(root_link_pose_w[..., 3:], root_link_pose_w[..., 3:]) + torch.testing.assert_close(body_link_pose_w[..., 3:], body_link_pose_w[..., 3:]) # lin_vel will not match # center of mass vel will be constant (i.e. spinning around com) - torch.testing.assert_close(torch.zeros_like(root_com_state_w[..., 7:10]), root_com_state_w[..., 7:10]) - torch.testing.assert_close(torch.zeros_like(body_com_state_w[..., 7:10]), body_com_state_w[..., 7:10]) + torch.testing.assert_close(torch.zeros_like(root_com_vel_w[..., :3]), root_com_vel_w[..., :3]) + torch.testing.assert_close(torch.zeros_like(body_com_vel_w[..., :3]), body_com_vel_w[..., :3]) # link frame will be moving, and should be equal to input angular velocity cross offset - lin_vel_rel_root_gt = quat_apply_inverse(root_link_state_w[..., 3:7], root_link_state_w[..., 7:10]) - lin_vel_rel_body_gt = quat_apply_inverse(body_link_state_w[..., 3:7], body_link_state_w[..., 7:10]) + lin_vel_rel_root_gt = quat_apply_inverse(root_link_pose_w[..., 3:], root_link_vel_w[..., :3]) + lin_vel_rel_body_gt = quat_apply_inverse(body_link_pose_w[..., 3:], body_link_vel_w[..., :3]) lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_cubes, 1)[..., 3:], -offset) torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_root_gt, atol=1e-4, rtol=1e-4) torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_body_gt.squeeze(-2), atol=1e-4, rtol=1e-4) # ang_vel will always match - torch.testing.assert_close(root_state_w[..., 10:], root_com_state_w[..., 10:]) - torch.testing.assert_close(root_state_w[..., 10:], root_link_state_w[..., 10:]) - torch.testing.assert_close(body_state_w[..., 10:], body_com_state_w[..., 10:]) - torch.testing.assert_close(body_state_w[..., 10:], body_link_state_w[..., 10:]) + torch.testing.assert_close(root_com_vel_w[..., 3:], root_com_vel_w[..., 3:]) + torch.testing.assert_close(root_com_vel_w[..., 3:], root_link_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_com_vel_w[..., 3:]) + torch.testing.assert_close(body_com_vel_w[..., 3:], body_link_vel_w[..., 3:]) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -1050,8 +1060,8 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): # check center of mass has been set torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_coms()), com) - rand_state = torch.zeros_like(wp.to_torch(cube_object.data.root_state_w)) - rand_state[..., :7] = wp.to_torch(cube_object.data.default_root_state)[..., :7] + rand_state = torch.zeros(num_cubes, 13, device=device) + rand_state[..., :7] = wp.to_torch(cube_object.data.default_root_pose) rand_state[..., :3] += env_pos # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -1065,19 +1075,27 @@ def test_write_root_state(num_cubes, device, with_offset, state_location): if state_location == "com": if i % 2 == 0: - cube_object.write_root_com_state_to_sim(rand_state) + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) else: - cube_object.write_root_com_state_to_sim(rand_state, env_ids=env_idx) + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:], env_ids=env_idx) elif state_location == "link": if i % 2 == 0: - cube_object.write_root_link_state_to_sim(rand_state) + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) else: - cube_object.write_root_link_state_to_sim(rand_state, env_ids=env_idx) + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7], env_ids=env_idx) + cube_object.write_root_link_velocity_to_sim_index( + root_velocity=rand_state[..., 7:], env_ids=env_idx + ) if state_location == "com": - torch.testing.assert_close(rand_state, wp.to_torch(cube_object.data.root_com_state_w)) + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.root_com_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.root_com_vel_w)) elif state_location == "link": - torch.testing.assert_close(rand_state, wp.to_torch(cube_object.data.root_link_state_w)) + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.root_link_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.root_link_vel_w)) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -1112,8 +1130,7 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, # check ceter of mass has been set torch.testing.assert_close(wp.to_torch(cube_object.root_view.get_coms()), com) - rand_state = torch.rand_like(wp.to_torch(cube_object.data.root_state_w)) - # rand_state[..., :7] = wp.to_torch(cube_object.data.default_root_state)[..., :7] + rand_state = torch.rand(num_cubes, 13, device=device) rand_state[..., :3] += env_pos # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -1126,67 +1143,73 @@ def test_write_state_functions_data_consistency(num_cubes, device, with_offset, cube_object.update(sim.cfg.dt) if state_location == "com": - cube_object.write_root_com_state_to_sim(rand_state) + cube_object.write_root_com_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_com_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) elif state_location == "link": - cube_object.write_root_link_state_to_sim(rand_state) + cube_object.write_root_link_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_link_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) elif state_location == "root": - cube_object.write_root_state_to_sim(rand_state) + cube_object.write_root_pose_to_sim_index(root_pose=rand_state[..., :7]) + cube_object.write_root_velocity_to_sim_index(root_velocity=rand_state[..., 7:]) if state_location == "com": - root_com_state_w = wp.to_torch(cube_object.data.root_com_state_w) + root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) + root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( - root_com_state_w[:, :3], - root_com_state_w[:, 3:7], + root_com_pose_w[:, :3], + root_com_pose_w[:, 3:], quat_rotate(quat_inv(body_com_pose_b[:, 0, 3:7]), -body_com_pose_b[:, 0, :3]), quat_inv(body_com_pose_b[:, 0, 3:7]), ) expected_root_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1) - root_link_state_w = wp.to_torch(cube_object.data.root_link_state_w) - # test both root_pose and root_link_state_w successfully updated when root_com_state_w updates - torch.testing.assert_close(expected_root_link_pose, root_link_state_w[:, :7]) - # skip 7:10 because they differs from link frame, this should be fine because we are only checking + root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) + root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + # test both root_pose and root_link successfully updated when root_com updates + torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity - torch.testing.assert_close(root_com_state_w[:, 10:], root_link_state_w[:, 10:]) - root_state_w = wp.to_torch(cube_object.data.root_state_w) - torch.testing.assert_close(expected_root_link_pose, root_state_w[:, :7]) - torch.testing.assert_close(root_com_state_w[:, 10:], root_state_w[:, 10:]) + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) + torch.testing.assert_close(expected_root_link_pose, root_link_pose_w) + torch.testing.assert_close(root_com_vel_w[:, 3:], wp.to_torch(cube_object.data.root_com_vel_w)[:, 3:]) elif state_location == "link": - root_link_state_w = wp.to_torch(cube_object.data.root_link_state_w) + root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) + root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) expected_com_pos, expected_com_quat = combine_frame_transforms( - root_link_state_w[:, :3], - root_link_state_w[:, 3:7], + root_link_pose_w[:, :3], + root_link_pose_w[:, 3:], body_com_pose_b[:, 0, :3], body_com_pose_b[:, 0, 3:7], ) expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) - root_com_state_w = wp.to_torch(cube_object.data.root_com_state_w) - # test both root_pose and root_com_state_w successfully updated when root_link_state_w updates - torch.testing.assert_close(expected_com_pose, root_com_state_w[:, :7]) - # skip 7:10 because they differs from link frame, this should be fine because we are only checking + root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) + root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) + # test both root_pose and root_com successfully updated when root_link updates + torch.testing.assert_close(expected_com_pose, root_com_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity - torch.testing.assert_close(root_link_state_w[:, 10:], root_com_state_w[:, 10:]) - root_state_w = wp.to_torch(cube_object.data.root_state_w) - torch.testing.assert_close(root_link_state_w[:, :7], root_state_w[:, :7]) - torch.testing.assert_close(root_link_state_w[:, 10:], root_state_w[:, 10:]) + torch.testing.assert_close(root_link_vel_w[:, 3:], root_com_vel_w[:, 3:]) + torch.testing.assert_close(root_link_pose_w, wp.to_torch(cube_object.data.root_link_pose_w)) + torch.testing.assert_close(root_link_vel_w[:, 3:], wp.to_torch(cube_object.data.root_com_vel_w)[:, 3:]) elif state_location == "root": - root_state_w = wp.to_torch(cube_object.data.root_state_w) + root_link_pose_w = wp.to_torch(cube_object.data.root_link_pose_w) + root_com_vel_w = wp.to_torch(cube_object.data.root_com_vel_w) body_com_pose_b = wp.to_torch(cube_object.data.body_com_pose_b) expected_com_pos, expected_com_quat = combine_frame_transforms( - root_state_w[:, :3], - root_state_w[:, 3:7], + root_link_pose_w[:, :3], + root_link_pose_w[:, 3:], body_com_pose_b[:, 0, :3], body_com_pose_b[:, 0, 3:7], ) expected_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1) - root_com_state_w = wp.to_torch(cube_object.data.root_com_state_w) - root_link_state_w = wp.to_torch(cube_object.data.root_link_state_w) - # test both root_com_state_w and root_link_state_w successfully updated when root_pose updates - torch.testing.assert_close(expected_com_pose, root_com_state_w[:, :7]) - torch.testing.assert_close(root_state_w[:, 7:], root_com_state_w[:, 7:]) - torch.testing.assert_close(root_state_w[:, :7], root_link_state_w[:, :7]) - torch.testing.assert_close(root_state_w[:, 10:], root_link_state_w[:, 10:]) + root_com_pose_w = wp.to_torch(cube_object.data.root_com_pose_w) + root_link_vel_w = wp.to_torch(cube_object.data.root_link_vel_w) + # test both root_com and root_link successfully updated when root_pose updates + torch.testing.assert_close(expected_com_pose, root_com_pose_w) + torch.testing.assert_close(root_com_vel_w, wp.to_torch(cube_object.data.root_com_vel_w)) + torch.testing.assert_close(root_link_pose_w, wp.to_torch(cube_object.data.root_link_pose_w)) + torch.testing.assert_close(root_com_vel_w[:, 3:], root_link_vel_w[:, 3:]) @pytest.mark.isaacsim_ci @@ -1235,3 +1258,4 @@ def test_warmup_attach_stage_not_called_for_cpu(): f"This indicates the CPU MBP broadphase double-initialization regression is present: " f"attach_stage() + force_load_physics_from_usd() must not be combined for CPU." ) + diff --git a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py index 7262add0ad8..977161c963e 100644 --- a/source/isaaclab_physx/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab_physx/test/assets/test_rigid_object_collection.py @@ -123,11 +123,11 @@ def test_initialization(sim, num_envs, num_cubes, device): # Check if object is initialized assert object_collection.is_initialized - assert len(object_collection.object_names) == num_cubes + assert len(object_collection.body_names) == num_cubes # Check buffers that exist and have correct shapes - assert wp.to_torch(object_collection.data.object_link_pos_w).shape == (num_envs, num_cubes, 3) - assert wp.to_torch(object_collection.data.object_link_quat_w).shape == (num_envs, num_cubes, 4) + assert wp.to_torch(object_collection.data.body_link_pos_w).shape == (num_envs, num_cubes, 3) + assert wp.to_torch(object_collection.data.body_link_quat_w).shape == (num_envs, num_cubes, 4) assert wp.to_torch(object_collection.data.body_mass).shape == (num_envs, num_cubes, 1) assert wp.to_torch(object_collection.data.body_inertia).shape == (num_envs, num_cubes, 9) @@ -190,20 +190,22 @@ def test_initialization_with_kinematic_enabled(sim, num_envs, num_cubes, device) # Check if object is initialized assert object_collection.is_initialized - assert len(object_collection.object_names) == num_cubes + assert len(object_collection.body_names) == num_cubes # Check buffers that exist and have correct shapes - assert wp.to_torch(object_collection.data.object_link_pos_w).shape == (num_envs, num_cubes, 3) - assert wp.to_torch(object_collection.data.object_link_quat_w).shape == (num_envs, num_cubes, 4) + assert wp.to_torch(object_collection.data.body_link_pos_w).shape == (num_envs, num_cubes, 3) + assert wp.to_torch(object_collection.data.body_link_quat_w).shape == (num_envs, num_cubes, 4) # Simulate physics for _ in range(2): sim.step() object_collection.update(sim.cfg.dt) # check that the object is kinematic - default_object_state = wp.to_torch(object_collection.data.default_object_state).clone() - default_object_state[..., :3] += origins.unsqueeze(1) - torch.testing.assert_close(wp.to_torch(object_collection.data.object_link_state_w), default_object_state) + default_body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + default_body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + default_body_pose[..., :3] += origins.unsqueeze(1) + torch.testing.assert_close(wp.to_torch(object_collection.data.body_link_pose_w), default_body_pose) + torch.testing.assert_close(wp.to_torch(object_collection.data.body_link_vel_w), default_body_vel) @pytest.mark.parametrize("num_cubes", [1, 2]) @@ -229,7 +231,7 @@ def test_external_force_buffer(sim, device): sim.reset() # find objects to apply the force - object_ids, object_names = object_collection.find_objects(".*") + object_ids, object_names = object_collection.find_bodies(".*") # reset object object_collection.reset() @@ -248,7 +250,7 @@ def test_external_force_buffer(sim, device): external_wrench_b[:, :, 0] = force external_wrench_b[:, :, 3] = force - object_collection.permanent_wrench_composer.set_forces_and_torques( + object_collection.permanent_wrench_composer.set_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], body_ids=object_ids, @@ -260,7 +262,7 @@ def test_external_force_buffer(sim, device): assert wp.to_torch(object_collection._permanent_wrench_composer.composed_force)[i, 0, 0].item() == force assert wp.to_torch(object_collection._permanent_wrench_composer.composed_torque)[i, 0, 0].item() == force - object_collection.instantaneous_wrench_composer.add_forces_and_torques( + object_collection.instantaneous_wrench_composer.add_forces_and_torques_index( body_ids=object_ids, forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], @@ -281,7 +283,7 @@ def test_external_force_on_single_body(sim, num_envs, num_cubes, device): sim.reset() # find objects to apply the force - object_ids, object_names = object_collection.find_objects(".*") + object_ids, object_names = object_collection.find_bodies(".*") # Sample a force equal to the weight of the object external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) @@ -290,22 +292,24 @@ def test_external_force_on_single_body(sim, num_envs, num_cubes, device): for i in range(5): # reset object state - object_state = wp.to_torch(object_collection.data.default_object_state).clone() + body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() # need to shift the position of the cubes otherwise they will be on top of each other - object_state[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_object_state_to_sim(object_state) + body_pose[..., :2] += origins.unsqueeze(1)[..., :2] + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) # reset object object_collection.reset() is_global = False if i % 2 == 0: - positions = wp.to_torch(object_collection.data.object_link_pos_w)[:, object_ids, :3] + positions = wp.to_torch(object_collection.data.body_link_pos_w)[:, object_ids, :3] is_global = True else: positions = None # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques( + object_collection.permanent_wrench_composer.set_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=positions, @@ -323,11 +327,11 @@ def test_external_force_on_single_body(sim, num_envs, num_cubes, device): # First object should still be at the same Z position (1.0) torch.testing.assert_close( - wp.to_torch(object_collection.data.object_link_pos_w)[:, 0::2, 2], - torch.ones_like(wp.to_torch(object_collection.data.object_pos_w)[:, 0::2, 2]), + wp.to_torch(object_collection.data.body_link_pos_w)[:, 0::2, 2], + torch.ones_like(wp.to_torch(object_collection.data.body_link_pos_w)[:, 0::2, 2]), ) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(object_collection.data.object_link_pos_w)[:, 1::2, 2] < 1.0) + assert torch.all(wp.to_torch(object_collection.data.body_link_pos_w)[:, 1::2, 2] < 1.0) @pytest.mark.parametrize("num_envs", [1, 2]) @@ -344,7 +348,7 @@ def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, dev sim.reset() # find objects to apply the force - object_ids, object_names = object_collection.find_objects(".*") + object_ids, object_names = object_collection.find_bodies(".*") # Sample a force equal to the weight of the object external_wrench_b = torch.zeros(object_collection.num_instances, len(object_ids), 6, device=sim.device) @@ -356,16 +360,18 @@ def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, dev # Desired force and torque for i in range(5): # reset object state - object_state = wp.to_torch(object_collection.data.default_object_state).clone() + body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() # need to shift the position of the cubes otherwise they will be on top of each other - object_state[..., :2] += origins.unsqueeze(1)[..., :2] - object_collection.write_object_state_to_sim(object_state) + body_pose[..., :2] += origins.unsqueeze(1)[..., :2] + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) # reset object object_collection.reset() is_global = False if i % 2 == 0: - body_com_pos_w = wp.to_torch(object_collection.data.object_link_pos_w)[:, object_ids, :3] + body_com_pos_w = wp.to_torch(object_collection.data.body_link_pos_w)[:, object_ids, :3] external_wrench_positions_b[..., 0] = 0.0 external_wrench_positions_b[..., 1] = 1.0 external_wrench_positions_b[..., 2] = 0.0 @@ -377,7 +383,7 @@ def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, dev external_wrench_positions_b[..., 2] = 0.0 # apply force - object_collection.permanent_wrench_composer.set_forces_and_torques( + object_collection.permanent_wrench_composer.set_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=external_wrench_positions_b, @@ -385,7 +391,7 @@ def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, dev env_ids=None, is_global=is_global, ) - object_collection.permanent_wrench_composer.add_forces_and_torques( + object_collection.permanent_wrench_composer.add_forces_and_torques_index( forces=external_wrench_b[..., :3], torques=external_wrench_b[..., 3:], positions=external_wrench_positions_b, @@ -402,9 +408,9 @@ def test_external_force_on_single_body_at_position(sim, num_envs, num_cubes, dev object_collection.update(sim.cfg.dt) # First object should be rotating around it's X axis - assert torch.all(wp.to_torch(object_collection.data.object_ang_vel_b)[:, 0::2, 0] > 0.1) + assert torch.all(wp.to_torch(object_collection.data.body_com_ang_vel_b)[:, 0::2, 0] > 0.1) # Second object should have fallen, so it's Z height should be less than initial height of 1.0 - assert torch.all(wp.to_torch(object_collection.data.object_link_pos_w)[:, 1::2, 2] < 1.0) + assert torch.all(wp.to_torch(object_collection.data.body_link_pos_w)[:, 1::2, 2] < 1.0) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -421,20 +427,20 @@ def test_set_object_state(sim, num_envs, num_cubes, device, gravity_enabled): object_collection, origins = generate_cubes_scene(num_envs=num_envs, num_cubes=num_cubes, device=device) sim.reset() - state_types = ["object_pos_w", "object_quat_w", "object_lin_vel_w", "object_ang_vel_w"] + state_types = ["body_link_pos_w", "body_link_quat_w", "body_com_lin_vel_w", "body_com_ang_vel_w"] # Set each state type individually as they are dependent on each other for state_type_to_randomize in state_types: state_dict = { - "object_pos_w": torch.zeros_like(wp.to_torch(object_collection.data.object_pos_w), device=sim.device), - "object_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( + "body_link_pos_w": torch.zeros_like(wp.to_torch(object_collection.data.body_link_pos_w), device=sim.device), + "body_link_quat_w": default_orientation(num=num_cubes * num_envs, device=sim.device).view( num_envs, num_cubes, 4 ), - "object_lin_vel_w": torch.zeros_like( - wp.to_torch(object_collection.data.object_lin_vel_w), device=sim.device + "body_com_lin_vel_w": torch.zeros_like( + wp.to_torch(object_collection.data.body_com_lin_vel_w), device=sim.device ), - "object_ang_vel_w": torch.zeros_like( - wp.to_torch(object_collection.data.object_ang_vel_w), device=sim.device + "body_com_ang_vel_w": torch.zeros_like( + wp.to_torch(object_collection.data.body_com_ang_vel_w), device=sim.device ), } @@ -443,29 +449,29 @@ def test_set_object_state(sim, num_envs, num_cubes, device, gravity_enabled): object_collection.reset() # Set random state - if state_type_to_randomize == "object_quat_w": + if state_type_to_randomize == "body_link_quat_w": state_dict[state_type_to_randomize] = random_orientation( num=num_cubes * num_envs, device=sim.device ).view(num_envs, num_cubes, 4) else: state_dict[state_type_to_randomize] = torch.randn(num_envs, num_cubes, 3, device=sim.device) # make sure objects do not overlap - if state_type_to_randomize == "object_pos_w": + if state_type_to_randomize == "body_link_pos_w": state_dict[state_type_to_randomize][..., :2] += origins.unsqueeze(1)[..., :2] # perform simulation for _ in range(5): - object_state = torch.cat( - [ - state_dict["object_pos_w"], - state_dict["object_quat_w"], - state_dict["object_lin_vel_w"], - state_dict["object_ang_vel_w"], - ], + body_pose = torch.cat( + [state_dict["body_link_pos_w"], state_dict["body_link_quat_w"]], + dim=-1, + ) + body_vel = torch.cat( + [state_dict["body_com_lin_vel_w"], state_dict["body_com_ang_vel_w"]], dim=-1, ) # reset object state - object_collection.write_object_state_to_sim(object_state=object_state) + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) sim.step() # assert that set object quantities are equal to the ones set in the state_dict @@ -515,60 +521,61 @@ def test_object_state_properties(sim, num_envs, num_cubes, device, with_offset, spin_twist[5] = torch.randn(1, device=device) # initial spawn point - init_com = wp.to_torch(cube_object.data.object_com_state_w)[..., :3] + init_com = wp.to_torch(cube_object.data.body_com_pose_w)[..., :3] for i in range(10): # spin the object around Z axis (com) - cube_object.write_object_velocity_to_sim(spin_twist.repeat(num_envs, num_cubes, 1)) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=spin_twist.repeat(num_envs, num_cubes, 1)) sim.step() cube_object.update(sim.cfg.dt) # get state properties - object_state_w = wp.to_torch(cube_object.data.object_state_w) - object_link_state_w = wp.to_torch(cube_object.data.object_link_state_w) - object_com_state_w = wp.to_torch(cube_object.data.object_com_state_w) + object_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + object_link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + object_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + object_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) # if offset is [0,0,0] all object_state_%_w will match and all body_%_w will match if not with_offset: - torch.testing.assert_close(object_state_w, object_com_state_w) - torch.testing.assert_close(object_state_w, object_link_state_w) + torch.testing.assert_close(object_link_pose_w, object_com_pose_w) + torch.testing.assert_close(object_com_vel_w, object_link_vel_w) else: # cubes are spinning around center of mass # position will not match # center of mass position will be constant (i.e. spinning around com) - torch.testing.assert_close(init_com, object_com_state_w[..., :3]) + torch.testing.assert_close(init_com, object_com_pose_w[..., :3]) # link position will be moving but should stay constant away from center of mass object_link_state_pos_rel_com = quat_apply_inverse( - object_link_state_w[..., 3:7], - object_link_state_w[..., :3] - object_com_state_w[..., :3], + object_link_pose_w[..., 3:], + object_link_pose_w[..., :3] - object_com_pose_w[..., :3], ) torch.testing.assert_close(-offset, object_link_state_pos_rel_com) # orientation of com will be a constant rotation from link orientation - com_quat_b = wp.to_torch(cube_object.data.object_com_quat_b) - com_quat_w = quat_mul(object_link_state_w[..., 3:7], com_quat_b) - torch.testing.assert_close(com_quat_w, object_com_state_w[..., 3:7]) + com_quat_b = wp.to_torch(cube_object.data.body_com_quat_b) + com_quat_w = quat_mul(object_link_pose_w[..., 3:], com_quat_b) + torch.testing.assert_close(com_quat_w, object_com_pose_w[..., 3:]) # orientation of link will match object state will always match - torch.testing.assert_close(object_state_w[..., 3:7], object_link_state_w[..., 3:7]) + torch.testing.assert_close(object_link_pose_w[..., 3:], object_link_pose_w[..., 3:]) # lin_vel will not match # center of mass vel will be constant (i.e. spinning around com) torch.testing.assert_close( - torch.zeros_like(object_com_state_w[..., 7:10]), - object_com_state_w[..., 7:10], + torch.zeros_like(object_com_vel_w[..., :3]), + object_com_vel_w[..., :3], ) # link frame will be moving, and should be equal to input angular velocity cross offset - lin_vel_rel_object_gt = quat_apply_inverse(object_link_state_w[..., 3:7], object_link_state_w[..., 7:10]) + lin_vel_rel_object_gt = quat_apply_inverse(object_link_pose_w[..., 3:], object_link_vel_w[..., :3]) lin_vel_rel_gt = torch.linalg.cross(spin_twist.repeat(num_envs, num_cubes, 1)[..., 3:], -offset) torch.testing.assert_close(lin_vel_rel_gt, lin_vel_rel_object_gt, atol=1e-4, rtol=1e-3) # ang_vel will always match - torch.testing.assert_close(object_state_w[..., 10:], object_com_state_w[..., 10:]) - torch.testing.assert_close(object_state_w[..., 10:], object_link_state_w[..., 10:]) + torch.testing.assert_close(object_com_vel_w[..., 3:], object_com_vel_w[..., 3:]) + torch.testing.assert_close(object_com_vel_w[..., 3:], object_link_vel_w[..., 3:]) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -608,9 +615,9 @@ def test_write_object_state(sim, num_envs, num_cubes, device, with_offset, state wp.to_torch(cube_object.reshape_view_to_data_2d(cube_object.root_view.get_coms().view(wp.transformf))), com ) - rand_state = torch.zeros_like(wp.to_torch(cube_object.data.object_link_state_w)) - rand_state[..., :7] = wp.to_torch(cube_object.data.default_object_state)[..., :7] - rand_state[..., :3] += wp.to_torch(cube_object.data.object_link_pos_w) + rand_state = torch.zeros(num_envs, num_cubes, 13, device=device) + rand_state[..., :7] = wp.to_torch(cube_object.data.default_body_pose) + rand_state[..., :3] += wp.to_torch(cube_object.data.body_link_pos_w) # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -622,19 +629,33 @@ def test_write_object_state(sim, num_envs, num_cubes, device, with_offset, state if state_location == "com": if i % 2 == 0: - cube_object.write_object_com_state_to_sim(rand_state) + cube_object.write_body_com_pose_to_sim_index(body_poses=rand_state[..., :7]) + cube_object.write_body_com_velocity_to_sim_index(body_velocities=rand_state[..., 7:]) else: - cube_object.write_object_com_state_to_sim(rand_state, env_ids=env_ids, object_ids=object_ids) + cube_object.write_body_com_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) elif state_location == "link": if i % 2 == 0: - cube_object.write_object_link_state_to_sim(rand_state) + cube_object.write_body_link_pose_to_sim_index(body_poses=rand_state[..., :7]) + cube_object.write_body_link_velocity_to_sim_index(body_velocities=rand_state[..., 7:]) else: - cube_object.write_object_link_state_to_sim(rand_state, env_ids=env_ids, object_ids=object_ids) + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_link_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) if state_location == "com": - torch.testing.assert_close(rand_state, wp.to_torch(cube_object.data.object_com_state_w)) + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.body_com_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.body_com_vel_w)) elif state_location == "link": - torch.testing.assert_close(rand_state, wp.to_torch(cube_object.data.object_link_state_w)) + torch.testing.assert_close(rand_state[..., :7], wp.to_torch(cube_object.data.body_link_pose_w)) + torch.testing.assert_close(rand_state[..., 7:], wp.to_torch(cube_object.data.body_link_vel_w)) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -650,11 +671,13 @@ def test_reset_object_collection(sim, num_envs, num_cubes, device): object_collection.update(sim.cfg.dt) # Move the object to a random position - object_state = wp.to_torch(object_collection.data.default_object_state).clone() - object_state[..., :3] = torch.randn(num_envs, num_cubes, 3, device=sim.device) + body_pose = wp.to_torch(object_collection.data.default_body_pose).clone() + body_pose[..., :3] = torch.randn(num_envs, num_cubes, 3, device=sim.device) # Random orientation - object_state[..., 3:7] = random_orientation(num=num_cubes, device=sim.device) - object_collection.write_object_state_to_sim(object_state) + body_pose[..., 3:7] = random_orientation(num=num_cubes, device=sim.device) + object_collection.write_body_link_pose_to_sim_index(body_poses=body_pose) + body_vel = wp.to_torch(object_collection.data.default_body_vel).clone() + object_collection.write_body_com_velocity_to_sim_index(body_velocities=body_vel) if i % 2 == 0: object_collection.reset() @@ -737,7 +760,7 @@ def test_gravity_vec_w(sim, num_envs, num_cubes, device, gravity_enabled): gravity[..., 2] = -9.81 # Check the body accelerations are correct - torch.testing.assert_close(wp.to_torch(object_collection.data.object_acc_w), gravity) + torch.testing.assert_close(wp.to_torch(object_collection.data.body_com_acc_w), gravity) @pytest.mark.parametrize("num_envs", [1, 3]) @@ -780,8 +803,8 @@ def test_write_object_state_functions_data_consistency( wp.to_torch(cube_object.reshape_view_to_data_2d(cube_object.root_view.get_coms().view(wp.transformf))), com ) - rand_state = torch.rand_like(wp.to_torch(cube_object.data.object_link_state_w)) - rand_state[..., :3] += wp.to_torch(cube_object.data.object_link_pos_w) + rand_state = torch.rand(num_envs, num_cubes, 13, device=device) + rand_state[..., :3] += wp.to_torch(cube_object.data.body_link_pos_w) # make quaternion a unit vector rand_state[..., 3:7] = torch.nn.functional.normalize(rand_state[..., 3:7], dim=-1) @@ -790,91 +813,95 @@ def test_write_object_state_functions_data_consistency( sim.step() cube_object.update(sim.cfg.dt) + body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + body_com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) object_link_to_com_pos, object_link_to_com_quat = subtract_frame_transforms( - wp.to_torch(cube_object.data.object_link_state_w)[..., :3].view(-1, 3), - wp.to_torch(cube_object.data.object_link_state_w)[..., 3:7].view(-1, 4), - wp.to_torch(cube_object.data.object_com_state_w)[..., :3].view(-1, 3), - wp.to_torch(cube_object.data.object_com_state_w)[..., 3:7].view(-1, 4), + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:7].view(-1, 4), + body_com_pose_w[..., :3].view(-1, 3), + body_com_pose_w[..., 3:7].view(-1, 4), ) if state_location == "com": - cube_object.write_object_com_state_to_sim(rand_state, env_ids=env_ids, object_ids=object_ids) + cube_object.write_body_com_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) elif state_location == "link": - cube_object.write_object_link_state_to_sim(rand_state, env_ids=env_ids, object_ids=object_ids) + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_link_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) elif state_location == "root": - cube_object.write_object_state_to_sim(rand_state, env_ids=env_ids, object_ids=object_ids) + cube_object.write_body_link_pose_to_sim_index( + body_poses=rand_state[..., :7], env_ids=env_ids, body_ids=object_ids + ) + cube_object.write_body_com_velocity_to_sim_index( + body_velocities=rand_state[..., 7:], env_ids=env_ids, body_ids=object_ids + ) if state_location == "com": + com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) expected_root_link_pos, expected_root_link_quat = combine_frame_transforms( - wp.to_torch(cube_object.data.object_com_state_w)[..., :3].view(-1, 3), - wp.to_torch(cube_object.data.object_com_state_w)[..., 3:7].view(-1, 4), + com_pose_w[..., :3].view(-1, 3), + com_pose_w[..., 3:].view(-1, 4), quat_rotate(quat_inv(object_link_to_com_quat), -object_link_to_com_pos), quat_inv(object_link_to_com_quat), ) - # torch.testing.assert_close(rand_state, wp.to_torch(cube_object.data.object_com_state_w)) expected_object_link_pose = torch.cat((expected_root_link_pos, expected_root_link_quat), dim=1).view( num_envs, -1, 7 ) - # test both root_pose and root_link_state_w successfully updated when root_com_state_w updates - torch.testing.assert_close( - expected_object_link_pose, wp.to_torch(cube_object.data.object_link_state_w)[..., :7] - ) - # skip 7:10 because they differs from link frame, this should be fine because we are only checking + link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + # test both root_pose and root_link successfully updated when root_com updates + torch.testing.assert_close(expected_object_link_pose, link_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity - torch.testing.assert_close( - wp.to_torch(cube_object.data.object_com_state_w)[..., 10:], - wp.to_torch(cube_object.data.object_link_state_w)[..., 10:], - ) - torch.testing.assert_close(expected_object_link_pose, wp.to_torch(cube_object.data.object_state_w)[..., :7]) - torch.testing.assert_close( - wp.to_torch(cube_object.data.object_com_state_w)[..., 10:], - wp.to_torch(cube_object.data.object_state_w)[..., 10:], - ) + torch.testing.assert_close(com_vel_w[..., 3:], link_vel_w[..., 3:]) + torch.testing.assert_close(expected_object_link_pose, link_pose_w) + torch.testing.assert_close(com_vel_w[..., 3:], wp.to_torch(cube_object.data.body_com_vel_w)[..., 3:]) elif state_location == "link": + link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) expected_com_pos, expected_com_quat = combine_frame_transforms( - wp.to_torch(cube_object.data.object_link_state_w)[..., :3].view(-1, 3), - wp.to_torch(cube_object.data.object_link_state_w)[..., 3:7].view(-1, 4), + link_pose_w[..., :3].view(-1, 3), + link_pose_w[..., 3:].view(-1, 4), object_link_to_com_pos, object_link_to_com_quat, ) expected_object_com_pose = torch.cat((expected_com_pos, expected_com_quat), dim=1).view(num_envs, -1, 7) - # test both root_pose and root_com_state_w successfully updated when root_link_state_w updates - torch.testing.assert_close(expected_object_com_pose, wp.to_torch(cube_object.data.object_com_state_w)[..., :7]) - # skip 7:10 because they differs from link frame, this should be fine because we are only checking + com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + # test both root_pose and root_com successfully updated when root_link updates + torch.testing.assert_close(expected_object_com_pose, com_pose_w) + # skip lin_vel because it differs from link frame, this should be fine because we are only checking # if velocity update is triggered, which can be determined by comparing angular velocity - torch.testing.assert_close( - wp.to_torch(cube_object.data.object_link_state_w)[..., 10:], - wp.to_torch(cube_object.data.object_com_state_w)[..., 10:], - ) - torch.testing.assert_close( - wp.to_torch(cube_object.data.object_link_state_w)[..., :7], - wp.to_torch(cube_object.data.object_state_w)[..., :7], - ) - torch.testing.assert_close( - wp.to_torch(cube_object.data.object_link_state_w)[..., 10:], - wp.to_torch(cube_object.data.object_state_w)[..., 10:], - ) + torch.testing.assert_close(link_vel_w[..., 3:], com_vel_w[..., 3:]) + torch.testing.assert_close(link_pose_w, wp.to_torch(cube_object.data.body_link_pose_w)) + torch.testing.assert_close(link_vel_w[..., 3:], wp.to_torch(cube_object.data.body_com_vel_w)[..., 3:]) elif state_location == "root": + body_link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + body_com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) expected_object_com_pos, expected_object_com_quat = combine_frame_transforms( - wp.to_torch(cube_object.data.object_state_w)[..., :3].view(-1, 3), - wp.to_torch(cube_object.data.object_state_w)[..., 3:7].view(-1, 4), + body_link_pose_w[..., :3].view(-1, 3), + body_link_pose_w[..., 3:].view(-1, 4), object_link_to_com_pos, object_link_to_com_quat, ) expected_object_com_pose = torch.cat((expected_object_com_pos, expected_object_com_quat), dim=1).view( num_envs, -1, 7 ) - # test both root_com_state_w and root_link_state_w successfully updated when root_pose updates - torch.testing.assert_close(expected_object_com_pose, wp.to_torch(cube_object.data.object_com_state_w)[..., :7]) - torch.testing.assert_close( - wp.to_torch(cube_object.data.object_state_w)[..., 7:], - wp.to_torch(cube_object.data.object_com_state_w)[..., 7:], - ) - torch.testing.assert_close( - wp.to_torch(cube_object.data.object_state_w)[..., :7], - wp.to_torch(cube_object.data.object_link_state_w)[..., :7], - ) - torch.testing.assert_close( - wp.to_torch(cube_object.data.object_state_w)[..., 10:], - wp.to_torch(cube_object.data.object_link_state_w)[..., 10:], - ) + com_pose_w = wp.to_torch(cube_object.data.body_com_pose_w) + com_vel_w = wp.to_torch(cube_object.data.body_com_vel_w) + link_pose_w = wp.to_torch(cube_object.data.body_link_pose_w) + link_vel_w = wp.to_torch(cube_object.data.body_link_vel_w) + # test both root_com and root_link successfully updated when root_pose updates + torch.testing.assert_close(expected_object_com_pose, com_pose_w) + torch.testing.assert_close(body_com_vel_w, com_vel_w) + torch.testing.assert_close(body_link_pose_w, link_pose_w) + torch.testing.assert_close(body_com_vel_w[..., 3:], link_vel_w[..., 3:]) diff --git a/source/isaaclab_physx/test/assets/test_surface_gripper.py b/source/isaaclab_physx/test/assets/test_surface_gripper.py index c569d601551..e85a4a8415c 100644 --- a/source/isaaclab_physx/test/assets/test_surface_gripper.py +++ b/source/isaaclab_physx/test/assets/test_surface_gripper.py @@ -31,7 +31,7 @@ ) from isaaclab.sim import build_simulation_context from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit # from isaacsim.robot.surface_gripper import GripperView @@ -171,7 +171,7 @@ def test_initialization(sim, num_articulations, device, add_ground_plane) -> Non device: The device to run the test on. add_ground_plane: Whether to add a ground plane to the simulation. """ - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: return surface_gripper_cfg, articulation_cfg = generate_surface_gripper_cfgs(kinematic_enabled=False) surface_gripper, articulation, _ = generate_surface_gripper( @@ -205,7 +205,7 @@ def test_initialization(sim, num_articulations, device, add_ground_plane) -> Non @pytest.mark.isaacsim_ci def test_raise_error_if_not_cpu(sim, device, add_ground_plane) -> None: """Test that the SurfaceGripper raises an error if the device is not CPU.""" - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: return num_articulations = 1 surface_gripper_cfg, articulation_cfg = generate_surface_gripper_cfgs(kinematic_enabled=False) From 38e1776275ddfe2e1304dd36ea8e11c7d700042c Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Mon, 23 Feb 2026 15:04:57 -0800 Subject: [PATCH 3/4] add hydra preset and group feature --- docs/source/features/hydra.rst | 205 +++++++ source/isaaclab_tasks/config/extension.toml | 2 +- source/isaaclab_tasks/docs/CHANGELOG.rst | 20 + .../franka/agents/rl_games_ppo_cfg.yaml | 12 + .../isaaclab_tasks/utils/hydra.py | 371 ++++++++++--- source/isaaclab_tasks/test/test_hydra.py | 501 +++++++++++++++--- 6 files changed, 973 insertions(+), 138 deletions(-) diff --git a/docs/source/features/hydra.rst b/docs/source/features/hydra.rst index 47e84fb328c..9e9a9ed2bda 100644 --- a/docs/source/features/hydra.rst +++ b/docs/source/features/hydra.rst @@ -127,3 +127,208 @@ the post init update is as follows: Here, when modifying ``env.decimation`` or ``env.sim.dt``, the user needs to give the updated ``env.sim.render_interval``, ``env.scene.height_scanner.update_period``, and ``env.scene.contact_forces.update_period`` as input as well. + + +Preset System +------------- + +The preset system lets you swap out entire config sections with a single command line argument. +Instead of overriding individual fields, you select a named preset that **completely replaces** the +config section (no field merging). + +Presets are defined directly on config classes using a ``presets`` attribute. The system recursively +discovers all presets from nested configs automatically. + + +Override Order +^^^^^^^^^^^^^^ + +The override order is strict and deterministic: + +**By type** (earlier stages are overwritten by later ones): + +1. **Presets** — Global preset selections (``presets=noise_less``) +2. **Groups** — Path preset selections (``env.actions.arm_action=relative_joint_position``) +3. **Single field** — Scalar overrides (``env.sim.dt=0.001``) + +**Within the same type:** + +- **Lower depth > higher depth** — Parent paths are applied before child paths (e.g. ``env.actions`` before ``env.actions.arm_action``) +- **Left > right** — Arguments are processed left-to-right; for the same target, the rightmost override wins + +**Concrete sequence:** + +1. Auto-default presets for paths not explicitly selected +2. Global presets (``presets=name``) +3. Path presets (``env.actions=name``), sorted by depth ascending +4. Scalar overrides (both preset-path and global) + + +Defining Presets +^^^^^^^^^^^^^^^^ + +There are three styles for defining presets: + +**Style 1: Inheritance** - Default values from base class, presets for alternatives: + +.. code-block:: python + + @configclass + class FrankaArmActionCfg(mdp.JointPositionActionCfg): + """Franka arm action config with presets for different action types.""" + + presets = { + "joint_position_to_limit": mdp.JointPositionToLimitsActionCfg( + asset_name="robot", joint_names=["panda_joint.*"] + ), + "relative_joint_position": mdp.RelativeJointPositionActionCfg( + asset_name="robot", joint_names=["panda_joint.*"], scale=0.2 + ), + } + +**Style 2: Inner class** - Self-contained with nested preset definitions: + +.. code-block:: python + + @configclass + class PhysxCfg: + """Simulation config with physics backend presets.""" + + backend: str = "physx" + + @configclass + class Newton: + backend: str = "newton" + + presets = {"newton": Newton()} + +**Style 3: Preset-only with auto-default** - Pure composition, no default fields: + +.. code-block:: python + + @configclass + class ObservationsCfg: + """Observation specifications with presets.""" + + presets = { + "default": DefaultObservationsCfg(), + "noise_less": NoiselessObservationsCfg(), + } + +With Style 3, the ``"default"`` preset is automatically applied when no preset is selected. + + +Using Presets +^^^^^^^^^^^^^ + +**Path presets** - Select a specific preset for one config path: + +.. code-block:: bash + + # Use relative joint position action + python train.py --task=Isaac-Reach-Franka-v0 \ + env.actions.arm_action=relative_joint_position + + # Use noiseless observations + python train.py --task=Isaac-Reach-Franka-v0 \ + env.observations=noise_less + +**Path preset + scalar override** - Select preset then modify a field: + +.. code-block:: bash + + python train.py --task=Isaac-Reach-Franka-v0 \ + env.actions.arm_action=relative_joint_position \ + env.actions.arm_action.scale=0.5 + +**Global presets** - Apply the same preset name everywhere it exists: + +.. code-block:: bash + + # Apply "noise_less" preset to all configs that define it + # (e.g., env.observations in Isaac-Reach-Franka-v0) + python train.py --task=Isaac-Reach-Franka-v0 \ + presets=noise_less + +**Multiple global presets** - Apply several non-conflicting presets (each applies to configs that define it): + +.. code-block:: bash + + # Apply noise_less to env.observations and relative_joint_position to env.actions.arm_action + python train.py --task=Isaac-Reach-Franka-v0 \ + presets=noise_less,relative_joint_position + +**Combined** - Global presets + path presets + scalar overrides: + +.. code-block:: bash + + python train.py --task=Isaac-Reach-Franka-v0 \ + presets=noise_less \ + env.actions.arm_action=relative_joint_position \ + env.actions.arm_action.scale=0.5 \ + env.sim.dt=0.002 + + +Global Preset Conflict Detection +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +If multiple global presets define the same path, an error is raised: + +.. code-block:: bash + + # ERROR: both "fast" and "noise_less" define env.observations + python train.py --task=Isaac-Reach-Franka-v0 \ + presets=fast,noise_less + + # ValueError: Conflicting global presets: 'fast' and 'noise_less' + # both define preset for 'env.observations' + + +Real-World Example +^^^^^^^^^^^^^^^^^^ + +The Franka Reach environment demonstrates presets in practice: + +.. literalinclude:: ../../../source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/joint_pos_env_cfg.py + :language: python + :start-at: @configclass + :end-before: ## + +This allows users to switch action types: + +.. code-block:: bash + + # Default: JointPositionActionCfg (from inheritance) + python train.py --task=Isaac-Reach-Franka-v0 + + # Switch to relative joint position + python train.py --task=Isaac-Reach-Franka-v0 \ + env.actions.arm_action=relative_joint_position + + # Switch to joint position with limits + python train.py --task=Isaac-Reach-Franka-v0 \ + env.actions.arm_action=joint_position_to_limit + + +Summary +^^^^^^^ + +.. list-table:: + :widths: 25 35 40 + :header-rows: 1 + + * - Override Type + - Syntax + - Effect + * - Scalar + - ``env.sim.dt=0.001`` + - Modify single field + * - Path preset + - ``env.actions.arm_action=relative`` + - Replace entire section + * - Global preset + - ``presets=noise_less`` + - Apply everywhere matching + * - Combined + - ``presets=noise_less env.sim.dt=0.001`` + - Global + scalar overrides diff --git a/source/isaaclab_tasks/config/extension.toml b/source/isaaclab_tasks/config/extension.toml index 28311023316..6bff2740334 100644 --- a/source/isaaclab_tasks/config/extension.toml +++ b/source/isaaclab_tasks/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "1.1.0" +version = "1.2.0" # Description title = "Isaac Lab Environments" diff --git a/source/isaaclab_tasks/docs/CHANGELOG.rst b/source/isaaclab_tasks/docs/CHANGELOG.rst index 6516d50ff99..569e114c081 100644 --- a/source/isaaclab_tasks/docs/CHANGELOG.rst +++ b/source/isaaclab_tasks/docs/CHANGELOG.rst @@ -1,6 +1,26 @@ Changelog --------- +1.2.0 (2026-02-22) +~~~~~~~~~~~~~~~~~~ + +Changed +^^^^^^^ + +* Documented the Hydra override order: **Presets > Groups > Single field**; within same type, + **lower depth > higher depth** and **left > right** (rightmost wins for same target). +* Updated Hydra documentation and examples to use real preset names (``noise_less``, + ``relative_joint_position``) instead of hypothetical ones (``inference``, ``newton``). + +Added +^^^^^ + +* Added override order tests in :mod:`isaaclab_tasks.test.test_hydra`: + * ``test_override_order_groups_override_global_presets`` — path preset overrides global preset + * ``test_override_order_single_field_overrides_preset`` — scalar overrides preset value + * ``test_override_order_left_right_scalar_wins`` — rightmost scalar wins for same target + + 1.1.0 (2026-02-13) ~~~~~~~~~~~~~~~~~~ diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/agents/rl_games_ppo_cfg.yaml b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/agents/rl_games_ppo_cfg.yaml index 09e4f9d48b5..c3109b4aad1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/agents/rl_games_ppo_cfg.yaml +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/config/franka/agents/rl_games_ppo_cfg.yaml @@ -81,3 +81,15 @@ params: clip_value: True clip_actions: False bounds_loss_coef: 0.0001 + +variants: + params.network.mlp: + large_network: + units: [256, 128, 64] + activation: elu + d2rl: False + + initializer: + name: default + regularizer: + name: None diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py b/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py index 525b425917f..16785f326f8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py @@ -3,105 +3,354 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Sub-module with utilities for the hydra configuration system.""" +"""Hydra utilities with REPLACE-only preset system. + +This module bypasses Hydra's default MERGE behavior for config groups. +Instead, when a preset is selected, the entire config section is REPLACED +with the preset - no field merging. + +Presets are defined on individual configclasses using a `presets` dict. +The system recursively discovers all presets and their paths automatically. + +Override order (strict): Presets > Groups > Single field. + Within same type: lower depth > higher depth; left > right (rightmost wins). + 1. Global presets: `presets=noise_less` -> apply everywhere matching + 2. Path presets: `env.actions=joint_control` -> REPLACE specific section + 3. Preset-path scalars: `env.actions.scale=5.0` -> handled by us + 4. Global scalars: `env.sim.dt=0.01` -> handled by Hydra + +Example usage: + # Apply "noise_less" preset everywhere it exists, then override scale + presets=noise_less env.actions.arm_action.scale=0.5 env.sim.dt=0.01 +""" import functools -from collections.abc import Callable +import sys +from collections.abc import Callable, Mapping try: import hydra from hydra.core.config_store import ConfigStore - from omegaconf import DictConfig, OmegaConf + from omegaconf import OmegaConf except ImportError: - raise ImportError("Hydra is not installed. Please install it by running 'pip install hydra-core'.") + raise ImportError("Hydra not installed. Run: pip install hydra-core") -from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg from isaaclab.envs.utils.spaces import replace_env_cfg_spaces_with_strings, replace_strings_with_env_cfg_spaces from isaaclab.utils import replace_slices_with_strings, replace_strings_with_slices from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry -def register_task_to_hydra( - task_name: str, agent_cfg_entry_point: str -) -> tuple[ManagerBasedRLEnvCfg | DirectRLEnvCfg, dict]: - """Register the task configuration to the Hydra configuration store. +def collect_presets(cfg, path: str = "") -> dict: + """Recursively walk config tree and collect presets with auto-discovered paths. - This function resolves the configuration file for the environment and agent based on the task's name. - It then registers the configurations to the Hydra configuration store. + Each configclass can define a `presets` dict attribute. Presets are collected + at the path where they're found. The dict can have two formats: + + 1. Simple format (presets for the current field): + presets = {"preset_name": ConfigInstance(), ...} + -> Collected at current path + + 2. Nested format (presets for child fields, legacy support): + presets = {"child.path": {"preset_name": ConfigInstance()}, ...} + -> Collected at specified sub-path Args: - task_name: The name of the task. - agent_cfg_entry_point: The entry point key to resolve the agent's configuration file. + cfg: A configclass instance to walk. + path: Current path prefix (used during recursion). Returns: - A tuple containing the parsed environment and agent configuration objects. + Dict mapping paths to preset dicts, e.g.: + {"actions.arm_action": {"relative_joint_position": , ...}} """ - # load the configurations - env_cfg = load_cfg_from_registry(task_name, "env_cfg_entry_point") - agent_cfg = None - if agent_cfg_entry_point: - agent_cfg = load_cfg_from_registry(task_name, agent_cfg_entry_point) - # replace gymnasium spaces with strings because OmegaConf does not support them. - # this must be done before converting the env configs to dictionary to avoid internal reinterpretations - env_cfg = replace_env_cfg_spaces_with_strings(env_cfg) - # convert the configs to dictionary - env_cfg_dict = env_cfg.to_dict() - if isinstance(agent_cfg, dict) or agent_cfg is None: - agent_cfg_dict = agent_cfg - else: - agent_cfg_dict = agent_cfg.to_dict() - cfg_dict = {"env": env_cfg_dict, "agent": agent_cfg_dict} - # replace slices with strings because OmegaConf does not support slices - cfg_dict = replace_slices_with_strings(cfg_dict) - # store the configuration to Hydra - ConfigStore.instance().store(name=task_name, node=cfg_dict) - return env_cfg, agent_cfg + result = {} + # Check if this config has presets + presets = getattr(cfg, "presets", None) + if presets and isinstance(presets, dict): + # Check format: simple (values are configs) or nested (values are dicts) + first_val = next(iter(presets.values()), None) if presets else None + is_nested_format = isinstance(first_val, dict) -def hydra_task_config(task_name: str, agent_cfg_entry_point: str) -> Callable: - """Decorator to handle the Hydra configuration for a task. + if is_nested_format: + # Nested format: {"field.path": {"name": cfg}} + for sub_path, sub_presets in presets.items(): + full_path = f"{path}.{sub_path}" if path else sub_path + result[full_path] = sub_presets + else: + # Simple format: {"name": cfg} - presets for current field + if path: + result[path] = presets + + # Recurse into nested configclass attributes + for name in dir(cfg): + if name.startswith("_") or name == "presets": + continue + try: + value = getattr(cfg, name) + except Exception: + continue - This decorator registers the task to Hydra and updates the environment and agent configurations from Hydra parsed - command line arguments. + # Check if it's a configclass (has __dataclass_fields__) + if hasattr(value, "__dataclass_fields__"): + child_path = f"{path}.{name}" if path else name + result.update(collect_presets(value, child_path)) + + return result + + +def hydra_task_config(task_name: str, agent_cfg_entry_point: str) -> Callable: + """Decorator for Hydra config with REPLACE-only preset semantics. Args: - task_name: The name of the task. - agent_cfg_entry_point: The entry point key to resolve the agent's configuration file. + task_name: Task name (e.g., "Isaac-Reach-Franka-v0") + agent_cfg_entry_point: Agent config entry point key Returns: - The decorated function with the envrionment's and agent's configurations updated from command line arguments. + Decorated function receiving (env_cfg, agent_cfg, *args, **kwargs) """ def decorator(func): @functools.wraps(func) def wrapper(*args, **kwargs): - # register the task to Hydra - env_cfg, agent_cfg = register_task_to_hydra(task_name.split(":")[-1], agent_cfg_entry_point) - - # define the new Hydra main function - @hydra.main(config_path=None, config_name=task_name.split(":")[-1], version_base="1.3") - def hydra_main(hydra_env_cfg: DictConfig, env_cfg=env_cfg, agent_cfg=agent_cfg): - # convert to a native dictionary - hydra_env_cfg = OmegaConf.to_container(hydra_env_cfg, resolve=True) - # replace string with slices because OmegaConf does not support slices - hydra_env_cfg = replace_strings_with_slices(hydra_env_cfg) - # update the configs with the Hydra command line arguments - env_cfg.from_dict(hydra_env_cfg["env"]) - # replace strings that represent gymnasium spaces because OmegaConf does not support them. - # this must be done after converting the env configs from dictionary to avoid internal reinterpretations + task = task_name.split(":")[-1] + env_cfg, agent_cfg, presets = register_task(task, agent_cfg_entry_point) + + # Split args: global presets, path presets, scalars + global_presets, preset_sel, preset_scalar, global_scalar = parse_overrides( + sys.argv[1:], presets + ) + + # Only pass global scalars to Hydra + original_argv, sys.argv = sys.argv, [sys.argv[0]] + global_scalar + + @hydra.main(config_path=None, config_name=task, version_base="1.3") + def hydra_main(hydra_cfg, env_cfg=env_cfg, agent_cfg=agent_cfg): + hydra_cfg = replace_strings_with_slices(OmegaConf.to_container(hydra_cfg, resolve=True)) + # Apply our preset handling (REPLACE, not merge) + apply_overrides( + env_cfg, agent_cfg, hydra_cfg, global_presets, preset_sel, preset_scalar, presets + ) + # Sync dict -> config objects + env_cfg.from_dict(hydra_cfg["env"]) env_cfg = replace_strings_with_env_cfg_spaces(env_cfg) - # get agent configs if isinstance(agent_cfg, dict) or agent_cfg is None: - agent_cfg = hydra_env_cfg["agent"] + agent_cfg = hydra_cfg["agent"] else: - agent_cfg.from_dict(hydra_env_cfg["agent"]) - # call the original function + agent_cfg.from_dict(hydra_cfg["agent"]) func(env_cfg, agent_cfg, *args, **kwargs) - # call the new Hydra main function - hydra_main() + try: + hydra_main() + finally: + sys.argv = original_argv return wrapper return decorator + + +def register_task(task_name: str, agent_entry: str) -> tuple: + """Load configs, collect presets recursively, register base config to Hydra. + + Presets are collected from nested configclasses and stored separately - + NOT registered as Hydra groups to avoid Hydra's merge behavior. + + Returns: + (env_cfg, agent_cfg, presets) where presets = + {"env": {"path": {"name": cfg}}, "agent": {...}} + """ + env_cfg = load_cfg_from_registry(task_name, "env_cfg_entry_point") + agent_cfg = None + if agent_entry: + agent_cfg = load_cfg_from_registry(task_name, agent_entry) + + # Collect presets recursively from the config tree + presets = { + "env": collect_presets(env_cfg), + "agent": collect_presets(agent_cfg) if agent_cfg else {}, + } + + # Convert to dict for Hydra (handle gym spaces and slices) + env_cfg = replace_env_cfg_spaces_with_strings(env_cfg) + if agent_cfg is not None and hasattr(agent_cfg, "to_dict"): + agent_dict = agent_cfg.to_dict() + else: + agent_dict = agent_cfg + env_dict = env_cfg.to_dict() # type: ignore[union-attr] + cfg_dict = replace_slices_with_strings({"env": env_dict, "agent": agent_dict}) + + # Register plain config (no groups) - Hydra only handles global scalars + ConfigStore.instance().store(name=task_name, node=OmegaConf.create(cfg_dict)) + return env_cfg, agent_cfg, presets + + +def parse_overrides(args: list[str], presets: dict) -> tuple: + """Categorize command line args by type. + + Args: + args: Command line args (without script name) + presets: {"env": {"path": {"name": cfg}}, "agent": {...}} + + Returns: + (global_presets, preset_sel, preset_scalar, global_scalar) where: + - global_presets: [name, ...] - apply to all matching configs + - preset_sel: [(section, path, name), ...] - REPLACE selections + - preset_scalar: [(full_path, value), ...] - scalars in preset paths + - global_scalar: [arg, ...] - pass to Hydra + """ + # Build lookup of preset group paths (e.g., "env.actions") + preset_paths = {f"{s}.{p}" for s, v in presets.items() for p in v} + global_presets, preset_sel, preset_scalar, global_scalar = [], [], [], [] + + for arg in args: + if "=" not in arg: + global_scalar.append(arg) + continue + + key, val = arg.split("=", 1) + if key == "presets": + # Global presets: presets=name1,name2 -> apply everywhere + global_presets.extend(v.strip() for v in val.split(",") if v.strip()) + elif key in preset_paths: + # Exact match -> preset selection + sec, path = key.split(".", 1) + preset_sel.append((sec, path, val)) + elif any(key.startswith(pp + ".") for pp in preset_paths): + # Prefix match -> scalar within preset path + preset_scalar.append((key, val)) + else: + # No match -> global scalar, let Hydra handle it + global_scalar.append(arg) + + # Sort preset selections: parents before children + preset_sel.sort(key=lambda x: x[1].count(".")) + return global_presets, preset_sel, preset_scalar, global_scalar + + +def apply_overrides( + env_cfg, + agent_cfg, + hydra_cfg: dict, + global_presets: list, + preset_sel: list, + preset_scalar: list, + presets: dict, +): + """Apply preset selections and scalar overrides with REPLACE semantics. + + This is the core function that implements REPLACE (not merge) behavior: + 0. Auto-apply "default" presets for paths not explicitly selected + 1. Global presets (presets=name) apply to ALL configs that have matching preset + 2. Path-specific presets (env.actions=name) replace specific sections + 3. Scalar overrides are applied on top + + Raises: + ValueError: If multiple global presets conflict on the same path. + """ + cfgs = {"env": env_cfg, "agent": agent_cfg} + + # Build set of paths that will be explicitly selected + selected_paths = {f"{sec}.{path}" for sec, path, _ in preset_sel} + for name in global_presets: + for sec in ("env", "agent"): + for path, path_presets in presets.get(sec, {}).items(): + if name in path_presets: + selected_paths.add(f"{sec}.{path}") + + # 0. Auto-apply "default" presets for paths not explicitly selected + for sec in ("env", "agent"): + for path, path_presets in presets.get(sec, {}).items(): + full_path = f"{sec}.{path}" + if full_path not in selected_paths and "default" in path_presets: + node = path_presets["default"] + if cfgs[sec]: + _setattr(cfgs[sec], path, node) + if hasattr(node, "to_dict"): + node_dict = node.to_dict() + else: + node_dict = dict(node) + _setattr(hydra_cfg, full_path, node_dict) + + # 1. Apply global presets - find all paths with matching preset name + # Track which paths are set by which preset to detect conflicts + applied_by: dict[str, str] = {} # full_path -> preset_name + for name in global_presets: + for sec in ("env", "agent"): + for path, path_presets in presets.get(sec, {}).items(): + if name in path_presets: + full_path = f"{sec}.{path}" + if full_path in applied_by: + raise ValueError( + f"Conflicting global presets: '{applied_by[full_path]}' and '{name}' " + f"both define preset for '{full_path}'" + ) + applied_by[full_path] = name + node = path_presets[name] + if cfgs[sec]: + _setattr(cfgs[sec], path, node) + if hasattr(node, "to_dict"): + node_dict = node.to_dict() + else: + node_dict = dict(node) + _setattr(hydra_cfg, f"{sec}.{path}", node_dict) + + # 2. Apply path-specific preset selections (REPLACE entire section) + for sec, path, name in preset_sel: + if path not in presets.get(sec, {}): + raise ValueError(f"Unknown preset group: {sec}.{path}") + if name not in presets[sec][path]: + avail = list(presets[sec][path].keys()) + raise ValueError(f"Unknown preset '{name}' for {sec}.{path}. Available: {avail}") + node = presets[sec][path][name] + if cfgs[sec]: + # REPLACE on config object + _setattr(cfgs[sec], path, node) + # REPLACE on dict (for from_dict sync) + if hasattr(node, "to_dict"): + node_dict = node.to_dict() + else: + node_dict = dict(node) + _setattr(hydra_cfg, f"{sec}.{path}", node_dict) + + # 2. Apply scalar overrides within preset paths + for full_path, val_str in preset_scalar: + if full_path.startswith("env."): + sec, path = "env", full_path[4:] + elif full_path.startswith("agent."): + sec, path = "agent", full_path[6:] + else: + continue + if cfgs[sec]: + val = _parse_val(val_str) + _setattr(cfgs[sec], path, val) + _setattr(hydra_cfg, full_path, val) + + +def _setattr(obj, path: str, val): + """Set nested attribute/key (e.g., "actions.arm_action.scale").""" + *parts, leaf = path.split(".") + for p in parts: + obj = obj[p] if isinstance(obj, Mapping) else getattr(obj, p) + if isinstance(obj, dict): + obj[leaf] = val + else: + setattr(obj, leaf, val) + + +def _parse_val(s: str): + """Parse string to Python value (bool, None, int, float, or str).""" + low = s.lower() + if low == "true": + return True + if low == "false": + return False + if low in ("none", "null"): + return None + try: + return float(s) if "." in s else int(s) + except ValueError: + # Strip quotes if present + if s[0] in "\"'" and s[-1] in "\"'": + return s[1:-1] + return s diff --git a/source/isaaclab_tasks/test/test_hydra.py b/source/isaaclab_tasks/test/test_hydra.py index 5c81cb3e650..f64847ed26f 100644 --- a/source/isaaclab_tasks/test/test_hydra.py +++ b/source/isaaclab_tasks/test/test_hydra.py @@ -5,8 +5,6 @@ """Launch Isaac Sim Simulator first.""" -import sys - from isaaclab.app import AppLauncher # launch the simulator @@ -14,92 +12,443 @@ simulation_app = app_launcher.app -"""Rest everything follows.""" +"""Self-contained tests for Hydra configuration utilities. + +These tests verify the REPLACE-only preset system without depending on +external environment configurations. +""" + +import pytest + +from isaaclab.utils import configclass + +from isaaclab_tasks.utils.hydra import ( + _parse_val, + _setattr, + apply_overrides, + collect_presets, + parse_overrides, +) + + +# ============================================================================= +# Test Configuration Classes with Presets +# ============================================================================= + + +@configclass +class JointPositionActionCfg: + """Default joint position action config.""" + + class_type: str = "JointPositionAction" + asset_name: str = "robot" + joint_names: list = [".*"] + scale: float = 1.0 + offset: float = 0.0 + + +@configclass +class RelativeJointPositionActionCfg: + """Relative joint position action config.""" + + class_type: str = "RelativeJointPositionAction" + asset_name: str = "robot" + joint_names: list = [".*"] + scale: float = 0.2 + offset: float = 0.0 + use_zero_offset: bool = True + + +@configclass +class VelocityActionCfg: + """Velocity action config.""" + + class_type: str = "VelocityAction" + asset_name: str = "robot" + velocity_scale: float = 1.0 + + +@configclass +class ArmActionCfgAutoDefault: + """Arm action config using auto-default pattern (no inheritance).""" + + presets = { + "default": JointPositionActionCfg(), + "joint_position": JointPositionActionCfg(), + "relative_joint_position": RelativeJointPositionActionCfg(), + } + + +@configclass +class ArmActionCfgInheritance(JointPositionActionCfg): + """Arm action config using inheritance pattern.""" + + presets = { + "joint_position": JointPositionActionCfg(), + "relative_joint_position": RelativeJointPositionActionCfg(), + } + + +# Use auto-default pattern as the default for tests +ArmActionCfg = ArmActionCfgAutoDefault + + +@configclass +class JointControlActionsCfg: + """Actions config with arm_action field.""" + + arm_action: ArmActionCfg = ArmActionCfg() + + +@configclass +class VelocityControlActionsCfg: + """Actions config with velocity_command (no arm_action).""" + + velocity_command: VelocityActionCfg = VelocityActionCfg() + + +@configclass +class NoiselessObservationsCfg: + """Noiseless observations.""" + + enable_corruption: bool = False + concatenate_terms: bool = True + noise_scale: float = 0.0 + + +@configclass +class FastObservationsCfg: + """Fast/inference observations - no corruption, simple.""" + + enable_corruption: bool = False + concatenate_terms: bool = False + noise_scale: float = 0.0 + + +@configclass +class ObservationsCfg: + """Observations config with presets.""" + + enable_corruption: bool = True + concatenate_terms: bool = True + noise_scale: float = 0.1 + + presets = { + "noise_less": NoiselessObservationsCfg(), + "fast": FastObservationsCfg(), + } + + +@configclass +class ActionsCfg: + """Actions config with presets for different control modes.""" + + arm_action: ArmActionCfg = ArmActionCfg() + + presets = { + "joint_control": JointControlActionsCfg(), + "velocity_control": VelocityControlActionsCfg(), + } + + +@configclass +class SmallPolicyCfg: + """Small policy network config.""" + + actor_hidden_dims: list = [64, 32] + + +@configclass +class LargePolicyCfg: + """Large policy network config.""" + + actor_hidden_dims: list = [512, 256, 128] + + +@configclass +class FastPolicyCfg: + """Fast/inference policy - smaller network for speed.""" + + actor_hidden_dims: list = [32, 16] + + +@configclass +class PolicyCfg: + """Policy config with presets.""" + + actor_hidden_dims: list = [256, 128] + + presets = { + "small_network": SmallPolicyCfg(), + "large_network": LargePolicyCfg(), + "fast": FastPolicyCfg(), + } + + +@configclass +class SampleEnvCfg: + """Sample environment config with nested configs.""" + + decimation: int = 4 + sim_dt: float = 0.005 + + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + + +@configclass +class SampleAgentCfg: + """Sample agent config.""" + + max_iterations: int = 1000 + learning_rate: float = 3e-4 + policy: PolicyCfg = PolicyCfg() + + +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture +def test_configs(): + """Create fresh test configs and collect presets recursively.""" + env_cfg = SampleEnvCfg() + agent_cfg = SampleAgentCfg() + + # Collect presets recursively (same as hydra.py does) + presets = { + "env": collect_presets(env_cfg), + "agent": collect_presets(agent_cfg), + } + + return env_cfg, agent_cfg, presets + + +# ============================================================================= +# Tests for collect_presets +# ============================================================================= + + +def test_collect_presets(): + """Test collecting presets from all levels of config tree.""" + env_cfg = SampleEnvCfg() + presets = collect_presets(env_cfg) + + # Top-level presets + assert "observations" in presets + assert "actions" in presets + assert "noise_less" in presets["observations"] + assert "velocity_control" in presets["actions"] + + # Nested presets from actions.arm_action + assert "actions.arm_action" in presets + assert "relative_joint_position" in presets["actions.arm_action"] + + +# ============================================================================= +# Tests for parse_overrides +# ============================================================================= + + +def test_parse_overrides_mixed(test_configs): + """Mix of all override types with proper categorization.""" + _, _, presets = test_configs + args = [ + "presets=fast", # global preset (applies to observations AND policy) + "env.decimation=10", # global scalar + "env.observations=noise_less", # path preset + "env.actions.arm_action=relative_joint_position", # nested preset + "env.actions.arm_action.scale=2.0", # preset scalar + ] + global_presets, preset_sel, preset_scalar, global_scalar = parse_overrides(args, presets) + + assert global_presets == ["fast"] + assert ("env", "observations", "noise_less") in preset_sel + assert ("env", "actions.arm_action", "relative_joint_position") in preset_sel + assert ("env.actions.arm_action.scale", "2.0") in preset_scalar + assert "env.decimation=10" in global_scalar + + +def test_parse_overrides_sorted_by_depth(test_configs): + """Parent presets should be applied before children.""" + _, _, presets = test_configs + args = [ + "env.actions.arm_action=relative_joint_position", # nested (depth 2) + "env.actions=joint_control", # parent (depth 1) + ] + _, preset_sel, _, _ = parse_overrides(args, presets) + + # joint_control (depth 1) should come before arm_action (depth 2) + assert preset_sel[0] == ("env", "actions", "joint_control") + assert preset_sel[1] == ("env", "actions.arm_action", "relative_joint_position") + + +# ============================================================================= +# Tests for apply_overrides +# ============================================================================= + + +def test_apply_overrides_global_preset(test_configs): + """Global preset should apply to all matching paths with same name.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + # Check defaults before applying + assert env_cfg.observations.enable_corruption is True + assert agent_cfg.policy.actor_hidden_dims == [256, 128] + + # Use global preset "fast" - should apply to BOTH observations AND policy + apply_overrides(env_cfg, agent_cfg, hydra_cfg, ["fast"], [], [], presets) + + # Should have applied FastObservationsCfg to env.observations + assert env_cfg.observations.enable_corruption is False + assert env_cfg.observations.concatenate_terms is False + + # Should have applied FastPolicyCfg to agent.policy + assert agent_cfg.policy.actor_hidden_dims == [32, 16] + + +def test_apply_overrides_multiple_global_presets(test_configs): + """Multiple non-conflicting global presets should all apply.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + # noise_less only affects observations, large_network only affects policy + apply_overrides(env_cfg, agent_cfg, hydra_cfg, ["noise_less", "large_network"], [], [], presets) + + # noise_less applied to observations + assert env_cfg.observations.enable_corruption is False + assert env_cfg.observations.noise_scale == 0.0 + + # large_network applied to policy + assert agent_cfg.policy.actor_hidden_dims == [512, 256, 128] + + +def test_apply_overrides_conflicting_global_presets(test_configs): + """Conflicting global presets should raise ValueError.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + # Both "fast" and "noise_less" define a preset for observations -> conflict + with pytest.raises(ValueError, match="Conflicting global presets"): + apply_overrides(env_cfg, agent_cfg, hydra_cfg, ["fast", "noise_less"], [], [], presets) + + +def test_apply_overrides_preset_with_scalars(test_configs): + """Preset selection + scalar overrides on new preset.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + preset_sel = [("env", "actions.arm_action", "relative_joint_position")] + preset_scalar = [ + ("env.actions.arm_action.scale", "5.0"), + ("env.actions.arm_action.use_zero_offset", "false"), + ] + + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], preset_sel, preset_scalar, presets) + + # Preset selected (RelativeJointPositionActionCfg has use_zero_offset) + assert hasattr(env_cfg.actions.arm_action, "use_zero_offset") + assert env_cfg.actions.arm_action.scale == 5.0 + assert env_cfg.actions.arm_action.use_zero_offset is False + + +def test_apply_overrides_nested_groups(test_configs): + """Select parent group, then nested group, then scalar.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + preset_sel = [ + ("env", "actions", "joint_control"), + ("env", "actions.arm_action", "relative_joint_position"), + ] + preset_scalar = [("env.actions.arm_action.scale", "7.0")] + + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], preset_sel, preset_scalar, presets) + + assert isinstance(env_cfg.actions, JointControlActionsCfg) + assert hasattr(env_cfg.actions.arm_action, "use_zero_offset") + assert env_cfg.actions.arm_action.scale == 7.0 + + +def test_apply_overrides_structural_replacement(test_configs): + """Selecting velocity_control replaces structure (removes arm_action).""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + preset_sel = [("env", "actions", "velocity_control")] + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], preset_sel, [], presets) + + assert isinstance(env_cfg.actions, VelocityControlActionsCfg) + assert hasattr(env_cfg.actions, "velocity_command") + assert not hasattr(env_cfg.actions, "arm_action") or env_cfg.actions.arm_action is None + + +def test_apply_overrides_unknown_raises(test_configs): + """Unknown preset group or name should raise ValueError.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + with pytest.raises(ValueError, match="Unknown preset group"): + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], [("env", "nonexistent", "x")], [], presets) + + with pytest.raises(ValueError, match="Unknown preset"): + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], [("env", "observations", "bad")], [], presets) -import functools -from collections.abc import Callable -import hydra -from hydra import compose, initialize -from omegaconf import OmegaConf +# ============================================================================= +# Override order tests (Presets > Groups > Single field; depth; left > right) +# ============================================================================= -from isaaclab.utils import replace_strings_with_slices -import isaaclab_tasks # noqa: F401 -from isaaclab_tasks.utils.hydra import register_task_to_hydra +def test_override_order_groups_override_global_presets(test_configs): + """Path preset (Groups) overrides global preset (Presets) for the same path.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + # Global preset "noise_less" sets observations; path preset "fast" overrides for same path + apply_overrides( + env_cfg, + agent_cfg, + hydra_cfg, + global_presets=["noise_less"], + preset_sel=[("env", "observations", "fast")], + preset_scalar=[], + presets=presets, + ) -def hydra_task_config_test(task_name: str, agent_cfg_entry_point: str) -> Callable: - """Copied from hydra.py hydra_task_config, since hydra.main requires a single point of entry, - which will not work with multiple tests. Here, we replace hydra.main with hydra initialize - and compose.""" + # Path preset (fast) wins: concatenate_terms=False, not noise_less's True + assert env_cfg.observations.concatenate_terms is False + assert env_cfg.observations.enable_corruption is False - def decorator(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - # register the task to Hydra - env_cfg, agent_cfg = register_task_to_hydra(task_name, agent_cfg_entry_point) - # replace hydra.main with initialize and compose - with initialize(config_path=None, version_base="1.3"): - hydra_env_cfg = compose(config_name=task_name, overrides=sys.argv[1:]) - # convert to a native dictionary - hydra_env_cfg = OmegaConf.to_container(hydra_env_cfg, resolve=True) - # replace string with slices because OmegaConf does not support slices - hydra_env_cfg = replace_strings_with_slices(hydra_env_cfg) - # update the configs with the Hydra command line arguments - env_cfg.from_dict(hydra_env_cfg["env"]) - if isinstance(agent_cfg, dict): - agent_cfg = hydra_env_cfg["agent"] - else: - agent_cfg.from_dict(hydra_env_cfg["agent"]) - # call the original function - func(env_cfg, agent_cfg, *args, **kwargs) +def test_override_order_single_field_overrides_preset(test_configs): + """Scalar (Single field) overrides preset value.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} - return wrapper + # RelativeJointPositionActionCfg has scale=0.2; scalar overrides to 9.0 + preset_sel = [("env", "actions.arm_action", "relative_joint_position")] + preset_scalar = [("env.actions.arm_action.scale", "9.0")] + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], preset_sel, preset_scalar, presets) - return decorator + assert env_cfg.actions.arm_action.scale == 9.0 -def test_hydra(): - """Test the hydra configuration system.""" +def test_override_order_left_right_scalar_wins(test_configs): + """For same target, rightmost scalar wins (left-to-right processing).""" + _, _, presets = test_configs + env_cfg = SampleEnvCfg() + agent_cfg = SampleAgentCfg() + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} - # set hardcoded command line arguments - sys.argv = [ - sys.argv[0], - "env.decimation=42", # test simple env modification - "env.events.physics_material.params.asset_cfg.joint_ids='slice(0 ,1, 2)'", # test slice setting - "env.scene.robot.init_state.joint_vel={.*: 4.0}", # test regex setting - "env.rewards.feet_air_time=null", # test setting to none - "agent.max_iterations=3", # test simple agent modification + # Two scalars for same field - rightmost (2.0) should win + preset_sel = [("env", "actions.arm_action", "relative_joint_position")] + preset_scalar = [ + ("env.actions.arm_action.scale", "1.0"), + ("env.actions.arm_action.scale", "2.0"), ] + apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], preset_sel, preset_scalar, presets) - @hydra_task_config_test("Isaac-Velocity-Flat-H1-v0", "rsl_rl_cfg_entry_point") - def main(env_cfg, agent_cfg): - # env - assert env_cfg.decimation == 42 - assert env_cfg.events.physics_material.params["asset_cfg"].joint_ids == slice(0, 1, 2) - assert env_cfg.scene.robot.init_state.joint_vel == {".*": 4.0} - assert env_cfg.rewards.feet_air_time is None - # agent - assert agent_cfg.max_iterations == 3 - - main() - # clean up - sys.argv = [sys.argv[0]] - hydra.core.global_hydra.GlobalHydra.instance().clear() - - -def test_nested_iterable_dict(): - """Test the hydra configuration system when dict is nested in an Iterable.""" - - @hydra_task_config_test("Isaac-Lift-Cube-Franka-v0", "rsl_rl_cfg_entry_point") - def main(env_cfg, agent_cfg): - # env - assert env_cfg.scene.ee_frame.target_frames[0].name == "end_effector" - assert env_cfg.scene.ee_frame.target_frames[0].offset.pos[2] == 0.1034 - - main() - # clean up - sys.argv = [sys.argv[0]] - hydra.core.global_hydra.GlobalHydra.instance().clear() + assert env_cfg.actions.arm_action.scale == 2.0 From e2239f527d107b71f9bac3c075cc728452624c3b Mon Sep 17 00:00:00 2001 From: Octi Zhang Date: Tue, 24 Feb 2026 21:23:13 -0800 Subject: [PATCH 4/4] save --- docs/source/refs/contributing.rst | 56 ++ scripts/benchmarks/benchmark_load_robot.py | 2 +- scripts/demos/arms.py | 2 +- scripts/demos/bipeds.py | 2 +- scripts/demos/hands.py | 2 +- scripts/demos/haply_teleoperation.py | 4 +- scripts/demos/multi_asset.py | 2 +- scripts/demos/pick_and_place.py | 2 +- scripts/demos/quadcopter.py | 2 +- scripts/demos/quadrupeds.py | 2 +- scripts/demos/sensors/cameras.py | 2 +- scripts/demos/sensors/contact_sensor.py | 2 +- .../demos/sensors/frame_transformer_sensor.py | 2 +- scripts/demos/sensors/imu_sensor.py | 2 +- scripts/demos/sensors/multi_mesh_raycaster.py | 2 +- .../sensors/multi_mesh_raycaster_camera.py | 2 +- scripts/demos/sensors/raycaster_sensor.py | 2 +- .../reinforcement_learning/rl_games/train.py | 39 +- .../reinforcement_learning/rsl_rl/train.py | 45 +- scripts/tutorials/01_assets/add_new_robot.py | 4 +- .../tutorials/01_assets/run_articulation.py | 2 +- .../01_assets/run_surface_gripper.py | 2 +- scripts/tutorials/02_scene/create_scene.py | 2 +- .../04_sensors/add_sensors_on_robot.py | 2 +- .../tutorials/05_controllers/run_diff_ik.py | 2 +- scripts/tutorials/05_controllers/run_osc.py | 2 +- .../isaaclab/isaaclab/actuators/__init__.py | 29 +- source/isaaclab/isaaclab/assets/__init__.py | 53 +- .../isaaclab/assets/articulation/__init__.py | 23 +- .../assets/articulation/articulation_cfg.py | 6 +- .../assets/articulation/base_articulation.py | 1 + source/isaaclab/isaaclab/assets/asset_base.py | 13 +- .../isaaclab/assets/asset_base_cfg.py | 9 +- .../isaaclab/assets/rigid_object/__init__.py | 23 +- .../assets/rigid_object/rigid_object_cfg.py | 5 +- .../rigid_object_collection/__init__.py | 23 +- .../rigid_object_collection_cfg.py | 6 +- .../isaaclab/isaaclab/cli/commands/install.py | 2 + source/isaaclab/isaaclab/cloner/__init__.py | 21 +- source/isaaclab/isaaclab/devices/__init__.py | 23 +- .../isaaclab/devices/openxr/__init__.py | 23 +- source/isaaclab/isaaclab/envs/__init__.py | 40 +- .../isaaclab/isaaclab/envs/direct_marl_env.py | 33 +- .../isaaclab/envs/direct_marl_env_cfg.py | 11 +- .../isaaclab/isaaclab/envs/direct_rl_env.py | 40 +- .../isaaclab/envs/direct_rl_env_cfg.py | 10 +- .../isaaclab/envs/manager_based_env.py | 47 +- .../isaaclab/envs/manager_based_env_cfg.py | 16 +- .../isaaclab/envs/manager_based_rl_env_cfg.py | 5 +- source/isaaclab/isaaclab/envs/mdp/__init__.py | 154 ++++- .../isaaclab/envs/mdp/actions/__init__.py | 21 +- .../isaaclab/envs/mdp/actions/actions_cfg.py | 43 +- .../mdp/actions/surface_gripper_actions.py | 2 +- .../isaaclab/envs/mdp/commands/__init__.py | 29 +- .../envs/mdp/commands/commands_cfg.py | 25 +- .../envs/mdp/commands/velocity_command.py | 7 + source/isaaclab/isaaclab/envs/mdp/events.py | 60 +- .../isaaclab/envs/mdp/observations.py | 4 +- source/isaaclab/isaaclab/envs/mdp/rewards.py | 4 +- .../isaaclab/envs/mdp/terminations.py | 4 +- source/isaaclab/isaaclab/envs/ui/__init__.py | 15 +- .../isaaclab/envs/ui/base_env_window.py | 7 +- .../isaaclab/isaaclab/envs/ui/empty_window.py | 2 - .../envs/ui/viewport_camera_controller.py | 5 +- source/isaaclab/isaaclab/envs/utils/marl.py | 13 +- source/isaaclab/isaaclab/managers/__init__.py | 43 +- .../isaaclab/managers/action_manager.py | 4 +- .../isaaclab/managers/command_manager.py | 8 +- .../isaaclab/managers/manager_base.py | 34 +- .../isaaclab/managers/manager_term_cfg.py | 36 +- .../isaaclab/managers/recorder_manager.py | 35 +- source/isaaclab/isaaclab/markers/__init__.py | 14 +- .../isaaclab/markers/config/__init__.py | 2 +- .../isaaclab/markers/visualization_markers.py | 44 +- .../markers/visualization_markers_cfg.py | 28 + .../isaaclab/physics/physics_manager.py | 2 - .../renderers/newton_warp_renderer.py | 60 +- source/isaaclab/isaaclab/scene/__init__.py | 11 +- .../isaaclab/scene/interactive_scene.py | 27 +- source/isaaclab/isaaclab/sensors/__init__.py | 59 +- .../isaaclab/sensors/camera/__init__.py | 20 +- .../isaaclab/sensors/camera/camera.py | 10 +- .../isaaclab/sensors/camera/camera_cfg.py | 5 +- .../isaaclab/sensors/camera/tiled_camera.py | 7 +- .../sensors/camera/tiled_camera_cfg.py | 5 +- .../sensors/contact_sensor/__init__.py | 17 +- .../contact_sensor/base_contact_sensor.py | 19 +- .../contact_sensor/contact_sensor_cfg.py | 46 +- .../sensors/frame_transformer/__init__.py | 17 +- .../frame_transformer_cfg.py | 5 +- .../isaaclab/isaaclab/sensors/imu/__init__.py | 17 +- .../isaaclab/isaaclab/sensors/imu/imu_cfg.py | 5 +- .../isaaclab/sensors/ray_caster/__init__.py | 31 +- .../ray_caster/multi_mesh_ray_caster.py | 4 +- .../multi_mesh_ray_caster_camera_cfg.py | 5 +- .../ray_caster/multi_mesh_ray_caster_cfg.py | 5 +- .../sensors/ray_caster/ray_cast_utils.py | 7 +- .../ray_caster/ray_caster_camera_cfg.py | 5 +- .../sensors/ray_caster/ray_caster_cfg.py | 5 +- .../isaaclab/isaaclab/sensors/sensor_base.py | 150 +++-- .../isaaclab/sensors/sensor_base_cfg.py | 6 +- source/isaaclab/isaaclab/sim/__init__.py | 60 +- .../isaaclab/sim/backend_detection.py | 35 ++ .../isaaclab/sim/converters/__init__.py | 17 +- .../isaaclab/sim/converters/mesh_converter.py | 6 +- .../isaaclab/sim/converters/urdf_converter.py | 12 +- .../sim/scene_data_providers/__init__.py | 2 + .../newton_scene_data_provider.py | 42 +- .../physx_scene_data_provider.py | 11 +- .../isaaclab/isaaclab/sim/schemas/__init__.py | 100 ++-- .../isaaclab/isaaclab/sim/schemas/schemas.py | 120 +++- .../isaaclab/sim/simulation_context.py | 39 +- .../isaaclab/sim/spawners/__init__.py | 36 +- .../sim/spawners/from_files/__init__.py | 21 +- .../sim/spawners/from_files/from_files.py | 69 ++- .../sim/spawners/from_files/from_files_cfg.py | 15 +- .../isaaclab/sim/spawners/lights/__init__.py | 11 +- .../sim/spawners/lights/lights_cfg.py | 7 +- .../sim/spawners/materials/__init__.py | 15 +- .../materials/physics_materials_cfg.py | 11 +- .../materials/visual_materials_cfg.py | 13 +- .../isaaclab/sim/spawners/meshes/__init__.py | 24 +- .../sim/spawners/meshes/meshes_cfg.py | 15 +- .../isaaclab/sim/spawners/sensors/__init__.py | 11 +- .../sim/spawners/sensors/sensors_cfg.py | 9 +- .../isaaclab/sim/spawners/shapes/__init__.py | 11 +- .../sim/spawners/shapes/shapes_cfg.py | 15 +- .../isaaclab/sim/spawners/spawner_cfg.py | 6 +- .../sim/spawners/wrappers/__init__.py | 11 +- .../sim/spawners/wrappers/wrappers_cfg.py | 8 +- .../isaaclab/isaaclab/sim/utils/__init__.py | 31 +- source/isaaclab/isaaclab/sim/utils/prims.py | 71 ++- source/isaaclab/isaaclab/sim/utils/queries.py | 2 + source/isaaclab/isaaclab/sim/utils/stage.py | 10 +- .../isaaclab/isaaclab/sim/views/__init__.py | 9 +- .../isaaclab/sim/views/xform_prim_view.py | 6 +- source/isaaclab/isaaclab/terrains/__init__.py | 48 +- .../terrains/height_field/__init__.py | 27 +- .../terrains/height_field/hf_terrains_cfg.py | 14 +- .../isaaclab/terrains/sub_terrain_cfg.py | 8 +- .../terrains/terrain_generator_cfg.py | 4 +- .../isaaclab/terrains/terrain_importer_cfg.py | 6 +- .../isaaclab/terrains/trimesh/__init__.py | 35 +- .../terrains/trimesh/mesh_terrains_cfg.py | 31 +- .../ui/widgets/manager_live_visualizer.py | 4 +- .../ui/xr_widgets/instruction_widget.py | 3 +- source/isaaclab/isaaclab/utils/__init__.py | 95 +++- source/isaaclab/isaaclab/utils/assets.py | 230 +++++--- .../isaaclab/isaaclab/utils/backend_utils.py | 21 +- .../isaaclab/utils/buffers/__init__.py | 15 +- source/isaaclab/isaaclab/utils/dict.py | 20 +- source/isaaclab/isaaclab/utils/helpers.py | 301 ++++++++++ source/isaaclab/isaaclab/utils/string.py | 59 ++ source/isaaclab/isaaclab/utils/timer.py | 112 ++-- source/isaaclab/isaaclab/utils/version.py | 10 + .../isaaclab/isaaclab/utils/warp/__init__.py | 3 +- source/isaaclab/isaaclab/utils/warp/utils.py | 30 + source/isaaclab/setup.py | 3 + .../test/assets/check_external_force.py | 2 +- .../test/assets/check_fixed_base_assets.py | 2 +- .../test/assets/check_ridgeback_franka.py | 2 +- .../test/controllers/test_differential_ik.py | 2 +- .../controllers/test_operational_space.py | 2 +- .../test/envs/test_color_randomization.py | 4 +- .../test/markers/check_markers_visibility.py | 2 +- .../test/scene/check_interactive_scene.py | 4 +- .../test/sim/test_simulation_render_config.py | 4 +- .../sim/test_simulation_stage_in_memory.py | 8 +- .../test/sim/test_spawn_from_files.py | 4 +- .../isaaclab/test/sim/test_urdf_converter.py | 4 +- source/isaaclab/test/utils/test_version.py | 10 +- source/isaaclab_assets/data/Robots | 1 + .../isaaclab_assets/sensors/__init__.py | 11 +- .../isaaclab_contrib/sensors/__init__.py | 14 +- .../sensors/tacsl_sensor/__init__.py | 13 +- .../tacsl_sensor/visuotactile_sensor_cfg.py | 8 +- .../isaaclab_newton/actuators/__init__.py | 29 + .../actuators/actuator_base.py | 275 +++++++++ .../isaaclab_newton/actuators/actuator_net.py | 203 +++++++ .../isaaclab_newton/actuators/actuator_pd.py | 462 +++++++++++++++ .../isaaclab_newton/actuators/kernels.py | 80 +++ .../assets/articulation/articulation.py | 36 +- .../isaaclab_newton/cloner/__init__.py | 9 +- .../cloner/newton_replicate.py | 112 +++- .../isaaclab_newton/physics/__init__.py | 24 + .../isaaclab_newton/physics/newton_manager.py | 509 +++++++++++++++++ .../physics/newton_manager_cfg.py | 215 +++++++ .../isaaclab_newton/sensors/__init__.py | 16 + .../sensors/contact_sensor/__init__.py | 17 + .../sensors/contact_sensor/contact_sensor.py | 529 ++++++++++++++++++ .../contact_sensor/contact_sensor_cfg.py | 45 ++ .../contact_sensor/contact_sensor_data.py | 257 +++++++++ .../contact_sensor/contact_sensor_kernels.py | 213 +++++++ .../sensors/tiled_camera_cfg.py | 25 + source/isaaclab_newton/setup.py | 4 +- .../isaaclab_physx/assets/__init__.py | 32 +- .../assets/articulation/__init__.py | 14 +- .../assets/articulation/articulation.py | 11 + .../assets/deformable_object/__init__.py | 17 +- .../deformable_object_cfg.py | 7 +- .../assets/surface_gripper/__init__.py | 14 +- .../surface_gripper/surface_gripper_cfg.py | 7 +- .../isaaclab_physx/cloner/__init__.py | 9 +- .../isaaclab_physx/physics/__init__.py | 16 +- .../isaaclab_physx/physics/physx_manager.py | 5 + .../physics/physx_manager_cfg.py | 8 +- .../isaaclab_physx/sensors/__init__.py | 21 +- .../sensors/contact_sensor/__init__.py | 12 +- .../sensors/contact_sensor/contact_sensor.py | 12 +- .../contact_sensor/contact_sensor_cfg.py | 41 ++ .../test/sensors/check_contact_sensor.py | 2 +- .../test/sensors/test_frame_transformer.py | 16 +- .../isaaclab_rl/rsl_rl/__init__.py | 30 +- .../isaaclab_rl/rsl_rl/vecenv_wrapper.py | 11 +- .../isaaclab_tasks/direct/ant/__init__.py | 2 +- .../isaaclab_tasks/direct/ant/ant_env.py | 60 +- .../isaaclab_tasks/direct/ant/ant_env_cfg.py | 66 +++ .../direct/anymal_c/anymal_c_env.py | 2 +- .../direct/automate/__init__.py | 4 +- .../direct/automate/assembly_env.py | 8 +- .../direct/automate/disassembly_env.py | 8 +- .../direct/cart_double_pendulum/__init__.py | 2 +- .../cart_double_pendulum_env.py | 53 +- .../cart_double_pendulum_env_cfg.py | 56 ++ .../direct/cartpole/__init__.py | 14 +- .../direct/cartpole/cartpole_camera_env.py | 186 ++---- .../cartpole/cartpole_camera_env_cfg.py | 193 +++++++ .../direct/cartpole/cartpole_env.py | 63 +-- .../direct/cartpole/cartpole_env_cfg.py | 49 ++ .../cartpole/cartpole_env_cfg.py | 2 +- .../cartpole_camera_env_cfg.py | 2 +- .../direct/factory/factory_env.py | 4 +- .../franka_cabinet/franka_cabinet_env.py | 11 +- .../direct/humanoid_amp/humanoid_amp_env.py | 2 +- .../inhand_manipulation_env.py | 2 +- .../direct/locomotion/locomotion_env.py | 2 +- .../direct/quadcopter/quadcopter_env.py | 2 +- .../shadow_hand_over/shadow_hand_over_env.py | 4 +- .../classic/cartpole/mdp/__init__.py | 6 +- .../classic/cartpole/mdp/rewards.py | 2 +- .../classic/humanoid/mdp/__init__.py | 6 +- .../classic/humanoid/mdp/observations.py | 2 +- .../classic/humanoid/mdp/rewards.py | 2 +- .../manager_based/drone_arl/mdp/__init__.py | 6 +- .../drone_arl/mdp/observations.py | 2 +- .../manager_based/drone_arl/mdp/rewards.py | 6 +- .../pick_place/mdp/__init__.py | 6 +- .../velocity/config/spot/mdp/events.py | 4 +- .../velocity/config/spot/mdp/rewards.py | 4 +- .../locomotion/velocity/mdp/__init__.py | 30 +- .../locomotion/velocity/mdp/curriculums.py | 4 +- .../locomotion/velocity/mdp/rewards.py | 2 +- .../locomotion/velocity/mdp/terminations.py | 2 +- .../locomotion/velocity/velocity_env_cfg.py | 114 ++-- .../manipulation/cabinet/mdp/__init__.py | 6 +- .../manipulation/cabinet/mdp/observations.py | 4 +- .../manipulation/deploy/mdp/__init__.py | 6 +- .../manipulation/deploy/mdp/events.py | 6 +- .../manipulation/deploy/mdp/observations.py | 2 +- .../manipulation/deploy/mdp/rewards.py | 2 +- .../manipulation/deploy/mdp/terminations.py | 11 +- .../manipulation/dexsuite/mdp/__init__.py | 6 +- .../mdp/commands/pose_commands_cfg.py | 9 +- .../manipulation/dexsuite/mdp/curriculums.py | 2 +- .../manipulation/dexsuite/mdp/observations.py | 2 +- .../manipulation/dexsuite/mdp/rewards.py | 4 +- .../manipulation/dexsuite/mdp/terminations.py | 2 +- .../manipulation/inhand/mdp/__init__.py | 6 +- .../manipulation/inhand/mdp/events.py | 4 +- .../manipulation/inhand/mdp/observations.py | 5 +- .../manipulation/inhand/mdp/rewards.py | 5 +- .../manipulation/lift/mdp/__init__.py | 6 +- .../manipulation/lift/mdp/observations.py | 2 +- .../manipulation/lift/mdp/rewards.py | 4 +- .../manipulation/lift/mdp/terminations.py | 2 +- .../manipulation/pick_place/mdp/__init__.py | 6 +- .../pick_place/mdp/terminations.py | 2 +- .../manipulation/place/mdp/__init__.py | 6 +- .../manipulation/place/mdp/observations.py | 4 +- .../manipulation/place/mdp/terminations.py | 2 +- .../manipulation/reach/mdp/__init__.py | 6 +- .../manipulation/reach/mdp/rewards.py | 2 +- .../manipulation/stack/mdp/__init__.py | 6 +- .../stack/mdp/franka_stack_events.py | 4 +- .../manipulation/stack/mdp/observations.py | 4 +- .../manipulation/stack/mdp/terminations.py | 2 +- .../manager_based/navigation/mdp/__init__.py | 6 +- .../mdp/pre_trained_policy_action.py | 2 +- .../isaaclab_tasks/utils/hydra.py | 112 ++++ .../isaaclab_tasks/utils/parse_cfg.py | 6 +- .../benchmarking/env_benchmark_test_utils.py | 5 +- source/isaaclab_tasks/test/env_test_utils.py | 4 +- .../test/test_config_load_no_pxr.py | 148 +++++ .../test_environments_with_stage_in_memory.py | 4 +- source/isaaclab_tasks/test/test_hydra.py | 78 +++ .../deprecated/openxr/manus_vive.py | 9 +- .../deprecated/openxr/manus_vive_utils.py | 3 +- .../templates/tasks/direct_multi-agent/env | 2 +- .../templates/tasks/direct_single-agent/env | 2 +- 299 files changed, 6911 insertions(+), 1852 deletions(-) create mode 100644 source/isaaclab/isaaclab/markers/visualization_markers_cfg.py create mode 100644 source/isaaclab/isaaclab/sim/backend_detection.py create mode 100644 source/isaaclab/isaaclab/utils/helpers.py create mode 100644 source/isaaclab/isaaclab/utils/warp/utils.py create mode 120000 source/isaaclab_assets/data/Robots create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/actuator_base.py create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/actuator_net.py create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/actuator_pd.py create mode 100644 source/isaaclab_newton/isaaclab_newton/actuators/kernels.py create mode 100644 source/isaaclab_newton/isaaclab_newton/physics/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py create mode 100644 source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py create mode 100644 source/isaaclab_newton/isaaclab_newton/sensors/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/__init__.py create mode 100644 source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py create mode 100644 source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_cfg.py create mode 100644 source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_data.py create mode 100644 source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_kernels.py create mode 100644 source/isaaclab_newton/isaaclab_newton/sensors/tiled_camera_cfg.py create mode 100644 source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env_cfg.py create mode 100644 source/isaaclab_tasks/test/test_config_load_no_pxr.py diff --git a/docs/source/refs/contributing.rst b/docs/source/refs/contributing.rst index 411742fd19e..c3b81a474c8 100644 --- a/docs/source/refs/contributing.rst +++ b/docs/source/refs/contributing.rst @@ -316,6 +316,62 @@ without triggering circular imports. It is important to note that this is the sole instance within our codebase where circular imports are used and are acceptable. In all other scenarios, we adhere to best practices and recommend that you do the same. +Lazy Loading and ``__init__.py`` Structure +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Isaac Lab uses the `lazy_loader `__ package +(``lazy.attach``) to defer the loading of heavy modules. This is critical because +certain backend modules (``pxr``, ``omni``, ``carb``, ``isaacsim``) must not be +imported before ``SimulationApp`` is initialized. Importing them too early causes +crashes such as ``free(): invalid pointer``. + +**Configuration vs. Implementation separation** + +The codebase follows a clear separation between *configuration classes* (pure data, +no runtime dependencies) and *implementation classes* (require backend modules at +runtime). In ``__init__.py`` files, this separation is expressed as: + +* **Cfg classes** are imported eagerly at the top of the file. +* **Implementation classes** are deferred via ``lazy.attach()``. + +This pattern makes the architectural intent explicit: configuration classes are +lightweight and always available, while implementation classes are only loaded when +actually used. + +.. code:: python + + # Example: isaaclab/assets/articulation/__init__.py + + import lazy_loader as lazy + + from .articulation_cfg import ArticulationCfg # safe -- pure data, no backend deps + + __getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "base_articulation": ["BaseArticulation"], + "base_articulation_data": ["BaseArticulationData"], + "articulation": ["Articulation"], # deferred -- may chain into backend + "articulation_data": ["ArticulationData"], + }, + ) + __all__ += ["ArticulationCfg"] + +**Rules for deciding eager vs. lazy** + +1. A ``Cfg`` class can be eagerly imported only if its module (and every module it + transitively imports) is free of backend imports (``pxr``, ``omni``, ``carb``, + ``isaacsim``). +2. If a ``Cfg`` module imports an implementation class (e.g. for ``class_type`` + defaults), verify that the implementation module is also backend-free before + making the ``Cfg`` eager. +3. If there is any doubt, keep the ``Cfg`` in the ``lazy.attach()`` block. Safety + over style. +4. For callable fields such as ``func`` in spawner configs, prefer string references + (e.g. ``"isaaclab.sim.spawners.shapes.shapes:spawn_sphere"``) over direct + function references. Strings are inert data and never trigger accidental imports. + + Type-hinting ^^^^^^^^^^^^ diff --git a/scripts/benchmarks/benchmark_load_robot.py b/scripts/benchmarks/benchmark_load_robot.py index 4d13ca4929d..d42d92941ce 100644 --- a/scripts/benchmarks/benchmark_load_robot.py +++ b/scripts/benchmarks/benchmark_load_robot.py @@ -144,7 +144,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(robot.data.default_joint_vel).clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() # Apply random action diff --git a/scripts/demos/arms.py b/scripts/demos/arms.py index ff5ff8c29f3..dbbfb2b372f 100644 --- a/scripts/demos/arms.py +++ b/scripts/demos/arms.py @@ -182,7 +182,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula robot.write_root_velocity_to_sim(root_state[:, 7:]) # set joint positions joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers robot.reset() print("[INFO]: Resetting robots state...") diff --git a/scripts/demos/bipeds.py b/scripts/demos/bipeds.py index daf5ef56e5a..c26bab66615 100644 --- a/scripts/demos/bipeds.py +++ b/scripts/demos/bipeds.py @@ -93,7 +93,7 @@ def run_simulator(sim: sim_utils.SimulationContext, robots: list[Articulation], wp.to_torch(robot.data.default_joint_pos), wp.to_torch(robot.data.default_joint_vel), ) - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) root_state = wp.to_torch(robot.data.default_root_state).clone() root_state[:, :3] += origins[index] robot.write_root_pose_to_sim(root_state[:, :7]) diff --git a/scripts/demos/hands.py b/scripts/demos/hands.py index 45e75c71143..42ebfca92fb 100644 --- a/scripts/demos/hands.py +++ b/scripts/demos/hands.py @@ -119,7 +119,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula wp.to_torch(robot.data.default_joint_pos).clone(), wp.to_torch(robot.data.default_joint_vel).clone(), ) - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset the internal state robot.reset() print("[INFO]: Resetting robots state...") diff --git a/scripts/demos/haply_teleoperation.py b/scripts/demos/haply_teleoperation.py index 6ee741e097b..723ac97eea4 100644 --- a/scripts/demos/haply_teleoperation.py +++ b/scripts/demos/haply_teleoperation.py @@ -195,7 +195,7 @@ def run_simulator( joint_pos = robot.data.default_joint_pos.clone() joint_pos[0, :7] = torch.tensor([0.0, -0.569, 0.0, -2.81, 0.0, 3.037, 0.741], device=robot.device) joint_vel = robot.data.default_joint_vel.clone() - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) for _ in range(10): scene.write_data_to_sim() @@ -253,7 +253,7 @@ def run_simulator( joint_pos = robot.data.default_joint_pos.clone() joint_pos[0, :7] = torch.tensor([0.0, -0.569, 0.0, -2.81, 0.0, 3.037, 0.741], device=robot.device) joint_vel = robot.data.default_joint_vel.clone() - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) cube_state = cube.data.default_root_state.clone() cube_state[:, :3] += scene.env_origins diff --git a/scripts/demos/multi_asset.py b/scripts/demos/multi_asset.py index 5e3b634c4cd..43b35361e92 100644 --- a/scripts/demos/multi_asset.py +++ b/scripts/demos/multi_asset.py @@ -260,7 +260,7 @@ def run_simulator(sim: SimulationContext, scene: InteractiveScene): wp.to_torch(robot.data.default_joint_pos).clone(), wp.to_torch(robot.data.default_joint_vel).clone(), ) - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting scene state...") diff --git a/scripts/demos/pick_and_place.py b/scripts/demos/pick_and_place.py index 8f15b284e7d..b1b6b78400b 100644 --- a/scripts/demos/pick_and_place.py +++ b/scripts/demos/pick_and_place.py @@ -390,7 +390,7 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel - self.pick_and_place.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self.pick_and_place.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, joint_ids=None, env_ids=env_ids) def _set_debug_vis_impl(self, debug_vis: bool): # create markers if necessary for the first tome diff --git a/scripts/demos/quadcopter.py b/scripts/demos/quadcopter.py index a418af6ff1d..9c73119a1fd 100644 --- a/scripts/demos/quadcopter.py +++ b/scripts/demos/quadcopter.py @@ -92,7 +92,7 @@ def main(): count = 0 # reset dof state joint_pos, joint_vel = wp.to_torch(robot.data.default_joint_pos), wp.to_torch(robot.data.default_joint_vel) - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) default_root_state = wp.to_torch(robot.data.default_root_state) robot.write_root_pose_to_sim(default_root_state[:, :7]) robot.write_root_velocity_to_sim(default_root_state[:, 7:]) diff --git a/scripts/demos/quadrupeds.py b/scripts/demos/quadrupeds.py index 3188a3fd309..ef77bd5ab10 100644 --- a/scripts/demos/quadrupeds.py +++ b/scripts/demos/quadrupeds.py @@ -145,7 +145,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula robot.write_root_velocity_to_sim(root_state[:, 7:]) # joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset the internal state robot.reset() print("[INFO]: Resetting robots state...") diff --git a/scripts/demos/sensors/cameras.py b/scripts/demos/sensors/cameras.py index b8d2e2b4842..540b9e11382 100644 --- a/scripts/demos/sensors/cameras.py +++ b/scripts/demos/sensors/cameras.py @@ -200,7 +200,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(scene["robot"].data.default_joint_vel).clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 - scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/demos/sensors/contact_sensor.py b/scripts/demos/sensors/contact_sensor.py index 0c03ac6ddbf..7b3655aa85e 100644 --- a/scripts/demos/sensors/contact_sensor.py +++ b/scripts/demos/sensors/contact_sensor.py @@ -117,7 +117,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(scene["robot"].data.default_joint_vel).clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 - scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/demos/sensors/frame_transformer_sensor.py b/scripts/demos/sensors/frame_transformer_sensor.py index fb643bf514e..7beb062672f 100644 --- a/scripts/demos/sensors/frame_transformer_sensor.py +++ b/scripts/demos/sensors/frame_transformer_sensor.py @@ -113,7 +113,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(scene["robot"].data.default_joint_vel).clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 - scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/demos/sensors/imu_sensor.py b/scripts/demos/sensors/imu_sensor.py index 8870647b27a..8a6c5c03fb7 100644 --- a/scripts/demos/sensors/imu_sensor.py +++ b/scripts/demos/sensors/imu_sensor.py @@ -84,7 +84,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(scene["robot"].data.default_joint_vel).clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 - scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/demos/sensors/multi_mesh_raycaster.py b/scripts/demos/sensors/multi_mesh_raycaster.py index 3fb5920f7ba..daae6dc0bdd 100644 --- a/scripts/demos/sensors/multi_mesh_raycaster.py +++ b/scripts/demos/sensors/multi_mesh_raycaster.py @@ -249,7 +249,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(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) + scene["asset"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting Asset state...") diff --git a/scripts/demos/sensors/multi_mesh_raycaster_camera.py b/scripts/demos/sensors/multi_mesh_raycaster_camera.py index 1b438038d1f..da73602608c 100644 --- a/scripts/demos/sensors/multi_mesh_raycaster_camera.py +++ b/scripts/demos/sensors/multi_mesh_raycaster_camera.py @@ -265,7 +265,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(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) + scene["asset"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting Asset state...") diff --git a/scripts/demos/sensors/raycaster_sensor.py b/scripts/demos/sensors/raycaster_sensor.py index 66411d6cd93..3cad4f07dd1 100644 --- a/scripts/demos/sensors/raycaster_sensor.py +++ b/scripts/demos/sensors/raycaster_sensor.py @@ -101,7 +101,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(scene["robot"].data.default_joint_vel).clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 - scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/reinforcement_learning/rl_games/train.py b/scripts/reinforcement_learning/rl_games/train.py index 5b85ba5b429..57e01e8b68d 100644 --- a/scripts/reinforcement_learning/rl_games/train.py +++ b/scripts/reinforcement_learning/rl_games/train.py @@ -56,9 +56,36 @@ # clear out sys.argv for Hydra sys.argv = [sys.argv[0]] + hydra_args -# launch omniverse app -app_launcher = AppLauncher(args_cli) -simulation_app = app_launcher.app +# Register tasks and resolve config with Hydra (presets applied) +import isaaclab_tasks # noqa: F401 +from isaaclab.sim.backend_detection import physics_backend_requires_kit +from isaaclab_tasks.utils.hydra import finalize_task_config, resolve_task_config + +env_cfg, agent_cfg, hydra_cfg = resolve_task_config(args_cli.task, args_cli.agent) + +# Only launch Isaac Sim Kit when the task's physics backend requires it +_requires_kit = physics_backend_requires_kit(env_cfg) +if _requires_kit: + app_launcher = AppLauncher(args_cli) + simulation_app = app_launcher.app +else: + class _DummySimulationApp: + def close(self): + pass + + class _AppLauncherStub: + """Minimal stand-in for AppLauncher when Kit is not launched (e.g. Newton backend).""" + + def __init__(self): + self.app = _DummySimulationApp() + self.local_rank = 0 + self.global_rank = 0 + + app_launcher = _AppLauncherStub() + simulation_app = app_launcher.app + +# Now that Kit is running (if needed), resolve callable strings in the configs +env_cfg, agent_cfg = finalize_task_config(env_cfg, agent_cfg, hydra_cfg) """Rest everything follows.""" @@ -87,16 +114,12 @@ from isaaclab_rl.rl_games import MultiObserver, PbtAlgoObserver, RlGamesGpuEnv, RlGamesVecEnvWrapper -import isaaclab_tasks # noqa: F401 -from isaaclab_tasks.utils.hydra import hydra_task_config - # import logger logger = logging.getLogger(__name__) # PLACEHOLDER: Extension template (do not remove this comment) -@hydra_task_config(args_cli.task, args_cli.agent) def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict): """Train with RL-Games agent.""" # override configurations with non-hydra CLI arguments @@ -256,6 +279,6 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen if __name__ == "__main__": # run the main function - main() + main(env_cfg, agent_cfg) # close sim app simulation_app.close() diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py index 0cce12d7eba..3975e32da05 100644 --- a/scripts/reinforcement_learning/rsl_rl/train.py +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -47,9 +47,36 @@ # clear out sys.argv for Hydra sys.argv = [sys.argv[0]] + hydra_args -# launch omniverse app -app_launcher = AppLauncher(args_cli) -simulation_app = app_launcher.app +# Register tasks and resolve config with Hydra (presets applied) +import isaaclab_tasks # noqa: F401 +from isaaclab.sim.backend_detection import physics_backend_requires_kit +from isaaclab_tasks.utils.hydra import finalize_task_config, resolve_task_config + +env_cfg, agent_cfg, hydra_cfg = resolve_task_config(args_cli.task, args_cli.agent) + +# Only launch Isaac Sim Kit when the task's physics backend requires it +_requires_kit = physics_backend_requires_kit(env_cfg) +if _requires_kit: + app_launcher = AppLauncher(args_cli) + simulation_app = app_launcher.app +else: + class _DummySimulationApp: + def close(self): + pass + + class _AppLauncherStub: + """Minimal stand-in for AppLauncher when Kit is not launched (e.g. Newton backend).""" + + def __init__(self): + self.app = _DummySimulationApp() + self.local_rank = 0 + self.global_rank = 0 + + app_launcher = _AppLauncherStub() + simulation_app = app_launcher.app + +# Now that Kit is running (if needed), resolve callable strings in the configs +env_cfg, agent_cfg = finalize_task_config(env_cfg, agent_cfg, hydra_cfg) """Check for minimum supported RSL-RL version.""" @@ -85,7 +112,6 @@ from rsl_rl.runners import DistillationRunner, OnPolicyRunner from isaaclab.envs import ( - DirectMARLEnv, DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg, @@ -94,11 +120,9 @@ from isaaclab.utils.dict import print_dict from isaaclab.utils.io import dump_yaml -from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper +from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper -import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils import get_checkpoint_path -from isaaclab_tasks.utils.hydra import hydra_task_config # import logger logger = logging.getLogger(__name__) @@ -111,8 +135,7 @@ torch.backends.cudnn.benchmark = False -@hydra_task_config(args_cli.task, args_cli.agent) -def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg): +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg): """Train with RSL-RL agent.""" # override configurations with non-hydra CLI arguments agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) @@ -170,7 +193,7 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) # convert to single-agent instance if required by the RL algorithm - if isinstance(env.unwrapped, DirectMARLEnv): + if type(env.unwrapped).__name__ == "DirectMARLEnv": env = multi_agent_to_single_agent(env) # save resume path before creating a new log_dir @@ -224,6 +247,6 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen if __name__ == "__main__": # run the main function - main() + main(env_cfg, agent_cfg) # close sim app simulation_app.close() diff --git a/scripts/tutorials/01_assets/add_new_robot.py b/scripts/tutorials/01_assets/add_new_robot.py index 9610f8eb05b..c31e81db309 100644 --- a/scripts/tutorials/01_assets/add_new_robot.py +++ b/scripts/tutorials/01_assets/add_new_robot.py @@ -126,12 +126,12 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(scene["Jetbot"].data.default_joint_pos).clone(), wp.to_torch(scene["Jetbot"].data.default_joint_vel).clone(), ) - scene["Jetbot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["Jetbot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) joint_pos, joint_vel = ( wp.to_torch(scene["Dofbot"].data.default_joint_pos).clone(), wp.to_torch(scene["Dofbot"].data.default_joint_vel).clone(), ) - scene["Dofbot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["Dofbot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting Jetbot and Dofbot state...") diff --git a/scripts/tutorials/01_assets/run_articulation.py b/scripts/tutorials/01_assets/run_articulation.py index 509bcaa0d76..db9ce6f1a47 100644 --- a/scripts/tutorials/01_assets/run_articulation.py +++ b/scripts/tutorials/01_assets/run_articulation.py @@ -98,7 +98,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula # set joint positions with some noise joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() joint_pos += torch.rand_like(joint_pos) * 0.1 - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers robot.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/tutorials/01_assets/run_surface_gripper.py b/scripts/tutorials/01_assets/run_surface_gripper.py index 97afe8dd1a2..f28c56211f9 100644 --- a/scripts/tutorials/01_assets/run_surface_gripper.py +++ b/scripts/tutorials/01_assets/run_surface_gripper.py @@ -118,7 +118,7 @@ def run_simulator( wp.to_torch(robot.data.default_joint_vel).clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers robot.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/tutorials/02_scene/create_scene.py b/scripts/tutorials/02_scene/create_scene.py index 8d5a71cee0d..fe234e59227 100644 --- a/scripts/tutorials/02_scene/create_scene.py +++ b/scripts/tutorials/02_scene/create_scene.py @@ -89,7 +89,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # set joint positions with some noise joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() joint_pos += torch.rand_like(joint_pos) * 0.1 - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/tutorials/04_sensors/add_sensors_on_robot.py b/scripts/tutorials/04_sensors/add_sensors_on_robot.py index f3d19736aa9..bd94e35a769 100644 --- a/scripts/tutorials/04_sensors/add_sensors_on_robot.py +++ b/scripts/tutorials/04_sensors/add_sensors_on_robot.py @@ -122,7 +122,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): wp.to_torch(scene["robot"].data.default_joint_vel).clone(), ) joint_pos += torch.rand_like(joint_pos) * 0.1 - scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting robot state...") diff --git a/scripts/tutorials/05_controllers/run_diff_ik.py b/scripts/tutorials/05_controllers/run_diff_ik.py index 58071b0e6af..b7a87402815 100644 --- a/scripts/tutorials/05_controllers/run_diff_ik.py +++ b/scripts/tutorials/05_controllers/run_diff_ik.py @@ -149,7 +149,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # reset joint state joint_pos = robot.data.default_joint_pos.clone() joint_vel = robot.data.default_joint_vel.clone() - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) robot.reset() # reset actions ik_commands[:] = ee_goals[current_goal_idx] diff --git a/scripts/tutorials/05_controllers/run_osc.py b/scripts/tutorials/05_controllers/run_osc.py index f0ad2221cd1..862f8779ff4 100644 --- a/scripts/tutorials/05_controllers/run_osc.py +++ b/scripts/tutorials/05_controllers/run_osc.py @@ -217,7 +217,7 @@ def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): # reset joint state to default default_joint_pos = robot.data.default_joint_pos.clone() default_joint_vel = robot.data.default_joint_vel.clone() - robot.write_joint_state_to_sim(default_joint_pos, default_joint_vel) + robot.write_joint_state_to_sim(position=default_joint_pos, velocity=default_joint_vel) robot.set_joint_effort_target(zero_joint_efforts) # Set zero torques in the initial step robot.write_data_to_sim() robot.reset() diff --git a/source/isaaclab/isaaclab/actuators/__init__.py b/source/isaaclab/isaaclab/actuators/__init__.py index db7d36b00a5..f96315ef7e0 100644 --- a/source/isaaclab/isaaclab/actuators/__init__.py +++ b/source/isaaclab/isaaclab/actuators/__init__.py @@ -22,15 +22,22 @@ and called by the :class:`isaaclab.assets.Articulation` class. """ -from .actuator_base import ActuatorBase -from .actuator_base_cfg import ActuatorBaseCfg -from .actuator_net import ActuatorNetLSTM, ActuatorNetMLP -from .actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg -from .actuator_pd import DCMotor, DelayedPDActuator, IdealPDActuator, ImplicitActuator, RemotizedPDActuator -from .actuator_pd_cfg import ( - DCMotorCfg, - DelayedPDActuatorCfg, - IdealPDActuatorCfg, - ImplicitActuatorCfg, - RemotizedPDActuatorCfg, +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "actuator_base": ["ActuatorBase"], + "actuator_base_cfg": ["ActuatorBaseCfg"], + "actuator_net": ["ActuatorNetLSTM", "ActuatorNetMLP"], + "actuator_net_cfg": ["ActuatorNetLSTMCfg", "ActuatorNetMLPCfg"], + "actuator_pd": ["DCMotor", "DelayedPDActuator", "IdealPDActuator", "ImplicitActuator", "RemotizedPDActuator"], + "actuator_pd_cfg": [ + "DCMotorCfg", + "DelayedPDActuatorCfg", + "IdealPDActuatorCfg", + "ImplicitActuatorCfg", + "RemotizedPDActuatorCfg", + ], + }, ) diff --git a/source/isaaclab/isaaclab/assets/__init__.py b/source/isaaclab/isaaclab/assets/__init__.py index 3eacce88028..ede8a06ef76 100644 --- a/source/isaaclab/isaaclab/assets/__init__.py +++ b/source/isaaclab/isaaclab/assets/__init__.py @@ -38,28 +38,33 @@ the corresponding actuator torques. """ -from .articulation import BaseArticulation, BaseArticulationData, Articulation, ArticulationCfg, ArticulationData -from .asset_base import AssetBase -from .asset_base_cfg import AssetBaseCfg -from .rigid_object import BaseRigidObject, BaseRigidObjectData, RigidObject, RigidObjectCfg, RigidObjectData -from .rigid_object_collection import BaseRigidObjectCollection, BaseRigidObjectCollectionData, RigidObjectCollection, RigidObjectCollectionCfg, RigidObjectCollectionData +import lazy_loader as lazy -__all__ = [ - "AssetBase", - "AssetBaseCfg", - "BaseArticulation", - "BaseArticulationData", - "Articulation", - "ArticulationCfg", - "ArticulationData", - "BaseRigidObject", - "BaseRigidObjectData", - "RigidObject", - "RigidObjectCfg", - "RigidObjectData", - "BaseRigidObjectCollection", - "BaseRigidObjectCollectionData", - "RigidObjectCollection", - "RigidObjectCollectionCfg", - "RigidObjectCollectionData", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "articulation": [ + "BaseArticulation", + "BaseArticulationData", + "Articulation", + "ArticulationCfg", + "ArticulationData", + ], + "asset_base": ["AssetBase"], + "asset_base_cfg": ["AssetBaseCfg"], + "rigid_object": [ + "BaseRigidObject", + "BaseRigidObjectData", + "RigidObject", + "RigidObjectCfg", + "RigidObjectData", + ], + "rigid_object_collection": [ + "BaseRigidObjectCollection", + "BaseRigidObjectCollectionData", + "RigidObjectCollection", + "RigidObjectCollectionCfg", + "RigidObjectCollectionData", + ], + }, +) diff --git a/source/isaaclab/isaaclab/assets/articulation/__init__.py b/source/isaaclab/isaaclab/assets/articulation/__init__.py index 792b7a5454a..8d4158b516b 100644 --- a/source/isaaclab/isaaclab/assets/articulation/__init__.py +++ b/source/isaaclab/isaaclab/assets/articulation/__init__.py @@ -5,16 +5,15 @@ """Sub-module for rigid articulated assets.""" -from .base_articulation import BaseArticulation -from .base_articulation_data import BaseArticulationData -from .articulation import Articulation -from .articulation_cfg import ArticulationCfg -from .articulation_data import ArticulationData +import lazy_loader as lazy -__all__ = [ - "BaseArticulation", - "BaseArticulationData", - "Articulation", - "ArticulationCfg", - "ArticulationData", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "base_articulation": ["BaseArticulation"], + "base_articulation_data": ["BaseArticulationData"], + "articulation": ["Articulation"], + "articulation_cfg": ["ArticulationCfg"], + "articulation_data": ["ArticulationData"], + }, +) diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py index f8c558302d3..c708a0c5a5e 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py @@ -6,12 +6,10 @@ from __future__ import annotations from dataclasses import MISSING - from isaaclab.actuators import ActuatorBaseCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from ..asset_base_cfg import AssetBaseCfg -from .articulation import Articulation @configclass @@ -38,7 +36,7 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg): # Initialize configurations. ## - class_type: type = Articulation + class_type: type | DeferredClass = DeferredClass("isaaclab.assets.articulation.articulation:Articulation") articulation_root_prim_path: str | None = None """Path to the articulation root prim under the :attr:`prim_path`. Defaults to None, in which case the class diff --git a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py index fbe240db505..bdc547935a0 100644 --- a/source/isaaclab/isaaclab/assets/articulation/base_articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/base_articulation.py @@ -2406,6 +2406,7 @@ def write_root_link_velocity_to_sim( @abstractmethod def write_joint_state_to_sim( self, + *, position: torch.Tensor | wp.array, velocity: torch.Tensor | wp.array, joint_ids: Sequence[int] | slice | None = None, diff --git a/source/isaaclab/isaaclab/assets/asset_base.py b/source/isaaclab/isaaclab/assets/asset_base.py index e6bd12220b3..099b693e2b2 100644 --- a/source/isaaclab/isaaclab/assets/asset_base.py +++ b/source/isaaclab/isaaclab/assets/asset_base.py @@ -312,18 +312,7 @@ def _initialize_callback(self, event): if not self._is_initialized: self._backend = PhysicsManager.get_backend() self._device = PhysicsManager.get_device() - try: - self._initialize_impl() - except Exception as e: - store_fn = getattr( - SimulationContext.instance().physics_manager, - "store_callback_exception", - None, - ) - if callable(store_fn): - store_fn(e) - else: - raise + self._initialize_impl() self._is_initialized = True def _invalidate_initialize_callback(self, event): diff --git a/source/isaaclab/isaaclab/assets/asset_base_cfg.py b/source/isaaclab/isaaclab/assets/asset_base_cfg.py index cb24b8f632d..8db5545bbbc 100644 --- a/source/isaaclab/isaaclab/assets/asset_base_cfg.py +++ b/source/isaaclab/isaaclab/assets/asset_base_cfg.py @@ -6,12 +6,13 @@ from __future__ import annotations from dataclasses import MISSING -from typing import Literal +from typing import TYPE_CHECKING, Literal from isaaclab.sim import SpawnerCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass -from .asset_base import AssetBase +if TYPE_CHECKING: + from .asset_base import AssetBase @configclass @@ -41,7 +42,7 @@ class InitialStateCfg: Defaults to (0.0, 0.0, 0.0, 1.0). """ - class_type: type[AssetBase] = None + class_type: type | DeferredClass | None = None """The associated asset class. Defaults to None, which means that the asset will be spawned but cannot be interacted with via the asset class. diff --git a/source/isaaclab/isaaclab/assets/rigid_object/__init__.py b/source/isaaclab/isaaclab/assets/rigid_object/__init__.py index 9fdde709c8b..aeafad4c5d1 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object/__init__.py +++ b/source/isaaclab/isaaclab/assets/rigid_object/__init__.py @@ -5,16 +5,15 @@ """Sub-module for rigid object assets.""" -from .base_rigid_object import BaseRigidObject -from .base_rigid_object_data import BaseRigidObjectData -from .rigid_object import RigidObject -from .rigid_object_cfg import RigidObjectCfg -from .rigid_object_data import RigidObjectData +import lazy_loader as lazy -__all__ = [ - "BaseRigidObject", - "BaseRigidObjectData", - "RigidObject", - "RigidObjectCfg", - "RigidObjectData", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "base_rigid_object": ["BaseRigidObject"], + "base_rigid_object_data": ["BaseRigidObjectData"], + "rigid_object": ["RigidObject"], + "rigid_object_cfg": ["RigidObjectCfg"], + "rigid_object_data": ["RigidObjectData"], + }, +) diff --git a/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_cfg.py b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_cfg.py index 8340aa45f76..f78c5c9557e 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_cfg.py +++ b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object_cfg.py @@ -3,10 +3,9 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from ..asset_base_cfg import AssetBaseCfg -from .rigid_object import RigidObject @configclass @@ -26,7 +25,7 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg): # Initialize configurations. ## - class_type: type = RigidObject + class_type: type | DeferredClass = DeferredClass("isaaclab.assets.rigid_object.rigid_object:RigidObject") init_state: InitialStateCfg = InitialStateCfg() """Initial state of the rigid object. Defaults to identity pose with zero velocity.""" diff --git a/source/isaaclab/isaaclab/assets/rigid_object_collection/__init__.py b/source/isaaclab/isaaclab/assets/rigid_object_collection/__init__.py index f0fd26363a0..aa73ac5213b 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object_collection/__init__.py +++ b/source/isaaclab/isaaclab/assets/rigid_object_collection/__init__.py @@ -5,16 +5,15 @@ """Sub-module for rigid object collection.""" -from .base_rigid_object_collection import BaseRigidObjectCollection -from .base_rigid_object_collection_data import BaseRigidObjectCollectionData -from .rigid_object_collection import RigidObjectCollection -from .rigid_object_collection_cfg import RigidObjectCollectionCfg -from .rigid_object_collection_data import RigidObjectCollectionData +import lazy_loader as lazy -__all__ = [ - "BaseRigidObjectCollection", - "BaseRigidObjectCollectionData", - "RigidObjectCollection", - "RigidObjectCollectionCfg", - "RigidObjectCollectionData", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "base_rigid_object_collection": ["BaseRigidObjectCollection"], + "base_rigid_object_collection_data": ["BaseRigidObjectCollectionData"], + "rigid_object_collection": ["RigidObjectCollection"], + "rigid_object_collection_cfg": ["RigidObjectCollectionCfg"], + "rigid_object_collection_data": ["RigidObjectCollectionData"], + }, +) diff --git a/source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection_cfg.py b/source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection_cfg.py index c99b20044dd..9b3f7d04803 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection_cfg.py +++ b/source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection_cfg.py @@ -6,16 +6,14 @@ from dataclasses import MISSING from isaaclab.assets.rigid_object import RigidObjectCfg -from isaaclab.utils import configclass - -from .rigid_object_collection import RigidObjectCollection +from isaaclab.utils import DeferredClass, configclass @configclass class RigidObjectCollectionCfg: """Configuration parameters for a rigid object collection.""" - class_type: type = RigidObjectCollection + class_type: type | DeferredClass = DeferredClass("isaaclab.assets.rigid_object_collection.rigid_object_collection:RigidObjectCollection") """The associated asset class. The class should inherit from :class:`isaaclab.assets.asset_base.AssetBase`. diff --git a/source/isaaclab/isaaclab/cli/commands/install.py b/source/isaaclab/isaaclab/cli/commands/install.py index 23265b00b55..e8582007d27 100644 --- a/source/isaaclab/isaaclab/cli/commands/install.py +++ b/source/isaaclab/isaaclab/cli/commands/install.py @@ -145,6 +145,8 @@ def _install_isaaclab_extensions() -> None: "install", "--editable", str(item), + "--extra-index-url", + "https://pypi.nvidia.com", ] ) diff --git a/source/isaaclab/isaaclab/cloner/__init__.py b/source/isaaclab/isaaclab/cloner/__init__.py index 8edaea840c8..b6af2229430 100644 --- a/source/isaaclab/isaaclab/cloner/__init__.py +++ b/source/isaaclab/isaaclab/cloner/__init__.py @@ -1,8 +1,21 @@ -# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sam/IsaacLab/blob/main/CONTRIBUTORS.md). # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause -from .cloner_cfg import TemplateCloneCfg -from .cloner_strategies import * -from .cloner_utils import * +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "cloner_cfg": ["TemplateCloneCfg"], + "cloner_strategies": ["random", "sequential"], + "cloner_utils": [ + "clone_from_template", + "make_clone_plan", + "usd_replicate", + "filter_collisions", + "grid_transforms", + ], + }, +) diff --git a/source/isaaclab/isaaclab/devices/__init__.py b/source/isaaclab/isaaclab/devices/__init__.py index b2605d39ca1..a50a847fa68 100644 --- a/source/isaaclab/isaaclab/devices/__init__.py +++ b/source/isaaclab/isaaclab/devices/__init__.py @@ -20,11 +20,18 @@ the peripheral device. """ -from .device_base import DeviceBase, DeviceCfg, DevicesCfg -from .gamepad import Se2Gamepad, Se2GamepadCfg, Se3Gamepad, Se3GamepadCfg -from .haply import HaplyDevice, HaplyDeviceCfg -from .keyboard import Se2Keyboard, Se2KeyboardCfg, Se3Keyboard, Se3KeyboardCfg -from .openxr import ManusVive, ManusViveCfg, OpenXRDevice, OpenXRDeviceCfg -from .retargeter_base import RetargeterBase, RetargeterCfg -from .spacemouse import Se2SpaceMouse, Se2SpaceMouseCfg, Se3SpaceMouse, Se3SpaceMouseCfg -from .teleop_device_factory import create_teleop_device +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "device_base": ["DeviceBase", "DeviceCfg", "DevicesCfg"], + "gamepad": ["Se2Gamepad", "Se2GamepadCfg", "Se3Gamepad", "Se3GamepadCfg"], + "haply": ["HaplyDevice", "HaplyDeviceCfg"], + "keyboard": ["Se2Keyboard", "Se2KeyboardCfg", "Se3Keyboard", "Se3KeyboardCfg"], + "openxr": ["ManusVive", "ManusViveCfg", "OpenXRDevice", "OpenXRDeviceCfg"], + "retargeter_base": ["RetargeterBase", "RetargeterCfg"], + "spacemouse": ["Se2SpaceMouse", "Se2SpaceMouseCfg", "Se3SpaceMouse", "Se3SpaceMouseCfg"], + "teleop_device_factory": ["create_teleop_device"], + }, +) diff --git a/source/isaaclab/isaaclab/devices/openxr/__init__.py b/source/isaaclab/isaaclab/devices/openxr/__init__.py index 9d44a65e4ab..67fa388e55d 100644 --- a/source/isaaclab/isaaclab/devices/openxr/__init__.py +++ b/source/isaaclab/isaaclab/devices/openxr/__init__.py @@ -15,17 +15,14 @@ :class:`DeprecationWarning` at instantiation time. """ -try: +_NAMES = ["ManusVive", "ManusViveCfg", "OpenXRDevice", "OpenXRDeviceCfg", "XrAnchorRotationMode", "XrCfg", "remove_camera_configs"] +__all__ = _NAMES - from isaaclab_teleop.deprecated.openxr import ( # noqa: F401 - ManusVive, - ManusViveCfg, - OpenXRDevice, - OpenXRDeviceCfg, - XrAnchorRotationMode, - XrCfg, - remove_camera_configs, - ) -except ImportError: - print("isaaclab_teleop is not installed. OpenXR teleoperation features will not be available.") - pass + +def __getattr__(name): + if name in _NAMES: + import importlib + + mod = importlib.import_module("isaaclab_teleop.deprecated.openxr") + return getattr(mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/source/isaaclab/isaaclab/envs/__init__.py b/source/isaaclab/isaaclab/envs/__init__.py index 543ff2ad4ba..133bda70334 100644 --- a/source/isaaclab/isaaclab/envs/__init__.py +++ b/source/isaaclab/isaaclab/envs/__init__.py @@ -42,16 +42,30 @@ .. _`Task Design Workflows`: https://docs.isaacsim.omniverse.nvidia.com/latest/introduction/workflows.html """ -from . import mdp, ui -from .common import VecEnvObs, VecEnvStepReturn, ViewerCfg -from .direct_marl_env import DirectMARLEnv -from .direct_marl_env_cfg import DirectMARLEnvCfg -from .direct_rl_env import DirectRLEnv -from .direct_rl_env_cfg import DirectRLEnvCfg -from .manager_based_env import ManagerBasedEnv -from .manager_based_env_cfg import ManagerBasedEnvCfg -from .manager_based_rl_env import ManagerBasedRLEnv -from .manager_based_rl_env_cfg import ManagerBasedRLEnvCfg -from .manager_based_rl_mimic_env import ManagerBasedRLMimicEnv -from .mimic_env_cfg import * -from .utils.marl import multi_agent_to_single_agent, multi_agent_with_one_agent +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["mdp", "ui"], + submod_attrs={ + "common": ["VecEnvObs", "VecEnvStepReturn", "ViewerCfg"], + "direct_marl_env": ["DirectMARLEnv"], + "direct_marl_env_cfg": ["DirectMARLEnvCfg"], + "direct_rl_env": ["DirectRLEnv"], + "direct_rl_env_cfg": ["DirectRLEnvCfg"], + "manager_based_env": ["ManagerBasedEnv"], + "manager_based_env_cfg": ["ManagerBasedEnvCfg"], + "manager_based_rl_env": ["ManagerBasedRLEnv"], + "manager_based_rl_env_cfg": ["ManagerBasedRLEnvCfg"], + "manager_based_rl_mimic_env": ["ManagerBasedRLMimicEnv"], + "mimic_env_cfg": [ + "DataGenConfig", + "SubTaskConfig", + "SubTaskConstraintType", + "SubTaskConstraintCoordinationScheme", + "SubTaskConstraintConfig", + "MimicEnvCfg", + ], + "utils.marl": ["multi_agent_to_single_agent", "multi_agent_with_one_agent"], + }, +) diff --git a/source/isaaclab/isaaclab/envs/direct_marl_env.py b/source/isaaclab/isaaclab/envs/direct_marl_env.py index 2848e9bd918..bfc283b5aad 100644 --- a/source/isaaclab/isaaclab/envs/direct_marl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_marl_env.py @@ -5,7 +5,6 @@ from __future__ import annotations -import builtins import inspect import logging import math @@ -19,9 +18,6 @@ import numpy as np import torch -import omni.kit.app -import omni.physx - from isaaclab.managers import EventManager from isaaclab.scene import InteractiveScene from isaaclab.sim import SimulationContext @@ -147,20 +143,13 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar self.event_manager.apply(mode="prestartup") # play the simulator to activate physics handles - # note: this activates the physics simulation view that exposes TensorAPIs - # note: when started in extension mode, first call sim.reset_async() and then initialize the managers - if builtins.ISAAC_LAUNCHED_FROM_TERMINAL is False: - print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...") - with Timer("[INFO]: Time taken for simulation start", "simulation_start"): - # since the reset can trigger callbacks which use the stage, - # we need to set the stage context here - with use_stage(self.sim.stage): - self.sim.reset() - # update scene to pre populate data buffers for assets and sensors. - # this is needed for the observation manager to get valid tensors for initialization. - # this shouldn't cause an issue since later on, users do a reset over all the environments - # so the lazy buffers would be reset. - self.scene.update(dt=self.physics_dt) + # note: this activates the physics simulation view that exposes TensorAPIs. + # sim.reset() must always be called so PHYSICS_READY fires and assets get initialized. + print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...") + with Timer("[INFO]: Time taken for simulation start", "simulation_start"): + with use_stage(self.sim.stage): + self.sim.reset() + self.scene.update(dt=self.physics_dt) # check if debug visualization is has been implemented by the environment source_code = inspect.getsource(self._set_debug_vis_impl) @@ -171,6 +160,12 @@ def __init__(self, cfg: DirectMARLEnvCfg, render_mode: str | None = None, **kwar # we need to do this here after all the managers are initialized # this is because they dictate the sensors and commands right now if self.sim.has_gui and self.cfg.ui_window_class_type is not None: + if hasattr(self.cfg.ui_window_class_type, "resolve"): + self.cfg.ui_window_class_type = self.cfg.ui_window_class_type.resolve() + elif isinstance(self.cfg.ui_window_class_type, str): + from isaaclab.utils.string import string_to_callable + + self.cfg.ui_window_class_type = string_to_callable(self.cfg.ui_window_class_type) self._window = self.cfg.ui_window_class_type(self, window_name="IsaacLab") else: # if no window, then we don't need to store the window @@ -574,6 +569,8 @@ def set_debug_vis(self, debug_vis: bool) -> bool: if debug_vis: # create a subscriber for the post update event if it doesn't exist if self._debug_vis_handle is None: + import omni.kit.app + app_interface = omni.kit.app.get_app_interface() self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) diff --git a/source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py b/source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py index 684f786ba34..fa8d85e963c 100644 --- a/source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py +++ b/source/isaaclab/isaaclab/envs/direct_marl_env_cfg.py @@ -8,16 +8,15 @@ from dataclasses import MISSING from typing import TYPE_CHECKING -if TYPE_CHECKING: - from isaaclab.devices.openxr import XrCfg - from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from isaaclab.utils.noise import NoiseModelCfg from .common import AgentID, SpaceType, ViewerCfg -from .ui import BaseEnvWindow + +if TYPE_CHECKING: + from isaaclab.devices.openxr import XrCfg @configclass @@ -35,7 +34,7 @@ class DirectMARLEnvCfg: """Physics simulation configuration. Default is SimulationCfg().""" # ui settings - ui_window_class_type: type | None = BaseEnvWindow + ui_window_class_type: type | DeferredClass | None = DeferredClass("isaaclab.envs.ui:BaseEnvWindow") """The class type of the UI window. Default is None. If None, then no UI window is created. diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env.py b/source/isaaclab/isaaclab/envs/direct_rl_env.py index 895923124a0..65b0b294040 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env.py @@ -5,7 +5,6 @@ from __future__ import annotations -import builtins import inspect import logging import math @@ -20,9 +19,6 @@ import numpy as np import torch -import omni.kit.app -import omni.physx - from isaaclab.managers import EventManager from isaaclab.scene import InteractiveScene from isaaclab.sim import SimulationContext @@ -153,20 +149,20 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs self.event_manager.apply(mode="prestartup") # play the simulator to activate physics handles - # note: this activates the physics simulation view that exposes TensorAPIs - # note: when started in extension mode, first call sim.reset_async() and then initialize the managers - if builtins.ISAAC_LAUNCHED_FROM_TERMINAL is False: - print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...") - with Timer("[INFO]: Time taken for simulation start", "simulation_start"): - # since the reset can trigger callbacks which use the stage, - # we need to set the stage context here - with use_stage(self.sim.stage): - self.sim.reset() - # update scene to pre populate data buffers for assets and sensors. - # this is needed for the observation manager to get valid tensors for initialization. - # this shouldn't cause an issue since later on, users do a reset over all the environments - # so the lazy buffers would be reset. - self.scene.update(dt=self.physics_dt) + # note: this activates the physics simulation view that exposes TensorAPIs. + # sim.reset() must always be called so PHYSICS_READY fires and assets (e.g. Newton + # articulations) get _root_view initialized. This is required for both PhysX and Newton. + print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...") + with Timer("[INFO]: Time taken for simulation start", "simulation_start"): + # since the reset can trigger callbacks which use the stage, + # we need to set the stage context here + with use_stage(self.sim.stage): + self.sim.reset() + # update scene to pre populate data buffers for assets and sensors. + # this is needed for the observation manager to get valid tensors for initialization. + # this shouldn't cause an issue since later on, users do a reset over all the environments + # so the lazy buffers would be reset. + self.scene.update(dt=self.physics_dt) # check if debug visualization is has been implemented by the environment source_code = inspect.getsource(self._set_debug_vis_impl) @@ -177,6 +173,12 @@ def __init__(self, cfg: DirectRLEnvCfg, render_mode: str | None = None, **kwargs # we need to do this here after all the managers are initialized # this is because they dictate the sensors and commands right now if self.sim.has_gui and self.cfg.ui_window_class_type is not None: + if hasattr(self.cfg.ui_window_class_type, "resolve"): + self.cfg.ui_window_class_type = self.cfg.ui_window_class_type.resolve() + elif isinstance(self.cfg.ui_window_class_type, str): + from isaaclab.utils.string import string_to_callable + + self.cfg.ui_window_class_type = string_to_callable(self.cfg.ui_window_class_type) self._window = self.cfg.ui_window_class_type(self, window_name="IsaacLab") else: # if no window, then we don't need to store the window @@ -543,6 +545,8 @@ def set_debug_vis(self, debug_vis: bool) -> bool: if debug_vis: # create a subscriber for the post update event if it doesn't exist if self._debug_vis_handle is None: + import omni.kit.app + app_interface = omni.kit.app.get_app_interface() self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) diff --git a/source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py b/source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py index 594e066c38b..bdc1d828aa2 100644 --- a/source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py +++ b/source/isaaclab/isaaclab/envs/direct_rl_env_cfg.py @@ -8,15 +8,15 @@ from dataclasses import MISSING from typing import TYPE_CHECKING -if TYPE_CHECKING: - from isaaclab.devices.openxr import XrCfg from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from isaaclab.utils.noise import NoiseModelCfg from .common import SpaceType, ViewerCfg -from .ui import BaseEnvWindow + +if TYPE_CHECKING: + from isaaclab.devices.openxr import XrCfg @configclass @@ -34,7 +34,7 @@ class DirectRLEnvCfg: """Physics simulation configuration. Default is SimulationCfg().""" # ui settings - ui_window_class_type: type | None = BaseEnvWindow + ui_window_class_type: type | DeferredClass | None = DeferredClass("isaaclab.envs.ui:BaseEnvWindow") """The class type of the UI window. Default is None. If None, then no UI window is created. diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 491ae9549d0..5c8aad9f1ec 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -5,7 +5,6 @@ from __future__ import annotations -import builtins import logging import warnings from collections.abc import Sequence @@ -20,12 +19,13 @@ from isaaclab.ui.widgets import ManagerLiveVisualizer from isaaclab.utils.seed import configure_seed from isaaclab.utils.timer import Timer +from source.isaaclab.isaaclab.utils.version import has_kit from .common import VecEnvObs from .manager_based_env_cfg import ManagerBasedEnvCfg -from .ui import ViewportCameraController from .utils.io_descriptors import export_articulations_data, export_scene_data + # import logger logger = logging.getLogger(__name__) @@ -101,10 +101,6 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # since it gets confused with Isaac Sim's SimulationContext class self.sim: SimulationContext = SimulationContext(self.cfg.sim) else: - # simulation context should only be created before the environment - # when in extension mode - if not builtins.ISAAC_LAUNCHED_FROM_TERMINAL: - raise RuntimeError("Simulation context already exists. Cannot create a new one.") self.sim: SimulationContext = SimulationContext.instance() # make sure torch is running on the correct device @@ -133,7 +129,7 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # allocate dictionary to store metrics self.extras = {} - # generate scene + with Timer("[INFO]: Time taken for scene creation", "scene_creation"): # set the stage context for scene creation steps which use the stage with use_stage(self.sim.stage): @@ -144,7 +140,8 @@ def __init__(self, cfg: ManagerBasedEnvCfg): # viewport is not available in other rendering modes so the function will throw a warning # FIXME: This needs to be fixed in the future when we unify the UI functionalities even for # non-rendering modes. - if self.sim.has_gui: + if has_kit() and self.sim.has_gui: + from .ui import ViewportCameraController self.viewport_camera_controller = ViewportCameraController(self, self.cfg.viewer) else: self.viewport_camera_controller = None @@ -158,28 +155,27 @@ def __init__(self, cfg: ManagerBasedEnvCfg): if "prestartup" in self.event_manager.available_modes: self.event_manager.apply(mode="prestartup") - # play the simulator to activate physics handles - # note: this activates the physics simulation view that exposes TensorAPIs - # note: when started in extension mode, first call sim.reset_async() and then initialize the managers - if builtins.ISAAC_LAUNCHED_FROM_TERMINAL is False: - print("[INFO]: Starting the simulation. This may take a few seconds. Please wait...") - with Timer("[INFO]: Time taken for simulation start", "simulation_start"): - # since the reset can trigger callbacks which use the stage, - # we need to set the stage context here - with use_stage(self.sim.stage): - self.sim.reset() - # update scene to pre populate data buffers for assets and sensors. - # this is needed for the observation manager to get valid tensors for initialization. - # this shouldn't cause an issue since later on, users do a reset over all the environments - # so the lazy buffers would be reset. - self.scene.update(dt=self.physics_dt) - # add timeline event to load managers - self.load_managers() + # play the simulator to activate physics handles and load managers + with Timer("[INFO]: Time taken for simulation start", "simulation_start"): + # since the reset can trigger callbacks which use the stage, + # we need to set the stage context here + with use_stage(self.sim.stage): + self.sim.reset() + # update scene to pre populate data buffers for assets and sensors. + # this is needed for the observation manager to get valid tensors for initialization. + self.scene.update(dt=self.physics_dt) + self.load_managers() # extend UI elements # we need to do this here after all the managers are initialized # this is because they dictate the sensors and commands right now if self.sim.has_gui and self.cfg.ui_window_class_type is not None: + if hasattr(self.cfg.ui_window_class_type, "resolve"): + self.cfg.ui_window_class_type = self.cfg.ui_window_class_type.resolve() + elif isinstance(self.cfg.ui_window_class_type, str): + from isaaclab.utils.string import string_to_callable + + self.cfg.ui_window_class_type = string_to_callable(self.cfg.ui_window_class_type) # setup live visualizers self.setup_manager_visualizers() self._window = self.cfg.ui_window_class_type(self, window_name="IsaacLab") @@ -325,7 +321,6 @@ def load_managers(self): def setup_manager_visualizers(self): """Creates live visualizers for manager terms.""" - self.manager_visualizers = { "action_manager": ManagerLiveVisualizer(manager=self.action_manager), "observation_manager": ManagerLiveVisualizer(manager=self.observation_manager), diff --git a/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py b/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py index 32a512aa6b4..cb17c5f47c7 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py @@ -14,19 +14,19 @@ from dataclasses import MISSING, field from typing import TYPE_CHECKING -import isaaclab.envs.mdp as mdp from isaaclab.devices.device_base import DevicesCfg - -if TYPE_CHECKING: - from isaaclab.devices.openxr import XrCfg from isaaclab.managers import EventTermCfg as EventTerm from isaaclab.managers import RecorderManagerBaseCfg as DefaultEmptyRecorderManagerCfg from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from .common import ViewerCfg -from .ui import BaseEnvWindow + +if TYPE_CHECKING: + from isaaclab.devices.openxr import XrCfg + + from .ui import BaseEnvWindow @configclass @@ -37,7 +37,7 @@ class DefaultEventManagerCfg: by the scene configuration. """ - reset_scene_to_default = EventTerm(func=mdp.reset_scene_to_default, mode="reset") + reset_scene_to_default = EventTerm(func="isaaclab.envs.mdp.events:reset_scene_to_default", mode="reset") @configclass @@ -52,7 +52,7 @@ class ManagerBasedEnvCfg: """Physics simulation configuration. Default is SimulationCfg().""" # ui settings - ui_window_class_type: type | None = BaseEnvWindow + ui_window_class_type: type | DeferredClass | None = DeferredClass("isaaclab.envs.ui:BaseEnvWindow") """The class type of the UI window. Default is None. If None, then no UI window is created. diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env_cfg.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env_cfg.py index 7fe29ed7003..4ff78305628 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env_cfg.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env_cfg.py @@ -7,10 +7,9 @@ from dataclasses import MISSING -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from .manager_based_env_cfg import ManagerBasedEnvCfg -from .ui import ManagerBasedRLEnvWindow @configclass @@ -18,7 +17,7 @@ class ManagerBasedRLEnvCfg(ManagerBasedEnvCfg): """Configuration for a reinforcement learning environment with the manager-based workflow.""" # ui settings - ui_window_class_type: type | None = ManagerBasedRLEnvWindow + ui_window_class_type: type | DeferredClass | None = DeferredClass("isaaclab.envs.ui:ManagerBasedRLEnvWindow") # general settings is_finite_horizon: bool = False diff --git a/source/isaaclab/isaaclab/envs/mdp/__init__.py b/source/isaaclab/isaaclab/envs/mdp/__init__.py index bdb507f3eb0..a84a9ad5402 100644 --- a/source/isaaclab/isaaclab/envs/mdp/__init__.py +++ b/source/isaaclab/isaaclab/envs/mdp/__init__.py @@ -15,11 +15,149 @@ """ -from .actions import * # noqa: F401, F403 -from .commands import * # noqa: F401, F403 -from .curriculums import * # noqa: F401, F403 -from .events import * # noqa: F401, F403 -from .observations import * # noqa: F401, F403 -from .recorders import * # noqa: F401, F403 -from .rewards import * # noqa: F401, F403 -from .terminations import * # noqa: F401, F403 +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["actions", "commands", "curriculums", "events", "observations", "recorders", "rewards", "terminations"], + submod_attrs={ + "actions": [ + "AbsBinaryJointPositionActionCfg", + "BinaryJointActionCfg", + "BinaryJointPositionActionCfg", + "BinaryJointVelocityActionCfg", + "DifferentialInverseKinematicsActionCfg", + "EMAJointPositionToLimitsActionCfg", + "JointActionCfg", + "JointEffortActionCfg", + "JointPositionActionCfg", + "JointPositionToLimitsActionCfg", + "JointVelocityActionCfg", + "NonHolonomicActionCfg", + "OperationalSpaceControllerActionCfg", + "RelativeJointPositionActionCfg", + "SurfaceGripperBinaryActionCfg", + ], + "commands": [ + "NormalVelocityCommandCfg", + "NullCommandCfg", + "TerrainBasedPose2dCommandCfg", + "UniformPose2dCommandCfg", + "UniformPoseCommandCfg", + "UniformVelocityCommandCfg", + "NullCommand", + "TerrainBasedPose2dCommand", + "UniformPose2dCommand", + "UniformPoseCommand", + "NormalVelocityCommand", + "UniformVelocityCommand", + ], + "curriculums": [ + "modify_env_param", + "modify_reward_weight", + "modify_term_cfg", + ], + "events": [ + "apply_external_force_torque", + "push_by_setting_velocity", + "randomize_actuator_gains", + "randomize_fixed_tendon_parameters", + "randomize_joint_parameters", + "randomize_physics_scene_gravity", + "randomize_rigid_body_collider_offsets", + "randomize_rigid_body_com", + "randomize_rigid_body_mass", + "randomize_rigid_body_material", + "randomize_rigid_body_scale", + "randomize_visual_color", + "randomize_visual_texture_material", + "reset_joints_by_offset", + "reset_joints_by_scale", + "reset_nodal_state_uniform", + "reset_root_state_from_terrain", + "reset_root_state_uniform", + "reset_root_state_with_random_orientation", + "reset_scene_to_default", + ], + "observations": [ + "base_ang_vel", + "base_lin_vel", + "base_pos_z", + "body_incoming_wrench", + "body_pose_w", + "body_projected_gravity_b", + "current_time_s", + "generated_commands", + "height_scan", + "image", + "image_features", + "imu_ang_vel", + "imu_lin_acc", + "imu_orientation", + "imu_projected_gravity", + "joint_effort", + "joint_pos", + "joint_pos_limit_normalized", + "joint_pos_rel", + "joint_vel", + "joint_vel_rel", + "last_action", + "projected_gravity", + "remaining_time_s", + "root_ang_vel_w", + "root_lin_vel_w", + "root_pos_w", + "root_quat_w", + ], + "recorders": [ + "ActionStateRecorderManagerCfg", + "InitialStateRecorder", + "InitialStateRecorderCfg", + "PostStepProcessedActionsRecorder", + "PostStepProcessedActionsRecorderCfg", + "PostStepStatesRecorder", + "PostStepStatesRecorderCfg", + "PreStepActionsRecorder", + "PreStepActionsRecorderCfg", + "PreStepFlatPolicyObservationsRecorder", + "PreStepFlatPolicyObservationsRecorderCfg", + ], + "rewards": [ + "action_l2", + "action_rate_l2", + "ang_vel_xy_l2", + "applied_torque_limits", + "base_height_l2", + "body_lin_acc_l2", + "contact_forces", + "desired_contacts", + "flat_orientation_l2", + "is_alive", + "is_terminated", + "is_terminated_term", + "joint_acc_l2", + "joint_deviation_l1", + "joint_pos_limits", + "joint_torques_l2", + "joint_vel_l1", + "joint_vel_l2", + "joint_vel_limits", + "lin_vel_z_l2", + "track_ang_vel_z_exp", + "track_lin_vel_xy_exp", + "undesired_contacts", + ], + "terminations": [ + "bad_orientation", + "command_resample", + "illegal_contact", + "joint_effort_out_of_limit", + "joint_pos_out_of_limit", + "joint_pos_out_of_manual_limit", + "joint_vel_out_of_limit", + "joint_vel_out_of_manual_limit", + "root_height_below_minimum", + "time_out", + ], + }, +) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/__init__.py b/source/isaaclab/isaaclab/envs/mdp/actions/__init__.py index 56b3ae25ff4..df7e1493dd3 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/__init__.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/__init__.py @@ -5,9 +5,18 @@ """Various action terms that can be used in the environment.""" -from .actions_cfg import * -from .binary_joint_actions import * -from .joint_actions import * -from .joint_actions_to_limits import * -from .non_holonomic_actions import * -from .surface_gripper_actions import * +import lazy_loader as lazy + +from .actions_cfg import * # noqa: F401, F403 + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "binary_joint_actions": ["BinaryJointAction", "BinaryJointPositionAction", "BinaryJointVelocityAction", "AbsBinaryJointPositionAction"], + "joint_actions": ["JointAction", "JointPositionAction", "RelativeJointPositionAction", "JointVelocityAction", "JointEffortAction"], + "joint_actions_to_limits": ["JointPositionToLimitsAction", "EMAJointPositionToLimitsAction"], + "non_holonomic_actions": ["NonHolonomicAction"], + "surface_gripper_actions": ["SurfaceGripperBinaryAction"], + "task_space_actions": ["DifferentialInverseKinematicsAction", "OperationalSpaceControllerAction"], + }, +) diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py index bb664445169..582426f4796 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/actions_cfg.py @@ -6,19 +6,14 @@ from __future__ import annotations from dataclasses import MISSING +from typing import TYPE_CHECKING from isaaclab.controllers import DifferentialIKControllerCfg, OperationalSpaceControllerCfg -from isaaclab.managers.action_manager import ActionTerm, ActionTermCfg -from isaaclab.utils import configclass - -from . import ( - binary_joint_actions, - joint_actions, - joint_actions_to_limits, - non_holonomic_actions, - surface_gripper_actions, - task_space_actions, -) +from isaaclab.managers.manager_term_cfg import ActionTermCfg +from isaaclab.utils import DeferredClass, configclass + +if TYPE_CHECKING: + from isaaclab.managers.action_manager import ActionTerm ## # Joint actions. @@ -49,7 +44,7 @@ class JointPositionActionCfg(JointActionCfg): See :class:`JointPositionAction` for more details. """ - class_type: type[ActionTerm] = joint_actions.JointPositionAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.joint_actions:JointPositionAction") use_default_offset: bool = True """Whether to use default joint positions configured in the articulation asset as offset. @@ -67,7 +62,7 @@ class RelativeJointPositionActionCfg(JointActionCfg): See :class:`RelativeJointPositionAction` for more details. """ - class_type: type[ActionTerm] = joint_actions.RelativeJointPositionAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.joint_actions:RelativeJointPositionAction") use_zero_offset: bool = True """Whether to ignore the offset defined in articulation asset. Defaults to True. @@ -83,7 +78,7 @@ class JointVelocityActionCfg(JointActionCfg): See :class:`JointVelocityAction` for more details. """ - class_type: type[ActionTerm] = joint_actions.JointVelocityAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.joint_actions:JointVelocityAction") use_default_offset: bool = True """Whether to use default joint velocities configured in the articulation asset as offset. @@ -100,7 +95,7 @@ class JointEffortActionCfg(JointActionCfg): See :class:`JointEffortAction` for more details. """ - class_type: type[ActionTerm] = joint_actions.JointEffortAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.joint_actions:JointEffortAction") ## @@ -115,7 +110,7 @@ class JointPositionToLimitsActionCfg(ActionTermCfg): See :class:`JointPositionToLimitsAction` for more details. """ - class_type: type[ActionTerm] = joint_actions_to_limits.JointPositionToLimitsAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.joint_actions_to_limits:JointPositionToLimitsAction") joint_names: list[str] = MISSING """List of joint names or regex expressions that the action will be mapped to.""" @@ -144,7 +139,7 @@ class EMAJointPositionToLimitsActionCfg(JointPositionToLimitsActionCfg): See :class:`EMAJointPositionToLimitsAction` for more details. """ - class_type: type[ActionTerm] = joint_actions_to_limits.EMAJointPositionToLimitsAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.joint_actions_to_limits:EMAJointPositionToLimitsAction") alpha: float | dict[str, float] = 1.0 """The weight for the moving average (float or dict of regex expressions). Defaults to 1.0. @@ -180,7 +175,7 @@ class BinaryJointPositionActionCfg(BinaryJointActionCfg): See :class:`BinaryJointPositionAction` for more details. """ - class_type: type[ActionTerm] = binary_joint_actions.BinaryJointPositionAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.binary_joint_actions:BinaryJointPositionAction") @configclass @@ -190,7 +185,7 @@ class BinaryJointVelocityActionCfg(BinaryJointActionCfg): See :class:`BinaryJointVelocityAction` for more details. """ - class_type: type[ActionTerm] = binary_joint_actions.BinaryJointVelocityAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.binary_joint_actions:BinaryJointVelocityAction") @configclass @@ -225,7 +220,7 @@ class AbsBinaryJointPositionActionCfg(ActionTermCfg): positive_threshold: bool = True """Whether to use positive (Open actions > Close actions) threshold. Defaults to True.""" - class_type: type[ActionTerm] = binary_joint_actions.AbsBinaryJointPositionAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.binary_joint_actions:AbsBinaryJointPositionAction") ## @@ -240,7 +235,7 @@ class NonHolonomicActionCfg(ActionTermCfg): See :class:`NonHolonomicAction` for more details. """ - class_type: type[ActionTerm] = non_holonomic_actions.NonHolonomicAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.non_holonomic_actions:NonHolonomicAction") body_name: str = MISSING """Name of the body which has the dummy mechanism connected to.""" @@ -283,7 +278,7 @@ class OffsetCfg: rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) """Quaternion rotation ``(x, y, z, w)`` w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0, 1.0).""" - class_type: type[ActionTerm] = task_space_actions.DifferentialInverseKinematicsAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.task_space_actions:DifferentialInverseKinematicsAction") joint_names: list[str] = MISSING """List of joint names or regex expressions that the action will be mapped to.""" @@ -319,7 +314,7 @@ class OffsetCfg: rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) """Quaternion rotation ``(x, y, z, w)`` w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0, 1.0).""" - class_type: type[ActionTerm] = task_space_actions.OperationalSpaceControllerAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.task_space_actions:OperationalSpaceControllerAction") joint_names: list[str] = MISSING """List of joint names or regex expressions that the action will be mapped to.""" @@ -378,4 +373,4 @@ class SurfaceGripperBinaryActionCfg(ActionTermCfg): close_command: float = 1.0 """The command value to close the gripper. Defaults to 1.0.""" - class_type: type[ActionTerm] = surface_gripper_actions.SurfaceGripperBinaryAction + class_type: type[ActionTerm] | DeferredClass = DeferredClass("isaaclab.envs.mdp.actions.surface_gripper_actions:SurfaceGripperBinaryAction") diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py index d2cd9ab581a..d839c5bd81b 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py @@ -10,12 +10,12 @@ from typing import TYPE_CHECKING import torch -from isaaclab_physx.assets import SurfaceGripper from isaaclab.managers.action_manager import ActionTerm if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv + from isaaclab_physx.assets import SurfaceGripper from . import actions_cfg diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/__init__.py b/source/isaaclab/isaaclab/envs/mdp/commands/__init__.py index bdfce40473b..441ee417d1b 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/__init__.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/__init__.py @@ -5,15 +5,22 @@ """Various command terms that can be used in the environment.""" -from .commands_cfg import ( - NormalVelocityCommandCfg, - NullCommandCfg, - TerrainBasedPose2dCommandCfg, - UniformPose2dCommandCfg, - UniformPoseCommandCfg, - UniformVelocityCommandCfg, +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "commands_cfg": [ + "NormalVelocityCommandCfg", + "NullCommandCfg", + "TerrainBasedPose2dCommandCfg", + "UniformPose2dCommandCfg", + "UniformPoseCommandCfg", + "UniformVelocityCommandCfg", + ], + "null_command": ["NullCommand"], + "pose_2d_command": ["TerrainBasedPose2dCommand", "UniformPose2dCommand"], + "pose_command": ["UniformPoseCommand"], + "velocity_command": ["NormalVelocityCommand", "UniformVelocityCommand"], + }, ) -from .null_command import NullCommand -from .pose_2d_command import TerrainBasedPose2dCommand, UniformPose2dCommand -from .pose_command import UniformPoseCommand -from .velocity_command import NormalVelocityCommand, UniformVelocityCommand diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/commands_cfg.py b/source/isaaclab/isaaclab/envs/mdp/commands/commands_cfg.py index 390980f3454..367f744061b 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/commands_cfg.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/commands_cfg.py @@ -3,25 +3,30 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + import math from dataclasses import MISSING +from typing import TYPE_CHECKING from isaaclab.managers import CommandTermCfg from isaaclab.markers import VisualizationMarkersCfg from isaaclab.markers.config import BLUE_ARROW_X_MARKER_CFG, FRAME_MARKER_CFG, GREEN_ARROW_X_MARKER_CFG from isaaclab.utils import configclass +from isaaclab.utils.string import DeferredClass -from .null_command import NullCommand -from .pose_2d_command import TerrainBasedPose2dCommand, UniformPose2dCommand -from .pose_command import UniformPoseCommand -from .velocity_command import NormalVelocityCommand, UniformVelocityCommand +if TYPE_CHECKING: + from .null_command import NullCommand + from .pose_2d_command import TerrainBasedPose2dCommand, UniformPose2dCommand + from .pose_command import UniformPoseCommand + from .velocity_command import NormalVelocityCommand, UniformVelocityCommand @configclass class NullCommandCfg(CommandTermCfg): """Configuration for the null command generator.""" - class_type: type = NullCommand + class_type: type | DeferredClass = DeferredClass("isaaclab.envs.mdp.commands.null_command:NullCommand") def __post_init__(self): """Post initialization.""" @@ -33,7 +38,7 @@ def __post_init__(self): class UniformVelocityCommandCfg(CommandTermCfg): """Configuration for the uniform velocity command generator.""" - class_type: type = UniformVelocityCommand + class_type: type | DeferredClass = DeferredClass("isaaclab.envs.mdp.commands.velocity_command:UniformVelocityCommand") asset_name: str = MISSING """Name of the asset in the environment for which the commands are generated.""" @@ -100,7 +105,7 @@ class Ranges: class NormalVelocityCommandCfg(UniformVelocityCommandCfg): """Configuration for the normal velocity command generator.""" - class_type: type = NormalVelocityCommand + class_type: type | DeferredClass = DeferredClass("isaaclab.envs.mdp.commands.velocity_command:NormalVelocityCommand") heading_command: bool = False # --> we don't use heading command for normal velocity command. @configclass @@ -133,7 +138,7 @@ class Ranges: class UniformPoseCommandCfg(CommandTermCfg): """Configuration for uniform pose command generator.""" - class_type: type = UniformPoseCommand + class_type: type | DeferredClass = DeferredClass("isaaclab.envs.mdp.commands.pose_command:UniformPoseCommand") asset_name: str = MISSING """Name of the asset in the environment for which the commands are generated.""" @@ -189,7 +194,7 @@ class Ranges: class UniformPose2dCommandCfg(CommandTermCfg): """Configuration for the uniform 2D-pose command generator.""" - class_type: type = UniformPose2dCommand + class_type: type | DeferredClass = DeferredClass("isaaclab.envs.mdp.commands.pose_2d_command:UniformPose2dCommand") asset_name: str = MISSING """Name of the asset in the environment for which the commands are generated.""" @@ -232,7 +237,7 @@ class Ranges: class TerrainBasedPose2dCommandCfg(UniformPose2dCommandCfg): """Configuration for the terrain-based position command generator.""" - class_type = TerrainBasedPose2dCommand + class_type: type | DeferredClass = DeferredClass("isaaclab.envs.mdp.commands.pose_2d_command:TerrainBasedPose2dCommand") @configclass class Ranges: diff --git a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py index eadc89af3af..7c146538e57 100644 --- a/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py +++ b/source/isaaclab/isaaclab/envs/mdp/commands/velocity_command.py @@ -18,6 +18,7 @@ from isaaclab.assets import Articulation from isaaclab.managers import CommandTerm from isaaclab.markers import VisualizationMarkers +from isaaclab.utils.version import has_kit if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -169,6 +170,9 @@ def _update_command(self): def _set_debug_vis_impl(self, debug_vis: bool): # set visibility of markers # note: parent only deals with callbacks. not their visibility + # Skip visualization markers in non-Kit environments (e.g., Newton standalone) + if not has_kit(): + return if debug_vis: # create markers if necessary for the first time if not hasattr(self, "goal_vel_visualizer"): @@ -189,6 +193,9 @@ def _debug_vis_callback(self, event): # note: this is needed in-case the robot is de-initialized. we can't access the data if not self.robot.is_initialized: return + # Skip if visualizers aren't available (e.g., non-Kit environments) + if not hasattr(self, "goal_vel_visualizer"): + return # get marker location # -- base state base_pos_w = wp.to_torch(self.robot.data.root_pos_w).clone() diff --git a/source/isaaclab/isaaclab/envs/mdp/events.py b/source/isaaclab/isaaclab/envs/mdp/events.py index 95e188a5836..c6ecbb0b3e2 100644 --- a/source/isaaclab/isaaclab/envs/mdp/events.py +++ b/source/isaaclab/isaaclab/envs/mdp/events.py @@ -21,24 +21,26 @@ import torch import warp as wp -from isaaclab_physx.assets import DeformableObject - -import carb -import omni.physics.tensors.impl.api as physx -from isaacsim.core.utils.extensions import enable_extension -from pxr import Gf, Sdf, UsdGeom, Vt import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils -from isaaclab.actuators import ImplicitActuator -from isaaclab.assets import Articulation, BaseArticulation, BaseRigidObject, RigidObject from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg -from isaaclab.sim.utils.stage import get_current_stage -from isaaclab.terrains import TerrainImporter from isaaclab.utils.version import compare_versions, get_isaac_sim_version if TYPE_CHECKING: + import carb + import omni.physics.tensors.impl.api as physx + from isaaclab.actuators import ImplicitActuator + from isaaclab.assets import Articulation, BaseArticulation, BaseRigidObject, RigidObject from isaaclab.envs import ManagerBasedEnv + from isaaclab.sim.utils.stage import get_current_stage + from isaaclab.terrains import TerrainImporter + from isaaclab_physx.assets import DeformableObject + + +def _is_instance_by_name(obj: object, *class_names: str) -> bool: + """Check isinstance without importing the class, by inspecting the MRO names.""" + return any(c.__name__ in class_names for c in type(obj).__mro__) # import logger logger = logging.getLogger(__name__) @@ -86,7 +88,7 @@ def randomize_rigid_body_scale( # extract the used quantities (to enable type-hinting) asset: RigidObject = env.scene[asset_cfg.name] - if isinstance(asset, Articulation): + if _is_instance_by_name(asset, "Articulation"): raise ValueError( "Scaling an articulation randomly is not supported, as it affects joint attributes and can cause" " unexpected behavior. To achieve different scales, we recommend generating separate USD files for" @@ -101,7 +103,7 @@ def randomize_rigid_body_scale( env_ids = env_ids.cpu() # acquire stage - stage = get_current_stage() + stage = env.sim.stage # resolve prim paths for spawning and cloning prim_paths = sim_utils.find_matching_prim_paths(asset.cfg.prim_path) @@ -124,6 +126,8 @@ def randomize_rigid_body_scale( relative_child_path = "/" + relative_child_path # use sdf changeblock for faster processing of USD properties + from pxr import Gf, Sdf, UsdGeom, Vt + with Sdf.ChangeBlock(): for i, env_id in enumerate(env_ids): # path to prim to randomize @@ -198,7 +202,7 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] self.asset: RigidObject | Articulation = env.scene[self.asset_cfg.name] - if not isinstance(self.asset, (BaseRigidObject, BaseArticulation)): + if not _is_instance_by_name(self.asset, "BaseRigidObject", "BaseArticulation"): raise ValueError( f"Randomization term 'randomize_rigid_body_material' not supported for asset: '{self.asset_cfg.name}'" f" with type: '{type(self.asset)}'." @@ -207,7 +211,7 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): # obtain number of shapes per body (needed for indexing the material properties correctly) # note: this is a workaround since the Articulation does not provide a direct way to obtain the number of shapes # per body. We use the physics simulation view to obtain the number of shapes per body. - if isinstance(self.asset, Articulation) and self.asset_cfg.body_ids != slice(None): + if _is_instance_by_name(self.asset, "Articulation") and self.asset_cfg.body_ids != slice(None): self.num_shapes_per_body = [] for link_path in self.asset.root_view.link_paths[0]: link_physx_view = self.asset._physics_sim_view.create_rigid_body_view(link_path) # type: ignore @@ -399,7 +403,7 @@ def __call__( inertias = wp.to_torch(self.asset.data.body_inertia).clone() print("inertias device: ", inertias.device) print("inertias shape: ", inertias.shape) - if isinstance(self.asset, BaseArticulation): + if _is_instance_by_name(self.asset, "BaseArticulation"): # inertia has shape: (num_envs, num_bodies, 9) for articulation inertias[env_ids[:, None], body_ids] = ( self.default_inertia[env_ids[:, None], body_ids] * ratios[..., None] @@ -549,8 +553,9 @@ def randomize_physics_scene_gravity( # unbatch the gravity tensor into a list gravity = gravity[0].tolist() - # set the gravity into the physics simulation - physics_sim_view: physx.SimulationView = sim_utils.SimulationContext.instance().physics_sim_view + import carb + + physics_sim_view: physx.SimulationView = env.sim.physics_sim_view physics_sim_view.set_gravity(carb.Float3(*gravity)) @@ -653,7 +658,7 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: stiffness[:, actuator_indices] = self.default_joint_stiffness[env_ids][:, global_indices].clone() randomize(stiffness, stiffness_distribution_params) actuator.stiffness[env_ids] = stiffness - if isinstance(actuator, ImplicitActuator): + if _is_instance_by_name(actuator, "ImplicitActuator"): self.asset.write_joint_stiffness_to_sim( stiffness, joint_ids=actuator.joint_indices, env_ids=env_ids ) @@ -663,7 +668,7 @@ def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor: damping[:, actuator_indices] = self.default_joint_damping[env_ids][:, global_indices].clone() randomize(damping, damping_distribution_params) actuator.damping[env_ids] = damping - if isinstance(actuator, ImplicitActuator): + if _is_instance_by_name(actuator, "ImplicitActuator"): self.asset.write_joint_damping_to_sim(damping, joint_ids=actuator.joint_indices, env_ids=env_ids) @@ -766,7 +771,7 @@ def __call__( static_friction_coeff = friction_coeff[env_ids_for_slice, joint_ids] # if isaacsim version is lower than 5.0.0 we can set only the static friction coefficient - if get_isaac_sim_version().major >= 5: + if has_kit() and get_isaac_sim_version().major >= 5: # Randomize raw tensors dynamic_friction_coeff = _randomize_prop_by_op( self.default_dynamic_joint_friction_coeff.clone(), @@ -1299,7 +1304,7 @@ def reset_joints_by_scale( joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits) # set into the physics simulation - asset.write_joint_state_to_sim(joint_pos, joint_vel, joint_ids=asset_cfg.joint_ids, env_ids=env_ids) + asset.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, joint_ids=asset_cfg.joint_ids, env_ids=env_ids) def reset_joints_by_offset( @@ -1339,7 +1344,7 @@ def reset_joints_by_offset( joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits) # set into the physics simulation - asset.write_joint_state_to_sim(joint_pos, joint_vel, joint_ids=asset_cfg.joint_ids, env_ids=env_ids) + asset.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, joint_ids=asset_cfg.joint_ids, env_ids=env_ids) def reset_nodal_state_uniform( @@ -1412,7 +1417,7 @@ def reset_scene_to_default(env: ManagerBasedEnv, env_ids: torch.Tensor, reset_jo default_joint_pos = wp.to_torch(articulation_asset.data.default_joint_pos)[env_ids].clone() default_joint_vel = wp.to_torch(articulation_asset.data.default_joint_vel)[env_ids].clone() # set into the physics simulation - articulation_asset.write_joint_state_to_sim(default_joint_pos, default_joint_vel, env_ids=env_ids) + articulation_asset.write_joint_state_to_sim(position=default_joint_pos, velocity=default_joint_vel, env_ids=env_ids) # reset joint targets if required if reset_joint_targets: articulation_asset.set_joint_position_target(default_joint_pos, env_ids=env_ids) @@ -1463,6 +1468,8 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): ) # enable replicator extension if not already enabled + from isaacsim.core.utils.extensions import enable_extension + enable_extension("omni.replicator.core") # we import the module here since we may not always need the replicator @@ -1531,8 +1538,7 @@ def rep_texture_randomization(): with rep.trigger.on_custom_event(event_name=event_name): rep_texture_randomization() else: - # acquire stage - stage = get_current_stage() + stage = env.sim.stage prims_group = rep.functional.get.prims(path_pattern=prim_path, stage=stage) num_prims = len(prims_group) @@ -1624,6 +1630,8 @@ def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): super().__init__(cfg, env) # enable replicator extension if not already enabled + from isaacsim.core.utils.extensions import enable_extension + enable_extension("omni.replicator.core") # we import the module here since we may not always need the replicator import omni.replicator.core as rep @@ -1680,7 +1688,7 @@ def rep_color_randomization(): with rep.trigger.on_custom_event(event_name=event_name): rep_color_randomization() else: - stage = get_current_stage() + stage = env.sim.stage prims_group = rep.functional.get.prims(path_pattern=mesh_prim_path, stage=stage) num_prims = len(prims_group) diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index 4a5b6b5f163..5bae5ec2d54 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -17,14 +17,14 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import SceneEntityCfg from isaaclab.managers.manager_base import ManagerTermBase from isaaclab.managers.manager_term_cfg import ObservationTermCfg -from isaaclab.sensors import Camera, Imu, RayCaster, RayCasterCamera, TiledCamera if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv + from isaaclab.sensors import Camera, Imu, RayCaster, RayCasterCamera, TiledCamera from isaaclab.envs.utils.io_descriptors import ( generic_io_descriptor, diff --git a/source/isaaclab/isaaclab/envs/mdp/rewards.py b/source/isaaclab/isaaclab/envs/mdp/rewards.py index 9e427cad836..5a53583f5e4 100644 --- a/source/isaaclab/isaaclab/envs/mdp/rewards.py +++ b/source/isaaclab/isaaclab/envs/mdp/rewards.py @@ -16,14 +16,14 @@ import torch import warp as wp -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import SceneEntityCfg from isaaclab.managers.manager_base import ManagerTermBase from isaaclab.managers.manager_term_cfg import RewardTermCfg -from isaaclab.sensors import ContactSensor, RayCaster if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import ContactSensor, RayCaster """ General. diff --git a/source/isaaclab/isaaclab/envs/mdp/terminations.py b/source/isaaclab/isaaclab/envs/mdp/terminations.py index 37a2bd3d2c0..1936f5dd50b 100644 --- a/source/isaaclab/isaaclab/envs/mdp/terminations.py +++ b/source/isaaclab/isaaclab/envs/mdp/terminations.py @@ -16,13 +16,13 @@ import torch import warp as wp -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import SceneEntityCfg -from isaaclab.sensors import ContactSensor if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv from isaaclab.managers.command_manager import CommandTerm + from isaaclab.sensors import ContactSensor """ MDP terminations. diff --git a/source/isaaclab/isaaclab/envs/ui/__init__.py b/source/isaaclab/isaaclab/envs/ui/__init__.py index 93db88399e4..b0e5782b733 100644 --- a/source/isaaclab/isaaclab/envs/ui/__init__.py +++ b/source/isaaclab/isaaclab/envs/ui/__init__.py @@ -10,7 +10,14 @@ toggling different debug visualization tools, and other user-defined functionalities. """ -from .base_env_window import BaseEnvWindow -from .empty_window import EmptyWindow -from .manager_based_rl_env_window import ManagerBasedRLEnvWindow -from .viewport_camera_controller import ViewportCameraController +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "base_env_window": ["BaseEnvWindow"], + "empty_window": ["EmptyWindow"], + "manager_based_rl_env_window": ["ManagerBasedRLEnvWindow"], + "viewport_camera_controller": ["ViewportCameraController"], + }, +) diff --git a/source/isaaclab/isaaclab/envs/ui/base_env_window.py b/source/isaaclab/isaaclab/envs/ui/base_env_window.py index 85eca9234fe..7f7d9ebfb7c 100644 --- a/source/isaaclab/isaaclab/envs/ui/base_env_window.py +++ b/source/isaaclab/isaaclab/envs/ui/base_env_window.py @@ -12,9 +12,6 @@ from typing import TYPE_CHECKING import isaacsim -import omni.kit.app -import omni.kit.commands -import omni.usd from pxr import Sdf, Usd, UsdGeom, UsdPhysics from isaaclab.sim.utils.stage import get_current_stage, resolve_paths @@ -64,6 +61,8 @@ def __init__(self, env: ManagerBasedEnv, window_name: str = "IsaacLab"): # get stage handle self.stage = get_current_stage() + import omni.ui + # Listeners for environment selection changes self._ui_listeners: list[ManagerLiveVisualizer] = [] @@ -300,6 +299,8 @@ def _visualize_manager(self, title: str, class_name: str) -> None: def _toggle_recording_animation_fn(self, value: bool): """Toggles the animation recording.""" + import omni.kit.commands + if value: # log directory to save the recording if not hasattr(self, "animation_log_dir"): diff --git a/source/isaaclab/isaaclab/envs/ui/empty_window.py b/source/isaaclab/isaaclab/envs/ui/empty_window.py index bc12f862d1b..26c8112dfe2 100644 --- a/source/isaaclab/isaaclab/envs/ui/empty_window.py +++ b/source/isaaclab/isaaclab/envs/ui/empty_window.py @@ -8,8 +8,6 @@ import asyncio from typing import TYPE_CHECKING -import omni.kit.app - if TYPE_CHECKING: import omni.ui diff --git a/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py b/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py index 0ba5368734b..90dbcc694d8 100644 --- a/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py +++ b/source/isaaclab/isaaclab/envs/ui/viewport_camera_controller.py @@ -13,9 +13,6 @@ import numpy as np import torch -import omni.kit.app -import omni.timeline - from isaaclab.assets.articulation.articulation import Articulation if TYPE_CHECKING: @@ -76,6 +73,8 @@ def __init__(self, env: ManagerBasedEnv | DirectRLEnv, cfg: ViewerCfg): self.update_view_to_world() # subscribe to post update event so that camera view can be updated at each rendering step + import omni.kit.app + app_interface = omni.kit.app.get_app_interface() app_event_stream = app_interface.get_post_update_event_stream() self._viewport_camera_update_handle = app_event_stream.create_subscription_to_pop( diff --git a/source/isaaclab/isaaclab/envs/utils/marl.py b/source/isaaclab/isaaclab/envs/utils/marl.py index 5901bfbc318..2774da6941e 100644 --- a/source/isaaclab/isaaclab/envs/utils/marl.py +++ b/source/isaaclab/isaaclab/envs/utils/marl.py @@ -6,15 +6,17 @@ from __future__ import annotations import math -from typing import Any +from typing import TYPE_CHECKING, Any import gymnasium as gym import numpy as np import torch from ..common import ActionType, AgentID, EnvStepReturn, ObsType, StateType, VecEnvObs, VecEnvStepReturn -from ..direct_marl_env import DirectMARLEnv -from ..direct_rl_env import DirectRLEnv + +if TYPE_CHECKING: + from ..direct_marl_env import DirectMARLEnv + from ..direct_rl_env import DirectRLEnv def multi_agent_to_single_agent(env: DirectMARLEnv, state_as_observation: bool = False) -> DirectRLEnv: @@ -45,6 +47,9 @@ def multi_agent_to_single_agent(env: DirectMARLEnv, state_as_observation: bool = as unconstructed (:attr:`DirectMARLEnvCfg.state_space`). """ + from ..direct_marl_env import DirectMARLEnv # noqa: F811 + from ..direct_rl_env import DirectRLEnv # noqa: F811 + class Env(DirectRLEnv): def __init__(self, env: DirectMARLEnv) -> None: self.env: DirectMARLEnv = env.unwrapped @@ -165,6 +170,8 @@ def multi_agent_with_one_agent(env: DirectMARLEnv, state_as_observation: bool = as unconstructed (:attr:`DirectMARLEnvCfg.state_space`). """ + from ..direct_marl_env import DirectMARLEnv # noqa: F811 + class Env(DirectMARLEnv): def __init__(self, env: DirectMARLEnv) -> None: self.env: DirectMARLEnv = env.unwrapped diff --git a/source/isaaclab/isaaclab/managers/__init__.py b/source/isaaclab/isaaclab/managers/__init__.py index 4a8266801b4..d3c2cf87218 100644 --- a/source/isaaclab/isaaclab/managers/__init__.py +++ b/source/isaaclab/isaaclab/managers/__init__.py @@ -10,25 +10,50 @@ designed to be modular and can be easily extended to support new functionality. """ -from .action_manager import ActionManager, ActionTerm -from .command_manager import CommandManager, CommandTerm -from .curriculum_manager import CurriculumManager -from .event_manager import EventManager -from .manager_base import ManagerBase, ManagerTermBase +import lazy_loader as lazy + from .manager_term_cfg import ( ActionTermCfg, CommandTermCfg, CurriculumTermCfg, + DatasetExportMode, EventTermCfg, ManagerTermBaseCfg, ObservationGroupCfg, ObservationTermCfg, + RecorderManagerBaseCfg, RecorderTermCfg, RewardTermCfg, TerminationTermCfg, ) -from .observation_manager import ObservationManager -from .recorder_manager import DatasetExportMode, RecorderManager, RecorderManagerBaseCfg, RecorderTerm -from .reward_manager import RewardManager from .scene_entity_cfg import SceneEntityCfg -from .termination_manager import TerminationManager + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "action_manager": ["ActionManager", "ActionTerm"], + "command_manager": ["CommandManager", "CommandTerm"], + "curriculum_manager": ["CurriculumManager"], + "event_manager": ["EventManager"], + "manager_base": ["ManagerBase", "ManagerTermBase"], + "observation_manager": ["ObservationManager"], + "recorder_manager": ["RecorderManager", "RecorderTerm"], + "reward_manager": ["RewardManager"], + "termination_manager": ["TerminationManager"], + }, +) +__all__ += [ + "ActionTermCfg", + "CommandTermCfg", + "CurriculumTermCfg", + "DatasetExportMode", + "EventTermCfg", + "ManagerTermBaseCfg", + "ObservationGroupCfg", + "ObservationTermCfg", + "RecorderManagerBaseCfg", + "RecorderTermCfg", + "RewardTermCfg", + "TerminationTermCfg", + "SceneEntityCfg", +] diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index 6d138451f99..bed47a2da41 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -17,8 +17,6 @@ import torch from prettytable import PrettyTable -import omni.kit.app - from isaaclab.envs.utils.io_descriptors import GenericActionIODescriptor from .manager_base import ManagerBase, ManagerTermBase @@ -132,6 +130,8 @@ def set_debug_vis(self, debug_vis: bool) -> bool: if debug_vis: # create a subscriber for the post update event if it doesn't exist if self._debug_vis_handle is None: + import omni.kit.app + app_interface = omni.kit.app.get_app_interface() self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) diff --git a/source/isaaclab/isaaclab/managers/command_manager.py b/source/isaaclab/isaaclab/managers/command_manager.py index 0fe66ff6b96..412de30fdff 100644 --- a/source/isaaclab/isaaclab/managers/command_manager.py +++ b/source/isaaclab/isaaclab/managers/command_manager.py @@ -16,7 +16,7 @@ import torch from prettytable import PrettyTable -import omni.kit.app +from isaaclab.utils.version import has_kit from .manager_base import ManagerBase, ManagerTermBase from .manager_term_cfg import CommandTermCfg @@ -102,10 +102,14 @@ def set_debug_vis(self, debug_vis: bool) -> bool: return False # toggle debug visualization objects self._set_debug_vis_impl(debug_vis) - # toggle debug visualization handles + # toggle debug visualization handles (only in Kit environments) + if not has_kit(): + return True if debug_vis: # create a subscriber for the post update event if it doesn't exist if self._debug_vis_handle is None: + import omni.kit.app + app_interface = omni.kit.app.get_app_interface() self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) diff --git a/source/isaaclab/isaaclab/managers/manager_base.py b/source/isaaclab/isaaclab/managers/manager_base.py index 158c713abfa..6da12a2c89e 100644 --- a/source/isaaclab/isaaclab/managers/manager_base.py +++ b/source/isaaclab/isaaclab/managers/manager_base.py @@ -5,6 +5,7 @@ from __future__ import annotations +import contextlib import copy import inspect import logging @@ -13,9 +14,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any -import omni.timeline - import isaaclab.utils.string as string_utils +from isaaclab.physics import PhysicsEvent from isaaclab.utils import class_to_dict, string_to_callable from .manager_term_cfg import ManagerTermBaseCfg @@ -156,18 +156,24 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): # if the simulation is not playing, we use callbacks to trigger the resolution of the scene # entities configuration. this is needed for cases where the manager is created after the # simulation, but before the simulation is playing. - # FIXME: Once Isaac Sim supports storing this information as USD schema, we can remove this - # callback and resolve the scene entities directly inside `_prepare_terms`. if not self._env.sim.is_playing(): # note: Use weakref on all callbacks to ensure that this object can be deleted when its destructor - # is called - # The order is set to 20 to allow asset/sensor initialization to complete before the scene entities - # are resolved. Those have the order 10. - timeline_event_stream = omni.timeline.get_timeline_interface().get_timeline_event_stream() - self._resolve_terms_handle = timeline_event_stream.create_subscription_to_pop_by_type( - int(omni.timeline.TimelineEventType.PLAY), - lambda event, obj=weakref.proxy(self): obj._resolve_terms_callback(event), - order=20, + # is called. The order is set to 20 to allow asset/sensor initialization to complete before the + # scene entities are resolved. Those have the order 10. + + def safe_callback(obj_ref): + """Safely invoke _resolve_terms_callback on a weakly-referenced object.""" + with contextlib.suppress(ReferenceError): + obj_ref._resolve_terms_callback(None) + + obj_ref = weakref.proxy(self) + + def callback_func(payload): + safe_callback(obj_ref) + + physics_mgr_cls = self._env.sim.physics_manager + self._resolve_terms_handle = physics_mgr_cls.register_callback( + callback_func, PhysicsEvent.PHYSICS_READY, order=20 ) else: self._resolve_terms_handle = None @@ -178,8 +184,8 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): def __del__(self): """Delete the manager.""" - if self._resolve_terms_handle: - self._resolve_terms_handle.unsubscribe() + if self._resolve_terms_handle is not None: + self._resolve_terms_handle.deregister() self._resolve_terms_handle = None """ diff --git a/source/isaaclab/isaaclab/managers/manager_term_cfg.py b/source/isaaclab/isaaclab/managers/manager_term_cfg.py index de7c23aa220..85124753946 100644 --- a/source/isaaclab/isaaclab/managers/manager_term_cfg.py +++ b/source/isaaclab/isaaclab/managers/manager_term_cfg.py @@ -7,13 +7,14 @@ from __future__ import annotations +import enum from collections.abc import Callable from dataclasses import MISSING from typing import TYPE_CHECKING, Any import torch -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from isaaclab.utils.modifiers import ModifierCfg from isaaclab.utils.noise import NoiseCfg, NoiseModelCfg @@ -353,3 +354,36 @@ class TerminationTermCfg(ManagerTermBaseCfg): Note: These usually correspond to tasks that have a fixed time limit. """ + + +class DatasetExportMode(enum.IntEnum): + """The mode to handle episode exports.""" + + EXPORT_NONE = 0 + EXPORT_ALL = 1 + EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES = 2 + EXPORT_SUCCEEDED_ONLY = 3 + + +@configclass +class RecorderManagerBaseCfg: + """Base class for configuring recorder manager terms.""" + + dataset_file_handler_class_type: type | DeferredClass = DeferredClass( + "isaaclab.utils.datasets:HDF5DatasetFileHandler" + ) + + dataset_export_dir_path: str = "/tmp/isaaclab/logs" + """The directory path where the recorded datasets are exported.""" + + dataset_filename: str = "dataset" + """Dataset file name without file extension.""" + + dataset_export_mode: DatasetExportMode = DatasetExportMode.EXPORT_ALL + """The mode to handle episode exports.""" + + export_in_record_pre_reset: bool = True + """Whether to export episodes in the record_pre_reset call.""" + + export_in_close: bool = False + """Whether to export episodes in the close call.""" diff --git a/source/isaaclab/isaaclab/managers/recorder_manager.py b/source/isaaclab/isaaclab/managers/recorder_manager.py index cc95d238156..4b30bde3abe 100644 --- a/source/isaaclab/isaaclab/managers/recorder_manager.py +++ b/source/isaaclab/isaaclab/managers/recorder_manager.py @@ -6,7 +6,6 @@ from __future__ import annotations -import enum import os from collections.abc import Sequence from typing import TYPE_CHECKING @@ -15,47 +14,15 @@ import warp as wp from prettytable import PrettyTable -from isaaclab.utils import configclass from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler from .manager_base import ManagerBase, ManagerTermBase -from .manager_term_cfg import RecorderTermCfg +from .manager_term_cfg import DatasetExportMode, RecorderManagerBaseCfg, RecorderTermCfg if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv -class DatasetExportMode(enum.IntEnum): - """The mode to handle episode exports.""" - - EXPORT_NONE = 0 # Export none of the episodes - EXPORT_ALL = 1 # Export all episodes to a single dataset file - EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES = 2 # Export succeeded and failed episodes in separate files - EXPORT_SUCCEEDED_ONLY = 3 # Export only succeeded episodes to a single dataset file - - -@configclass -class RecorderManagerBaseCfg: - """Base class for configuring recorder manager terms.""" - - dataset_file_handler_class_type: type = HDF5DatasetFileHandler - - dataset_export_dir_path: str = "/tmp/isaaclab/logs" - """The directory path where the recorded datasets are exported.""" - - dataset_filename: str = "dataset" - """Dataset file name without file extension.""" - - dataset_export_mode: DatasetExportMode = DatasetExportMode.EXPORT_ALL - """The mode to handle episode exports.""" - - export_in_record_pre_reset: bool = True - """Whether to export episodes in the record_pre_reset call.""" - - export_in_close: bool = False - """Whether to export episodes in the close call.""" - - class RecorderTerm(ManagerTermBase): """Base class for recorder terms. diff --git a/source/isaaclab/isaaclab/markers/__init__.py b/source/isaaclab/isaaclab/markers/__init__.py index eb7b6976100..348feb8b5c2 100644 --- a/source/isaaclab/isaaclab/markers/__init__.py +++ b/source/isaaclab/isaaclab/markers/__init__.py @@ -21,5 +21,15 @@ """ -from .config import * # noqa: F401, F403 -from .visualization_markers import VisualizationMarkers, VisualizationMarkersCfg +import lazy_loader as lazy + +from .visualization_markers_cfg import VisualizationMarkersCfg + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["config"], + submod_attrs={ + "visualization_markers": ["VisualizationMarkers"], + }, +) +__all__ += ["VisualizationMarkersCfg"] diff --git a/source/isaaclab/isaaclab/markers/config/__init__.py b/source/isaaclab/isaaclab/markers/config/__init__.py index 54d30051259..54713d65935 100644 --- a/source/isaaclab/isaaclab/markers/config/__init__.py +++ b/source/isaaclab/isaaclab/markers/config/__init__.py @@ -4,7 +4,7 @@ # SPDX-License-Identifier: BSD-3-Clause import isaaclab.sim as sim_utils -from isaaclab.markers.visualization_markers import VisualizationMarkersCfg +from isaaclab.markers.visualization_markers_cfg import VisualizationMarkersCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR ## diff --git a/source/isaaclab/isaaclab/markers/visualization_markers.py b/source/isaaclab/isaaclab/markers/visualization_markers.py index b87765c547d..dad0e23f0d4 100644 --- a/source/isaaclab/isaaclab/markers/visualization_markers.py +++ b/source/isaaclab/isaaclab/markers/visualization_markers.py @@ -20,37 +20,18 @@ from __future__ import annotations import logging -from dataclasses import MISSING import numpy as np import torch -import omni.physx.scripts.utils as physx_utils -from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics, Vt - import isaaclab.sim as sim_utils -from isaaclab.sim.spawners import SpawnerCfg -from isaaclab.utils.configclass import configclass + +from .visualization_markers_cfg import VisualizationMarkersCfg # import logger logger = logging.getLogger(__name__) -@configclass -class VisualizationMarkersCfg: - """A class to configure a :class:`VisualizationMarkers`.""" - - prim_path: str = MISSING - """The prim path where the :class:`UsdGeom.PointInstancer` will be created.""" - - markers: dict[str, SpawnerCfg] = MISSING - """The dictionary of marker configurations. - - The key is the name of the marker, and the value is the configuration of the marker. - The key is used to identify the marker in the class. - """ - - class VisualizationMarkers: """A class to coordinate groups of visual markers (loaded from USD). @@ -144,6 +125,8 @@ def __init__(self, cfg: VisualizationMarkersCfg): Raises: ValueError: When no markers are provided in the :obj:`cfg`. """ + from pxr import Gf, UsdGeom + # get next free path for the prim prim_path = sim_utils.get_next_free_prim_path(cfg.prim_path) # create a new prim @@ -202,6 +185,8 @@ def set_visibility(self, visible: bool): Args: visible: flag to set the visibility. """ + from pxr import UsdGeom + imageable = UsdGeom.Imageable(self._instancer_manager) if visible: imageable.MakeVisible() @@ -214,6 +199,8 @@ def is_visible(self) -> bool: Returns: True if the markers are visible, False otherwise. """ + from pxr import UsdGeom + return self._instancer_manager.GetVisibilityAttr().Get() != UsdGeom.Tokens.invisible def visualize( @@ -270,6 +257,8 @@ def visualize( ValueError: When input arrays do not follow the expected shapes. ValueError: When the function is called with all None arguments. """ + from pxr import Vt + # check if it is visible (if not then let's not waste time) if not self.is_visible(): return @@ -370,6 +359,9 @@ def _process_prototype_prim(self, prim: Usd.Prim): Args: prim: The prim to check. """ + import omni.physx.scripts.utils as physx_utils + from pxr import Sdf, UsdGeom, UsdPhysics + # check if prim is valid if not prim.IsValid(): raise ValueError(f"Prim at path '{prim.GetPrimAtPath()}' is not valid.") @@ -402,5 +394,11 @@ def _process_prototype_prim(self, prim: Usd.Prim): # add children to list all_prims += child_prim.GetChildren() - # remove any physics on the markers because they are only for visualization! - physx_utils.removeRigidBodySubtree(prim) + # remove any physics on the markers because they are only for visualization (PhysX backend only) + from isaaclab.sim.simulation_context import SimulationContext + + sim_ctx = SimulationContext.instance() + if sim_ctx is not None and "Physx" in sim_ctx.physics_manager.__name__: + import omni.physx.scripts.utils as physx_utils + + physx_utils.removeRigidBodySubtree(prim) diff --git a/source/isaaclab/isaaclab/markers/visualization_markers_cfg.py b/source/isaaclab/isaaclab/markers/visualization_markers_cfg.py new file mode 100644 index 00000000000..b7cc2002ccf --- /dev/null +++ b/source/isaaclab/isaaclab/markers/visualization_markers_cfg.py @@ -0,0 +1,28 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for visualization markers.""" + +from __future__ import annotations + +from dataclasses import MISSING + +from isaaclab.sim.spawners import SpawnerCfg +from isaaclab.utils.configclass import configclass + + +@configclass +class VisualizationMarkersCfg: + """A class to configure a :class:`VisualizationMarkers`.""" + + prim_path: str = MISSING + """The prim path where the :class:`UsdGeom.PointInstancer` will be created.""" + + markers: dict[str, SpawnerCfg] = MISSING + """The dictionary of marker configurations. + + The key is the name of the marker, and the value is the configuration of the marker. + The key is used to identify the marker in the class. + """ diff --git a/source/isaaclab/isaaclab/physics/physics_manager.py b/source/isaaclab/isaaclab/physics/physics_manager.py index 13132cc2062..fe79c1d2cd5 100644 --- a/source/isaaclab/isaaclab/physics/physics_manager.py +++ b/source/isaaclab/isaaclab/physics/physics_manager.py @@ -156,8 +156,6 @@ def dispatch_event(cls, event: PhysicsEvent, payload: Any = None) -> None: callback(payload) except ReferenceError: cls.deregister_callback(cid) - except Exception as e: - logger.error(f"Callback {cid} for {event.value} failed: {e}") @classmethod def clear_callbacks(cls) -> None: diff --git a/source/isaaclab/isaaclab/renderers/newton_warp_renderer.py b/source/isaaclab/isaaclab/renderers/newton_warp_renderer.py index 6185a1e4eda..9216c626f3d 100644 --- a/source/isaaclab/isaaclab/renderers/newton_warp_renderer.py +++ b/source/isaaclab/isaaclab/renderers/newton_warp_renderer.py @@ -13,16 +13,11 @@ import torch import warp as wp -from isaaclab.sim import SimulationContext from isaaclab.utils.math import convert_camera_frame_orientation_convention -from ..visualizers import VisualizerCfg - if TYPE_CHECKING: from isaaclab.sensors import SensorBase - from ..sim.scene_data_providers import SceneDataProvider - logger = logging.getLogger(__name__) @@ -86,14 +81,17 @@ def get_output(self, output_name: str) -> wp.array: return None def update(self, positions: torch.Tensor, orientations: torch.Tensor, intrinsics: torch.Tensor): + if self.render_context is None: + return + converted_orientations = convert_camera_frame_orientation_convention( orientations, origin="world", target="opengl" ) - self.camera_transforms = wp.empty((1, self.render_context.world_count), dtype=wp.transformf) + self.camera_transforms = wp.empty((1, self.render_context.num_worlds), dtype=wp.transformf) wp.launch( RenderData._update_transforms, - self.render_context.world_count, + self.render_context.num_worlds, [positions, converted_orientations, self.camera_transforms], ) @@ -111,14 +109,14 @@ def _from_torch(self, tensor: torch.Tensor, dtype) -> wp.array: return wp.array( ptr=torch_array.ptr, dtype=dtype, - shape=(self.render_context.world_count, self.num_cameras, self.height, self.width), + shape=(self.render_context.num_worlds, self.num_cameras, self.height, self.width), device=torch_array.device, copy=False, ) logger.warning("NewtonWarpRenderer - torch output array is non-contiguous") return wp.zeros( - (self.render_context.world_count, self.num_cameras, self.height, self.width), + (self.render_context.num_worlds, self.num_cameras, self.height, self.width), dtype=dtype, device=torch_array.device, ) @@ -137,16 +135,41 @@ class NewtonWarpRenderer: RenderData = RenderData def __init__(self): - self.newton_sensor = newton.sensors.SensorTiledCamera(self.get_scene_data_provider().get_newton_model()) + # Defer SensorTiledCamera creation until the Newton model is available (after physics init). + # Newton is not yet started when __init__ runs (called from _setup_scene), so we cannot + # access NewtonManager.get_model() here — it returns None until MODEL_INIT callbacks fire. + self.newton_sensor: newton.sensors.SensorTiledCamera | None = None + logger.info("NewtonWarpRenderer: created (sensor creation deferred until first use)") + + def _ensure_sensor(self) -> None: + """Create the Newton sensor the first time it is needed (after Newton physics is ready).""" + if self.newton_sensor is not None: + return + from isaaclab_newton.physics import NewtonManager + + model = NewtonManager.get_model() + if model is None: + raise RuntimeError( + "NewtonWarpRenderer: Newton model is not available yet. " + "Ensure that the renderer is used only after Newton physics has been initialized." + ) + logger.info( + f"NewtonWarpRenderer: creating SensorTiledCamera with Newton model " + f"(num_worlds={model.num_worlds})" + ) + self.newton_sensor = newton.sensors.SensorTiledCamera(model) def create_render_data(self, sensor: SensorBase) -> RenderData: + self._ensure_sensor() + assert self.newton_sensor is not None return RenderData(self.newton_sensor.render_context, sensor) def set_outputs(self, render_data: RenderData, output_data: dict[str, torch.Tensor]): render_data.set_outputs(output_data) def update_transforms(self): - SimulationContext.instance().update_scene_data_provider(True) + # No-op: state is read live from NewtonManager each render call, no caching to flush. + pass def update_camera( self, render_data: RenderData, positions: torch.Tensor, orientations: torch.Tensor, intrinsics: torch.Tensor @@ -154,8 +177,18 @@ def update_camera( render_data.update(positions, orientations, intrinsics) def render(self, render_data: RenderData): + from isaaclab_newton.physics import NewtonManager + + state = NewtonManager.get_state_0() + if state is None: + logger.warning("NewtonWarpRenderer.render: Newton state is None, skipping render") + return + + logger.debug(f"NewtonWarpRenderer.render: using live Newton state (type={type(state).__name__})") + self._ensure_sensor() + assert self.newton_sensor is not None self.newton_sensor.render( - self.get_scene_data_provider().get_newton_state(), + state, render_data.camera_transforms, render_data.camera_rays, color_image=render_data.outputs.color_image, @@ -170,6 +203,3 @@ def write_output(self, render_data: RenderData, output_name: str, output_data: t if image_data is not None: if image_data.ptr != output_data.data_ptr(): wp.copy(wp.from_torch(output_data), image_data) - - def get_scene_data_provider(self) -> SceneDataProvider: - return SimulationContext.instance().initialize_scene_data_provider([VisualizerCfg(visualizer_type="newton")]) diff --git a/source/isaaclab/isaaclab/scene/__init__.py b/source/isaaclab/isaaclab/scene/__init__.py index 4b614ea4bce..ce8de0a53f8 100644 --- a/source/isaaclab/isaaclab/scene/__init__.py +++ b/source/isaaclab/isaaclab/scene/__init__.py @@ -25,5 +25,14 @@ :mod:`isaaclab.managers` sub-package for more details. """ -from .interactive_scene import InteractiveScene +import lazy_loader as lazy + from .interactive_scene_cfg import InteractiveSceneCfg + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "interactive_scene": ["InteractiveScene"], + }, +) +__all__ += ["InteractiveSceneCfg"] diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 54fbcc068af..bf82dc2912f 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -7,14 +7,17 @@ import logging from collections.abc import Sequence -from typing import Any +from typing import TYPE_CHECKING, Any import torch import warp as wp -from isaaclab_physx.assets import DeformableObject, DeformableObjectCfg, SurfaceGripper, SurfaceGripperCfg +from isaaclab_physx.assets import DeformableObjectCfg, SurfaceGripperCfg from pxr import Sdf +if TYPE_CHECKING: + from isaaclab_physx.assets import DeformableObject, SurfaceGripper + import isaaclab.sim as sim_utils from isaaclab import cloner from isaaclab.assets import ( @@ -463,7 +466,7 @@ def reset_to( # joint state joint_position = asset_state["joint_position"].clone() joint_velocity = asset_state["joint_velocity"].clone() - articulation.write_joint_state_to_sim(joint_position, joint_velocity, env_ids=env_ids) + articulation.write_joint_state_to_sim(position=joint_position, velocity=joint_velocity, env_ids=env_ids) # FIXME: This is not generic as it assumes PD control over the joints. # This assumption does not hold for effort controlled joints. articulation.set_joint_position_target(joint_position, env_ids=env_ids) @@ -694,8 +697,11 @@ def _add_entities_from_cfg(self): # noqa: C901 self._terrain = asset_cfg.class_type(asset_cfg) elif isinstance(asset_cfg, ArticulationCfg): self._articulations[asset_name] = asset_cfg.class_type(asset_cfg) - elif isinstance(asset_cfg, DeformableObjectCfg): - self._deformable_objects[asset_name] = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, DeformableObjectCfg) or isinstance(asset_cfg, SurfaceGripperCfg): + if isinstance(asset_cfg, DeformableObjectCfg): + self._deformable_objects[asset_name] = asset_cfg.class_type(asset_cfg) + elif isinstance(asset_cfg, SurfaceGripperCfg): + self._surface_grippers[asset_name] = asset_cfg.class_type(asset_cfg) elif isinstance(asset_cfg, RigidObjectCfg): self._rigid_objects[asset_name] = asset_cfg.class_type(asset_cfg) elif isinstance(asset_cfg, RigidObjectCollectionCfg): @@ -716,9 +722,6 @@ def _add_entities_from_cfg(self): # noqa: C901 if hasattr(rigid_object_cfg, "collision_group") and rigid_object_cfg.collision_group == -1: asset_paths = sim_utils.find_matching_prim_paths(rigid_object_cfg.prim_path) self._global_prim_paths += asset_paths - elif isinstance(asset_cfg, SurfaceGripperCfg): - # add surface grippers to scene - self._surface_grippers[asset_name] = asset_cfg.class_type(asset_cfg) elif isinstance(asset_cfg, SensorBaseCfg): # Update target frame path(s)' regex name space for FrameTransformer if isinstance(asset_cfg, FrameTransformerCfg): @@ -728,10 +731,10 @@ def _add_entities_from_cfg(self): # noqa: C901 updated_target_frames.append(target_frame) asset_cfg.target_frames = updated_target_frames elif isinstance(asset_cfg, ContactSensorCfg): - updated_filter_prim_paths_expr = [] - for filter_prim_path in asset_cfg.filter_prim_paths_expr: - updated_filter_prim_paths_expr.append(filter_prim_path.format(ENV_REGEX_NS=self.env_regex_ns)) - asset_cfg.filter_prim_paths_expr = updated_filter_prim_paths_expr + for attr_name in ("filter_prim_paths_expr", "shape_path", "filter_shape_paths_expr"): + paths = getattr(asset_cfg, attr_name, None) + if paths is not None: + setattr(asset_cfg, attr_name, [p.format(ENV_REGEX_NS=self.env_regex_ns) for p in paths]) elif isinstance(asset_cfg, VisuoTactileSensorCfg): if hasattr(asset_cfg, "camera_cfg") and asset_cfg.camera_cfg is not None: asset_cfg.camera_cfg.prim_path = asset_cfg.camera_cfg.prim_path.format( diff --git a/source/isaaclab/isaaclab/sensors/__init__.py b/source/isaaclab/isaaclab/sensors/__init__.py index f0a5719e505..7f030607551 100644 --- a/source/isaaclab/isaaclab/sensors/__init__.py +++ b/source/isaaclab/isaaclab/sensors/__init__.py @@ -35,10 +35,55 @@ """ -from .camera import * # noqa: F401, F403 -from .contact_sensor import * # noqa: F401, F403 -from .frame_transformer import * # noqa: F401 -from .imu import * # noqa: F401, F403 -from .ray_caster import * # noqa: F401, F403 -from .sensor_base import SensorBase # noqa: F401 -from .sensor_base_cfg import SensorBaseCfg # noqa: F401 +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["camera", "contact_sensor", "frame_transformer", "imu", "ray_caster"], + submod_attrs={ + "sensor_base": ["SensorBase"], + "sensor_base_cfg": ["SensorBaseCfg"], + "camera": ["Camera", "CameraCfg", "CameraData", "TiledCamera", "TiledCameraCfg", "save_images_to_file"], + "contact_sensor": [ + "BaseContactSensor", + "BaseContactSensorData", + "ContactSensor", + "ContactSensorCfg", + "ContactSensorData", + ], + "frame_transformer": [ + "BaseFrameTransformer", + "BaseFrameTransformerData", + "FrameTransformer", + "FrameTransformerCfg", + "FrameTransformerData", + "OffsetCfg", + ], + "imu": ["BaseImu", "BaseImuData", "Imu", "ImuCfg", "ImuData"], + "ray_caster": [ + "MultiMeshRayCaster", + "MultiMeshRayCasterCamera", + "MultiMeshRayCasterCameraCfg", + "MultiMeshRayCasterCameraData", + "MultiMeshRayCasterCfg", + "MultiMeshRayCasterData", + "RayCaster", + "RayCasterCamera", + "RayCasterCameraCfg", + "RayCasterCfg", + "RayCasterData", + ], + }, +) + +# Re-export patterns submodule from ray_caster for backward compat +# (from isaaclab.sensors import patterns) +_lazy_getattr = __getattr__ + + +def __getattr__(name): + if name == "patterns": + from isaaclab.sensors.ray_caster import patterns + + return patterns + return _lazy_getattr(name) diff --git a/source/isaaclab/isaaclab/sensors/camera/__init__.py b/source/isaaclab/isaaclab/sensors/camera/__init__.py index f2318067b58..c4defad2ca9 100644 --- a/source/isaaclab/isaaclab/sensors/camera/__init__.py +++ b/source/isaaclab/isaaclab/sensors/camera/__init__.py @@ -5,9 +5,17 @@ """Sub-module for camera wrapper around USD camera prim.""" -from .camera import Camera -from .camera_cfg import CameraCfg -from .camera_data import CameraData -from .tiled_camera import TiledCamera -from .tiled_camera_cfg import TiledCameraCfg -from .utils import * # noqa: F401, F403 +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["utils"], + submod_attrs={ + "camera": ["Camera"], + "camera_cfg": ["CameraCfg"], + "camera_data": ["CameraData"], + "tiled_camera": ["TiledCamera"], + "tiled_camera_cfg": ["TiledCameraCfg"], + "utils": ["save_images_to_file"], + }, +) diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index f2252f0acd5..8bf1184d93b 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -27,7 +27,7 @@ create_rotation_matrix_from_view, quat_from_matrix, ) -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit from ..sensor_base import SensorBase from .camera_data import CameraData @@ -130,8 +130,8 @@ def __init__(self, cfg: CameraCfg): settings.set_bool("/isaaclab/render/rtx_sensors", True) # This is only introduced in isaac sim 6.0 - isaac_sim_version = get_isaac_sim_version() - if isaac_sim_version.major >= 6: + isaac_sim_version = get_isaac_sim_version() if has_kit() else None + if has_kit() and isaac_sim_version is not None and isaac_sim_version.major >= 6: # Set RTX flag to enable fast path if only depth or albedo is requested supported_fast_types = {"distance_to_camera", "distance_to_image_plane", "depth", "albedo"} if all(data_type in supported_fast_types for data_type in self.cfg.data_types): @@ -140,7 +140,7 @@ def __init__(self, cfg: CameraCfg): # If we have GUI / viewport enabled, we turn off fast path so that the viewport is not black if settings.get("/isaaclab/has_gui"): settings.set_bool("/rtx/sdg/force/disableColorRender", False) - else: + elif has_kit() and isaac_sim_version is not None and isaac_sim_version.major < 6: if "albedo" in self.cfg.data_types: logger.warning( "Albedo annotator is only supported in Isaac Sim 6.0+. The albedo data type will be ignored." @@ -194,7 +194,7 @@ def __init__(self, cfg: CameraCfg): # HACK: We need to disable instancing for semantic_segmentation and instance_segmentation_fast to work # checks for Isaac Sim v4.5 as this issue exists there - if get_isaac_sim_version() == version.parse("4.5"): + if has_kit() and get_isaac_sim_version() == version.parse("4.5"): if "semantic_segmentation" in self.cfg.data_types or "instance_segmentation_fast" in self.cfg.data_types: logger.warning( "Isaac Sim 4.5 introduced a bug in Camera and TiledCamera when outputting instance and semantic" diff --git a/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py b/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py index a589b0d170e..78cde4e2c9c 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera_cfg.py @@ -9,10 +9,9 @@ from typing import Literal from isaaclab.sim import FisheyeCameraCfg, PinholeCameraCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from ..sensor_base_cfg import SensorBaseCfg -from .camera import Camera @configclass @@ -39,7 +38,7 @@ class OffsetCfg: """ - class_type: type = Camera + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.camera.camera:Camera") offset: OffsetCfg = OffsetCfg() """The offset pose of the sensor's frame from the sensor's parent frame. Defaults to identity. diff --git a/source/isaaclab/isaaclab/sensors/camera/tiled_camera.py b/source/isaaclab/isaaclab/sensors/camera/tiled_camera.py index 60f4974d639..59505dda0ce 100644 --- a/source/isaaclab/isaaclab/sensors/camera/tiled_camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/tiled_camera.py @@ -16,6 +16,7 @@ from pxr import UsdGeom +from isaaclab.utils.version import has_kit from isaaclab.app.settings_manager import get_settings_manager from isaaclab.sim.views import XformPrimView from isaaclab.utils.warp.kernels import reshape_tiled_image @@ -149,14 +150,12 @@ def _initialize_impl(self): RuntimeError: If the number of camera prims in the view does not match the number of environments. RuntimeError: If replicator was not found. """ - if not get_settings_manager().get("/isaaclab/cameras_enabled"): + if has_kit() and not get_settings_manager().get("/isaaclab/cameras_enabled"): raise RuntimeError( "A camera was spawned without the --enable_cameras flag. Please use --enable_cameras to enable" " rendering." ) - import omni.replicator.core as rep - # Initialize parent class SensorBase._initialize_impl(self) # Create a view for the sensor @@ -189,6 +188,8 @@ def _initialize_impl(self): self.render_data = self.renderer.create_render_data(self) else: + # Kit/RTX path: replicator and annotators are only available when Kit is running + import omni.replicator.core as rep # Create replicator tiled render product rp = rep.create.render_product_tiled( cameras=cam_prim_paths, tile_resolution=(self.cfg.width, self.cfg.height) diff --git a/source/isaaclab/isaaclab/sensors/camera/tiled_camera_cfg.py b/source/isaaclab/isaaclab/sensors/camera/tiled_camera_cfg.py index 2241a0648fd..f82514acd7a 100644 --- a/source/isaaclab/isaaclab/sensors/camera/tiled_camera_cfg.py +++ b/source/isaaclab/isaaclab/sensors/camera/tiled_camera_cfg.py @@ -3,14 +3,13 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from .camera_cfg import CameraCfg -from .tiled_camera import TiledCamera @configclass class TiledCameraCfg(CameraCfg): """Configuration for a tiled rendering-based camera sensor.""" - class_type: type = TiledCamera + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.camera.tiled_camera:TiledCamera") diff --git a/source/isaaclab/isaaclab/sensors/contact_sensor/__init__.py b/source/isaaclab/isaaclab/sensors/contact_sensor/__init__.py index 5f2d3df68a0..c0d18f9a3dd 100644 --- a/source/isaaclab/isaaclab/sensors/contact_sensor/__init__.py +++ b/source/isaaclab/isaaclab/sensors/contact_sensor/__init__.py @@ -5,10 +5,15 @@ """Sub-module for rigid contact sensor.""" -from .base_contact_sensor import BaseContactSensor -from .base_contact_sensor_data import BaseContactSensorData -from .contact_sensor import ContactSensor -from .contact_sensor_cfg import ContactSensorCfg -from .contact_sensor_data import ContactSensorData +import lazy_loader as lazy -__all__ = ["BaseContactSensor", "BaseContactSensorData", "ContactSensor", "ContactSensorCfg", "ContactSensorData"] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "base_contact_sensor": ["BaseContactSensor"], + "base_contact_sensor_data": ["BaseContactSensorData"], + "contact_sensor": ["ContactSensor"], + "contact_sensor_cfg": ["ContactSensorCfg"], + "contact_sensor_data": ["ContactSensorData"], + }, +) diff --git a/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor.py b/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor.py index f0f592597a1..3884d14e869 100644 --- a/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor.py +++ b/source/isaaclab/isaaclab/sensors/contact_sensor/base_contact_sensor.py @@ -182,22 +182,5 @@ def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: return wp.from_torch(currently_detached * less_than_dt_detached) """ - Implementation - Abstract methods to be implemented by backend-specific subclasses. + Implementation - Methods to be implemented by backend-specific subclasses. """ - - @abstractmethod - def _initialize_impl(self): - """Initializes the sensor handles and internal buffers. - - Subclasses should call ``super()._initialize_impl()`` first to initialize - the common sensor infrastructure from :class:`SensorBase`. - """ - super()._initialize_impl() - - @abstractmethod - def _update_buffers_impl(self, env_ids: Sequence[int]): - raise NotImplementedError - - def _invalidate_initialize_callback(self, event): - """Invalidates the scene elements.""" - super()._invalidate_initialize_callback(event) diff --git a/source/isaaclab/isaaclab/sensors/contact_sensor/contact_sensor_cfg.py b/source/isaaclab/isaaclab/sensors/contact_sensor/contact_sensor_cfg.py index c811a7ca63d..2620a1db617 100644 --- a/source/isaaclab/isaaclab/sensors/contact_sensor/contact_sensor_cfg.py +++ b/source/isaaclab/isaaclab/sensors/contact_sensor/contact_sensor_cfg.py @@ -5,53 +5,38 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.markers.config import CONTACT_SENSOR_MARKER_CFG -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from ..sensor_base_cfg import SensorBaseCfg -from .contact_sensor import ContactSensor @configclass class ContactSensorCfg(SensorBaseCfg): - """Configuration for the contact sensor.""" + """Base configuration for the contact sensor. - class_type: type = ContactSensor + This contains fields common to all physics backends. For backend-specific fields, use: + - :class:`isaaclab_physx.sensors.ContactSensorCfg` for PhysX + - :class:`isaaclab_newton.sensors.ContactSensorCfg` for Newton + """ + + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.contact_sensor.contact_sensor:ContactSensor") track_pose: bool = False """Whether to track the pose of the sensor's origin. Defaults to False.""" - track_contact_points: bool = False - """Whether to track the contact point locations. Defaults to False.""" - - track_friction_forces: bool = False - """Whether to track the friction forces at the contact points. Defaults to False.""" - - max_contact_data_count_per_prim: int = 4 - """The maximum number of contacts across all batches of the sensor to keep track of. Default is 4. - - This parameter sets the total maximum counts of the simulation across all bodies and environments. The total number - of contacts allowed is max_contact_data_count_per_prim*num_envs*num_sensor_bodies. - - .. note:: - - If the environment is very contact rich it is suggested to increase this parameter to avoid out of bounds memory - errors and loss of contact data leading to inaccurate measurements. - - """ - track_air_time: bool = False """Whether to track the air/contact time of the bodies (time between contacts). Defaults to False.""" - force_threshold: float = 1.0 + force_threshold: float = 0.0 """The threshold on the norm of the contact force that determines whether two bodies are in collision or not. This value is only used for tracking the mode duration (the time in contact or in air), if :attr:`track_air_time` is True. """ - filter_prim_paths_expr: list[str] = list() - """The list of primitive paths (or expressions) to filter contacts with. Defaults to an empty list, in which case - no filtering is applied. + filter_prim_paths_expr: list[str] | None = None + """A list of body path expressions to filter contacts with. Defaults to None, meaning contacts + with all bodies are reported. The contact sensor allows reporting contacts between the primitive specified with :attr:`prim_path` and other primitives in the scene. For instance, in a scene containing a robot, a ground plane and an object, @@ -62,13 +47,6 @@ class ContactSensorCfg(SensorBaseCfg): will be replaced with the environment namespace. Example: ``{ENV_REGEX_NS}/Object`` will be replaced with ``/World/envs/env_.*/Object``. - - .. attention:: - The reporting of filtered contacts only works when the sensor primitive :attr:`prim_path` corresponds to a - single primitive in that environment. If the sensor primitive corresponds to multiple primitives, the - filtering will not work as expected. Please check :class:`~isaaclab.sensors.contact_sensor.ContactSensor` - for more details. - If track_contact_points is true, then filter_prim_paths_expr cannot be an empty list! """ visualizer_cfg: VisualizationMarkersCfg = CONTACT_SENSOR_MARKER_CFG.replace(prim_path="/Visuals/ContactSensor") diff --git a/source/isaaclab/isaaclab/sensors/frame_transformer/__init__.py b/source/isaaclab/isaaclab/sensors/frame_transformer/__init__.py index 664a91e4eb0..a9827e9b830 100644 --- a/source/isaaclab/isaaclab/sensors/frame_transformer/__init__.py +++ b/source/isaaclab/isaaclab/sensors/frame_transformer/__init__.py @@ -5,10 +5,15 @@ """Sub-module for frame transformer sensor.""" -from .base_frame_transformer import BaseFrameTransformer -from .base_frame_transformer_data import BaseFrameTransformerData -from .frame_transformer import FrameTransformer -from .frame_transformer_cfg import FrameTransformerCfg, OffsetCfg -from .frame_transformer_data import FrameTransformerData +import lazy_loader as lazy -__all__ = ["BaseFrameTransformer", "BaseFrameTransformerData", "FrameTransformer", "FrameTransformerCfg", "FrameTransformerData", "OffsetCfg"] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "base_frame_transformer": ["BaseFrameTransformer"], + "base_frame_transformer_data": ["BaseFrameTransformerData"], + "frame_transformer": ["FrameTransformer"], + "frame_transformer_cfg": ["FrameTransformerCfg", "OffsetCfg"], + "frame_transformer_data": ["FrameTransformerData"], + }, +) diff --git a/source/isaaclab/isaaclab/sensors/frame_transformer/frame_transformer_cfg.py b/source/isaaclab/isaaclab/sensors/frame_transformer/frame_transformer_cfg.py index 389f22c505c..e3281923126 100644 --- a/source/isaaclab/isaaclab/sensors/frame_transformer/frame_transformer_cfg.py +++ b/source/isaaclab/isaaclab/sensors/frame_transformer/frame_transformer_cfg.py @@ -8,10 +8,9 @@ from dataclasses import MISSING from isaaclab.markers.config import FRAME_MARKER_CFG, VisualizationMarkersCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from ..sensor_base_cfg import SensorBaseCfg -from .frame_transformer import FrameTransformer @configclass @@ -54,7 +53,7 @@ class FrameCfg: offset: OffsetCfg = OffsetCfg() """The pose offset from the parent prim frame.""" - class_type: type = FrameTransformer + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.frame_transformer.frame_transformer:FrameTransformer") prim_path: str = MISSING """The prim path of the body to transform from (source frame).""" diff --git a/source/isaaclab/isaaclab/sensors/imu/__init__.py b/source/isaaclab/isaaclab/sensors/imu/__init__.py index 2f372fade3a..eec5be6cc0a 100644 --- a/source/isaaclab/isaaclab/sensors/imu/__init__.py +++ b/source/isaaclab/isaaclab/sensors/imu/__init__.py @@ -7,10 +7,15 @@ Imu Sensor """ -from .base_imu import BaseImu -from .base_imu_data import BaseImuData -from .imu import Imu -from .imu_cfg import ImuCfg -from .imu_data import ImuData +import lazy_loader as lazy -__all__ = ["BaseImu", "BaseImuData", "Imu", "ImuCfg", "ImuData"] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "base_imu": ["BaseImu"], + "base_imu_data": ["BaseImuData"], + "imu": ["Imu"], + "imu_cfg": ["ImuCfg"], + "imu_data": ["ImuData"], + }, +) diff --git a/source/isaaclab/isaaclab/sensors/imu/imu_cfg.py b/source/isaaclab/isaaclab/sensors/imu/imu_cfg.py index 46c0e68cbc7..7045f500fa9 100644 --- a/source/isaaclab/isaaclab/sensors/imu/imu_cfg.py +++ b/source/isaaclab/isaaclab/sensors/imu/imu_cfg.py @@ -7,17 +7,16 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.markers.config import RED_ARROW_X_MARKER_CFG -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from ..sensor_base_cfg import SensorBaseCfg -from .imu import Imu @configclass class ImuCfg(SensorBaseCfg): """Configuration for an Inertial Measurement Unit (IMU) sensor.""" - class_type: type = Imu + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.imu.imu:Imu") @configclass class OffsetCfg: diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py b/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py index 06f479ed2ee..55d5b3085b6 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/__init__.py @@ -15,15 +15,22 @@ the same ray-casting operations as the sensor implementations, but return the results as images. """ -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 -from .ray_caster_cfg import RayCasterCfg -from .ray_caster_data import RayCasterData +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["patterns"], + submod_attrs={ + "multi_mesh_ray_caster": ["MultiMeshRayCaster"], + "multi_mesh_ray_caster_camera": ["MultiMeshRayCasterCamera"], + "multi_mesh_ray_caster_camera_cfg": ["MultiMeshRayCasterCameraCfg"], + "multi_mesh_ray_caster_camera_data": ["MultiMeshRayCasterCameraData"], + "multi_mesh_ray_caster_cfg": ["MultiMeshRayCasterCfg"], + "multi_mesh_ray_caster_data": ["MultiMeshRayCasterData"], + "ray_caster": ["RayCaster"], + "ray_caster_camera": ["RayCasterCamera"], + "ray_caster_camera_cfg": ["RayCasterCameraCfg"], + "ray_caster_cfg": ["RayCasterCfg"], + "ray_caster_data": ["RayCasterData"], + }, +) diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py index 39be0d7ca0d..9bc26764fee 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/multi_mesh_ray_caster.py @@ -15,8 +15,6 @@ import trimesh import warp as wp -import omni.physics.tensors.impl.api as physx - import isaaclab.sim as sim_utils from isaaclab.sim.views import XformPrimView from isaaclab.utils.math import matrix_from_quat, quat_mul @@ -29,6 +27,8 @@ if TYPE_CHECKING: from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg + import omni.physics.tensors.impl.api as physx + # import logger logger = logging.getLogger(__name__) 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 index 45df51ce6d8..f51be978c0c 100644 --- 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 @@ -7,9 +7,8 @@ import logging -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass -from .multi_mesh_ray_caster_camera import MultiMeshRayCasterCamera from .multi_mesh_ray_caster_cfg import MultiMeshRayCasterCfg from .ray_caster_camera_cfg import RayCasterCameraCfg @@ -21,7 +20,7 @@ class MultiMeshRayCasterCameraCfg(RayCasterCameraCfg, MultiMeshRayCasterCfg): """Configuration for the multi-mesh ray-cast camera sensor.""" - class_type: type = MultiMeshRayCasterCamera + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.ray_caster.multi_mesh_ray_caster_camera:MultiMeshRayCasterCamera") def __post_init__(self): super().__post_init__() 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 index f5393920162..b569f5b6c09 100644 --- 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 @@ -8,9 +8,8 @@ from dataclasses import MISSING -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass -from .multi_mesh_ray_caster import MultiMeshRayCaster from .ray_caster_cfg import RayCasterCfg @@ -49,7 +48,7 @@ class RaycastTargetCfg: Not tracking the mesh transformations is recommended when the meshes are static to increase performance. """ - class_type: type = MultiMeshRayCaster + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.ray_caster.multi_mesh_ray_caster:MultiMeshRayCaster") mesh_prim_paths: list[str | RaycastTargetCfg] = MISSING """The list of mesh primitive paths to ray cast against. diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_cast_utils.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_cast_utils.py index ac503b28bf5..73949b2705b 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_cast_utils.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_cast_utils.py @@ -6,14 +6,15 @@ """Utility functions for ray-cast sensors.""" from __future__ import annotations - +from typing import TYPE_CHECKING import torch import warp as wp -import omni.physics.tensors.impl.api as physx - from isaaclab.sim.views import XformPrimView +if TYPE_CHECKING: + import omni.physics.tensors.impl.api as physx + def obtain_world_pose_from_view( physx_view: XformPrimView | physx.ArticulationView | physx.RigidBodyView, diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera_cfg.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera_cfg.py index d8e43381708..c900e29d7a0 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera_cfg.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_camera_cfg.py @@ -8,10 +8,9 @@ from dataclasses import MISSING from typing import Literal -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from .patterns import PinholeCameraPatternCfg -from .ray_caster_camera import RayCasterCamera from .ray_caster_cfg import RayCasterCfg @@ -39,7 +38,7 @@ class OffsetCfg: """ - class_type: type = RayCasterCamera + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.ray_caster.ray_caster_camera:RayCasterCamera") offset: OffsetCfg = OffsetCfg() """The offset pose of the sensor's frame from the sensor's parent frame. Defaults to identity.""" diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py index 23981b98ee9..88ea98f2cf4 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/ray_caster_cfg.py @@ -12,11 +12,10 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.markers.config import RAY_CASTER_MARKER_CFG -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from ..sensor_base_cfg import SensorBaseCfg from .patterns.patterns_cfg import PatternBaseCfg -from .ray_caster import RayCaster @configclass @@ -32,7 +31,7 @@ class OffsetCfg: rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) """Quaternion rotation (x, y, z, w) w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0, 1.0).""" - class_type: type = RayCaster + class_type: type | DeferredClass = DeferredClass("isaaclab.sensors.ray_caster.ray_caster:RayCaster") mesh_prim_paths: list[str] = MISSING """The list of mesh primitive paths to ray cast against. diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py index 9bead099e8c..f9318f748d7 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base.py @@ -19,13 +19,12 @@ from typing import TYPE_CHECKING, Any import torch -from isaaclab_physx.physics import IsaacEvents, PhysxManager - -import omni.kit.app +import warp as wp import isaaclab.sim as sim_utils from isaaclab.physics import PhysicsEvent from isaaclab.sim.utils.stage import get_current_stage +from isaaclab.utils.version import has_kit if TYPE_CHECKING: from .sensor_base_cfg import SensorBaseCfg @@ -41,6 +40,11 @@ class SensorBase(ABC): The sensor is updated at the specified update period. If the update period is zero, then the sensor is updated at every simulation step. + + Backend-specific implementations should inherit from this class and implement: + - :meth:`_initialize_impl` - Initialize sensor handles and buffers + - :meth:`_update_outdated_buffers` - Update sensor data for outdated environments + - Either :meth:`_update_buffers_impl_index` (PhysX) or :meth:`_update_buffers_impl_mask` (Newton) """ def __init__(self, cfg: SensorBaseCfg): @@ -149,10 +153,12 @@ def set_debug_vis(self, debug_vis: bool) -> bool: self._set_debug_vis_impl(debug_vis) # toggle debug visualization flag self._is_visualizing = debug_vis - # toggle debug visualization handles + # toggle debug visualization handles (Kit/omni only for PhysX backend) if debug_vis: # create a subscriber for the post update event if it doesn't exist - if self._debug_vis_handle is None: + if self._debug_vis_handle is None and has_kit(): + import omni.kit.app + app_interface = omni.kit.app.get_app_interface() self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) @@ -165,28 +171,28 @@ def set_debug_vis(self, debug_vis: bool) -> bool: # return success return True - def reset(self, env_ids: Sequence[int] | None = None): + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): """Resets the sensor internals. Args: - env_ids: The sensor ids to reset. Defaults to None. + env_ids: The sensor ids to reset. Defaults to None (all instances). + env_mask: The sensor mask to reset. Defaults to None (all instances). """ - # Resolve sensor ids if env_ids is None: env_ids = slice(None) - # Reset the timestamp for the sensors self._timestamp[env_ids] = 0.0 self._timestamp_last_update[env_ids] = 0.0 - # Set all reset sensors to outdated so that they are updated when data is called the next time. self._is_outdated[env_ids] = True def update(self, dt: float, force_recompute: bool = False): - # Update the timestamp for the sensors + """Updates the sensor data. + + Args: + dt: The time step for the update. + force_recompute: Whether to force recomputation of sensor data. + """ self._timestamp += dt self._is_outdated |= self._timestamp - self._timestamp_last_update + 1e-6 >= self.cfg.update_period - # Update the buffers - # TODO (from @mayank): Why is there a history length here when it doesn't mean anything in the sensor base?!? - # It is only for the contact sensor but there we should redefine the update function IMO. if force_recompute or self._is_visualizing or (self.cfg.history_length > 0): self._update_outdated_buffers() @@ -194,39 +200,48 @@ def update(self, dt: float, force_recompute: bool = False): Implementation specific. """ + _device: str + """Memory device for computation.""" + _backend: str + """Simulation backend (e.g., "newton", "physx").""" + _num_envs: int + """Number of environments.""" + @abstractmethod def _initialize_impl(self): - """Initializes the sensor-related handles and internal buffers.""" - # Obtain Simulation Context + """Initializes the sensor-related handles and internal buffers. + + Subclasses should call ``super()._initialize_impl()`` first to obtain + device, backend, num_envs, timestamps, and parent prims from the + simulation context, then add sensor-specific setup. + """ sim = sim_utils.SimulationContext.instance() if sim is None: raise RuntimeError("Simulation Context is not initialized!") - # Obtain device and backend self._device = sim.device self._backend = sim.backend self._sim_physics_dt = sim.get_physics_dt() - # Count number of environments env_prim_path_expr = self.cfg.prim_path.rsplit("/", 1)[0] self._parent_prims = sim_utils.find_matching_prims(env_prim_path_expr) self._num_envs = len(self._parent_prims) - # Boolean tensor indicating whether the sensor data has to be refreshed self._is_outdated = torch.ones(self._num_envs, dtype=torch.bool, device=self._device) - # Current timestamp (in seconds) self._timestamp = torch.zeros(self._num_envs, device=self._device) - # Timestamp from last update self._timestamp_last_update = torch.zeros_like(self._timestamp) - # Initialize debug visualization handle - if self._debug_vis_handle is None: - # set initial state of debug visualization - self.set_debug_vis(self.cfg.debug_vis) + def _update_outdated_buffers(self): + """Fills the sensor data for the outdated sensors.""" + outdated_env_ids = self._is_outdated.nonzero().squeeze(-1) + if len(outdated_env_ids) > 0: + self._update_buffers_impl(outdated_env_ids) + self._timestamp_last_update[outdated_env_ids] = self._timestamp[outdated_env_ids] + self._is_outdated[outdated_env_ids] = False - @abstractmethod def _update_buffers_impl(self, env_ids: Sequence[int]): """Fills the sensor data for provided environment ids. - This function does not perform any time-based checks and directly fills the data into the - data container. + Called by the default :meth:`_update_outdated_buffers`. Subclasses that + override ``_update_outdated_buffers`` (e.g. mask-based backends) do not + need to implement this method. Args: env_ids: The indices of the sensors that are ready to capture. @@ -254,52 +269,50 @@ def _debug_vis_callback(self, event): """ def _register_callbacks(self): - """Registers physics lifecycle and prim deletion callbacks.""" + """Registers physics lifecycle callbacks via the current backend's physics manager.""" + physics_mgr_cls = sim_utils.SimulationContext.instance().physics_manager - # register simulator callbacks (with weakref safety to avoid crashes on deletion) def safe_callback(callback_name, event, obj_ref): """Safely invoke a callback on a weakly-referenced object, ignoring ReferenceError if deleted.""" try: obj = obj_ref getattr(obj, callback_name)(event) except ReferenceError: - # Object has been deleted; ignore. pass - # note: use weakref on callbacks to ensure that this object can be deleted when its destructor is called. obj_ref = weakref.proxy(self) - # Register PHYSICS_READY callback for initialization (order=10 for lower priority) - self._initialize_handle = PhysxManager.register_callback( + # Backend-agnostic: PHYSICS_READY (init) and STOP (invalidate) + self._initialize_handle = physics_mgr_cls.register_callback( lambda payload, obj_ref=obj_ref: safe_callback("_initialize_callback", payload, obj_ref), PhysicsEvent.PHYSICS_READY, order=10, ) - # Register TIMELINE_STOP callback for invalidation (PhysX-specific) - self._invalidate_initialize_handle = PhysxManager.register_callback( - lambda event, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", event, obj_ref), - IsaacEvents.TIMELINE_STOP, + self._invalidate_initialize_handle = physics_mgr_cls.register_callback( + lambda payload, obj_ref=obj_ref: safe_callback("_invalidate_initialize_callback", payload, obj_ref), + PhysicsEvent.STOP, order=10, ) - # Register PRIM_DELETION callback (PhysX-specific) - self._prim_deletion_handle = PhysxManager.register_callback( - lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref), - IsaacEvents.PRIM_DELETION, - ) + # Optional: prim deletion (only supported by PhysX backend) + self._prim_deletion_handle = None + physics_backend = physics_mgr_cls.__name__ + if "Physx" in physics_backend: + from isaaclab_physx.physics import IsaacEvents + + self._prim_deletion_handle = physics_mgr_cls.register_callback( + lambda event, obj_ref=obj_ref: safe_callback("_on_prim_deletion", event, obj_ref), + IsaacEvents.PRIM_DELETION, + ) def _initialize_callback(self, event): """Initializes the scene elements. .. note:: - PhysX handles are only enabled once the simulator starts playing. Hence, this function needs to be - called whenever the simulator "plays" from a "stop" state. + Physics handles are only valid once the simulation is ready. This callback runs when + :attr:`PhysicsEvent.PHYSICS_READY` is dispatched by the current backend. """ if not self._is_initialized: - try: - self._initialize_impl() - except Exception as e: - # Store exception to be raised after callback completes - PhysxManager.store_callback_exception(e) + self._initialize_impl() self._is_initialized = True def _invalidate_initialize_callback(self, event): @@ -310,15 +323,12 @@ def _invalidate_initialize_callback(self, event): self._debug_vis_handle = None def _on_prim_deletion(self, event) -> None: - """Invalidates and deletes the callbacks when the prim is deleted. - - Args: - event: The prim deletion event containing the prim path in payload. + """Invalidates and clears callbacks when the prim is deleted. - Note: - This function is called when the prim is deleted. + Only used when the backend supports prim deletion events (e.g. PhysX). """ - prim_path = event.payload["prim_path"] + payload = getattr(event, "payload", event) if not isinstance(event, dict) else event + prim_path = payload.get("prim_path", "") if isinstance(payload, dict) else "" if prim_path == "/": self._clear_callbacks() return @@ -329,32 +339,16 @@ def _on_prim_deletion(self, event) -> None: self._clear_callbacks() def _clear_callbacks(self) -> None: - """Clears the callbacks.""" - if self._initialize_handle is not None: + """Clears all registered callbacks.""" + if getattr(self, "_initialize_handle", None) is not None: self._initialize_handle.deregister() self._initialize_handle = None - if self._invalidate_initialize_handle is not None: + if getattr(self, "_invalidate_initialize_handle", None) is not None: self._invalidate_initialize_handle.deregister() self._invalidate_initialize_handle = None - if self._prim_deletion_handle is not None: + if getattr(self, "_prim_deletion_handle", None) is not None: self._prim_deletion_handle.deregister() self._prim_deletion_handle = None - # Clear debug visualization - if self._debug_vis_handle is not None: + if getattr(self, "_debug_vis_handle", None) is not None: self._debug_vis_handle.unsubscribe() self._debug_vis_handle = None - - """ - Helper functions. - """ - - def _update_outdated_buffers(self): - """Fills the sensor data for the outdated sensors.""" - outdated_env_ids = self._is_outdated.nonzero().squeeze(-1) - if len(outdated_env_ids) > 0: - # obtain new data - self._update_buffers_impl(outdated_env_ids) - # update the timestamp from last update - self._timestamp_last_update[outdated_env_ids] = self._timestamp[outdated_env_ids] - # set outdated flag to false for the updated sensors - self._is_outdated[outdated_env_ids] = False diff --git a/source/isaaclab/isaaclab/sensors/sensor_base_cfg.py b/source/isaaclab/isaaclab/sensors/sensor_base_cfg.py index 85875b2e499..1fa94dab26a 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base_cfg.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base_cfg.py @@ -3,11 +3,15 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + from dataclasses import MISSING +from typing import TYPE_CHECKING from isaaclab.utils import configclass -from .sensor_base import SensorBase +if TYPE_CHECKING: + from .sensor_base import SensorBase @configclass diff --git a/source/isaaclab/isaaclab/sim/__init__.py b/source/isaaclab/isaaclab/sim/__init__.py index 9cb5fdc13aa..7623f317b07 100644 --- a/source/isaaclab/isaaclab/sim/__init__.py +++ b/source/isaaclab/isaaclab/sim/__init__.py @@ -26,32 +26,34 @@ """ -import warnings - -from .converters import * # noqa: F401, F403 -from .schemas import * # noqa: F401, F403 -from .simulation_cfg import RenderCfg, SimulationCfg # noqa: F401, F403 -from .simulation_context import SimulationContext, build_simulation_context # noqa: F401, F403 -from .spawners import * # noqa: F401, F403 -from .utils import * # noqa: F401, F403 -from .views import * # noqa: F401, F403 - -# Deprecated alias for PhysxCfg -> PhysxCfg -# This supports old code that uses `from isaaclab.sim import PhysxCfg` -try: - from isaaclab_physx.physics import PhysxCfg as _PhysxCfg - - class PhysxCfg(_PhysxCfg): - """DEPRECATED: Use PhysxCfg from isaaclab_physx.physics instead.""" - - def __init__(self, *args, **kwargs): - warnings.warn( - "PhysxCfg is deprecated. Use PhysxCfg from isaaclab_physx.physics instead.", - DeprecationWarning, - stacklevel=2, - ) - super().__init__(*args, **kwargs) - -except ImportError: - # isaaclab_physx not installed - PhysxCfg = None # type: ignore +import lazy_loader as lazy + +from .simulation_cfg import RenderCfg, SimulationCfg + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["converters", "schemas", "spawners", "utils", "views"], + submod_attrs={ + "simulation_context": ["SimulationContext", "build_simulation_context"], + }, +) +__all__ += ["RenderCfg", "SimulationCfg"] + +_lazy_getattr = __getattr__ +_SUBPACKAGES = ("converters", "schemas", "spawners", "utils", "views") + + +def __getattr__(name): + try: + return _lazy_getattr(name) + except AttributeError: + pass + import importlib + + for subpkg in _SUBPACKAGES: + try: + submod = importlib.import_module(f"{__name__}.{subpkg}") + return getattr(submod, name) + except (ImportError, AttributeError): + continue + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/source/isaaclab/isaaclab/sim/backend_detection.py b/source/isaaclab/isaaclab/sim/backend_detection.py new file mode 100644 index 00000000000..f4f029bccac --- /dev/null +++ b/source/isaaclab/isaaclab/sim/backend_detection.py @@ -0,0 +1,35 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Utilities to detect physics backend from config without importing Kit or backends.""" + + +def physics_backend_requires_kit(env_cfg) -> bool: + """Return True if the env's physics backend requires Isaac Sim Kit (e.g. PhysX), False otherwise. + + Used to decide whether to launch AppLauncher before creating the environment. + Newton backend does not require Kit; PhysX and default (None) do. + + This only inspects the config and type name of ``env_cfg.sim.physics``, so it is safe to call + before any Kit/Omniverse imports. + + Args: + env_cfg: Environment config with a ``sim`` attribute that has a ``physics`` attribute + (e.g. ManagerBasedRLEnvCfg, DirectRLEnvCfg). + + Returns: + True if Kit must be launched (PhysX or physics is None); False if Newton (no Kit). + """ + sim = getattr(env_cfg, "sim", None) + if sim is None: + return True + physics = getattr(sim, "physics", None) + if physics is None: + return True # SimulationContext defaults to PhysxCfg when None + name = type(physics).__name__ + # Newton backend does not require Isaac Sim Kit + if name == "NewtonCfg": + return False + return True diff --git a/source/isaaclab/isaaclab/sim/converters/__init__.py b/source/isaaclab/isaaclab/sim/converters/__init__.py index 7503c53bdd8..29d1b386a86 100644 --- a/source/isaaclab/isaaclab/sim/converters/__init__.py +++ b/source/isaaclab/isaaclab/sim/converters/__init__.py @@ -16,11 +16,20 @@ """ -from .asset_converter_base import AssetConverterBase +import lazy_loader as lazy + from .asset_converter_base_cfg import AssetConverterBaseCfg -from .mesh_converter import MeshConverter from .mesh_converter_cfg import MeshConverterCfg -from .mjcf_converter import MjcfConverter from .mjcf_converter_cfg import MjcfConverterCfg -from .urdf_converter import UrdfConverter from .urdf_converter_cfg import UrdfConverterCfg + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "asset_converter_base": ["AssetConverterBase"], + "mesh_converter": ["MeshConverter"], + "mjcf_converter": ["MjcfConverter"], + "urdf_converter": ["UrdfConverter"], + }, +) +__all__ += ["AssetConverterBaseCfg", "MeshConverterCfg", "MjcfConverterCfg", "UrdfConverterCfg"] diff --git a/source/isaaclab/isaaclab/sim/converters/mesh_converter.py b/source/isaaclab/isaaclab/sim/converters/mesh_converter.py index 6d2df6285e0..3392b31dea2 100644 --- a/source/isaaclab/isaaclab/sim/converters/mesh_converter.py +++ b/source/isaaclab/isaaclab/sim/converters/mesh_converter.py @@ -7,9 +7,7 @@ import logging import os -import omni -import omni.kit.commands -from isaacsim.core.utils.extensions import enable_extension + from pxr import Gf, Tf, Usd, UsdGeom, UsdPhysics, UsdUtils from isaaclab.sim.converters.asset_converter_base import AssetConverterBase @@ -137,6 +135,7 @@ def _convert_asset(self, cfg: MeshConverterCfg): # Delete the old Xform and make the new Xform the default prim stage.SetDefaultPrim(xform_prim) # Apply default Xform rotation to mesh -> enable to set rotation and scale + import omni.kit.commands omni.kit.commands.execute( "CreateDefaultXformOnPrimCommand", prim_path=xform_prim.GetPath(), @@ -217,6 +216,7 @@ async def _convert_mesh_to_usd(in_file: str, out_file: str, load_materials: bool Returns: True if the conversion succeeds. """ + from isaacsim.core.utils.extensions import enable_extension enable_extension("omni.kit.asset_converter") import omni.kit.asset_converter diff --git a/source/isaaclab/isaaclab/sim/converters/urdf_converter.py b/source/isaaclab/isaaclab/sim/converters/urdf_converter.py index 5d9c9b7ce2e..a7ca7c7280e 100644 --- a/source/isaaclab/isaaclab/sim/converters/urdf_converter.py +++ b/source/isaaclab/isaaclab/sim/converters/urdf_converter.py @@ -11,10 +11,7 @@ from packaging.version import Version -import omni.kit.app -import omni.kit.commands - -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit from .asset_converter_base import AssetConverterBase from .urdf_converter_cfg import UrdfConverterCfg @@ -61,10 +58,12 @@ def __init__(self, cfg: UrdfConverterCfg): Args: cfg: The configuration instance for URDF to USD conversion. """ + import omni.kit.app + # switch to older version of the URDF importer extension manager = omni.kit.app.get_app().get_extension_manager() - if get_isaac_sim_version() == Version("5.1"): + if has_kit() and get_isaac_sim_version() == Version("5.1"): if not manager.is_extension_enabled("isaacsim.asset.importer.urdf-2.4.31"): manager.set_extension_enabled_immediate("isaacsim.asset.importer.urdf-2.4.31", True) else: @@ -87,6 +86,7 @@ def _convert_asset(self, cfg: UrdfConverterCfg): Args: cfg: The URDF conversion configuration. """ + import omni.kit.commands import_config = self._get_urdf_import_config() # parse URDF file @@ -124,6 +124,8 @@ def _get_urdf_import_config(self) -> isaacsim.asset.importer.urdf._urdf.ImportCo Returns: The constructed ``ImportConfig`` object containing the desired settings. """ + import omni.kit.commands + # create a new import config _, import_config = omni.kit.commands.execute("URDFCreateImportConfig") diff --git a/source/isaaclab/isaaclab/sim/scene_data_providers/__init__.py b/source/isaaclab/isaaclab/sim/scene_data_providers/__init__.py index 1e8e7c0ac3f..73a6e641c9e 100644 --- a/source/isaaclab/isaaclab/sim/scene_data_providers/__init__.py +++ b/source/isaaclab/isaaclab/sim/scene_data_providers/__init__.py @@ -5,10 +5,12 @@ """Scene data providers for visualizers and renderers.""" +from .newton_scene_data_provider import NewtonSceneDataProvider from .physx_scene_data_provider import PhysxSceneDataProvider from .scene_data_provider import SceneDataProvider __all__ = [ "SceneDataProvider", "PhysxSceneDataProvider", + "NewtonSceneDataProvider", ] diff --git a/source/isaaclab/isaaclab/sim/scene_data_providers/newton_scene_data_provider.py b/source/isaaclab/isaaclab/sim/scene_data_providers/newton_scene_data_provider.py index 94092ae974d..413ca121728 100644 --- a/source/isaaclab/isaaclab/sim/scene_data_providers/newton_scene_data_provider.py +++ b/source/isaaclab/isaaclab/sim/scene_data_providers/newton_scene_data_provider.py @@ -13,24 +13,50 @@ class NewtonSceneDataProvider: """TODO: implement when the newton physics backend is added.""" - def __init__(self, visualizer_cfgs: list[Any] | None) -> None: + def __init__(self, visualizer_cfgs: list[Any] | None, simulation_context) -> None: + self._simulation_context = simulation_context self._metadata = {"physics_backend": "newton"} def update(self, env_ids: list[int] | None = None) -> None: - """TODO: implement when the newton physics backend is added.""" + """Refresh any cached scene data. + + For Newton backend, this is a no-op as the model and state are managed by NewtonManager. + """ pass def get_newton_model(self) -> Any | None: - """TODO: implement when the newton physics backend is added.""" - return None + """Return Newton model from NewtonManager.""" + from isaaclab_newton.physics import NewtonManager + + return NewtonManager.get_model() def get_newton_state(self, env_ids: list[int] | None = None) -> Any | None: - """TODO: implement when the newton physics backend is added.""" - return None + """Return Newton state from NewtonManager. + + Args: + env_ids: Optional list of environment IDs. Currently not supported for filtering. + Returns the full state for all environments. + + Returns: + The current Newton state (state_0) from NewtonManager. + """ + from isaaclab_newton.physics import NewtonManager + + # For now, return state_0 (current state) for all environments + # TODO: Implement env_ids filtering if needed + return NewtonManager.get_state_0() + + def get_model(self) -> Any | None: + """Return Newton model (alias for get_newton_model for visualizer compatibility).""" + return self.get_newton_model() + + def get_state(self, env_ids: list[int] | None = None) -> Any | None: + """Return Newton state (alias for get_newton_state for visualizer compatibility).""" + return self.get_newton_state(env_ids) def get_usd_stage(self) -> Any | None: - """TODO: implement when the newton physics backend is added.""" - return None + """Return the USD stage handle from simulation context.""" + return getattr(self._simulation_context, "stage", None) def get_metadata(self) -> dict[str, Any]: return dict(self._metadata) diff --git a/source/isaaclab/isaaclab/sim/scene_data_providers/physx_scene_data_provider.py b/source/isaaclab/isaaclab/sim/scene_data_providers/physx_scene_data_provider.py index aab6371b464..a84eb843aad 100644 --- a/source/isaaclab/isaaclab/sim/scene_data_providers/physx_scene_data_provider.py +++ b/source/isaaclab/isaaclab/sim/scene_data_providers/physx_scene_data_provider.py @@ -65,11 +65,9 @@ def _determine_num_envs_in_scene(self) -> int: return max_env_id + 1 if max_env_id >= 0 else 0 def __init__(self, visualizer_cfgs: list[Any] | None, stage, simulation_context) -> None: - from isaacsim.core.simulation_manager import SimulationManager - self._simulation_context = simulation_context self._stage = stage - self._physics_sim_view = SimulationManager.get_physics_sim_view() + self._physics_sim_view = simulation_context.physics_sim_view self._rigid_body_view = None self._articulation_view = None self._xform_views: dict[str, Any] = {} @@ -152,8 +150,11 @@ def _build_newton_model_from_usd(self) -> None: self._newton_state = self._newton_model.state() # Extract scene structure from Newton model (single source of truth) - self._rigid_body_paths = list(self._newton_model.body_label) - self._articulation_paths = list(self._newton_model.articulation_label) + # Newton >= 0.x renamed body_label -> body_key, articulation_label -> articulation_key + body_attr = "body_key" if hasattr(self._newton_model, "body_key") else "body_label" + artic_attr = "articulation_key" if hasattr(self._newton_model, "articulation_key") else "articulation_label" + self._rigid_body_paths = list(getattr(self._newton_model, body_attr)) + self._articulation_paths = list(getattr(self._newton_model, artic_attr)) self._xform_views.clear() self._view_body_index_map = {} diff --git a/source/isaaclab/isaaclab/sim/schemas/__init__.py b/source/isaaclab/isaaclab/sim/schemas/__init__.py index c8402fdb13c..63ae5035574 100644 --- a/source/isaaclab/isaaclab/sim/schemas/__init__.py +++ b/source/isaaclab/isaaclab/sim/schemas/__init__.py @@ -32,27 +32,8 @@ """ -from .schemas import ( - MESH_APPROXIMATION_TOKENS, - PHYSX_MESH_COLLISION_CFGS, - USD_MESH_COLLISION_CFGS, - activate_contact_sensors, - define_articulation_root_properties, - define_collision_properties, - define_deformable_body_properties, - define_mass_properties, - define_mesh_collision_properties, - define_rigid_body_properties, - modify_articulation_root_properties, - modify_collision_properties, - modify_deformable_body_properties, - modify_fixed_tendon_properties, - modify_joint_drive_properties, - modify_mass_properties, - modify_mesh_collision_properties, - modify_rigid_body_properties, - modify_spatial_tendon_properties, -) +import lazy_loader as lazy + from .schemas_cfg import ( ArticulationRootPropertiesCfg, BoundingCubePropertiesCfg, @@ -72,56 +53,47 @@ TriangleMeshSimplificationPropertiesCfg, ) -__all__ = [ - # articulation root +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "schemas": [ + "MESH_APPROXIMATION_TOKENS", + "PHYSX_MESH_COLLISION_CFGS", + "USD_MESH_COLLISION_CFGS", + "activate_contact_sensors", + "define_articulation_root_properties", + "define_collision_properties", + "define_deformable_body_properties", + "define_mass_properties", + "define_mesh_collision_properties", + "define_rigid_body_properties", + "modify_articulation_root_properties", + "modify_collision_properties", + "modify_deformable_body_properties", + "modify_fixed_tendon_properties", + "modify_joint_drive_properties", + "modify_mass_properties", + "modify_mesh_collision_properties", + "modify_rigid_body_properties", + "modify_spatial_tendon_properties", + ], + }, +) +__all__ += [ "ArticulationRootPropertiesCfg", - "define_articulation_root_properties", - "modify_articulation_root_properties", - # rigid bodies - "RigidBodyPropertiesCfg", - "define_rigid_body_properties", - "modify_rigid_body_properties", - "activate_contact_sensors", - # colliders + "BoundingCubePropertiesCfg", + "BoundingSpherePropertiesCfg", "CollisionPropertiesCfg", - "define_collision_properties", - "modify_collision_properties", - # deformables + "ConvexDecompositionPropertiesCfg", + "ConvexHullPropertiesCfg", "DeformableBodyPropertiesCfg", - "define_deformable_body_properties", - "modify_deformable_body_properties", - # joints + "FixedTendonPropertiesCfg", "JointDrivePropertiesCfg", - "modify_joint_drive_properties", - # mass "MassPropertiesCfg", - "define_mass_properties", - "modify_mass_properties", - # mesh colliders "MeshCollisionPropertiesCfg", - "define_mesh_collision_properties", - "modify_mesh_collision_properties", - # bounding cube - "BoundingCubePropertiesCfg", - # bounding sphere - "BoundingSpherePropertiesCfg", - # convex decomposition - "ConvexDecompositionPropertiesCfg", - # convex hull - "ConvexHullPropertiesCfg", - # sdf mesh + "RigidBodyPropertiesCfg", "SDFMeshPropertiesCfg", - # triangle mesh + "SpatialTendonPropertiesCfg", "TriangleMeshPropertiesCfg", - # triangle mesh simplification "TriangleMeshSimplificationPropertiesCfg", - # tendons - "FixedTendonPropertiesCfg", - "SpatialTendonPropertiesCfg", - "modify_fixed_tendon_properties", - "modify_spatial_tendon_properties", - # Constants for configs that use PhysX vs USD API - "PHYSX_MESH_COLLISION_CFGS", - "USD_MESH_COLLISION_CFGS", - "MESH_APPROXIMATION_TOKENS", ] diff --git a/source/isaaclab/isaaclab/sim/schemas/schemas.py b/source/isaaclab/isaaclab/sim/schemas/schemas.py index 04791e5b8ec..97d7e08cd2a 100644 --- a/source/isaaclab/isaaclab/sim/schemas/schemas.py +++ b/source/isaaclab/isaaclab/sim/schemas/schemas.py @@ -10,9 +10,7 @@ import math from typing import Any -import omni.physx.scripts.utils as physx_utils -from omni.physx.scripts import deformableUtils as deformable_utils -from pxr import Usd, UsdPhysics +from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics from isaaclab.sim.utils.stage import get_current_stage from isaaclab.utils.string import to_camel_case @@ -199,7 +197,7 @@ def modify_articulation_root_properties( ) # create a fixed joint between the root link and the world frame - physx_utils.createJoint(stage=stage, joint_type="Fixed", from_prim=None, to_prim=articulation_prim) + create_joint(stage=stage, joint_type="Fixed", from_prim=None, to_prim=articulation_prim) # Having a fixed joint on a rigid body is not treated as "fixed base articulation". # instead, it is treated as a part of the maximal coordinate tree. @@ -576,7 +574,6 @@ def activate_contact_sensors(prim_path: str, threshold: float = 0.0, stage: Usd. safe_set_attribute_on_usd_prim(child_prim, "physxRigidBody:sleepThreshold", 0.0, camel_case=False) # add contact report API with threshold of zero if "PhysxContactReportAPI" not in child_applied: - logger.debug(f"Adding contact report API to prim: '{child_prim.GetPrimPath()}'") child_prim.AddAppliedSchema("PhysxContactReportAPI") safe_set_attribute_on_usd_prim(child_prim, "physxContactReport:threshold", threshold, camel_case=False) # increment number of contact sensors @@ -971,6 +968,7 @@ def modify_deformable_body_properties( "self_collision_filter_distance", ] } + from omni.physx.scripts import deformableUtils as deformable_utils status = deformable_utils.add_physx_deformable_body(stage, prim_path=prim_path, **attr_kwargs) # check if the deformable body was successfully added if not status: @@ -1158,3 +1156,115 @@ def modify_mesh_collision_properties( # success return True + + +""" +Joint creation utilities. +""" + +MAX_FLOAT = 3.40282347e38 + + +def create_unused_path(stage: Usd.Stage, base_path: str, path: str) -> str: + """Create an unused path for a joint.""" + if stage.GetPrimAtPath(base_path + "/" + path).IsValid(): + uniquifier = 0 + while stage.GetPrimAtPath(base_path + "/" + path + str(uniquifier)).IsValid(): + uniquifier += 1 + path = path + str(uniquifier) + return path + + +def create_joint( + stage: Usd.Stage, + joint_type: str, + from_prim: Usd.Prim | None, + to_prim: Usd.Prim | None, + joint_name: str | None = None, + joint_base_path: str | None = None, +) -> Usd.Prim: + """Create a joint between two prims.""" + if to_prim is None: + to_prim = from_prim + from_prim = None + + from_path = from_prim.GetPath().pathString if from_prim is not None and from_prim.IsValid() else "" + to_path = to_prim.GetPath().pathString if to_prim is not None and to_prim.IsValid() else "" + single_selection = from_path == "" or to_path == "" + + if joint_base_path is None: + joint_base_path = to_path + else: + if not stage.GetPrimAtPath(joint_base_path): + stage.DefinePrim(joint_base_path) + + base_prim = stage.GetPrimAtPath(joint_base_path) + while base_prim != stage.GetPseudoRoot(): + if base_prim.IsInPrototype() or base_prim.IsInstanceProxy() or base_prim.IsInstanceable(): + base_prim = base_prim.GetParent() + else: + break + joint_base_path = str(base_prim.GetPrimPath()) + if joint_base_path == "/": + joint_base_path = "" + if joint_name is None: + joint_name = "/" + create_unused_path(stage, joint_base_path, joint_type + "Joint") + else: + joint_name = "/" + create_unused_path(stage, joint_base_path, joint_name) + joint_path = joint_base_path + joint_name + + if joint_type == "Fixed": + component = UsdPhysics.FixedJoint.Define(stage, joint_path) + elif joint_type == "Revolute": + component = UsdPhysics.RevoluteJoint.Define(stage, joint_path) + component.CreateAxisAttr("X") + elif joint_type == "Prismatic": + component = UsdPhysics.PrismaticJoint.Define(stage, joint_path) + component.CreateAxisAttr("X") + elif joint_type == "Spherical": + component = UsdPhysics.SphericalJoint.Define(stage, joint_path) + component.CreateAxisAttr("X") + elif joint_type == "Distance": + component = UsdPhysics.DistanceJoint.Define(stage, joint_path) + component.CreateMinDistanceAttr(0.0) + component.CreateMaxDistanceAttr(0.0) + else: + component = UsdPhysics.Joint.Define(stage, joint_path) + prim = component.GetPrim() + for limit_name in ["transX", "transY", "transZ", "rotX", "rotY", "rotZ"]: + limit_api = UsdPhysics.LimitAPI.Apply(prim, limit_name) + limit_api.CreateLowAttr(1.0) + limit_api.CreateHighAttr(-1.0) + + xfCache = UsdGeom.XformCache() + + if not single_selection: + to_pose = xfCache.GetLocalToWorldTransform(to_prim) + from_pose = xfCache.GetLocalToWorldTransform(from_prim) + rel_pose = to_pose * from_pose.GetInverse() + rel_pose = rel_pose.RemoveScaleShear() + pos1 = Gf.Vec3f(rel_pose.ExtractTranslation()) + rot1 = Gf.Quatf(rel_pose.ExtractRotationQuat()) + + component.CreateBody0Rel().SetTargets([Sdf.Path(from_path)]) + component.CreateBody1Rel().SetTargets([Sdf.Path(to_path)]) + component.CreateLocalPos0Attr().Set(pos1) + component.CreateLocalRot0Attr().Set(rot1) + component.CreateLocalPos1Attr().Set(Gf.Vec3f(0.0)) + component.CreateLocalRot1Attr().Set(Gf.Quatf(1.0)) + else: + to_pose = xfCache.GetLocalToWorldTransform(to_prim) + to_pose = to_pose.RemoveScaleShear() + pos1 = Gf.Vec3f(to_pose.ExtractTranslation()) + rot1 = Gf.Quatf(to_pose.ExtractRotationQuat()) + + component.CreateBody1Rel().SetTargets([Sdf.Path(to_path)]) + component.CreateLocalPos0Attr().Set(pos1) + component.CreateLocalRot0Attr().Set(rot1) + component.CreateLocalPos1Attr().Set(Gf.Vec3f(0.0)) + component.CreateLocalRot1Attr().Set(Gf.Quatf(1.0)) + + component.CreateBreakForceAttr().Set(MAX_FLOAT) + component.CreateBreakTorqueAttr().Set(MAX_FLOAT) + + return stage.GetPrimAtPath(joint_base_path + joint_name) diff --git a/source/isaaclab/isaaclab/sim/simulation_context.py b/source/isaaclab/isaaclab/sim/simulation_context.py index 5dc9d13e663..6a6c005882b 100644 --- a/source/isaaclab/isaaclab/sim/simulation_context.py +++ b/source/isaaclab/isaaclab/sim/simulation_context.py @@ -27,7 +27,6 @@ from .scene_data_providers import SceneDataProvider from .simulation_cfg import SimulationCfg -from .spawners import DomeLightCfg, GroundPlaneCfg logger = logging.getLogger(__name__) @@ -143,7 +142,14 @@ def __init__(self, cfg: SimulationCfg | None = None): self.cfg.physics = PhysxCfg() self._physics = self.cfg.physics - self.physics_manager: type[PhysicsManager] = self._physics.class_type + _class_type = self._physics.class_type + if isinstance(_class_type, str): + from isaaclab.utils.string import string_to_callable + + _class_type = string_to_callable(_class_type) + elif hasattr(_class_type, "resolve"): + _class_type = _class_type.resolve() + self.physics_manager: type[PhysicsManager] = _class_type # type: ignore[assignment] self.physics_manager.initialize(self) self._apply_render_cfg_settings() @@ -193,9 +199,9 @@ def _apply_render_cfg_settings(self) -> None: ) isaaclab_app_exp_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), *[".."] * 4, "apps") - from isaaclab.utils.version import get_isaac_sim_version + from isaaclab.utils.version import get_isaac_sim_version, has_kit - if get_isaac_sim_version().major < 6: + if has_kit() and get_isaac_sim_version().major < 6: isaaclab_app_exp_path = os.path.join(isaaclab_app_exp_path, "isaacsim_5") preset_filename = os.path.join(isaaclab_app_exp_path, f"rendering_modes/{rendering_mode}.kit") @@ -425,13 +431,22 @@ def initialize_visualizers(self) -> None: def initialize_scene_data_provider(self, visualizer_cfgs: list[Any]) -> SceneDataProvider: if self._scene_data_provider is None: - from .scene_data_providers import PhysxSceneDataProvider + from .scene_data_providers import NewtonSceneDataProvider, PhysxSceneDataProvider - # TODO: When Newton/Warp backend scene data provider is implemented and validated, - # switch provider selection to route by physics backend: - # - Omni/PhysX -> PhysxSceneDataProvider - # - Newton/Warp -> NewtonSceneDataProvider - self._scene_data_provider = PhysxSceneDataProvider(visualizer_cfgs, self.stage, self) + # Route by physics backend: Newton -> NewtonSceneDataProvider (reads live + # model/state from NewtonManager), PhysX -> PhysxSceneDataProvider. + _use_newton = False + try: + from isaaclab_newton.physics import NewtonManager + + _use_newton = NewtonManager.get_model() is not None + except ImportError: + pass + + if _use_newton: + self._scene_data_provider = NewtonSceneDataProvider(visualizer_cfgs, self) + else: + self._scene_data_provider = PhysxSceneDataProvider(visualizer_cfgs, self.stage, self) return self._scene_data_provider @property @@ -682,10 +697,14 @@ def build_simulation_context( sim = SimulationContext(sim_cfg) if add_ground_plane: + from .spawners import GroundPlaneCfg + cfg = GroundPlaneCfg() cfg.func("/World/defaultGroundPlane", cfg) if add_lighting or (auto_add_lighting and sim.get_setting("/isaaclab/has_gui")): + from .spawners import DomeLightCfg + cfg = DomeLightCfg( color=(0.1, 0.1, 0.1), enable_color_temperature=True, color_temperature=5500, intensity=10000 ) diff --git a/source/isaaclab/isaaclab/sim/spawners/__init__.py b/source/isaaclab/isaaclab/sim/spawners/__init__.py index 916141906e1..339fab982aa 100644 --- a/source/isaaclab/isaaclab/sim/spawners/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/__init__.py @@ -54,11 +54,31 @@ class and the function call in a single line of code. """ -from .from_files import * # noqa: F401, F403 -from .lights import * # noqa: F401, F403 -from .materials import * # noqa: F401, F403 -from .meshes import * # noqa: F401, F403 -from .sensors import * # noqa: F401, F403 -from .shapes import * # noqa: F401, F403 -from .spawner_cfg import * # noqa: F401, F403 -from .wrappers import * # noqa: F401, F403 +import lazy_loader as lazy + +from .spawner_cfg import DeformableObjectSpawnerCfg, RigidObjectSpawnerCfg, SpawnerCfg + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["from_files", "lights", "materials", "meshes", "sensors", "shapes", "wrappers"], +) +__all__ += ["SpawnerCfg", "RigidObjectSpawnerCfg", "DeformableObjectSpawnerCfg"] + +_lazy_getattr = __getattr__ +_SUBPACKAGES = ("from_files", "lights", "materials", "meshes", "sensors", "shapes", "wrappers") + + +def __getattr__(name): + try: + return _lazy_getattr(name) + except AttributeError: + pass + import importlib + + for subpkg in _SUBPACKAGES: + try: + submod = importlib.import_module(f"{__name__}.{subpkg}") + return getattr(submod, name) + except (ImportError, AttributeError): + continue + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/__init__.py b/source/isaaclab/isaaclab/sim/spawners/from_files/__init__.py index a95ac491b0a..f2cc26caaa8 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/__init__.py @@ -13,11 +13,18 @@ """ -from .from_files import ( - spawn_from_mjcf, - spawn_from_urdf, - spawn_from_usd, - spawn_from_usd_with_compliant_contact_material, - spawn_ground_plane, +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "from_files": [ + "spawn_from_mjcf", + "spawn_from_urdf", + "spawn_from_usd", + "spawn_from_usd_with_compliant_contact_material", + "spawn_ground_plane", + ], + "from_files_cfg": ["GroundPlaneCfg", "MjcfFileCfg", "UrdfFileCfg", "UsdFileCfg", "UsdFileWithCompliantContactCfg"], + }, ) -from .from_files_cfg import GroundPlaneCfg, MjcfFileCfg, UrdfFileCfg, UsdFileCfg, UsdFileWithCompliantContactCfg diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py index 31c3c01abf0..1c1f96ff780 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files.py @@ -8,9 +8,9 @@ import logging from typing import TYPE_CHECKING -import omni.kit.commands -from pxr import Gf, Sdf, Usd +from pxr import Gf, Sdf, Usd, UsdGeom +from isaaclab.utils.version import has_kit from isaaclab.sim import converters, schemas from isaaclab.sim.spawners.materials import RigidBodyMaterialCfg from isaaclab.sim.utils import ( @@ -25,7 +25,7 @@ select_usd_variants, set_prim_visibility, ) -from isaaclab.utils.assets import check_usd_path_with_timeout +from isaaclab.utils.assets import check_file_path, retrieve_file_path if TYPE_CHECKING: from . import from_files_cfg @@ -238,26 +238,33 @@ def spawn_ground_plane( stage=stage, type_to_create_if_not_exist=Sdf.ValueTypeNames.Color3f, ) - # Remove the light from the ground plane + # Remove the light from the ground plane (USD API, works without Kit/Newton) # It isn't bright enough and messes up with the user's lighting settings - omni.kit.commands.execute("ToggleVisibilitySelectedPrims", selected_paths=[f"{prim_path}/SphereLight"], stage=stage) - - prim = stage.GetPrimAtPath(prim_path) - # Apply semantic tags - if hasattr(cfg, "semantic_tags") and cfg.semantic_tags is not None: - # note: taken from replicator scripts.utils.utils.py - for semantic_type, semantic_value in cfg.semantic_tags: - # deal with spaces by replacing them with underscores - semantic_type_sanitized = semantic_type.replace(" ", "_") - semantic_value_sanitized = semantic_value.replace(" ", "_") - # add labels to the prim - add_labels(prim, labels=[semantic_value_sanitized], instance_name=semantic_type_sanitized) - - # Apply visibility - set_prim_visibility(prim, cfg.visible) - - # return the prim - return prim + light_prim = stage.GetPrimAtPath(f"{prim_path}/SphereLight") + if light_prim.IsValid(): + imageable = UsdGeom.Imageable(light_prim) + imageable.MakeInvisible() + if has_kit(): + import omni.kit.commands + + omni.kit.commands.execute("ToggleVisibilitySelectedPrims", selected_paths=[f"{prim_path}/SphereLight"], stage=stage) + + prim = stage.GetPrimAtPath(prim_path) + # Apply semantic tags + if hasattr(cfg, "semantic_tags") and cfg.semantic_tags is not None: + # note: taken from replicator scripts.utils.utils.py + for semantic_type, semantic_value in cfg.semantic_tags: + # deal with spaces by replacing them with underscores + semantic_type_sanitized = semantic_type.replace(" ", "_") + semantic_value_sanitized = semantic_value.replace(" ", "_") + # add labels to the prim + add_labels(prim, labels=[semantic_value_sanitized], instance_name=semantic_type_sanitized) + + # Apply visibility + set_prim_visibility(prim, cfg.visible) + + # return the prim + return prim """ @@ -296,15 +303,15 @@ def _spawn_from_usd_file( Raises: FileNotFoundError: If the USD file does not exist at the given path. """ - # check if usd path exists with periodic logging until timeout - if not check_usd_path_with_timeout(usd_path): - if "4.5" in usd_path: - usd_5_0_path = usd_path.replace("http", "https").replace("/4.5", "/5.0") - if not check_usd_path_with_timeout(usd_5_0_path): - raise FileNotFoundError(f"USD file not found at path at either: '{usd_path}' or '{usd_5_0_path}'.") - usd_path = usd_5_0_path - else: - raise FileNotFoundError(f"USD file not found at path at: '{usd_path}'.") + # check file path exists (supports local paths, S3, HTTP/HTTPS URLs) + # check_file_path returns: 0 (not found), 1 (local), 2 (remote) + file_status = check_file_path(usd_path) + if file_status == 0: + raise FileNotFoundError(f"USD file not found at path: '{usd_path}'.") + + # Download remote files (S3, HTTP, HTTPS) to local cache + if file_status == 2: + usd_path = retrieve_file_path(usd_path, force_download=False) # Obtain current stage stage = get_current_stage() diff --git a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py index ad9f5585904..6e89733038a 100644 --- a/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/from_files/from_files_cfg.py @@ -5,17 +5,14 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import MISSING from isaaclab.sim import converters, schemas from isaaclab.sim.spawners import materials from isaaclab.sim.spawners.spawner_cfg import DeformableObjectSpawnerCfg, RigidObjectSpawnerCfg, SpawnerCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR -from . import from_files - @configclass class FileCfg(RigidObjectSpawnerCfg, DeformableObjectSpawnerCfg): @@ -96,7 +93,7 @@ class UsdFileCfg(FileCfg): This is done by calling the respective function with the specified properties. """ - func: Callable = from_files.spawn_from_usd + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.from_files.from_files:spawn_from_usd") usd_path: str = MISSING """Path to the USD file to spawn asset from.""" @@ -129,7 +126,7 @@ class UrdfFileCfg(FileCfg, converters.UrdfConverterCfg): """ - func: Callable = from_files.spawn_from_urdf + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.from_files.from_files:spawn_from_urdf") @configclass @@ -151,7 +148,7 @@ class MjcfFileCfg(FileCfg, converters.MjcfConverterCfg): """ - func: Callable = from_files.spawn_from_mjcf + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.from_files.from_files:spawn_from_mjcf") """ @@ -169,7 +166,7 @@ class UsdFileWithCompliantContactCfg(UsdFileCfg): material application. """ - func: Callable = from_files.spawn_from_usd_with_compliant_contact_material + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.from_files.from_files:spawn_from_usd_with_compliant_contact_material") compliant_contact_stiffness: float | None = None """Stiffness of the compliant contact. Defaults to None. @@ -201,7 +198,7 @@ class GroundPlaneCfg(SpawnerCfg): This uses the USD for the standard grid-world ground plane from Isaac Sim by default. """ - func: Callable = from_files.spawn_ground_plane + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.from_files.from_files:spawn_ground_plane") usd_path: str = f"{ISAAC_NUCLEUS_DIR}/Environments/Grid/default_environment.usd" """Path to the USD file to spawn asset from. Defaults to the grid-world ground plane.""" diff --git a/source/isaaclab/isaaclab/sim/spawners/lights/__init__.py b/source/isaaclab/isaaclab/sim/spawners/lights/__init__.py index df0c638f58f..193f1e839e5 100644 --- a/source/isaaclab/isaaclab/sim/spawners/lights/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/lights/__init__.py @@ -10,5 +10,12 @@ `_. """ -from .lights import spawn_light -from .lights_cfg import CylinderLightCfg, DiskLightCfg, DistantLightCfg, DomeLightCfg, LightCfg, SphereLightCfg +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "lights": ["spawn_light"], + "lights_cfg": ["CylinderLightCfg", "DiskLightCfg", "DistantLightCfg", "DomeLightCfg", "LightCfg", "SphereLightCfg"], + }, +) diff --git a/source/isaaclab/isaaclab/sim/spawners/lights/lights_cfg.py b/source/isaaclab/isaaclab/sim/spawners/lights/lights_cfg.py index 4768ffcba57..896278f6fd1 100644 --- a/source/isaaclab/isaaclab/sim/spawners/lights/lights_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/lights/lights_cfg.py @@ -5,14 +5,11 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import MISSING from typing import Literal from isaaclab.sim.spawners.spawner_cfg import SpawnerCfg -from isaaclab.utils import configclass - -from . import lights +from isaaclab.utils import DeferredClass, configclass @configclass @@ -26,7 +23,7 @@ class LightCfg(SpawnerCfg): The default values for the attributes are those specified in the their official documentation. """ - func: Callable = lights.spawn_light + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.lights.lights:spawn_light") prim_type: str = MISSING """The prim type name for the light prim.""" diff --git a/source/isaaclab/isaaclab/sim/spawners/materials/__init__.py b/source/isaaclab/isaaclab/sim/spawners/materials/__init__.py index 0b0e7851494..048024a4bca 100644 --- a/source/isaaclab/isaaclab/sim/spawners/materials/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/materials/__init__.py @@ -52,7 +52,14 @@ .. _Physics Scene: https://openusd.org/dev/api/usd_physics_page_front.html """ -from .physics_materials import spawn_deformable_body_material, spawn_rigid_body_material -from .physics_materials_cfg import DeformableBodyMaterialCfg, PhysicsMaterialCfg, RigidBodyMaterialCfg -from .visual_materials import spawn_from_mdl_file, spawn_preview_surface -from .visual_materials_cfg import GlassMdlCfg, MdlFileCfg, PreviewSurfaceCfg, VisualMaterialCfg +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "physics_materials": ["spawn_deformable_body_material", "spawn_rigid_body_material"], + "physics_materials_cfg": ["DeformableBodyMaterialCfg", "PhysicsMaterialCfg", "RigidBodyMaterialCfg"], + "visual_materials": ["spawn_from_mdl_file", "spawn_preview_surface"], + "visual_materials_cfg": ["GlassMdlCfg", "MdlFileCfg", "PreviewSurfaceCfg", "VisualMaterialCfg"], + }, +) diff --git a/source/isaaclab/isaaclab/sim/spawners/materials/physics_materials_cfg.py b/source/isaaclab/isaaclab/sim/spawners/materials/physics_materials_cfg.py index be63a9dd4a4..8c036932a4f 100644 --- a/source/isaaclab/isaaclab/sim/spawners/materials/physics_materials_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/materials/physics_materials_cfg.py @@ -5,13 +5,10 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import MISSING from typing import Literal -from isaaclab.utils import configclass - -from . import physics_materials +from isaaclab.utils import DeferredClass, configclass @configclass @@ -24,7 +21,7 @@ class PhysicsMaterialCfg: `PhysX documentation `__. """ - func: Callable = MISSING + func: DeferredClass = MISSING """Function to use for creating the material.""" @@ -35,7 +32,7 @@ class RigidBodyMaterialCfg(PhysicsMaterialCfg): See :meth:`spawn_rigid_body_material` for more information. """ - func: Callable = physics_materials.spawn_rigid_body_material + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.materials.physics_materials:spawn_rigid_body_material") static_friction: float = 0.5 """The static friction coefficient. Defaults to 0.5.""" @@ -89,7 +86,7 @@ class DeformableBodyMaterialCfg(PhysicsMaterialCfg): """ - func: Callable = physics_materials.spawn_deformable_body_material + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.materials.physics_materials:spawn_deformable_body_material") density: float | None = None """The material density. Defaults to None, in which case the simulation decides the default density.""" diff --git a/source/isaaclab/isaaclab/sim/spawners/materials/visual_materials_cfg.py b/source/isaaclab/isaaclab/sim/spawners/materials/visual_materials_cfg.py index 09e022c95cb..3accaaa8c7b 100644 --- a/source/isaaclab/isaaclab/sim/spawners/materials/visual_materials_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/materials/visual_materials_cfg.py @@ -5,19 +5,16 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import MISSING -from isaaclab.utils import configclass - -from . import visual_materials +from isaaclab.utils import DeferredClass, configclass @configclass class VisualMaterialCfg: """Configuration parameters for creating a visual material.""" - func: Callable = MISSING + func: DeferredClass = MISSING """The function to use for creating the material.""" @@ -28,7 +25,7 @@ class PreviewSurfaceCfg(VisualMaterialCfg): See :meth:`spawn_preview_surface` for more information. """ - func: Callable = visual_materials.spawn_preview_surface + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.materials.visual_materials:spawn_preview_surface") diffuse_color: tuple[float, float, float] = (0.18, 0.18, 0.18) """The RGB diffusion color. This is the base color of the surface. Defaults to a dark gray.""" @@ -53,7 +50,7 @@ class MdlFileCfg(VisualMaterialCfg): See :meth:`spawn_from_mdl_file` for more information. """ - func: Callable = visual_materials.spawn_from_mdl_file + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.materials.visual_materials:spawn_from_mdl_file") mdl_path: str = MISSING """The path to the MDL material. @@ -95,7 +92,7 @@ class GlassMdlCfg(VisualMaterialCfg): The default values are taken from the glass material in the NVIDIA Nucleus. """ - func: Callable = visual_materials.spawn_from_mdl_file + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.materials.visual_materials:spawn_from_mdl_file") mdl_path: str = "OmniGlass.mdl" """The path to the MDL material. Defaults to the glass material in the NVIDIA Nucleus.""" diff --git a/source/isaaclab/isaaclab/sim/spawners/meshes/__init__.py b/source/isaaclab/isaaclab/sim/spawners/meshes/__init__.py index 49836dc5cbd..a4ad5d16009 100644 --- a/source/isaaclab/isaaclab/sim/spawners/meshes/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/meshes/__init__.py @@ -20,5 +20,25 @@ .. _USDGeomMesh: https://openusd.org/release/api/class_usd_geom_mesh.html """ -from .meshes import spawn_mesh_capsule, spawn_mesh_cone, spawn_mesh_cuboid, spawn_mesh_cylinder, spawn_mesh_sphere -from .meshes_cfg import MeshCapsuleCfg, MeshCfg, MeshConeCfg, MeshCuboidCfg, MeshCylinderCfg, MeshSphereCfg +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "meshes": [ + "spawn_mesh_capsule", + "spawn_mesh_cone", + "spawn_mesh_cuboid", + "spawn_mesh_cylinder", + "spawn_mesh_sphere", + ], + "meshes_cfg": [ + "MeshCapsuleCfg", + "MeshCfg", + "MeshConeCfg", + "MeshCuboidCfg", + "MeshCylinderCfg", + "MeshSphereCfg", + ], + }, +) diff --git a/source/isaaclab/isaaclab/sim/spawners/meshes/meshes_cfg.py b/source/isaaclab/isaaclab/sim/spawners/meshes/meshes_cfg.py index d5c39a505b8..c98a545dd5f 100644 --- a/source/isaaclab/isaaclab/sim/spawners/meshes/meshes_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/meshes/meshes_cfg.py @@ -5,15 +5,12 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import MISSING from typing import Literal from isaaclab.sim.spawners import materials from isaaclab.sim.spawners.spawner_cfg import DeformableObjectSpawnerCfg, RigidObjectSpawnerCfg -from isaaclab.utils import configclass - -from . import meshes +from isaaclab.utils import DeferredClass, configclass @configclass @@ -74,7 +71,7 @@ class MeshSphereCfg(MeshCfg): See :meth:`spawn_mesh_sphere` for more information. """ - func: Callable = meshes.spawn_mesh_sphere + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.meshes.meshes:spawn_mesh_sphere") radius: float = MISSING """Radius of the sphere (in m).""" @@ -87,7 +84,7 @@ class MeshCuboidCfg(MeshCfg): See :meth:`spawn_mesh_cuboid` for more information. """ - func: Callable = meshes.spawn_mesh_cuboid + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.meshes.meshes:spawn_mesh_cuboid") size: tuple[float, float, float] = MISSING """Size of the cuboid (in m).""" @@ -100,7 +97,7 @@ class MeshCylinderCfg(MeshCfg): See :meth:`spawn_cylinder` for more information. """ - func: Callable = meshes.spawn_mesh_cylinder + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.meshes.meshes:spawn_mesh_cylinder") radius: float = MISSING """Radius of the cylinder (in m).""" @@ -117,7 +114,7 @@ class MeshCapsuleCfg(MeshCfg): See :meth:`spawn_capsule` for more information. """ - func: Callable = meshes.spawn_mesh_capsule + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.meshes.meshes:spawn_mesh_capsule") radius: float = MISSING """Radius of the capsule (in m).""" @@ -134,7 +131,7 @@ class MeshConeCfg(MeshCfg): See :meth:`spawn_cone` for more information. """ - func: Callable = meshes.spawn_mesh_cone + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.meshes.meshes:spawn_mesh_cone") radius: float = MISSING """Radius of the cone (in m).""" diff --git a/source/isaaclab/isaaclab/sim/spawners/sensors/__init__.py b/source/isaaclab/isaaclab/sim/spawners/sensors/__init__.py index ac61868c025..6583ae01ec0 100644 --- a/source/isaaclab/isaaclab/sim/spawners/sensors/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/sensors/__init__.py @@ -11,5 +11,12 @@ """ -from .sensors import spawn_camera -from .sensors_cfg import FisheyeCameraCfg, PinholeCameraCfg +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "sensors": ["spawn_camera"], + "sensors_cfg": ["FisheyeCameraCfg", "PinholeCameraCfg"], + }, +) diff --git a/source/isaaclab/isaaclab/sim/spawners/sensors/sensors_cfg.py b/source/isaaclab/isaaclab/sim/spawners/sensors/sensors_cfg.py index 44e5eb06173..6c2bbf26121 100644 --- a/source/isaaclab/isaaclab/sim/spawners/sensors/sensors_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/sensors/sensors_cfg.py @@ -5,14 +5,11 @@ from __future__ import annotations -from collections.abc import Callable from typing import Literal import isaaclab.utils.sensors as sensor_utils from isaaclab.sim.spawners.spawner_cfg import SpawnerCfg -from isaaclab.utils import configclass - -from . import sensors +from isaaclab.utils import DeferredClass, configclass @configclass @@ -26,7 +23,7 @@ class PinholeCameraCfg(SpawnerCfg): world unit is Meter s.t. all of these values are set in cm. """ - func: Callable = sensors.spawn_camera + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.sensors.sensors:spawn_camera") projection_type: str = "pinhole" """Type of projection to use for the camera. Defaults to "pinhole". @@ -172,7 +169,7 @@ class FisheyeCameraCfg(PinholeCameraCfg): .. _fish-eye camera: https://en.wikipedia.org/wiki/Fisheye_lens """ - func: Callable = sensors.spawn_camera + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.sensors.sensors:spawn_camera") projection_type: Literal[ "fisheyePolynomial", diff --git a/source/isaaclab/isaaclab/sim/spawners/shapes/__init__.py b/source/isaaclab/isaaclab/sim/spawners/shapes/__init__.py index 8f6cab9439c..97cd6dcf9f0 100644 --- a/source/isaaclab/isaaclab/sim/spawners/shapes/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/shapes/__init__.py @@ -14,5 +14,12 @@ """ -from .shapes import spawn_capsule, spawn_cone, spawn_cuboid, spawn_cylinder, spawn_sphere -from .shapes_cfg import CapsuleCfg, ConeCfg, CuboidCfg, CylinderCfg, ShapeCfg, SphereCfg +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "shapes": ["spawn_capsule", "spawn_cone", "spawn_cuboid", "spawn_cylinder", "spawn_sphere"], + "shapes_cfg": ["CapsuleCfg", "ConeCfg", "CuboidCfg", "CylinderCfg", "ShapeCfg", "SphereCfg"], + }, +) diff --git a/source/isaaclab/isaaclab/sim/spawners/shapes/shapes_cfg.py b/source/isaaclab/isaaclab/sim/spawners/shapes/shapes_cfg.py index 921b1ede29c..6411d24bb5a 100644 --- a/source/isaaclab/isaaclab/sim/spawners/shapes/shapes_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/shapes/shapes_cfg.py @@ -5,15 +5,12 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import MISSING from typing import Literal from isaaclab.sim.spawners import materials from isaaclab.sim.spawners.spawner_cfg import RigidObjectSpawnerCfg -from isaaclab.utils import configclass - -from . import shapes +from isaaclab.utils import DeferredClass, configclass @configclass @@ -54,7 +51,7 @@ class SphereCfg(ShapeCfg): See :meth:`spawn_sphere` for more information. """ - func: Callable = shapes.spawn_sphere + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.shapes.shapes:spawn_sphere") radius: float = MISSING """Radius of the sphere (in m).""" @@ -67,7 +64,7 @@ class CuboidCfg(ShapeCfg): See :meth:`spawn_cuboid` for more information. """ - func: Callable = shapes.spawn_cuboid + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.shapes.shapes:spawn_cuboid") size: tuple[float, float, float] = MISSING """Size of the cuboid.""" @@ -80,7 +77,7 @@ class CylinderCfg(ShapeCfg): See :meth:`spawn_cylinder` for more information. """ - func: Callable = shapes.spawn_cylinder + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.shapes.shapes:spawn_cylinder") radius: float = MISSING """Radius of the cylinder (in m).""" @@ -97,7 +94,7 @@ class CapsuleCfg(ShapeCfg): See :meth:`spawn_capsule` for more information. """ - func: Callable = shapes.spawn_capsule + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.shapes.shapes:spawn_capsule") radius: float = MISSING """Radius of the capsule (in m).""" @@ -114,7 +111,7 @@ class ConeCfg(ShapeCfg): See :meth:`spawn_cone` for more information. """ - func: Callable = shapes.spawn_cone + func: DeferredClass = DeferredClass("isaaclab.sim.spawners.shapes.shapes:spawn_cone") radius: float = MISSING """Radius of the cone (in m).""" diff --git a/source/isaaclab/isaaclab/sim/spawners/spawner_cfg.py b/source/isaaclab/isaaclab/sim/spawners/spawner_cfg.py index d4ef0da9b60..197dd8124e9 100644 --- a/source/isaaclab/isaaclab/sim/spawners/spawner_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/spawner_cfg.py @@ -7,12 +7,14 @@ from collections.abc import Callable from dataclasses import MISSING - -from pxr import Usd +from typing import TYPE_CHECKING from isaaclab.sim import schemas from isaaclab.utils import configclass +if TYPE_CHECKING: + from pxr import Usd + @configclass class SpawnerCfg: diff --git a/source/isaaclab/isaaclab/sim/spawners/wrappers/__init__.py b/source/isaaclab/isaaclab/sim/spawners/wrappers/__init__.py index 4006fa1a6ab..3b37134df75 100644 --- a/source/isaaclab/isaaclab/sim/spawners/wrappers/__init__.py +++ b/source/isaaclab/isaaclab/sim/spawners/wrappers/__init__.py @@ -10,5 +10,12 @@ different configurations. """ -from .wrappers import spawn_multi_asset, spawn_multi_usd_file -from .wrappers_cfg import MultiAssetSpawnerCfg, MultiUsdFileCfg +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "wrappers": ["spawn_multi_asset", "spawn_multi_usd_file"], + "wrappers_cfg": ["MultiAssetSpawnerCfg", "MultiUsdFileCfg"], + }, +) diff --git a/source/isaaclab/isaaclab/sim/spawners/wrappers/wrappers_cfg.py b/source/isaaclab/isaaclab/sim/spawners/wrappers/wrappers_cfg.py index f3d7b7986ec..8bdd9aa6ca8 100644 --- a/source/isaaclab/isaaclab/sim/spawners/wrappers/wrappers_cfg.py +++ b/source/isaaclab/isaaclab/sim/spawners/wrappers/wrappers_cfg.py @@ -7,9 +7,7 @@ from isaaclab.sim.spawners.from_files import UsdFileCfg from isaaclab.sim.spawners.spawner_cfg import DeformableObjectSpawnerCfg, RigidObjectSpawnerCfg, SpawnerCfg -from isaaclab.utils import configclass - -from . import wrappers +from isaaclab.utils import DeferredClass, configclass @configclass @@ -29,7 +27,7 @@ class MultiAssetSpawnerCfg(RigidObjectSpawnerCfg, DeformableObjectSpawnerCfg): """ - func = wrappers.spawn_multi_asset + func = DeferredClass("isaaclab.sim.spawners.wrappers.wrappers:spawn_multi_asset") assets_cfg: list[SpawnerCfg] = MISSING """List of asset configurations to spawn.""" @@ -57,7 +55,7 @@ class MultiUsdFileCfg(UsdFileCfg): """ - func = wrappers.spawn_multi_usd_file + func = DeferredClass("isaaclab.sim.spawners.wrappers.wrappers:spawn_multi_usd_file") usd_path: str | list[str] = MISSING """Path or a list of paths to the USD files to spawn asset from.""" diff --git a/source/isaaclab/isaaclab/sim/utils/__init__.py b/source/isaaclab/isaaclab/sim/utils/__init__.py index 3a85ae44c2f..0934f290e1d 100644 --- a/source/isaaclab/isaaclab/sim/utils/__init__.py +++ b/source/isaaclab/isaaclab/sim/utils/__init__.py @@ -5,9 +5,28 @@ """Utilities built around USD operations.""" -from .legacy import * # noqa: F401, F403 -from .prims import * # noqa: F401, F403 -from .queries import * # noqa: F401, F403 -from .semantics import * # noqa: F401, F403 -from .stage import * # noqa: F401, F403 -from .transforms import * # noqa: F401, F403 +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["legacy", "prims", "queries", "semantics", "stage", "transforms"], +) + +_lazy_getattr = __getattr__ +_SUBMODULES = ("legacy", "prims", "queries", "semantics", "stage", "transforms") + + +def __getattr__(name): + try: + return _lazy_getattr(name) + except AttributeError: + pass + import importlib + + for submod_name in _SUBMODULES: + try: + submod = importlib.import_module(f"{__name__}.{submod_name}") + return getattr(submod, name) + except (ImportError, AttributeError): + continue + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/source/isaaclab/isaaclab/sim/utils/prims.py b/source/isaaclab/isaaclab/sim/utils/prims.py index 174484bb532..0b9dade3a67 100644 --- a/source/isaaclab/isaaclab/sim/utils/prims.py +++ b/source/isaaclab/isaaclab/sim/utils/prims.py @@ -16,11 +16,9 @@ import torch -import omni.kit.commands from pxr import Sdf, Usd, UsdGeom, UsdPhysics, UsdShade, UsdUtils - from isaaclab.utils.string import to_camel_case -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit from .queries import find_matching_prim_paths from .semantics import add_labels @@ -213,13 +211,17 @@ def delete_prim(prim_path: str | Sequence[str], stage: Usd.Stage | None = None) stage_id = stage_cache.GetId(stage).ToLongInt() if stage_id < 0: stage_id = stage_cache.Insert(stage).ToLongInt() - # delete prims - success, _ = omni.kit.commands.execute( - "DeletePrimsCommand", - paths=prim_path, - stage=stage, - ) - return success + + if has_kit(): + # delete prims + import omni.kit.commands + + success, _ = omni.kit.commands.execute( + "DeletePrimsCommand", + paths=prim_path, + stage=stage, + ) + return success """ @@ -680,21 +682,24 @@ def wrapper(prim_path: str | Sdf.Path, cfg: SpawnerCfg, *args, **kwargs): imageable.MakeVisible() else: imageable.MakeInvisible() - # set the semantic annotations + # set the semantic annotations (custom USD attributes, no SemanticsAPI import) if hasattr(cfg, "semantic_tags") and cfg.semantic_tags is not None: - # note: taken from replicator scripts.utils.utils.py for semantic_type, semantic_value in cfg.semantic_tags: - # deal with spaces by replacing them with underscores semantic_type_sanitized = semantic_type.replace(" ", "_") semantic_value_sanitized = semantic_value.replace(" ", "_") - # set the semantic API for the instance instance_name = f"{semantic_type_sanitized}_{semantic_value_sanitized}" - sem = Semantics.SemanticsAPI.Apply(prim, instance_name) - # create semantic type and data attributes - sem.CreateSemanticTypeAttr() - sem.CreateSemanticDataAttr() - sem.GetSemanticTypeAttr().Set(semantic_type) - sem.GetSemanticDataAttr().Set(semantic_value) + + type_attr_name = f"semantic:{instance_name}:semantic:type" + type_attr = prim.GetAttribute(type_attr_name) + if not type_attr: + type_attr = prim.CreateAttribute(type_attr_name, Sdf.ValueTypeNames.String) + type_attr.Set(semantic_type) + + data_attr_name = f"semantic:{instance_name}:semantic:data" + data_attr = prim.GetAttribute(data_attr_name) + if not data_attr: + data_attr = prim.CreateAttribute(data_attr_name, Sdf.ValueTypeNames.String) + data_attr.Set(semantic_value) # activate rigid body contact sensors (lazy import to avoid circular import with schemas) if hasattr(cfg, "activate_contact_sensors") and cfg.activate_contact_sensors: from ..schemas import schemas as _schemas @@ -751,6 +756,8 @@ def bind_visual_material( Raises: ValueError: If the provided prim paths do not exist on stage. """ + if not has_kit(): + return # get stage handle if stage is None: stage = get_current_stage() @@ -769,6 +776,7 @@ def bind_visual_material( binding_strength = "weakerThanDescendants" # obtain material binding API # note: we prefer using the command here as it is more robust than the USD API + import omni.kit.commands success, _ = omni.kit.commands.execute( "BindMaterialCommand", prim_path=prim_path, @@ -882,6 +890,18 @@ def add_usd_reference( Raises: FileNotFoundError: When the input USD file is not found at the specified path. """ + # resolve remote USD paths to local (same as Newton / add_reference_to_stage) + from isaaclab.utils.assets import check_file_path, retrieve_file_path + + file_status = check_file_path(usd_path) + if file_status == 0: + raise FileNotFoundError(f"Unable to open the usd file at path: {usd_path}") + if file_status == 2: + try: + usd_path = retrieve_file_path(usd_path, force_download=False) + except Exception as e: + raise FileNotFoundError(f"Failed to retrieve USD file from {usd_path}") from e + # get current stage stage = get_current_stage() if stage is None else stage # get prim at path @@ -899,16 +919,19 @@ def _add_reference_to_prim(prim: Usd.Prim) -> Usd.Prim: return prim # Compatibility with Isaac Sim 4.5 where omni.metrics is not available - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: + return _add_reference_to_prim(prim) + + # When Kit is not running (e.g. Newton), use USD-only path (no omni.metrics) + if not has_kit(): return _add_reference_to_prim(prim) - # check if the USD file is valid and add reference to the prim + # check if the USD file is valid and add reference to the prim (Isaac Sim 5+ with Kit) sdf_layer = Sdf.Layer.FindOrOpen(usd_path) if not sdf_layer: raise FileNotFoundError(f"Unable to open the usd file at path: {usd_path}") - # import metrics assembler interface - # note: this is only available in Isaac Sim 5.0 and above + # import metrics assembler interface (Kit/Isaac Sim only) from omni.metrics.assembler.core import get_metrics_assembler_interface # obtain the stage ID diff --git a/source/isaaclab/isaaclab/sim/utils/queries.py b/source/isaaclab/isaaclab/sim/utils/queries.py index 1929e70941c..7039d4e0b00 100644 --- a/source/isaaclab/isaaclab/sim/utils/queries.py +++ b/source/isaaclab/isaaclab/sim/utils/queries.py @@ -43,6 +43,8 @@ def get_next_free_prim_path(path: str, stage: Usd.Stage | None = None) -> str: >>> sim_utils.get_next_free_prim_path("/World/Cube") /World/Cube_02 """ + import omni.usd + # get current stage stage = get_current_stage() if stage is None else stage diff --git a/source/isaaclab/isaaclab/sim/utils/stage.py b/source/isaaclab/isaaclab/sim/utils/stage.py index 88b4ebeeca4..2cb24698f88 100644 --- a/source/isaaclab/isaaclab/sim/utils/stage.py +++ b/source/isaaclab/isaaclab/sim/utils/stage.py @@ -7,7 +7,6 @@ from __future__ import annotations -import builtins import contextlib import logging import os @@ -251,7 +250,7 @@ def use_stage(stage: Usd.Stage) -> Generator[None, None, None]: ... pass >>> # operate on the default stage attached to the USD context """ - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: logger.warning("Isaac Sim < 5.0 does not support thread-local stage contexts. Skipping use_stage().") yield # no-op else: @@ -321,6 +320,8 @@ def update_stage() -> None: >>> for _ in range(100): ... sim.step() # Handles updates internally """ + import omni.kit.app + omni.kit.app.get_app_interface().update() @@ -473,7 +474,9 @@ def _predicate_from_path(prim: Usd.Prim) -> bool: # delete prims delete_prim(prim_paths_to_delete) - if builtins.ISAAC_LAUNCHED_FROM_TERMINAL is False: # type: ignore + if has_kit(): + import omni.kit.app + omni.kit.app.get_app_interface().update() @@ -508,6 +511,7 @@ def get_current_stage(fabric: bool = False) -> Usd.Stage: return stage + def get_current_stage_id() -> int: """Get the current open stage ID. diff --git a/source/isaaclab/isaaclab/sim/views/__init__.py b/source/isaaclab/isaaclab/sim/views/__init__.py index eb5bea7690c..b3139c5cd58 100644 --- a/source/isaaclab/isaaclab/sim/views/__init__.py +++ b/source/isaaclab/isaaclab/sim/views/__init__.py @@ -5,4 +5,11 @@ """Views for manipulating USD prims.""" -from .xform_prim_view import XformPrimView +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "xform_prim_view": ["XformPrimView"], + }, +) diff --git a/source/isaaclab/isaaclab/sim/views/xform_prim_view.py b/source/isaaclab/isaaclab/sim/views/xform_prim_view.py index c6d8009b64a..6652f50cbe2 100644 --- a/source/isaaclab/isaaclab/sim/views/xform_prim_view.py +++ b/source/isaaclab/isaaclab/sim/views/xform_prim_view.py @@ -12,10 +12,10 @@ import torch import warp as wp -import carb from pxr import Gf, Sdf, Usd, UsdGeom, Vt import isaaclab.sim as sim_utils +from isaaclab.app.settings_manager import get_settings_manager from isaaclab.utils.warp import fabric as fabric_utils logger = logging.getLogger(__name__) @@ -136,8 +136,8 @@ def __init__( " Use sim_utils.standardize_xform_ops() to prepare the prim." ) - # Determine if Fabric is supported on the device - self._use_fabric = carb.settings.get_settings().get("/physics/fabricEnabled") + # Determine if Fabric is supported on the device (via SettingsManager for Newton/standalone compatibility) + self._use_fabric = bool(get_settings_manager().get("/physics/fabricEnabled", False)) logger.debug(f"Using Fabric for the XFormPrimView over '{self._prim_path}' on device '{self._device}'.") # Check for unsupported Fabric + CPU combination diff --git a/source/isaaclab/isaaclab/terrains/__init__.py b/source/isaaclab/isaaclab/terrains/__init__.py index 6f0b5018557..baa0a486e7d 100644 --- a/source/isaaclab/isaaclab/terrains/__init__.py +++ b/source/isaaclab/isaaclab/terrains/__init__.py @@ -19,11 +19,43 @@ * :meth:`TerrainImporter.import_usd`: spawn a prim as reference to input USD file. """ -from .height_field import * # noqa: F401, F403 -from .sub_terrain_cfg import FlatPatchSamplingCfg, SubTerrainBaseCfg -from .terrain_generator import TerrainGenerator -from .terrain_generator_cfg import TerrainGeneratorCfg -from .terrain_importer import TerrainImporter -from .terrain_importer_cfg import TerrainImporterCfg -from .trimesh import * # noqa: F401, F403 -from .utils import color_meshes_by_height, create_prim_from_mesh +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=["height_field", "trimesh"], + submod_attrs={ + "sub_terrain_cfg": ["FlatPatchSamplingCfg", "SubTerrainBaseCfg"], + "terrain_generator": ["TerrainGenerator"], + "terrain_generator_cfg": ["TerrainGeneratorCfg"], + "terrain_importer": ["TerrainImporter"], + "terrain_importer_cfg": ["TerrainImporterCfg"], + "utils": ["color_meshes_by_height", "create_prim_from_mesh"], + "trimesh": [ + "MeshBoxTerrainCfg", + "MeshFloatingRingTerrainCfg", + "MeshGapTerrainCfg", + "MeshInvertedPyramidStairsTerrainCfg", + "MeshPitTerrainCfg", + "MeshPlaneTerrainCfg", + "MeshPyramidStairsTerrainCfg", + "MeshRailsTerrainCfg", + "MeshRandomGridTerrainCfg", + "MeshRepeatedBoxesTerrainCfg", + "MeshRepeatedCylindersTerrainCfg", + "MeshRepeatedPyramidsTerrainCfg", + "MeshStarTerrainCfg", + ], + "height_field": [ + "HfDiscreteObstaclesTerrainCfg", + "HfInvertedPyramidSlopedTerrainCfg", + "HfInvertedPyramidStairsTerrainCfg", + "HfPyramidSlopedTerrainCfg", + "HfPyramidStairsTerrainCfg", + "HfRandomUniformTerrainCfg", + "HfSteppingStonesTerrainCfg", + "HfTerrainBaseCfg", + "HfWaveTerrainCfg", + ], + }, +) diff --git a/source/isaaclab/isaaclab/terrains/height_field/__init__.py b/source/isaaclab/isaaclab/terrains/height_field/__init__.py index 3bc28ba3ccf..492d1bd112a 100644 --- a/source/isaaclab/isaaclab/terrains/height_field/__init__.py +++ b/source/isaaclab/isaaclab/terrains/height_field/__init__.py @@ -25,14 +25,21 @@ """ -from .hf_terrains_cfg import ( - HfDiscreteObstaclesTerrainCfg, - HfInvertedPyramidSlopedTerrainCfg, - HfInvertedPyramidStairsTerrainCfg, - HfPyramidSlopedTerrainCfg, - HfPyramidStairsTerrainCfg, - HfRandomUniformTerrainCfg, - HfSteppingStonesTerrainCfg, - HfTerrainBaseCfg, - HfWaveTerrainCfg, +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "hf_terrains_cfg": [ + "HfDiscreteObstaclesTerrainCfg", + "HfInvertedPyramidSlopedTerrainCfg", + "HfInvertedPyramidStairsTerrainCfg", + "HfPyramidSlopedTerrainCfg", + "HfPyramidStairsTerrainCfg", + "HfRandomUniformTerrainCfg", + "HfSteppingStonesTerrainCfg", + "HfTerrainBaseCfg", + "HfWaveTerrainCfg", + ], + }, ) diff --git a/source/isaaclab/isaaclab/terrains/height_field/hf_terrains_cfg.py b/source/isaaclab/isaaclab/terrains/height_field/hf_terrains_cfg.py index df1d6dcc21a..b70c72ebc89 100644 --- a/source/isaaclab/isaaclab/terrains/height_field/hf_terrains_cfg.py +++ b/source/isaaclab/isaaclab/terrains/height_field/hf_terrains_cfg.py @@ -6,9 +6,9 @@ from dataclasses import MISSING from isaaclab.utils import configclass +from isaaclab.utils.string import DeferredClass from ..sub_terrain_cfg import SubTerrainBaseCfg -from . import hf_terrains @configclass @@ -42,7 +42,7 @@ class HfTerrainBaseCfg(SubTerrainBaseCfg): class HfRandomUniformTerrainCfg(HfTerrainBaseCfg): """Configuration for a random uniform height field terrain.""" - function = hf_terrains.random_uniform_terrain + function = DeferredClass("isaaclab.terrains.height_field.hf_terrains:random_uniform_terrain") noise_range: tuple[float, float] = MISSING """The minimum and maximum height noise (i.e. along z) of the terrain (in m).""" @@ -63,7 +63,7 @@ class HfRandomUniformTerrainCfg(HfTerrainBaseCfg): class HfPyramidSlopedTerrainCfg(HfTerrainBaseCfg): """Configuration for a pyramid sloped height field terrain.""" - function = hf_terrains.pyramid_sloped_terrain + function = DeferredClass("isaaclab.terrains.height_field.hf_terrains:pyramid_sloped_terrain") slope_range: tuple[float, float] = MISSING """The slope of the terrain (in radians).""" @@ -95,7 +95,7 @@ class HfInvertedPyramidSlopedTerrainCfg(HfPyramidSlopedTerrainCfg): class HfPyramidStairsTerrainCfg(HfTerrainBaseCfg): """Configuration for a pyramid stairs height field terrain.""" - function = hf_terrains.pyramid_stairs_terrain + function = DeferredClass("isaaclab.terrains.height_field.hf_terrains:pyramid_stairs_terrain") step_height_range: tuple[float, float] = MISSING """The minimum and maximum height of the steps (in m).""" @@ -130,7 +130,7 @@ class HfInvertedPyramidStairsTerrainCfg(HfPyramidStairsTerrainCfg): class HfDiscreteObstaclesTerrainCfg(HfTerrainBaseCfg): """Configuration for a discrete obstacles height field terrain.""" - function = hf_terrains.discrete_obstacles_terrain + function = DeferredClass("isaaclab.terrains.height_field.hf_terrains:discrete_obstacles_terrain") obstacle_height_mode: str = "choice" """The mode to use for the obstacle height. Defaults to "choice". @@ -155,7 +155,7 @@ class HfDiscreteObstaclesTerrainCfg(HfTerrainBaseCfg): class HfWaveTerrainCfg(HfTerrainBaseCfg): """Configuration for a wave height field terrain.""" - function = hf_terrains.wave_terrain + function = DeferredClass("isaaclab.terrains.height_field.hf_terrains:wave_terrain") amplitude_range: tuple[float, float] = MISSING """The minimum and maximum amplitude of the wave (in m).""" @@ -168,7 +168,7 @@ class HfWaveTerrainCfg(HfTerrainBaseCfg): class HfSteppingStonesTerrainCfg(HfTerrainBaseCfg): """Configuration for a stepping stones height field terrain.""" - function = hf_terrains.stepping_stones_terrain + function = DeferredClass("isaaclab.terrains.height_field.hf_terrains:stepping_stones_terrain") stone_height_max: float = MISSING """The maximum height of the stones (in m).""" diff --git a/source/isaaclab/isaaclab/terrains/sub_terrain_cfg.py b/source/isaaclab/isaaclab/terrains/sub_terrain_cfg.py index 0dab3ea8f3c..5167afd5eb6 100644 --- a/source/isaaclab/isaaclab/terrains/sub_terrain_cfg.py +++ b/source/isaaclab/isaaclab/terrains/sub_terrain_cfg.py @@ -8,12 +8,14 @@ from collections.abc import Callable from dataclasses import MISSING - -import numpy as np -import trimesh +from typing import TYPE_CHECKING from isaaclab.utils import configclass +if TYPE_CHECKING: + import numpy as np + import trimesh + @configclass class FlatPatchSamplingCfg: diff --git a/source/isaaclab/isaaclab/terrains/terrain_generator_cfg.py b/source/isaaclab/isaaclab/terrains/terrain_generator_cfg.py index 6a3238c7cb4..03272e32a93 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_generator_cfg.py +++ b/source/isaaclab/isaaclab/terrains/terrain_generator_cfg.py @@ -18,16 +18,16 @@ from typing import Literal from isaaclab.utils import configclass +from isaaclab.utils.string import DeferredClass from .sub_terrain_cfg import SubTerrainBaseCfg -from .terrain_generator import TerrainGenerator @configclass class TerrainGeneratorCfg: """Configuration for the terrain generator.""" - class_type: type = TerrainGenerator + class_type: type | DeferredClass = DeferredClass("isaaclab.terrains.terrain_generator:TerrainGenerator") """The class to use for the terrain generator. Defaults to :class:`isaaclab.terrains.terrain_generator.TerrainGenerator`. diff --git a/source/isaaclab/isaaclab/terrains/terrain_importer_cfg.py b/source/isaaclab/isaaclab/terrains/terrain_importer_cfg.py index 7b42e06caaf..25f589a729c 100644 --- a/source/isaaclab/isaaclab/terrains/terrain_importer_cfg.py +++ b/source/isaaclab/isaaclab/terrains/terrain_importer_cfg.py @@ -9,9 +9,7 @@ from typing import TYPE_CHECKING, Literal import isaaclab.sim as sim_utils -from isaaclab.utils import configclass - -from .terrain_importer import TerrainImporter +from isaaclab.utils import DeferredClass, configclass if TYPE_CHECKING: from .terrain_generator_cfg import TerrainGeneratorCfg @@ -21,7 +19,7 @@ class TerrainImporterCfg: """Configuration for the terrain manager.""" - class_type: type = TerrainImporter + class_type: type | DeferredClass = DeferredClass("isaaclab.terrains.terrain_importer:TerrainImporter") """The class to use for the terrain importer. Defaults to :class:`isaaclab.terrains.terrain_importer.TerrainImporter`. diff --git a/source/isaaclab/isaaclab/terrains/trimesh/__init__.py b/source/isaaclab/isaaclab/terrains/trimesh/__init__.py index b27b7a92110..f6b77813c1a 100644 --- a/source/isaaclab/isaaclab/terrains/trimesh/__init__.py +++ b/source/isaaclab/isaaclab/terrains/trimesh/__init__.py @@ -12,18 +12,25 @@ efficient than the height-field representation, but it is not as flexible. """ -from .mesh_terrains_cfg import ( - MeshBoxTerrainCfg, - MeshFloatingRingTerrainCfg, - MeshGapTerrainCfg, - MeshInvertedPyramidStairsTerrainCfg, - MeshPitTerrainCfg, - MeshPlaneTerrainCfg, - MeshPyramidStairsTerrainCfg, - MeshRailsTerrainCfg, - MeshRandomGridTerrainCfg, - MeshRepeatedBoxesTerrainCfg, - MeshRepeatedCylindersTerrainCfg, - MeshRepeatedPyramidsTerrainCfg, - MeshStarTerrainCfg, +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "mesh_terrains_cfg": [ + "MeshBoxTerrainCfg", + "MeshFloatingRingTerrainCfg", + "MeshGapTerrainCfg", + "MeshInvertedPyramidStairsTerrainCfg", + "MeshPitTerrainCfg", + "MeshPlaneTerrainCfg", + "MeshPyramidStairsTerrainCfg", + "MeshRailsTerrainCfg", + "MeshRandomGridTerrainCfg", + "MeshRepeatedBoxesTerrainCfg", + "MeshRepeatedCylindersTerrainCfg", + "MeshRepeatedPyramidsTerrainCfg", + "MeshStarTerrainCfg", + ], + }, ) diff --git a/source/isaaclab/isaaclab/terrains/trimesh/mesh_terrains_cfg.py b/source/isaaclab/isaaclab/terrains/trimesh/mesh_terrains_cfg.py index 4247e21486b..3a037c419c5 100644 --- a/source/isaaclab/isaaclab/terrains/trimesh/mesh_terrains_cfg.py +++ b/source/isaaclab/isaaclab/terrains/trimesh/mesh_terrains_cfg.py @@ -7,9 +7,8 @@ from dataclasses import MISSING from typing import Literal -import isaaclab.terrains.trimesh.mesh_terrains as mesh_terrains -import isaaclab.terrains.trimesh.utils as mesh_utils_terrains from isaaclab.utils import configclass +from isaaclab.utils.string import DeferredClass from ..sub_terrain_cfg import SubTerrainBaseCfg @@ -22,14 +21,14 @@ class MeshPlaneTerrainCfg(SubTerrainBaseCfg): """Configuration for a plane mesh terrain.""" - function = mesh_terrains.flat_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:flat_terrain") @configclass class MeshPyramidStairsTerrainCfg(SubTerrainBaseCfg): """Configuration for a pyramid stair mesh terrain.""" - function = mesh_terrains.pyramid_stairs_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:pyramid_stairs_terrain") border_width: float = 0.0 """The width of the border around the terrain (in m). Defaults to 0.0. @@ -63,14 +62,14 @@ class MeshInvertedPyramidStairsTerrainCfg(MeshPyramidStairsTerrainCfg): This is the same as :class:`MeshPyramidStairsTerrainCfg` except that the steps are inverted. """ - function = mesh_terrains.inverted_pyramid_stairs_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:inverted_pyramid_stairs_terrain") @configclass class MeshRandomGridTerrainCfg(SubTerrainBaseCfg): """Configuration for a random grid mesh terrain.""" - function = mesh_terrains.random_grid_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:random_grid_terrain") grid_width: float = MISSING """The width of the grid cells (in m).""" @@ -93,7 +92,7 @@ class MeshRandomGridTerrainCfg(SubTerrainBaseCfg): class MeshRailsTerrainCfg(SubTerrainBaseCfg): """Configuration for a terrain with box rails as extrusions.""" - function = mesh_terrains.rails_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:rails_terrain") rail_thickness_range: tuple[float, float] = MISSING """The thickness of the inner and outer rails (in m).""" @@ -109,7 +108,7 @@ class MeshRailsTerrainCfg(SubTerrainBaseCfg): class MeshPitTerrainCfg(SubTerrainBaseCfg): """Configuration for a terrain with a pit that leads out of the pit.""" - function = mesh_terrains.pit_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:pit_terrain") pit_depth_range: tuple[float, float] = MISSING """The minimum and maximum height of the pit (in m).""" @@ -125,7 +124,7 @@ class MeshPitTerrainCfg(SubTerrainBaseCfg): class MeshBoxTerrainCfg(SubTerrainBaseCfg): """Configuration for a terrain with boxes (similar to a pyramid).""" - function = mesh_terrains.box_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:box_terrain") box_height_range: tuple[float, float] = MISSING """The minimum and maximum height of the box (in m).""" @@ -141,7 +140,7 @@ class MeshBoxTerrainCfg(SubTerrainBaseCfg): class MeshGapTerrainCfg(SubTerrainBaseCfg): """Configuration for a terrain with a gap around the platform.""" - function = mesh_terrains.gap_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:gap_terrain") gap_width_range: tuple[float, float] = MISSING """The minimum and maximum width of the gap (in m).""" @@ -154,7 +153,7 @@ class MeshGapTerrainCfg(SubTerrainBaseCfg): class MeshFloatingRingTerrainCfg(SubTerrainBaseCfg): """Configuration for a terrain with a floating ring around the center.""" - function = mesh_terrains.floating_ring_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:floating_ring_terrain") ring_width_range: tuple[float, float] = MISSING """The minimum and maximum width of the ring (in m).""" @@ -173,7 +172,7 @@ class MeshFloatingRingTerrainCfg(SubTerrainBaseCfg): class MeshStarTerrainCfg(SubTerrainBaseCfg): """Configuration for a terrain with a star pattern.""" - function = mesh_terrains.star_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:star_terrain") num_bars: int = MISSING """The number of bars per-side the star. Must be greater than 2.""" @@ -201,7 +200,7 @@ class ObjectCfg: height: float = MISSING """The height (along z) of the object (in m).""" - function = mesh_terrains.repeated_objects_terrain + function = DeferredClass("isaaclab.terrains.trimesh.mesh_terrains:repeated_objects_terrain") object_type: Literal["cylinder", "box", "cone"] | callable = MISSING """The type of object to generate. @@ -263,7 +262,7 @@ class ObjectCfg(MeshRepeatedObjectsTerrainCfg.ObjectCfg): degrees: bool = True """Whether the angle is in degrees. Defaults to True.""" - object_type = mesh_utils_terrains.make_cone + object_type = DeferredClass("isaaclab.terrains.trimesh.utils:make_cone") object_params_start: ObjectCfg = MISSING """The object curriculum parameters at the start of the curriculum.""" @@ -287,7 +286,7 @@ class ObjectCfg(MeshRepeatedObjectsTerrainCfg.ObjectCfg): degrees: bool = True """Whether the angle is in degrees. Defaults to True.""" - object_type = mesh_utils_terrains.make_box + object_type = DeferredClass("isaaclab.terrains.trimesh.utils:make_box") object_params_start: ObjectCfg = MISSING """The box curriculum parameters at the start of the curriculum.""" @@ -311,7 +310,7 @@ class ObjectCfg(MeshRepeatedObjectsTerrainCfg.ObjectCfg): degrees: bool = True """Whether the angle is in degrees. Defaults to True.""" - object_type = mesh_utils_terrains.make_cylinder + object_type = DeferredClass("isaaclab.terrains.trimesh.utils:make_cylinder") object_params_start: ObjectCfg = MISSING """The box curriculum parameters at the start of the curriculum.""" diff --git a/source/isaaclab/isaaclab/ui/widgets/manager_live_visualizer.py b/source/isaaclab/isaaclab/ui/widgets/manager_live_visualizer.py index 0b20b10dabc..cc0a2186c89 100644 --- a/source/isaaclab/isaaclab/ui/widgets/manager_live_visualizer.py +++ b/source/isaaclab/isaaclab/ui/widgets/manager_live_visualizer.py @@ -12,8 +12,6 @@ import numpy -import omni.kit.app - from isaaclab.managers import ManagerBase from isaaclab.sim import SimulationContext from isaaclab.utils import configclass @@ -196,6 +194,8 @@ def _set_debug_vis_impl(self, debug_vis: bool): if debug_vis: # if enabled create a subscriber for the post update event if it doesn't exist if not hasattr(self, "_debug_vis_handle") or self._debug_vis_handle is None: + import omni.kit.app + app_interface = omni.kit.app.get_app_interface() self._debug_vis_handle = app_interface.get_post_update_event_stream().create_subscription_to_pop( lambda event, obj=weakref.proxy(self): obj._debug_vis_callback(event) diff --git a/source/isaaclab/isaaclab/ui/xr_widgets/instruction_widget.py b/source/isaaclab/isaaclab/ui/xr_widgets/instruction_widget.py index 45d2d0a5daf..1ee9751bc00 100644 --- a/source/isaaclab/isaaclab/ui/xr_widgets/instruction_widget.py +++ b/source/isaaclab/isaaclab/ui/xr_widgets/instruction_widget.py @@ -10,7 +10,6 @@ import textwrap from typing import TYPE_CHECKING, Any, TypeAlias -import omni.kit.commands import omni.ui as ui from pxr import Gf @@ -194,6 +193,8 @@ def show_instruction( widget_args=[wrapped_text, {"font_size": font_size, "color": text_color}, width], ) + import omni.kit.commands + copied_prim = omni.kit.commands.execute( "CopyPrim", path_from=prim_path_source, diff --git a/source/isaaclab/isaaclab/utils/__init__.py b/source/isaaclab/isaaclab/utils/__init__.py index 1295715857f..cd4233ae1c9 100644 --- a/source/isaaclab/isaaclab/utils/__init__.py +++ b/source/isaaclab/isaaclab/utils/__init__.py @@ -5,15 +5,86 @@ """Sub-package containing utilities for common operations and helper functions.""" -from .array import * -from .buffers import * -from .configclass import configclass -from .dict import * -from .interpolation import * -from .logger import * -from .mesh import * -from .modifiers import * -from .string import * -from .timer import Timer -from .types import * -from .version import * +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=[], + submod_attrs={ + "array": [ + "TensorData", + "TENSOR_TYPES", + "TENSOR_TYPE_CONVERSIONS", + "convert_to_torch", + ], + "buffers": [ + "CircularBuffer", + "DelayBuffer", + "TimestampedBuffer", + "TimestampedBufferWarp", + ], + "configclass": [ + "configclass", + ], + "dict": [ + "class_to_dict", + "update_class_from_dict", + "dict_to_md5_hash", + "convert_dict_to_backend", + "update_dict", + "replace_slices_with_strings", + "replace_strings_with_slices", + "print_dict", + ], + "interpolation": [ + "LinearInterpolation", + ], + "logger": [ + "configure_logging", + "ColoredFormatter", + "RateLimitFilter", + ], + "mesh": [ + "create_trimesh_from_geom_mesh", + "create_trimesh_from_geom_shape", + "convert_faces_to_triangles", + "PRIMITIVE_MESH_TYPES", + ], + "modifiers": [ + "ModifierCfg", + "ModifierBase", + "DigitalFilter", + "DigitalFilterCfg", + "Integrator", + "IntegratorCfg", + "bias", + "clip", + "scale", + ], + "string": [ + "DeferredClass", + "to_camel_case", + "to_snake_case", + "string_to_slice", + "is_lambda_expression", + "callable_to_string", + "string_to_callable", + "resolve_matching_names", + "resolve_matching_names_values", + "find_unique_string_name", + "find_root_prim_path_from_regex", + ], + "timer": [ + "Timer", + "TimerError", + ], + "types": [ + "ArticulationActions", + ], + "version": [ + "has_kit", + "get_isaac_sim_version", + "compare_versions", + ], + }, +) diff --git a/source/isaaclab/isaaclab/utils/assets.py b/source/isaaclab/isaaclab/utils/assets.py index f0704d17f4c..a8b190ed444 100644 --- a/source/isaaclab/isaaclab/utils/assets.py +++ b/source/isaaclab/isaaclab/utils/assets.py @@ -13,15 +13,14 @@ .. _Omniverse Nucleus: https://docs.omniverse.nvidia.com/nucleus/latest/overview/overview.html """ -import asyncio import io import logging import os +import posixpath import tempfile -import time +from pathlib import Path from typing import Literal - -import omni.client +from urllib.parse import urlparse logger = logging.getLogger(__name__) @@ -52,6 +51,8 @@ def _parse_kit_asset_root() -> str: ISAACLAB_NUCLEUS_DIR: str = f"{ISAAC_NUCLEUS_DIR}/IsaacLab" """Path to the ``Isaac/IsaacLab`` directory on the NVIDIA Nucleus Server.""" +USD_EXTENSIONS = {".usd", ".usda", ".usdz"} + def check_file_path(path: str) -> Literal[0, 1, 2]: """Checks if a file exists on the Nucleus Server or locally. @@ -68,8 +69,9 @@ def check_file_path(path: str) -> Literal[0, 1, 2]: """ if os.path.isfile(path): return 1 - # we need to convert backslash to forward slash on Windows for omni.client API - elif omni.client.stat(path.replace(os.sep, "/"))[0] == omni.client.Result.OK: + import omni.client + + if omni.client.stat(path.replace(os.sep, "/"))[0] == omni.client.Result.OK: return 2 else: return 0 @@ -102,6 +104,8 @@ def retrieve_file_path(path: str, download_dir: str | None = None, force_downloa if file_status == 1: return os.path.abspath(path) elif file_status == 2: + import omni.client + # resolve download directory if download_dir is None: download_dir = tempfile.gettempdir() @@ -110,16 +114,38 @@ def retrieve_file_path(path: str, download_dir: str | None = None, force_downloa # create download directory if it does not exist if not os.path.exists(download_dir): os.makedirs(download_dir) - # download file in temp directory using os - file_name = os.path.basename(omni.client.break_url(path.replace(os.sep, "/")).path) - target_path = os.path.join(download_dir, file_name) - # check if file already exists locally - if not os.path.isfile(target_path) or force_download: - # copy file to local machine - result = omni.client.copy(path.replace(os.sep, "/"), target_path, omni.client.CopyBehavior.OVERWRITE) - if result != omni.client.Result.OK and force_download: - raise RuntimeError(f"Unable to copy file: '{path}'. Is the Nucleus Server running?") - return os.path.abspath(target_path) + # recursive download: mirror remote tree under download_dir + remote_url = path.replace(os.sep, "/") + to_visit = [remote_url] + visited = set() + local_root = None + + while to_visit: + cur_url = to_visit.pop() + if cur_url in visited: + continue + visited.add(cur_url) + + cur_rel = urlparse(cur_url).path.lstrip("/") + target_path = os.path.join(download_dir, cur_rel) + os.makedirs(os.path.dirname(target_path), exist_ok=True) + + if not os.path.isfile(target_path) or force_download: + result = omni.client.copy(cur_url, target_path, omni.client.CopyBehavior.OVERWRITE) + if result != omni.client.Result.OK and force_download: + raise RuntimeError(f"Unable to copy file: '{cur_url}'. Is the Nucleus Server running?") + + if local_root is None: + local_root = target_path + + # recurse into USD dependencies and referenced assets + if Path(target_path).suffix.lower() in USD_EXTENSIONS: + for ref in _find_usd_references(target_path): + ref_url = _resolve_reference_url(cur_url, ref) + if ref_url and ref_url not in visited: + to_visit.append(ref_url) + + return os.path.abspath(local_root) else: raise FileNotFoundError(f"Unable to find the file: {path}") @@ -142,83 +168,117 @@ def read_file(path: str) -> io.BytesIO: with open(path, "rb") as f: return io.BytesIO(f.read()) elif file_status == 2: + import omni.client + file_content = omni.client.read_file(path.replace(os.sep, "/"))[2] return io.BytesIO(memoryview(file_content).tobytes()) else: raise FileNotFoundError(f"Unable to find the file: {path}") -""" -Nucleus Connection. -""" - - -def check_usd_path_with_timeout(usd_path: str, timeout: float = 300, log_interval: float = 30) -> bool: - """Checks whether the given USD file path is available on the NVIDIA Nucleus server. - - This function synchronously runs an asynchronous USD path availability check, - logging progress periodically until it completes. The file is available on the server - if the HTTP status code is 200. Otherwise, the file is not available on the server. - - This is useful for checking server responsiveness before attempting to load a remote - asset. It will block execution until the check completes or times out. - - Args: - usd_path: The remote USD file path to check. - timeout: Maximum time (in seconds) to wait for the server check. - log_interval: Interval (in seconds) at which progress is logged. - - Returns: - Whether the given USD path is available on the server. - """ - start_time = time.time() - loop = asyncio.get_event_loop() - - coroutine = _is_usd_path_available(usd_path, timeout) - task = asyncio.ensure_future(coroutine) - - next_log_time = start_time + log_interval - - first_log = True - while not task.done(): - now = time.time() - if now >= next_log_time: - elapsed = int(now - start_time) - if first_log: - logger.warning(f"Checking server availability for USD path: {usd_path} (timeout: {timeout}s)") - first_log = False - logger.warning(f"Waiting for server response... ({elapsed}s elapsed)") - next_log_time += log_interval - loop.run_until_complete(asyncio.sleep(0.1)) # Yield to allow async work - - return task.result() +def _is_downloadable_asset(path: str) -> bool: + """Return True for USD or other asset types we mirror locally (textures, etc.).""" + clean = path.split("?", 1)[0].split("#", 1)[0] + suffix = Path(clean).suffix.lower() + if suffix == ".mdl": + # MDL modules (OmniPBR.mdl, OmniSurface.mdl, ...) come from MDL search paths + return False + if not suffix: + return False + if suffix not in {".usd", ".usda", ".usdz", ".png", ".jpg", ".jpeg", ".exr", ".hdr", ".tif", ".tiff"}: + return False + return True -""" -Helper functions. -""" - - -async def _is_usd_path_available(usd_path: str, timeout: float) -> bool: - """Checks whether the given USD path is available on the Omniverse Nucleus server. - - This function is a asynchronous routine to check the availability of the given USD path on - the Omniverse Nucleus server. It will return True if the USD path is available on the server, - False otherwise. - Args: - usd_path: The remote or local USD file path to check. - timeout: Timeout in seconds for the async stat call. +def _find_usd_references(local_usd_path: str) -> set[str]: + """Use Sdf API to collect referenced assets from a USD layer.""" + from pxr import Sdf - Returns: - Whether the given USD path is available on the server. - """ try: - result, _ = await asyncio.wait_for(omni.client.stat_async(usd_path), timeout=timeout) - return result == omni.client.Result.OK - except asyncio.TimeoutError: - logger.warning(f"Timed out after {timeout}s while checking for USD: {usd_path}") - return False - except Exception as ex: - logger.warning(f"Exception during USD file check: {type(ex).__name__}: {ex}") - return False + layer = Sdf.Layer.FindOrOpen(local_usd_path) + except Exception: + logger.warning("Failed to open USD layer: %s", local_usd_path, exc_info=True) + return set() + + if layer is None: + return set() + + refs: set[str] = set() + + # Sublayers + for sub_path in getattr(layer, "subLayerPaths", []) or []: + if sub_path and _is_downloadable_asset(sub_path): + refs.add(str(sub_path)) + + def _walk_prim(prim_spec: Sdf.PrimSpec) -> None: + # References + ref_list = prim_spec.referenceList + for field in ("addedItems", "prependedItems", "appendedItems", "explicitItems"): + items = getattr(ref_list, field, None) + if not items: + continue + for ref in items: + asset_path = getattr(ref, "assetPath", None) + if asset_path and _is_downloadable_asset(asset_path): + refs.add(str(asset_path)) + + # Payloads + payload_list = prim_spec.payloadList + for field in ("addedItems", "prependedItems", "appendedItems", "explicitItems"): + items = getattr(payload_list, field, None) + if not items: + continue + for payload in items: + asset_path = getattr(payload, "assetPath", None) + if asset_path and _is_downloadable_asset(asset_path): + refs.add(str(asset_path)) + + # AssetPath-valued attributes (this is where OmniPBR.mdl, textures, etc. show up) + for attr_spec in prim_spec.attributes.values(): + default = attr_spec.default + if isinstance(default, Sdf.AssetPath): + if default.path and _is_downloadable_asset(default.path): + refs.add(default.path) + elif isinstance(default, Sdf.AssetPathArray): + for ap in default: + if ap.path and _is_downloadable_asset(ap.path): + refs.add(ap.path) + + # Variants - each variant set can have multiple variants with their own prim content + for variant_set_spec in prim_spec.variantSets.values(): + for variant_spec in variant_set_spec.variants.values(): + variant_prim_spec = variant_spec.primSpec + if variant_prim_spec is not None: + _walk_prim(variant_prim_spec) + + for child in prim_spec.nameChildren.values(): + _walk_prim(child) + + for root_prim in layer.rootPrims.values(): + _walk_prim(root_prim) + + return refs + + +def _resolve_reference_url(base_url: str, ref: str) -> str: + """Resolve a USD asset reference against a base URL (http/local).""" + ref = ref.strip() + if not ref: + return ref + + parsed_ref = urlparse(ref) + if parsed_ref.scheme: + return ref + + base = urlparse(base_url) + if base.scheme == "": + base_dir = os.path.dirname(base_url) + return os.path.normpath(os.path.join(base_dir, ref)) + + base_dir = posixpath.dirname(base.path) + if ref.startswith("/"): + new_path = posixpath.normpath(ref) + else: + new_path = posixpath.normpath(posixpath.join(base_dir, ref)) + return f"{base.scheme}://{base.netloc}{new_path}" diff --git a/source/isaaclab/isaaclab/utils/backend_utils.py b/source/isaaclab/isaaclab/utils/backend_utils.py index 7ed5444efcd..0443a05dfcd 100644 --- a/source/isaaclab/isaaclab/utils/backend_utils.py +++ b/source/isaaclab/isaaclab/utils/backend_utils.py @@ -5,6 +5,7 @@ import importlib import logging +from isaaclab.sim import SimulationContext logger = logging.getLogger(__name__) @@ -35,38 +36,40 @@ def register(cls, name: str, sub_class) -> None: def __new__(cls, *args, **kwargs): """Create a new instance of an implementation based on the backend.""" + from isaaclab.sim import SimulationContext - # TODO: Make the backend configurable. - backend = "physx" + physics_mgr = SimulationContext.instance().physics_manager + mgr_name = physics_mgr.__name__ # e.g. "PhysxManager", "NewtonManager" + physics_backend = mgr_name.replace("Manager", "").lower() # "physx", "newton" if cls == FactoryBase: raise TypeError("FactoryBase cannot be instantiated directly. Please subclass it.") # If backend is not in registry, try to import it and register the class. # This is done to only import the module once. - if backend not in cls._registry: - # Construct the module name from the backend and the determined subpath. - module_name = f"isaaclab_{backend}.{cls._module_subpath}" + if physics_backend not in cls._registry: + # Construct the module name from the physics_backend and the determined subpath. + module_name = f"isaaclab_{physics_backend}.{cls._module_subpath}" try: module = importlib.import_module(module_name) module_class = getattr(module, cls.__name__) # Manually register the class - cls.register(backend, module_class) + cls.register(physics_backend, module_class) except ImportError as e: raise ValueError( - f"Could not import module for backend {backend!r} for factory {cls.__name__}. " + f"Could not import module for backend {physics_backend!r} for factory {cls.__name__}. " f"Attempted to import from '{module_name}'.\n" f"Original error: {e}" ) from e # Now check registry again. The import should have registered the class. try: - impl = cls._registry[backend] + impl = cls._registry[physics_backend] except KeyError: available = list(cls.get_registry_keys()) raise ValueError( - f"Unknown backend {backend!r} for {cls.__name__}. " + f"Unknown backend {physics_backend!r} for {cls.__name__}. " f"A module was found at '{module_name}', but it did not contain a class with the name {cls.__name__}.\n" f"Currently available backends: {available}." ) from None diff --git a/source/isaaclab/isaaclab/utils/buffers/__init__.py b/source/isaaclab/isaaclab/utils/buffers/__init__.py index fbba57b89fc..9c4987e0687 100644 --- a/source/isaaclab/isaaclab/utils/buffers/__init__.py +++ b/source/isaaclab/isaaclab/utils/buffers/__init__.py @@ -5,7 +5,14 @@ """Sub-module containing different buffers.""" -from .circular_buffer import CircularBuffer -from .delay_buffer import DelayBuffer -from .timestamped_buffer import TimestampedBuffer -from .timestamped_buffer_warp import TimestampedBufferWarp +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "circular_buffer": ["CircularBuffer"], + "delay_buffer": ["DelayBuffer"], + "timestamped_buffer": ["TimestampedBuffer"], + "timestamped_buffer_warp": ["TimestampedBufferWarp"], + }, +) diff --git a/source/isaaclab/isaaclab/utils/dict.py b/source/isaaclab/isaaclab/utils/dict.py index de2062d6697..4aa78145994 100644 --- a/source/isaaclab/isaaclab/utils/dict.py +++ b/source/isaaclab/isaaclab/utils/dict.py @@ -11,10 +11,7 @@ from collections.abc import Iterable, Mapping, Sized from typing import Any -import torch - -from .array import TENSOR_TYPE_CONVERSIONS, TENSOR_TYPES -from .string import callable_to_string, string_to_callable, string_to_slice +from .string import DeferredClass, callable_to_string, string_to_callable, string_to_slice """ Dictionary <-> Class operations. @@ -42,10 +39,9 @@ def class_to_dict(obj: object) -> dict[str, Any]: # convert object to dictionary if isinstance(obj, dict): obj_dict = obj - elif isinstance(obj, torch.Tensor): - # We have to treat torch tensors specially because `torch.tensor.__dict__` returns an empty - # dict, which would mean that a torch.tensor would be stored as an empty dict. Instead we - # want to store it directly as the tensor. + elif (getattr(type(obj), "__module__", "") or "").startswith(("torch", "numpy", "warp")): + # Tensor-like objects (torch.Tensor, np.ndarray, wp.array) have __dict__ that returns + # an empty/misleading dict. Return them directly instead of recursing. return obj elif hasattr(obj, "__dict__"): obj_dict = obj.__dict__ @@ -142,8 +138,10 @@ def update_class_from_dict(obj, data: dict[str, Any], _ns: str = "") -> None: # -- 3) callable attribute → resolve string -------------- elif callable(obj_mem): - # update function name - value = string_to_callable(value) + try: + value = string_to_callable(value) + except (ImportError, ValueError): + value = DeferredClass(value) # -- 4) simple scalar / explicit None --------------------- elif value is None or isinstance(value, type(obj_mem)): @@ -227,6 +225,8 @@ def convert_dict_to_backend( The updated dict with the data converted to the desired backend. """ # THINK: Should we also support converting to a specific device, e.g. "cuda:0"? + from .array import TENSOR_TYPE_CONVERSIONS, TENSOR_TYPES + # Check the backend is valid. if backend not in TENSOR_TYPE_CONVERSIONS: raise ValueError(f"Unknown backend '{backend}'. Supported backends are 'numpy', 'torch', and 'warp'.") diff --git a/source/isaaclab/isaaclab/utils/helpers.py b/source/isaaclab/isaaclab/utils/helpers.py new file mode 100644 index 00000000000..54020d2b05b --- /dev/null +++ b/source/isaaclab/isaaclab/utils/helpers.py @@ -0,0 +1,301 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Utility functions for functions.""" + +from __future__ import annotations + +import functools +import inspect +import logging +import torch +import warnings +from collections.abc import Callable + +import warp as wp + +logger = logging.getLogger(__name__) +warnings.simplefilter("once", UserWarning) +logging.captureWarnings(True) + + +def deprecated( + *replacement_function_names: str, + message: str = "", + since: str | None = None, + remove_in: str | None = None, +): + """A decorator to mark functions as deprecated. + + It will result in a warning being emitted when the function is used. + + Args: + replacement_function_names: The names of the functions to use instead. + message: A custom message to append to the warning. + since: The version in which the function was deprecated. + remove_in: The version in which the function will be removed. + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + # Form deprecation message. + deprecation_message = f"Call to deprecated function '{func.__name__}'." + # Add version information if provided. + if since: + deprecation_message += f" It was deprecated in version {since}." + if remove_in: + deprecation_message += f" It will be removed in version {remove_in}." + else: + deprecation_message += " It will be removed in a future version." + # Add replacement function information if provided. + if replacement_function_names: + deprecation_message += f" Use {', '.join(replacement_function_names)} instead." + # Add custom message if provided. + if message: + deprecation_message += f" {message}" + + # Emit warning. + warnings.warn( + deprecation_message, + UserWarning, + ) + # Call the original function. + return func(*args, **kwargs) + + return wrapper + + return decorator + + +def warn_overhead_cost( + *replacement_function_names: str, + message: str = "", +): + """A decorator to mark functions as having a high overhead cost. + + It will result in a warning being emitted when the function is used. + + Args: + replacement_function_names: The names of the functions to use instead. + message: A custom message to append to the warning. + """ + + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + # Form deprecation message. + warning_message = f"Call to '{func.__name__}' which is a high overhead operation." + # Add replacement function information if provided. + if replacement_function_names: + warning_message += f" Use {', '.join(replacement_function_names)} instead." + # Add custom message if provided. + if message: + warning_message += f" {message}" + + # Emit warning. + warnings.warn( + warning_message, + UserWarning, + ) + # Call the original function. + return func(*args, **kwargs) + + return wrapper + + return decorator + + +def _analyze_and_convert_args(self, func, *args, **kwargs): + """A helper to analyze and convert arguments from PyTorch to Warp.""" + sig = inspect.signature(func) + bound_args = sig.bind(self, *args, **kwargs) + bound_args.apply_defaults() + arguments = bound_args.arguments + + # -- Device conversion + device = getattr(self, "device", "cpu") + + # -- Tensor conversion + spec = getattr(func, "_torch_frontend_spec", {}) + tensor_args = spec.get("tensor_args", {}) + first_torch_tensor = None + for arg_name in tensor_args: + arg_value = arguments.get(arg_name) + if isinstance(arg_value, torch.Tensor): + if first_torch_tensor is None: + first_torch_tensor = arg_value + tensor = arguments[arg_name] + dtype = tensor_args[arg_name] + arguments[arg_name] = wp.from_torch(tensor, dtype=dtype) + + # -- Mask conversion + mask_configs = { + "env_mask": {"id_arg": "env_ids", "shape_attrs": ["num_instances", "num_envs"]}, + "joint_mask": {"id_arg": "joint_ids", "shape_attrs": ["num_joints"]}, + "body_mask": {"id_arg": "body_ids", "shape_attrs": ["num_bodies"]}, + } + + for mask_name, config in mask_configs.items(): + id_arg_name = config["id_arg"] + if mask_name in sig.parameters and id_arg_name in arguments and arguments[id_arg_name] is not None: + indices = arguments.pop(id_arg_name) + shape_val = 0 + for attr in config["shape_attrs"]: + val = getattr(self, attr, None) + if val is not None: + shape_val = val + break + if shape_val == 0: + raise ValueError( + f"Cannot convert '{id_arg_name}' to '{mask_name}'. The instance is missing one of the " + f"following attributes: {config['shape_attrs']}." + ) + + mask_torch = torch.zeros(shape_val, dtype=torch.bool, device=device) + + if isinstance(indices, slice): + mask_torch[indices] = True + elif isinstance(indices, (list, tuple, torch.Tensor)): + mask_torch[torch.as_tensor(indices, device=device)] = True + else: + raise TypeError(f"Unsupported type for indices '{type(indices)}'.") + + arguments[mask_name] = wp.from_torch(mask_torch) + + arguments.pop("self") + return arguments + + +def _convert_output_to_torch(data): + """Recursively convert warp arrays in a data structure to torch tensors.""" + if isinstance(data, wp.array): + return wp.to_torch(data) + elif isinstance(data, (list, tuple)): + return type(data)(_convert_output_to_torch(item) for item in data) + elif isinstance(data, dict): + return {key: _convert_output_to_torch(value) for key, value in data.items()} + else: + return data + + +def torch_frontend_method(tensor_args: dict[str, any] | None = None, *, convert_output: bool = False): + """A method decorator to specify tensor conversion rules for the torch frontend. + + This decorator attaches metadata to a method, which is then used by the + `torch_frontend_class` decorator to apply the correct data conversions. + + Args: + tensor_args: A dictionary mapping tensor argument names to their + target `warp.dtype`. Defaults to None. + convert_output: If True, the output of the decorated function will be + converted from warp arrays to torch tensors. Defaults to False. + + Example: + >>> @torch_frontend_class + ... class MyAsset: + ... @torch_frontend_method({"root_state": wp.transformf}) + ... def write_root_state_to_sim(self, root_state: wp.array, env_mask: wp.array | None = None): + ... pass + ... + ... @torch_frontend_method(convert_output=True) + ... def get_root_state(self) -> wp.array: + ... pass + """ + if tensor_args is None: + tensor_args = {} + + def decorator(func: Callable) -> Callable: + setattr(func, "_torch_frontend_spec", {"tensor_args": tensor_args, "convert_output": convert_output}) + return func + + return decorator + + +def torch_frontend_class(cls=None, *, indices_arg: str = "env_ids", mask_arg: str = "env_mask"): + """A class decorator to add a PyTorch frontend to a class that uses a Warp backend. + + This decorator patches the ``__init__`` method of a class. After the original ``__init__`` is called, + it checks for a `frontend` attribute on the instance. If `self.frontend` is "torch", it inspects + the class for methods decorated with `@torch_frontend_method` and wraps them to make them + accept PyTorch tensors. + + The wrapped methods will: + - Convert specified PyTorch tensor arguments to Warp arrays based on the rules defined + in the `@torch_frontend_method` decorator. + - Convert an argument with indices (e.g., `env_ids`) to a boolean mask (e.g., `env_mask`). + + If `self.frontend` is not "torch", the class's methods remain unchanged, ensuring zero + overhead for the default Warp backend. + + Args: + cls: The class to decorate. This is handled automatically by Python. + indices_arg: The name of the argument that may contain indices for conversion to a mask. + Defaults to "env_ids". + mask_arg: The name of the argument that will receive the generated boolean mask. + Defaults to "env_mask". + + Example: + >>> import warp as wp + >>> import torch + >>> + >>> @torch_frontend_class + ... class MyAsset: + ... def __init__(self, num_envs, device, frontend="warp"): + ... self.num_instances = num_envs + ... self.device = device + ... self.frontend = frontend + ... wp.init() + ... + ... @torch_frontend_method({"root_state": wp.transformf}) + ... def write_root_state_to_sim(self, root_state: wp.array, env_mask: wp.array | None = None): + ... print(f"root_state type: {type(root_state)}") + ... if env_mask: + ... print(f"env_mask type: {type(env_mask)}") + ... + >>> # -- Using warp frontend (no overhead) + >>> asset_wp = MyAsset(num_envs=4, device="cpu", frontend="warp") + >>> root_state_wp = wp.zeros(4, dtype=wp.transformf) + >>> asset_wp.write_root_state_to_sim(root_state_wp) + root_state type: + >>> + >>> # -- Using torch frontend (methods are patched) + >>> asset_torch = MyAsset(num_envs=4, device="cpu", frontend="torch") + >>> root_state_torch = torch.rand(4, 7) + >>> asset_torch.write_root_state_to_sim(root_state_torch, env_ids=[0, 2]) + root_state type: + env_mask type: + """ + # This allows using the decorator with or without parentheses: + # @torch_frontend_class or @torch_frontend_class(indices_arg="...") + if cls is None: + return functools.partial(torch_frontend_class, indices_arg=indices_arg, mask_arg=mask_arg) + + original_init = cls.__init__ + + @functools.wraps(original_init) + def new_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + + if getattr(self, "frontend", "warp") == "torch": + for name, method in inspect.getmembers(cls, predicate=inspect.isfunction): + if hasattr(method, "_torch_frontend_spec"): + spec = getattr(method, "_torch_frontend_spec") + convert_output = spec.get("convert_output", False) + + @functools.wraps(method) + def adapted_method_wrapper(self, *args, method=method, **kwargs): + converted_args = _analyze_and_convert_args(self, method, *args, **kwargs) + output = method(self, **converted_args) + + if convert_output: + return _convert_output_to_torch(output) + else: + return output + + setattr(self, name, adapted_method_wrapper.__get__(self, cls)) + + cls.__init__ = new_init + return cls diff --git a/source/isaaclab/isaaclab/utils/string.py b/source/isaaclab/isaaclab/utils/string.py index dc1cdaf5347..86e4af36d98 100644 --- a/source/isaaclab/isaaclab/utils/string.py +++ b/source/isaaclab/isaaclab/utils/string.py @@ -393,6 +393,65 @@ def find_unique_string_name(initial_name: str, is_unique_fn: Callable[[str], boo return result +class DeferredClass: + """A callable wrapper that lazily resolves a class from a ``"module:ClassName"`` string. + + This allows configuration classes to specify their associated implementation class as a + string path, deferring the import until the class is first called. Consumers can call + instances of this class exactly like the resolved class itself -- no ``isinstance`` checks + or ``string_to_callable`` calls are needed at the call site. + + Example usage in a cfg class: + + .. code-block:: python + + class ArticulationCfg(AssetBaseCfg): + class_type: type | DeferredClass = DeferredClass( + "isaaclab.assets.articulation.articulation:Articulation" + ) + + The resolved class is cached after the first call. + """ + + def __init__(self, class_path: str): + """Initialize with a ``"module:ClassName"`` string path. + + Args: + class_path: Dotted module path and class name separated by ``:``. + For example: ``"isaaclab.assets.articulation.articulation:Articulation"``. + """ + self._class_path = class_path + self._resolved: type | None = None + + def resolve(self) -> type: + """Resolve and return the class, importing it if not yet done.""" + if self._resolved is None: + self._resolved = string_to_callable(self._class_path) + return self._resolved + + def __call__(self, *args, **kwargs): + """Instantiate the resolved class, forwarding all arguments.""" + return self.resolve()(*args, **kwargs) + + @property + def __name__(self) -> str: + """Return the class name portion of the path, resolving lazily if needed.""" + if self._resolved is not None: + return self._resolved.__name__ + return self._class_path.split(":")[-1] + + @property + def __module__(self) -> str: + """Return the module portion of the path, resolving lazily if needed.""" + if self._resolved is not None: + return self._resolved.__module__ + return self._class_path.split(":")[0] + + def __repr__(self) -> str: + status = repr(self._resolved) if self._resolved is not None else f"'{self._class_path}'" + return f"DeferredClass({status})" + + def find_root_prim_path_from_regex(prim_path_regex: str) -> tuple[str, int]: """Find the first prim above the regex pattern prim and its position. Args: diff --git a/source/isaaclab/isaaclab/utils/timer.py b/source/isaaclab/isaaclab/utils/timer.py index 8851710e064..39e0ceef42f 100644 --- a/source/isaaclab/isaaclab/utils/timer.py +++ b/source/isaaclab/isaaclab/utils/timer.py @@ -10,7 +10,9 @@ import math import time from contextlib import ContextDecorator -from typing import Any, ClassVar +from typing import Any, ClassVar, Literal + +import warp as wp class TimerError(Exception): @@ -67,16 +69,21 @@ class Timer(ContextDecorator): This dictionary logs the timer information. The keys are the names given to the timer class at its initialization. If no :attr:`name` is passed to the constructor, no time is recorded in the dictionary. - - In each of the dictionaries, we store the following information: - - last: The last elapsed time - - m2: The sum of squares of differences from the current mean (Intermediate value in Welford's Algorithm) - - mean: The mean of the elapsed time - - std: The standard deviation of the elapsed time - - n: The number of samples """ - def __init__(self, msg: str | None = None, name: str | None = None): + enable = True + """Whether to enable the timer.""" + + enable_display_output = True + """Whether to enable the display output.""" + + def __init__( + self, + msg: str | None = None, + name: str | None = None, + enable: bool = True, + format: Literal["s", "ms", "us", "ns"] = "s", + ): """Initializes the timer. Args: @@ -84,12 +91,26 @@ def __init__(self, msg: str | None = None, name: str | None = None): class in a context manager. Defaults to None. name: The name to use for logging times in a global dictionary. Defaults to None. + enable: Whether to enable the timer. Defaults to True. + format: The format to use for the elapsed time. Defaults to "s". """ self._msg = msg self._name = name self._start_time = None self._stop_time = None self._elapsed_time = None + self._enable = enable if Timer.enable else False + self._format = format + + # Check if the format is valid + assert format in ["s", "ms", "us", "ns"], f"Invalid format, {format} is not in [s, ms, us, ns]" + # Convert the format to a multiplier + self._multiplier = { + "s": 1.0, + "ms": 1000.0, + "us": 1000000.0, + "ns": 1000000000.0, + }[format] # Online welford's algorithm to compute the mean and std of the elapsed time self._mean = 0.0 @@ -103,7 +124,7 @@ def __str__(self) -> str: Returns: A string containing the elapsed time. """ - return f"{self.time_elapsed:0.6f} seconds" + return f"{(self.time_elapsed * self._multiplier):0.6f} {self._format}" """ Properties @@ -129,6 +150,9 @@ def total_run_time(self) -> float: def start(self): """Start timing.""" + if not self._enable: + return + if self._start_time is not None: raise TimerError("Timer is running. Use .stop() to stop it") @@ -136,17 +160,25 @@ def start(self): def stop(self): """Stop timing.""" + if not self._enable: + return + if self._start_time is None: raise TimerError("Timer is not running. Use .start() to start it") + # Synchronize the device to make sure we time the whole operation + wp.synchronize_device() + + # Get the elapsed time self._stop_time = time.perf_counter() self._elapsed_time = self._stop_time - self._start_time self._start_time = None - if self._name: + if (self._name is not None) and (self._enable): # Update the welford's algorithm - self._update_welford(self._elapsed_time) - # Store the information in the global dictionary + self.update_welford(self._elapsed_time) + + # Update the timing info Timer.timing_info[self._name] = { "last": self._elapsed_time, "m2": self._m2, @@ -156,21 +188,12 @@ def stop(self): } """ - Internal helpers. + Online welford's algorithm to compute the mean and std of the elapsed time """ - def _update_welford(self, value: float): - """Update the welford's algorithm with a new value. - - This algorithm computes the mean and standard deviation of the elapsed time in a numerically stable way. - It may become numerically unstable if n becomes very large. + def update_welford(self, value: float): + """Update the welford's algorithm with a new value.""" - Note: We use the global dictionary to retrieve the current values. We do this to make the timer - instances stateful. - - Args: - value: The new value to add to the statistics. - """ try: self._n = Timer.timing_info[self._name]["n"] + 1 delta = value - Timer.timing_info[self._name]["mean"] @@ -182,11 +205,8 @@ def _update_welford(self, value: float): self._mean = value self._m2 = 0.0 - # Update the std (sample standard deviation with Bessel's correction) - if self._n > 1: - self._std = math.sqrt(self._m2 / (self._n - 1)) - else: - self._std = 0.0 + # Update the std + self._std = math.sqrt(self._m2 / self._n) """ Context managers @@ -201,8 +221,15 @@ def __exit__(self, *exc_info: Any): """Stop timing.""" self.stop() # print message - if self._msg is not None: - print(self._msg, f": {self._elapsed_time:0.6f} seconds") + if self._enable: + if (self._msg is not None) and (Timer.enable_display_output): + print( + self._msg, + f"Last: {(self._elapsed_time * self._multiplier):0.6f} {self._format}, " + f"Mean: {(self._mean * self._multiplier):0.6f} {self._format}, " + f"Std: {(self._std * self._multiplier):0.6f} {self._format}, " + f"N: {self._n}", + ) """ Static Methods @@ -224,25 +251,12 @@ def get_timer_info(name: str) -> float: """ if name not in Timer.timing_info: raise TimerError(f"Timer {name} does not exist") - # Non-breaking change: return the last elapsed time return Timer.timing_info.get(name)["last"] @staticmethod def get_timer_statistics(name: str) -> dict[str, float]: - """Retrieves the statistics of the time logged in the global dictionary based on name. - - Returns a dictionary containing the mean, standard deviation, and number of samples as - well as the last measurement. Available keys are: - - mean: The mean of the elapsed time - - std: The standard deviation of the elapsed time - - n: The number of samples - - last: The last elapsed time - - Args: - name: Name of the the entry to be retrieved. - - Raises: - TimerError: If name doesn't exist in the log. + """Retrieves the time logged in the global dictionary + based on name. Returns: A dictionary containing the time logged for all timers. @@ -251,6 +265,6 @@ def get_timer_statistics(name: str) -> dict[str, float]: if name not in Timer.timing_info: raise TimerError(f"Timer {name} does not exist") - keys = ["mean", "std", "n", "last"] + keys = ["mean", "std", "n"] - return {k: Timer.timing_info.get(name)[k] for k in keys} + return {k: Timer.timing_info[name][k] for k in keys} diff --git a/source/isaaclab/isaaclab/utils/version.py b/source/isaaclab/isaaclab/utils/version.py index d32a458f480..5fc92900cc5 100644 --- a/source/isaaclab/isaaclab/utils/version.py +++ b/source/isaaclab/isaaclab/utils/version.py @@ -30,6 +30,16 @@ def has_kit() -> bool: return False +@functools.lru_cache(maxsize=1) +def has_kit() -> bool: + """Check if the Omniverse/Isaac Sim app is launched and running. Result is cached.""" + try: + import omni.usd + return True + except (ImportError, AttributeError): + return False + + @functools.lru_cache(maxsize=1) def get_isaac_sim_version() -> Version: """Get the Isaac Sim version as a Version object, cached for performance. diff --git a/source/isaaclab/isaaclab/utils/warp/__init__.py b/source/isaaclab/isaaclab/utils/warp/__init__.py index 12c9a20f8bb..b4d82b8ed7f 100644 --- a/source/isaaclab/isaaclab/utils/warp/__init__.py +++ b/source/isaaclab/isaaclab/utils/warp/__init__.py @@ -5,5 +5,6 @@ """Sub-module containing operations based on warp.""" -from . import fabric # noqa: F401 +from . import fabric, utils # noqa: F401 from .ops import convert_to_warp_mesh, raycast_dynamic_meshes, raycast_mesh, raycast_single_mesh +from .utils import make_mask_from_torch_ids diff --git a/source/isaaclab/isaaclab/utils/warp/utils.py b/source/isaaclab/isaaclab/utils/warp/utils.py new file mode 100644 index 00000000000..0a7d60080b1 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/warp/utils.py @@ -0,0 +1,30 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Warp utility functions.""" + +from collections.abc import Sequence + +import torch +import warp as wp + + +def make_mask_from_torch_ids(num_envs: int, env_ids: Sequence[int] | torch.Tensor, device: str) -> wp.array: + """Create a warp boolean mask array from environment indices. + + Args: + num_envs: Total number of environments. + env_ids: Sequence or tensor of environment indices to set as True. + device: Device to create the mask on. + + Returns: + A warp array of shape (num_envs,) with dtype wp.bool, where True indicates + the environment is in env_ids. + """ + # Create a torch bool tensor + mask = torch.zeros(num_envs, dtype=torch.bool, device=device) + mask[env_ids] = True + # Convert to warp array + return wp.from_torch(mask.contiguous(), dtype=wp.bool) diff --git a/source/isaaclab/setup.py b/source/isaaclab/setup.py index fc200e614f5..63058f77cad 100644 --- a/source/isaaclab/setup.py +++ b/source/isaaclab/setup.py @@ -23,6 +23,7 @@ "onnx>=1.18.0", # 1.16.2 throws access violation on Windows "prettytable==3.3.0", "toml", + "lazy_loader", # devices "hidapi==0.14.0.post2", # reinforcement learning @@ -41,6 +42,8 @@ "pillow==12.0.0", # livestream "starlette==0.49.1", + # assets + "omniverseclient", # testing "pytest", "pytest-mock", diff --git a/source/isaaclab/test/assets/check_external_force.py b/source/isaaclab/test/assets/check_external_force.py index f038b907a1f..e07ca1bac38 100644 --- a/source/isaaclab/test/assets/check_external_force.py +++ b/source/isaaclab/test/assets/check_external_force.py @@ -103,7 +103,7 @@ def main(): robot.write_root_velocity_to_sim(root_state[:, 7:]) # reset dof state joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) robot.reset() # apply force robot.permanent_wrench_composer.set_forces_and_torques( diff --git a/source/isaaclab/test/assets/check_fixed_base_assets.py b/source/isaaclab/test/assets/check_fixed_base_assets.py index c62c5a3334d..7b609c2a3d8 100644 --- a/source/isaaclab/test/assets/check_fixed_base_assets.py +++ b/source/isaaclab/test/assets/check_fixed_base_assets.py @@ -115,7 +115,7 @@ def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, Articula robot.write_root_velocity_to_sim(root_state[:, 7:]) # joint state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset the internal state robot.reset() print("[INFO]: Resetting robots state...") diff --git a/source/isaaclab/test/assets/check_ridgeback_franka.py b/source/isaaclab/test/assets/check_ridgeback_franka.py index 724cd8b2a08..24c409bd2e8 100644 --- a/source/isaaclab/test/assets/check_ridgeback_franka.py +++ b/source/isaaclab/test/assets/check_ridgeback_franka.py @@ -88,7 +88,7 @@ def run_simulator(sim: sim_utils.SimulationContext, robot: Articulation): ep_step_count = 0 # reset dof state joint_pos, joint_vel = robot.data.default_joint_pos.clone(), robot.data.default_joint_vel.clone() - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset internals robot.reset() # reset command diff --git a/source/isaaclab/test/controllers/test_differential_ik.py b/source/isaaclab/test/controllers/test_differential_ik.py index 34ae911ac91..85b19c6e7ff 100644 --- a/source/isaaclab/test/controllers/test_differential_ik.py +++ b/source/isaaclab/test/controllers/test_differential_ik.py @@ -182,7 +182,7 @@ def _run_ik_controller( joint_pos = wp.to_torch(robot.data.default_joint_pos).clone() joint_vel = wp.to_torch(robot.data.default_joint_vel).clone() # joint_pos *= sample_uniform(0.9, 1.1, joint_pos.shape, joint_pos.device) - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) robot.set_joint_position_target(joint_pos) robot.write_data_to_sim() # randomize root state yaw, ik should work regardless base rotation diff --git a/source/isaaclab/test/controllers/test_operational_space.py b/source/isaaclab/test/controllers/test_operational_space.py index 771cde6e5c0..e4712db4137 100644 --- a/source/isaaclab/test/controllers/test_operational_space.py +++ b/source/isaaclab/test/controllers/test_operational_space.py @@ -1351,7 +1351,7 @@ def _run_op_space_controller( # reset joint state to default default_joint_pos = wp.to_torch(robot.data.default_joint_pos).clone() default_joint_vel = wp.to_torch(robot.data.default_joint_vel).clone() - robot.write_joint_state_to_sim(default_joint_pos, default_joint_vel) + robot.write_joint_state_to_sim(position=default_joint_pos, velocity=default_joint_vel) robot.set_joint_effort_target(zero_joint_efforts) # Set zero torques in the initial step robot.write_data_to_sim() robot.reset() diff --git a/source/isaaclab/test/envs/test_color_randomization.py b/source/isaaclab/test/envs/test_color_randomization.py index a7edda7036c..a944cd63e4e 100644 --- a/source/isaaclab/test/envs/test_color_randomization.py +++ b/source/isaaclab/test/envs/test_color_randomization.py @@ -30,7 +30,7 @@ from isaaclab.managers import ObservationTermCfg as ObsTerm from isaaclab.managers import SceneEntityCfg from isaaclab.utils import configclass -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab_tasks.manager_based.classic.cartpole.cartpole_env_cfg import CartpoleSceneCfg @@ -138,7 +138,7 @@ def __post_init__(self): def test_color_randomization(device): """Test color randomization for cartpole environment.""" # skip test if stage in memory is not supported - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: pytest.skip("Color randomization test hangs in this version of Isaac Sim") # Create a new stage diff --git a/source/isaaclab/test/markers/check_markers_visibility.py b/source/isaaclab/test/markers/check_markers_visibility.py index 98dbee8ddcd..aefbd7cb5ef 100644 --- a/source/isaaclab/test/markers/check_markers_visibility.py +++ b/source/isaaclab/test/markers/check_markers_visibility.py @@ -105,7 +105,7 @@ def run_simulator( scene["robot"].data.default_joint_pos.clone(), scene["robot"].data.default_joint_vel.clone(), ) - scene["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # clear internal buffers scene.reset() print("[INFO]: Resetting robot state...") diff --git a/source/isaaclab/test/scene/check_interactive_scene.py b/source/isaaclab/test/scene/check_interactive_scene.py index 5b2463b315a..1682f229fb0 100644 --- a/source/isaaclab/test/scene/check_interactive_scene.py +++ b/source/isaaclab/test/scene/check_interactive_scene.py @@ -132,12 +132,12 @@ def main(): # -- robot 1 scene.articulations["robot_1"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot_1"].write_root_velocity_to_sim(root_state[:, 7:]) - scene.articulations["robot_1"].write_joint_state_to_sim(joint_pos, joint_vel) + scene.articulations["robot_1"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # -- robot 2 root_state[:, 1] += 1.0 scene.articulations["robot_2"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot_2"].write_root_velocity_to_sim(root_state[:, 7:]) - scene.articulations["robot_2"].write_joint_state_to_sim(joint_pos, joint_vel) + scene.articulations["robot_2"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset buffers scene.reset() print(">>>>>>>> Reset!") diff --git a/source/isaaclab/test/sim/test_simulation_render_config.py b/source/isaaclab/test/sim/test_simulation_render_config.py index e5ad4274d18..89011e84cac 100644 --- a/source/isaaclab/test/sim/test_simulation_render_config.py +++ b/source/isaaclab/test/sim/test_simulation_render_config.py @@ -23,7 +23,7 @@ from isaaclab.app.settings_manager import get_settings_manager from isaaclab.sim.simulation_cfg import RenderCfg, SimulationCfg from isaaclab.sim.simulation_context import SimulationContext -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit @pytest.mark.skip(reason="Timeline not stopped") @@ -106,7 +106,7 @@ def test_render_cfg_presets(): # grab isaac lab apps path isaaclab_app_exp_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), *[".."] * 4, "apps") # for Isaac Sim 5 compatibility, we use the 5 rendering mode app files in a different folder - if get_isaac_sim_version().major < 6: + if has_kit() and get_isaac_sim_version().major < 6: isaaclab_app_exp_path = os.path.join(isaaclab_app_exp_path, "isaacsim_5") # grab preset settings diff --git a/source/isaaclab/test/sim/test_simulation_stage_in_memory.py b/source/isaaclab/test/sim/test_simulation_stage_in_memory.py index ebe8158750f..9bd7658e809 100644 --- a/source/isaaclab/test/sim/test_simulation_stage_in_memory.py +++ b/source/isaaclab/test/sim/test_simulation_stage_in_memory.py @@ -26,7 +26,7 @@ import isaaclab.sim as sim_utils from isaaclab.sim.simulation_context import SimulationCfg, SimulationContext from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit @pytest.fixture @@ -50,7 +50,7 @@ def test_stage_in_memory_with_shapes(sim): """Test spawning of shapes with stage in memory.""" # skip test if stage in memory is not supported - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: pytest.skip("Stage in memory is not supported in this version of Isaac Sim") # grab stage in memory and set as current stage via the with statement @@ -122,7 +122,7 @@ def test_stage_in_memory_with_usds(sim): """Test spawning of USDs with stage in memory.""" # skip test if stage in memory is not supported - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: pytest.skip("Stage in memory is not supported in this version of Isaac Sim") # define parameters @@ -177,7 +177,7 @@ def test_stage_in_memory_with_clone_in_fabric(sim): """Test cloning in fabric with stage in memory.""" # skip test if stage in memory is not supported - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: pytest.skip("Stage in memory is not supported in this version of Isaac Sim") # define parameters diff --git a/source/isaaclab/test/sim/test_spawn_from_files.py b/source/isaaclab/test/sim/test_spawn_from_files.py index ccfef35837c..3a74713fcb2 100644 --- a/source/isaaclab/test/sim/test_spawn_from_files.py +++ b/source/isaaclab/test/sim/test_spawn_from_files.py @@ -20,7 +20,7 @@ import isaaclab.sim as sim_utils from isaaclab.sim import SimulationCfg, SimulationContext from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit @pytest.fixture @@ -69,7 +69,7 @@ def test_spawn_urdf(sim): """Test loading prim from URDF file.""" # pin the urdf importer extension to the older version manager = omni.kit.app.get_app().get_extension_manager() - if get_isaac_sim_version() == Version("5.1"): + if has_kit() and get_isaac_sim_version() == Version("5.1"): pinned_urdf_extension_name = "isaacsim.asset.importer.urdf-2.4.31" else: pinned_urdf_extension_name = "isaacsim.asset.importer.urdf" diff --git a/source/isaaclab/test/sim/test_urdf_converter.py b/source/isaaclab/test/sim/test_urdf_converter.py index 6e0fe50e045..64f967b6d76 100644 --- a/source/isaaclab/test/sim/test_urdf_converter.py +++ b/source/isaaclab/test/sim/test_urdf_converter.py @@ -24,7 +24,7 @@ import isaaclab.sim as sim_utils from isaaclab.sim import SimulationCfg, SimulationContext from isaaclab.sim.converters import UrdfConverter, UrdfConverterCfg -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit # Create a fixture for setup and teardown @@ -34,7 +34,7 @@ def sim_config(): sim_utils.create_new_stage() # pin the urdf importer extension to the older version manager = omni.kit.app.get_app().get_extension_manager() - if get_isaac_sim_version() == Version("5.1"): + if has_kit() and get_isaac_sim_version() == Version("5.1"): pinned_urdf_extension_name = "isaacsim.asset.importer.urdf-2.4.31" else: pinned_urdf_extension_name = "isaacsim.asset.importer.urdf" diff --git a/source/isaaclab/test/utils/test_version.py b/source/isaaclab/test/utils/test_version.py index ba737b53643..9e422b4a2ef 100644 --- a/source/isaaclab/test/utils/test_version.py +++ b/source/isaaclab/test/utils/test_version.py @@ -17,11 +17,13 @@ import pytest from packaging.version import Version -from isaaclab.utils.version import compare_versions, get_isaac_sim_version +from isaaclab.utils.version import compare_versions, get_isaac_sim_version, has_kit def test_get_isaac_sim_version(): """Test that get_isaac_sim_version returns cached Version object.""" + if not has_kit(): + pytest.skip("Isaac Sim app not running") # Call twice to ensure caching works version1 = get_isaac_sim_version() version2 = get_isaac_sim_version() @@ -42,6 +44,8 @@ def test_get_isaac_sim_version(): def test_get_isaac_sim_version_format(): """Test that get_isaac_sim_version returns correct format.""" + if not has_kit(): + pytest.skip("Isaac Sim app not running") isaac_version = get_isaac_sim_version() # Should be able to convert to string @@ -60,6 +64,8 @@ def test_get_isaac_sim_version_format(): def test_version_caching_performance(): """Test that caching improves performance for version checks.""" + if not has_kit(): + pytest.skip("Isaac Sim app not running") # First call (will cache) version1 = get_isaac_sim_version() @@ -72,6 +78,8 @@ def test_version_caching_performance(): def test_version_comparison_operators(): """Test that Version objects support natural comparisons.""" + if not has_kit(): + pytest.skip("Isaac Sim app not running") isaac_version = get_isaac_sim_version() # Should support comparison operators diff --git a/source/isaaclab_assets/data/Robots b/source/isaaclab_assets/data/Robots new file mode 120000 index 00000000000..3b4d952b100 --- /dev/null +++ b/source/isaaclab_assets/data/Robots @@ -0,0 +1 @@ +/home/zhengyuz/Documents/Assets/Robots/ \ No newline at end of file diff --git a/source/isaaclab_assets/isaaclab_assets/sensors/__init__.py b/source/isaaclab_assets/isaaclab_assets/sensors/__init__.py index f5f6c6ac116..f61e05c4ffa 100644 --- a/source/isaaclab_assets/isaaclab_assets/sensors/__init__.py +++ b/source/isaaclab_assets/isaaclab_assets/sensors/__init__.py @@ -7,5 +7,12 @@ # Configuration for different assets. ## -from .gelsight import * -from .velodyne import * +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "gelsight": ["GELSIGHT_R15_CFG", "GELSIGHT_MINI_CFG"], + "velodyne": ["VELODYNE_VLP_16_RAYCASTER_CFG"], + }, +) diff --git a/source/isaaclab_contrib/isaaclab_contrib/sensors/__init__.py b/source/isaaclab_contrib/isaaclab_contrib/sensors/__init__.py index a7ea884318a..ad0a13574fc 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/sensors/__init__.py +++ b/source/isaaclab_contrib/isaaclab_contrib/sensors/__init__.py @@ -22,4 +22,16 @@ """ -from .tacsl_sensor import * +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "tacsl_sensor": [ + "GelSightRenderCfg", + "VisuoTactileSensor", + "VisuoTactileSensorCfg", + "VisuoTactileSensorData", + ], + }, +) diff --git a/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/__init__.py b/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/__init__.py index 869b233d166..3058a26201a 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/__init__.py +++ b/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/__init__.py @@ -5,6 +5,13 @@ """TacSL Tactile Sensor implementation for IsaacLab.""" -from .visuotactile_sensor import VisuoTactileSensor -from .visuotactile_sensor_cfg import GelSightRenderCfg, VisuoTactileSensorCfg -from .visuotactile_sensor_data import VisuoTactileSensorData +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "visuotactile_sensor": ["VisuoTactileSensor"], + "visuotactile_sensor_cfg": ["GelSightRenderCfg", "VisuoTactileSensorCfg"], + "visuotactile_sensor_data": ["VisuoTactileSensorData"], + }, +) diff --git a/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor_cfg.py b/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor_cfg.py index 5aaf9cd6731..6603451b501 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor_cfg.py +++ b/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor_cfg.py @@ -12,11 +12,9 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.markers.config import VISUO_TACTILE_SENSOR_MARKER_CFG from isaaclab.sensors import SensorBaseCfg, TiledCameraCfg -from isaaclab.utils import configclass +from isaaclab.utils import DeferredClass, configclass from isaaclab.utils.assets import ISAACLAB_NUCLEUS_DIR -from .visuotactile_sensor import VisuoTactileSensor - ## # GelSight Render Configuration ## @@ -109,7 +107,9 @@ class VisuoTactileSensorCfg(SensorBaseCfg): It can capture tactile RGB/depth images and compute penalty-based contact forces. """ - class_type: type = VisuoTactileSensor + class_type: type | DeferredClass = DeferredClass( + "isaaclab_contrib.sensors.tacsl_sensor.visuotactile_sensor:VisuoTactileSensor" + ) # Sensor type and capabilities render_cfg: GelSightRenderCfg = MISSING diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/__init__.py b/source/isaaclab_newton/isaaclab_newton/actuators/__init__.py new file mode 100644 index 00000000000..584f21cd13c --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/__init__.py @@ -0,0 +1,29 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package for different actuator models. + +Actuator models are used to model the behavior of the actuators in an articulation. These +are usually meant to be used in simulation to model different actuator dynamics and delays. + +There are two main categories of actuator models that are supported: + +- **Implicit**: Motor model with ideal PD from the physics engine. This is similar to having a continuous time + PD controller. The motor model is implicit in the sense that the motor model is not explicitly defined by the user. +- **Explicit**: Motor models based on physical drive models. + + - **Physics-based**: Derives the motor models based on first-principles. + - **Neural Network-based**: Learned motor models from actuator data. + +Every actuator model inherits from the :class:`isaaclab.actuators.ActuatorBase` class, +which defines the common interface for all actuator models. The actuator models are handled +and called by the :class:`isaaclab.assets.Articulation` class. +""" + +from .actuator_base import ActuatorBase +from .actuator_net import ActuatorNetLSTM, ActuatorNetMLP + +# from .actuator_pd import DCMotor, DelayedPDActuator, IdealPDActuator, ImplicitActuator, RemotizedPDActuator +from .actuator_pd import DCMotor, IdealPDActuator, ImplicitActuator diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/actuator_base.py b/source/isaaclab_newton/isaaclab_newton/actuators/actuator_base.py new file mode 100644 index 00000000000..24a28ae5f2c --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/actuator_base.py @@ -0,0 +1,275 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from abc import ABC, abstractmethod +from enum import IntEnum +from typing import TYPE_CHECKING, ClassVar + +import warp as wp +from isaaclab_newton.actuators.kernels import clip_efforts_with_limits + +import isaaclab.utils.string as string_utils +from isaaclab.actuators.actuator_cfg import ActuatorBaseCfg +from isaaclab.utils.warp.update_kernels import update_array2D_with_array1D_indexed, update_array2D_with_value_indexed + +if TYPE_CHECKING: + from isaaclab_newton.assets.articulation.articulation_data import ArticulationData + + +class ActuatorBase(ABC): + """Base class for actuator models over a collection of actuated joints in an articulation. + + Actuator models augment the simulated articulation joints with an external drive dynamics model. + The model is used to convert the user-provided joint commands (positions, velocities and efforts) + into the desired joint positions, velocities and efforts that are applied to the simulated articulation. + + The base class provides the interface for the actuator models. It is responsible for parsing the + actuator parameters from the configuration and storing them as buffers. It also provides the + interface for resetting the actuator state and computing the desired joint commands for the simulation. + + For each actuator model, a corresponding configuration class is provided. The configuration class + is used to parse the actuator parameters from the configuration. It also specifies the joint names + for which the actuator model is applied. These names can be specified as regular expressions, which + are matched against the joint names in the articulation. + + To see how the class is used, check the :class:`isaaclab.assets.Articulation` class. + """ + + __backend_name__: str = "newton" + """The name of the backend for the actuator.""" + + is_implicit_model: ClassVar[bool] = False + """Flag indicating if the actuator is an implicit or explicit actuator model. + + If a class inherits from :class:`ImplicitActuator`, then this flag should be set to :obj:`True`. + """ + + data: ArticulationData + """The data for the articulation.""" + + _DEFAULT_MAX_EFFORT_SIM: ClassVar[float] = 1.0e9 + """The default maximum effort for the actuator joints in the simulation. Defaults to 1.0e9. + + If the :attr:`ActuatorBaseCfg.effort_limit_sim` is not specified and the actuator is an explicit + actuator, then this value is used. + """ + + def __init__( + self, + cfg: ActuatorBaseCfg, + joint_names: list[str], + joint_mask: wp.array, + joint_indices: list[int], + env_mask: wp.array, + data: ArticulationData, + device: str, + ): + """Initialize the actuator. + + The actuator parameters are parsed from the configuration and stored as buffers. If the parameters + are not specified in the configuration, then their values provided in the constructor are used. + + .. note:: + The values in the constructor are typically obtained through the USD schemas corresponding + to the joints in the actuator model. + + Args: + cfg: The configuration of the actuator model. + joint_names: The joint names in the articulation. + joint_mask: The mask of joints to use. + env_mask: The mask of environments to use. + articulation_data: The data for the articulation. + device: Device used for processing. + """ + # save parameters + self.cfg = cfg + self._device = device + self._joint_names = joint_names + self._joint_mask = joint_mask + self._joint_indices = joint_indices + self._env_mask = env_mask + # Get the number of environments and joints from the articulation data + self._num_envs = env_mask.shape[0] + self._num_joints = joint_mask.shape[0] + # Get the data from the articulation + self.data = data + + # For explicit models, we do not want to enforce the effort limit through the solver + # (unless it is explicitly set) + if not self.is_implicit_model and self.cfg.effort_limit_sim is None: + self.cfg.effort_limit_sim = self._DEFAULT_MAX_EFFORT_SIM + + # resolve usd, actuator configuration values + # case 1: if usd_value == actuator_cfg_value: all good, + # case 2: if usd_value != actuator_cfg_value: we use actuator_cfg_value + # case 3: if actuator_cfg_value is None: we use usd_value + + to_check = [ + ("velocity_limit_sim", self.data._sim_bind_joint_vel_limits_sim), + ("effort_limit_sim", self.data._sim_bind_joint_effort_limits_sim), + ("stiffness", self.data._actuator_stiffness), + ("damping", self.data._actuator_damping), + ("armature", self.data._sim_bind_joint_armature), + ("friction", self.data._sim_bind_joint_friction_coeff), + ("dynamic_friction", self.data._joint_dynamic_friction), + ("viscous_friction", self.data._joint_viscous_friction), + ] + for param_name, newton_val in to_check: + cfg_val = getattr(self.cfg, param_name) + self._parse_joint_parameter(cfg_val, newton_val) + + def __str__(self) -> str: + """Returns: A string representation of the actuator group.""" + # resolve model type (implicit or explicit) + model_type = "implicit" if self.is_implicit_model else "explicit" + + return ( + f" object:\n" + f"\tModel type : {model_type}\n" + f"\tNumber of joints : {self.num_joints}\n" + f"\tJoint names expression: {self.cfg.joint_names_expr}\n" + f"\tJoint names : {self.joint_names}\n" + f"\tJoint mask : {self.joint_mask}\n" + ) + + """ + Properties. + """ + + @property + def num_joints(self) -> int: + """Number of actuators in the group.""" + return len(self._joint_names) + + @property + def joint_names(self) -> list[str]: + """Articulation's joint names that are part of the group.""" + return self._joint_names + + @property + def joint_mask(self) -> wp.array: + """Articulation's masked indices that denote which joints are part of the group.""" + return self._joint_mask + + """ + Operations. + """ + + @abstractmethod + def reset(self, env_mask: wp.array): + """Reset the internals within the group. + + Args: + env_mask: Mask of environments to reset. + """ + raise NotImplementedError + + @abstractmethod + def compute(self): + """Process the actuator group actions and compute the articulation actions. + + It computes the articulation actions based on the actuator model type. To do so, it reads + the following quantities from the articulation data: + - sim_bind_joint_pos, the current joint positions + - sim_bind_joint_vel, the current joint velocities + - joint_control_mode, the current joint control mode + - joint_stiffness, the current joint stiffness + - joint_damping, the current joint damping + + With these, it updates the following quantities: + + """ + raise NotImplementedError + + """ + Helper functions. + """ + + def _parse_joint_parameter(self, cfg_value: float | dict[str, float] | None, original_value: wp.array | None): + """Parse the joint parameter from the configuration. + + Args: + cfg_value: The parameter value from the configuration. If None, then use the default value. + default_value: The default value to use if the parameter is None. If it is also None, + then an error is raised. + + Returns: + The parsed parameter value. + + Raises: + TypeError: If the parameter value is not of the expected type. + TypeError: If the default value is not of the expected type. + ValueError: If the parameter value is None and no default value is provided. + ValueError: If the default value tensor is the wrong shape. + """ + # parse the parameter + if cfg_value is not None: + if isinstance(cfg_value, (float, int)): + if isinstance(cfg_value, IntEnum): + cfg_value = int(cfg_value.value) + if original_value.dtype is wp.float32: + cfg_value = float(cfg_value) + elif original_value.dtype is wp.int32: + cfg_value = int(cfg_value) + # if int, then use the same value for all joints + wp.launch( + update_array2D_with_value_indexed, + dim=(self._num_envs, len(self._joint_indices)), + device=self._device, + inputs=[ + cfg_value, + original_value, + None, + wp.array(self._joint_indices, dtype=wp.int32, device=self._device), + ], + ) + elif isinstance(cfg_value, dict): + # if dict, then parse the regular expression + indices, _, values = string_utils.resolve_matching_names_values(cfg_value, self.joint_names) + indices_global = [self._joint_indices[i] for i in indices] + wp.launch( + update_array2D_with_array1D_indexed, + dim=(self._num_envs, len(indices_global)), + device=self._device, + inputs=[ + wp.array(values, dtype=wp.float32, device=self._device), + original_value, + None, + wp.array(indices_global, dtype=wp.int32, device=self._device), + ], + ) + else: + raise TypeError( + f"Invalid type for parameter value: {type(cfg_value)} for " + + f"actuator on joints {self.joint_names}. Expected float or dict." + ) + elif original_value is None: + raise ValueError("The parameter value is None and no newton value is provided.") + + def _clip_effort(self, effort: wp.array, clipped_effort: wp.array) -> None: + """Clip the desired torques based on the motor limits. + + .. note:: The array is modified in place. + + Args: + desired_torques: The desired torques to clip. Expected shape is (num_envs, num_joints). + + Returns: + The clipped torques. + """ + wp.launch( + clip_efforts_with_limits, + dim=(self._num_envs, self._num_joints), + inputs=[ + self.data._sim_bind_joint_effort_limits_sim, + effort, + clipped_effort, + self._env_mask, + self.joint_mask, + ], + device=self._device, + ) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/actuator_net.py b/source/isaaclab_newton/isaaclab_newton/actuators/actuator_net.py new file mode 100644 index 00000000000..f9a86d527bf --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/actuator_net.py @@ -0,0 +1,203 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Neural network models for actuators. + +Currently, the following models are supported: + +* Multi-Layer Perceptron (MLP) +* Long Short-Term Memory (LSTM) + +""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab.utils.assets import read_file +from isaaclab.utils.warp.update_kernels import update_array2D_with_array2D_masked + +from .actuator_pd import DCMotor + +if TYPE_CHECKING: + from .actuator_net_cfg import ActuatorNetLSTMCfg, ActuatorNetMLPCfg + + +class ActuatorNetLSTM(DCMotor): + """Actuator model based on recurrent neural network (LSTM). + + Unlike the MLP implementation :cite:t:`hwangbo2019learning`, this class implements + the learned model as a temporal neural network (LSTM) based on the work from + :cite:t:`rudin2022learning`. This removes the need of storing a history as the + hidden states of the recurrent network captures the history. + + Note: + Only the desired joint positions are used as inputs to the network. + """ + + cfg: ActuatorNetLSTMCfg + """The configuration of the actuator model.""" + + def __init__(self, cfg: ActuatorNetLSTMCfg, *args, **kwargs): + super().__init__(cfg, *args, **kwargs) + + assert self.cfg.control_mode == "position", "ActuatorNetLSTM only supports position control" + + # load the model from JIT file + file_bytes = read_file(self.cfg.network_file) + self.network = torch.jit.load(file_bytes, map_location=self._device).eval() + + # extract number of lstm layers and hidden dim from the shape of weights + num_layers = len(self.network.lstm.state_dict()) // 4 + hidden_dim = self.network.lstm.state_dict()["weight_hh_l0"].shape[1] + # create buffers for storing LSTM inputs + self.sea_input = torch.zeros(self._num_envs * self.num_joints, 1, 2, device=self._device) + self.sea_hidden_state = torch.zeros( + num_layers, self._num_envs * self.num_joints, hidden_dim, device=self._device + ) + self.sea_cell_state = torch.zeros(num_layers, self._num_envs * self.num_joints, hidden_dim, device=self._device) + # reshape via views (doesn't change the actual memory layout) + layer_shape_per_env = (num_layers, self._num_envs, self.num_joints, hidden_dim) + self.sea_hidden_state_per_env = self.sea_hidden_state.view(layer_shape_per_env) + self.sea_cell_state_per_env = self.sea_cell_state.view(layer_shape_per_env) + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int]): + # reset the hidden and cell states for the specified environments + with torch.no_grad(): + self.sea_hidden_state_per_env[:, env_ids] = 0.0 + self.sea_cell_state_per_env[:, env_ids] = 0.0 + + def compute(self): + # compute network inputs + self.sea_input[:, 0, 0] = wp.to_torch( + self.data._actuator_position_target - self.data._sim_bind_joint_pos + ).flatten() + self.sea_input[:, 0, 1] = wp.to_torch(self.data._sim_bind_joint_vel).flatten() + + # run network inference + with torch.inference_mode(): + torques, (self.sea_hidden_state[:], self.sea_cell_state[:]) = self.network( + self.sea_input, (self.sea_hidden_state, self.sea_cell_state) + ) + self.data._computed_effort = wp.from_torch(torques.reshape(self._num_envs, self.num_joints)) + # clip the computed effort based on the motor limits + self._clip_effort(self.data._computed_effort, self.data._applied_effort) + # update the joint effort + wp.launch( + update_array2D_with_array2D_masked, + dim=(self._num_envs, self.num_joints), + inputs=[ + self.data._applied_effort, + self.data.joint_effort, + self._env_mask, + self._joint_mask, + ], + device=self._device, + ) + + +class ActuatorNetMLP(DCMotor): + """Actuator model based on multi-layer perceptron and joint history. + + Many times the analytical model is not sufficient to capture the actuator dynamics, the + delay in the actuator response, or the non-linearities in the actuator. In these cases, + a neural network model can be used to approximate the actuator dynamics. This model is + trained using data collected from the physical actuator and maps the joint state and the + desired joint command to the produced torque by the actuator. + + This class implements the learned model as a neural network based on the work from + :cite:t:`hwangbo2019learning`. The class stores the history of the joint positions errors + and velocities which are used to provide input to the neural network. The model is loaded + as a TorchScript. + + Note: + Only the desired joint positions are used as inputs to the network. + + """ + + cfg: ActuatorNetMLPCfg + """The configuration of the actuator model.""" + + def __init__(self, cfg: ActuatorNetMLPCfg, *args, **kwargs): + super().__init__(cfg, *args, **kwargs) + + assert self.cfg.control_mode == "position", "ActuatorNetMLP only supports position control" + + # load the model from JIT file + file_bytes = read_file(self.cfg.network_file) + self.network = torch.jit.load(file_bytes, map_location=self._device).eval() + + # create buffers for MLP history + history_length = max(self.cfg.input_idx) + 1 + self._joint_pos_error_history = torch.zeros( + self._num_envs, history_length, self.num_joints, device=self._device + ) + self._joint_vel_history = torch.zeros(self._num_envs, history_length, self.num_joints, device=self._device) + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int]): + # reset the history for the specified environments + self._joint_pos_error_history[env_ids] = 0.0 + self._joint_vel_history[env_ids] = 0.0 + + def compute(self): + # move history queue by 1 and update top of history + # -- positions + self._joint_pos_error_history = self._joint_pos_error_history.roll(1, 1) + self._joint_pos_error_history[:, 0] = wp.to_torch( + self.data._actuator_position_target - self.data._sim_bind_joint_pos + ) + # -- velocity + self._joint_vel_history = self._joint_vel_history.roll(1, 1) + self._joint_vel_history[:, 0] = wp.to_torch(self.data._sim_bind_joint_vel) + + # compute network inputs + # -- positions + pos_input = torch.cat([self._joint_pos_error_history[:, i].unsqueeze(2) for i in self.cfg.input_idx], dim=2) + pos_input = pos_input.view(self._num_envs * self.num_joints, -1) + # -- velocity + vel_input = torch.cat([self._joint_vel_history[:, i].unsqueeze(2) for i in self.cfg.input_idx], dim=2) + vel_input = vel_input.view(self._num_envs * self.num_joints, -1) + # -- scale and concatenate inputs + if self.cfg.input_order == "pos_vel": + network_input = torch.cat([pos_input * self.cfg.pos_scale, vel_input * self.cfg.vel_scale], dim=1) + elif self.cfg.input_order == "vel_pos": + network_input = torch.cat([vel_input * self.cfg.vel_scale, pos_input * self.cfg.pos_scale], dim=1) + else: + raise ValueError( + f"Invalid input order for MLP actuator net: {self.cfg.input_order}. Must be 'pos_vel' or 'vel_pos'." + ) + + # run network inference + with torch.inference_mode(): + torques = self.network(network_input).view(self._num_envs, self.num_joints) + computed_effort = torques.view(self._num_envs, self.num_joints) * self.cfg.torque_scale + + # clip the computed effort based on the motor limits + self.data._computed_effort = wp.from_torch(computed_effort) + self._clip_effort(self.data._computed_effort, self.data._applied_effort) + # update the joint effort + wp.launch( + update_array2D_with_array2D_masked, + dim=(self._num_envs, self.num_joints), + inputs=[ + self.data._applied_effort, + self.data.joint_effort, + self._env_mask, + self._joint_mask, + ], + device=self._device, + ) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/actuator_pd.py b/source/isaaclab_newton/isaaclab_newton/actuators/actuator_pd.py new file mode 100644 index 00000000000..3694d8ddab5 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/actuator_pd.py @@ -0,0 +1,462 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import warp as wp + +from isaaclab.utils.warp.update_kernels import update_array2D_with_array2D_masked + +from .actuator_base import ActuatorBase +from .kernels import clip_efforts_dc_motor, compute_pd_actuator + +# from isaaclab.utils import DelayBuffer, LinearInterpolation + + +if TYPE_CHECKING: + from isaaclab.actuators.actuator_cfg import ( # DelayedPDActuatorCfg,; RemotizedPDActuatorCfg, + DCMotorCfg, + IdealPDActuatorCfg, + ImplicitActuatorCfg, + ) + +# import logger +logger = logging.getLogger(__name__) + + +""" +Implicit Actuator Models. +""" + + +class ImplicitActuator(ActuatorBase): + """Implicit actuator model that is handled by the simulation. + + This performs a similar function as the :class:`IdealPDActuator` class. However, the PD control is handled + implicitly by the simulation which performs continuous-time integration of the PD control law. This is + generally more accurate than the explicit PD control law used in :class:`IdealPDActuator` when the simulation + time-step is large. + + The articulation class sets the stiffness and damping parameters from the implicit actuator configuration + into the simulation. Thus, the class does not perform its own computations on the joint action that + needs to be applied to the simulation. However, it computes the approximate torques for the actuated joint + since PhysX does not expose this quantity explicitly. + + .. caution:: + + The class is only provided for consistency with the other actuator models. It does not implement any + functionality and should not be used. All values should be set to the simulation directly. + """ + + cfg: ImplicitActuatorCfg + """The configuration for the actuator model.""" + + def __init__(self, cfg: ImplicitActuatorCfg, *args, **kwargs): + # effort limits + if cfg.effort_limit_sim is None and cfg.effort_limit is not None: + # throw a warning that we have a replacement for the deprecated parameter + logger.warning( + "The object has a value for 'effort_limit'." + " This parameter will be removed in the future." + " To set the effort limit, please use 'effort_limit_sim' instead." + ) + cfg.effort_limit_sim = cfg.effort_limit + elif cfg.effort_limit_sim is not None and cfg.effort_limit is None: + # TODO: Eventually we want to get rid of 'effort_limit' for implicit actuators. + # We should do this once all parameters have an "_sim" suffix. + cfg.effort_limit = cfg.effort_limit_sim + elif cfg.effort_limit_sim is not None and cfg.effort_limit is not None: + if cfg.effort_limit_sim != cfg.effort_limit: + raise ValueError( + "The object has set both 'effort_limit_sim' and 'effort_limit'" + f" and they have different values {cfg.effort_limit_sim} != {cfg.effort_limit}." + " Please only set 'effort_limit_sim' for implicit actuators." + ) + + # velocity limits + if cfg.velocity_limit_sim is None and cfg.velocity_limit is not None: + # throw a warning that previously this was not set + # it leads to different simulation behavior so we want to remain backwards compatible + logger.warning( + "The object has a value for 'velocity_limit'." + " Previously, although this value was specified, it was not getting used by implicit" + " actuators. Since this parameter affects the simulation behavior, we continue to not" + " use it. This parameter will be removed in the future." + " To set the velocity limit, please use 'velocity_limit_sim' instead." + ) + cfg.velocity_limit = None + elif cfg.velocity_limit_sim is not None and cfg.velocity_limit is None: + # TODO: Eventually we want to get rid of 'velocity_limit' for implicit actuators. + # We should do this once all parameters have an "_sim" suffix. + cfg.velocity_limit = cfg.velocity_limit_sim + elif cfg.velocity_limit_sim is not None and cfg.velocity_limit is not None: + if cfg.velocity_limit_sim != cfg.velocity_limit: + raise ValueError( + "The object has set both 'velocity_limit_sim' and 'velocity_limit'" + f" and they have different values {cfg.velocity_limit_sim} != {cfg.velocity_limit}." + " Please only set 'velocity_limit_sim' for implicit actuators." + ) + + # set implicit actuator model flag + ImplicitActuator.is_implicit_model = True + # call the base class + super().__init__(cfg, *args, **kwargs) + + """ + Operations. + """ + + def reset(self, *args, **kwargs): + # This is a no-op. There is no state to reset for implicit actuators. + pass + + def compute(self): + """Process the actuator group actions and compute the articulation actions. + + In case of implicit actuator, the control action is directly returned as the computed action. + This function is a no-op and does not perform any computation on the input control action. + However, it computes the approximate torques for the actuated joint since PhysX does not compute + this quantity explicitly. + + Args: + control_action: The joint action instance comprising of the desired joint positions, joint velocities + and (feed-forward) joint efforts. + joint_pos: The current joint positions of the joints in the group. Shape is (num_envs, num_joints). + joint_vel: The current joint velocities of the joints in the group. Shape is (num_envs, num_joints). + + Returns: + The computed desired joint positions, joint velocities and joint efforts. + """ + wp.launch( + compute_pd_actuator, + dim=(self._num_envs, self._num_joints), + inputs=[ + self.data._actuator_position_target, + self.data._actuator_velocity_target, + self.data._actuator_effort_target, + self.data._sim_bind_joint_pos, + self.data._sim_bind_joint_vel, + self.data._actuator_stiffness, + self.data._actuator_damping, + self.data._computed_effort, + self._env_mask, + self._joint_mask, + ], + device=self._device, + ) + self._clip_effort(self.data._computed_effort, self.data._applied_effort) + # update the joint effort + wp.launch( + update_array2D_with_array2D_masked, + dim=(self._num_envs, self.num_joints), + inputs=[ + self.data._actuator_effort_target, + self.data.joint_effort, + self._env_mask, + self._joint_mask, + ], + device=self._device, + ) + + +""" +Explicit Actuator Models. +""" + + +class IdealPDActuator(ActuatorBase): + r"""Ideal torque-controlled actuator model with a simple saturation model. + + It employs the following model for computing torques for the actuated joint :math:`j`: + + .. math:: + + \tau_{j, computed} = k_p * (q - q_{des}) + k_d * (\dot{q} - \dot{q}_{des}) + \tau_{ff} + + where, :math:`k_p` and :math:`k_d` are joint stiffness and damping gains, :math:`q` and :math:`\dot{q}` + are the current joint positions and velocities, :math:`q_{des}`, :math:`\dot{q}_{des}` and :math:`\tau_{ff}` + are the desired joint positions, velocities and torques commands. + + The clipping model is based on the maximum torque applied by the motor. It is implemented as: + + .. math:: + + \tau_{j, max} & = \gamma \times \tau_{motor, max} \\ + \tau_{j, applied} & = clip(\tau_{computed}, -\tau_{j, max}, \tau_{j, max}) + + where the clipping function is defined as :math:`clip(x, x_{min}, x_{max}) = min(max(x, x_{min}), x_{max})`. + The parameters :math:`\gamma` is the gear ratio of the gear box connecting the motor and the actuated joint ends, + and :math:`\tau_{motor, max}` is the maximum motor effort possible. These parameters are read from + the configuration instance passed to the class. + """ + + cfg: IdealPDActuatorCfg + """The configuration for the actuator model.""" + + """ + Operations. + """ + + def reset(self, env_ids: Sequence[int]): + pass + + def compute(self): + wp.launch( + compute_pd_actuator, + dim=(self._num_envs, self._num_joints), + inputs=[ + self.data._actuator_position_target, + self.data._actuator_velocity_target, + self.data._actuator_effort_target, + self.data._sim_bind_joint_pos, + self.data._sim_bind_joint_vel, + self.data._actuator_stiffness, + self.data._actuator_damping, + self.data._computed_effort, + self._env_mask, + self._joint_mask, + ], + device=self._device, + ) + self._clip_effort(self.data._computed_effort, self.data._applied_effort) + # update the joint effort + wp.launch( + update_array2D_with_array2D_masked, + dim=(self._num_envs, self.num_joints), + inputs=[ + self.data._applied_effort, + self.data.joint_effort, + self._env_mask, + self._joint_mask, + ], + device=self._device, + ) + + +class DCMotor(IdealPDActuator): + r"""Direct control (DC) motor actuator model with velocity-based saturation model. + + It uses the same model as the :class:`IdealActuator` for computing the torques from input commands. + However, it implements a saturation model defined by DC motor characteristics. + + A DC motor is a type of electric motor that is powered by direct current electricity. In most cases, + the motor is connected to a constant source of voltage supply, and the current is controlled by a rheostat. + Depending on various design factors such as windings and materials, the motor can draw a limited maximum power + from the electronic source, which limits the produced motor torque and speed. + + A DC motor characteristics are defined by the following parameters: + + * Continuous-rated speed (:math:`\dot{q}_{motor, max}`) : The maximum-rated speed of the motor. + * Continuous-stall torque (:math:`\tau_{motor, max}`): The maximum-rated torque produced at 0 speed. + * Saturation torque (:math:`\tau_{motor, sat}`): The maximum torque that can be outputted for a short period. + + Based on these parameters, the instantaneous minimum and maximum torques are defined as follows: + + .. math:: + + \tau_{j, max}(\dot{q}) & = clip \left (\tau_{j, sat} \times \left(1 - + \frac{\dot{q}}{\dot{q}_{j, max}}\right), 0.0, \tau_{j, max} \right) \\ + \tau_{j, min}(\dot{q}) & = clip \left (\tau_{j, sat} \times \left( -1 - + \frac{\dot{q}}{\dot{q}_{j, max}}\right), - \tau_{j, max}, 0.0 \right) + + where :math:`\gamma` is the gear ratio of the gear box connecting the motor and the actuated joint ends, + :math:`\dot{q}_{j, max} = \gamma^{-1} \times \dot{q}_{motor, max}`, :math:`\tau_{j, max} = + \gamma \times \tau_{motor, max}` and :math:`\tau_{j, peak} = \gamma \times \tau_{motor, peak}` + are the maximum joint velocity, maximum joint torque and peak torque, respectively. These parameters + are read from the configuration instance passed to the class. + + Using these values, the computed torques are clipped to the minimum and maximum values based on the + instantaneous joint velocity: + + .. math:: + + \tau_{j, applied} = clip(\tau_{computed}, \tau_{j, min}(\dot{q}), \tau_{j, max}(\dot{q})) + + """ + + cfg: DCMotorCfg + """The configuration for the actuator model.""" + + def __init__(self, cfg: DCMotorCfg, *args, **kwargs): + super().__init__(cfg, *args, **kwargs) + # parse configuration + if self.cfg.saturation_effort is not None: + self._saturation_effort = self.cfg.saturation_effort + else: + self._saturation_effort = wp.inf + # check that quantities are provided + if self.cfg.velocity_limit is None: + raise ValueError("The velocity limit must be provided for the DC motor actuator model.") + + """ + Operations. + """ + + def compute(self): + # calculate the desired joint torques + return super().compute() + + """ + Helper functions. + """ + + def _clip_effort(self, effort: wp.array, clipped_effort: wp.array) -> None: + wp.launch( + clip_efforts_dc_motor, + dim=(self._num_envs, self._num_joints), + inputs=[ + self._saturation_effort, + self.data._sim_bind_joint_effort_limits_sim, + self.data._sim_bind_joint_vel_limits_sim, + self.data._sim_bind_joint_vel, + effort, + clipped_effort, + self._env_mask, + self._joint_mask, + ], + device=self._device, + ) + + +# class DelayedPDActuator(IdealPDActuator): +# """Ideal PD actuator with delayed command application. +# +# This class extends the :class:`IdealPDActuator` class by adding a delay to the actuator commands. The delay +# is implemented using a circular buffer that stores the actuator commands for a certain number of physics steps. +# The most recent actuation value is pushed to the buffer at every physics step, but the final actuation value +# applied to the simulation is lagged by a certain number of physics steps. +# +# The amount of time lag is configurable and can be set to a random value between the minimum and maximum time +# lag bounds at every reset. The minimum and maximum time lag values are set in the configuration instance passed +# to the class. +# """ +# +# cfg: DelayedPDActuatorCfg +# """The configuration for the actuator model.""" +# +# def __init__(self, cfg: DelayedPDActuatorCfg, *args, **kwargs): +# super().__init__(cfg, *args, **kwargs) +# # instantiate the delay buffers +# self.joint_targets_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) +# self.efforts_delay_buffer = DelayBuffer(cfg.max_delay, self._num_envs, device=self._device) +# # all of the envs +# self._ALL_INDICES = torch.arange(self._num_envs, dtype=torch.long, device=self._device) +# +# def reset(self, env_ids: Sequence[int]): +# super().reset(env_ids) +# # number of environments (since env_ids can be a slice) +# if env_ids is None or env_ids == slice(None): +# num_envs = self._num_envs +# else: +# num_envs = len(env_ids) +# # set a new random delay for environments in env_ids +# time_lags = torch.randint( +# low=self.cfg.min_delay, +# high=self.cfg.max_delay + 1, +# size=(num_envs,), +# dtype=torch.int, +# device=self._device, +# ) +# # set delays +# self.joint_targets_delay_buffer.set_time_lag(time_lags, env_ids) +# self.efforts_delay_buffer.set_time_lag(time_lags, env_ids) +# # reset buffers +# self.joint_targets_delay_buffer.reset(env_ids) +# self.efforts_delay_buffer.reset(env_ids) +# +# def compute( +# self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor +# ) -> ArticulationActions: +# # apply delay based on the delay the model for all the setpoints +# control_action.joint_targets = self.joint_targets_delay_buffer.compute(control_action.joint_targets) +# control_action.joint_efforts = self.efforts_delay_buffer.compute(control_action.joint_efforts) +# # compte actuator model +# return super().compute(control_action, joint_pos, joint_vel) +# +# +# class RemotizedPDActuator(DelayedPDActuator): +# """Ideal PD actuator with angle-dependent torque limits. +# +# This class extends the :class:`DelayedPDActuator` class by adding angle-dependent torque limits to the actuator. +# The torque limits are applied by querying a lookup table describing the relationship between the joint angle +# and the maximum output torque. The lookup table is provided in the configuration instance passed to the class. +# +# The torque limits are interpolated based on the current joint positions and applied to the actuator commands. +# """ +# +# def __init__( +# self, +# cfg: RemotizedPDActuatorCfg, +# joint_names: list[str], +# joint_ids: Sequence[int], +# num_envs: int, +# device: str, +# control_mode: str = "position", +# stiffness: torch.Tensor | float = 0.0, +# damping: torch.Tensor | float = 0.0, +# armature: torch.Tensor | float = 0.0, +# friction: torch.Tensor | float = 0.0, +# effort_limit: torch.Tensor | float = torch.inf, +# velocity_limit: torch.Tensor | float = torch.inf, +# ): +# # remove effort and velocity box constraints from the base class +# cfg.effort_limit = torch.inf +# cfg.velocity_limit = torch.inf +# # call the base method and set default effort_limit and velocity_limit to inf +# super().__init__( +# cfg, +# joint_names, +# joint_ids, +# num_envs, +# device, +# control_mode, +# stiffness, +# damping, +# armature, +# friction, +# torch.inf, +# torch.inf, +# ) +# self._joint_parameter_lookup = torch.tensor(cfg.joint_parameter_lookup, device=device) +# # define remotized joint torque limit +# self._torque_limit = LinearInterpolation(self.angle_samples, self.max_torque_samples, device=device) +# +# """ +# Properties. +# """ +# +# @property +# def angle_samples(self) -> torch.Tensor: +# return self._joint_parameter_lookup[:, 0] +# +# @property +# def transmission_ratio_samples(self) -> torch.Tensor: +# return self._joint_parameter_lookup[:, 1] +# +# @property +# def max_torque_samples(self) -> torch.Tensor: +# return self._joint_parameter_lookup[:, 2] +# +# """ +# Operations. +# """ +# +# def compute( +# self, control_action: ArticulationActions, joint_pos: torch.Tensor, joint_vel: torch.Tensor +# ) -> ArticulationActions: +# # call the base method +# control_action = super().compute(control_action, joint_pos, joint_vel) +# # compute the absolute torque limits for the current joint positions +# abs_torque_limits = self._torque_limit.compute(joint_pos) +# # apply the limits +# control_action.joint_efforts = torch.clamp( +# control_action.joint_efforts, min=-abs_torque_limits, max=abs_torque_limits +# ) +# self.applied_effort = control_action.joint_efforts +# return control_action diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py new file mode 100644 index 00000000000..f311673fdee --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py @@ -0,0 +1,80 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import warp as wp + + +@wp.kernel +def compute_pd_actuator( + pos_tgts: wp.array2d(dtype=wp.float32), + vel_tgts: wp.array2d(dtype=wp.float32), + added_effort: wp.array2d(dtype=wp.float32), + joint_pos: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + stiffness: wp.array2d(dtype=wp.float32), + damping: wp.array2d(dtype=wp.float32), + computed_effort: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), +) -> None: + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + computed_effort[env_index, joint_index] = ( + stiffness[env_index, joint_index] * (pos_tgts[env_index, joint_index] - joint_pos[env_index, joint_index]) + + damping[env_index, joint_index] * (vel_tgts[env_index, joint_index] - joint_vel[env_index, joint_index]) + + added_effort[env_index, joint_index] + ) + + +@wp.kernel +def clip_efforts_with_limits( + limits: wp.array2d(dtype=wp.float32), + joint_array: wp.array2d(dtype=wp.float32), + clipped_joint_array: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), +): + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + clipped_joint_array[env_index, joint_index] = wp.clamp( + joint_array[env_index, joint_index], -limits[env_index, joint_index], limits[env_index, joint_index] + ) + + +@wp.func +def clip_effort_dc_motor( + saturation_effort: float, + vel_limit: float, + effort_limit: float, + joint_vel: float, + effort: float, +): + max_effort = saturation_effort * (1.0 - joint_vel / vel_limit) + min_effort = saturation_effort * (-1.0 - joint_vel / vel_limit) + max_effort = wp.clamp(max_effort, 0.0, effort_limit) + min_effort = wp.clamp(min_effort, -effort_limit, 0.0) + return wp.clamp(effort, min_effort, max_effort) + + +@wp.kernel +def clip_efforts_dc_motor( + saturation_effort: float, + effort_limit: wp.array2d(dtype=wp.float32), + vel_limit: wp.array2d(dtype=wp.float32), + joint_vel: wp.array2d(dtype=wp.float32), + joint_array: wp.array2d(dtype=wp.float32), + clipped_joint_array: wp.array2d(dtype=wp.float32), + env_mask: wp.array(dtype=wp.bool), + joint_mask: wp.array(dtype=wp.bool), +): + env_index, joint_index = wp.tid() + if env_mask[env_index] and joint_mask[joint_index]: + clipped_joint_array[env_index, joint_index] = clip_effort_dc_motor( + saturation_effort, + vel_limit[env_index, joint_index], + effort_limit[env_index, joint_index], + joint_vel[env_index, joint_index], + joint_array[env_index, joint_index], + ) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index c4dcdd6cace..3e386617eee 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -952,6 +952,18 @@ def write_root_link_velocity_to_sim_mask( if self.data._root_com_state_w is not None: self.data._root_com_state_w.timestamp = -1.0 + def write_joint_state_to_sim_index( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + """Write joint positions and velocities over selected environment indices into the simulation.""" + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) + def write_joint_state_to_sim( self, *, @@ -968,7 +980,29 @@ def write_joint_state_to_sim( DeprecationWarning, stacklevel=2, ) - # set into simulation + self.write_joint_state_to_sim_index( + position=position, velocity=velocity, joint_ids=joint_ids, env_ids=env_ids + ) + + def write_joint_state_to_sim_index( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) + + def write_joint_state_to_sim_index( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/__init__.py b/source/isaaclab_newton/isaaclab_newton/cloner/__init__.py index 3460eea6864..0915dc36e88 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/__init__.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/__init__.py @@ -3,6 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause -from .newton_replicate import newton_replicate +import lazy_loader as lazy -__all__ = ["newton_replicate"] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "newton_replicate": ["newton_replicate"], + }, +) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index 36a3e6ef22b..b6c77d9038b 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -7,8 +7,10 @@ import torch -from pxr import Usd - +import warp as wp +from pxr import Usd, UsdGeom +from isaaclab_newton.physics import NewtonManager +from newton import ModelBuilder, solvers def newton_replicate( stage: Usd.Stage, @@ -22,22 +24,90 @@ def newton_replicate( up_axis: str = "Z", simplify_meshes: bool = True, ): - """Replicate prims for Newton physics backend. - - Newton does not require explicit physics replication like PhysX. - This is a no-op placeholder that maintains API compatibility. - - Args: - stage: USD stage. - sources: Source prim paths. - destinations: Destination templates with ``"{}"`` for env index. - env_ids: Environment indices. - mapping: Bool mask selecting envs per source. - use_fabric: Unused (for API compatibility). - device: Unused (for API compatibility). - - Returns: - None - """ - # Newton doesn't need explicit physics replication - USD replication is sufficient - pass + """Replicate prims into a Newton ``ModelBuilder`` using a per-source mapping.""" + + if positions is None: + positions = torch.zeros((mapping.size(1), 3), device=mapping.device, dtype=torch.float32) + if quaternions is None: + quaternions = torch.zeros((mapping.size(1), 4), device=mapping.device, dtype=torch.float32) + quaternions[:, 3] = 1.0 + + # load empty stage + builder = ModelBuilder(up_axis=up_axis) + stage_info = builder.add_usd(stage, ignore_paths=["/World/envs"] + sources) + + # build a prototype for each source + protos: dict[str, ModelBuilder] = {} + for src_path in sources: + p = ModelBuilder(up_axis=up_axis) + solvers.SolverMuJoCo.register_custom_attributes(p) + inverse_env_xform = get_inverse_env_xform(stage, src_path) + p.add_usd( + stage, + root_path=src_path, + load_visual_shapes=True, + skip_mesh_approximation=True, + xform=inverse_env_xform, + ) + if simplify_meshes: + p.approximate_meshes("convex_hull", keep_visual_shapes=True) + protos[src_path] = p + + # create a separate world for each environment (heterogeneous spawning) + # Newton assigns sequential world IDs (0, 1, 2, ...), so we need to track the mapping + newton_world_to_env_id = {} + for col, env_id in enumerate(env_ids.tolist()): + # begin a new world context (Newton assigns world ID = col) + builder.begin_world() + newton_world_to_env_id[col] = env_id + + # add all active sources for this world + for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist(): + builder.add_builder( + protos[sources[row]], + xform=wp.transform(positions[col].tolist(), quaternions[col].tolist()), + ) + + # end the world context + builder.end_world() + + # per-source, per-world renaming (strict prefix swap), compact style preserved + for i, src_path in enumerate(sources): + src_prefix_len = len(src_path.rstrip("/")) + swap = lambda name, new_root: new_root + name[src_prefix_len:] # noqa: E731 + world_cols = torch.nonzero(mapping[i], as_tuple=True)[0].tolist() + # Map Newton world IDs (sequential) to destination paths using env_ids + world_roots = {int(env_ids[c]): destinations[i].format(int(env_ids[c])) for c in world_cols} + + for t in ("body", "joint", "shape", "articulation"): + keys, worlds_arr = getattr(builder, f"{t}_key"), getattr(builder, f"{t}_world") + for k, w in enumerate(worlds_arr): + if w in world_roots and keys[k].startswith(src_path): + keys[k] = swap(keys[k], world_roots[w]) + + NewtonManager.set_builder(builder) + NewtonManager._num_envs = mapping.size(1) + return builder, stage_info + + +def get_inverse_env_xform(stage, src_path: str): + """Get the inverse transform of src_path to convert world→local.""" + xform_cache = UsdGeom.XformCache() + world_xform = xform_cache.GetLocalToWorldTransform(stage.GetPrimAtPath(src_path)) + + # Get the inverse of the world transform + inv_xform = world_xform.GetInverse() + + # Extract translation and rotation from inverse + inv_translation = inv_xform.ExtractTranslation() + inv_rotation = inv_xform.ExtractRotationQuat() + + inv_pos = (inv_translation[0], inv_translation[1], inv_translation[2]) + inv_quat = ( + inv_rotation.GetImaginary()[0], + inv_rotation.GetImaginary()[1], + inv_rotation.GetImaginary()[2], + inv_rotation.GetReal(), + ) + + return wp.transform(inv_pos, inv_quat) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/__init__.py b/source/isaaclab_newton/isaaclab_newton/physics/__init__.py new file mode 100644 index 00000000000..136f17d53c1 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation backends for simulation interfaces.""" + +import lazy_loader as lazy + +from .newton_manager_cfg import ( + FeatherstoneSolverCfg, + MJWarpSolverCfg, + NewtonCfg, + NewtonSolverCfg, + XPBDSolverCfg, +) + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "newton_manager": ["NewtonManager"], + }, +) +__all__ += ["FeatherstoneSolverCfg", "MJWarpSolverCfg", "NewtonCfg", "NewtonSolverCfg", "XPBDSolverCfg"] diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py new file mode 100644 index 00000000000..238581eecce --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -0,0 +1,509 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Newton physics manager for Isaac Lab.""" + +from __future__ import annotations + +import logging +import numpy as np +import re +from typing import TYPE_CHECKING + +import numpy as np +import warp as wp +from newton import Axis, BroadPhaseMode, CollisionPipeline, Contacts, Control, Model, ModelBuilder, State, eval_fk +from newton.sensors import SensorContact as NewtonContactSensor +from newton.solvers import SolverBase, SolverFeatherstone, SolverMuJoCo, SolverNotifyFlags, SolverXPBD + +from isaaclab.physics import PhysicsEvent, PhysicsManager +from isaaclab.sim.utils.stage import get_current_stage +from isaaclab.utils.timer import Timer + +if TYPE_CHECKING: + from isaaclab.sim.simulation_context import SimulationContext + +logger = logging.getLogger(__name__) + + +def flipped_match(x: str, y: str) -> bool: + """Flipped match function for contact partner matching. + + Args: + x: The body/shape name in the simulation. + y: The body/shape name in the contact view. + + Returns: + True if the body/shape name matches the contact view pattern. + """ + return re.match(y, x) is not None + + +class NewtonManager(PhysicsManager): + """Newton physics manager for Isaac Lab. + + This is a class-level (singleton-like) manager for the Newton simulation. + It handles solver configuration, physics stepping, and reset. + + Lifecycle: initialize() -> reset() -> step() (repeated) -> close() + """ + + _solver_dt: float = 1.0 / 200.0 + _num_substeps: int = 1 + _num_envs: int | None = None + + # Newton model and state + _builder: ModelBuilder = None + _model: Model = None + _solver: SolverBase = None + _solver_type: str = "mujoco_warp" + _use_single_state: bool | None = None + """Use only one state for both input and output for solver stepping. Requires solver support.""" + _state_0: State = None + _state_1: State = None + _control: Control = None + + # Physics settings + _gravity_vector: tuple[float, float, float] = (0.0, 0.0, -9.81) + _up_axis: str = "Z" + + # Collision and contacts + _contacts: Contacts | None = None + _needs_collision_pipeline: bool = False + _collision_pipeline = None + _newton_contact_sensors: dict = {} # Maps sensor_key to NewtonContactSensor + _report_contacts: bool = False + + # CUDA graphing + _graph = None + + # USD/Fabric sync + _newton_stage_path = None + _usdrt_stage = None + _newton_index_attr = "newton:index" + _clone_physics_only = False + + # Model changes (callbacks use unified system from PhysicsManager) + _model_changes: set[int] = set() + + # Views list for assets to register their views + _views: list = [] + + @classmethod + def initialize(cls, sim_context: SimulationContext) -> None: + """Initialize the manager with simulation context. + + Args: + sim_context: Parent simulation context. + """ + super().initialize(sim_context) + + # Newton-specific setup: get gravity from SimulationCfg (not physics manager cfg) + sim = PhysicsManager._sim + if sim is not None: + cls._gravity_vector = sim.cfg.gravity # type: ignore[union-attr] + + # USD fabric sync only needed for OV rendering + viz_str = sim.get_setting("/isaaclab/visualizer") or "" + requested = [v.strip() for v in viz_str.split(",") if v.strip()] + cls._clone_physics_only = "omniverse" not in requested + + @classmethod + def reset(cls, soft: bool = False) -> None: + """Reset physics simulation. + + Args: + soft: If True, skip full reinitialization. + """ + if not soft: + cls.start_simulation() + cls.initialize_solver() + + @classmethod + def forward(cls) -> None: + """Update articulation kinematics without stepping physics.""" + eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None) + + @classmethod + def step(cls) -> None: + """Step the physics simulation.""" + sim = PhysicsManager._sim + if sim is None or not sim.is_playing(): + return + + # Notify solver of model changes + if cls._model_changes: + for change in cls._model_changes: + cls._solver.notify_model_changed(change) + cls._model_changes = set() + + # Step simulation (graphed or not) + cfg = PhysicsManager._cfg + if cfg is not None and cfg.use_cuda_graph: # type: ignore[union-attr] + wp.capture_launch(cls._graph) # type: ignore[arg-type] + else: + cls._simulate() + + # Debug convergence info + if cfg is not None and cfg.debug_mode: # type: ignore[union-attr] + convergence_data = cls.get_solver_convergence_steps() + if convergence_data["max"] == cls._solver.mjw_model.opt.iterations: + logger.warning(f"Solver didn't converge! max_iter={convergence_data['max']}") + + PhysicsManager._sim_time += cls._solver_dt * cls._num_substeps + + @classmethod + def close(cls) -> None: + """Clean up Newton physics resources.""" + cls.clear() + super().close() + + @classmethod + def get_physics_sim_view(cls) -> list: + """Get the list of registered views. + + Assets can append their views to this list, and sensors can access them. + Returns a list that callers can append to. + + Returns: + List of registered views (e.g., NewtonArticulationView instances). + """ + return cls._views + + @classmethod + def is_fabric_enabled(cls) -> bool: + """Check if fabric interface is enabled (not applicable for Newton).""" + return False + + @classmethod + def clear(cls): + """Clear all Newton-specific state (callbacks cleared by super().close()).""" + cls._builder = None + cls._model = None + cls._solver = None + cls._use_single_state = None + cls._state_0 = None + cls._state_1 = None + cls._control = None + cls._contacts = None + cls._needs_collision_pipeline = False + cls._collision_pipeline = None + cls._newton_contact_sensors = {} + cls._report_contacts = False + cls._graph = None + cls._newton_stage_path = None + cls._usdrt_stage = None + cls._up_axis = "Z" + cls._model_changes = set() + cls._views = [] + + @classmethod + def set_builder(cls, builder: ModelBuilder) -> None: + """Set the Newton model builder.""" + cls._builder = builder + + @classmethod + def add_model_change(cls, change: SolverNotifyFlags) -> None: + """Register a model change to notify the solver.""" + cls._model_changes.add(change) + + @classmethod + def start_simulation(cls) -> None: + """Start simulation by finalizing model and initializing state. + + This function finalizes the model and initializes the simulation state. + Note: Collision pipeline is initialized later in initialize_solver() after + we determine whether the solver needs external collision detection. + """ + logger.debug(f"Builder: {cls._builder}") + + # Create builder from USD stage if not provided + if cls._builder is None: + cls.instantiate_builder_from_stage() + + logger.info("Dispatching MODEL_INIT callbacks") + cls.dispatch_event(PhysicsEvent.MODEL_INIT) + + device = PhysicsManager._device + logger.info(f"Finalizing model on device: {device}") + cls._builder.up_axis = Axis.from_string(cls._up_axis) + # Set smaller contact margin for manipulation examples (default 10cm is too large) + cls._builder.default_shape_cfg.contact_margin = 0.01 + with Timer(name="newton_finalize_builder", msg="Finalize builder took:", enable=True, format="ms"): + cls._model = cls._builder.finalize(device=device) + cls._model.set_gravity(cls._gravity_vector) + cls._model.num_envs = cls._num_envs + + cls._state_0 = cls._model.state() + cls._state_1 = cls._model.state() + cls._control = cls._model.control() + eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None) + + logger.info("Dispatching PHYSICS_READY callbacks") + cls.dispatch_event(PhysicsEvent.PHYSICS_READY) + + # Setup USD/Fabric sync for Omniverse rendering + if not cls._clone_physics_only: + import usdrt + + cls._usdrt_stage = get_current_stage(fabric=True) + for i, prim_path in enumerate(cls._model.body_key): + prim = cls._usdrt_stage.GetPrimAtPath(prim_path) + prim.CreateAttribute(cls._newton_index_attr, usdrt.Sdf.ValueTypeNames.UInt, True) + prim.GetAttribute(cls._newton_index_attr).Set(i) + xformable_prim = usdrt.Rt.Xformable(prim) + if not xformable_prim.HasWorldXform(): + xformable_prim.SetWorldXformFromUsd() + + @classmethod + def instantiate_builder_from_stage(cls): + """Create builder from USD stage.""" + from pxr import UsdGeom + + stage = get_current_stage() + up_axis = UsdGeom.GetStageUpAxis(stage) + builder = ModelBuilder(up_axis=up_axis) + builder.add_usd(stage) + cls.set_builder(builder) + + @classmethod + def _initialize_contacts(cls) -> None: + """Unified method to initialize contacts and collision pipeline. + + This method handles both Newton collision pipeline and MuJoCo contact modes. + It ensures contacts are properly initialized with force attributes if sensors are registered. + """ + if cls._needs_collision_pipeline: + # Newton collision pipeline: create pipeline and generate contacts + if cls._collision_pipeline is None: + cls._collision_pipeline = CollisionPipeline(cls._model, broad_phase_mode=BroadPhaseMode.EXPLICIT) + if cls._contacts is None: + cls._contacts = cls._collision_pipeline.contacts() + + elif cls._solver is not None and isinstance(cls._solver, SolverMuJoCo): + # MuJoCo contacts mode: create properly sized Contacts object + # The solver's update_contacts() will populate this from MuJoCo data + rigid_contact_max = cls._solver.get_max_contact_count() + cls._contacts = Contacts( + rigid_contact_max=rigid_contact_max, + soft_contact_max=0, + device=PhysicsManager._device, + requested_attributes=cls._model.get_requested_contact_attributes(), + ) + + @classmethod + def initialize_solver(cls) -> None: + """Initialize the solver and collision pipeline. + + This function initializes the solver based on the specified solver type. Currently, only XPBD and MuJoCoWarp + are supported. If the solver requires external collision detection (i.e., not using MuJoCo's internal + contacts), a unified collision pipeline is created. + + The graphing of the simulation is performed in this function if the simulation is ran using + a CUDA enabled device. + + .. warning:: + When using a CUDA enabled device, the simulation will be graphed. This means that this function steps the + simulation once to capture the graph. Hence, this function should only be called after everything else in + the simulation is initialized. + """ + cfg = PhysicsManager._cfg + if cfg is None: + return + + with Timer(name="newton_initialize_solver", msg="Initialize solver took:", enable=True, format="ms"): + cls._num_substeps = cfg.num_substeps # type: ignore[union-attr] + cls._solver_dt = cls.get_physics_dt() / cls._num_substeps + + # Create solver from config + solver_cfg = cfg.solver_cfg # type: ignore[union-attr] + cfg_dict = solver_cfg.to_dict() if hasattr(solver_cfg, "to_dict") else {} + cls._solver_type = cfg_dict.pop("solver_type", "mujoco_warp") + + if cls._solver_type == "mujoco_warp": + # SolverMuJoCo does not require distinct input & output states + cls._use_single_state = True + cls._solver = SolverMuJoCo(cls._model, **cfg_dict) + elif cls._solver_type == "xpbd": + cls._use_single_state = False + cls._solver = SolverXPBD(cls._model, **cfg_dict) + elif cls._solver_type == "featherstone": + cls._use_single_state = False + cls._solver = SolverFeatherstone(cls._model, **cfg_dict) + else: + raise ValueError(f"Invalid solver type: {cls._solver_type}") + + # Determine if we need external collision detection + # - SolverMuJoCo with use_mujoco_contacts=True: uses internal MuJoCo collision detection + # - SolverMuJoCo with use_mujoco_contacts=False: needs Newton's unified collision pipeline + # - Other solvers (XPBD, Featherstone): always need Newton's unified collision pipeline + if isinstance(cls._solver, SolverMuJoCo): + # Handle both dict and object configs + if hasattr(solver_cfg, "use_mujoco_contacts"): + use_mujoco_contacts = solver_cfg.use_mujoco_contacts + elif isinstance(solver_cfg, dict): + use_mujoco_contacts = solver_cfg.get("use_mujoco_contacts", False) + else: + use_mujoco_contacts = getattr(solver_cfg, "use_mujoco_contacts", False) + cls._needs_collision_pipeline = not use_mujoco_contacts + else: + cls._needs_collision_pipeline = True + + # Initialize contacts and collision pipeline + cls._initialize_contacts() + + # Ensure we are using a CUDA enabled device + device = PhysicsManager._device + assert device.startswith("cuda"), "NewtonManager only supports CUDA enabled devices" + + with Timer(name="newton_cuda_graph", msg="CUDA graph took:", enable=True, format="ms"): + if cfg.use_cuda_graph: # type: ignore[union-attr] + with wp.ScopedCapture() as capture: + cls._simulate() + cls._graph = capture.graph + + @classmethod + def _simulate(cls) -> None: + """Run one simulation step with substeps.""" + + # MJWarp can use its internal collision pipeline. + if cls._needs_collision_pipeline: + cls._collision_pipeline.collide(cls._state_0, cls._contacts) + contacts = cls._contacts + else: + contacts = None + + def step_fn(state_0, state_1): + cls._solver.step(state_0, state_1, cls._control, contacts, cls._solver_dt) + + if cls._use_single_state: + for i in range(cls._num_substeps): + step_fn(cls._state_0, cls._state_0) + cls._state_0.clear_forces() + else: + cfg = PhysicsManager._cfg + need_copy_on_last_substep = (cfg is not None and cfg.use_cuda_graph) and cls._num_substeps % 2 == 1 # type: ignore[union-attr] + + for i in range(cls._num_substeps): + step_fn(cls._state_0, cls._state_1) + if need_copy_on_last_substep and i == cls._num_substeps - 1: + cls._state_0.assign(cls._state_1) + else: + cls._state_0, cls._state_1 = cls._state_1, cls._state_0 + cls._state_0.clear_forces() + + # Populate contacts for contact sensors + if cls._report_contacts: + # For newton_contacts (unified pipeline): use locally computed contacts + # For mujoco_contacts: use class-level _contacts, solver populates it from MuJoCo data + eval_contacts = contacts if contacts is not None else cls._contacts + cls._solver.update_contacts(eval_contacts, cls._state_0) + for sensor in cls._newton_contact_sensors.values(): + sensor.eval(eval_contacts) + + @classmethod + def get_solver_convergence_steps(cls) -> dict[str, float | int]: + """Get solver convergence statistics.""" + niter = cls._solver.mjw_data.solver_niter.numpy() + return { + "max": np.max(niter), + "mean": np.mean(niter), + "min": np.min(niter), + "std": np.std(niter), + } + + # State accessors (used extensively by articulation/rigid object data) + @classmethod + def get_model(cls) -> Model: + """Get the Newton model.""" + return cls._model + + @classmethod + def get_state_0(cls) -> State: + """Get the current state.""" + return cls._state_0 + + @classmethod + def get_state_1(cls) -> State: + """Get the next state.""" + return cls._state_1 + + @classmethod + def get_control(cls) -> Control: + """Get the control object.""" + return cls._control + + @classmethod + def get_dt(cls) -> float: + """Get the physics timestep. Alias for get_physics_dt().""" + return cls.get_physics_dt() + + @classmethod + def get_solver_dt(cls) -> float: + """Get the solver substep timestep.""" + return cls._solver_dt + + @classmethod + def add_contact_sensor( + cls, + body_names_expr: str | list[str] | None = None, + shape_names_expr: str | list[str] | None = None, + contact_partners_body_expr: str | list[str] | None = None, + contact_partners_shape_expr: str | list[str] | None = None, + prune_noncolliding: bool = True, + verbose: bool = False, + ) -> None: + """Add a contact sensor for reporting contacts between bodies/shapes. + + Note: Only one contact sensor can be active at a time. + + Args: + body_names_expr: Expression for body names to sense. + shape_names_expr: Expression for shape names to sense. + contact_partners_body_expr: Expression for contact partner body names. + contact_partners_shape_expr: Expression for contact partner shape names. + prune_noncolliding: Make force matrix sparse using collision pairs. + verbose: Print verbose information. + """ + # Validate inputs + if body_names_expr is None and shape_names_expr is None: + raise ValueError("At least one of body_names_expr or shape_names_expr must be provided") + if body_names_expr is not None and shape_names_expr is not None: + raise ValueError("Only one of body_names_expr or shape_names_expr must be provided") + if contact_partners_body_expr is not None and contact_partners_shape_expr is not None: + raise ValueError("Only one of contact_partners_body_expr or contact_partners_shape_expr must be provided") + + # Log sensor configuration + sensor_target = body_names_expr or shape_names_expr + partner_filter = contact_partners_body_expr or contact_partners_shape_expr or "all bodies/shapes" + logger.info(f"Adding contact sensor for {sensor_target} with filter {partner_filter}") + + # Create unique key for this sensor + sensor_key = (body_names_expr, shape_names_expr, contact_partners_body_expr, contact_partners_shape_expr) + + # Create and store the sensor + # Note: SensorContact constructor requests 'force' attribute from the model + newton_sensor = NewtonContactSensor( + cls._model, + sensing_obj_bodies=body_names_expr, + sensing_obj_shapes=shape_names_expr, + counterpart_bodies=contact_partners_body_expr, + counterpart_shapes=contact_partners_shape_expr, + match_fn=flipped_match, + include_total=True, + prune_noncolliding=prune_noncolliding, + verbose=verbose, + ) + cls._newton_contact_sensors[sensor_key] = newton_sensor + cls._report_contacts = True + + # Regenerate contacts only if they were already created without force attribute + # If solver is not initialized, contacts will be created with force in initialize_solver() + if cls._solver is not None and cls._contacts is not None: + # Only regenerate if contacts don't have force attribute (sensor.eval() requires it) + if cls._contacts.force is None: + cls._initialize_contacts() + + return sensor_key diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py new file mode 100644 index 00000000000..b4253c8e884 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py @@ -0,0 +1,215 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for Newton physics manager.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.physics import PhysicsCfg +from isaaclab.utils import DeferredClass, configclass + +if TYPE_CHECKING: + from isaaclab.physics import PhysicsManager + + from .newton_manager import NewtonManager + + +@configclass +class NewtonSolverCfg: + """Configuration for Newton solver-related parameters. + + These parameters are used to configure the Newton solver. For more information, see the `Newton documentation`_. + + .. _Newton documentation: https://newton.readthedocs.io/en/latest/ + """ + + solver_type: str = "None" + """Solver type. + + Used to select the right solver class. + """ + + +@configclass +class MJWarpSolverCfg(NewtonSolverCfg): + """Configuration for MuJoCo Warp solver-related parameters. + + These parameters are used to configure the MuJoCo Warp solver. For more information, see the + `MuJoCo Warp documentation`_. + + .. _MuJoCo Warp documentation: https://github.com/google-deepmind/mujoco_warp + """ + + solver_type: str = "mujoco_warp" + """Solver type. Can be "mujoco_warp".""" + + njmax: int = 300 + """Number of constraints per environment (world).""" + + nconmax: int | None = None + """Number of contact points per environment (world).""" + + iterations: int = 100 + """Number of solver iterations.""" + + ls_iterations: int = 50 + """Number of line search iterations for the solver.""" + + solver: str = "newton" + """Solver type. Can be "cg" or "newton", or their corresponding MuJoCo integer constants.""" + + integrator: str = "euler" + """Integrator type. Can be "euler", "rk4", or "implicitfast", or their corresponding MuJoCo integer constants.""" + + use_mujoco_cpu: bool = False + """Whether to use the pure MuJoCo backend instead of `mujoco_warp`.""" + + disable_contacts: bool = False + """Whether to disable contact computation in MuJoCo.""" + + default_actuator_gear: float | None = None + """Default gear ratio for all actuators.""" + + actuator_gears: dict[str, float] | None = None + """Dictionary mapping joint names to specific gear ratios, overriding the `default_actuator_gear`.""" + + update_data_interval: int = 1 + """Frequency (in simulation steps) at which to update the MuJoCo Data object from the Newton state. + + If 0, Data is never updated after initialization. + """ + + save_to_mjcf: str | None = None + """Optional path to save the generated MJCF model file. + + If None, the MJCF model is not saved. + """ + + impratio: float = 1.0 + """Frictional-to-normal constraint impedance ratio.""" + + cone: str = "pyramidal" + """The type of contact friction cone. Can be "pyramidal" or "elliptic".""" + + ls_parallel: bool = False + """Whether to use parallel line search.""" + + use_mujoco_contacts: bool = True + """Whether to use MuJoCo's contact solver.""" + + +@configclass +class XPBDSolverCfg(NewtonSolverCfg): + """An implicit integrator using eXtended Position-Based Dynamics (XPBD) for rigid and soft body simulation. + + References: + - Miles Macklin, Matthias Müller, and Nuttapong Chentanez. 2016. XPBD: position-based simulation of compliant + constrained dynamics. In Proceedings of the 9th International Conference on Motion in Games (MIG '16). + Association for Computing Machinery, New York, NY, USA, 49-54. https://doi.org/10.1145/2994258.2994272 + - Matthias Müller, Miles Macklin, Nuttapong Chentanez, Stefan Jeschke, and Tae-Yong Kim. 2020. Detailed rigid + body simulation with extended position based dynamics. In Proceedings of the ACM SIGGRAPH/Eurographics + Symposium on Computer Animation (SCA '20). Eurographics Association, Goslar, DEU, + Article 10, 1-12. https://doi.org/10.1111/cgf.14105 + + """ + + solver_type: str = "xpbd" + """Solver type. Can be "xpbd".""" + + iterations: int = 2 + """Number of solver iterations.""" + + soft_body_relaxation: float = 0.9 + """Relaxation parameter for soft body simulation.""" + + soft_contact_relaxation: float = 0.9 + """Relaxation parameter for soft contact simulation.""" + + joint_linear_relaxation: float = 0.7 + """Relaxation parameter for joint linear simulation.""" + + joint_angular_relaxation: float = 0.4 + """Relaxation parameter for joint angular simulation.""" + + joint_linear_compliance: float = 0.0 + """Compliance parameter for joint linear simulation.""" + + joint_angular_compliance: float = 0.0 + """Compliance parameter for joint angular simulation.""" + + rigid_contact_relaxation: float = 0.8 + """Relaxation parameter for rigid contact simulation.""" + + rigid_contact_con_weighting: bool = True + """Whether to use contact constraint weighting for rigid contact simulation.""" + + angular_damping: float = 0.0 + """Angular damping parameter for rigid contact simulation.""" + + enable_restitution: bool = False + """Whether to enable restitution for rigid contact simulation.""" + + +@configclass +class FeatherstoneSolverCfg(NewtonSolverCfg): + """A semi-implicit integrator using symplectic Euler. + + It operates on reduced (also called generalized) coordinates to simulate articulated rigid body dynamics + based on Featherstone's composite rigid body algorithm (CRBA). + + See: Featherstone, Roy. Rigid Body Dynamics Algorithms. Springer US, 2014. + + Semi-implicit time integration is a variational integrator that + preserves energy, however it not unconditionally stable, and requires a time-step + small enough to support the required stiffness and damping forces. + + See: https://en.wikipedia.org/wiki/Semi-implicit_Euler_method + """ + + solver_type: str = "featherstone" + """Solver type. Can be "featherstone".""" + + angular_damping: float = 0.05 + """Angular damping parameter for rigid contact simulation.""" + + update_mass_matrix_interval: int = 1 + """Frequency (in simulation steps) at which to update the mass matrix.""" + + friction_smoothing: float = 1.0 + """Friction smoothing parameter.""" + + use_tile_gemm: bool = False + """Whether to use tile-based GEMM for the mass matrix.""" + + fuse_cholesky: bool = True + """Whether to fuse the Cholesky decomposition.""" + + +@configclass +class NewtonCfg(PhysicsCfg): + """Configuration for Newton physics manager. + + This configuration includes Newton-specific simulation settings and solver configuration. + """ + + class_type: type[PhysicsManager] | DeferredClass = DeferredClass("isaaclab_newton.physics.newton_manager:NewtonManager") + """The class type of the NewtonManager.""" + + num_substeps: int = 1 + """Number of substeps to use for the solver.""" + + debug_mode: bool = False + """Whether to enable debug mode for the solver.""" + + use_cuda_graph: bool = True + """Whether to use CUDA graphing when simulating. + + If set to False, the simulation performance will be severely degraded. + """ + + solver_cfg: NewtonSolverCfg = MJWarpSolverCfg() + """Solver configuration. Default is MJWarpSolverCfg().""" diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/__init__.py b/source/isaaclab_newton/isaaclab_newton/sensors/__init__.py new file mode 100644 index 00000000000..4d942654837 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-package containing Newton-specific sensor implementations.""" + +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "contact_sensor": ["ContactSensor", "ContactSensorCfg", "ContactSensorData"], + "tiled_camera_cfg": ["TiledCameraCfgNewtonWarpRenderer"], + }, +) diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/__init__.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/__init__.py new file mode 100644 index 00000000000..84392c0100b --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Sub-module for contact sensor based on :class:`newton.SensorContact`.""" + +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "contact_sensor": ["ContactSensor"], + "contact_sensor_cfg": ["ContactSensorCfg"], + "contact_sensor_data": ["ContactSensorData"], + }, +) diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py new file mode 100644 index 00000000000..8d77a72df39 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor.py @@ -0,0 +1,529 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Ignore optional memory usage warning globally +# pyright: reportOptionalSubscript=false + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import warp as wp +from isaaclab_newton.physics import NewtonManager +from newton.sensors import MatchKind +from newton.sensors import SensorContact as NewtonContactSensor + +import isaaclab.utils.string as string_utils +from isaaclab.markers import VisualizationMarkers +from isaaclab.sensors.contact_sensor.base_contact_sensor import BaseContactSensor +from isaaclab.utils.helpers import deprecated + +from .contact_sensor_data import ContactSensorData +from .contact_sensor_kernels import ( + compute_first_transition_kernel, + copy_from_newton_kernel, + reset_contact_sensor_kernel, + reset_sensor_state_kernel, + update_contact_sensor_kernel, + update_outdated_envs_kernel, + update_timestamp_kernel, +) + +if TYPE_CHECKING: + from .contact_sensor_cfg import ContactSensorCfg + +logger = logging.getLogger(__name__) + + +class ContactSensor(BaseContactSensor): + """A contact reporting sensor. + + The contact sensor reports the normal contact forces on a rigid body or shape in the world frame. + + The sensor can be configured to report the contact forces on a set of sensors (bodies or shapes) + against specific filter objects using the :attr:`ContactSensorCfg.filter_prim_paths_expr`. This is + useful when you want to report the contact forces between the sensors and a specific set of objects + in the scene. The data can be accessed using the :attr:`ContactSensorData.force_matrix_w`. + + .. _Newton SensorContact: https://newton-physics.github.io/newton/api/_generated/newton.sensors.SensorContact.html + """ + + cfg: ContactSensorCfg + """The configuration parameters.""" + + def __init__(self, cfg: ContactSensorCfg): + """Initializes the contact sensor object. + + Args: + cfg: The configuration parameters. + """ + # initialize base class + super().__init__(cfg) + + # Create empty variables for storing output data + self._data: ContactSensorData = ContactSensorData() + + def __str__(self) -> str: + """Returns: A string containing information about the instance.""" + return ( + f"Contact sensor @ '{self.cfg.prim_path}': \n" + f"\tupdate period (s) : {self.cfg.update_period}\n" + f"\tnumber of sensors : {self.num_sensors}\n" + f"\tsensor names : {self.sensor_names}\n" + ) + + """ + Properties + """ + + @property + def num_instances(self) -> int | None: + return self._num_sensors + + @property + def data(self) -> ContactSensorData: + # update sensors if needed + self._update_outdated_buffers() + # return the data + return self._data + + @property + def num_sensors(self) -> int: + """Number of sensors (bodies or shapes with contact sensing attached).""" + return self._num_sensors + + @property + def sensor_names(self) -> list[str] | None: + """Ordered names of sensors (shapes or bodies with contact sensing attached).""" + return self._sensor_names + + @property + def filter_object_names(self) -> list[str] | None: + return self._filter_object_names + + @property + def num_filter_objects(self) -> int: + """Number of filter objects (counterparts) for contact filtering.""" + return self._num_filter_objects + + @property + def contact_view(self) -> NewtonContactSensor: + """View for the contact forces captured (Newton).""" + return NewtonManager._newton_contact_sensors[self._sensor_key] + + """ + Operations + """ + + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None): + # Resolve env_mask + if env_ids is None and env_mask is None: + env_mask = wp.full(shape=self._num_envs, value=True, dtype=wp.bool, device=self._device) + elif env_mask is None and env_ids is not None: + from isaaclab.utils.warp.utils import make_mask_from_torch_ids + + env_mask = make_mask_from_torch_ids(self._num_envs, env_ids, device=self._device) + + # Reset timestamps and outdated flags + wp.launch( + reset_sensor_state_kernel, + dim=self._num_envs, + inputs=[env_mask, self._is_outdated, self._timestamp, self._timestamp_last_update], + device=self._device, + ) + + # Compute num_filter_objects + num_filter_objects = self._num_filter_objects + + # Reset contact sensor buffers via kernel + wp.launch( + reset_contact_sensor_kernel, + dim=(self._num_envs, self._num_sensors), + inputs=[ + self.cfg.history_length, + num_filter_objects, + env_mask, + self._data._net_forces_w, + self._data._net_forces_w_history, + self._data._force_matrix_w, + ], + outputs=[ + self._data._current_air_time, + self._data._last_air_time, + self._data._current_contact_time, + self._data._last_contact_time, + ], + device=self._device, + ) + + def find_sensors( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[wp.array, list[str], list[int]]: + """Find sensors based on the name keys. + + Args: + name_keys: A regular expression or a list of regular expressions to match the sensor names. + preserve_order: Whether to preserve the order of the name keys in the output. Defaults to False. + + Returns: + A tuple containing the sensor mask (wp.array), names (list[str]), and indices (list[int]). + + Raises: + ValueError: If sensor_names is None (sensor not initialized). + """ + if self.sensor_names is None: + raise ValueError("Sensor names not available. Is the sensor initialized?") + indices, names = string_utils.resolve_matching_names(name_keys, self.sensor_names, preserve_order) + mask = wp.array([name in names for name in self.sensor_names], dtype=wp.bool, device=self._device) + return mask, names, indices + + def compute_first_contact(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: + """Checks if sensors that have established contact within the last :attr:`dt` seconds. + + This function checks if the sensors have established contact within the last :attr:`dt` seconds + by comparing the current contact time with the given time period. If the contact time is less + than the given time period, then the sensors are considered to be in contact. + + Note: + The function assumes that :attr:`dt` is a factor of the sensor update time-step. In other + words :math:`dt / dt_sensor = n`, where :math:`n` is a natural number. This is always true + if the sensor is updated by the physics or the environment stepping time-step and the sensor + is read by the environment stepping time-step. + + Args: + dt: The time period since the contact was established. + abs_tol: The absolute tolerance for the comparison. + + Returns: + A boolean array indicating the sensors that have established contact within the last + :attr:`dt` seconds. Shape is (N, S), where N is the number of environments and S is the + number of sensors. + + Raises: + RuntimeError: If the sensor is not configured to track contact time. + """ + # check if the sensor is configured to track contact time + if not self.cfg.track_air_time: + raise RuntimeError( + "The contact sensor is not configured to track contact time." + "Please enable the 'track_air_time' in the sensor configuration." + ) + wp.launch( + compute_first_transition_kernel, + dim=(self._num_envs, self._num_sensors), + inputs=[float(dt + abs_tol), self._data._current_contact_time], + outputs=[self._data._first_transition], + device=self._device, + ) + return self._data._first_transition + + def compute_first_air(self, dt: float, abs_tol: float = 1.0e-8) -> wp.array: + """Checks if sensors that have broken contact within the last :attr:`dt` seconds. + + This function checks if the sensors have broken contact within the last :attr:`dt` seconds + by comparing the current air time with the given time period. If the air time is less + than the given time period, then the sensors are considered to not be in contact. + + Note: + It assumes that :attr:`dt` is a factor of the sensor update time-step. In other words, + :math:`dt / dt_sensor = n`, where :math:`n` is a natural number. This is always true if + the sensor is updated by the physics or the environment stepping time-step and the sensor + is read by the environment stepping time-step. + + Args: + dt: The time period since the contract is broken. + abs_tol: The absolute tolerance for the comparison. + + Returns: + A boolean array indicating the sensors that have broken contact within the last :attr:`dt` seconds. + Shape is (N, S), where N is the number of environments and S is the number of sensors. + + Raises: + RuntimeError: If the sensor is not configured to track contact time. + """ + # check if the sensor is configured to track contact time + if not self.cfg.track_air_time: + raise RuntimeError( + "The contact sensor is not configured to track contact time." + "Please enable the 'track_air_time' in the sensor configuration." + ) + + wp.launch( + compute_first_transition_kernel, + dim=(self._num_envs, self._num_sensors), + inputs=[float(dt + abs_tol), self._data._current_air_time], + outputs=[self._data._first_transition], + device=self._device, + ) + return self._data._first_transition + + def update(self, dt: float, force_recompute: bool = False): + """Updates the sensor data.""" + if not self._is_initialized: + return + # Update the timestamp for the sensors + wp.launch( + update_timestamp_kernel, + dim=self._num_envs, + inputs=[ + self._is_outdated, + self._timestamp, + self._timestamp_last_update, + self._sim_physics_dt, + self.cfg.update_period, + ], + device=self._device, + ) + # Update the buffers + if force_recompute or self._is_visualizing: + self._update_outdated_buffers() + + def _update_outdated_buffers(self): + """Fills the sensor data for the outdated sensors.""" + self._update_buffers_impl_mask(self._is_outdated) + # Update timestamps and clear outdated flags + wp.launch( + update_outdated_envs_kernel, + dim=self._num_envs, + inputs=[self._is_outdated, self._timestamp, self._timestamp_last_update], + device=self._device, + ) + + """ + Implementation. + """ + + def _initialize_impl(self): + """Initializes the sensor-related handles and internal buffers.""" + # Obtain Simulation Context + import isaaclab.sim as sim_utils + + sim = sim_utils.SimulationContext.instance() + if sim is None: + raise RuntimeError("Simulation Context is not initialized!") + # Obtain device and backend + self._device = sim.device + self._backend = sim.backend + num_envs = NewtonManager._num_envs + if num_envs is None: + raise RuntimeError("NewtonManager._num_envs is None. Is the simulation initialized?") + self._num_envs = num_envs + self._sim_physics_dt = wp.full( + shape=self._num_envs, value=sim.get_physics_dt(), dtype=wp.float32, device=self._device + ) + self._is_outdated = wp.full(shape=self._num_envs, value=True, dtype=wp.bool, device=self._device) + self._timestamp = wp.zeros(shape=self._num_envs, dtype=wp.float32, device=self._device) + self._timestamp_last_update = wp.zeros_like(self._timestamp) + + # construct regex expression for the sensor names + + if self.cfg.filter_prim_paths_expr is not None or self.cfg.filter_shape_paths_expr is not None: + self._generate_force_matrix = True + else: + self._generate_force_matrix = False + + sensor_body_regex = self.cfg.prim_path + if self.cfg.shape_path is not None: + sensor_shape_regex = "(" + "|".join(self.cfg.shape_path) + ")" + else: + sensor_shape_regex = None + if self.cfg.filter_prim_paths_expr is not None: + filter_object_body_regex = "(" + "|".join(self.cfg.filter_prim_paths_expr) + ")" + else: + filter_object_body_regex = None + if self.cfg.filter_shape_paths_expr is not None: + filter_object_shape_regex = "(" + "|".join(self.cfg.filter_shape_paths_expr) + ")" + else: + filter_object_shape_regex = None + + # Store the sensor key for later lookup + self._sensor_key = ( + sensor_body_regex, + sensor_shape_regex, + filter_object_body_regex, + filter_object_shape_regex, + ) + + NewtonManager.add_contact_sensor( + body_names_expr=sensor_body_regex, + shape_names_expr=sensor_shape_regex, + contact_partners_body_expr=filter_object_body_regex, + contact_partners_shape_expr=filter_object_shape_regex, + prune_noncolliding=True, + ) + + self._create_buffers() + + def _create_buffers(self): + # Get Newton sensor shape: (n_sensors * n_envs, n_counterparts) + newton_shape = self.contact_view.shape + + # resolve the true count of sensors + self._num_sensors = newton_shape[0] // self._num_envs + + # Check that number of sensors is an integer + if newton_shape[0] % self._num_envs != 0: + raise RuntimeError( + "Number of sensors is not an integer multiple of the number of environments. Received:" + f" {self._num_sensors} sensors and {self._num_envs} environments." + ) + logger.info(f"Contact sensor initialized with {self._num_sensors} sensors.") + + # Assume homogeneous envs, i.e. all envs have the same number of sensors + # Only get the names for the first env. Expected structure: /World/envs/env_.*/... + def get_name(idx, match_kind): + if match_kind == MatchKind.BODY: + return NewtonManager._model.body_key[idx].split("/")[-1] + if match_kind == MatchKind.SHAPE: + return NewtonManager._model.shape_key[idx].split("/")[-1] + return "MATCH_ANY" + + self._sensor_names = [get_name(idx, kind) for idx, kind in self.contact_view.sensing_objs] + # Assumes the environments are processed in order. + self._sensor_names = self._sensor_names[: self._num_sensors] + self._filter_object_names = [get_name(idx, kind) for idx, kind in self.contact_view.counterparts] + + # Number of filter objects (counterparts minus the total column) + self._num_filter_objects = max(newton_shape[1] - 1, 0) + + # Store reshaped Newton net_force view for copying data + # Newton net_force shape: (n_sensors * n_envs, n_counterparts) + # Reshaped to: (n_envs, n_sensors, n_counterparts) + self._newton_forces_view = self.contact_view.net_force.reshape((self._num_envs, self._num_sensors, -1)) + + # prepare data buffers + logger.info( + f"Creating buffers for contact sensor data with num_envs: {self._num_envs}, num_sensors:" + f" {self._num_sensors}, num_filter_objects: {self._num_filter_objects}, history_length:" + f" {self.cfg.history_length}, generate_force_matrix: {self._generate_force_matrix}, track_air_time:" + f" {self.cfg.track_air_time}, track_pose: {self.cfg.track_pose}, device: {self._device}" + ) + self._data.create_buffers( + self._num_envs, + self._num_sensors, + self._num_filter_objects, + self.cfg.history_length, + self._generate_force_matrix, + self.cfg.track_air_time, + self.cfg.track_pose, + self._device, + ) + + def _update_buffers_impl_mask(self, env_mask: wp.array): + """Fills the buffers of the sensor data (mask-based API for Newton). + + Args: + env_mask: Boolean mask array indicating which environments to update. + """ + # Copy data from Newton into owned buffers (respecting env_mask) + # Launch with 3D for coalescing: dim=(num_envs, num_sensors, max(num_filter_objects, 1)) + wp.launch( + copy_from_newton_kernel, + dim=(self._num_envs, self._num_sensors, max(self._num_filter_objects, 1)), + inputs=[ + env_mask, + self._newton_forces_view, + ], + outputs=[ + self._data._net_forces_w, + self._data._force_matrix_w, + ], + device=self._device, + ) + + # Update history and air/contact time tracking + wp.launch( + update_contact_sensor_kernel, + dim=(self._num_envs, self._num_sensors), + inputs=[ + self.cfg.history_length, + self.cfg.force_threshold, + env_mask, + self._data._net_forces_w, + self._timestamp, + self._timestamp_last_update, + self._data._net_forces_w_history, + self._data._current_air_time, + self._data._current_contact_time, + self._data._last_air_time, + self._data._last_contact_time, + ], + device=self._device, + ) + + # FIXME: Re-enable this when we have a non-physx rigid body view? + # (tracked in https://github.com/newton-physics/newton/issues/1489) + # obtain the pose of the sensor origin + # if self.cfg.track_pose: + # pose = self.body_physx_view.get_transforms().view(-1, self._num_sensors, 7)[env_ids] + # pose[..., 3:] = convert_quat(pose[..., 3:], to="wxyz") + # self._data.pos_w[env_ids], self._data.quat_w[env_ids] = pose.split([3, 4], dim=-1) + + def _set_debug_vis_impl(self, debug_vis: bool): + # set visibility of markers + # note: parent only deals with callbacks. not their visibility + if debug_vis: + # create markers if necessary for the first tome + if not hasattr(self, "contact_visualizer"): + self.contact_visualizer = VisualizationMarkers(self.cfg.visualizer_cfg) + # set their visibility to true + self.contact_visualizer.set_visibility(True) + else: + if hasattr(self, "contact_visualizer"): + self.contact_visualizer.set_visibility(False) + + def _debug_vis_callback(self, event): + # safely return if view becomes invalid + return + # note: this invalidity happens because of isaac sim view callbacks + # if self.body_physx_view is None: + # return + # marker indices + # 0: contact, 1: no contact + # net_contact_force_w = torch.norm(self._data.net_forces_w, dim=-1) + # marker_indices = torch.where(net_contact_force_w > self.cfg.force_threshold, 0, 1) + # check if prim is visualized + # if self.cfg.track_pose: + # frame_origins: torch.Tensor = self._data.pos_w + # else: + # pose = self.body_physx_view.get_transforms() + # frame_origins = pose.view(-1, self._num_sensors, 7)[:, :, :3] + # visualize + # self.contact_visualizer.visualize(frame_origins.view(-1, 3), marker_indices=marker_indices.view(-1)) + + """ + Internal simulation callbacks. + """ + + def _invalidate_initialize_callback(self, event): + """Invalidates the scene elements.""" + # call parent + super()._invalidate_initialize_callback(event) + # set all existing views to None to invalidate them + # TODO: invalidate NewtonManager if necessary + + """ + Renamed + """ + + @property + @deprecated("use num_sensors") + def num_bodies(self) -> int: + return self.num_sensors + + @property + @deprecated("use sensor_names") + def body_names(self) -> list[str] | None: + return self.sensor_names + + @deprecated("use find_sensors") + def find_bodies( + self, name_keys: str | Sequence[str], preserve_order: bool = False + ) -> tuple[list[int], list[str]]: + """Find bodies (sensors) by name. Returns (indices, names) for compatibility with base class.""" + _, names, indices = self.find_sensors(name_keys, preserve_order) + return indices, names diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_cfg.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_cfg.py new file mode 100644 index 00000000000..e423848ecdf --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_cfg.py @@ -0,0 +1,45 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.sensors.contact_sensor import ContactSensorCfg as BaseContactSensorCfg +from isaaclab.utils import DeferredClass, configclass + + +@configclass +class ContactSensorCfg(BaseContactSensorCfg): + """Configuration for the Newton contact sensor. + + Extends :class:`isaaclab.sensors.ContactSensorCfg` with Newton-specific fields. + """ + + class_type: type | DeferredClass = DeferredClass("isaaclab_newton.sensors.contact_sensor.contact_sensor:ContactSensor") + + shape_path: list[str] | None = None + """A list of shape path expressions to filter which shapes to sense contacts ON. Defaults to None, + meaning all shapes on the bodies specified by :attr:`prim_path` are sensed. + + A body (rigid body) can have multiple collision shapes. This option allows more fine-grained + control over which specific shapes to sense contacts on. + + .. note:: + The expression in the list can contain the environment namespace regex ``{ENV_REGEX_NS}`` which + will be replaced with the environment namespace. + """ + + filter_shape_paths_expr: list[str] | None = None + """A list of shape path expressions to filter contacts WITH. Defaults to None, meaning contacts + with all shapes are reported. + + This is similar to :attr:`filter_prim_paths_expr` but filters at the shape level instead of body level. + A body (rigid body) can have multiple collision shapes. This option allows more fine-grained + control over which specific shapes to detect contacts with. + + .. note:: + The expression in the list can contain the environment namespace regex ``{ENV_REGEX_NS}`` which + will be replaced with the environment namespace. + + Example: ``{ENV_REGEX_NS}/Object/CollisionMesh`` will be replaced with + ``/World/envs/env_.*/Object/CollisionMesh``. + """ diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_data.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_data.py new file mode 100644 index 00000000000..82ace1d0faa --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_data.py @@ -0,0 +1,257 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# needed to import for allowing type-hinting: torch.Tensor | None +from __future__ import annotations + +import logging + +import warp as wp + +from isaaclab.sensors.contact_sensor.base_contact_sensor_data import BaseContactSensorData + +logger = logging.getLogger(__name__) + + +class ContactSensorData(BaseContactSensorData): + """Data container for the contact reporting sensor.""" + + _pos_w: wp.array | None + _quat_w: wp.array | None + + _net_forces_w: wp.array | None + _net_forces_w_history: wp.array | None + _force_matrix_w: wp.array | None + _force_matrix_w_history: wp.array | None + _last_air_time: wp.array | None + _current_air_time: wp.array | None + _last_contact_time: wp.array | None + _current_contact_time: wp.array | None + _first_transition: wp.array | None + + @property + def pos_w(self) -> wp.array | None: + """Position of the sensor origin in world frame. + + `wp.vec3f` array whose shape is (N,) where N is the number of sensors. Note, that when casted to as a + `torch.Tensor`, the shape will be (N, 3). + + Note: + If the :attr:`ContactSensorCfg.track_pose` is False, then this quantity is None. + """ + return self._pos_w + + @property + def quat_w(self) -> wp.array | None: + """Orientation of the sensor origin in quaternion (x, y, z, w) in world frame. + + `wp.quatf` whose shape is (N,) where N is the number of sensors. Note, that when casted to as a `torch.Tensor`, + the shape will be (N, 4). + + Note: + If the :attr:`ContactSensorCfg.track_pose` is False, then this quantity is None. + """ + return self._quat_w + + @property + def pose_w(self) -> wp.array | None: + """Pose of the sensor origin in world frame. Shape is (N, 7). Quaternion in xyzw order. + + Not implemented for Newton backend; use :attr:`pos_w` and :attr:`quat_w` instead. Returns None. + """ + return NotImplementedError() + + @property + def contact_pos_w(self) -> wp.array | None: + """Average position of contact points. Shape is (N, B, M, 3). + + Not yet implemented for Newton backend. Returns None. + """ + return NotImplementedError() + + @property + def friction_forces_w(self) -> wp.array | None: + """Sum of friction forces. Shape is (N, B, M, 3). + + Not yet implemented for Newton backend. Returns None. + """ + return NotImplementedError() + + @property + def net_forces_w(self) -> wp.array2d | None: + """The net (total) contact forces in world frame. + + `wp.vec3f` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. + Note, that when casted to as a `torch.Tensor`, the shape will be (N, S, 3). + + Note: + This quantity is the sum of the contact forces acting on each sensor. It must not be confused + with the total contact forces acting on the sensors (which also includes the tangential forces). + """ + return self._net_forces_w + + @property + def net_forces_w_history(self) -> wp.array3d | None: + """The net (total) contact forces in world frame. + + `wp.vec3f` array whose shape is (N, T, S) where N is the number of environments, T is the configured history + length and S is the number of sensors. Note, that when casted to as a `torch.Tensor`, the shape will be + (N, T, S, 3). + + In the history dimension, the first index is the most recent and the last index is the oldest. + + Note: + This quantity is the sum of the contact forces acting on each sensor. It must not be confused + with the total contact forces acting on the sensors (which also includes the tangential forces). + """ + return self._net_forces_w_history + + @property + def force_matrix_w(self) -> wp.array3d | None: + """The contact forces between sensors and filter objects in world frame. + + `wp.vec3f` array whose shape is (N, S, F) where N is the number of environments, S is number of sensors + and F is the number of filter objects. Note, that when casted to as a `torch.Tensor`, the shape will be + (N, S, F, 3). + + Note: + If the :attr:`ContactSensorCfg.filter_prim_paths_expr` is empty, then this quantity is None. + """ + return self._force_matrix_w + + @property + def force_matrix_w_history(self) -> wp.array4d | None: + """The contact forces between sensors and filter objects in world frame. + + `wp.vec3f` array whose shape is (N, T, S, F) where N is the number of environments, T is the configured history + length, S is number of sensors and F is the number of filter objects. Note, that when casted to as a + `torch.Tensor`, the shape will be (N, T, S, F, 3). + + In the history dimension, the first index is the most recent and the last index is the oldest. + + Note: + If the :attr:`ContactSensorCfg.filter_prim_paths_expr` is empty, then this quantity is None. + """ + return self._force_matrix_w_history + + @property + def last_air_time(self) -> wp.array2d | None: + """Time spent (in s) in the air before the last contact. + + `wp.float32` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. + Note, that when casted to as a `torch.Tensor`, the shape will be (N, S). + + Note: + If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. + """ + return self._last_air_time + + @property + def current_air_time(self) -> wp.array2d | None: + """Time spent (in s) in the air since the last detach. + + `wp.float32` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. + Note, that when casted to as a `torch.Tensor`, the shape will be (N, S). + + Note: + If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. + """ + return self._current_air_time + + @property + def last_contact_time(self) -> wp.array2d | None: + """Time spent (in s) in contact before the last detach. + + `wp.float32` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. + Note, that when casted to as a `torch.Tensor`, the shape will be (N, S). + + Note: + If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. + """ + return self._last_contact_time + + @property + def current_contact_time(self) -> wp.array2d | None: + """Time spent (in s) in contact since the last contact. + + `wp.float32` array whose shape is (N, S) where N is the number of environments and S is the number of sensors. + Note, that when casted to as a `torch.Tensor`, the shape will be (N, S). + + Note: + If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None. + """ + return self._current_contact_time + + def create_buffers( + self, + num_envs: int, + num_sensors: int, + num_filter_objects: int, + history_length: int, + generate_force_matrix: bool, + track_air_time: bool, + track_pose: bool, + device: str, + ) -> None: + """Creates the buffers for the contact sensor data. + + Args: + num_envs: The number of environments. + num_sensors: The number of sensors. + num_filter_objects: The number of filter objects (counterparts). + history_length: The history length. + generate_force_matrix: Whether to generate the force matrix. + track_air_time: Whether to track the air time. + track_pose: Whether to track the pose. + device: The device to use. + """ + logger.info( + f"Creating buffers for contact sensor data with num_envs: {num_envs}, num_sensors: {num_sensors}," + f" num_filter_objects: {num_filter_objects}, history_length: {history_length}, generate_force_matrix:" + f" {generate_force_matrix}, track_air_time: {track_air_time}, track_pose: {track_pose}, device: {device}" + ) + # Track pose if requested + if track_pose: + self._pose = wp.zeros((num_envs,), dtype=wp.transformf, device=device) + pos_scalars = wp.array(self._pose, dtype=wp.float32, device=device, copy=False) + self._pos_w = wp.array(pos_scalars[:, :3], dtype=wp.vec3f, device=device, copy=False) + self._quat_w = wp.array(pos_scalars[:, 3:], dtype=wp.quatf, device=device, copy=False) + else: + self._pose = None + self._pos_w = None + self._quat_w = None + + # Create owned buffer for net (total) forces - shape: (num_envs, num_sensors) + self._net_forces_w = wp.zeros((num_envs, num_sensors), dtype=wp.vec3f, device=device) + # Track net forces history if requested + if history_length > 0: + self._net_forces_w_history = wp.zeros( + (num_envs, history_length, num_sensors), dtype=wp.vec3f, device=device + ) + self._force_matrix_w_history = None # TODO: implement force matrix history if needed + else: + self._net_forces_w_history = None + self._force_matrix_w_history = None + + # Create owned buffer for force matrix - shape: (num_envs, num_sensors, num_filter_objects) + # None if no filter objects configured + if num_filter_objects > 0: + self._force_matrix_w = wp.zeros((num_envs, num_sensors, num_filter_objects), dtype=wp.vec3f, device=device) + else: + self._force_matrix_w = None + + # Track air time if requested + if track_air_time: + self._last_air_time = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) + self._current_air_time = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) + self._last_contact_time = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) + self._current_contact_time = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) + self._first_transition = wp.zeros((num_envs, num_sensors), dtype=wp.float32, device=device) + else: + self._last_air_time = None + self._current_air_time = None + self._last_contact_time = None + self._current_contact_time = None + self._first_transition = None diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_kernels.py b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_kernels.py new file mode 100644 index 00000000000..55824b51193 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/contact_sensor/contact_sensor_kernels.py @@ -0,0 +1,213 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Ignore optional memory usage warning globally +# pyright: reportOptionalSubscript=false + +import warp as wp + + +@wp.kernel +def copy_from_newton_kernel( + # in + env_mask: wp.array(dtype=wp.bool), + newton_forces: wp.array3d(dtype=wp.vec3f), # (n_envs, n_sensors, n_counterparts) + # outputs + net_force_total: wp.array2d(dtype=wp.vec3f), # (n_envs, n_sensors) + force_matrix: wp.array3d(dtype=wp.vec3f), # (n_envs, n_sensors, n_filter_objects) or None +): + """Copy contact force data from Newton sensor into owned buffers. + + Launch with dim=(num_envs, num_sensors, max(num_filter_objects, 1)) for coalescing. + When num_filter_objects == 0, trailing dim is 1 and only total is copied. + """ + env, sensor, f_idx = wp.tid() + + if env_mask: + if not env_mask[env]: + return + + # Copy total force (column 0) - only thread with f_idx == 0 does this + if f_idx == 0: + net_force_total[env, sensor] = newton_forces[env, sensor, 0] + + # Copy per-filter-object forces (columns 1+) + # Guard with `if force_matrix:` to handle None case (no filter objects) + if force_matrix: + force_matrix[env, sensor, f_idx] = newton_forces[env, sensor, f_idx + 1] + + +@wp.kernel +def reset_contact_sensor_kernel( + # in + history_length: int, + num_filter_objects: int, + env_mask: wp.array(dtype=wp.bool), + # in-out + net_forces_w: wp.array2d(dtype=wp.vec3f), + net_forces_w_history: wp.array3d(dtype=wp.vec3f), + force_matrix_w: wp.array3d(dtype=wp.vec3f), + # outputs + current_air_time: wp.array2d(dtype=wp.float32), + last_air_time: wp.array2d(dtype=wp.float32), + current_contact_time: wp.array2d(dtype=wp.float32), + last_contact_time: wp.array2d(dtype=wp.float32), +): + """Reset the contact sensor data for specified environments. + + Launch with dim=(num_envs, num_sensors). + """ + env, sensor = wp.tid() + + if env_mask: + if not env_mask[env]: + return + + # Reset net forces + net_forces_w[env, sensor] = wp.vec3f(0.0) + + # Reset history + if net_forces_w_history: + for i in range(history_length): + net_forces_w_history[env, i, sensor] = wp.vec3f(0.0) + + # Reset force matrix (guard for None case) + if force_matrix_w: + for f in range(num_filter_objects): + force_matrix_w[env, sensor, f] = wp.vec3f(0.0) + + # Reset air/contact time tracking + if current_air_time: + current_air_time[env, sensor] = 0.0 + last_air_time[env, sensor] = 0.0 + current_contact_time[env, sensor] = 0.0 + last_contact_time[env, sensor] = 0.0 + + +@wp.kernel +def update_contact_sensor_kernel( + # in + history_length: int, + contact_force_threshold: wp.float32, + env_mask: wp.array(dtype=wp.bool), + net_forces: wp.array2d(dtype=wp.vec3f), + timestamp: wp.array(dtype=wp.float32), + timestamp_last_update: wp.array(dtype=wp.float32), + # in-out + net_forces_history: wp.array3d(dtype=wp.vec3f), + current_air_time: wp.array2d(dtype=wp.float32), + current_contact_time: wp.array2d(dtype=wp.float32), + # out + last_air_time: wp.array2d(dtype=wp.float32), + last_contact_time: wp.array2d(dtype=wp.float32), +): + """Update the contact sensor data (history and air/contact time tracking). + + Launch with dim=(num_envs, num_sensors). + """ + env, sensor = wp.tid() + + if env_mask: + if not env_mask[env]: + return + + # Update history + if net_forces_history: + for i in range(history_length - 1, 0, -1): + net_forces_history[env, i, sensor] = net_forces_history[env, i - 1, sensor] + net_forces_history[env, 0, sensor] = net_forces[env, sensor] + + # Update air/contact time tracking + if current_air_time: + elapsed_time = timestamp[env] - timestamp_last_update[env] + in_contact = wp.length_sq(net_forces[env, sensor]) > contact_force_threshold * contact_force_threshold + + cat = current_air_time[env, sensor] + cct = current_contact_time[env, sensor] + is_first_contact = in_contact and (cat > 0.0) + is_first_detached = not in_contact and (cct > 0.0) + + if is_first_contact: + last_air_time[env, sensor] = cat + elapsed_time + elif is_first_detached: + last_contact_time[env, sensor] = cct + elapsed_time + + current_contact_time[env, sensor] = wp.where(in_contact, cct + elapsed_time, 0.0) + current_air_time[env, sensor] = wp.where(in_contact, 0.0, cat + elapsed_time) + + +@wp.kernel +def compute_first_transition_kernel( + # in + threshold: wp.float32, + time: wp.array2d(dtype=wp.float32), + # out + result: wp.array2d(dtype=wp.float32), +): + """Compute boolean mask (as float) for sensors whose time is in (0, threshold). + + Used by both compute_first_contact (with current_contact_time) and + compute_first_air (with current_air_time). + + Launch with dim=(num_envs, num_sensors). + """ + env, sensor = wp.tid() + t = time[env, sensor] + if t > 0.0 and t < threshold: + result[env, sensor] = 1.0 + else: + result[env, sensor] = 0.0 + + +@wp.kernel +def reset_sensor_state_kernel( + env_mask: wp.array(dtype=wp.bool), + is_outdated: wp.array(dtype=wp.bool), + timestamp: wp.array(dtype=wp.float32), + timestamp_last_update: wp.array(dtype=wp.float32), +): + """Reset sensor state (timestamps and outdated flags). + + Launch with dim=(num_envs,). + """ + i = wp.tid() + if env_mask[i]: + is_outdated[i] = True + timestamp[i] = 0.0 + timestamp_last_update[i] = 0.0 + + +@wp.kernel +def update_timestamp_kernel( + is_outdated: wp.array(dtype=wp.bool), + timestamp: wp.array(dtype=wp.float32), + timestamp_last_update: wp.array(dtype=wp.float32), + sim_physics_dt: wp.array(dtype=wp.float32), + update_period: wp.float32, +): + """Update timestamps and mark environments as outdated if needed. + + Launch with dim=(num_envs,). + """ + i = wp.tid() + timestamp[i] = timestamp[i] + sim_physics_dt[i] + if timestamp[i] - timestamp_last_update[i] >= update_period: + is_outdated[i] = True + + +@wp.kernel +def update_outdated_envs_kernel( + is_outdated: wp.array(dtype=wp.bool), + timestamp: wp.array(dtype=wp.float32), + timestamp_last_update: wp.array(dtype=wp.float32), +): + """Update timestamp_last_update and clear outdated flag for outdated environments. + + Launch with dim=(num_envs,). + """ + i = wp.tid() + if is_outdated[i]: + timestamp_last_update[i] = timestamp[i] + is_outdated[i] = False diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/tiled_camera_cfg.py b/source/isaaclab_newton/isaaclab_newton/sensors/tiled_camera_cfg.py new file mode 100644 index 00000000000..dae1292a2a2 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sensors/tiled_camera_cfg.py @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""TiledCamera configuration that wires in the Newton warp renderer.""" + +from __future__ import annotations + +from isaaclab.sensors.camera.tiled_camera_cfg import TiledCameraCfg +from isaaclab.utils import configclass + + +@configclass +class TiledCameraCfgNewtonWarpRenderer(TiledCameraCfg): + """TiledCameraCfg variant that uses :class:`NewtonWarpRenderer`. + + The ``renderer_class`` field holds a ``"module:ClassName"`` string that is + resolved lazily by the environment's ``_setup_scene`` via + ``isaaclab.utils.string.string_to_callable``. Using a plain string (rather + than :class:`DeferredClass`) keeps the field serialisation-friendly and + prevents ``cfg.validate()`` from recursing into warp type internals. + """ + + renderer_class: str = "isaaclab.renderers.newton_warp_renderer:NewtonWarpRenderer" diff --git a/source/isaaclab_newton/setup.py b/source/isaaclab_newton/setup.py index 706f4fa8938..e2b68bfebd2 100644 --- a/source/isaaclab_newton/setup.py +++ b/source/isaaclab_newton/setup.py @@ -21,8 +21,8 @@ "prettytable==3.3.0", # newton "mujoco==3.5.0", - "mujoco-warp @ git+https://github.com/google-deepmind/mujoco_warp.git@7bb099c082d50803fd63bfe57217025d9f5cd2df", - "newton @ git+https://github.com/newton-physics/newton.git@c2c8588602951238ce07e34edaf983afccf23b5f", + "mujoco-warp==3.5.0", + "newton @ git+https://github.com/newton-physics/newton.git@51ce35e8def843377546764033edc33a0b479d65", "imgui-bundle==1.92.0", "PyOpenGL-accelerate==3.1.10", ] diff --git a/source/isaaclab_physx/isaaclab_physx/assets/__init__.py b/source/isaaclab_physx/isaaclab_physx/assets/__init__.py index 3b31699c1f7..a02b5b58792 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/__init__.py @@ -38,22 +38,16 @@ the corresponding actuator torques. """ -from .articulation import Articulation, ArticulationData -from .deformable_object import DeformableObject, DeformableObjectCfg, DeformableObjectData -from .rigid_object import RigidObject, RigidObjectData -from .rigid_object_collection import RigidObjectCollection, RigidObjectCollectionData -from .surface_gripper import SurfaceGripper, SurfaceGripperCfg - -__all__ = [ - "Articulation", - "ArticulationData", - "DeformableObject", - "DeformableObjectCfg", - "DeformableObjectData", - "RigidObject", - "RigidObjectData", - "RigidObjectCollection", - "RigidObjectCollectionData", - "SurfaceGripper", - "SurfaceGripperCfg", -] +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submodules=[], + submod_attrs={ + "articulation": ["Articulation", "ArticulationData"], + "deformable_object": ["DeformableObject", "DeformableObjectCfg", "DeformableObjectData"], + "rigid_object": ["RigidObject", "RigidObjectData"], + "rigid_object_collection": ["RigidObjectCollection", "RigidObjectCollectionData"], + "surface_gripper": ["SurfaceGripper", "SurfaceGripperCfg"], + }, +) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/__init__.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/__init__.py index 6efcec972dc..c8a01a32cb6 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/__init__.py @@ -5,10 +5,12 @@ """Sub-module for rigid articulated assets.""" -from .articulation import Articulation -from .articulation_data import ArticulationData +import lazy_loader as lazy -__all__ = [ - "Articulation", - "ArticulationData", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "articulation": ["Articulation"], + "articulation_data": ["ArticulationData"], + }, +) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 892e20d2937..ace71f9aab2 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -850,6 +850,17 @@ def write_joint_state_to_sim( self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) + def write_joint_state_to_sim_index( + self, + *, + position: torch.Tensor | wp.array, + velocity: torch.Tensor | wp.array, + joint_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + ) -> None: + self.write_joint_position_to_sim_index(position=position, joint_ids=joint_ids, env_ids=env_ids) + self.write_joint_velocity_to_sim_index(velocity=velocity, joint_ids=joint_ids, env_ids=env_ids) + def write_joint_state_to_sim_mask( self, *, diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/__init__.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/__init__.py index 34ebd7892ff..4bd285da835 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/__init__.py @@ -5,12 +5,13 @@ """Sub-module for deformable object assets.""" -from .deformable_object import DeformableObject -from .deformable_object_cfg import DeformableObjectCfg -from .deformable_object_data import DeformableObjectData +import lazy_loader as lazy -__all__ = [ - "DeformableObject", - "DeformableObjectCfg", - "DeformableObjectData", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "deformable_object": ["DeformableObject"], + "deformable_object_cfg": ["DeformableObjectCfg"], + "deformable_object_data": ["DeformableObjectData"], + }, +) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object_cfg.py b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object_cfg.py index 280577baefb..964a75a776f 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object_cfg.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/deformable_object/deformable_object_cfg.py @@ -9,15 +9,16 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.markers.config import DEFORMABLE_TARGET_MARKER_CFG from isaaclab.utils import configclass - -from .deformable_object import DeformableObject +from isaaclab.utils.string import DeferredClass @configclass class DeformableObjectCfg(AssetBaseCfg): """Configuration parameters for a deformable object.""" - class_type: type = DeformableObject + class_type: type | DeferredClass = DeferredClass( + "isaaclab_physx.assets.deformable_object.deformable_object:DeformableObject" + ) visualizer_cfg: VisualizationMarkersCfg = DEFORMABLE_TARGET_MARKER_CFG.replace( prim_path="/Visuals/DeformableTarget" diff --git a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/__init__.py b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/__init__.py index 1c7771cc5bc..e269687dc2e 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/__init__.py @@ -5,10 +5,12 @@ """Sub-module for surface_gripper assets.""" -from .surface_gripper import SurfaceGripper -from .surface_gripper_cfg import SurfaceGripperCfg +import lazy_loader as lazy -__all__ = [ - "SurfaceGripper", - "SurfaceGripperCfg", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "surface_gripper": ["SurfaceGripper"], + "surface_gripper_cfg": ["SurfaceGripperCfg"], + }, +) diff --git a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper_cfg.py b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper_cfg.py index a6e3bdb2ebb..0b82d08fc15 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper_cfg.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/surface_gripper/surface_gripper_cfg.py @@ -7,8 +7,7 @@ from isaaclab.assets.asset_base_cfg import AssetBaseCfg from isaaclab.utils import configclass - -from .surface_gripper import SurfaceGripper +from isaaclab.utils.string import DeferredClass @configclass @@ -30,4 +29,6 @@ class SurfaceGripperCfg(AssetBaseCfg): retry_interval: float | None = None """The amount of time the gripper will spend trying to grasp an object.""" - class_type: type = SurfaceGripper + class_type: type | DeferredClass = DeferredClass( + "isaaclab_physx.assets.surface_gripper.surface_gripper:SurfaceGripper" + ) diff --git a/source/isaaclab_physx/isaaclab_physx/cloner/__init__.py b/source/isaaclab_physx/isaaclab_physx/cloner/__init__.py index df60869211e..2dcf79b8f3b 100644 --- a/source/isaaclab_physx/isaaclab_physx/cloner/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/cloner/__init__.py @@ -3,6 +3,11 @@ # # SPDX-License-Identifier: BSD-3-Clause -from .physx_replicate import physx_replicate +import lazy_loader as lazy -__all__ = ["physx_replicate"] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "physx_replicate": ["physx_replicate"], + }, +) diff --git a/source/isaaclab_physx/isaaclab_physx/physics/__init__.py b/source/isaaclab_physx/isaaclab_physx/physics/__init__.py index cc84bb7e38a..c783ca0da8a 100644 --- a/source/isaaclab_physx/isaaclab_physx/physics/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/physics/__init__.py @@ -5,12 +5,14 @@ """Implementation backends for simulation interfaces.""" -from .physx_manager import PhysxManager, IsaacEvents -from .physx_manager_cfg import PhysxCfg +import lazy_loader as lazy +from .physx_manager_cfg import PhysxCfg -__all__ = [ - "PhysxManager", - "IsaacEvents", - "PhysxCfg", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "physx_manager": ["PhysxManager", "IsaacEvents"], + }, +) +__all__ += ["PhysxCfg"] diff --git a/source/isaaclab_physx/isaaclab_physx/physics/physx_manager.py b/source/isaaclab_physx/isaaclab_physx/physics/physx_manager.py index d03b0abda0b..e797e5a21e7 100644 --- a/source/isaaclab_physx/isaaclab_physx/physics/physx_manager.py +++ b/source/isaaclab_physx/isaaclab_physx/physics/physx_manager.py @@ -67,6 +67,10 @@ class IsaacEvents(Enum): PhysicsEvent.STOP: IsaacEvents.TIMELINE_STOP, } +_PHYSICS_EVENT_VALUE_TO_ISAAC_EVENT: dict[str, IsaacEvents] = { + e.value: ie for e, ie in _PHYSICS_EVENT_TO_ISAAC_EVENT.items() +} + class AnimationRecorder: """Handles animation recording using PhysX PVD interface.""" @@ -371,6 +375,7 @@ def _subscribe_to_event( ) -> Any: """Subscribe to PhysX events. Maps PhysicsEvent → IsaacEvents.""" isaac_event = _PHYSICS_EVENT_TO_ISAAC_EVENT.get(event) + isaac_event = _PHYSICS_EVENT_VALUE_TO_ISAAC_EVENT.get(event.value) return cls._subscribe_isaac(callback, isaac_event, order, name) if isaac_event else None @classmethod diff --git a/source/isaaclab_physx/isaaclab_physx/physics/physx_manager_cfg.py b/source/isaaclab_physx/isaaclab_physx/physics/physx_manager_cfg.py index d30252c7877..9e3941512b3 100644 --- a/source/isaaclab_physx/isaaclab_physx/physics/physx_manager_cfg.py +++ b/source/isaaclab_physx/isaaclab_physx/physics/physx_manager_cfg.py @@ -10,13 +10,13 @@ from typing import TYPE_CHECKING, Literal from isaaclab.physics import PhysicsCfg -from isaaclab.utils import configclass - -from .physx_manager import PhysxManager +from isaaclab.utils import DeferredClass, configclass if TYPE_CHECKING: from isaaclab.physics import PhysicsManager + from .physx_manager import PhysxManager + @configclass class PhysxCfg(PhysicsCfg): @@ -41,7 +41,7 @@ class PhysxCfg(PhysicsCfg): # PhysX Scene Settings # ------------------------------------------------------------------ - class_type: type[PhysicsManager] = PhysxManager + class_type: type[PhysicsManager] | DeferredClass = DeferredClass("isaaclab_physx.physics.physx_manager:PhysxManager") """The class type of the PhysxManager.""" # ------------------------------------------------------------------ diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/__init__.py b/source/isaaclab_physx/isaaclab_physx/sensors/__init__.py index dde43d8ce3f..6d3198e068b 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/__init__.py @@ -5,15 +5,14 @@ """Sub-package containing PhysX-specific sensor implementations.""" -from .contact_sensor import ContactSensor, ContactSensorData -from .frame_transformer import FrameTransformer, FrameTransformerData -from .imu import Imu, ImuData +import lazy_loader as lazy -__all__ = [ - "ContactSensor", - "ContactSensorData", - "FrameTransformer", - "FrameTransformerData", - "Imu", - "ImuData", -] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "contact_sensor": ["ContactSensor", "ContactSensorData"], + "contact_sensor.contact_sensor_cfg": ["ContactSensorCfg"], + "frame_transformer": ["FrameTransformer", "FrameTransformerData"], + "imu": ["Imu", "ImuData"], + }, +) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/__init__.py b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/__init__.py index e2f5f2b7772..8aa2002bb90 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/__init__.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/__init__.py @@ -5,7 +5,13 @@ """Sub-module for PhysX rigid contact sensor.""" -from .contact_sensor import ContactSensor -from .contact_sensor_data import ContactSensorData +import lazy_loader as lazy -__all__ = ["ContactSensor", "ContactSensorData"] +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "contact_sensor": ["ContactSensor"], + "contact_sensor_cfg": ["ContactSensorCfg"], + "contact_sensor_data": ["ContactSensorData"], + }, +) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py index 4d9ab046017..49d44c80189 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor.py @@ -159,9 +159,7 @@ def contact_view(self) -> physx.RigidContactView: """ def reset(self, env_ids: Sequence[int] | None = None): - # reset the timers and counters super().reset(env_ids) - # resolve env_ids to warp array if env_ids is None: env_ids_wp = self._ALL_ENV_INDICES num_env_ids = self._num_envs @@ -188,7 +186,7 @@ def reset(self, env_ids: Sequence[int] | None = None): ) # reset force matrix - if len(self.cfg.filter_prim_paths_expr) != 0: + if self.cfg.filter_prim_paths_expr: M = self._num_filter_shapes wp.launch( reset_vec3f_3d, @@ -280,7 +278,7 @@ def _initialize_impl(self): body_names_regex = f"{self.cfg.prim_path.rsplit('/', 1)[0]}/{body_names_regex}" # convert regex expressions to glob expressions for PhysX body_names_glob = body_names_regex.replace(".*", "*") - filter_prim_paths_glob = [expr.replace(".*", "*") for expr in self.cfg.filter_prim_paths_expr] + filter_prim_paths_glob = [expr.replace(".*", "*") for expr in (self.cfg.filter_prim_paths_expr or [])] # create a rigid prim view for the sensor self._body_physx_view = self._physics_sim_view.create_rigid_body_view(body_names_glob) @@ -301,7 +299,7 @@ def _initialize_impl(self): # check if filter paths are valid if self.cfg.track_contact_points or self.cfg.track_friction_forces: - if len(self.cfg.filter_prim_paths_expr) == 0: + if not self.cfg.filter_prim_paths_expr: raise ValueError( "The 'filter_prim_paths_expr' is empty. Please specify a valid filter pattern to track" f" {'contact points' if self.cfg.track_contact_points else 'friction forces'}." @@ -314,7 +312,7 @@ def _initialize_impl(self): ) # Store filter shapes count - self._num_filter_shapes = self.contact_view.filter_count if len(self.cfg.filter_prim_paths_expr) != 0 else 0 + self._num_filter_shapes = self.contact_view.filter_count if self.cfg.filter_prim_paths_expr else 0 # Store effective history length (always >= 1 for consistent buffer shapes) self._history_length = max(self.cfg.history_length, 1) @@ -366,7 +364,7 @@ def _update_buffers_impl(self, env_ids: Sequence[int]): ) # -- Force matrix -- - if len(self.cfg.filter_prim_paths_expr) != 0: + if self.cfg.filter_prim_paths_expr: M = self._num_filter_shapes # PhysX returns (N*B, M, 3) float32 -> (N*B, M) vec3f force_matrix_flat = self.contact_view.get_contact_force_matrix(dt=self._sim_physics_dt).view(wp.vec3f) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor_cfg.py b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor_cfg.py new file mode 100644 index 00000000000..be91f5c064c --- /dev/null +++ b/source/isaaclab_physx/isaaclab_physx/sensors/contact_sensor/contact_sensor_cfg.py @@ -0,0 +1,41 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.sensors.contact_sensor import ContactSensorCfg as BaseContactSensorCfg +from isaaclab.utils import DeferredClass, configclass + + +@configclass +class ContactSensorCfg(BaseContactSensorCfg): + """Configuration for the PhysX contact sensor. + + Extends :class:`isaaclab.sensors.ContactSensorCfg` with PhysX-specific fields. + """ + + class_type: type | DeferredClass = DeferredClass("isaaclab_physx.sensors.contact_sensor.contact_sensor:ContactSensor") + + track_contact_points: bool = False + """Whether to track the contact point locations. Defaults to False. + + .. note:: + Requires :attr:`filter_prim_paths_expr` to be set. + """ + + track_friction_forces: bool = False + """Whether to track the friction forces at the contact points. Defaults to False. + + .. note:: + Requires :attr:`filter_prim_paths_expr` to be set. + """ + + max_contact_data_count_per_prim: int = 4 + """The maximum number of contacts per primitive to track. Default is 4. + + The total number of contacts allowed is ``max_contact_data_count_per_prim * num_envs * num_sensor_bodies``. + + .. note:: + If the environment is contact-rich, increase this parameter to avoid out of bounds memory + errors and loss of contact data leading to inaccurate measurements. + """ diff --git a/source/isaaclab_physx/test/sensors/check_contact_sensor.py b/source/isaaclab_physx/test/sensors/check_contact_sensor.py index 035518fdf4d..47bed15243d 100644 --- a/source/isaaclab_physx/test/sensors/check_contact_sensor.py +++ b/source/isaaclab_physx/test/sensors/check_contact_sensor.py @@ -149,7 +149,7 @@ def main(): # reset dof state joint_pos, joint_vel = robot.data.default_joint_pos, robot.data.default_joint_vel - robot.write_joint_state_to_sim(joint_pos, joint_vel) + robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) robot.reset() dt = [] diff --git a/source/isaaclab_physx/test/sensors/test_frame_transformer.py b/source/isaaclab_physx/test/sensors/test_frame_transformer.py index 150e0c47faf..4f2badab40f 100644 --- a/source/isaaclab_physx/test/sensors/test_frame_transformer.py +++ b/source/isaaclab_physx/test/sensors/test_frame_transformer.py @@ -162,7 +162,7 @@ def test_frame_transformer_feet_wrt_base(sim): # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:]) - scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene.articulations["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset buffers scene.reset() @@ -260,7 +260,7 @@ def test_frame_transformer_feet_wrt_thigh(sim): # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:]) - scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene.articulations["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset buffers scene.reset() @@ -338,7 +338,7 @@ def test_frame_transformer_robot_body_to_external_cube(sim): # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:]) - scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene.articulations["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset buffers scene.reset() @@ -522,7 +522,7 @@ def test_frame_transformer_all_bodies(sim): # -- robot scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:]) - scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel) + scene.articulations["robot"].write_joint_state_to_sim(position=joint_pos, velocity=joint_vel) # reset buffers scene.reset() @@ -714,8 +714,8 @@ class MultiRobotSceneCfg(InteractiveSceneCfg): scene.articulations["robot"].write_root_pose_to_sim(root_state[:, :7]) scene.articulations["robot"].write_root_velocity_to_sim(root_state[:, 7:]) scene.articulations["robot"].write_joint_state_to_sim( - scene.articulations["robot"].data.default_joint_pos, - scene.articulations["robot"].data.default_joint_vel, + position=scene.articulations["robot"].data.default_joint_pos, + velocity=scene.articulations["robot"].data.default_joint_vel, ) # Reset robot_1 root_state_1 = wp.to_torch(scene.articulations["robot_1"].data.default_root_state).clone() @@ -723,8 +723,8 @@ class MultiRobotSceneCfg(InteractiveSceneCfg): scene.articulations["robot_1"].write_root_pose_to_sim(root_state_1[:, :7]) scene.articulations["robot_1"].write_root_velocity_to_sim(root_state_1[:, 7:]) scene.articulations["robot_1"].write_joint_state_to_sim( - scene.articulations["robot_1"].data.default_joint_pos, - scene.articulations["robot_1"].data.default_joint_vel, + position=scene.articulations["robot_1"].data.default_joint_pos, + velocity=scene.articulations["robot_1"].data.default_joint_vel, ) scene.reset() diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.py b/source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.py index 3d8c6296899..4bdd3d4ece2 100644 --- a/source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.py +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/__init__.py @@ -15,9 +15,27 @@ """ -from .distillation_cfg import * -from .exporter import export_policy_as_jit, export_policy_as_onnx -from .rl_cfg import * -from .rnd_cfg import RslRlRndCfg -from .symmetry_cfg import RslRlSymmetryCfg -from .vecenv_wrapper import RslRlVecEnvWrapper +import lazy_loader as lazy + +__getattr__, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "distillation_cfg": [ + "RslRlDistillationRunnerCfg", + "RslRlDistillationAlgorithmCfg", + "RslRlDistillationStudentTeacherCfg", + "RslRlDistillationStudentTeacherRecurrentCfg", + ], + "exporter": ["export_policy_as_jit", "export_policy_as_onnx"], + "rl_cfg": [ + "RslRlBaseRunnerCfg", + "RslRlOnPolicyRunnerCfg", + "RslRlPpoActorCriticCfg", + "RslRlPpoActorCriticRecurrentCfg", + "RslRlPpoAlgorithmCfg", + ], + "rnd_cfg": ["RslRlRndCfg"], + "symmetry_cfg": ["RslRlSymmetryCfg"], + "vecenv_wrapper": ["RslRlVecEnvWrapper"], + }, +) diff --git a/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py b/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py index dde20f2bb16..2af29e92b77 100644 --- a/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py +++ b/source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py @@ -3,12 +3,17 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + +from typing import TYPE_CHECKING + import gymnasium as gym import torch from rsl_rl.env import VecEnv from tensordict import TensorDict -from isaaclab.envs import DirectRLEnv, ManagerBasedRLEnv +if TYPE_CHECKING: + from isaaclab.envs import DirectRLEnv, ManagerBasedRLEnv class RslRlVecEnvWrapper(VecEnv): @@ -38,7 +43,9 @@ def __init__(self, env: ManagerBasedRLEnv | DirectRLEnv, clip_actions: float | N """ # check that input is valid - if not isinstance(env.unwrapped, ManagerBasedRLEnv) and not isinstance(env.unwrapped, DirectRLEnv): + _valid_bases = {"ManagerBasedRLEnv", "DirectRLEnv", "ManagerBasedEnv", "DirectMARLEnv"} + _mro_names = {c.__name__ for c in type(env.unwrapped).__mro__} + if not _valid_bases & _mro_names: raise ValueError( "The environment must be inherited from ManagerBasedRLEnv or DirectRLEnv. Environment type:" f" {type(env)}" diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/ant/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/__init__.py index 9881cd66ca7..bd8e1bba2b5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/ant/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/__init__.py @@ -20,7 +20,7 @@ entry_point=f"{__name__}.ant_env:AntEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.ant_env:AntEnvCfg", + "env_cfg_entry_point": f"{__name__}.ant_env_cfg:AntEnvCfg", "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:AntPPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env.py index 39ae57b2967..1facc07440d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env.py @@ -5,67 +5,9 @@ from __future__ import annotations -import isaaclab.sim as sim_utils -from isaaclab.assets import ArticulationCfg -from isaaclab.envs import DirectRLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg -from isaaclab.terrains import TerrainImporterCfg -from isaaclab.utils import configclass - from isaaclab_tasks.direct.locomotion.locomotion_env import LocomotionEnv -from isaaclab_assets.robots.ant import ANT_CFG - - -@configclass -class AntEnvCfg(DirectRLEnvCfg): - # env - episode_length_s = 15.0 - decimation = 2 - action_scale = 0.5 - action_space = 8 - observation_space = 36 - state_space = 0 - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) - terrain = TerrainImporterCfg( - prim_path="/World/ground", - terrain_type="plane", - collision_group=-1, - physics_material=sim_utils.RigidBodyMaterialCfg( - friction_combine_mode="average", - restitution_combine_mode="average", - static_friction=1.0, - dynamic_friction=1.0, - restitution=0.0, - ), - debug_vis=False, - ) - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True - ) - - # robot - robot: ArticulationCfg = ANT_CFG.replace(prim_path="/World/envs/env_.*/Robot") - joint_gears: list = [15, 15, 15, 15, 15, 15, 15, 15] - - heading_weight: float = 0.5 - up_weight: float = 0.1 - - energy_cost_scale: float = 0.05 - actions_cost_scale: float = 0.005 - alive_reward_scale: float = 0.5 - dof_vel_scale: float = 0.2 - - death_cost: float = -2.0 - termination_height: float = 0.31 - - angular_velocity_scale: float = 1.0 - contact_force_scale: float = 0.1 +from .ant_env_cfg import AntEnvCfg class AntEnv(LocomotionEnv): diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env_cfg.py new file mode 100644 index 00000000000..1078e817755 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/ant/ant_env_cfg.py @@ -0,0 +1,66 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg +from isaaclab.envs import DirectRLEnvCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.terrains import TerrainImporterCfg +from isaaclab.utils import configclass + +from isaaclab_assets.robots.ant import ANT_CFG + + +@configclass +class AntEnvCfg(DirectRLEnvCfg): + # env + episode_length_s = 15.0 + decimation = 2 + action_scale = 0.5 + action_space = 8 + observation_space = 36 + state_space = 0 + + # simulation + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="plane", + collision_group=-1, + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="average", + restitution_combine_mode="average", + static_friction=1.0, + dynamic_friction=1.0, + restitution=0.0, + ), + debug_vis=False, + ) + + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True + ) + + # robot + robot: ArticulationCfg = ANT_CFG.replace(prim_path="/World/envs/env_.*/Robot") + joint_gears: list = [15, 15, 15, 15, 15, 15, 15, 15] + + heading_weight: float = 0.5 + up_weight: float = 0.1 + + energy_cost_scale: float = 0.05 + actions_cost_scale: float = 0.005 + alive_reward_scale: float = 0.5 + dof_vel_scale: float = 0.2 + + death_cost: float = -2.0 + termination_height: float = 0.31 + + angular_velocity_scale: float = 1.0 + contact_force_scale: float = 0.1 diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py index 838d4110c35..61fa3c11246 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py @@ -190,7 +190,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): default_root_state[:, :3] += self._terrain.env_origins[env_ids] self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) - self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self._robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) # Logging extras = dict() for key in self._episode_sums.keys(): diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/__init__.py index 9da64598c0c..ef9160ac3e0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/__init__.py @@ -16,7 +16,7 @@ entry_point=f"{__name__}.assembly_env:AssemblyEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.assembly_env:AssemblyEnvCfg", + "env_cfg_entry_point": f"{__name__}.assembly_env_cfg:AssemblyEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", }, ) @@ -27,7 +27,7 @@ entry_point=f"{__name__}.disassembly_env:DisassemblyEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.disassembly_env:DisassemblyEnvCfg", + "env_cfg_entry_point": f"{__name__}.disassembly_env_cfg:DisassemblyEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", }, ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py index c6eb6f1ab56..e38f32c4678 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py @@ -10,8 +10,6 @@ import torch import warp as wp -import carb - import isaaclab.sim as sim_utils from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import DirectRLEnv @@ -723,7 +721,7 @@ def set_pos_inverse_kinematics(self, env_ids): self.ctrl_target_joint_pos[env_ids, 0:7] = self.joint_pos[env_ids, 0:7] # Update dof state. - self._robot.write_joint_state_to_sim(self.joint_pos, self.joint_vel) + self._robot.write_joint_state_to_sim(position=self.joint_pos, velocity=self.joint_vel) self._robot.reset() self._robot.set_joint_position_target(self.ctrl_target_joint_pos) @@ -743,7 +741,7 @@ def _set_franka_to_default_pose(self, joints, env_ids): joint_effort = torch.zeros_like(joint_pos) self.ctrl_target_joint_pos[env_ids, :] = joint_pos self._robot.set_joint_position_target(self.ctrl_target_joint_pos[env_ids], env_ids=env_ids) - self._robot.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + self._robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) self._robot.reset() self._robot.set_joint_effort_target(joint_effort, env_ids=env_ids) @@ -824,6 +822,8 @@ def randomize_held_initial_state(self, env_ids, pre_grasp): def randomize_initial_state(self, env_ids): """Randomize initial state and perform any episode-level randomization.""" + import carb + # Disable gravity. physics_sim_view = sim_utils.SimulationContext.instance().physics_sim_view physics_sim_view.set_gravity(carb.Float3(0.0, 0.0, 0.0)) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py index 2ad750001f1..0bb8a095b59 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/automate/disassembly_env.py @@ -10,8 +10,6 @@ import torch import warp as wp -import carb - import isaaclab.sim as sim_utils from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import DirectRLEnv @@ -549,7 +547,7 @@ def set_pos_inverse_kinematics(self, env_ids): self.ctrl_target_joint_pos[env_ids, 0:7] = self.joint_pos[env_ids, 0:7] # Update dof state. - self._robot.write_joint_state_to_sim(self.joint_pos, self.joint_vel) + self._robot.write_joint_state_to_sim(position=self.joint_pos, velocity=self.joint_vel) self._robot.reset() self._robot.set_joint_position_target(self.ctrl_target_joint_pos) @@ -611,7 +609,7 @@ def _set_franka_to_default_pose(self, joints, env_ids): joint_effort = torch.zeros_like(joint_pos) self.ctrl_target_joint_pos[env_ids, :] = joint_pos self._robot.set_joint_position_target(self.ctrl_target_joint_pos[env_ids], env_ids=env_ids) - self._robot.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + self._robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) self._robot.reset() self._robot.set_joint_effort_target(joint_effort, env_ids=env_ids) @@ -692,6 +690,8 @@ def close_gripper(self, env_ids): def randomize_initial_state(self, env_ids): """Randomize initial state and perform any episode-level randomization.""" + import carb + # Disable gravity. physics_sim_view = sim_utils.SimulationContext.instance().physics_sim_view physics_sim_view.set_gravity(carb.Float3(0.0, 0.0, 0.0)) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/__init__.py index b4913fb2c3a..cfdb8af38f4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/__init__.py @@ -20,7 +20,7 @@ entry_point=f"{__name__}.cart_double_pendulum_env:CartDoublePendulumEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.cart_double_pendulum_env:CartDoublePendulumEnvCfg", + "env_cfg_entry_point": f"{__name__}.cart_double_pendulum_env_cfg:CartDoublePendulumEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", "skrl_ippo_cfg_entry_point": f"{agents.__name__}:skrl_ippo_cfg.yaml", diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py index f9d372ef98c..a4e3bd02131 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env.py @@ -12,57 +12,12 @@ import warp as wp import isaaclab.sim as sim_utils -from isaaclab.assets import Articulation, ArticulationCfg -from isaaclab.envs import DirectMARLEnv, DirectMARLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg +from isaaclab.assets import Articulation +from isaaclab.envs import DirectMARLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane -from isaaclab.utils import configclass from isaaclab.utils.math import sample_uniform -from isaaclab_assets.robots.cart_double_pendulum import CART_DOUBLE_PENDULUM_CFG - - -@configclass -class CartDoublePendulumEnvCfg(DirectMARLEnvCfg): - # env - decimation = 2 - episode_length_s = 5.0 - possible_agents = ["cart", "pendulum"] - action_spaces = {"cart": 1, "pendulum": 1} - observation_spaces = {"cart": 4, "pendulum": 3} - state_space = -1 - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) - - # robot - robot_cfg: ArticulationCfg = CART_DOUBLE_PENDULUM_CFG.replace(prim_path="/World/envs/env_.*/Robot") - cart_dof_name = "slider_to_cart" - pole_dof_name = "cart_to_pole" - pendulum_dof_name = "pole_to_pendulum" - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=4096, env_spacing=4.0, replicate_physics=True) - - # reset - max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] - initial_pole_angle_range = [-0.25, 0.25] # the range in which the pole angle is sampled from on reset [rad] - initial_pendulum_angle_range = [-0.25, 0.25] # the range in which the pendulum angle is sampled from on reset [rad] - - # action scales - cart_action_scale = 100.0 # [N] - pendulum_action_scale = 50.0 # [Nm] - - # reward scales - rew_scale_alive = 1.0 - rew_scale_terminated = -2.0 - rew_scale_cart_pos = 0 - rew_scale_cart_vel = -0.01 - rew_scale_pole_pos = -1.0 - rew_scale_pole_vel = -0.01 - rew_scale_pendulum_pos = -1.0 - rew_scale_pendulum_vel = -0.01 +from .cart_double_pendulum_env_cfg import CartDoublePendulumEnvCfg class CartDoublePendulumEnv(DirectMARLEnv): @@ -188,7 +143,7 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) - self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self.robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) @torch.jit.script diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env_cfg.py new file mode 100644 index 00000000000..048750158cd --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cart_double_pendulum/cart_double_pendulum_env_cfg.py @@ -0,0 +1,56 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.assets import ArticulationCfg +from isaaclab.envs import DirectMARLEnvCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.utils import configclass + +from isaaclab_assets.robots.cart_double_pendulum import CART_DOUBLE_PENDULUM_CFG + + +@configclass +class CartDoublePendulumEnvCfg(DirectMARLEnvCfg): + # env + decimation = 2 + episode_length_s = 5.0 + possible_agents = ["cart", "pendulum"] + action_spaces = {"cart": 1, "pendulum": 1} + observation_spaces = {"cart": 4, "pendulum": 3} + state_space = -1 + + # simulation + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) + + # robot + robot_cfg: ArticulationCfg = CART_DOUBLE_PENDULUM_CFG.replace(prim_path="/World/envs/env_.*/Robot") + cart_dof_name = "slider_to_cart" + pole_dof_name = "cart_to_pole" + pendulum_dof_name = "pole_to_pendulum" + + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=4096, env_spacing=4.0, replicate_physics=True) + + # reset + max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] + initial_pole_angle_range = [-0.25, 0.25] # the range in which the pole angle is sampled from on reset [rad] + initial_pendulum_angle_range = [-0.25, 0.25] # the range in which the pendulum angle is sampled from on reset [rad] + + # action scales + cart_action_scale = 100.0 # [N] + pendulum_action_scale = 50.0 # [Nm] + + # reward scales + rew_scale_alive = 1.0 + rew_scale_terminated = -2.0 + rew_scale_cart_pos = 0 + rew_scale_cart_vel = -0.01 + rew_scale_pole_pos = -1.0 + rew_scale_pole_vel = -0.01 + rew_scale_pendulum_pos = -1.0 + rew_scale_pendulum_vel = -0.01 diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/__init__.py index dc8cab42492..a560af74206 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/__init__.py @@ -20,7 +20,7 @@ entry_point=f"{__name__}.cartpole_env:CartpoleEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_env:CartpoleEnvCfg", + "env_cfg_entry_point": f"{__name__}.cartpole_env_cfg:CartpoleEnvCfg", "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:CartpolePPORunnerCfg", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml", @@ -33,7 +33,7 @@ entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_camera_env:CartpoleRGBCameraEnvCfg", + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleRGBCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_camera_ppo_cfg.yaml", }, @@ -44,7 +44,7 @@ entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_camera_env:CartpoleAlbedoCameraEnvCfg", + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleAlbedoCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_camera_ppo_cfg.yaml", }, @@ -55,7 +55,7 @@ entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_camera_env:CartpoleSimpleShadingConstantCameraEnvCfg", + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleSimpleShadingConstantCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_camera_ppo_cfg.yaml", }, @@ -66,7 +66,7 @@ entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_camera_env:CartpoleSimpleShadingDiffuseCameraEnvCfg", + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleSimpleShadingDiffuseCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_camera_ppo_cfg.yaml", }, @@ -77,7 +77,7 @@ entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_camera_env:CartpoleSimpleShadingFullCameraEnvCfg", + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleSimpleShadingFullCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_camera_ppo_cfg.yaml", }, @@ -88,7 +88,7 @@ entry_point=f"{__name__}.cartpole_camera_env:CartpoleCameraEnv", disable_env_checker=True, kwargs={ - "env_cfg_entry_point": f"{__name__}.cartpole_camera_env:CartpoleDepthCameraEnvCfg", + "env_cfg_entry_point": f"{__name__}.cartpole_camera_env_cfg:CartpoleDepthCameraEnvCfg", "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml", "skrl_cfg_entry_point": f"{agents.__name__}:skrl_camera_ppo_cfg.yaml", }, diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py index 119bdad47d0..7b288afc2b9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env.py @@ -12,15 +12,19 @@ import warp as wp import isaaclab.sim as sim_utils -from isaaclab.assets import Articulation, ArticulationCfg -from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg, ViewerCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sensors import TiledCamera, TiledCameraCfg, save_images_to_file -from isaaclab.sim import SimulationCfg -from isaaclab.utils import configclass +from isaaclab.assets import Articulation +from isaaclab.envs import DirectRLEnv +from isaaclab.sensors import TiledCamera, save_images_to_file from isaaclab.utils.math import sample_uniform -from isaaclab_assets.robots.cartpole import CARTPOLE_CFG +from .cartpole_camera_env_cfg import ( + CartpoleAlbedoCameraEnvCfg, + CartpoleDepthCameraEnvCfg, + CartpoleRGBCameraEnvCfg, + CartpoleSimpleShadingConstantCameraEnvCfg, + CartpoleSimpleShadingDiffuseCameraEnvCfg, + CartpoleSimpleShadingFullCameraEnvCfg, +) SIMPLE_SHADING_TYPES = { "simple_shading_constant_diffuse", @@ -29,147 +33,6 @@ } -@configclass -class CartpoleRGBCameraEnvCfg(DirectRLEnvCfg): - # env - decimation = 2 - episode_length_s = 5.0 - action_scale = 100.0 # [N] - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) - - # robot - robot_cfg: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="/World/envs/env_.*/Robot") - cart_dof_name = "slider_to_cart" - pole_dof_name = "cart_to_pole" - - # camera - tiled_camera: TiledCameraCfg = TiledCameraCfg( - prim_path="/World/envs/env_.*/Camera", - offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(0.0, 0.0, 0.0, 1.0), convention="world"), - data_types=["rgb"], - spawn=sim_utils.PinholeCameraCfg( - focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) - ), - width=100, - height=100, - ) - write_image_to_file = False - - # spaces - action_space = 1 - state_space = 0 - observation_space = [tiled_camera.height, tiled_camera.width, 3] - - # change viewer settings - viewer = ViewerCfg(eye=(20.0, 20.0, 20.0)) - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=512, env_spacing=20.0, replicate_physics=True) - - # reset - max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] - initial_pole_angle_range = [-0.125, 0.125] # the range in which the pole angle is sampled from on reset [rad] - - # reward scales - rew_scale_alive = 1.0 - rew_scale_terminated = -2.0 - rew_scale_pole_pos = -1.0 - rew_scale_cart_vel = -0.01 - rew_scale_pole_vel = -0.005 - - -@configclass -class CartpoleDepthCameraEnvCfg(CartpoleRGBCameraEnvCfg): - # camera - tiled_camera: TiledCameraCfg = TiledCameraCfg( - prim_path="/World/envs/env_.*/Camera", - offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(0.0, 0.0, 0.0, 1.0), convention="world"), - data_types=["depth"], - spawn=sim_utils.PinholeCameraCfg( - focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) - ), - width=100, - height=100, - ) - - # spaces - observation_space = [tiled_camera.height, tiled_camera.width, 1] - - -@configclass -class CartpoleAlbedoCameraEnvCfg(CartpoleRGBCameraEnvCfg): - # camera - tiled_camera: TiledCameraCfg = TiledCameraCfg( - prim_path="/World/envs/env_.*/Camera", - offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), - data_types=["albedo"], - spawn=sim_utils.PinholeCameraCfg( - focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) - ), - width=100, - height=100, - ) - - # spaces - observation_space = [tiled_camera.height, tiled_camera.width, 3] - - -@configclass -class CartpoleSimpleShadingConstantCameraEnvCfg(CartpoleRGBCameraEnvCfg): - # camera - tiled_camera: TiledCameraCfg = TiledCameraCfg( - prim_path="/World/envs/env_.*/Camera", - offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), - data_types=["simple_shading_constant_diffuse"], - spawn=sim_utils.PinholeCameraCfg( - focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) - ), - width=100, - height=100, - ) - - # spaces - observation_space = [tiled_camera.height, tiled_camera.width, 3] - - -@configclass -class CartpoleSimpleShadingDiffuseCameraEnvCfg(CartpoleRGBCameraEnvCfg): - # camera - tiled_camera: TiledCameraCfg = TiledCameraCfg( - prim_path="/World/envs/env_.*/Camera", - offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), - data_types=["simple_shading_diffuse_mdl"], - spawn=sim_utils.PinholeCameraCfg( - focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) - ), - width=100, - height=100, - ) - - # spaces - observation_space = [tiled_camera.height, tiled_camera.width, 3] - - -@configclass -class CartpoleSimpleShadingFullCameraEnvCfg(CartpoleRGBCameraEnvCfg): - # camera - tiled_camera: TiledCameraCfg = TiledCameraCfg( - prim_path="/World/envs/env_.*/Camera", - offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), - data_types=["simple_shading_full_mdl"], - spawn=sim_utils.PinholeCameraCfg( - focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) - ), - width=100, - height=100, - ) - - # spaces - observation_space = [tiled_camera.height, tiled_camera.width, 3] - - class CartpoleCameraEnv(DirectRLEnv): """Cartpole Camera Environment.""" @@ -207,9 +70,12 @@ def close(self): def _setup_scene(self): """Setup the scene with the cartpole and camera.""" self._cartpole = Articulation(self.cfg.robot_cfg) + # Create camera without renderer first; renderer is attached after the + # scene is cloned so that the Newton model has been built by the time + # NewtonWarpRenderer.__init__ calls initialize_scene_data_provider(). self._tiled_camera = TiledCamera(self.cfg.tiled_camera) - # clone and replicate + # clone and replicate (Newton model is built during this step) self.scene.clone_environments(copy_from_source=False) if self.device == "cpu": # we need to explicitly filter collisions for CPU simulation @@ -222,6 +88,19 @@ def _setup_scene(self): light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75)) light_cfg.func("/World/Light", light_cfg) + # Now that the scene is fully built, create the renderer and attach it. + # TiledCamera._initialize_impl() checks self.renderer when it first runs, + # so the renderer must be set before the first scene.update() call. + renderer_class_path = getattr(self.cfg.tiled_camera, "renderer_class", None) + if renderer_class_path is not None: + from isaaclab.utils.string import string_to_callable + + renderer = string_to_callable(renderer_class_path)() + self._tiled_camera.renderer = renderer + else: + renderer = None + self._renderer = renderer + def _pre_physics_step(self, actions: torch.Tensor) -> None: self.actions = self.action_scale * actions.clone() @@ -229,6 +108,11 @@ def _apply_action(self) -> None: self._cartpole.set_joint_effort_target(self.actions, joint_ids=self._cart_dof_idx) def _get_observations(self) -> dict: + # Update Newton renderer transforms after scene.update() so the camera + # positions reflect the latest physics state before rendering. + if self._renderer is not None: + self._renderer.update_transforms() + data_type = self.cfg.tiled_camera.data_types[0] if "rgb" in self.cfg.tiled_camera.data_types: camera_data = self._tiled_camera.data.output[data_type] / 255.0 @@ -251,7 +135,7 @@ def _get_observations(self) -> dict: observations = {"policy": camera_data.clone()} if self.cfg.write_image_to_file: - save_images_to_file(observations["policy"], f"cartpole_{data_type}.png") + save_images_to_file(self._tiled_camera.data.output[data_type] / 255.0, f"cartpole_{data_type}.png") return observations @@ -301,7 +185,7 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self._cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) self._cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) - self._cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self._cartpole.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) @torch.jit.script diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env_cfg.py new file mode 100644 index 00000000000..a5b2a2e9a56 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_camera_env_cfg.py @@ -0,0 +1,193 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg +from isaaclab.envs import DirectRLEnvCfg, ViewerCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors import TiledCameraCfg +from isaaclab.sim import SimulationCfg +from isaaclab.utils import configclass +from isaaclab_newton.physics import NewtonCfg +from isaaclab_newton.sensors import TiledCameraCfgNewtonWarpRenderer as NewtonTiledCameraCfg +from isaaclab_physx.physics import PhysxCfg + +from isaaclab_assets.robots.cartpole import CARTPOLE_CFG + + +@configclass +class PhysicsCfg: + presets = { + "default": PhysxCfg(), + "newton": NewtonCfg(), + } + + +@configclass +class CameraCfg: + """Preset selector for the tiled camera. + + - ``default`` — standard Isaac Sim RTX :class:`TiledCamera` + - ``newton`` — :class:`TiledCamera` backed by :class:`NewtonWarpRenderer` + + Select at runtime with ``env.tiled_camera=newton``. + """ + + presets = { + "default": TiledCameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(0.0, 0.0, 0.0, 1.0), convention="world"), + data_types=["rgb"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ), + width=100, + height=100, + ), + "newton": NewtonTiledCameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(0.0, 0.0, 0.0, 1.0), convention="world"), + data_types=["rgb"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ), + width=100, + height=100, + ), + } + + +@configclass +class CartpoleRGBCameraEnvCfg(DirectRLEnvCfg): + # env + decimation = 2 + episode_length_s = 5.0 + action_scale = 100.0 # [N] + + # simulation + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation, physics=PhysicsCfg()) + + # robot + robot_cfg: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="/World/envs/env_.*/Robot") + cart_dof_name = "slider_to_cart" + pole_dof_name = "cart_to_pole" + + # camera — resolved by preset system; defaults to TiledCameraCfg (RTX pipeline) + tiled_camera: CameraCfg = CameraCfg() + write_image_to_file = False + + # spaces — dimensions fixed at 100 x 100 for all camera presets + action_space = 1 + state_space = 0 + observation_space = [100, 100, 3] + + # change viewer settings + viewer = ViewerCfg(eye=(20.0, 20.0, 20.0)) + + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=512, env_spacing=20.0, replicate_physics=True) + + # reset + max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] + initial_pole_angle_range = [-0.125, 0.125] # the range in which the pole angle is sampled from on reset [rad] + + # reward scales + rew_scale_alive = 1.0 + rew_scale_terminated = -2.0 + rew_scale_pole_pos = -1.0 + rew_scale_cart_vel = -0.01 + rew_scale_pole_vel = -0.005 + + +@configclass +class CartpoleDepthCameraEnvCfg(CartpoleRGBCameraEnvCfg): + # camera + tiled_camera: TiledCameraCfg = TiledCameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(0.0, 0.0, 0.0, 1.0), convention="world"), + data_types=["depth"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ), + width=100, + height=100, + ) + + # spaces + observation_space = [100, 100, 1] + + +@configclass +class CartpoleAlbedoCameraEnvCfg(CartpoleRGBCameraEnvCfg): + # camera + tiled_camera: TiledCameraCfg = TiledCameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), + data_types=["albedo"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ), + width=100, + height=100, + ) + + # spaces + observation_space = [100, 100, 3] + + +@configclass +class CartpoleSimpleShadingConstantCameraEnvCfg(CartpoleRGBCameraEnvCfg): + # camera + tiled_camera: TiledCameraCfg = TiledCameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), + data_types=["simple_shading_constant_diffuse"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ), + width=100, + height=100, + ) + + # spaces + observation_space = [100, 100, 3] + + +@configclass +class CartpoleSimpleShadingDiffuseCameraEnvCfg(CartpoleRGBCameraEnvCfg): + # camera + tiled_camera: TiledCameraCfg = TiledCameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), + data_types=["simple_shading_diffuse_mdl"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ), + width=100, + height=100, + ) + + # spaces + observation_space = [100, 100, 3] + + +@configclass +class CartpoleSimpleShadingFullCameraEnvCfg(CartpoleRGBCameraEnvCfg): + # camera + tiled_camera: TiledCameraCfg = TiledCameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=TiledCameraCfg.OffsetCfg(pos=(-5.0, 0.0, 2.0), rot=(1.0, 0.0, 0.0, 0.0), convention="world"), + data_types=["simple_shading_full_mdl"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 20.0) + ), + width=100, + height=100, + ) + + # spaces + observation_space = [100, 100, 3] diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py index 03ad3ff4422..b565705c54a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env.py @@ -6,56 +6,18 @@ from __future__ import annotations import math +import torch from collections.abc import Sequence -import torch import warp as wp import isaaclab.sim as sim_utils -from isaaclab.assets import Articulation, ArticulationCfg -from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg -from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sim import SimulationCfg +from isaaclab.assets import Articulation +from isaaclab.envs import DirectRLEnv from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane -from isaaclab.utils import configclass from isaaclab.utils.math import sample_uniform -from isaaclab_assets.robots.cartpole import CARTPOLE_CFG - - -@configclass -class CartpoleEnvCfg(DirectRLEnvCfg): - # env - decimation = 2 - episode_length_s = 5.0 - action_scale = 100.0 # [N] - action_space = 1 - observation_space = 4 - state_space = 0 - - # simulation - sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) - - # robot - robot_cfg: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="/World/envs/env_.*/Robot") - cart_dof_name = "slider_to_cart" - pole_dof_name = "cart_to_pole" - - # scene - scene: InteractiveSceneCfg = InteractiveSceneCfg( - num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True - ) - - # reset - max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] - initial_pole_angle_range = [-0.25, 0.25] # the range in which the pole angle is sampled from on reset [rad] - - # reward scales - rew_scale_alive = 1.0 - rew_scale_terminated = -2.0 - rew_scale_pole_pos = -1.0 - rew_scale_cart_vel = -0.01 - rew_scale_pole_vel = -0.005 +from .cartpole_env_cfg import CartpoleEnvCfg class CartpoleEnv(DirectRLEnv): @@ -134,24 +96,25 @@ def _reset_idx(self, env_ids: Sequence[int] | None): env_ids = self.cartpole._ALL_INDICES super()._reset_idx(env_ids) - joint_pos = wp.to_torch(self.cartpole.data.default_joint_pos)[env_ids] + joint_pos = wp.to_torch(self.cartpole.data.default_joint_pos)[env_ids].clone() joint_pos[:, self._pole_dof_idx] += sample_uniform( self.cfg.initial_pole_angle_range[0] * math.pi, self.cfg.initial_pole_angle_range[1] * math.pi, joint_pos[:, self._pole_dof_idx].shape, joint_pos.device, ) - joint_vel = wp.to_torch(self.cartpole.data.default_joint_vel)[env_ids] + joint_vel = wp.to_torch(self.cartpole.data.default_joint_vel)[env_ids].clone() - default_root_state = wp.to_torch(self.cartpole.data.default_root_state)[env_ids] - default_root_state[:, :3] += self.scene.env_origins[env_ids] + default_root_pose = wp.to_torch(self.cartpole.data.default_root_pose)[env_ids].clone() + default_root_pose[:, :3] += self.scene.env_origins[env_ids] self.joint_pos[env_ids] = joint_pos self.joint_vel[env_ids] = joint_vel - self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids) - self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) - self.cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self.cartpole.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids) + self.cartpole.write_joint_state_to_sim_index( + position=joint_pos, velocity=joint_vel, joint_ids=None, env_ids=env_ids + ) @torch.jit.script @@ -173,4 +136,4 @@ def compute_rewards( rew_cart_vel = rew_scale_cart_vel * torch.sum(torch.abs(cart_vel).unsqueeze(dim=1), dim=-1) rew_pole_vel = rew_scale_pole_vel * torch.sum(torch.abs(pole_vel).unsqueeze(dim=1), dim=-1) total_reward = rew_alive + rew_termination + rew_pole_pos + rew_cart_vel + rew_pole_vel - return total_reward + return total_reward \ No newline at end of file diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env_cfg.py new file mode 100644 index 00000000000..d1c31d13d87 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole/cartpole_env_cfg.py @@ -0,0 +1,49 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.assets import ArticulationCfg +from isaaclab.envs import DirectRLEnvCfg +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim import SimulationCfg +from isaaclab.utils import configclass + +from isaaclab_assets.robots.cartpole import CARTPOLE_CFG + + +@configclass +class CartpoleEnvCfg(DirectRLEnvCfg): + # env + decimation = 2 + episode_length_s = 5.0 + action_scale = 100.0 # [N] + action_space = 1 + observation_space = 4 + state_space = 0 + + # simulation + sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation) + + # robot + robot_cfg: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="/World/envs/env_.*/Robot") + cart_dof_name = "slider_to_cart" + pole_dof_name = "cart_to_pole" + + # scene + scene: InteractiveSceneCfg = InteractiveSceneCfg( + num_envs=4096, env_spacing=4.0, replicate_physics=True, clone_in_fabric=True + ) + + # reset + max_cart_pos = 3.0 # the cart is reset if it exceeds that position [m] + initial_pole_angle_range = [-0.25, 0.25] # the range in which the pole angle is sampled from on reset [rad] + + # reward scales + rew_scale_alive = 1.0 + rew_scale_terminated = -2.0 + rew_scale_pole_pos = -1.0 + rew_scale_cart_vel = -0.01 + rew_scale_pole_vel = -0.005 diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole_showcase/cartpole/cartpole_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole_showcase/cartpole/cartpole_env_cfg.py index e6e8169b1ba..9384305a07e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole_showcase/cartpole/cartpole_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole_showcase/cartpole/cartpole_env_cfg.py @@ -9,7 +9,7 @@ from isaaclab.utils import configclass -from isaaclab_tasks.direct.cartpole.cartpole_env import CartpoleEnvCfg +from isaaclab_tasks.direct.cartpole.cartpole_env_cfg import CartpoleEnvCfg ### # Observation space as Box diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole_showcase/cartpole_camera/cartpole_camera_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole_showcase/cartpole_camera/cartpole_camera_env_cfg.py index cb37d20e047..0625bb2a172 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole_showcase/cartpole_camera/cartpole_camera_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/cartpole_showcase/cartpole_camera/cartpole_camera_env_cfg.py @@ -11,7 +11,7 @@ from isaaclab.sensors import TiledCameraCfg from isaaclab.utils import configclass -from isaaclab_tasks.direct.cartpole.cartpole_camera_env import CartpoleRGBCameraEnvCfg as CartpoleCameraEnvCfg +from isaaclab_tasks.direct.cartpole.cartpole_camera_env_cfg import CartpoleRGBCameraEnvCfg as CartpoleCameraEnvCfg def get_tiled_camera_cfg(data_type: str, width: int = 100, height: int = 100) -> TiledCameraCfg: diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py index 2faffb91746..de6f30dc8ec 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/factory/factory_env.py @@ -543,7 +543,7 @@ def set_pos_inverse_kinematics( self.ctrl_target_joint_pos[env_ids, 0:7] = self.joint_pos[env_ids, 0:7] # Update dof state. - self._robot.write_joint_state_to_sim(self.joint_pos, self.joint_vel) + self._robot.write_joint_state_to_sim(position=self.joint_pos, velocity=self.joint_vel) self._robot.set_joint_position_target(self.ctrl_target_joint_pos) # Simulate and update tensors. @@ -596,7 +596,7 @@ def _set_franka_to_default_pose(self, joints, env_ids): joint_effort = torch.zeros_like(joint_pos) self.ctrl_target_joint_pos[env_ids, :] = joint_pos self._robot.set_joint_position_target(self.ctrl_target_joint_pos[env_ids], env_ids=env_ids) - self._robot.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + self._robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) self._robot.reset() self._robot.set_joint_effort_target(joint_effort, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py index 554e76cd339..a815bc9c44a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/franka_cabinet/franka_cabinet_env.py @@ -8,15 +8,12 @@ import torch import warp as wp -from pxr import UsdGeom - import isaaclab.sim as sim_utils from isaaclab.actuators import ImplicitActuatorCfg from isaaclab.assets import Articulation, ArticulationCfg from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg from isaaclab.scene import InteractiveSceneCfg from isaaclab.sim import SimulationCfg -from isaaclab.sim.utils.stage import get_current_stage from isaaclab.terrains import TerrainImporterCfg from isaaclab.utils import configclass from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR @@ -206,6 +203,10 @@ def get_env_local_pose(env_pos: torch.Tensor, xformable: UsdGeom.Xformable, devi self.robot_dof_targets = torch.zeros((self.num_envs, self._robot.num_joints), device=self.device) + from pxr import UsdGeom + + from isaaclab.sim.utils.stage import get_current_stage + stage = get_current_stage() hand_pose = get_env_local_pose( self.scene.env_origins[0], @@ -342,11 +343,11 @@ def _reset_idx(self, env_ids: torch.Tensor | None): joint_pos = torch.clamp(joint_pos, self.robot_dof_lower_limits, self.robot_dof_upper_limits) joint_vel = torch.zeros_like(joint_pos) self._robot.set_joint_position_target(joint_pos, env_ids=env_ids) - self._robot.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + self._robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) # cabinet state zeros = torch.zeros((len(env_ids), self._cabinet.num_joints), device=self.device) - self._cabinet.write_joint_state_to_sim(zeros, zeros, env_ids=env_ids) + self._cabinet.write_joint_state_to_sim(position=zeros, velocity=zeros, env_ids=env_ids) # Need to refresh the intermediate values so that _get_observations() can use the latest values self._compute_intermediate_values(env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py index 34e64d609d8..9ac9e56b834 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/humanoid_amp/humanoid_amp_env.py @@ -136,7 +136,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): self.robot.write_root_link_pose_to_sim(root_state[:, :7], env_ids) self.robot.write_root_com_velocity_to_sim(root_state[:, 7:], env_ids) - self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self.robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) # reset strategies diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py index 7ec39aef9a1..38bb94b4b28 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/inhand_manipulation/inhand_manipulation_env.py @@ -242,7 +242,7 @@ def _reset_idx(self, env_ids: Sequence[int] | None): self.hand_dof_targets[env_ids] = dof_pos self.hand.set_joint_position_target(dof_pos, env_ids=env_ids) - self.hand.write_joint_state_to_sim(dof_pos, dof_vel, env_ids=env_ids) + self.hand.write_joint_state_to_sim(position=dof_pos, velocity=dof_vel, env_ids=env_ids) self.successes[env_ids] = 0 self._compute_intermediate_values() diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py index 5eaac627897..893740f3898 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/locomotion/locomotion_env.py @@ -220,7 +220,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): self.robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) - self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self.robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) to_target = self.targets[env_ids] - default_root_state[:, :3] to_target[:, 2] = 0.0 diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py index 12506d0fe87..e8803227e1f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/quadcopter/quadcopter_env.py @@ -236,7 +236,7 @@ def _reset_idx(self, env_ids: torch.Tensor | None): default_root_state[:, :3] += self._terrain.env_origins[env_ids] self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) - self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self._robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) def _set_debug_vis_impl(self, debug_vis: bool): # create markers if necessary for the first time diff --git a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py index a03b4e050c5..7573c2d5901 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/direct/shadow_hand_over/shadow_hand_over_env.py @@ -347,7 +347,7 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): self.right_hand_dof_targets[env_ids] = dof_pos self.right_hand.set_joint_position_target(dof_pos, env_ids=env_ids) - self.right_hand.write_joint_state_to_sim(dof_pos, dof_vel, env_ids=env_ids) + self.right_hand.write_joint_state_to_sim(position=dof_pos, velocity=dof_vel, env_ids=env_ids) # reset left hand delta_max = self.hand_dof_upper_limits[env_ids] - wp.to_torch(self.left_hand.data.default_joint_pos)[env_ids] @@ -369,7 +369,7 @@ def _reset_idx(self, env_ids: Sequence[int] | torch.Tensor | None): self.left_hand_dof_targets[env_ids] = dof_pos self.left_hand.set_joint_position_target(dof_pos, env_ids=env_ids) - self.left_hand.write_joint_state_to_sim(dof_pos, dof_vel, env_ids=env_ids) + self.left_hand.write_joint_state_to_sim(position=dof_pos, velocity=dof_vel, env_ids=env_ids) self._compute_intermediate_values() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.py index 155079c558f..041e7cccadd 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/__init__.py @@ -5,6 +5,10 @@ """This sub-module contains the functions that are specific to the cartpole environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .rewards import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/rewards.py index e20bd178c6e..3bdb71d5606 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/cartpole/mdp/rewards.py @@ -10,11 +10,11 @@ import torch import warp as wp -from isaaclab.assets import Articulation from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import wrap_to_pi if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.py index 9fd05f55635..68bebf6077e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/__init__.py @@ -5,7 +5,11 @@ """This sub-module contains the functions that are specific to the humanoid environment.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .observations import * from .rewards import * + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/observations.py index c7fe81a2f27..5445328cc17 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/observations.py @@ -11,10 +11,10 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/rewards.py index 209fbad7a4a..ceaa1371b91 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/classic/humanoid/mdp/rewards.py @@ -12,12 +12,12 @@ import isaaclab.utils.math as math_utils import isaaclab.utils.string as string_utils -from isaaclab.assets import Articulation from isaaclab.managers import ManagerTermBase, RewardTermCfg, SceneEntityCfg from . import observations as obs if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv 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 index bc4d65f5b1f..1803c9d4336 100644 --- 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 @@ -5,10 +5,14 @@ """This sub-module contains the functions that are specific to the drone ARL environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from isaaclab_contrib.mdp import * # noqa: F401, F403 from .commands import * # noqa: F401, F403 from .observations import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) 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 index 634ad6d2226..61e5908b1e6 100644 --- 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 @@ -18,12 +18,12 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation from isaaclab.managers import SceneEntityCfg from isaaclab_contrib.assets import Multirotor if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv from isaaclab.envs.utils.io_descriptors import generic_io_descriptor, record_shape 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 index c7851154118..d508e982f0c 100644 --- 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 @@ -12,12 +12,12 @@ import isaaclab.utils.math as math_utils +from isaaclab.managers import SceneEntityCfg + if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv -from isaaclab.assets import RigidObject -from isaaclab.managers import SceneEntityCfg - """ Drone control rewards. """ diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.py index 7e559309b5c..55d2e3aaedb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomanipulation/pick_place/mdp/__init__.py @@ -6,7 +6,11 @@ """This sub-module contains the functions that are specific to the locomanipulation environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .actions import * # noqa: F401, F403 from .observations import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/events.py index a807e33d3f2..c48b37df931 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/events.py @@ -17,11 +17,11 @@ import torch import warp as wp -from isaaclab.assets import Articulation from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import sample_uniform if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv @@ -57,4 +57,4 @@ def reset_joints_around_default( joint_pos = sample_uniform(joint_min_pos, joint_max_pos, joint_min_pos.shape, joint_min_pos.device) joint_vel = sample_uniform(joint_min_vel, joint_max_vel, joint_min_vel.shape, joint_min_vel.device) # set into the physics simulation - asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + asset.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py index 5f20b91acb4..f4961b9b4e2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/spot/mdp/rewards.py @@ -16,13 +16,13 @@ import torch import warp as wp -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import ManagerTermBase, SceneEntityCfg -from isaaclab.sensors import ContactSensor if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv from isaaclab.managers import RewardTermCfg + from isaaclab.sensors import ContactSensor ## diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/__init__.py index 6f6cad00712..cc266156107 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/__init__.py @@ -5,8 +5,30 @@ """This sub-module contains the functions that are specific to the locomotion environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import lazy_loader as lazy -from .curriculums import * # noqa: F401, F403 -from .rewards import * # noqa: F401, F403 -from .terminations import * # noqa: F401, F403 +_lazy_getattr, __dir__, __all__ = lazy.attach( + __name__, + submod_attrs={ + "curriculums": ["terrain_levels_vel"], + "rewards": [ + "feet_air_time", + "feet_air_time_positive_biped", + "feet_slide", + "track_lin_vel_xy_yaw_frame_exp", + "track_ang_vel_z_world_exp", + "stand_still_joint_deviation_l1", + ], + "terminations": ["terrain_out_of_bounds"], + }, +) + + +def __getattr__(name): + try: + return _lazy_getattr(name) + except AttributeError: + pass + import isaaclab.envs.mdp as _parent_mdp + + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py index 750aafefd9c..ece0a616115 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/curriculums.py @@ -17,12 +17,12 @@ import torch import warp as wp -from isaaclab.assets import Articulation from isaaclab.managers import SceneEntityCfg -from isaaclab.terrains import TerrainImporter if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.terrains import TerrainImporter def terrain_levels_vel( diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/rewards.py index 64c197e9f96..2198cff0cea 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/rewards.py @@ -18,11 +18,11 @@ from isaaclab.envs import mdp from isaaclab.managers import SceneEntityCfg -from isaaclab.sensors import ContactSensor from isaaclab.utils.math import quat_apply_inverse, yaw_quat if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import ContactSensor def feet_air_time( diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/terminations.py index 6bb6e778986..63add623c46 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/mdp/terminations.py @@ -16,10 +16,10 @@ import torch import warp as wp -from isaaclab.assets import RigidObject from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/velocity_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/velocity_env_cfg.py index 63cb6f086f9..fe72d63d83e 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/velocity_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/velocity_env_cfg.py @@ -7,9 +7,9 @@ from dataclasses import MISSING from isaaclab_physx.physics import PhysxCfg - -import isaaclab.sim as sim_utils +from isaaclab_newton.physics import NewtonCfg, MJWarpSolverCfg from isaaclab.assets import ArticulationCfg, AssetBaseCfg +import isaaclab.sim as sim_utils from isaaclab.envs import ManagerBasedRLEnvCfg from isaaclab.managers import CurriculumTermCfg as CurrTerm from isaaclab.managers import EventTermCfg as EventTerm @@ -19,7 +19,10 @@ from isaaclab.managers import SceneEntityCfg from isaaclab.managers import TerminationTermCfg as DoneTerm from isaaclab.scene import InteractiveSceneCfg -from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns +from isaaclab.sensors import RayCasterCfg, patterns +from isaaclab.sensors import ContactSensorCfg +from isaaclab_newton.sensors import ContactSensorCfg as NewtonContactSensorCfg +from isaaclab_physx.sensors import ContactSensorCfg as PhysXContactSensorCfg from isaaclab.terrains import TerrainImporterCfg from isaaclab.utils import configclass from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR @@ -37,6 +40,12 @@ # Scene definition ## +@configclass +class VelocityEnvContactSensorCfg: + presets: dict[str, ContactSensorCfg] = { + "default": PhysXContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True), + "newton": NewtonContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True) + } @configclass class MySceneCfg(InteractiveSceneCfg): @@ -73,7 +82,7 @@ class MySceneCfg(InteractiveSceneCfg): debug_vis=False, mesh_prim_paths=["/World/ground"], ) - contact_forces = ContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True) + contact_forces = VelocityEnvContactSensorCfg() # lights sky_light = AssetBaseCfg( prim_path="/World/skyLight", @@ -149,41 +158,8 @@ def __post_init__(self): @configclass -class EventCfg: +class VelocityEnvEventsCfg: """Configuration for events.""" - - # startup - physics_material = EventTerm( - func=mdp.randomize_rigid_body_material, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names=".*"), - "static_friction_range": (0.8, 0.8), - "dynamic_friction_range": (0.6, 0.6), - "restitution_range": (0.0, 0.0), - "num_buckets": 64, - }, - ) - - add_base_mass = EventTerm( - func=mdp.randomize_rigid_body_mass, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names="base"), - "mass_distribution_params": (-5.0, 5.0), - "operation": "add", - }, - ) - - base_com = EventTerm( - func=mdp.randomize_rigid_body_com, - mode="startup", - params={ - "asset_cfg": SceneEntityCfg("robot", body_names="base"), - "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)}, - }, - ) - # reset base_external_force_torque = EventTerm( func=mdp.apply_external_force_torque, @@ -228,6 +204,44 @@ class EventCfg: params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}}, ) +@configclass +class EventsCfg(VelocityEnvEventsCfg): + # startup + physics_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "static_friction_range": (0.8, 0.8), + "dynamic_friction_range": (0.6, 0.6), + "restitution_range": (0.0, 0.0), + "num_buckets": 64, + }, + ) + + add_base_mass = EventTerm( + func=mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names="base"), + "mass_distribution_params": (-5.0, 5.0), + "operation": "add", + }, + ) + + base_com = EventTerm( + func=mdp.randomize_rigid_body_com, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names="base"), + "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)}, + }, + ) + + presets: dict[str, VelocityEnvEventsCfg] = { + "newton": VelocityEnvEventsCfg(), + } + @configclass class RewardsCfg: @@ -287,6 +301,26 @@ class CurriculumCfg: # Environment configuration ## +@configclass +class PhysicsCfg: + presets = { + "default" : PhysxCfg(gpu_max_rigid_patch_count=10 * 2**15), + "newton" : NewtonCfg( + solver_cfg=MJWarpSolverCfg( + njmax=45, + nconmax=30, + ls_iterations=30, + cone="pyramidal", + impratio=1, + ls_parallel=True, + integrator="implicitfast", + ), + num_substeps=1, + debug_mode=False, + ), + } + + @configclass class LocomotionVelocityRoughEnvCfg(ManagerBasedRLEnvCfg): @@ -301,7 +335,7 @@ class LocomotionVelocityRoughEnvCfg(ManagerBasedRLEnvCfg): # MDP settings rewards: RewardsCfg = RewardsCfg() terminations: TerminationsCfg = TerminationsCfg() - events: EventCfg = EventCfg() + events: EventsCfg = EventsCfg() curriculum: CurriculumCfg = CurriculumCfg() def __post_init__(self): @@ -313,7 +347,7 @@ def __post_init__(self): self.sim.dt = 0.005 self.sim.render_interval = self.decimation self.sim.physics_material = self.scene.terrain.physics_material - self.sim.physics = PhysxCfg(gpu_max_rigid_patch_count=10 * 2**15) + self.sim.physics = PhysicsCfg() # update sensor update periods # we tick all the sensors based on the smallest update period (physics update period) if self.scene.height_scanner is not None: diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.py index 79a9af2f736..1585f49053c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/__init__.py @@ -5,7 +5,11 @@ """This sub-module contains the functions that are specific to the cabinet environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .observations import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/observations.py index a6a081ca3f1..75e1ff44157 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cabinet/mdp/observations.py @@ -11,11 +11,11 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import ArticulationData -from isaaclab.sensors import FrameTransformerData if TYPE_CHECKING: + from isaaclab.assets import ArticulationData from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import FrameTransformerData def rel_ee_object_distance(env: ManagerBasedRLEnv) -> torch.Tensor: diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.py index 10ab3ea7e7f..61765a2eb7d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/__init__.py @@ -5,10 +5,14 @@ """This sub-module contains the functions that are specific to the locomotion environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .events import * # noqa: F401, F403 from .noise_models import * # noqa: F401, F403 from .observations import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 from .terminations import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py index 80c1827bf20..6ea83b10ef4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/events.py @@ -14,12 +14,12 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg from isaaclab_tasks.direct.automate import factory_control as fc if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedEnv @@ -345,7 +345,7 @@ def __call__( # Write to sim self.robot_asset.set_joint_position_target(joint_pos, env_ids=env_ids) self.robot_asset.set_joint_velocity_target(joint_vel, env_ids=env_ids) - self.robot_asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + self.robot_asset.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) # Set gripper to grasp position joint_pos = wp.to_torch(self.robot_asset.data.joint_pos)[env_ids].clone() @@ -358,7 +358,7 @@ def __call__( self.gripper_joint_setter_func(joint_pos, [row_idx], self.finger_joints, hand_grasp_width) self.robot_asset.set_joint_position_target(joint_pos, joint_ids=self.all_joints, env_ids=env_ids) - self.robot_asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + self.robot_asset.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) # Set gripper to closed position for row_idx, env_id in enumerate(env_ids.tolist()): diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py index bacf41c110e..e70b77d7d8f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/observations.py @@ -12,11 +12,11 @@ import torch import warp as wp -from isaaclab.assets import RigidObject from isaaclab.managers import ManagerTermBase, ObservationTermCfg, SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv from .events import randomize_gear_type diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/rewards.py index ea8e9a63756..2204a0186bc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/rewards.py @@ -13,11 +13,11 @@ import warp as wp from isaaclab.managers import ManagerTermBase, RewardTermCfg, SceneEntityCfg -from isaaclab.sensors.frame_transformer.frame_transformer import FrameTransformer from isaaclab.utils.math import combine_frame_transforms if TYPE_CHECKING: from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors.frame_transformer.frame_transformer import FrameTransformer from .events import randomize_gear_type diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/terminations.py index 7eee20fc81d..4d6c8a643ce 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/deploy/mdp/terminations.py @@ -7,18 +7,19 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING import torch import warp as wp -import carb - import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation + +logger = logging.getLogger(__name__) from isaaclab.managers import ManagerTermBase, SceneEntityCfg, TerminationTermCfg if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv from .events import randomize_gear_type @@ -117,7 +118,7 @@ def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedEnv): # Find end effector index once eef_indices, _ = self.robot_asset.find_bodies([self.end_effector_body_name]) if len(eef_indices) == 0: - carb.log_warn( + logger.warning( f"{self.end_effector_body_name} not found in robot body names. Cannot check gear drop condition." ) self.eef_idx = None @@ -250,7 +251,7 @@ def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedEnv): # Find end effector index once eef_indices, _ = self.robot_asset.find_bodies([self.end_effector_body_name]) if len(eef_indices) == 0: - carb.log_warn( + logger.warning( f"{self.end_effector_body_name} not found in robot body names. Cannot check gear orientation condition." ) self.eef_idx = None diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py index a6537b1a5e1..3f56ae4414b 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/__init__.py @@ -3,10 +3,14 @@ # # SPDX-License-Identifier: BSD-3-Clause -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .commands import * # noqa: F401, F403 from .curriculums import * # noqa: F401, F403 from .observations import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 from .terminations import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands_cfg.py index e3c83882a3f..385760e1aa3 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/commands/pose_commands_cfg.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + from dataclasses import MISSING import isaaclab.sim as sim_utils @@ -10,8 +12,7 @@ from isaaclab.markers import VisualizationMarkersCfg from isaaclab.utils import configclass from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR - -from . import pose_commands as dex_cmd +from isaaclab.utils.string import DeferredClass ALIGN_MARKER_CFG = VisualizationMarkersCfg( markers={ @@ -35,7 +36,9 @@ class ObjectUniformPoseCommandCfg(CommandTermCfg): """Configuration for uniform pose command generator.""" - class_type: type = dex_cmd.ObjectUniformPoseCommand + class_type: type = DeferredClass( + "isaaclab_tasks.manager_based.manipulation.dexsuite.mdp.commands.pose_commands:ObjectUniformPoseCommand" + ) asset_name: str = MISSING """Name of the coordinate referencing asset in the environment for which the commands are generated respect to.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py index 0b42e2f38ae..c6b64556bf8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/curriculums.py @@ -11,12 +11,12 @@ import torch import warp as wp -from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import mdp from isaaclab.managers import ManagerTermBase, SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms, compute_pose_error if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py index 84faaace851..63ffd23081a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/observations.py @@ -10,13 +10,13 @@ import torch import warp as wp -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import ManagerTermBase, SceneEntityCfg from isaaclab.utils.math import quat_apply, quat_apply_inverse, quat_inv, quat_mul, subtract_frame_transforms from .utils import sample_object_point_cloud if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py index 47dde4eee17..c2a7a86f723 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/rewards.py @@ -10,14 +10,14 @@ import torch import warp as wp -from isaaclab.assets import RigidObject from isaaclab.managers import SceneEntityCfg -from isaaclab.sensors import ContactSensor from isaaclab.utils import math as math_utils from isaaclab.utils.math import combine_frame_transforms, compute_pose_error if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import ContactSensor def action_rate_l2_clamped(env: ManagerBasedRLEnv) -> torch.Tensor: diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py index 86ef893a5c3..96eea355a5c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/terminations.py @@ -16,10 +16,10 @@ import torch import warp as wp -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.py index 0fafe450036..86f345b84fd 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/__init__.py @@ -5,10 +5,14 @@ """This sub-module contains the functions that are specific to the in-hand manipulation environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .commands import * # noqa: F401, F403 from .events import * # noqa: F401, F403 from .observations import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 from .terminations import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/events.py index d36c2aa6823..6e68011d5ca 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/events.py @@ -12,11 +12,11 @@ import torch import warp as wp -from isaaclab.assets import Articulation from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg from isaaclab.utils.math import sample_uniform if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedEnv @@ -181,4 +181,4 @@ def __call__( joint_vel = joint_vel.clamp(-joint_vel_limits, joint_vel_limits) # set into the physics simulation - self._asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + self._asset.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/observations.py index 3d392e1d348..8baf070b074 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/observations.py @@ -5,17 +5,20 @@ """Functions specific to the in-hand dexterous manipulation environments.""" +from __future__ import annotations + from typing import TYPE_CHECKING import torch import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import RigidObject + from .commands import InHandReOrientationCommand diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/rewards.py index 1e104f283bf..a2f84c47e19 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/inhand/mdp/rewards.py @@ -5,17 +5,20 @@ """Functions specific to the in-hand dexterous manipulation environments.""" +from __future__ import annotations + from typing import TYPE_CHECKING import torch import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import RigidObject + from .commands import InHandReOrientationCommand diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.py index f3dd0fecdf8..1c1364db41f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/__init__.py @@ -5,8 +5,12 @@ """This sub-module contains the functions that are specific to the lift environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .observations import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 from .terminations import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py index 15d2a6a660f..1ccf6096981 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/observations.py @@ -10,11 +10,11 @@ import torch import warp as wp -from isaaclab.assets import RigidObject from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import subtract_frame_transforms if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py index f558a3ec652..2f1ccaf0287 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/rewards.py @@ -10,13 +10,13 @@ import torch import warp as wp -from isaaclab.assets import RigidObject from isaaclab.managers import SceneEntityCfg -from isaaclab.sensors import FrameTransformer from isaaclab.utils.math import combine_frame_transforms if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import FrameTransformer def object_is_lifted( diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py index 60e2675d8fe..756cbae9049 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/lift/mdp/terminations.py @@ -15,11 +15,11 @@ import torch -from isaaclab.assets import RigidObject from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.py index 555bfb7cbe8..b50c74cb7f5 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/__init__.py @@ -5,8 +5,12 @@ """This sub-module contains the functions that are specific to the lift environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .observations import * # noqa: F401, F403 from .pick_place_events import * # noqa: F401, F403 from .terminations import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/terminations.py index 453e4732477..6c1f9161042 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/pick_place/mdp/terminations.py @@ -16,10 +16,10 @@ import torch import warp as wp -from isaaclab.assets import RigidObject from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.py index 41f76fcdb1b..0864271e738 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/__init__.py @@ -5,7 +5,11 @@ """This sub-module contains the functions that are specific to the pick and place environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .observations import * # noqa: F401, F403 from .terminations import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/observations.py index e3748c0b4c1..70c4e9391ad 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/observations.py @@ -11,12 +11,12 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import SceneEntityCfg -from isaaclab.sensors import FrameTransformer if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import FrameTransformer def object_poses_in_base_frame( diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/terminations.py index a6dba6bca6b..139617d3050 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/place/mdp/terminations.py @@ -17,10 +17,10 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.py index 3fec83fe70a..0f112ed4894 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/__init__.py @@ -5,6 +5,10 @@ """This sub-module contains the functions that are specific to the locomotion environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .rewards import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py index f51db1dcfad..6ba3176b535 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/reach/mdp/rewards.py @@ -10,11 +10,11 @@ import torch import warp as wp -from isaaclab.assets import RigidObject from isaaclab.managers import SceneEntityCfg from isaaclab.utils.math import combine_frame_transforms, quat_error_magnitude, quat_mul if TYPE_CHECKING: + from isaaclab.assets import RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.py index ea04fcc468e..b68181b3371 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/__init__.py @@ -5,7 +5,11 @@ """This sub-module contains the functions that are specific to the lift environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .observations import * # noqa: F401, F403 from .terminations import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py index b20600f59a7..8ef8cb6eb64 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/franka_stack_events.py @@ -17,10 +17,10 @@ from isaacsim.core.utils.extensions import enable_extension import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation, AssetBase from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import Articulation, AssetBase from isaaclab.envs import ManagerBasedEnv @@ -73,7 +73,7 @@ def randomize_joint_by_gaussian_offset( # Set into the physics simulation asset.set_joint_position_target(joint_pos, env_ids=env_ids) asset.set_joint_velocity_target(joint_vel, env_ids=env_ids) - asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids) + asset.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, env_ids=env_ids) def sample_random_color(base=(0.75, 0.75, 0.75), variation=0.1): diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py index 8cd03a2133f..dd6da5177e8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py @@ -11,12 +11,12 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation, RigidObject, RigidObjectCollection from isaaclab.managers import SceneEntityCfg -from isaaclab.sensors import FrameTransformer if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject, RigidObjectCollection from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.sensors import FrameTransformer def cube_positions_in_world_frame( diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py index 9da0c6abcbf..bacc3b9d563 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py @@ -16,10 +16,10 @@ import torch import warp as wp -from isaaclab.assets import Articulation, RigidObject from isaaclab.managers import SceneEntityCfg if TYPE_CHECKING: + from isaaclab.assets import Articulation, RigidObject from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py index 213391d362b..66255a3a3f1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/__init__.py @@ -5,7 +5,11 @@ """This sub-module contains the functions that are specific to the locomotion environments.""" -from isaaclab.envs.mdp import * # noqa: F401, F403 +import isaaclab.envs.mdp as _parent_mdp from .pre_trained_policy_action import * # noqa: F401, F403 from .rewards import * # noqa: F401, F403 + + +def __getattr__(name): + return getattr(_parent_mdp, name) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py index a5956bdc56e..a16699888b1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/navigation/mdp/pre_trained_policy_action.py @@ -12,7 +12,6 @@ import warp as wp import isaaclab.utils.math as math_utils -from isaaclab.assets import Articulation from isaaclab.managers import ActionTerm, ActionTermCfg, ObservationGroupCfg, ObservationManager from isaaclab.markers import VisualizationMarkers from isaaclab.markers.config import BLUE_ARROW_X_MARKER_CFG, GREEN_ARROW_X_MARKER_CFG @@ -20,6 +19,7 @@ from isaaclab.utils.assets import check_file_path, read_file if TYPE_CHECKING: + from isaaclab.assets import Articulation from isaaclab.envs import ManagerBasedRLEnv diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py b/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py index 16785f326f8..6ba12b85315 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/hydra.py @@ -99,6 +99,79 @@ def collect_presets(cfg, path: str = "") -> dict: return result +def resolve_task_config(task_name: str, agent_cfg_entry_point: str): + """Resolve env and agent configs with Hydra overrides and presets applied. + + This performs the full Hydra resolution and returns the final configs. + It is safe to call before Kit is launched because it defers + ``from_dict`` (which triggers ``string_to_callable`` and imports + implementation modules) until :func:`finalize_task_config` is called. + + Args: + task_name: Task name (e.g., "Isaac-Velocity-Flat-Anymal-C-v0"). + agent_cfg_entry_point: Agent config entry point key (e.g., "rsl_rl_cfg_entry_point"). + + Returns: + Tuple of (env_cfg, agent_cfg) with Hydra presets applied to the + config objects. Scalar Hydra overrides are captured but not yet + synced to the config objects — call :func:`finalize_task_config` + after Kit is running to complete that step. + """ + task = task_name.split(":")[-1] + env_cfg, agent_cfg, presets = register_task(task, agent_cfg_entry_point) + + global_presets, preset_sel, preset_scalar, global_scalar = parse_overrides( + sys.argv[1:], presets + ) + + original_argv, sys.argv = sys.argv, [sys.argv[0]] + global_scalar + + resolved = {} + + @hydra.main(config_path=None, config_name=task, version_base="1.3") + def hydra_main(hydra_cfg, env_cfg=env_cfg, agent_cfg=agent_cfg): + hydra_cfg = replace_strings_with_slices(OmegaConf.to_container(hydra_cfg, resolve=True)) + apply_overrides( + env_cfg, agent_cfg, hydra_cfg, global_presets, preset_sel, preset_scalar, presets + ) + _cleanup_presets(env_cfg) + _cleanup_presets(agent_cfg) + _cleanup_presets_dict(hydra_cfg) + resolved["env_cfg"] = env_cfg + resolved["agent_cfg"] = agent_cfg + resolved["hydra_cfg"] = hydra_cfg + + try: + hydra_main() + finally: + sys.argv = original_argv + + return resolved["env_cfg"], resolved["agent_cfg"], resolved["hydra_cfg"] + + +def finalize_task_config(env_cfg, agent_cfg, hydra_cfg): + """Sync Hydra scalar overrides into config objects. + + Call this **after** Kit is launched so that ``from_dict`` can safely + resolve callable strings (which import implementation modules). + + Args: + env_cfg: Environment config returned by :func:`resolve_task_config`. + agent_cfg: Agent config returned by :func:`resolve_task_config`. + hydra_cfg: Raw Hydra dict returned by :func:`resolve_task_config`. + + Returns: + Tuple of (env_cfg, agent_cfg) fully resolved. + """ + env_cfg.from_dict(hydra_cfg["env"]) + env_cfg = replace_strings_with_env_cfg_spaces(env_cfg) + if isinstance(agent_cfg, dict) or agent_cfg is None: + agent_cfg = hydra_cfg["agent"] + else: + agent_cfg.from_dict(hydra_cfg["agent"]) + return env_cfg, agent_cfg + + def hydra_task_config(task_name: str, agent_cfg_entry_point: str) -> Callable: """Decorator for Hydra config with REPLACE-only preset semantics. @@ -131,6 +204,10 @@ def hydra_main(hydra_cfg, env_cfg=env_cfg, agent_cfg=agent_cfg): apply_overrides( env_cfg, agent_cfg, hydra_cfg, global_presets, preset_sel, preset_scalar, presets ) + # Clean up presets from both config objects and dict before syncing + _cleanup_presets(env_cfg) + _cleanup_presets(agent_cfg) + _cleanup_presets_dict(hydra_cfg) # Sync dict -> config objects env_cfg.from_dict(hydra_cfg["env"]) env_cfg = replace_strings_with_env_cfg_spaces(env_cfg) @@ -327,6 +404,41 @@ def apply_overrides( _setattr(hydra_cfg, full_path, val) +def _cleanup_presets(cfg): + """Recursively remove ``presets`` fields from a config tree after resolution. + + Once all presets have been selected and applied, the ``presets`` metadata + is no longer needed and would confuse downstream managers that expect every + field to be a term config. + """ + if cfg is None: + return + if hasattr(cfg, "presets"): + try: + delattr(cfg, "presets") + except Exception: + pass + for name in list(getattr(cfg, "__dataclass_fields__", {}).keys()): + if name == "presets": + continue + try: + value = getattr(cfg, name) + except Exception: + continue + if hasattr(value, "__dataclass_fields__"): + _cleanup_presets(value) + + +def _cleanup_presets_dict(d: dict): + """Recursively remove ``presets`` keys from a nested dict.""" + if not isinstance(d, dict): + return + d.pop("presets", None) + for v in d.values(): + if isinstance(v, dict): + _cleanup_presets_dict(v) + + def _setattr(obj, path: str, val): """Set nested attribute/key (e.g., "actions.arm_action.scale").""" *parts, leaf = path.split(".") diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py index 0002c5d58d9..c340181e866 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/parse_cfg.py @@ -5,16 +5,20 @@ """Sub-module with utilities for parsing and loading configurations.""" +from __future__ import annotations + import collections import importlib import inspect import os import re +from typing import TYPE_CHECKING import gymnasium as gym import yaml -from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg +if TYPE_CHECKING: + from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg def load_cfg_from_registry(task_name: str, entry_point_key: str) -> dict | object: diff --git a/source/isaaclab_tasks/test/benchmarking/env_benchmark_test_utils.py b/source/isaaclab_tasks/test/benchmarking/env_benchmark_test_utils.py index 89d32f844c5..f8373b8f15a 100644 --- a/source/isaaclab_tasks/test/benchmarking/env_benchmark_test_utils.py +++ b/source/isaaclab_tasks/test/benchmarking/env_benchmark_test_utils.py @@ -11,7 +11,7 @@ import numpy as np import yaml - +from isaaclab.utils.version import get_isaac_sim_version, has_kit def _get_repo_path(): """Get the repository root by searching for marker files. @@ -176,9 +176,8 @@ def _retrieve_logs(workflow, task): repo_path = _get_repo_path() # Defer Isaac Sim version import to avoid preloading USD before SimulationApp starts. - from isaaclab.utils.version import get_isaac_sim_version - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: repo_path = os.path.join(repo_path, "..") if workflow == "rl_games": log_files_path = os.path.join(repo_path, f"logs/{workflow}/{task}/*/summaries/*") diff --git a/source/isaaclab_tasks/test/env_test_utils.py b/source/isaaclab_tasks/test/env_test_utils.py index 7dc0008360a..c8ad07e8302 100644 --- a/source/isaaclab_tasks/test/env_test_utils.py +++ b/source/isaaclab_tasks/test/env_test_utils.py @@ -16,7 +16,7 @@ from isaaclab.app.settings_manager import get_settings_manager from isaaclab.envs.utils.spaces import sample_space from isaaclab.sim import SimulationContext -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit from isaaclab_tasks.utils.parse_cfg import parse_env_cfg @@ -111,7 +111,7 @@ def _run_environments( """ # skip test if stage in memory is not supported - if get_isaac_sim_version().major < 5 and create_stage_in_memory: + if has_kit() and get_isaac_sim_version().major < 5 and create_stage_in_memory: pytest.skip("Stage in memory is not supported in this version of Isaac Sim") # skip suction gripper environments as they require CPU simulation and cannot be run with GPU simulation diff --git a/source/isaaclab_tasks/test/test_config_load_no_pxr.py b/source/isaaclab_tasks/test/test_config_load_no_pxr.py new file mode 100644 index 00000000000..6e60d9f868d --- /dev/null +++ b/source/isaaclab_tasks/test/test_config_load_no_pxr.py @@ -0,0 +1,148 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Test that config loading does not import modules that crash Kit's startup. + +Config loading (``load_cfg_from_registry``) runs BEFORE SimulationApp to +inspect the physics backend and decide whether Kit is needed at all. +Config classes are pure data — they must be constructable without any +runtime dependencies that conflict with Kit's internal ``fork()``. + +Forbidden categories: + 1. Backend / simulator modules (pxr, omni, carb, isaacsim) — require + SimulationApp / Kit to be initialized first. + 2. SciPy — loads OpenBLAS which registers atfork handlers that crash + Kit's internal fork() during startup. + +The only allowed mechanism for deferring heavy imports is ``lazy_loader`` +(``lazy.attach``). No other lazy-load hacks (manual ``__getattr__``, +wrapper callables, etc.) should be introduced. +""" + +import subprocess +import sys +import textwrap + +import gymnasium +import pytest + +import isaaclab_tasks # noqa: F401 -- triggers task registration + +# Forbidden module prefixes -- these must NOT appear in sys.modules after +# config loading because they require SimulationApp / a specific physics +# backend to be started first, or because they are heavyweight runtime +# libraries that should never be needed to construct pure-data config objects. +_FORBIDDEN_PREFIXES = ( + # Backend / simulator runtime (require SimulationApp / Kit) + "pxr", # USD Python bindings + "omni", # Omniverse runtime + "carb", # Carbonite framework + "isaacsim", # Isaac Sim modules + # SciPy loads OpenBLAS which crashes Kit's fork() + "scipy", +) + +_ALL_ISAAC_TASKS = sorted(name for name in gymnasium.registry if name.startswith("Isaac-")) + + +def _build_check_script(task_name: str) -> str: + """Return a self-contained Python script that loads a task config and + checks for forbidden imports. + + Uses ``__builtins__.__import__`` hook which intercepts every single + import at the lowest level — more reliable than ``sys.meta_path`` + finders which can be bypassed by C extensions, cached modules, or + custom loaders. + + The hook does NOT block or stub imports. It lets the real import + proceed and simply records the first occurrence of each forbidden + top-level module together with the stack trace that triggered it. + This way the full import chain is preserved for diagnostics without + altering runtime behaviour. + """ + return textwrap.dedent(f"""\ + import sys + import traceback + + _FORBIDDEN = {_FORBIDDEN_PREFIXES!r} + _violations = {{}} # top_module -> stack trace (first occurrence only) + _original_import = __builtins__.__import__ + + def _tracing_import(name, *args, **kwargs): + top = name.split(".")[0] + if top in _FORBIDDEN and top not in _violations: + _violations[top] = traceback.format_stack() + return _original_import(name, *args, **kwargs) + + __builtins__.__import__ = _tracing_import + + # ---- load the task config ---- + _load_error = None + try: + import isaaclab_tasks # noqa: F401 + from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + cfg = load_cfg_from_registry("{task_name}", "env_cfg_entry_point") + except Exception as exc: + _load_error = str(exc) + + # ---- report ---- + if _load_error: + print(f"ERROR: config load crashed: {{_load_error}}", file=sys.stderr) + + if _violations: + print(f"FAIL: {{len(_violations)}} forbidden top-level module(s) imported:", file=sys.stderr) + for top_mod, stack_frames in sorted(_violations.items()): + print(f"\\n=== {{top_mod}} ===", file=sys.stderr) + # show only the interesting frames (skip importlib internals) + for frame in stack_frames: + if "importlib" not in frame and " tuple[int, str, str]: + """Run the check in a subprocess so the current process stays clean.""" + result = subprocess.run( + [sys.executable, "-c", _build_check_script(task_name)], + capture_output=True, + text=True, + timeout=120, + ) + return result.returncode, result.stdout, result.stderr + + +@pytest.mark.parametrize("task_name", _ALL_ISAAC_TASKS) +def test_config_load_does_not_import_backend_modules(task_name: str): + """Config loading must not import forbidden runtime modules. + + Config classes are pure data. They must not pull in backend modules + (pxr, omni, carb, isaacsim), heavyweight compute libraries (torch, + scipy, numpy, warp), or I/O libraries (h5py, trimesh, PIL, lxml, rtree). + + Fix: use lazy_loader (lazy.attach) in __init__.py files, TYPE_CHECKING + guards, and string references / DeferredClass for implementation classes. + """ + rc, stdout, stderr = _run_config_load_check(task_name) + assert rc == 0, ( + f"Config loading for '{task_name}' imported forbidden backend modules.\n" + f"Forbidden prefixes: {_FORBIDDEN_PREFIXES}\n" + f"The import chain(s) for each violation are shown below.\n" + f"--- stderr ---\n{stderr}\n--- stdout ---\n{stdout}\n" + "Fix: use lazy_loader (lazy.attach) in the offending __init__.py, " + "or move the import under TYPE_CHECKING and use a string reference / " + "_is_instance_by_name() for isinstance checks." + ) diff --git a/source/isaaclab_tasks/test/test_environments_with_stage_in_memory.py b/source/isaaclab_tasks/test/test_environments_with_stage_in_memory.py index 8a99436b91c..2a87b7e9586 100644 --- a/source/isaaclab_tasks/test/test_environments_with_stage_in_memory.py +++ b/source/isaaclab_tasks/test/test_environments_with_stage_in_memory.py @@ -19,7 +19,7 @@ app_launcher = AppLauncher(headless=True, enable_cameras=True) simulation_app = app_launcher.app -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit """Rest everything follows.""" @@ -51,7 +51,7 @@ @pytest.mark.parametrize("task_name", setup_environment(include_play=False, factory_envs=False, multi_agent=False)) def test_environments_with_stage_in_memory_and_clone_in_fabric_disabled(task_name, num_envs, device): # skip test if stage in memory is not supported - if get_isaac_sim_version().major < 5: + if has_kit() and get_isaac_sim_version().major < 5: pytest.skip("Stage in memory is not supported in this version of Isaac Sim") # run environments with stage in memory diff --git a/source/isaaclab_tasks/test/test_hydra.py b/source/isaaclab_tasks/test/test_hydra.py index f64847ed26f..a911c15c881 100644 --- a/source/isaaclab_tasks/test/test_hydra.py +++ b/source/isaaclab_tasks/test/test_hydra.py @@ -23,6 +23,8 @@ from isaaclab.utils import configclass from isaaclab_tasks.utils.hydra import ( + _cleanup_presets, + _cleanup_presets_dict, _parse_val, _setattr, apply_overrides, @@ -452,3 +454,79 @@ def test_override_order_left_right_scalar_wins(test_configs): apply_overrides(env_cfg, agent_cfg, hydra_cfg, [], preset_sel, preset_scalar, presets) assert env_cfg.actions.arm_action.scale == 2.0 + + +# ============================================================================= +# Tests for preset cleanup +# ============================================================================= + + +def test_cleanup_presets_removes_from_config(test_configs): + """After cleanup, no config object in the tree should have a 'presets' attribute.""" + env_cfg, agent_cfg, presets = test_configs + + # Verify presets exist before cleanup + assert hasattr(env_cfg.observations, "presets") + assert hasattr(env_cfg.actions, "presets") + assert hasattr(env_cfg.actions.arm_action, "presets") + assert hasattr(agent_cfg.policy, "presets") + + _cleanup_presets(env_cfg) + _cleanup_presets(agent_cfg) + + # All presets should be gone + assert not hasattr(env_cfg.observations, "presets") + assert not hasattr(env_cfg.actions, "presets") + assert not hasattr(env_cfg.actions.arm_action, "presets") + assert not hasattr(agent_cfg.policy, "presets") + + # Non-preset fields should be untouched + assert env_cfg.decimation == 4 + assert env_cfg.observations.enable_corruption is True + assert agent_cfg.policy.actor_hidden_dims == [256, 128] + + +def test_cleanup_presets_dict_removes_from_dict(test_configs): + """After cleanup, no nested dict should contain a 'presets' key.""" + env_cfg, agent_cfg, _ = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + # Verify presets exist in dict before cleanup + assert "presets" in hydra_cfg["env"]["observations"] + assert "presets" in hydra_cfg["env"]["actions"] + + _cleanup_presets_dict(hydra_cfg) + + # Walk the entire dict tree to ensure no presets remain + def _check_no_presets(d, path=""): + if isinstance(d, dict): + assert "presets" not in d, f"Found 'presets' key at {path}" + for k, v in d.items(): + _check_no_presets(v, f"{path}.{k}") + + _check_no_presets(hydra_cfg) + + +def test_cleanup_after_apply_overrides(test_configs): + """Full flow: apply overrides then cleanup — config is clean for managers.""" + env_cfg, agent_cfg, presets = test_configs + hydra_cfg = {"env": env_cfg.to_dict(), "agent": agent_cfg.to_dict()} + + # Apply a global preset + apply_overrides(env_cfg, agent_cfg, hydra_cfg, ["fast"], [], [], presets) + + # Cleanup + _cleanup_presets(env_cfg) + _cleanup_presets(agent_cfg) + _cleanup_presets_dict(hydra_cfg) + + # Preset was applied + assert env_cfg.observations.enable_corruption is False + assert agent_cfg.policy.actor_hidden_dims == [32, 16] + + # No presets remain anywhere + assert not hasattr(env_cfg.observations, "presets") + assert not hasattr(env_cfg.actions, "presets") + assert not hasattr(agent_cfg.policy, "presets") + assert "presets" not in hydra_cfg["env"].get("observations", {}) + assert "presets" not in hydra_cfg["agent"].get("policy", {}) diff --git a/source/isaaclab_teleop/isaaclab_teleop/deprecated/openxr/manus_vive.py b/source/isaaclab_teleop/isaaclab_teleop/deprecated/openxr/manus_vive.py index 0f76d31c7b6..08c1ce6b047 100644 --- a/source/isaaclab_teleop/isaaclab_teleop/deprecated/openxr/manus_vive.py +++ b/source/isaaclab_teleop/isaaclab_teleop/deprecated/openxr/manus_vive.py @@ -24,7 +24,7 @@ from isaaclab.devices.device_base import DeviceBase, DeviceCfg from isaaclab.devices.retargeter_base import RetargeterBase -from isaaclab.utils.version import get_isaac_sim_version +from isaaclab.utils.version import get_isaac_sim_version, has_kit from .common import HAND_JOINT_NAMES from .xr_cfg import XrCfg @@ -83,9 +83,10 @@ def __init__(self, cfg: ManusViveCfg, retargeters: list[RetargeterBase] | None = ) super().__init__(retargeters) # Enforce minimum Isaac Sim version (>= 5.1) - isaac_sim_version = get_isaac_sim_version() - if isaac_sim_version < version.parse("5.1"): - raise RuntimeError(f"ManusVive requires Isaac Sim >= 5.1. Detected version: '{isaac_sim_version}'.") + if has_kit(): + isaac_sim_version = get_isaac_sim_version() + if isaac_sim_version < version.parse("5.1"): + raise RuntimeError(f"ManusVive requires Isaac Sim >= 5.1. Detected version: '{isaac_sim_version}'.") self._xr_cfg = cfg.xr_cfg or XrCfg() self._additional_callbacks = dict() self._vc_subscription = ( diff --git a/source/isaaclab_teleop/isaaclab_teleop/deprecated/openxr/manus_vive_utils.py b/source/isaaclab_teleop/isaaclab_teleop/deprecated/openxr/manus_vive_utils.py index 97056b21ac6..f9248913d66 100644 --- a/source/isaaclab_teleop/isaaclab_teleop/deprecated/openxr/manus_vive_utils.py +++ b/source/isaaclab_teleop/isaaclab_teleop/deprecated/openxr/manus_vive_utils.py @@ -9,8 +9,6 @@ import numpy as np -from isaacsim.core.utils.extensions import enable_extension - # For testing purposes, we need to mock the XRCore XRCore, XRPoseValidityFlags = None, None @@ -62,6 +60,7 @@ class ManusViveIntegration: def __init__(self): + from isaacsim.core.utils.extensions import enable_extension enable_extension("isaacsim.xr.input_devices") from isaacsim.xr.input_devices.impl.manus_vive_integration import get_manus_vive_integration diff --git a/tools/template/templates/tasks/direct_multi-agent/env b/tools/template/templates/tasks/direct_multi-agent/env index eec2331722e..cecd3fbdbdf 100644 --- a/tools/template/templates/tasks/direct_multi-agent/env +++ b/tools/template/templates/tasks/direct_multi-agent/env @@ -141,7 +141,7 @@ class {{ task.classname }}Env(DirectMARLEnv): self.robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) - self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self.robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, joint_ids=None, env_ids=env_ids) @torch.jit.script diff --git a/tools/template/templates/tasks/direct_single-agent/env b/tools/template/templates/tasks/direct_single-agent/env index e6f47fd3366..fe5f6d731af 100644 --- a/tools/template/templates/tasks/direct_single-agent/env +++ b/tools/template/templates/tasks/direct_single-agent/env @@ -110,7 +110,7 @@ class {{ task.classname }}Env(DirectRLEnv): self.robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids) self.robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids) - self.robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids) + self.robot.write_joint_state_to_sim(position=joint_pos, velocity=joint_vel, joint_ids=None, env_ids=env_ids) @torch.jit.script