diff --git a/docs/source/how-to/cloning.rst b/docs/source/how-to/cloning.rst index 4377923bfa45..16d5bcece155 100644 --- a/docs/source/how-to/cloning.rst +++ b/docs/source/how-to/cloning.rst @@ -129,7 +129,10 @@ Clone Plans For one source row, passing ``sources``, ``destinations``, and ``mask`` by hand is simple. For heterogeneous scenes, the mapping is easier to build with -:func:`~isaaclab.cloner.make_clone_plan`. +:func:`~isaaclab.cloner.make_clone_plan`, which returns the raw flat components. Composing +those components into a :class:`~isaaclab.cloner.ClonePlan` together with the per-environment +pose buffer is the caller's responsibility — keeping pose authority on the side that owns the +buffer (typically the scene) avoids duplicating tensors. :class:`~isaaclab.cloner.ClonePlan` stores the same flat contract used by direct cloning: @@ -180,7 +183,7 @@ The plan maps those source rows to all environments: from isaaclab.cloner import make_clone_plan, sequential - plan = make_clone_plan( + sources, destinations, clone_mask = make_clone_plan( sources=[ [ "/World/envs/env_0/Object", @@ -196,12 +199,12 @@ The plan maps those source rows to all environments: # source row used by env: 0, 1, 2, 0, 1, 2, 0, 1 -Direct code can use the plan exactly like the hand-written direct example: +Direct code can use the components exactly like the hand-written direct example: .. code-block:: python - physx_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask, device="cuda:0") - usd_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask) + physx_replicate(stage, sources, destinations, env_ids, clone_mask, device="cuda:0") + usd_replicate(stage, sources, destinations, env_ids, clone_mask) When variants span multiple groups, such as robot variants and object variants, ``make_clone_plan`` enumerates the Cartesian product of the groups and assigns one diff --git a/source/isaaclab/changelog.d/newton-clone-plan.rst b/source/isaaclab/changelog.d/newton-clone-plan.rst new file mode 100644 index 000000000000..7248966e36fc --- /dev/null +++ b/source/isaaclab/changelog.d/newton-clone-plan.rst @@ -0,0 +1,6 @@ +Fixed +^^^^^ + +* Fixed Newton replicated-scene cloning so source clone plans are available + before sensor construction and asset USD replication is skipped for Newton + physics replication. diff --git a/source/isaaclab/changelog.d/string-compile-template-pattern.minor.rst b/source/isaaclab/changelog.d/string-compile-template-pattern.minor.rst new file mode 100644 index 000000000000..68dec4512ce3 --- /dev/null +++ b/source/isaaclab/changelog.d/string-compile-template-pattern.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added :func:`~isaaclab.utils.string.compile_template_pattern` to compile a slotted template + (with ``{}`` slots expanding to a caller-provided regex fragment) into a regular expression. diff --git a/source/isaaclab/changelog.d/string-strip-templated-prefix.minor.rst b/source/isaaclab/changelog.d/string-strip-templated-prefix.minor.rst new file mode 100644 index 000000000000..e9c8cf338894 --- /dev/null +++ b/source/isaaclab/changelog.d/string-strip-templated-prefix.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added :func:`~isaaclab.utils.string.strip_templated_prefix` to strip a templated prefix + (with ``{}`` slots matching one path segment each) from a string and return the remainder. diff --git a/source/isaaclab/isaaclab/cloner/__init__.pyi b/source/isaaclab/isaaclab/cloner/__init__.pyi index 1ee123e7cf56..f3fc2801c6d1 100644 --- a/source/isaaclab/isaaclab/cloner/__init__.pyi +++ b/source/isaaclab/isaaclab/cloner/__init__.pyi @@ -6,12 +6,14 @@ __all__ = [ "CloneCfg", "ClonePlan", - "random", - "sequential", + "cfg_source_path", "disabled_fabric_change_notifies", "filter_collisions", "grid_transforms", "make_clone_plan", + "path_source_path", + "random", + "sequential", "usd_replicate", ] @@ -19,9 +21,11 @@ from .clone_plan import ClonePlan from .cloner_cfg import CloneCfg from .cloner_strategies import random, sequential from .cloner_utils import ( + cfg_source_path, disabled_fabric_change_notifies, filter_collisions, grid_transforms, make_clone_plan, + path_source_path, usd_replicate, ) diff --git a/source/isaaclab/isaaclab/cloner/cloner_utils.py b/source/isaaclab/isaaclab/cloner/cloner_utils.py index 337fad42f45f..86d7232176f5 100644 --- a/source/isaaclab/isaaclab/cloner/cloner_utils.py +++ b/source/isaaclab/isaaclab/cloner/cloner_utils.py @@ -9,6 +9,7 @@ import itertools import logging import math +import re from collections.abc import Iterator, Sequence import torch @@ -21,6 +22,137 @@ logger = logging.getLogger(__name__) +def get_suffix(path_expr: str, destination_template: str) -> str | None: + """Return the part of ``path_expr`` below a destination template's env-instance root. + + The template's ``"{}"`` slot matches exactly one path segment (a concrete id like ``env_3`` + or a wildcard like ``env_.*``). + + Example: + >>> tmpl = "/World/scenes/{}/Robot" + >>> get_suffix("/World/scenes/env_3/Robot/base", tmpl) + '/base' + >>> get_suffix("/World/scenes/.*/Robot/base", tmpl) + '/base' + >>> get_suffix("/World/scenes/env_3/Robot", tmpl) + '' + >>> get_suffix("/World/scenes/env_3/Sensor", tmpl) is None + True + >>> get_suffix("/World/scenes/env_3/RobotArm", tmpl) is None + True + >>> get_suffix("/World/scenes/env_3/sub/Robot/base", tmpl) is None + True + """ + pattern = re.compile(r"[^/]+".join(re.escape(part) for part in destination_template.split("{}"))) + match = pattern.match(path_expr) + if match is None: + return None + suffix = path_expr[match.end() :] + return None if suffix and not suffix.startswith("/") else suffix + + +def resolve_clone_plan_source(path_expr: str, plan: ClonePlan) -> tuple[str, str, str] | None: + """Resolve a destination path expression to its row's source path, destination glob, and asset suffix. + + Finds the rows whose destination template owns ``path_expr`` (same matching + logic as :func:`iter_clone_plan_matches`), OR-merges their + :attr:`~isaaclab.cloner.ClonePlan.clone_mask` rows, and splits the + expression at the row's destination template so the asset-relative suffix is + returned for downstream walks. + + Args: + path_expr: Destination-side path expression (e.g., a sensor's ``prim_path``, + with ``.*`` env wildcard). + plan: Active clone plan. + + Returns: + Three-tuple of ``(source_asset_path, dest_glob_prefix, asset_suffix)``. The + ``asset_suffix`` is the part of ``path_expr`` beyond the matching row's + destination template (empty when ``path_expr`` equals the row's template). + Returns ``None`` when ``path_expr`` matches no row in the plan, letting + callers fall back to direct stage resolution (e.g. for sensor frames + mounted at the env root rather than under a planned asset). + + Raises: + ValueError: When ``path_expr``'s matching rows span multiple distinct + destination templates. + NotImplementedError: When the union of matching rows' clone masks does not + cover every env (partial-env heterogeneous coverage is unsupported). + """ + matching_template: str | None = None + matching_rows: list[int] = [] + matching_suffix: str | None = None + for source_index, destination_template in enumerate(plan.destinations): + if "{}" not in destination_template: + continue + suffix = get_suffix(path_expr, destination_template) + if suffix is None: + continue + if matching_template is None: + matching_template = destination_template + matching_suffix = suffix + elif destination_template != matching_template: + raise ValueError( + f"path_expr {path_expr!r}: matches multiple destination templates" + f" {matching_template!r} and {destination_template!r}." + ) + matching_rows.append(source_index) + if matching_template is None: + return None + if not plan.clone_mask[matching_rows].any(dim=0).all(): + raise NotImplementedError( + f"path_expr {path_expr!r}: partial-env heterogeneous coverage is unsupported;" + " matching rows must collectively cover all envs." + ) + return plan.sources[matching_rows[0]], matching_template.replace("{}", "*"), matching_suffix or "" + + +def iter_clone_plan_matches(plan: ClonePlan, path_expr: str) -> Iterator[tuple[str, str, str, tuple[int, ...]]]: + """Yield clone-plan entries whose destinations own a path expression. + + Example: + For an entry with source root ``"/World/source/Robot"``, destination + template ``"/World/scenes/{}/Robot"``, and populated env ids + ``(0, 2)``, querying ``"/World/scenes/.*/Robot/base"`` yields + ``("/World/source/Robot", "/World/scenes/{}/Robot", + "/World/source/Robot/base", (0, 2))``. + + Args: + plan: Clone plan to query. + path_expr: Destination prim path or path expression. Expressions are + matched against each clone-plan destination template by treating + the template's ``"{}"`` field as the populated environment slot. + + Yields: + Tuples ``(source_root, destination_template, source_path, env_ids)`` + for the nearest matching destination root. Multiple source variants + with the same destination root are preserved. + """ + matches: list[tuple[str, str, str, tuple[int, ...]]] = [] + for source_index, (source_root, destination_template) in enumerate(zip(plan.sources, plan.destinations)): + if "{}" not in destination_template: + continue + + env_ids = tuple(int(i) for i in plan.clone_mask[source_index].nonzero(as_tuple=False).flatten().tolist()) + if not env_ids: + continue + + source_root = source_root.rstrip("/") or "/" + destination_template = destination_template.rstrip("/") or "/" + + suffix = get_suffix(path_expr, destination_template) + if suffix is None: + continue + source_path = source_root + suffix if source_root != "/" else suffix or "/" + + matches.append((source_root, destination_template, source_path, env_ids)) + + matches.sort(key=lambda match: len(match[1].format(match[3][0])), reverse=True) + if matches: + owner_length = len(matches[0][1].format(matches[0][3][0])) + yield from (match for match in matches if len(match[1].format(match[3][0])) == owner_length) + + @contextlib.contextmanager def disabled_fabric_change_notifies(stage: Usd.Stage, *, restore: bool = True) -> Iterator[None]: """Suspend the ``IFabricUsd`` USD notice listener for the body of the ``with`` block. @@ -104,23 +236,26 @@ def make_clone_plan( num_clones: int, clone_strategy: callable, device: str = "cpu", -) -> ClonePlan: - """Construct a cloning plan mapping prototype prims to per-environment destinations. +) -> tuple[tuple[str, ...], tuple[str, ...], torch.Tensor]: + """Compute the flat source/destination/mask components of a clone plan. - The plan enumerates all combinations of prototypes, selects a combination per environment using ``clone_strategy``, - and builds a boolean masking matrix indicating which prototype populates each environment slot. + Enumerates all combinations of prototypes, selects a combination per environment using + ``clone_strategy``, and builds the boolean masking matrix that indicates which prototype + populates each environment slot. The caller composes the returned tuple into a + :class:`ClonePlan`. Args: - sources: Prototype prim paths grouped by asset type (e.g., [[robot_a, robot_b], [obj_x]]). + sources: Prototype prim paths grouped by asset type (e.g., ``[[robot_a, robot_b], [obj_x]]``). destinations: Destination path templates (one per group) with ``"{}"`` placeholder for env id. num_clones: Number of environments to populate. clone_strategy: Function that picks a prototype combo per environment; signature ``clone_strategy(combos: Tensor, num_clones: int, device: str) -> Tensor[num_clones, num_groups]``. - device: Torch device for tensors in the plan. Defaults to ``"cpu"``. + device: Torch device for the returned mask. Defaults to ``"cpu"``. Returns: - A :class:`ClonePlan` whose ``sources`` and ``destinations`` are flattened per-source rows and - whose ``clone_mask`` is a ``[num_src, num_clones]`` boolean tensor. + A tuple ``(sources, destinations, clone_mask)`` where ``sources`` and ``destinations`` + are flattened per-source entries (one entry per prototype) and ``clone_mask`` is a + ``[num_src, num_clones]`` boolean tensor on ``device``. """ if len(sources) != len(destinations): raise ValueError(f"Expected one destination per source group, got {len(destinations)} and {len(sources)}.") @@ -150,7 +285,7 @@ def make_clone_plan( masking = torch.zeros((sum(group_sizes), num_clones), dtype=torch.bool, device=device) masking[rows, cols] = True - return ClonePlan(sources=src, destinations=dest, clone_mask=masking) + return src, dest, masking def usd_replicate( diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index da013ba0a591..c1f6a4c855c6 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -20,7 +20,7 @@ from isaaclab.controllers.operational_space import OperationalSpaceController from isaaclab.managers.action_manager import ActionTerm from isaaclab.sensors import ContactSensor, ContactSensorCfg, FrameTransformer, FrameTransformerCfg -from isaaclab.sim.utils import find_matching_prims +from isaaclab.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source if TYPE_CHECKING: from isaaclab.envs import ManagerBasedEnv @@ -336,8 +336,21 @@ def __init__(self, cfg: actions_cfg.OperationalSpaceControllerActionCfg, env: Ma # is provided. if self.cfg.task_frame_rel_path is not None: # The source RigidObject can be any child of the articulation asset (we will not use it), - # hence, we will use the first RigidObject child. - root_rigidbody_path = self._first_RigidObject_child_path() + # hence, we will use the first RigidObject descendant. + def has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) + + matches = resolve_matching_prims_from_source(self._asset.cfg.prim_path) + if not matches: + raise ValueError(f"No prim found at '{self._asset.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString + rigid_prims = get_all_matching_child_prims( + walk_root, predicate=has_rigid_body_api, traverse_instance_prims=False + ) + if not rigid_prims: + raise ValueError(f"No descendant rigid body found under the expression: '{self._asset.cfg.prim_path}'.") + root_rigidbody_path = root_expr + rigid_prims[0].GetPath().pathString[len(walk_root) :] task_frame_transformer_path = "/World/envs/env_.*/" + self.cfg.task_frame_rel_path task_frame_transformer_cfg = FrameTransformerCfg( prim_path=root_rigidbody_path, @@ -560,29 +573,6 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: """ - def _first_RigidObject_child_path(self): - """Finds the first ``RigidObject`` child under the articulation asset. - - Raises: - ValueError: If no child ``RigidObject`` is found under the articulation asset. - - Returns: - str: The path to the first ``RigidObject`` child under the articulation asset. - """ - child_prims = find_matching_prims(self._asset.cfg.prim_path + "/.*") - rigid_child_prim = None - # Loop through the list and stop at the first RigidObject found - for prim in child_prims: - if prim.HasAPI(UsdPhysics.RigidBodyAPI): - rigid_child_prim = prim - break - if rigid_child_prim is None: - raise ValueError("No child rigid body found under the expression: '{self._asset.cfg.prim_path}'/.") - rigid_child_prim_path = rigid_child_prim.GetPath().pathString - # Remove the specific env index from the path string - rigid_child_prim_path = self._asset.cfg.prim_path + "/" + rigid_child_prim_path.split("/")[-1] - return rigid_child_prim_path - def _resolve_command_indexes(self): """Resolves the indexes for the various command elements within the command tensor. diff --git a/source/isaaclab/isaaclab/scene/interactive_scene.py b/source/isaaclab/isaaclab/scene/interactive_scene.py index 1d1afa2cfad9..a5621b444ace 100644 --- a/source/isaaclab/isaaclab/scene/interactive_scene.py +++ b/source/isaaclab/isaaclab/scene/interactive_scene.py @@ -38,6 +38,7 @@ from isaaclab.sim.utils.stage import get_current_stage, get_current_stage_id from isaaclab.sim.views import FrameView from isaaclab.terrains import TerrainImporter, TerrainImporterCfg +from isaaclab.utils.version import has_kit # Note: This is a temporary import for the VisuoTactileSensorCfg class. # It will be removed once the VisuoTactileSensor class is added to the core Isaac Lab framework. @@ -162,20 +163,14 @@ def __init__(self, cfg: InteractiveSceneCfg): self._physics_scene_path = None # prepare cloner for environment replication self.env_prim_paths = [f"{self.env_ns}/env_{i}" for i in range(self.cfg.num_envs)] + is_newton_replicated_scene = self.cfg.replicate_physics and self.physics_backend.startswith("newton") self.cloner_cfg = cloner.CloneCfg( clone_regex=self.env_regex_ns, clone_in_fabric=self.cfg.clone_in_fabric, device=self.device, physics_clone_fn=physics_clone_fn, - # USD replication runs for every backend. PhysX/Newton need per-env - # USD prims for sensor discovery. For OVPhysX, the per-env USD - # subtrees are layered on TOP of the physics-side ``physx.clone()`` - # replicas -- PhysX is indifferent to additional USD content and - # the two layers don't conflict. Probing whether this assumption - # holds in practice; revert to ``not startswith("ovphysx")`` if - # ``physx.clone()`` errors on already-populated targets. - clone_usd=True, + clone_usd=not is_newton_replicated_scene or has_kit(), ) # create source prim @@ -183,33 +178,29 @@ def __init__(self, cfg: InteractiveSceneCfg): self.env_fmt = self.env_regex_ns.replace(".*", "{}") # allocate env indices self._ALL_INDICES = torch.arange(self.cfg.num_envs, dtype=torch.long, device=self.device) - self._default_env_origins, _ = cloner.grid_transforms(self.num_envs, self.cfg.env_spacing, device=self.device) - # copy empty prim of env_0 to env_1, env_2, ..., env_{num_envs-1} with correct location. - # Suspend Fabric's USD notice listener: scene-init is followed by ``SimulationContext.reset``, - # which does the Fabric resync naturally — re-enabling here would just trigger a redundant batch. - # Note: ``restore=False`` means the listener stays disabled past this ``with`` block — through - # ``_add_entities_from_cfg`` and ``clone_environments`` below — until ``SimulationContext.reset`` - # re-enables it. The nested suspension inside ``clone_environments`` becomes a no-op as a result. + pos, quat = cloner.grid_transforms(self.num_envs, self.cfg.env_spacing, device=self.device) + self._default_env_pose = torch.cat([pos, quat], dim=-1) + + homo_mask = torch.ones((1, self.num_envs), device=self.device, dtype=torch.bool) + # Suspend Fabric's USD notice listener enable fast usd cloning with cloner.disabled_fabric_change_notifies(self.stage, restore=False): - cloner.usd_replicate( - self.stage, - [self.env_fmt.format(0)], - [self.env_fmt], - self._ALL_INDICES, - positions=self._default_env_origins, - ) + # copy empty prim of env_0 to env_1, env_2, ..., env_{num_envs-1} with correct location. + rep_args = (self.stage, [self.env_fmt.format(0)], [self.env_fmt], self._ALL_INDICES, homo_mask, pos, quat) + cloner.usd_replicate(*rep_args) self._global_prim_paths = list() has_scene_cfg_entities = self._is_scene_setup_from_cfg() if has_scene_cfg_entities: self._clone_plan = self._build_clone_plan_from_cfg() + self.sim.set_clone_plan(self._clone_plan) self._add_entities_from_cfg() else: self._clone_plan = cloner.ClonePlan( sources=(self.env_fmt.format(0),), destinations=(self.env_fmt,), - clone_mask=torch.ones((1, self.num_envs), device=self.device, dtype=torch.bool), + clone_mask=homo_mask, ) + self.sim.set_clone_plan(self._clone_plan) # Aggregate scene-data requirements from declared visualizers and constructed sensors, # then publish to ``SimulationContext`` so downstream providers (constructed later by @@ -250,7 +241,7 @@ def set_spawn_paths(spawn_cfg, paths: list[str | None]) -> None: ordered_items = [item for item in items if not isinstance(item[1], SensorBaseCfg)] ordered_items += [item for item in items if isinstance(item[1], SensorBaseCfg)] - # One group is one prim path template plus its spawn variants. + # One group is one cfg's prim path template plus its spawn variants. groups = [] for _, asset_cfg in ordered_items: cfgs = asset_cfg.rigid_objects.values() if isinstance(asset_cfg, RigidObjectCollectionCfg) else [asset_cfg] @@ -258,48 +249,47 @@ def set_spawn_paths(spawn_cfg, paths: list[str | None]) -> None: prim_path = cfg.prim_path.format(ENV_REGEX_NS=self.env_regex_ns) if not hasattr(cfg, "spawn") or cfg.spawn is None or self.env_ns not in prim_path: continue - if (count := num_variants(cfg.spawn)) <= 0: - raise ValueError(f"Spawner at '{prim_path}' must have at least one variant.") - groups.append((cfg.spawn, prim_path.replace(self.env_regex_ns, self.env_fmt), count)) + if (count := num_variants(cfg.spawn)) > 0: + groups.append((cfg, cfg.spawn, prim_path.replace(self.env_regex_ns, self.env_fmt), count)) if not groups: return None # Homogeneous scenes still spawn sources at env_0, but publish the simpler env-root plan. - if all(count == 1 for _, _, count in groups): - for spawn_cfg, destination, _ in groups: + if all(count == 1 for _, _, _, count in groups): + for _, spawn_cfg, destination, _ in groups: set_spawn_paths(spawn_cfg, [destination.format(0)]) + clone_mask = torch.ones((1, self.num_envs), device=self.device, dtype=torch.bool) return cloner.ClonePlan( sources=(self.env_fmt.format(0),), destinations=(self.env_fmt,), - clone_mask=torch.ones((1, self.num_envs), device=self.device, dtype=torch.bool), + clone_mask=clone_mask, ) - plan = cloner.make_clone_plan( - [[destination.format(i) for i in range(count)] for _, destination, count in groups], - [destination for _, destination, _ in groups], - self.num_envs, - self.cloner_cfg.clone_strategy, - self.device, + sources, destinations, clone_mask = cloner.make_clone_plan( + sources=[[destination.format(i) for i in range(count)] for _, _, destination, count in groups], + destinations=[destination for _, _, destination, _ in groups], + num_clones=self.num_envs, + clone_strategy=self.cloner_cfg.clone_strategy, + device=self.device, ) - # Move each planned source row to the first environment that actually uses it. - row = 0 - sources = list(plan.sources) - for spawn_cfg, destination, count in groups: - mask = plan.clone_mask[row : row + count] - env_ids = mask.to(torch.int).argmax(dim=1).tolist() - active = mask.any(dim=1).tolist() - paths = [destination.format(env_id) if is_active else None for env_id, is_active in zip(env_ids, active)] - for i, path in zip(range(row, row + count), paths): + # Move each planned source entry to the first environment that actually uses it. + source_start = 0 + sources = list(sources) + for cfg, spawn_cfg, destination, count in groups: + submask = clone_mask[source_start : source_start + count] + env_ids = submask.to(torch.int).argmax(dim=1).tolist() + active = submask.any(dim=1).tolist() + paths = [destination.format(eid) if a else None for eid, a in zip(env_ids, active)] + for offset, path in enumerate(paths): if path is not None: - sources[i] = path + sources[source_start + offset] = path set_spawn_paths(spawn_cfg, paths) - row += count + source_start += count - plan = cloner.ClonePlan(sources=tuple(sources), destinations=plan.destinations, clone_mask=plan.clone_mask) - logger.debug("Built heterogeneous ClonePlan with %d source rows.", len(plan.sources)) - return plan + logger.debug("Built heterogeneous ClonePlan with %d source entries.", len(sources)) + return cloner.ClonePlan(sources=tuple(sources), destinations=destinations, clone_mask=clone_mask) def clone_environments(self, copy_from_source: bool = False): """Creates clones of the environment ``/World/envs/env_0``. @@ -331,16 +321,12 @@ def clone_environments(self, copy_from_source: bool = False): self.cloner_cfg.physics_clone_fn( self.stage, *replicate_args, - positions=self._default_env_origins, + positions=self._default_env_pose[:, :3], device=self.cloner_cfg.device, ) if self.cloner_cfg.clone_usd: - is_env_root_plan = ( - len(plan.sources) == 1 - and plan.sources[0] == self.env_fmt.format(0) - and plan.destinations == (self.env_fmt,) - ) - usd_positions = self._default_env_origins if is_env_root_plan else None + is_env_root_plan = len(plan.sources) == 1 and plan.destinations == (self.env_fmt,) + usd_positions = self._default_env_pose[:, :3] if is_env_root_plan else None cloner.usd_replicate(self.stage, *replicate_args, positions=usd_positions) # Publish to ``SimulationContext`` (the canonical owner). The :attr:`clone_plan` @@ -514,7 +500,7 @@ def env_origins(self) -> torch.Tensor: if self._terrain is not None: return self._terrain.env_origins else: - return self._default_env_origins + return self._default_env_pose[:, :3] @property def terrain(self) -> TerrainImporter | None: diff --git a/source/isaaclab/isaaclab/sensors/camera/camera.py b/source/isaaclab/isaaclab/sensors/camera/camera.py index a26961c6fd75..d13a560600a5 100644 --- a/source/isaaclab/isaaclab/sensors/camera/camera.py +++ b/source/isaaclab/isaaclab/sensors/camera/camera.py @@ -13,7 +13,7 @@ import torch import warp as wp -from pxr import UsdGeom +from pxr import UsdGeom, UsdPhysics import isaaclab.sim as sim_utils import isaaclab.utils.sensors as sensor_utils @@ -155,7 +155,22 @@ def __init__(self, cfg: CameraCfg): rot_offset = rot_offset.squeeze(0).cpu().numpy() if self.cfg.spawn is not None and self.cfg.spawn.vertical_aperture is None: self.cfg.spawn.vertical_aperture = self.cfg.spawn.horizontal_aperture * self.cfg.height / self.cfg.width - self._resolve_and_spawn("camera", translation=self.cfg.offset.pos, orientation=rot_offset) + # Resolve the camera prim path and spawn it, redirecting to a child if prim_path is a physics body. + spawn = self.cfg.spawn + if spawn is not None: + probe_path = (spawn.spawn_path or self.cfg.prim_path) if spawn is not None else self.cfg.prim_path + probe_matches = sim_utils.resolve_matching_prims_from_source(probe_path) + source_prim, _source_destination_expr = probe_matches[0] if probe_matches else (None, None) + if source_prim is not None and source_prim.IsValid(): + if source_prim.HasAPI(UsdPhysics.ArticulationRootAPI) or source_prim.HasAPI(UsdPhysics.RigidBodyAPI): + logger.info(f" Spawning camera at '{self.cfg.prim_path}/camera'.") + self.cfg.prim_path = spawn.spawn_path = f"{self.cfg.prim_path}/camera" + + spawn_target = spawn.spawn_path or self.cfg.prim_path + if sim_utils.find_first_matching_prim(spawn_target) is None: + spawn.func(spawn_target, spawn, translation=self.cfg.offset.pos, orientation=rot_offset) + if not sim_utils.find_matching_prims(spawn_target): + raise RuntimeError(f"Could not find prim with path {spawn_target!r}.") # An ISP (any ``isp_cfg`` other than ``None``) requires the HDR AOV; # an explicit ``"rgb_hdr"`` in ``data_types`` also requires the @@ -483,9 +498,9 @@ def _initialize_impl(self): render_spec = CameraRenderSpec( cfg=self.cfg, device=device_str, - num_instances=len(cam_paths), + num_instances=self._num_envs, camera_prim_paths=cam_paths, - view_count=len(cam_paths), + view_count=self._num_envs, camera_path_relative_to_env_0=rel_under_env0, ) @@ -512,8 +527,12 @@ def _initialize_impl(self): # Create frame count buffer self._frame = ProxyArray(wp.zeros(self._view.count, device=self._device, dtype=wp.int64)) - # Convert all encapsulated prims to Camera - for cam_prim in self._view.prims: + # Convert all encapsulated prims to Camera. Newton keeps only source USD camera prims. + self._sensor_prims.clear() + view_prims = list(self._view.prims) + if not view_prims and cam_paths: + view_prims = [self.stage.GetPrimAtPath(cam_paths[0])] * self._view.count + for cam_prim in view_prims: # Obtain the prim path cam_prim_path = cam_prim.GetPath().pathString # Check if prim is a camera diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster.py index 43920af58338..25a47a91c1ed 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_multi_mesh_ray_caster.py @@ -6,16 +6,16 @@ from __future__ import annotations import logging -import re from typing import TYPE_CHECKING import numpy as np import trimesh import warp as wp -from pxr import UsdPhysics +from pxr import Usd, UsdPhysics import isaaclab.sim as sim_utils +from isaaclab.cloner.cloner_utils import iter_clone_plan_matches from isaaclab.sim.simulation_context import SimulationContext from isaaclab.utils.mesh import PRIMITIVE_MESH_TYPES, create_trimesh_from_geom_mesh, create_trimesh_from_geom_shape from isaaclab.utils.warp import ProxyArray, convert_to_warp_mesh @@ -206,62 +206,62 @@ def _build_mesh_records( records_per_env = [[] for _ in range(self._num_envs)] target_in_plan = False tracked_target_exprs: list[str] = [target_cfg.prim_expr] - + has_rigid_body_api = lambda p: p.HasAPI(UsdPhysics.RigidBodyAPI) # noqa: E731 # Prefer ClonePlan data for env-scoped targets; destination USD prims may not exist. if plan is not None and target_cfg.track_mesh_transforms: - target_path = re.sub(r"env_\.\*", "env_0", target_cfg.prim_expr) plan_tracked_target_exprs: list[str] = [] - for row, (source_root, destination_template) in enumerate(zip(plan.sources, plan.destinations)): - if "{}" not in destination_template: - continue - - dest_path = destination_template.format(0) - suffix = target_path.removeprefix(dest_path) - if suffix == target_path or (suffix and not suffix.startswith("/")): - continue - + prim_expr = target_cfg.prim_expr + for source_root, destination_template, source_path, env_ids in iter_clone_plan_matches(plan, prim_expr): target_in_plan = True - env_ids = plan.clone_mask[row].nonzero(as_tuple=False).squeeze(-1) - if env_ids.numel() == 0: - continue - # Load meshes from the authored source row. - source_prims = sim_utils.find_matching_prims(source_root + suffix) + # Load meshes from the authored source entry. + source_prims = sim_utils.find_matching_prims(source_path) if not source_prims: - raise RuntimeError(f"No ClonePlan source prims matched '{source_root + suffix}'.") + raise RuntimeError(f"No ClonePlan source prims matched '{source_path}'.") mesh_ids: list[int] = [] row_tracked_target_exprs: list[str] = [] for source_prim in source_prims: - owner_prim = source_prim - while owner_prim and owner_prim.IsValid() and str(owner_prim.GetPath()) != "/": - if owner_prim.HasAPI(UsdPhysics.RigidBodyAPI): - break - owner_prim = owner_prim.GetParent() - if owner_prim is None or not owner_prim.IsValid() or not owner_prim.HasAPI(UsdPhysics.RigidBodyAPI): - raise RuntimeError( - f"Cannot track ClonePlan target '{target_cfg.prim_expr}' because source prim " - f"'{source_prim.GetPath()}' has no rigid-body ancestor." - ) - mesh_id = self._load_target_prim_warp_mesh(source_prim, target_cfg, reference_prim=owner_prim) - dummy_mesh_id = mesh_id if dummy_mesh_id is None else dummy_mesh_id - mesh_ids.append(mesh_id) - owner_path = str(owner_prim.GetPath()) - if owner_path == source_root: - owner_suffix = "" - elif owner_path.startswith(source_root + "/"): - owner_suffix = owner_path[len(source_root) :] + source_prim_path = str(source_prim.GetPath()) + # Use a bounded rigid-body ancestor when the match is below a body; otherwise + # enumerate rigid bodies under the match, including the matched prim itself. + owner_prim = None + rigid_body_records: list[tuple[Usd.Prim, Usd.Prim]] + if not source_prim.HasAPI(UsdPhysics.ArticulationRootAPI): + owner_prim = sim_utils.get_first_matching_ancestor_prim(source_prim_path, has_rigid_body_api) + owner_path = None if owner_prim is None else str(owner_prim.GetPath()) + if owner_path and (owner_path == source_root or owner_path.startswith(f"{source_root}/")): + rigid_body_records = [(source_prim, owner_prim)] else: + rigid_body_records = [ + (p, p) for p in sim_utils.get_all_matching_child_prims(source_prim_path, has_rigid_body_api) + ] + if not rigid_body_records: raise RuntimeError( - f"Tracked target owner '{owner_path}' is not under ClonePlan source root '{source_root}'." + f"Cannot track ClonePlan target '{target_cfg.prim_expr}' because source prim " + f"'{source_prim.GetPath()}' has no rigid-body ancestor or descendant." ) - row_tracked_target_exprs.append(destination_template.replace("{}", ".*") + owner_suffix) + for geometry_prim, owner_prim in rigid_body_records: + mesh_id = self._load_target_prim_warp_mesh(geometry_prim, target_cfg, reference_prim=owner_prim) + dummy_mesh_id = mesh_id if dummy_mesh_id is None else dummy_mesh_id + mesh_ids.append(mesh_id) + owner_path = str(owner_prim.GetPath()) + if owner_path == source_root: + owner_suffix = "" + elif owner_path.startswith(source_root + "/"): + owner_suffix = owner_path[len(source_root) :] + else: + raise RuntimeError( + f"Tracked target owner '{owner_path}' is not under ClonePlan source root " + f"'{source_root}'." + ) + row_tracked_target_exprs.append(destination_template.format(".*") + owner_suffix) if len(row_tracked_target_exprs) > len(plan_tracked_target_exprs): plan_tracked_target_exprs = row_tracked_target_exprs # Geometry is selected by ClonePlan; live pose is supplied by backend body/site views. - for env_id in env_ids.tolist(): + for env_id in env_ids: for mesh_id in mesh_ids: records_per_env[env_id].append((mesh_id, (1.0e9, 1.0e9, 1.0e9), (0.0, 0.0, 0.0, 1.0))) @@ -286,11 +286,7 @@ def _build_mesh_records( if reference_prim.HasAPI(UsdPhysics.RigidBodyAPI): break reference_prim = reference_prim.GetParent() - if ( - reference_prim is None - or not reference_prim.IsValid() - or not reference_prim.HasAPI(UsdPhysics.RigidBodyAPI) - ): + if reference_prim is None or not reference_prim.IsValid() or not has_rigid_body_api(reference_prim): raise RuntimeError( f"Cannot track non-physics ray-cast target '{target_cfg.prim_expr}'. " "Set track_mesh_transforms=False for static targets, or apply RigidBodyAPI to dynamic targets." @@ -299,10 +295,7 @@ def _build_mesh_records( mesh_id = self._load_target_prim_warp_mesh(target_prim, target_cfg, reference_prim=reference_prim) dummy_mesh_id = mesh_id if dummy_mesh_id is None else dummy_mesh_id - pos, quat = sim_utils.resolve_prim_pose(reference_prim) - pos = (float(pos[0]), float(pos[1]), float(pos[2])) - quat = (float(quat[0]), float(quat[1]), float(quat[2]), float(quat[3])) - records.append((mesh_id, pos, quat)) + records.append((mesh_id, *sim_utils.resolve_prim_pose(reference_prim))) if len(records) == 1: return [list(records) for _ in range(self._num_envs)], dummy_mesh_id, tracked_target_exprs diff --git a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py index 27af7fa8b0be..b9b6aa8b7bc6 100644 --- a/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py +++ b/source/isaaclab/isaaclab/sensors/ray_caster/base_ray_caster.py @@ -69,9 +69,6 @@ def __init__(self, cfg: RayCasterCfg): """ BaseRayCaster._instance_count += 1 super().__init__(cfg) - # Resolve physics-body paths and spawn the sensor Xform child if needed. - self._requested_prim_path = self.cfg.prim_path - self._resolve_and_spawn("raycaster") self._data = RayCasterData() def __str__(self) -> str: diff --git a/source/isaaclab/isaaclab/sensors/sensor_base.py b/source/isaaclab/isaaclab/sensors/sensor_base.py index 15fca5ee4ad4..f2d6e3fbac1c 100644 --- a/source/isaaclab/isaaclab/sensors/sensor_base.py +++ b/source/isaaclab/isaaclab/sensors/sensor_base.py @@ -21,12 +21,19 @@ import warp as wp +from pxr import UsdPhysics + import isaaclab.sim as sim_utils +from isaaclab.cloner.cloner_utils import iter_clone_plan_matches from isaaclab.physics import PhysicsEvent, PhysicsManager +from isaaclab.sim.utils.queries import get_first_matching_ancestor_prim +from isaaclab.sim.utils.transforms import resolve_prim_pose from .kernels import reset_envs_kernel, update_outdated_envs_kernel, update_timestamp_kernel if TYPE_CHECKING: + from isaaclab.cloner import ClonePlan + from .sensor_base_cfg import SensorBaseCfg logger = logging.getLogger(__name__) @@ -58,6 +65,8 @@ def __init__(self, cfg: SensorBaseCfg): self._is_initialized = False # flag for whether the sensor is in visualization mode self._is_visualizing = False + # clone plan used for this sensor's latest initialization + self._clone_plan: ClonePlan | None = None self.stage = sim_utils.get_current_stage() # register various callback functions @@ -216,10 +225,19 @@ def _initialize_impl(self): 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) + # Count number of environments. + self._clone_plan = sim.get_clone_plan() + clone_plan = self._clone_plan + clone_plan_matches = () + if clone_plan is not None: + clone_plan_matches = tuple(iter_clone_plan_matches(clone_plan, self.cfg.prim_path)) + if clone_plan_matches: + self._parent_prims = [] + self._num_envs = int(clone_plan.clone_mask.shape[1]) + else: + 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) # Create warp env mask arrays for "all envs" cases and resets. # Note: We use wp.to_torch() to create zero-copy torch tensor views of warp arrays. # This allows warp arrays to be passed to warp kernels while the corresponding torch @@ -322,6 +340,7 @@ def _initialize_callback(self, event): def _invalidate_initialize_callback(self, event): """Invalidates the scene elements.""" self._is_initialized = False + self._clone_plan = None sim_ctx = sim_utils.SimulationContext.instance() if sim_ctx is not None: sim_ctx.vis_marker_registry.clear_debug_vis_callback(self) @@ -393,70 +412,63 @@ def _resolve_indices_and_mask( self._reset_mask_torch[env_ids] = True return self._reset_mask - def _resolve_and_spawn(self, sensor_name: str, **spawn_kwargs) -> None: - """Resolve physics-body prim paths and spawn the sensor prim if needed. - - Behavior matrix (``spawn`` refers to ``cfg.spawn``): - - +----------------+------------------+--------------------------------------------+ - | ``spawn`` | ``prim_path`` | Action | - +================+==================+============================================+ - | not ``None`` | physics body | Append ``/``, spawn child. | - +----------------+------------------+--------------------------------------------+ - | not ``None`` | non-physics prim | Use existing prim, skip spawn. | - | | (already exists) | | - +----------------+------------------+--------------------------------------------+ - | not ``None`` | does not exist | Spawn prim at ``prim_path``. | - +----------------+------------------+--------------------------------------------+ - | ``None`` | physics body | Raise ``ValueError``. | - +----------------+------------------+--------------------------------------------+ - | ``None`` | non-physics prim | Use as-is (no spawn). | - +----------------+------------------+--------------------------------------------+ - - Args: - sensor_name: Short identifier (e.g. ``"raycaster"``, ``"camera"``). - **spawn_kwargs: Extra keyword arguments forwarded to ``cfg.spawn.func`` - (e.g. ``translation``, ``orientation``). + def _resolve_rigid_body_ancestor_expr( + self, + ) -> tuple[str, tuple[float, float, float] | None, tuple[float, float, float, float] | None]: + """Resolve the rigid-body ancestor view expression and the sensor-to-body offset. + + The sensor's :attr:`SensorBaseCfg.prim_path` may point to any frame + inside the asset. To create a physics view, this helper walks ancestors + from that prim until it finds one with ``UsdPhysics.RigidBodyAPI``, + builds the corresponding destination-side expression, and computes the + fixed transform from that body to the configured sensor frame. + + Combines two resolution paths: + + 1. When an active :class:`~isaaclab.cloner.ClonePlan` exists, the + source-side env path is taken from the plan via + :func:`~isaaclab.cloner.resolve_clone_plan_source`, the rigid-body ancestor + is located on that source env, and the destination expression is + reconstructed by trimming the sensor-relative suffix from the plan's + destination glob. + 2. Otherwise (stage scan fallback for non-cloned setups), the first + matching env is located via + :func:`~isaaclab.sim.utils.queries.find_first_matching_prim`, the + rigid-body ancestor is located on that env, and the destination + expression is the configured :attr:`SensorBaseCfg.prim_path` minus + the sensor-relative suffix. + + The returned expression may still contain regex-style wildcards (e.g. + ``.*``); callers are responsible for converting to glob form for their + physics view (e.g. ``.replace(".*", "*")``). - Raises: - ValueError: If ``spawn`` is ``None`` and ``prim_path`` is a physics body. - RuntimeError: If the prim does not exist after the spawn attempt. + Returns: + A tuple of: + + * ``rigid_parent_expr``: destination-side view expression that + matches the rigid-body ancestor across envs. + * ``fixed_pos_b``: sensor-relative-to-body translation [m] (xyz), + or ``None`` when the sensor is mounted directly at the body + origin. + * ``fixed_quat_b``: sensor-relative-to-body rotation as a + quaternion ``(x, y, z, w)``, or ``None`` when the sensor is + mounted directly at the body origin. """ - from pxr import UsdPhysics # noqa: PLC0415 - - spawn = getattr(self.cfg, "spawn", None) - has_spawn = spawn is not None - - # Determine the path to probe for physics-body redirect - spawn_path = (getattr(spawn, "spawn_path", None) or self.cfg.prim_path) if has_spawn else None - probe_path = spawn_path if spawn_path is not None else self.cfg.prim_path - - prim = sim_utils.find_first_matching_prim(probe_path) - if prim is not None and prim.IsValid(): - is_physics = prim.HasAPI(UsdPhysics.ArticulationRootAPI) or prim.HasAPI(UsdPhysics.RigidBodyAPI) - if is_physics: - if not has_spawn: - raise ValueError( - f"Sensor prim_path '{self.cfg.prim_path}' resolves to a physics body but" - f" no spawner is configured (spawn=None). Either set spawn or point" - f" prim_path at a non-physics child (e.g. '{self.cfg.prim_path}/{sensor_name}')." - ) - logger.info( - f"Sensor prim_path '{self.cfg.prim_path}' points at a physics body." - f" Redirecting to '{self.cfg.prim_path}/{sensor_name}'." - ) - self.cfg.prim_path = f"{self.cfg.prim_path}/{sensor_name}" - if getattr(spawn, "spawn_path", None) is not None: - spawn.spawn_path = f"{spawn.spawn_path}/{sensor_name}" - - if not has_spawn: - return + matches = sim_utils.resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + prim, target_expr = matches[0] + + ancestor_prim = get_first_matching_ancestor_prim( + prim.GetPath(), predicate=lambda _prim: _prim.HasAPI(UsdPhysics.RigidBodyAPI) + ) + if ancestor_prim is None: + raise RuntimeError(f"Failed to find a rigid body ancestor prim at path expression: {self.cfg.prim_path}") - spawn_target = getattr(spawn, "spawn_path", None) or self.cfg.prim_path - prim = sim_utils.find_first_matching_prim(spawn_target) - if prim is None or not prim.IsValid(): - spawn.func(spawn_target, spawn, **spawn_kwargs) + if ancestor_prim == prim: + return target_expr, None, None - check_path = getattr(spawn, "spawn_path", None) or self.cfg.prim_path - if len(sim_utils.find_matching_prims(check_path)) == 0: - raise RuntimeError(f"Could not find prim with path {check_path!r}.") + relative_path = prim.GetPath().MakeRelativePath(ancestor_prim.GetPath()).pathString + rigid_parent_expr = target_expr.replace("/" + relative_path, "") + fixed_pos_b, fixed_quat_b = resolve_prim_pose(prim, ancestor_prim) + return rigid_parent_expr, fixed_pos_b, fixed_quat_b diff --git a/source/isaaclab/isaaclab/sim/__init__.pyi b/source/isaaclab/isaaclab/sim/__init__.pyi index 0c787cc64c67..a18eb04c56fd 100644 --- a/source/isaaclab/isaaclab/sim/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/__init__.pyi @@ -159,6 +159,7 @@ __all__ = [ "get_all_matching_child_prims", "find_first_matching_prim", "find_matching_prims", + "resolve_matching_prims_from_source", "find_matching_prim_paths", "find_global_fixed_joint_prim", "add_labels", @@ -351,6 +352,7 @@ from .utils import ( get_all_matching_child_prims, find_first_matching_prim, find_matching_prims, + resolve_matching_prims_from_source, find_matching_prim_paths, find_global_fixed_joint_prim, add_labels, diff --git a/source/isaaclab/isaaclab/sim/utils/__init__.pyi b/source/isaaclab/isaaclab/sim/utils/__init__.pyi index de4ab9c93266..5927f93c0a7f 100644 --- a/source/isaaclab/isaaclab/sim/utils/__init__.pyi +++ b/source/isaaclab/isaaclab/sim/utils/__init__.pyi @@ -34,6 +34,7 @@ __all__ = [ "get_all_matching_child_prims", "find_first_matching_prim", "find_matching_prims", + "resolve_matching_prims_from_source", "find_matching_prim_paths", "find_global_fixed_joint_prim", "add_labels", @@ -94,6 +95,7 @@ from .queries import ( get_all_matching_child_prims, find_first_matching_prim, find_matching_prims, + resolve_matching_prims_from_source, find_matching_prim_paths, find_global_fixed_joint_prim, ) diff --git a/source/isaaclab/isaaclab/sim/utils/queries.py b/source/isaaclab/isaaclab/sim/utils/queries.py index 1929e70941c7..ce3c95b64cb4 100644 --- a/source/isaaclab/isaaclab/sim/utils/queries.py +++ b/source/isaaclab/isaaclab/sim/utils/queries.py @@ -13,6 +13,9 @@ from pxr import Sdf, Usd, UsdPhysics +from isaaclab.cloner.cloner_utils import resolve_clone_plan_source +from isaaclab.sim.simulation_context import SimulationContext + from .stage import get_current_stage # import logger @@ -354,6 +357,90 @@ def find_matching_prims(prim_path_regex: str, stage: Usd.Stage | None = None) -> return output_prims +def resolve_matching_prims_from_source( + path_expr: str, + *, + predicate: Callable[[Usd.Prim], bool] | None = None, + env_regex_ns: str = "/World/envs/env_.*", +) -> list[tuple[Usd.Prim, str]]: + """Resolve prims matching ``path_expr`` (regex) under the first instance. + + Identify the env-id segment, concretize it to the *first instance* (the authored source + template in clone-plan mode, env-0 in legacy mode), then evaluate the remainder of + ``path_expr`` as a path-segment regex via :func:`find_matching_prims`. Downstream regex + tokens (e.g. ``LF_.*``, ``.*_foot``) are preserved verbatim and matched there. + + Args: + path_expr: Destination-side path expression (e.g. a ``prim_path``), which may contain + regex wildcards in the env-id and/or asset-relative segments. + predicate: Optional callable accepting a :class:`Usd.Prim` and returning ``True`` for + prims to keep. ``None`` keeps every match. + env_regex_ns: Instance-root namespace regex, defaulting to the standard + ``"/World/envs/env_.*"``. In legacy (no-clone-plan) resolution, when ``path_expr`` + sits under this namespace its path depth fixes the per-instance ("env") boundary. + Otherwise the boundary falls back to the first regex segment of ``path_expr`` (the + first ``.*`` is treated as the env id), which covers ad-hoc roots such as + ``"/World/Table_.*/Object"``. Layouts the fallback would mis-split (e.g. more than + one wildcard level) must pass an explicit namespace here. Ignored when a clone plan + owns ``path_expr``. + + Returns: + List of ``(matched_prim, destination_expr)`` pairs, where ``destination_expr`` is the + multi-instance path expression (not a single concrete instance) so callers can build + views spanning every instance. Empty when ``path_expr`` matches no prim. + """ + plan = SimulationContext.instance().get_clone_plan() + resolved = resolve_clone_plan_source(path_expr, plan) if plan is not None else None + if resolved is not None: + source_path, dest_glob, asset_suffix = resolved + walk_root = source_path + asset_suffix + results = [ + (prim, dest_glob + prim.GetPath().pathString[len(source_path) :]) for prim in find_matching_prims(walk_root) + ] + else: + # No clone plan, or ``path_expr`` is not owned by any plan row. Resolve from the stage + # in two phases (mirroring the clone-plan branch above): (1) locate ONE instance root to + # search from, (2) collect the bodies of interest within just that instance and map each + # back to the multi-instance pattern. Phase 1 stops at the first match and phase 2 walks + # under a concrete instance prefix, so only a single instance subtree is traversed. + segments = path_expr.strip("/").split("/") + ns_segments = env_regex_ns.strip("/").split("/") + # Instance ("env") boundary. Assume the standard namespace ``env_regex_ns`` and put the + # boundary at its depth when ``path_expr`` sits under it -- literal ns segments must + # match, wildcard ns segments (e.g. ``env_.*``) accept any segment. Otherwise fall back + # to the first regex segment of ``path_expr`` (treat the first ``.*`` as the env id), + # covering ad-hoc roots like ``/World/Table_.*/Object``. A layout the fallback would + # mis-split (e.g. multiple wildcard levels) must pass ``env_regex_ns``. + under_ns = len(segments) >= len(ns_segments) and all( + ns_seg == seg or not ns_seg.isidentifier() for ns_seg, seg in zip(ns_segments, segments) + ) + if under_ns: + instance_seg = len(ns_segments) - 1 + else: + instance_seg = next((i for i, seg in enumerate(segments) if not seg.isidentifier()), None) + first = find_first_matching_prim(path_expr) + if first is None: + results = [] + elif instance_seg is None: + # Fully concrete path: a single instance, mapped to itself. + results = [(first, first.GetPath().pathString)] + else: + instance_expr = "/" + "/".join(segments[: instance_seg + 1]) + match_segments = first.GetPath().pathString.strip("/").split("/") + instance_root = "/" + "/".join(match_segments[: instance_seg + 1]) + trailing = segments[instance_seg + 1 :] + walk_root = instance_root + ("/" + "/".join(trailing) if trailing else "") + results = [ + (prim, instance_expr + prim.GetPath().pathString[len(instance_root) :]) + for prim in find_matching_prims(walk_root) + if prim.GetPath().pathString == instance_root + or prim.GetPath().pathString.startswith(instance_root + "/") + ] + if predicate is not None: + results = [(prim, dest) for prim, dest in results if predicate(prim)] + return results + + def find_matching_prim_paths(prim_path_regex: str, stage: Usd.Stage | None = None) -> list[str]: """Find all the matching prim paths in the stage based on input regex expression. diff --git a/source/isaaclab/test/scene/test_interactive_scene.py b/source/isaaclab/test/scene/test_interactive_scene.py index f56803ef5cf6..ce8b2a6cb26d 100644 --- a/source/isaaclab/test/scene/test_interactive_scene.py +++ b/source/isaaclab/test/scene/test_interactive_scene.py @@ -155,7 +155,8 @@ def _set_clone_plan(plan): ) scene.env_fmt = "/World/envs/env_{}" scene._ALL_INDICES = torch.arange(3, dtype=torch.long) - scene._default_env_origins = torch.zeros((3, 3), dtype=torch.float32) + scene._default_env_pose = torch.zeros((3, 7), dtype=torch.float32) + scene._default_env_pose[:, 6] = 1.0 # identity quaternion (xyzw) scene._clone_plan = ClonePlan( sources=(scene.env_fmt.format(0),), destinations=(scene.env_fmt,), @@ -195,8 +196,14 @@ def _usd_replicate(stage, *args, **kwargs): mapping = physics_calls[0][1][3] assert mapping.dtype == torch.bool assert mapping.shape == (1, scene.num_envs) - assert physics_calls[0][2]["positions"] is scene._default_env_origins - assert usd_calls[0][2]["positions"] is scene._default_env_origins + # Positions are a zero-copy ``[:, :3]`` view of ``_default_env_pose``: same storage, sliced shape. + pose_storage_ptr = scene._default_env_pose.untyped_storage().data_ptr() + physics_positions = physics_calls[0][2]["positions"] + usd_positions = usd_calls[0][2]["positions"] + assert physics_positions.untyped_storage().data_ptr() == pose_storage_ptr + assert physics_positions.shape == (scene.num_envs, 3) + assert usd_positions.untyped_storage().data_ptr() == pose_storage_ptr + assert usd_positions.shape == (scene.num_envs, 3) assert len(set_plan_calls) == 1 plan = set_plan_calls[-1] assert isinstance(plan, ClonePlan) @@ -237,7 +244,8 @@ def test_clone_environments_executes_asset_level_plan_without_usd_positions(monk scene._sensors = {} scene.env_fmt = "/World/envs/env_{}" scene._ALL_INDICES = torch.arange(2, dtype=torch.long) - scene._default_env_origins = torch.ones((2, 3), dtype=torch.float32) + scene._default_env_pose = torch.ones((2, 7), dtype=torch.float32) + scene._default_env_pose[:, 3:6] = 0.0 # identity quaternion (xyzw) scene._clone_plan = ClonePlan( sources=("/World/envs/env_0/Object", "/World/envs/env_1/Object"), destinations=("/World/envs/env_{}/Object", "/World/envs/env_{}/Object"), @@ -269,7 +277,9 @@ def _noop_fabric_notices(stage, *, restore=True): scene.clone_environments(copy_from_source=False) assert len(physics_calls) == 1 - assert physics_calls[0][1]["positions"] is scene._default_env_origins + physics_positions = physics_calls[0][1]["positions"] + assert physics_positions.untyped_storage().data_ptr() == scene._default_env_pose.untyped_storage().data_ptr() + assert physics_positions.shape == (scene.num_envs, 3) assert len(usd_calls) == 1 assert usd_calls[0][1]["positions"] is None assert set_plan_calls == [scene._clone_plan] diff --git a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py index 8f5014b3a860..99a44bc6958a 100644 --- a/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py +++ b/source/isaaclab/test/sensors/test_multi_mesh_ray_caster_camera.py @@ -465,6 +465,7 @@ def _create_heterogeneous_clone_scene(sim: sim_utils.SimulationContext, num_envs sim_utils.create_prim("/World/envs", "Xform", stage=stage) for env_id, origin in enumerate(env_origins.cpu().tolist()): sim_utils.create_prim(env_fmt.format(env_id), "Xform", translation=tuple(origin), stage=stage) + sim_utils.create_prim(env_fmt.format(env_id) + "/RayCasterCamera", "Xform", stage=stage) robot_mask = torch.zeros((2, num_envs), dtype=torch.bool, device=sim.device) robot_mask[0, 0::2] = True diff --git a/source/isaaclab/test/sensors/test_ray_caster_integration.py b/source/isaaclab/test/sensors/test_ray_caster_integration.py index 91fa9c12fedc..d197bf5f558f 100644 --- a/source/isaaclab/test/sensors/test_ray_caster_integration.py +++ b/source/isaaclab/test/sensors/test_ray_caster_integration.py @@ -387,7 +387,7 @@ def _create_object_body(path: str) -> None: env2_part = stage.GetPrimAtPath("/World/envs/env_2/Object/part_0") assert env2_part is None or not env2_part.IsValid() - # Geometry is selected from ClonePlan rows, but poses come from the batched object view. + # Geometry is selected from ClonePlan entries, but poses come from the batched object view. mesh_ids = wp.to_torch(sensor._mesh_ids_wp).cpu() mesh_positions = wp.to_torch(sensor._mesh_positions_w).cpu() assert sensor._mesh_ids_wp.shape == (num_envs, 1) diff --git a/source/isaaclab/test/sim/test_cloner.py b/source/isaaclab/test/sim/test_cloner.py index 1f526ac74584..db6568b8884d 100644 --- a/source/isaaclab/test/sim/test_cloner.py +++ b/source/isaaclab/test/sim/test_cloner.py @@ -20,7 +20,8 @@ from pxr import UsdGeom import isaaclab.sim as sim_utils -from isaaclab.cloner import make_clone_plan, sequential, usd_replicate +from isaaclab.cloner import ClonePlan, make_clone_plan, sequential, usd_replicate +from isaaclab.cloner.cloner_utils import iter_clone_plan_matches from isaaclab.sim import build_simulation_context pytestmark = pytest.mark.isaacsim_ci @@ -223,7 +224,7 @@ def test_clone_decorator_wildcard_patterns( def test_make_clone_plan_returns_flat_source_rows(sim): """make_clone_plan exposes the flat source-to-env mask used by scene cloning.""" - plan = make_clone_plan( + sources, destinations, clone_mask = make_clone_plan( [["/World/envs/env_0/Object", "/World/envs/env_1/Object"]], ["/World/envs/env_{}/Object"], num_clones=4, @@ -231,10 +232,107 @@ def test_make_clone_plan_returns_flat_source_rows(sim): device=sim.cfg.device, ) - assert plan.sources == ("/World/envs/env_0/Object", "/World/envs/env_1/Object") - assert plan.destinations == ("/World/envs/env_{}/Object", "/World/envs/env_{}/Object") - assert plan.clone_mask.shape == (2, 4) - assert plan.clone_mask.dtype == torch.bool - assert torch.all(plan.clone_mask.sum(dim=0) == 1) - actual_source_idx = plan.clone_mask.to(torch.int).argmax(dim=0).cpu() + assert sources == ("/World/envs/env_0/Object", "/World/envs/env_1/Object") + assert destinations == ("/World/envs/env_{}/Object", "/World/envs/env_{}/Object") + assert clone_mask.shape == (2, 4) + assert clone_mask.dtype == torch.bool + assert torch.all(clone_mask.sum(dim=0) == 1) + actual_source_idx = clone_mask.to(torch.int).argmax(dim=0).cpu() assert torch.equal(actual_source_idx, torch.tensor([0, 1, 0, 1])) + + +def test_iter_clone_plan_matches(sim): + """ClonePlan entries can be matched by destination path expression.""" + sources, destinations, clone_mask = make_clone_plan( + [["/World/envs/env_0/Object", "/World/envs/env_1/Object"]], + ["/World/envs/env_{}/Object"], + num_clones=4, + clone_strategy=sequential, + device=sim.cfg.device, + ) + plan = ClonePlan(sources=sources, destinations=destinations, clone_mask=clone_mask) + + matches = list(iter_clone_plan_matches(plan, "/World/envs/env_.*/Object/Body/Camera")) + + assert matches == [ + ( + "/World/envs/env_0/Object", + "/World/envs/env_{}/Object", + "/World/envs/env_0/Object/Body/Camera", + (0, 2), + ), + ( + "/World/envs/env_1/Object", + "/World/envs/env_{}/Object", + "/World/envs/env_1/Object/Body/Camera", + (1, 3), + ), + ] + + plan = ClonePlan( + sources=("/World/envs/env_3/Object",), + destinations=("/World/envs/env_{}/Object",), + clone_mask=torch.tensor([[False, False, True, True]], device=sim.cfg.device), + ) + + matches = list(iter_clone_plan_matches(plan, "/World/envs/env_.*/Object/Body/Camera")) + + assert matches == [ + ( + "/World/envs/env_3/Object", + "/World/envs/env_{}/Object", + "/World/envs/env_3/Object/Body/Camera", + (2, 3), + ) + ] + + plan = ClonePlan( + sources=("/World/source/Object",), + destinations=("/World/scenes/{}/Object",), + clone_mask=torch.tensor([[True, True]], device=sim.cfg.device), + ) + + matches = list(iter_clone_plan_matches(plan, "/World/scenes/.*/Object/Body/Camera")) + + assert matches == [ + ( + "/World/source/Object", + "/World/scenes/{}/Object", + "/World/source/Object/Body/Camera", + (0, 1), + ) + ] + + plan = ClonePlan( + sources=("/World/source",), + destinations=("/World/scenes/{}",), + clone_mask=torch.tensor([[True, True]], device=sim.cfg.device), + ) + + matches = list(iter_clone_plan_matches(plan, "/World/scenes/.*/Object/Body/Camera")) + + assert matches == [ + ( + "/World/source", + "/World/scenes/{}", + "/World/source/Object/Body/Camera", + (0, 1), + ) + ] + + plan = ClonePlan( + sources=("/World/envs/env_0", "/World/envs/env_0/Object"), + destinations=("/World/envs/env_{}", "/World/envs/env_{}/Object"), + clone_mask=torch.tensor([[True, True], [True, True]], device=sim.cfg.device), + ) + + matches = list(iter_clone_plan_matches(plan, "/World/envs/env_.*/Object/Body/Camera")) + + assert matches == [ + ( + "/World/envs/env_0/Object", + "/World/envs/env_{}/Object", + "/World/envs/env_0/Object/Body/Camera", + (0, 1), + ) + ] diff --git a/source/isaaclab_contrib/changelog.d/newton-clone-plan.rst b/source/isaaclab_contrib/changelog.d/newton-clone-plan.rst new file mode 100644 index 000000000000..ecfebcd81efc --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/newton-clone-plan.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed Newton replicated-scene setup for deformable VBD managers to use + clone-plan source prims. diff --git a/source/isaaclab_contrib/isaaclab_contrib/deformable/coupled_featherstone_vbd_manager.py b/source/isaaclab_contrib/isaaclab_contrib/deformable/coupled_featherstone_vbd_manager.py index 312d9fc69f10..19e5aa62cee6 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/deformable/coupled_featherstone_vbd_manager.py +++ b/source/isaaclab_contrib/isaaclab_contrib/deformable/coupled_featherstone_vbd_manager.py @@ -243,7 +243,7 @@ def instantiate_builder_from_stage(cls): ) # Inject registered sites into the proto before replication - global_sites, proto_sites = cls._cl_inject_sites(builder, {proto_path: proto}) + global_sites, proto_sites, world_sites = cls._cl_inject_sites(builder, {proto_path: proto}) global_site_map: dict[str, tuple[int, None]] = {label: (idx, None) for label, idx in global_sites.items()} num_worlds = len(env_paths) local_site_map: dict[str, list[list[int]]] = {} @@ -264,7 +264,13 @@ def instantiate_builder_from_stage(cls): rotation.GetImaginary()[2], rotation.GetReal(), ) - builder.add_builder(proto, xform=wp.transform(pos, quat)) + env_xform = wp.transform(pos, quat) + builder.add_builder(proto, xform=env_xform) + for label, xform in world_sites.items(): + if label not in local_site_map: + local_site_map[label] = [[] for _ in range(num_worlds)] + site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(env_xform, xform), label=label) + local_site_map[label][col].append(site_idx) for label, proto_shape_indices in site_entries.items(): if label not in local_site_map: local_site_map[label] = [[] for _ in range(num_worlds)] diff --git a/source/isaaclab_contrib/isaaclab_contrib/deformable/coupled_mjwarp_vbd_manager.py b/source/isaaclab_contrib/isaaclab_contrib/deformable/coupled_mjwarp_vbd_manager.py index 5097284bafc7..dd288f4347f3 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/deformable/coupled_mjwarp_vbd_manager.py +++ b/source/isaaclab_contrib/isaaclab_contrib/deformable/coupled_mjwarp_vbd_manager.py @@ -243,7 +243,7 @@ def instantiate_builder_from_stage(cls): ) # Inject registered sites into the proto before replication - global_sites, proto_sites = cls._cl_inject_sites(builder, {proto_path: proto}) + global_sites, proto_sites, world_sites = cls._cl_inject_sites(builder, {proto_path: proto}) global_site_map: dict[str, tuple[int, None]] = {label: (idx, None) for label, idx in global_sites.items()} num_worlds = len(env_paths) local_site_map: dict[str, list[list[int]]] = {} @@ -264,7 +264,13 @@ def instantiate_builder_from_stage(cls): rotation.GetImaginary()[2], rotation.GetReal(), ) - builder.add_builder(proto, xform=wp.transform(pos, quat)) + env_xform = wp.transform(pos, quat) + builder.add_builder(proto, xform=env_xform) + for label, xform in world_sites.items(): + if label not in local_site_map: + local_site_map[label] = [[] for _ in range(num_worlds)] + site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(env_xform, xform), label=label) + local_site_map[label][col].append(site_idx) for label, proto_shape_indices in site_entries.items(): if label not in local_site_map: local_site_map[label] = [[] for _ in range(num_worlds)] diff --git a/source/isaaclab_contrib/isaaclab_contrib/deformable/vbd_manager.py b/source/isaaclab_contrib/isaaclab_contrib/deformable/vbd_manager.py index 88380f078673..48c31b06e4f1 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/deformable/vbd_manager.py +++ b/source/isaaclab_contrib/isaaclab_contrib/deformable/vbd_manager.py @@ -218,7 +218,7 @@ def instantiate_builder_from_stage(cls): ) # Inject registered sites into the proto before replication - global_sites, proto_sites = cls._cl_inject_sites(builder, {proto_path: proto}) + global_sites, proto_sites, world_sites = cls._cl_inject_sites(builder, {proto_path: proto}) global_site_map: dict[str, tuple[int, None]] = {label: (idx, None) for label, idx in global_sites.items()} num_worlds = len(env_paths) local_site_map: dict[str, list[list[int]]] = {} @@ -239,7 +239,13 @@ def instantiate_builder_from_stage(cls): rotation.GetImaginary()[2], rotation.GetReal(), ) - builder.add_builder(proto, xform=wp.transform(pos, quat)) + env_xform = wp.transform(pos, quat) + builder.add_builder(proto, xform=env_xform) + for label, xform in world_sites.items(): + if label not in local_site_map: + local_site_map[label] = [[] for _ in range(num_worlds)] + site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(env_xform, xform), label=label) + local_site_map[label][col].append(site_idx) for label, proto_shape_indices in site_entries.items(): if label not in local_site_map: local_site_map[label] = [[] for _ in range(num_worlds)] diff --git a/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor.py b/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor.py index 12ff6cfd3e8d..970449712f37 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor.py +++ b/source/isaaclab_contrib/isaaclab_contrib/sensors/tacsl_sensor/visuotactile_sensor.py @@ -319,7 +319,15 @@ def _create_physx_views(self) -> None: c. Creates rigid body view for object """ - elastomer_pattern = self._parent_prims[0].GetPath().pathString.replace("env_0", "env_*") + # Resolve the elastomer's destination expression (multi-env glob form for PhysX views). + # The sensor's cfg.prim_path lives under the elastomer; the parent expression is the + # elastomer body itself (matching :attr:`SensorBase._parent_prims`). + elastomer_expr = self.cfg.prim_path.rsplit("/", 1)[0] + matches = sim_utils.resolve_matching_prims_from_source(elastomer_expr) + if not matches: + raise RuntimeError(f"No prim found at '{elastomer_expr}'.") + elastomer_dest_expr = matches[0][1] + elastomer_pattern = elastomer_dest_expr.replace(".*", "*") self._elastomer_body_view = self._physics_sim_view.create_rigid_body_view([elastomer_pattern]) # Get elastomer COM for velocity correction self._elastomer_com_b = ( @@ -416,8 +424,14 @@ def _generate_tactile_points(self, num_divs: list, margin: float, visualize: boo """ - # Get the elastomer prim path - elastomer_prim_path = self._parent_prims[0].GetPath().pathString + # Resolve the elastomer's source-side env prim and use it as the walk root. + # The sensor's cfg.prim_path lives under the elastomer; the parent expression is the + # elastomer body itself (matching :attr:`SensorBase._parent_prims`). + elastomer_expr = self.cfg.prim_path.rsplit("/", 1)[0] + matches = sim_utils.resolve_matching_prims_from_source(elastomer_expr) + if not matches: + raise RuntimeError(f"No prim found at '{elastomer_expr}'.") + elastomer_prim_path = matches[0][0].GetPath().pathString def is_visual_mesh(prim) -> bool: """Check if a mesh prim has visual properties (visual mesh, not collision mesh).""" diff --git a/source/isaaclab_newton/changelog.d/newton-clone-plan.rst b/source/isaaclab_newton/changelog.d/newton-clone-plan.rst new file mode 100644 index 000000000000..e40f9e37b0b9 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/newton-clone-plan.rst @@ -0,0 +1,7 @@ +Fixed +^^^^^ + +* Fixed Newton frame-view and ray-caster sensor resolution to use clone-plan + source paths and Newton model labels instead of cloned destination USD prims. +* Fixed Newton Warp camera preparation to avoid requiring PPISP when camera ISP + is disabled. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index e5b85927d471..90ced0daed8b 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py @@ -26,11 +26,11 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation +from isaaclab.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source _HAS_NEWTON_ACTUATORS = importlib.util.find_spec("isaaclab_newton.actuators") is not None from isaaclab.physics import PhysicsEvent -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 @@ -3493,41 +3493,27 @@ def _initialize_impl(self): 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, + + def has_articulation_root_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.ArticulationRootAPI)) + + matches = resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString + root_prims = get_all_matching_child_prims( + walk_root, predicate=has_articulation_root_api, 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: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] 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." + f"Expected exactly one ArticulationRootAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - - # 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 - + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] # -- articulation self._root_view = ArticulationView( SimulationManager.get_model(), 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 index b93c9075393d..a56e4fb1a636 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object/rigid_object.py @@ -17,10 +17,10 @@ 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.physics import PhysicsEvent +from isaaclab.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_newton.assets import kernels as shared_kernels @@ -988,47 +988,24 @@ def set_inertias_mask( """ def _initialize_impl(self): - # 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, + def has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) + + matches = resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString + root_prims = get_all_matching_child_prims( + walk_root, predicate=has_rigid_body_api, traverse_instance_prims=False ) - if len(root_prims) == 0: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] raise RuntimeError( - f"Failed to find a rigid body when resolving '{self.cfg.prim_path}'." - " Please ensure that the prim has 'USD RigidBodyAPI' applied." + f"Expected exactly one RigidBodyAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - 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) :] + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] # -- object view self._root_view = ArticulationView( SimulationManager.get_model(), diff --git a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py index b11415d48231..ee586f351b86 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/rigid_object_collection/rigid_object_collection.py @@ -1074,51 +1074,28 @@ def set_inertias_mask( def _initialize_impl(self): # clear body names list to prevent double counting on re-initialization self._body_names_list.clear() + + def has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) + root_prim_path_exprs: list[str] = [] for name, rigid_body_cfg in self.cfg.rigid_objects.items(): - # 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(rigid_body_cfg.prim_path) - if template_prim is None: - raise RuntimeError(f"Failed to find prim for expression: '{rigid_body_cfg.prim_path}'.") - template_prim_path = template_prim.GetPath().pathString - - # find rigid root prims + matches = sim_utils.resolve_matching_prims_from_source(rigid_body_cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{rigid_body_cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString root_prims = sim_utils.get_all_matching_child_prims( - template_prim_path, - predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI), - traverse_instance_prims=False, + walk_root, predicate=has_rigid_body_api, traverse_instance_prims=False ) - if len(root_prims) == 0: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] raise RuntimeError( - f"Failed to find a rigid body when resolving '{rigid_body_cfg.prim_path}'." - " Please ensure that the prim has 'USD RigidBodyAPI' applied." + f"Expected exactly one RigidBodyAPI prim under '{walk_root}'" + f" (resolved from '{rigid_body_cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - if len(root_prims) > 1: - raise RuntimeError( - f"Failed to find a single rigid body when resolving '{rigid_body_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." - ) - - # check that no rigid object has an articulation root API, which decreases simulation performance - 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 '{rigid_body_cfg.prim_path}' in the rigid object" - f" collection. These are 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 = rigid_body_cfg.prim_path + root_prim_path[len(template_prim_path) :] + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] root_prim_path_exprs.append(root_prim_path_expr.replace(".*", "*")) self._body_names_list.append(name) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py index e99b1eb7abdd..e868607c05c5 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py @@ -98,7 +98,7 @@ def _build_newton_builder_from_mapping( protos[src_path] = p # Inject registered sites into prototypes (and global sites into main builder) - global_sites, proto_sites = NewtonManager._cl_inject_sites(builder, protos) + global_sites, proto_sites, world_sites = NewtonManager._cl_inject_sites(builder, protos) # Global sites: (int, None) global_site_map: dict[str, tuple[int, None]] = {label: (idx, None) for label, idx in global_sites.items()} @@ -114,6 +114,12 @@ def _build_newton_builder_from_mapping( builder.begin_world() # add all active sources for this world delta_pos = (positions[col] - env0_pos).tolist() + env_xform = wp.transform(positions[col].tolist(), quaternions[col].tolist()) + for label, xform in world_sites.items(): + if label not in local_site_map: + local_site_map[label] = [[] for _ in range(num_worlds)] + site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(env_xform, xform), label=label) + local_site_map[label][col].append(site_idx) for row in torch.nonzero(mapping[:, col], as_tuple=True)[0].tolist(): proto = protos[sources[row]] offset = builder.shape_count @@ -285,6 +291,12 @@ def newton_physics_replicate( Returns: Tuple of the populated Newton model builder and stage metadata. """ + 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 + builder, stage_info, site_index_map = _build_newton_builder_from_mapping( stage=stage, sources=sources, @@ -297,6 +309,9 @@ def newton_physics_replicate( ) _rename_builder_labels(builder, sources, destinations, env_ids, mapping) NewtonManager._cl_site_index_map = site_index_map + NewtonManager._world_xforms = [ + wp.transform(positions[col].tolist(), quaternions[col].tolist()) for col in range(mapping.size(1)) + ] NewtonManager.set_builder(builder) NewtonManager._num_envs = mapping.size(1) return builder, stage_info diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 21a37d45d0b7..d5991ed395ab 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -275,15 +275,16 @@ class NewtonManager(PhysicsManager): # CL: Cloning / Replication logic # TODO: These attributes support cloning-specific logic and should be moved into a cloner class # Pending site requests from sensors. - # Key: (body_pattern, xform_floats), Value: (label, wp.transform) - # identical (body_pattern, transform) reuses the same site. - _cl_pending_sites: dict[tuple[str | None, tuple[float, ...]], tuple[str, wp.transform]] = {} + # Key: (body_pattern, per_world, xform_floats), Value: (label, wp.transform) + # identical (body_pattern, per_world, transform) reuses the same site. + _cl_pending_sites: dict[tuple[str | None, bool, tuple[float, ...]], tuple[str, wp.transform]] = {} # Maps each site label to its resolved global or local site entry. _GlobalSite = tuple[int, None] _LocalSite = tuple[None, list[list[int]]] _SiteEntry = _GlobalSite | _LocalSite _cl_site_index_map: dict[str, _SiteEntry] = {} + _world_xforms: list[wp.transform] | None = None @classmethod def initialize(cls, sim_context: SimulationContext) -> None: @@ -691,6 +692,7 @@ def clear(cls): NewtonManager._scene_data_backend = None NewtonManager._cl_pending_sites = {} NewtonManager._cl_site_index_map = {} + NewtonManager._world_xforms = None NewtonManager._pending_extended_state_attributes = set() NewtonManager._pending_extended_contact_attributes = set() NewtonManager._views = [] @@ -728,14 +730,14 @@ def create_builder(cls, up_axis: str | None = None, **kwargs) -> ModelBuilder: return builder @classmethod - def cl_register_site(cls, body_pattern: str | None, xform: wp.transform) -> str: + def cl_register_site(cls, body_pattern: str | None, xform: wp.transform, *, per_world: bool = False) -> str: """Register a site request for injection into prototypes before replication. Sensors call this during ``__init__``. Sites are injected into prototype builders by :meth:`_cl_inject_sites` (called from ``newton_replicate``) before ``add_builder``, so they replicate correctly per-world. - Identical ``(body_pattern, transform)`` registrations share sites. + Identical ``(body_pattern, per_world, transform)`` registrations share sites. The *body_pattern* is matched against prototype-local body labels (e.g. ``"Robot/link.*"``) when replication is active, or against the @@ -748,12 +750,16 @@ def cl_register_site(cls, body_pattern: str | None, xform: wp.transform) -> str: for multi-body wildcards), or ``None`` for global sites (world-origin reference, etc.). xform: Site transform relative to body. + per_world: When ``True``, ``body_pattern`` must be ``None`` and one + bodyless site is created in each cloned world's frame. Returns: Assigned site label suffix. """ + if per_world and body_pattern is not None: + raise ValueError("per_world site registration requires body_pattern=None.") xform_key = tuple(xform) - key = (body_pattern, xform_key) + key = (body_pattern, per_world, xform_key) if key in cls._cl_pending_sites: return cls._cl_pending_sites[key][0] label = f"ft_{len(cls._cl_pending_sites)}" @@ -791,7 +797,7 @@ def _cl_inject_sites( cls, main_builder: ModelBuilder, proto_builders: dict[str, ModelBuilder], - ) -> tuple[dict[str, int], dict[int, dict[str, list[int]]]]: + ) -> tuple[dict[str, int], dict[int, dict[str, list[int]]], dict[str, wp.transform]]: """Inject registered sites into prototype builders before replication. Non-global sites are matched against prototype body labels using @@ -809,14 +815,20 @@ def _cl_inject_sites( proto_builders: ``{src_path: ModelBuilder}`` prototype builders. Returns: - Tuple of ``(global_sites, proto_sites)`` where *global_sites* maps - ``{label: main_builder_shape_idx}`` and *proto_sites* maps - ``{id(proto): {label: [proto_local_shape_idx, ...]}}``. + Tuple of ``(global_sites, proto_sites, world_sites)`` where + *global_sites* maps ``{label: main_builder_shape_idx}``, + *proto_sites* maps ``{id(proto): {label: [proto_local_shape_idx, ...]}}``, + and *world_sites* maps ``{label: env_root_relative_transform}``. """ global_sites: dict[str, int] = {} proto_sites: dict[int, dict[str, list[int]]] = {} - for (body_pattern, _xform_key), (label, xform) in cls._cl_pending_sites.items(): + world_sites: dict[str, wp.transform] = {} + + for (body_pattern, per_world, _xform_key), (label, xform) in cls._cl_pending_sites.items(): + if per_world: + world_sites[label] = xform + continue if body_pattern is None: site_idx = main_builder.add_site(body=-1, xform=xform, label=label) global_sites[label] = site_idx @@ -849,7 +861,7 @@ def _cl_inject_sites( ) cls._cl_pending_sites.clear() - return global_sites, proto_sites + return global_sites, proto_sites, world_sites @classmethod def _cl_inject_sites_fallback(cls) -> None: @@ -858,12 +870,16 @@ def _cl_inject_sites_fallback(cls) -> None: Populates :attr:`_cl_site_index_map` with the unified per-world structure: - Global sites (``body_pattern is None``): ``(shape_idx, None)`` - - Local sites: ``(None, [[idx, ...]])`` — one sublist for the single world. + - Local and world sites: ``(None, [[idx, ...]])`` — one sublist for the single world. """ builder = cls._builder body_labels = list(builder.body_label) - for (body_pattern, _xform_key), (label, xform) in cls._cl_pending_sites.items(): + for (body_pattern, per_world, _xform_key), (label, xform) in cls._cl_pending_sites.items(): + if per_world: + site_idx = builder.add_site(body=-1, xform=xform, label=label) + cls._cl_site_index_map[label] = (None, [[site_idx]]) + continue if body_pattern is None: site_idx = builder.add_site(body=-1, xform=xform, label=label) cls._cl_site_index_map[label] = (site_idx, None) @@ -1060,6 +1076,7 @@ def instantiate_builder_from_stage(cls): if not env_paths: # No env Xforms — flat loading builder.add_usd(stage, schema_resolvers=schema_resolvers) + NewtonManager._world_xforms = [wp.transform()] else: # Load everything except the env subtrees (ground plane, lights, etc.) ignore_paths = [path for _, path in env_paths] @@ -1075,11 +1092,12 @@ def instantiate_builder_from_stage(cls): ) # Inject registered sites into the proto before replication - global_sites, proto_sites = cls._cl_inject_sites(builder, {proto_path: proto}) + global_sites, proto_sites, world_sites = cls._cl_inject_sites(builder, {proto_path: proto}) global_site_map: dict[str, tuple[int, None]] = {label: (idx, None) for label, idx in global_sites.items()} num_worlds = len(env_paths) local_site_map: dict[str, list[list[int]]] = {} site_entries = proto_sites.get(id(proto), {}) + world_xforms: list[wp.transform] = [] # Add each env as a separate Newton world xform_cache = UsdGeom.XformCache() @@ -1096,7 +1114,14 @@ def instantiate_builder_from_stage(cls): rotation.GetImaginary()[2], rotation.GetReal(), ) - builder.add_builder(proto, xform=wp.transform(pos, quat)) + env_xform = wp.transform(pos, quat) + world_xforms.append(env_xform) + builder.add_builder(proto, xform=env_xform) + for label, xform in world_sites.items(): + if label not in local_site_map: + local_site_map[label] = [[] for _ in range(num_worlds)] + site_idx = builder.add_site(body=-1, xform=wp.transform_multiply(env_xform, xform), label=label) + local_site_map[label][col].append(site_idx) for label, proto_shape_indices in site_entries.items(): if label not in local_site_map: local_site_map[label] = [[] for _ in range(num_worlds)] @@ -1108,6 +1133,7 @@ def instantiate_builder_from_stage(cls): **global_site_map, **{label: (None, per_world) for label, per_world in local_site_map.items()}, } + NewtonManager._world_xforms = world_xforms NewtonManager._num_envs = len(env_paths) cls.set_builder(builder) diff --git a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py index 27574c3907d9..04745e894677 100644 --- a/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py +++ b/source/isaaclab_newton/isaaclab_newton/renderers/newton_warp_renderer.py @@ -261,7 +261,7 @@ def prepare_cameras(self, stage: Any, spec: CameraRenderSpec) -> None: owns the sentinel-resolution + cfg-normalization step. Newton has no USD-side overrides to author beyond this. """ - if not spec.camera_prim_paths: + if spec.cfg.isp_cfg is None or not spec.camera_prim_paths: return from isaaclab_ppisp import resolve_and_normalize diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/joint_wrench/joint_wrench_sensor.py b/source/isaaclab_newton/isaaclab_newton/sensors/joint_wrench/joint_wrench_sensor.py index d07b4809c6ea..7b8721538913 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/joint_wrench/joint_wrench_sensor.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/joint_wrench/joint_wrench_sensor.py @@ -16,7 +16,7 @@ from pxr import UsdPhysics from isaaclab.sensors.joint_wrench import BaseJointWrenchSensor -from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims +from isaaclab.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source from isaaclab_newton.physics import NewtonManager @@ -37,10 +37,11 @@ class JointWrenchSensor(BaseJointWrenchSensor): (child-side joint frame, child-side joint anchor as reference point) before storing it in per-joint force / torque buffers. - :attr:`~isaaclab.sensors.SensorBaseCfg.prim_path` must point at either - the articulation root prim or a parent prim containing a single - articulation root in every environment. ``FREE`` and ``FIXED`` joints are - excluded — neither has a meaningful joint anchor. + :attr:`~isaaclab.sensors.SensorBaseCfg.prim_path` may point at either an + articulation root expression or an env-scoped parent prefix. Newton label + matching selects the articulations owned by that prefix. ``FREE`` and + ``FIXED`` joints are excluded because neither has a meaningful joint + anchor. """ cfg: JointWrenchSensorCfg @@ -127,7 +128,24 @@ def _initialize_impl(self) -> None: model = NewtonManager.get_model() state_0 = NewtonManager.get_state_0() - root_prim_path_expr = self._resolve_articulation_root_prim_path() + def has_articulation_root_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.ArticulationRootAPI)) + + matches = resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString + root_prims = get_all_matching_child_prims( + walk_root, predicate=has_articulation_root_api, traverse_instance_prims=False + ) + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] + raise RuntimeError( + f"Expected exactly one ArticulationRootAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." + ) + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] self._root_view = ArticulationView( model, root_prim_path_expr.replace(".*", "*"), @@ -170,35 +188,6 @@ def _initialize_impl(self) -> None: logger.info(f"Joint wrench sensor initialized: {self._num_envs} envs, {self._num_joints} joints") - def _resolve_articulation_root_prim_path(self) -> str: - """Resolve the articulation root prim path expression from the configured asset prim path.""" - 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 - - first_env_root_prims = get_all_matching_child_prims( - first_env_matching_prim_path, - predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI) - and prim.GetAttribute("physxArticulation:articulationEnabled").Get() is not False, - 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." - ) - - 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) :] - return self.cfg.prim_path + root_prim_path_relative_to_prim_path - def _update_buffers_impl(self, env_mask: wp.array) -> None: """Convert Newton's body_parent_f into INCOMING_JOINT_FRAME force and torque buffers. diff --git a/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster.py b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster.py index e10076d884c8..b514d52ca685 100644 --- a/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab_newton/isaaclab_newton/sensors/ray_caster/ray_caster.py @@ -15,6 +15,7 @@ from pxr import UsdPhysics import isaaclab.sim as sim_utils +from isaaclab.cloner import resolve_clone_plan_source from isaaclab.sensors.ray_caster.base_ray_caster import BaseRayCaster from isaaclab.sensors.ray_caster.kernels import ( ALIGNMENT_BASE, @@ -63,14 +64,8 @@ def _gather_pose_by_index_kernel( quat_dst[i] = quat_src[src_idx] -def _find_physics_ancestor(prim): - """Return the nearest rigid-body ancestor for a sensor or target prim.""" - ancestor = prim - while ancestor and ancestor.IsValid() and ancestor.GetPath().pathString != "/": - if ancestor.HasAPI(UsdPhysics.RigidBodyAPI): - return ancestor - ancestor = ancestor.GetParent() - return None +def _has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) def _newton_body_pattern(body_path: str) -> str: @@ -105,48 +100,66 @@ def __init__(self: Any, cfg): """Register sensor and dynamic target sites before cloning occurs.""" super().__init__(cfg) # pyright: ignore[reportCallIssue] self._sensor_site_labels = self._register_sites_for_expr(self.cfg.prim_path) - self._tracked_site_labels_by_expr: dict[str | tuple[str, ...], list[str]] = {} + self._tracked_site_labels_by_target: dict[tuple[str, ...], list[str]] = {} for target_cfg in getattr(self, "_raycast_targets_cfg", []): if target_cfg.track_mesh_transforms: owner_exprs = self._resolve_target_owner_exprs(target_cfg.prim_expr) labels = self._register_target_sites_for_exprs(owner_exprs) - self._tracked_site_labels_by_expr[target_cfg.prim_expr] = labels - self._tracked_site_labels_by_expr[tuple(owner_exprs)] = labels + self._tracked_site_labels_by_target[tuple(owner_exprs)] = labels def _register_sites_for_expr(self, prim_expr: str) -> list[str]: """Register Newton sites for a prim expression and return site labels.""" - prims = sim_utils.find_matching_prims(prim_expr) - labels: list[str] = [] - if len(prims) == 0: - identity = wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat(0.0, 0.0, 0.0, 1.0)) - return [NewtonManager.cl_register_site(_newton_body_pattern(prim_expr), identity)] + identity = wp.transform(wp.vec3(0.0, 0.0, 0.0), wp.quat(0.0, 0.0, 0.0, 1.0)) + attach_expr = prim_expr + if prim_expr.rsplit("/", 1)[-1].lower() in ("camera", "raycaster"): + attach_expr = prim_expr.rsplit("/", 1)[0] - for prim in prims: - body = _find_physics_ancestor(prim) - if body is None: - pos, quat = sim_utils.resolve_prim_pose(prim) - xform = wp.transform(wp.vec3(*[float(v) for v in pos]), wp.quat(*[float(v) for v in quat])) - labels.append(NewtonManager.cl_register_site(None, xform)) - else: - pos, quat = sim_utils.resolve_prim_pose(prim, body) - xform = wp.transform(wp.vec3(*[float(v) for v in pos]), wp.quat(*[float(v) for v in quat])) - labels.append(NewtonManager.cl_register_site(_newton_body_pattern(str(body.GetPath())), xform)) - # Keep the first copy of each label; cloned envs can report the same prototype site more than once. - return list(dict.fromkeys(labels)) + plan = sim_utils.SimulationContext.instance().get_clone_plan() + if plan is not None: + for destination_template in plan.destinations: + if "{}" not in destination_template: + continue + destination_prefix, _ = destination_template.split("{}", 1) + if attach_expr.startswith(destination_prefix) and "/" not in attach_expr[len(destination_prefix) :]: + return [NewtonManager.cl_register_site(None, identity, per_world=True)] + + return [NewtonManager.cl_register_site(_newton_body_pattern(attach_expr), identity)] def _resolve_target_owner_exprs(self, prim_expr: str) -> list[str]: """Resolve mesh target expressions to owning rigid-body expressions.""" + plan = sim_utils.SimulationContext.instance().get_clone_plan() + resolved = resolve_clone_plan_source(prim_expr, plan) if plan is not None else None + if resolved is not None: + source_path, dest_glob, asset_suffix = resolved + walk_root = source_path + asset_suffix + source_prims = sim_utils.find_matching_prims(walk_root) + if not source_prims: + raise RuntimeError(f"No ClonePlan source prims matched '{walk_root}'.") + owner_exprs: list[str] = [] + for source_prim in source_prims: + body = sim_utils.get_first_matching_ancestor_prim(source_prim.GetPath(), predicate=_has_rigid_body_api) + if body is None: + raise RuntimeError( + f"Cannot track non-physics ray-cast target '{prim_expr}' with Newton. " + "Set track_mesh_transforms=False for static targets, or apply RigidBodyAPI" + " to dynamic targets." + ) + owner_prim_path = str(body.GetPath()) + owner_exprs.append(dest_glob + owner_prim_path[len(source_path) :]) + return list(dict.fromkeys(owner_exprs)) + + # Legacy fallback: no clone plan, or the target is not owned by any plan row. prims = sim_utils.find_matching_prims(prim_expr) if len(prims) == 0: return [_newton_body_pattern(prim_expr)] - - owner_exprs: list[str] = [] + owner_exprs = [] for prim in prims: - body = _find_physics_ancestor(prim) + body = sim_utils.get_first_matching_ancestor_prim(prim.GetPath(), predicate=_has_rigid_body_api) if body is None: raise RuntimeError( f"Cannot track non-physics ray-cast target '{prim_expr}' with Newton. " - "Set track_mesh_transforms=False for static targets, or apply RigidBodyAPI to dynamic targets." + "Set track_mesh_transforms=False for static targets, or apply RigidBodyAPI" + " to dynamic targets." ) owner_exprs.append(_newton_body_pattern(str(body.GetPath()))) return list(dict.fromkeys(owner_exprs)) @@ -227,12 +240,8 @@ def get_world_poses(self: Any, indices=None): def _create_tracked_target_view(self: Any, target_prim_path: str | list[str]): """Resolve dynamic multi-mesh target sites to raw Newton site indices.""" - target_key = tuple(target_prim_path) if isinstance(target_prim_path, list) else target_prim_path - labels = self._tracked_site_labels_by_expr.get(target_key) - if labels is None: - target_exprs = target_prim_path if isinstance(target_prim_path, list) else [target_prim_path] - labels = self._register_target_sites_for_exprs([_newton_body_pattern(expr) for expr in target_exprs]) - self._tracked_site_labels_by_expr[target_key] = labels + target_exprs = target_prim_path if isinstance(target_prim_path, list) else [target_prim_path] + labels = self._tracked_site_labels_by_target[tuple(target_exprs)] site_indices = self._resolve_site_indices(labels, str(target_prim_path), self._num_envs) return wp.array(site_indices, dtype=wp.int32, device=self._device) diff --git a/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py index c7758f0cbe42..d22d6537f3a0 100644 --- a/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py +++ b/source/isaaclab_newton/isaaclab_newton/sim/views/newton_site_frame_view.py @@ -3,7 +3,7 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Newton-backed FrameView — Warp-native, GPU-resident pose queries.""" +"""Newton-backed FrameView using Newton body labels and injected sites.""" from __future__ import annotations @@ -11,11 +11,13 @@ import warp as wp -from pxr import Gf, Usd, UsdGeom +from pxr import UsdPhysics import isaaclab.sim as sim_utils +from isaaclab.cloner.cloner_utils import iter_clone_plan_matches from isaaclab.physics import PhysicsEvent from isaaclab.sim.views.base_frame_view import BaseFrameView +from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.warp import ProxyArray from isaaclab_newton.physics.newton_manager import NewtonManager @@ -25,47 +27,8 @@ WORLD_BODY_INDEX = -1 -# ------------------------------------------------------------------ -# Warp kernels -# ------------------------------------------------------------------ - - @wp.kernel def _compute_site_world_transforms( - body_q: wp.array(dtype=wp.transformf), - site_body: wp.array(dtype=wp.int32), - site_local: wp.array(dtype=wp.transformf), - out_pos: wp.array(dtype=wp.vec3f), - out_quat: wp.array(dtype=wp.vec4f), -): - """Compute world-space transforms for every site in the view. - - For each site *i*, computes ``world = body_q[site_body[i]] * site_local[i]`` - and splits the result into position and quaternion outputs. When - ``site_body[i] == -1`` the site is world-attached and ``site_local[i]`` is - returned directly. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - A value of ``-1`` indicates a world-attached site. - site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. - out_pos: Output world positions [m], shape ``[num_sites]``. - out_quat: Output world orientations as ``(qx, qy, qz, qw)``, shape ``[num_sites]``. - """ - i = wp.tid() - bid = site_body[i] - if bid == -1: - world = site_local[i] - else: - world = wp.transform_multiply(body_q[bid], site_local[i]) - out_pos[i] = wp.transform_get_translation(world) - q = wp.transform_get_rotation(world) - out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) - - -@wp.kernel -def _compute_site_world_transforms_indexed( body_q: wp.array(dtype=wp.transformf), site_body: wp.array(dtype=wp.int32), site_local: wp.array(dtype=wp.transformf), @@ -73,24 +36,11 @@ def _compute_site_world_transforms_indexed( out_pos: wp.array(dtype=wp.vec3f), out_quat: wp.array(dtype=wp.vec4f), ): - """Indexed variant of :func:`_compute_site_world_transforms`. - - Only computes world transforms for the subset of sites selected by - ``indices``. Thread *i* reads ``indices[i]`` to obtain the site index, - then writes the result to ``out_pos[i]`` / ``out_quat[i]``. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. - indices: Site indices to query, shape ``[M]``. - out_pos: Output world positions [m], shape ``[M]``. - out_quat: Output world orientations as ``(qx, qy, qz, qw)``, shape ``[M]``. - """ + """Compute world-space transforms for selected sites.""" i = wp.tid() si = indices[i] bid = site_body[si] - if bid == -1: + if bid == WORLD_BODY_INDEX: world = site_local[si] else: world = wp.transform_multiply(body_q[bid], site_local[si]) @@ -100,169 +50,23 @@ def _compute_site_world_transforms_indexed( @wp.kernel -def _gather_scales( - shape_scale: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - site_body: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - out_scales: wp.array(dtype=wp.vec3f), -): - """Gather per-site scales from collision shapes on the same body. - - For each site *i*, linearly scans all shapes to find the first one whose - ``shape_body`` matches ``site_body[i]`` and copies its scale. Falls back - to ``(1, 1, 1)`` if no shape is found on that body. - - Args: - shape_scale: Per-shape scale vectors from the Newton model, shape ``[num_shapes]``. - shape_body: Per-shape parent body index, shape ``[num_shapes]``. - site_body: Per-site body index, shape ``[num_sites]``. - num_shapes: Total number of shapes in the model. - out_scales: Output scale per site, shape ``[num_sites]``. - """ - i = wp.tid() - bid = site_body[i] - found = int(0) - for s in range(num_shapes): - if shape_body[s] == bid and found == 0: - out_scales[i] = shape_scale[s] - found = 1 - if found == 0: - out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) - - -@wp.kernel -def _gather_scales_indexed( - shape_scale: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - site_body: wp.array(dtype=wp.int32), - indices: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - out_scales: wp.array(dtype=wp.vec3f), -): - """Indexed variant of :func:`_gather_scales`. - - Args: - shape_scale: Per-shape scale vectors from the Newton model, shape ``[num_shapes]``. - shape_body: Per-shape parent body index, shape ``[num_shapes]``. - site_body: Per-site body index, shape ``[num_sites]``. - indices: Site indices to query, shape ``[M]``. - num_shapes: Total number of shapes in the model. - out_scales: Output scale per queried site, shape ``[M]``. - """ - i = wp.tid() - si = indices[i] - bid = site_body[si] - found = int(0) - for s in range(num_shapes): - if shape_body[s] == bid and found == 0: - out_scales[i] = shape_scale[s] - found = 1 - if found == 0: - out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) - - -@wp.kernel -def _scatter_scales( - site_body: wp.array(dtype=wp.int32), - new_scales: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - shape_scale: wp.array(dtype=wp.vec3f), -): - """Scatter per-site scales to all collision shapes on the same body. - - For each site *i*, writes ``new_scales[i]`` to every shape whose - ``shape_body`` matches ``site_body[i]``. Multiple shapes on the same - body all receive the same scale. - - Args: - site_body: Per-site body index, shape ``[num_sites]``. - new_scales: New scale to apply per site, shape ``[num_sites]``. - shape_body: Per-shape parent body index, shape ``[num_shapes]``. - num_shapes: Total number of shapes in the model. - shape_scale: Per-shape scale vectors to write into (modified in-place), - shape ``[num_shapes]``. - """ - i = wp.tid() - bid = site_body[i] - for s in range(num_shapes): - if shape_body[s] == bid: - shape_scale[s] = new_scales[i] - - -@wp.kernel -def _scatter_scales_indexed( - site_body: wp.array(dtype=wp.int32), +def _gather_site_local_transforms( + site_local: wp.array(dtype=wp.transformf), indices: wp.array(dtype=wp.int32), - new_scales: wp.array(dtype=wp.vec3f), - shape_body: wp.array(dtype=wp.int32), - num_shapes: wp.int32, - shape_scale: wp.array(dtype=wp.vec3f), + out_pos: wp.array(dtype=wp.vec3f), + out_quat: wp.array(dtype=wp.vec4f), ): - """Indexed variant of :func:`_scatter_scales`. - - Args: - site_body: Per-site body index, shape ``[num_sites]``. - indices: Site indices to update, shape ``[M]``. - new_scales: New scale to apply per selected site, shape ``[M]``. - shape_body: Per-shape parent body index, shape ``[num_shapes]``. - num_shapes: Total number of shapes in the model. - shape_scale: Per-shape scale vectors to write into (modified in-place), - shape ``[num_shapes]``. - """ + """Gather local transforms for selected sites.""" i = wp.tid() si = indices[i] - bid = site_body[si] - for s in range(num_shapes): - if shape_body[s] == bid: - shape_scale[s] = new_scales[i] - - -# ------------------------------------------------------------------ -# World-pose site_local write kernels -# ------------------------------------------------------------------ + local_tf = site_local[si] + out_pos[i] = wp.transform_get_translation(local_tf) + q = wp.transform_get_rotation(local_tf) + out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) @wp.kernel def _write_site_local_from_world_poses( - body_q: wp.array(dtype=wp.transformf), - site_body: wp.array(dtype=wp.int32), - world_pos: wp.array(dtype=wp.vec3f), - world_quat: wp.array(dtype=wp.vec4f), - site_local: wp.array(dtype=wp.transformf), -): - """Update site local offsets so that the sites reach desired world poses. - - For each site *i*, computes - ``site_local[i] = inv(body_q[site_body[i]]) * desired_world`` so that - a subsequent ``body_q[bid] * site_local[i]`` yields the requested world - pose. For world-attached sites (``site_body[i] == -1``) the desired world - transform is written directly into ``site_local[i]``. - - Does **not** modify ``body_q``. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - world_pos: Desired world positions [m], shape ``[num_sites]``. - world_quat: Desired world orientations as ``(qx, qy, qz, qw)``, shape ``[num_sites]``. - site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. - """ - i = wp.tid() - w_pos = world_pos[i] - w_q = world_quat[i] - desired_world = wp.transform(w_pos, wp.quatf(w_q[0], w_q[1], w_q[2], w_q[3])) - - bid = site_body[i] - if bid == -1: - site_local[i] = desired_world - else: - site_local[i] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) - - -@wp.kernel -def _write_site_local_from_world_poses_indexed( body_q: wp.array(dtype=wp.transformf), site_body: wp.array(dtype=wp.int32), indices: wp.array(dtype=wp.int32), @@ -270,16 +74,7 @@ def _write_site_local_from_world_poses_indexed( world_quat: wp.array(dtype=wp.vec4f), site_local: wp.array(dtype=wp.transformf), ): - """Indexed variant of :func:`_write_site_local_from_world_poses`. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - indices: Site indices to update, shape ``[M]``. - world_pos: Desired world positions [m], shape ``[M]``. - world_quat: Desired world orientations as ``(qx, qy, qz, qw)``, shape ``[M]``. - site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. - """ + """Update local offsets so selected sites reach desired world poses.""" i = wp.tid() si = indices[i] w_pos = world_pos[i] @@ -287,435 +82,358 @@ def _write_site_local_from_world_poses_indexed( desired_world = wp.transform(w_pos, wp.quatf(w_q[0], w_q[1], w_q[2], w_q[3])) bid = site_body[si] - if bid == -1: + if bid == WORLD_BODY_INDEX: site_local[si] = desired_world else: site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) -# ------------------------------------------------------------------ -# Local-pose Warp kernels -# ------------------------------------------------------------------ - - @wp.kernel -def _compute_site_local_transforms( - body_q: wp.array(dtype=wp.transformf), - site_body: wp.array(dtype=wp.int32), +def _write_site_local_from_local_poses( + indices: wp.array(dtype=wp.int32), + local_pos: wp.array(dtype=wp.vec3f), + local_quat: wp.array(dtype=wp.vec4f), site_local: wp.array(dtype=wp.transformf), - parent_site_body: wp.array(dtype=wp.int32), - parent_site_local: wp.array(dtype=wp.transformf), - out_pos: wp.array(dtype=wp.vec3f), - out_quat: wp.array(dtype=wp.vec4f), ): - """Compute parent-relative transforms for every site in the view. - - For each site *i*, computes the world pose of both the site and its USD - parent, then returns ``inv(parent_world) * prim_world``. When - ``site_body[i] == -1`` the site is world-attached and ``site_local[i]`` - is used as the world transform directly. The same convention applies to - the parent arrays. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. - parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. - parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. - out_pos: Output parent-relative positions [m], shape ``[num_sites]``. - out_quat: Output parent-relative orientations as ``(qx, qy, qz, qw)``, - shape ``[num_sites]``. - """ + """Update local offsets for selected sites.""" i = wp.tid() - prim_bid = site_body[i] - if prim_bid == -1: - prim_world = site_local[i] - else: - prim_world = wp.transform_multiply(body_q[prim_bid], site_local[i]) - - parent_bid = parent_site_body[i] - if parent_bid == -1: - parent_world = parent_site_local[i] - else: - parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[i]) - - local_tf = wp.transform_multiply(wp.transform_inverse(parent_world), prim_world) - out_pos[i] = wp.transform_get_translation(local_tf) - q = wp.transform_get_rotation(local_tf) - out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) + si = indices[i] + l_pos = local_pos[i] + l_q = local_quat[i] + site_local[si] = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) @wp.kernel -def _compute_site_local_transforms_indexed( - body_q: wp.array(dtype=wp.transformf), +def _gather_scales( + shape_scale: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), site_body: wp.array(dtype=wp.int32), - site_local: wp.array(dtype=wp.transformf), - parent_site_body: wp.array(dtype=wp.int32), - parent_site_local: wp.array(dtype=wp.transformf), indices: wp.array(dtype=wp.int32), - out_pos: wp.array(dtype=wp.vec3f), - out_quat: wp.array(dtype=wp.vec4f), + num_shapes: wp.int32, + out_scales: wp.array(dtype=wp.vec3f), ): - """Indexed variant of :func:`_compute_site_local_transforms`. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - site_local: Per-site local offset relative to its parent body, shape ``[num_sites]``. - parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. - parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. - indices: Site indices to query, shape ``[M]``. - out_pos: Output parent-relative positions [m], shape ``[M]``. - out_quat: Output parent-relative orientations as ``(qx, qy, qz, qw)``, - shape ``[M]``. - """ + """Gather per-site scales from collision shapes on the same body.""" i = wp.tid() si = indices[i] - prim_bid = site_body[si] - if prim_bid == -1: - prim_world = site_local[si] - else: - prim_world = wp.transform_multiply(body_q[prim_bid], site_local[si]) - - parent_bid = parent_site_body[si] - if parent_bid == -1: - parent_world = parent_site_local[si] - else: - parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[si]) - - local_tf = wp.transform_multiply(wp.transform_inverse(parent_world), prim_world) - out_pos[i] = wp.transform_get_translation(local_tf) - q = wp.transform_get_rotation(local_tf) - out_quat[i] = wp.vec4f(q[0], q[1], q[2], q[3]) - - -@wp.kernel -def _write_site_local_from_local_poses( - body_q: wp.array(dtype=wp.transformf), - site_body: wp.array(dtype=wp.int32), - parent_site_body: wp.array(dtype=wp.int32), - parent_site_local: wp.array(dtype=wp.transformf), - local_pos: wp.array(dtype=wp.vec3f), - local_quat: wp.array(dtype=wp.vec4f), - site_local: wp.array(dtype=wp.transformf), -): - """Update site local offsets so that sites reach desired parent-relative poses. - - For each site *i*, reconstructs the desired world pose as - ``parent_world * desired_local``, then solves for the body-relative offset: - ``site_local[i] = inv(body_q[bid]) * desired_world``. For world-attached - sites (``site_body[i] == -1``) the world transform is written directly. - - Does **not** modify ``body_q``. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. - parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. - local_pos: Desired parent-relative positions [m], shape ``[num_sites]``. - local_quat: Desired parent-relative orientations as ``(qx, qy, qz, qw)``, - shape ``[num_sites]``. - site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. - """ - i = wp.tid() - parent_bid = parent_site_body[i] - if parent_bid == -1: - parent_world = parent_site_local[i] - else: - parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[i]) - - l_pos = local_pos[i] - l_q = local_quat[i] - local_tf = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) - desired_world = wp.transform_multiply(parent_world, local_tf) - - bid = site_body[i] - if bid == -1: - site_local[i] = desired_world - else: - site_local[i] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) + bid = site_body[si] + found = int(0) + for s in range(num_shapes): + if shape_body[s] == bid and found == 0: + out_scales[i] = shape_scale[s] + found = 1 + if found == 0: + out_scales[i] = wp.vec3f(1.0, 1.0, 1.0) @wp.kernel -def _write_site_local_from_local_poses_indexed( - body_q: wp.array(dtype=wp.transformf), +def _scatter_scales( site_body: wp.array(dtype=wp.int32), - parent_site_body: wp.array(dtype=wp.int32), - parent_site_local: wp.array(dtype=wp.transformf), indices: wp.array(dtype=wp.int32), - local_pos: wp.array(dtype=wp.vec3f), - local_quat: wp.array(dtype=wp.vec4f), - site_local: wp.array(dtype=wp.transformf), + new_scales: wp.array(dtype=wp.vec3f), + shape_body: wp.array(dtype=wp.int32), + num_shapes: wp.int32, + shape_scale: wp.array(dtype=wp.vec3f), ): - """Indexed variant of :func:`_write_site_local_from_local_poses`. - - Args: - body_q: Rigid-body world transforms from the Newton state, shape ``[num_bodies]``. - site_body: Per-site body index (flat model-level), shape ``[num_sites]``. - parent_site_body: Per-site USD-parent body index, shape ``[num_sites]``. - parent_site_local: Per-site USD-parent local offset, shape ``[num_sites]``. - indices: Site indices to update, shape ``[M]``. - local_pos: Desired parent-relative positions [m], shape ``[M]``. - local_quat: Desired parent-relative orientations as ``(qx, qy, qz, qw)``, - shape ``[M]``. - site_local: Per-site local offset (modified in-place), shape ``[num_sites]``. - """ + """Scatter per-site scales to collision shapes on the same body.""" i = wp.tid() si = indices[i] - parent_bid = parent_site_body[si] - if parent_bid == -1: - parent_world = parent_site_local[si] - else: - parent_world = wp.transform_multiply(body_q[parent_bid], parent_site_local[si]) - - l_pos = local_pos[i] - l_q = local_quat[i] - local_tf = wp.transform(l_pos, wp.quatf(l_q[0], l_q[1], l_q[2], l_q[3])) - desired_world = wp.transform_multiply(parent_world, local_tf) - bid = site_body[si] - if bid == -1: - site_local[si] = desired_world - else: - site_local[si] = wp.transform_multiply(wp.transform_inverse(body_q[bid]), desired_world) - - -# ------------------------------------------------------------------ -# View class -# ------------------------------------------------------------------ + for s in range(num_shapes): + if shape_body[s] == bid: + shape_scale[s] = new_scales[i] class NewtonSiteFrameView(BaseFrameView): - """Batched prim view for non-physics prims tracked as sites on Newton bodies. - - Each matched USD prim must be a **non-physics** prim (camera, sensor, - Xform marker, etc.) that sits as a child of a Newton rigid body in the - USD hierarchy. The prim path must **not** resolve directly to a physics - body or collision shape -- those are owned by Newton and should be - accessed through :class:`~isaaclab_newton.assets.Articulation` or - :class:`~isaaclab_newton.assets.RigidObject` instead. - - At init time each prim is resolved to a ``(body_index, site_local)`` - pair via ancestor walk: the nearest ancestor that appears in - ``model.body_label`` becomes the attachment body, and the relative USD - transform becomes the site offset. If no body ancestor exists the prim - is attached to the world frame (``body_index = -1``). - - World poses are computed on GPU as - ``body_q[body_index] * site_local`` via a Warp kernel. Both - ``set_world_poses`` and ``set_local_poses`` update ``site_local`` -- - neither touches ``body_q``. - - Pose getters return :class:`~isaaclab.utils.warp.ProxyArray`. Setters accept ``wp.array``. - - Raises: - ValueError: If any matched prim resolves to a Newton physics body - or collision shape. - """ + """Batched Newton site view for non-physics frames. - def __init__(self, prim_path: str, device: str = "cpu", stage: Usd.Stage | None = None, **kwargs): - """Initialize the Newton site-based frame view. - - Resolves all USD prims matching ``prim_path`` and, for each one, walks - the USD ancestor hierarchy to find the nearest Newton rigid body. The - relative transform between the prim and its ancestor body becomes the - site's local offset. + The public construction contract matches the generic :class:`FrameView`: + callers provide a prim expression and the backend resolves the source prim + into Newton body-local or world-local sites. + """ - If the Newton model is already finalized the view initializes - immediately; otherwise initialization is deferred to a - :attr:`PhysicsEvent.PHYSICS_READY` callback. + def __init__( + self, + prim_path: str | list[str], + device: str = "cpu", + validate_xform_ops: bool = True, + stage: object | None = None, + **kwargs, + ): + """Initialize the Newton site frame view. Args: - prim_path: USD prim path pattern (may contain regex). - device: Warp device for GPU arrays (e.g. ``"cuda:0"``). - stage: USD stage to search. Defaults to the current stage. - **kwargs: Unused; accepted for interface compatibility with other - :class:`~isaaclab.sim.views.BaseFrameView` backends. + prim_path: User-facing frame path pattern, or list of patterns. + device: Warp device for GPU arrays. + validate_xform_ops: Whether to validate source USD xform ops. + stage: USD stage that contains the source prims. + **kwargs: Unused. """ - self._prim_path = prim_path + del kwargs + + self._prim_paths = [prim_path] if isinstance(prim_path, str) else list(prim_path) + self._prim_path = prim_path if isinstance(prim_path, str) else ", ".join(self._prim_paths) self._device = device + self._prims = [] stage = sim_utils.get_current_stage() if stage is None else stage - self._prims: list[Usd.Prim] = sim_utils.find_matching_prims(prim_path, stage=stage) + self._site_specs = self._resolve_site_specs(stage, validate_xform_ops) + self._site_labels: list[str] = [] + self._site_body: wp.array | None = None + self._site_local: wp.array | None = None + self._site_indices: wp.array | None = None + self._pos_buf: wp.array | None = None + self._quat_buf: wp.array | None = None + self._local_pos_buf: wp.array | None = None + self._local_quat_buf: wp.array | None = None + self._pos_ta: ProxyArray | None = None + self._quat_ta: ProxyArray | None = None + self._local_pos_ta: ProxyArray | None = None + self._local_quat_ta: ProxyArray | None = None + self._count = 0 model = NewtonManager.get_model() if model is not None: - self._initialize_impl(model) + self._initialize_from_specs(model) else: + for body_patterns, xform, per_world, _env_ids in self._site_specs: + if body_patterns is None: + self._site_labels.append(NewtonManager.cl_register_site(None, xform, per_world=per_world)) + else: + for body_pattern in body_patterns: + self._site_labels.append(NewtonManager.cl_register_site(body_pattern, xform)) self._physics_ready_handle = NewtonManager.register_callback( - self._on_physics_ready, PhysicsEvent.PHYSICS_READY, name=f"site_view_{prim_path}" + self._on_physics_ready, PhysicsEvent.PHYSICS_READY, name=f"site_view_{self._prim_path}" ) - def _on_physics_ready(self, _event) -> None: - """Callback invoked when the Newton model becomes available.""" - self._initialize_impl(NewtonManager.get_model()) + def _resolve_site_specs( + self, stage, validate_xform_ops: bool + ) -> list[tuple[tuple[str, ...] | None, wp.transform, bool, tuple[int, ...] | None]]: + """Resolve source prims into Newton site registration specs.""" + plan = sim_utils.SimulationContext.instance().get_clone_plan() + model = NewtonManager.get_model() + body_labels = list(model.body_label) if model is not None else () + shape_labels = list(model.shape_label) if model is not None else () + use_clone_body_pattern = model is None + specs: list[tuple[tuple[str, ...] | None, wp.transform, bool, tuple[int, ...] | None]] = [] - def _initialize_impl(self, model) -> None: - """Resolve USD prims to Newton body indices and allocate GPU buffers.""" - body_labels = list(model.body_label) - body_label_set = set(body_labels) - body_label_to_idx = {path: idx for idx, path in enumerate(body_labels)} - shape_label_set = set(model.shape_label) + for path_expr in self._prim_paths: + if resolve_matching_names(path_expr, body_labels, raise_when_no_match=False)[1]: + raise ValueError( + f"FrameView prim '{path_expr}' is a Newton physics body. " + "FrameView should only be used for non-physics frames." + ) + if resolve_matching_names(path_expr, shape_labels, raise_when_no_match=False)[1]: + raise ValueError( + f"FrameView prim '{path_expr}' is a Newton collision shape. " + "FrameView should only be used for non-physics frames." + ) + matches = tuple(iter_clone_plan_matches(plan, path_expr)) if plan is not None else () + if matches: + for source_root, destination_template, source_path, env_ids in matches: + source_prim = None + if not any(token in source_path for token in "*[]()+?|\\"): + source_prim = stage.GetPrimAtPath(source_path) + if source_prim is None or not source_prim.IsValid(): + source_prim = sim_utils.find_first_matching_prim(source_path, stage) + if source_prim is None or not source_prim.IsValid(): + raise RuntimeError(f"FrameView '{path_expr}' could not resolve source prim '{source_path}'.") + specs.append( + self._resolve_source_prim( + source_prim, + validate_xform_ops, + source_root, + destination_template, + env_ids, + use_clone_body_pattern, + stage, + ) + ) + continue + + prim = sim_utils.find_first_matching_prim(path_expr, stage) + if prim is None or not prim.IsValid(): + raise RuntimeError(f"FrameView '{path_expr}' could not resolve a source prim.") + specs.append( + self._resolve_source_prim(prim, validate_xform_ops, None, None, None, use_clone_body_pattern, stage) + ) + + return specs + + def _resolve_source_prim( + self, + prim, + validate_xform_ops: bool, + source_root: str | None, + destination_template: str | None, + env_ids: tuple[int, ...] | None, + use_clone_body_pattern: bool, + stage, + ) -> tuple[tuple[str, ...] | None, wp.transform, bool, tuple[int, ...] | None]: + """Resolve one source prim into body patterns and a local frame.""" + prim_path = prim.GetPath().pathString + if prim.HasAPI(UsdPhysics.RigidBodyAPI) or prim.HasAPI(UsdPhysics.ArticulationRootAPI): + raise ValueError( + f"FrameView prim '{prim_path}' is a Newton physics body. " + "FrameView should only be used for non-physics frames." + ) + if validate_xform_ops: + sim_utils.standardize_xform_ops(prim) + if not sim_utils.validate_standard_xform_ops(prim): + raise ValueError(f"FrameView prim '{prim_path}' does not have standard xform ops.") + + body_prim = prim.GetParent() + while body_prim and body_prim.IsValid(): + if body_prim.HasAPI(UsdPhysics.RigidBodyAPI) or body_prim.HasAPI(UsdPhysics.ArticulationRootAPI): + pos, quat = sim_utils.resolve_prim_pose(prim, body_prim) + body_path = body_prim.GetPath().pathString + if source_root is not None and destination_template is not None: + assert env_ids is not None + if body_path == source_root: + suffix = "" + elif body_path.startswith(source_root + "/"): + suffix = body_path[len(source_root) :] + elif source_root.startswith(body_path + "/"): + suffix = source_root[len(body_path) :] + if use_clone_body_pattern: + destination_root = destination_template.format(".*") + if not destination_root.endswith(suffix): + raise RuntimeError( + f"FrameView destination root '{destination_root}' does not end with '{suffix}'." + ) + return (destination_root[: -len(suffix)],), wp.transform(pos, quat), False, env_ids + body_patterns = [] + for env_id in env_ids: + destination_root = destination_template.format(env_id) + if not destination_root.endswith(suffix): + raise RuntimeError( + f"FrameView destination root '{destination_root}' does not end with '{suffix}'." + ) + body_patterns.append(destination_root[: -len(suffix)]) + return tuple(body_patterns), wp.transform(pos, quat), False, env_ids + else: + raise RuntimeError(f"FrameView source body '{body_path}' is not under '{source_root}'.") + if use_clone_body_pattern: + body_patterns = (destination_template.format(".*") + suffix,) + else: + body_patterns = tuple(destination_template.format(env_id) + suffix for env_id in env_ids) + else: + body_patterns = (body_path,) + return body_patterns, wp.transform(pos, quat), False, env_ids + body_prim = body_prim.GetParent() + + ref_prim = stage.GetPrimAtPath(source_root) if source_root is not None else None + pos, quat = sim_utils.resolve_prim_pose(prim, ref_prim if ref_prim and ref_prim.IsValid() else None) + return None, wp.transform(pos, quat), source_root is not None, env_ids - xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + def _on_physics_ready(self, _event) -> None: + """Callback invoked when the Newton model becomes available.""" + self._initialize_from_site_map(NewtonManager.get_model()) + def _initialize_from_site_map(self, model) -> None: + """Initialize arrays from injected Newton sites.""" + site_map = NewtonManager._cl_site_index_map + body_t = wp.to_torch(model.shape_body) + xform_t = wp.to_torch(model.shape_transform) site_bodies: list[int] = [] site_locals: list[list[float]] = [] - parent_bodies: list[int] = [] - parent_locals: list[list[float]] = [] - identity_xform = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] - resolve_cache: dict[str, tuple[int, list[float]]] = {} + for site_label in self._site_labels: + global_idx, per_world = site_map[site_label] + site_indices = ( + [global_idx] if per_world is None else [site_idx for sites in per_world for site_idx in sites] + ) + for site_idx in site_indices: + site_bodies.append(int(body_t[site_idx].item())) + site_locals.append([float(v) for v in xform_t[site_idx].tolist()]) - for prim in self._prims: - pp = prim.GetPath().pathString - if pp in body_label_set: - raise ValueError( - f"FrameView prim '{pp}' is a Newton physics body. " - "FrameView should only be used for non-physics prims (cameras, sensors, Xform markers). " - "Use Articulation or RigidObject APIs to control physics bodies." - ) - if pp in shape_label_set: - raise ValueError( - f"FrameView prim '{pp}' is a Newton collision shape. " - "FrameView should only be used for non-physics prims (cameras, sensors, Xform markers). " - "Use Articulation or RigidObject APIs to control collision shapes." - ) + self._create_buffers(site_bodies, site_locals) - body_idx, local_xform = self._resolve_ancestor_body(prim, body_label_to_idx, xform_cache) - site_bodies.append(body_idx) - site_locals.append(local_xform) - - parent = prim.GetParent() - if not parent or not parent.IsValid() or parent.GetPath().pathString == "/": - parent_bodies.append(WORLD_BODY_INDEX) - parent_locals.append(identity_xform) - else: - parent_path = parent.GetPath().pathString - if parent_path in resolve_cache: - pb_idx, pb_local = resolve_cache[parent_path] - elif parent_path in body_label_to_idx: - pb_idx = body_label_to_idx[parent_path] - pb_local = identity_xform - resolve_cache[parent_path] = (pb_idx, pb_local) - else: - pb_idx, pb_local = self._resolve_ancestor_body(parent, body_label_to_idx, xform_cache) - resolve_cache[parent_path] = (pb_idx, pb_local) - parent_bodies.append(pb_idx) - parent_locals.append(pb_local) + def _initialize_from_specs(self, model) -> None: + """Initialize arrays directly from resolved specs and Newton body labels.""" + body_labels = list(model.body_label) + site_bodies: list[int] = [] + site_locals: list[list[float]] = [] + for body_patterns, xform, per_world, env_ids in self._site_specs: + if body_patterns is None: + if per_world: + if NewtonManager._world_xforms is None: + raise RuntimeError(f"FrameView '{self._prim_path}' needs Newton cloned-world transforms.") + world_ids = range(len(NewtonManager._world_xforms)) if env_ids is None else env_ids + for world_id in world_ids: + world_xform = NewtonManager._world_xforms[world_id] + site_bodies.append(WORLD_BODY_INDEX) + site_locals.append([float(v) for v in wp.transform_multiply(world_xform, xform)]) + else: + site_bodies.append(WORLD_BODY_INDEX) + site_locals.append([float(v) for v in xform]) + continue + + for body_pattern in body_patterns: + matched_indices, _ = resolve_matching_names(body_pattern, body_labels, raise_when_no_match=False) + if not matched_indices: + raise ValueError( + f"FrameView '{self._prim_path}' body pattern '{body_pattern}' matched no Newton bodies." + ) + + for body_idx in matched_indices: + site_bodies.append(body_idx) + site_locals.append([float(v) for v in xform]) + + self._create_buffers(site_bodies, site_locals) + + def _create_buffers(self, site_bodies: list[int], site_locals: list[list[float]]) -> None: + """Allocate view buffers from body indices and local transforms.""" + self._count = len(site_bodies) device = self._device self._site_body = wp.array(site_bodies, dtype=wp.int32, device=device) - self._site_local = wp.array( - [wp.transform(*x) for x in site_locals], - dtype=wp.transformf, - device=device, - ) - self._parent_site_body = wp.array(parent_bodies, dtype=wp.int32, device=device) - self._parent_site_local = wp.array( - [wp.transform(*x) for x in parent_locals], - dtype=wp.transformf, - device=device, - ) - - self._pos_buf = wp.zeros(self.count, dtype=wp.vec3f, device=device) - self._quat_buf = wp.zeros(self.count, dtype=wp.vec4f, device=device) - self._local_pos_buf = wp.zeros(self.count, dtype=wp.vec3f, device=device) - self._local_quat_buf = wp.zeros(self.count, dtype=wp.vec4f, device=device) + self._site_local = wp.array([wp.transform(*x) for x in site_locals], dtype=wp.transformf, device=device) + self._site_indices = wp.array(list(range(self._count)), dtype=wp.int32, device=device) + self._pos_buf = wp.zeros(self._count, dtype=wp.vec3f, device=device) + self._quat_buf = wp.zeros(self._count, dtype=wp.vec4f, device=device) + self._local_pos_buf = wp.zeros(self._count, dtype=wp.vec3f, device=device) + self._local_quat_buf = wp.zeros(self._count, dtype=wp.vec4f, device=device) self._pos_ta = ProxyArray(self._pos_buf) self._quat_ta = ProxyArray(self._quat_buf) self._local_pos_ta = ProxyArray(self._local_pos_buf) self._local_quat_ta = ProxyArray(self._local_quat_buf) - @staticmethod - def _resolve_ancestor_body( - prim: Usd.Prim, - body_label_to_idx: dict[str, int], - xform_cache: UsdGeom.XformCache, - ) -> tuple[int, list[float]]: - """Walk USD ancestors to find the nearest Newton body and compute the relative local transform. - - Args: - prim: The USD prim to resolve. - body_label_to_idx: Dict mapping body prim paths to their Newton body indices. - xform_cache: USD xform cache for efficient transform lookups. - - Returns: - A tuple ``(body_index, local_xform_7)`` where *local_xform_7* is - ``[tx, ty, tz, qx, qy, qz, qw]``. If no body ancestor exists, - ``body_index`` is :data:`WORLD_BODY_INDEX` and the local transform - is the prim's world transform. - """ - prim_world_tf = xform_cache.GetLocalToWorldTransform(prim) - prim_world_tf.Orthonormalize() - - ancestor = prim.GetParent() - while ancestor and ancestor.IsValid() and ancestor.GetPath().pathString != "/": - ancestor_path = ancestor.GetPath().pathString - body_idx = body_label_to_idx.get(ancestor_path) - if body_idx is not None: - ancestor_world_tf = xform_cache.GetLocalToWorldTransform(ancestor) - ancestor_world_tf.Orthonormalize() - local_tf = prim_world_tf * ancestor_world_tf.GetInverse() - return body_idx, _gf_matrix_to_xform7(local_tf) - ancestor = ancestor.GetParent() - - return WORLD_BODY_INDEX, _gf_matrix_to_xform7(prim_world_tf) - @property def prims(self) -> list: - """List of USD prims being managed by this view.""" + """List of USD prims being managed by this view. + + Newton site views do not retain USD prim handles. + """ return self._prims @property def count(self) -> int: - """Number of prims in this view.""" - return len(self._prims) + """Number of frames in this view.""" + return self._count @property def device(self) -> str: - """Device where arrays are allocated (cpu or cuda).""" + """Device where arrays are allocated.""" return self._device - # ------------------------------------------------------------------ - # World poses - # ------------------------------------------------------------------ - def get_world_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get world-space positions and orientations. - - Args: - indices: Subset of sites to query. ``None`` means all sites. - - Returns: - A tuple ``(positions, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` - wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. - """ + """Get world-space positions and orientations.""" state = NewtonManager.get_state_0() - - if indices is not None: - n = len(indices) - pos_buf = wp.zeros(n, dtype=wp.vec3f, device=self._device) - quat_buf = wp.zeros(n, dtype=wp.vec4f, device=self._device) - wp.launch( - _compute_site_world_transforms_indexed, - dim=n, - inputs=[state.body_q, self._site_body, self._site_local, indices], - outputs=[pos_buf, quat_buf], - device=self._device, - ) - return ProxyArray(pos_buf), ProxyArray(quat_buf) + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + pos_buf = self._pos_buf if indices is None else wp.zeros(n, dtype=wp.vec3f, device=self._device) + quat_buf = self._quat_buf if indices is None else wp.zeros(n, dtype=wp.vec4f, device=self._device) wp.launch( _compute_site_world_transforms, - dim=self.count, - inputs=[state.body_q, self._site_body, self._site_local], - outputs=[self._pos_buf, self._quat_buf], + dim=n, + inputs=[state.body_q, self._site_body, self._site_local, site_indices], + outputs=[pos_buf, quat_buf], device=self._device, ) - return self._pos_ta, self._quat_ta + if indices is None: + return self._pos_ta, self._quat_ta + return ProxyArray(pos_buf), ProxyArray(quat_buf) def set_world_poses( self, @@ -723,24 +441,11 @@ def set_world_poses( orientations: wp.array | None = None, indices: wp.array | None = None, ) -> None: - """Set world-space positions and/or orientations. - - Updates the internal ``site_local`` offsets so that - ``body_q[body] * new_site_local`` yields the desired world pose. - Does **not** modify ``body_q``. - - Args: - positions: Desired world positions ``(M, 3)``. ``None`` leaves - positions unchanged. - orientations: Desired world quaternions ``(M, 4)`` as - ``(qx, qy, qz, qw)``. ``None`` leaves orientations unchanged. - indices: Subset of sites to update. ``None`` means all sites. - """ + """Set world-space positions and/or orientations.""" if positions is None and orientations is None: return state = NewtonManager.get_state_0() - if positions is None or orientations is None: cur_pos_ta, cur_quat_ta = self.get_world_poses(indices) if positions is None: @@ -748,74 +453,32 @@ def set_world_poses( if orientations is None: orientations = cur_quat_ta.warp - if indices is not None: - wp.launch( - _write_site_local_from_world_poses_indexed, - dim=len(indices), - inputs=[state.body_q, self._site_body, indices, positions, orientations, self._site_local], - device=self._device, - ) - else: - wp.launch( - _write_site_local_from_world_poses, - dim=self.count, - inputs=[state.body_q, self._site_body, positions, orientations, self._site_local], - device=self._device, - ) - - # ------------------------------------------------------------------ - # Local poses (parent-relative) - # ------------------------------------------------------------------ + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + wp.launch( + _write_site_local_from_world_poses, + dim=n, + inputs=[state.body_q, self._site_body, site_indices, positions, orientations, self._site_local], + device=self._device, + ) def get_local_poses(self, indices: wp.array | None = None) -> tuple[ProxyArray, ProxyArray]: - """Get parent-relative positions and orientations. - - Computes ``inv(parent_world) * prim_world`` for each site. - - Args: - indices: Subset of sites to query. ``None`` means all sites. - - Returns: - A tuple ``(translations, orientations)`` of :class:`~isaaclab.utils.warp.ProxyArray` - wrappers. Use ``.warp`` for the underlying ``wp.array`` or ``.torch`` for a - cached zero-copy ``torch.Tensor`` view. - """ - state = NewtonManager.get_state_0() - - if indices is not None: - n = len(indices) - pos_buf = wp.zeros(n, dtype=wp.vec3f, device=self._device) - quat_buf = wp.zeros(n, dtype=wp.vec4f, device=self._device) - wp.launch( - _compute_site_local_transforms_indexed, - dim=n, - inputs=[ - state.body_q, - self._site_body, - self._site_local, - self._parent_site_body, - self._parent_site_local, - indices, - ], - outputs=[pos_buf, quat_buf], - device=self._device, - ) - return ProxyArray(pos_buf), ProxyArray(quat_buf) + """Get body-local positions and orientations.""" + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + pos_buf = self._local_pos_buf if indices is None else wp.zeros(n, dtype=wp.vec3f, device=self._device) + quat_buf = self._local_quat_buf if indices is None else wp.zeros(n, dtype=wp.vec4f, device=self._device) wp.launch( - _compute_site_local_transforms, - dim=self.count, - inputs=[ - state.body_q, - self._site_body, - self._site_local, - self._parent_site_body, - self._parent_site_local, - ], - outputs=[self._local_pos_buf, self._local_quat_buf], + _gather_site_local_transforms, + dim=n, + inputs=[self._site_local, site_indices], + outputs=[pos_buf, quat_buf], device=self._device, ) - return self._local_pos_ta, self._local_quat_ta + if indices is None: + return self._local_pos_ta, self._local_quat_ta + return ProxyArray(pos_buf), ProxyArray(quat_buf) def set_local_poses( self, @@ -823,24 +486,10 @@ def set_local_poses( orientations: wp.array | None = None, indices: wp.array | None = None, ) -> None: - """Set parent-relative translations and/or orientations. - - Updates the internal ``site_local`` offsets so that - ``inv(parent_world) * (body_q[bid] * site_local)`` yields the desired - local pose. Does **not** modify ``body_q``. - - Args: - translations: Desired parent-relative translations ``(M, 3)``. - ``None`` leaves translations unchanged. - orientations: Desired parent-relative quaternions ``(M, 4)`` as - ``(qx, qy, qz, qw)``. ``None`` leaves orientations unchanged. - indices: Subset of sites to update. ``None`` means all sites. - """ + """Set body-local translations and/or orientations.""" if translations is None and orientations is None: return - state = NewtonManager.get_state_0() - if translations is None or orientations is None: cur_pos_ta, cur_quat_ta = self.get_local_poses(indices) if translations is None: @@ -848,104 +497,40 @@ def set_local_poses( if orientations is None: orientations = cur_quat_ta.warp - if indices is not None: - wp.launch( - _write_site_local_from_local_poses_indexed, - dim=len(indices), - inputs=[ - state.body_q, - self._site_body, - self._parent_site_body, - self._parent_site_local, - indices, - translations, - orientations, - self._site_local, - ], - device=self._device, - ) - else: - wp.launch( - _write_site_local_from_local_poses, - dim=self.count, - inputs=[ - state.body_q, - self._site_body, - self._parent_site_body, - self._parent_site_local, - translations, - orientations, - self._site_local, - ], - device=self._device, - ) - - # ------------------------------------------------------------------ - # Scales - # ------------------------------------------------------------------ + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + wp.launch( + _write_site_local_from_local_poses, + dim=n, + inputs=[site_indices, translations, orientations, self._site_local], + device=self._device, + ) def get_scales(self, indices: wp.array | None = None) -> wp.array: - """Get per-site scales by reading from the first collision shape on the same body. - - Args: - indices: Subset of sites to query. ``None`` means all sites. - - Returns: - A ``wp.array`` of shape ``(M, 3)``. - """ + """Get per-site scales by reading from the first collision shape on the same body.""" model = NewtonManager.get_model() num_shapes = model.shape_count - - if indices is not None: - n = len(indices) - out = wp.zeros(n, dtype=wp.vec3f, device=self._device) - wp.launch( - _gather_scales_indexed, - dim=n, - inputs=[model.shape_scale, model.shape_body, self._site_body, indices, num_shapes], - outputs=[out], - device=self._device, - ) - else: - out = wp.zeros(self.count, dtype=wp.vec3f, device=self._device) - wp.launch( - _gather_scales, - dim=self.count, - inputs=[model.shape_scale, model.shape_body, self._site_body, num_shapes], - outputs=[out], - device=self._device, - ) + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + out = wp.zeros(n, dtype=wp.vec3f, device=self._device) + wp.launch( + _gather_scales, + dim=n, + inputs=[model.shape_scale, model.shape_body, self._site_body, site_indices, num_shapes], + outputs=[out], + device=self._device, + ) return out def set_scales(self, scales: wp.array, indices: wp.array | None = None) -> None: - """Set per-site scales by writing to all collision shapes on the same body. - - Args: - scales: New scales ``(M, 3)`` as ``wp.array``. - indices: Subset of sites to update. ``None`` means all sites. - """ + """Set per-site scales by writing to all collision shapes on the same body.""" model = NewtonManager.get_model() num_shapes = model.shape_count - - if indices is not None: - wp.launch( - _scatter_scales_indexed, - dim=len(indices), - inputs=[self._site_body, indices, scales, model.shape_body, num_shapes, model.shape_scale], - device=self._device, - ) - else: - wp.launch( - _scatter_scales, - dim=self.count, - inputs=[self._site_body, scales, model.shape_body, num_shapes, model.shape_scale], - device=self._device, - ) - - -def _gf_matrix_to_xform7(mat: Gf.Matrix4d) -> list[float]: - """Convert a ``Gf.Matrix4d`` to ``[tx, ty, tz, qx, qy, qz, qw]``.""" - t = mat.ExtractTranslation() - q = mat.ExtractRotationQuat() - imag = q.GetImaginary() - return [float(t[0]), float(t[1]), float(t[2]), float(imag[0]), float(imag[1]), float(imag[2]), float(q.GetReal())] + site_indices = self._site_indices if indices is None else indices + n = self.count if indices is None else len(indices) + wp.launch( + _scatter_scales, + dim=n, + inputs=[self._site_body, site_indices, scales, model.shape_body, num_shapes, model.shape_scale], + device=self._device, + ) diff --git a/source/isaaclab_newton/test/sensors/test_joint_wrench_sensor.py b/source/isaaclab_newton/test/sensors/test_joint_wrench_sensor.py index c1f15291947e..793ca32de440 100644 --- a/source/isaaclab_newton/test/sensors/test_joint_wrench_sensor.py +++ b/source/isaaclab_newton/test/sensors/test_joint_wrench_sensor.py @@ -183,7 +183,7 @@ def test_multi_body_articulation(sim): def test_nested_articulation_root_resolution(sim): - """Sensor accepts an asset prim path whose articulation root is nested in the USD asset.""" + """Sensor covers a nested articulation root from the configured asset prefix.""" scene = InteractiveScene(_NestedRootAntSceneCfg(num_envs=1)) sim.reset() diff --git a/source/isaaclab_newton/test/sensors/test_site_injection.py b/source/isaaclab_newton/test/sensors/test_site_injection.py index 6bf40635a6c0..0b9fec3c8adc 100644 --- a/source/isaaclab_newton/test/sensors/test_site_injection.py +++ b/source/isaaclab_newton/test/sensors/test_site_injection.py @@ -86,7 +86,7 @@ def setup_method(self): def test_global_site_entry_is_int_none_tuple(self): xform = wp.transform() - NewtonManager._cl_pending_sites = {(None, tuple(xform)): ("ft_0", xform)} + NewtonManager._cl_pending_sites = {(None, False, tuple(xform)): ("ft_0", xform)} NewtonManager._cl_inject_sites_fallback() entry = NewtonManager._cl_site_index_map["ft_0"] @@ -96,7 +96,7 @@ def test_global_site_entry_is_int_none_tuple(self): def test_global_site_pending_cleared(self): xform = wp.transform() - NewtonManager._cl_pending_sites = {(None, tuple(xform)): ("ft_0", xform)} + NewtonManager._cl_pending_sites = {(None, False, tuple(xform)): ("ft_0", xform)} NewtonManager._cl_inject_sites_fallback() assert len(NewtonManager._cl_pending_sites) == 0 @@ -111,7 +111,7 @@ def setup_method(self): def test_single_body_entry_shape(self): xform = wp.transform() - NewtonManager._cl_pending_sites = {("Robot/base", tuple(xform)): ("ft_0", xform)} + NewtonManager._cl_pending_sites = {("Robot/base", False, tuple(xform)): ("ft_0", xform)} NewtonManager._cl_inject_sites_fallback() entry = NewtonManager._cl_site_index_map["ft_0"] @@ -132,7 +132,7 @@ def setup_method(self): def test_wildcard_entry_shape(self): xform = wp.transform() - NewtonManager._cl_pending_sites = {("Robot/.*_foot", tuple(xform)): ("ft_0", xform)} + NewtonManager._cl_pending_sites = {("Robot/.*_foot", False, tuple(xform)): ("ft_0", xform)} NewtonManager._cl_inject_sites_fallback() entry = NewtonManager._cl_site_index_map["ft_0"] @@ -143,11 +143,47 @@ def test_wildcard_entry_shape(self): def test_no_match_raises(self): xform = wp.transform() - NewtonManager._cl_pending_sites = {("Robot/nonexistent", tuple(xform)): ("ft_0", xform)} + NewtonManager._cl_pending_sites = {("Robot/nonexistent", False, tuple(xform)): ("ft_0", xform)} with pytest.raises(ValueError): NewtonManager._cl_inject_sites_fallback() +class TestWorldSite: + """World-local sites are per-world, not global.""" + + def setup_method(self): + NewtonManager.clear() + NewtonManager._builder = MockBuilder([]) + + def test_world_site_reuses_label(self): + xform = wp.transform((1.0, 2.0, 3.0), wp.quat_identity()) + label_0 = NewtonManager.cl_register_site(None, xform, per_world=True) + label_1 = NewtonManager.cl_register_site(None, xform, per_world=True) + + assert label_0 == label_1 + + def test_world_site_fallback_entry_is_local(self): + xform = wp.transform((1.0, 2.0, 3.0), wp.quat_identity()) + label = NewtonManager.cl_register_site(None, xform, per_world=True) + NewtonManager._cl_inject_sites_fallback() + + global_idx, per_world = NewtonManager._cl_site_index_map[label] + assert global_idx is None + assert isinstance(per_world, list) + assert len(per_world) == 1 + assert len(per_world[0]) == 1 + + def test_inject_sites_returns_world_sites(self): + xform = wp.transform((1.0, 2.0, 3.0), wp.quat_identity()) + label = NewtonManager.cl_register_site(None, xform, per_world=True) + global_sites, proto_sites, world_sites = NewtonManager._cl_inject_sites(MockBuilder([]), {}) + + assert global_sites == {} + assert proto_sites == {} + assert world_sites[label] == xform + assert NewtonManager._cl_pending_sites == {} + + # --------------------------------------------------------------------------- # FrameTransformer._validate_site_map # --------------------------------------------------------------------------- diff --git a/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py b/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py index 594a55e13584..d114a1da2a80 100644 --- a/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py +++ b/source/isaaclab_newton/test/sim/test_views_xform_prim_newton.py @@ -25,8 +25,6 @@ from isaaclab_newton.physics.newton_manager import NewtonManager from isaaclab_newton.sim.views import NewtonSiteFrameView as FrameView -from pxr import Gf - import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg from isaaclab.scene import InteractiveScene, InteractiveSceneCfg @@ -85,16 +83,9 @@ def factory(num_envs: int, device: str) -> ViewBundle: sim = ctx.__enter__() sim._app_control_on_stop_handle = None InteractiveScene(_SceneCfg(num_envs=num_envs, env_spacing=2.0)) - - stage = sim_utils.get_current_stage() - for i in range(num_envs): - prim = stage.DefinePrim(f"/World/envs/env_{i}/Cube/CameraMount", "Xform") - sim_utils.standardize_xform_ops(prim) - prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(*CHILD_OFFSET)) - prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) - - sim.reset() + sim_utils.create_prim("/World/envs/env_0/Cube/CameraMount", translation=CHILD_OFFSET) view = FrameView("/World/envs/env_.*/Cube/CameraMount", device=device) + sim.reset() return ViewBundle( view=view, @@ -143,6 +134,50 @@ def test_reject_shape_path(device): ctx.__exit__(None, None, None) +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_clone_plan_view_uses_source_child_without_destination_usd(device): + """FrameView expands a registered body-local site through the ClonePlan.""" + num_envs = 3 + ctx = _sim_context(device, num_envs=num_envs) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_SceneCfg(num_envs=num_envs, env_spacing=2.0)) + + stage = sim_utils.get_current_stage() + assert stage.GetPrimAtPath("/World/envs/env_0/Cube").IsValid() + assert not stage.GetPrimAtPath("/World/envs/env_1/Cube").IsValid() + sim_utils.create_prim("/World/envs/env_0/Cube/CameraMount", translation=CHILD_OFFSET) + + view = FrameView("/World/envs/env_.*/Cube/CameraMount", device=device) + sim.reset() + + assert view.count == num_envs + assert not stage.GetPrimAtPath("/World/envs/env_1/Cube/CameraMount").IsValid() + pos = view.get_world_poses()[0].torch + expected = _get_body_positions(num_envs, device) + torch.tensor(CHILD_OFFSET, device=device) + torch.testing.assert_close(pos, expected, atol=1e-5, rtol=0) + ctx.__exit__(None, None, None) + + +@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) +def test_view_can_resolve_from_body_labels_after_reset(device): + """FrameView can resolve a body-local frame directly from Newton body labels.""" + num_envs = 3 + ctx = _sim_context(device, num_envs=num_envs) + sim = ctx.__enter__() + sim._app_control_on_stop_handle = None + InteractiveScene(_SceneCfg(num_envs=num_envs, env_spacing=2.0)) + sim_utils.create_prim("/World/envs/env_0/Cube/CameraMount", translation=CHILD_OFFSET) + + sim.reset() + view = FrameView("/World/envs/env_.*/Cube/CameraMount", device=device) + + pos = view.get_world_poses()[0].torch + expected = _get_body_positions(num_envs, device) + torch.tensor(CHILD_OFFSET, device=device) + torch.testing.assert_close(pos, expected, atol=1e-5, rtol=0) + ctx.__exit__(None, None, None) + + # ================================================================== # Newton edge case: world-attached prim (body=-1) # ================================================================== @@ -150,19 +185,14 @@ def test_reject_shape_path(device): @pytest.mark.parametrize("device", ["cpu", "cuda:0"]) def test_world_attached_returns_initial_pose(device): - """A world-rooted Xform returns its USD-authored position.""" + """A world-rooted frame returns its configured position.""" ctx = _sim_context(device, num_envs=2) sim = ctx.__enter__() sim._app_control_on_stop_handle = None InteractiveScene(_SceneCfg(num_envs=2, env_spacing=2.0)) - stage = sim_utils.get_current_stage() - prim = stage.DefinePrim("/World/StaticMarker", "Xform") - sim_utils.standardize_xform_ops(prim) - prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(*WORLD_MARKER_POS)) - prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) - sim.reset() + sim_utils.create_prim("/World/StaticMarker", translation=WORLD_MARKER_POS) view = FrameView("/World/StaticMarker", device=device) pos = view.get_world_poses()[0].torch @@ -179,13 +209,8 @@ def test_world_attached_set_world_roundtrip(device): sim._app_control_on_stop_handle = None InteractiveScene(_SceneCfg(num_envs=2, env_spacing=2.0)) - stage = sim_utils.get_current_stage() - prim = stage.DefinePrim("/World/StaticMarker", "Xform") - sim_utils.standardize_xform_ops(prim) - prim.GetAttribute("xformOp:translate").Set(Gf.Vec3d(*WORLD_MARKER_POS)) - prim.GetAttribute("xformOp:orient").Set(Gf.Quatd(1.0, 0.0, 0.0, 0.0)) - sim.reset() + sim_utils.create_prim("/World/StaticMarker", translation=WORLD_MARKER_POS) view = FrameView("/World/StaticMarker", device=device) new_pos = _wp_vec3f([[10.0, 20.0, 30.0]], device=device) diff --git a/source/isaaclab_ovphysx/changelog.d/newton-clone-plan.skip b/source/isaaclab_ovphysx/changelog.d/newton-clone-plan.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py index fdb836fc08b9..a59a49e65189 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/articulation/articulation.py @@ -3474,61 +3474,43 @@ def _initialize_impl(self) -> None: self._ovphysx = physx_instance self._device = OvPhysxManager.get_device() - # IsaacLab uses two conventions for env-glob prim paths: - # /World/envs/env_.*/Robot -- regex dot-star for "any env index" - # /World/envs/{ENV_REGEX_NS}/... -- explicit placeholder - # ovphysx ``create_tensor_binding`` expects fnmatch-style globs, so both map to '*'. - prim_path = self.cfg.prim_path - pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", prim_path) - pattern = re.sub(r"\.\*", "*", pattern) - - # ``PhysicsArticulationRootAPI`` may live on a CHILD prim rather than on - # the cfg prim itself. ``create_tensor_binding`` only matches prims that - # have the API applied, so the pattern must be extended to the actual - # articulation root. - stage = PhysicsManager._sim.stage + # Resolve the articulation root expression. if self.cfg.articulation_root_prim_path is not None: - # explicit subpath: skip auto-discovery but validate the prim exists - root_relative = self.cfg.articulation_root_prim_path - self._articulation_root_path = prim_path + root_relative - if sim_utils.find_first_matching_prim(self._articulation_root_path, stage=stage) is None: - raise RuntimeError( - f"Failed to find articulation root prim at '{self._articulation_root_path}'." - " Check that ``cfg.articulation_root_prim_path`` points at a prim that exists" - " in the USD stage." - ) - pattern = pattern + root_relative - logger.info("OvPhysxManager: explicit articulation root '%s' (pattern '%s')", root_relative, pattern) + root_prim_path_expr = self.cfg.prim_path + self.cfg.articulation_root_prim_path else: - first_prim = sim_utils.find_first_matching_prim(prim_path, stage=stage) - if first_prim is None: - raise RuntimeError(f"Failed to find prim for expression: '{prim_path}'.") - first_prim_path = first_prim.GetPath().pathString + def has_articulation_root_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.ArticulationRootAPI)) + + matches = sim_utils.resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString root_prims = sim_utils.get_all_matching_child_prims( - first_prim_path, - predicate=lambda p: p.HasAPI(UsdPhysics.ArticulationRootAPI), - traverse_instance_prims=False, + walk_root, predicate=has_articulation_root_api, traverse_instance_prims=False ) - if len(root_prims) == 0: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] raise RuntimeError( - f"Failed to find an articulation root when resolving '{prim_path}'." - " Ensure the prim has 'USD ArticulationRootAPI' applied." + f"Expected exactly one ArticulationRootAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - if len(root_prims) > 1: - raise RuntimeError( - f"Failed to find a single articulation root when resolving '{prim_path}'." - f" Found multiple under '{first_prim_path}'." - ) - - self._articulation_root_path = root_prims[0].GetPath().pathString - root_relative = self._articulation_root_path[len(first_prim_path) :] - if root_relative: - pattern = pattern + root_relative - logger.info( - "OvPhysxManager: articulation root at '%s' (pattern extended to '%s')", root_relative, pattern - ) - + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] + # Validate the prim exists on the live stage -- ``create_tensor_binding`` silently + # returns a 0-count binding when the pattern matches nothing, surfacing as obscure + # AttributeErrors deep in property accessors. Also stash the concrete source-side + # root path for tendon discovery downstream. + stage = PhysicsManager._sim.stage + first_match = sim_utils.find_first_matching_prim(root_prim_path_expr, stage=stage) + if first_match is None: + raise RuntimeError(f"Failed to find articulation root prim at '{root_prim_path_expr}'.") + self._articulation_root_path = first_match.GetPath().pathString + + # IsaacLab paths may use ``.*`` regex or ``{ENV_REGEX_NS}`` placeholder; ovphysx + # ``create_tensor_binding`` expects fnmatch globs. + pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", root_prim_path_expr) + pattern = re.sub(r"\.\*", "*", pattern) self._binding_pattern = pattern # eagerly create every binding the data container reads at init, so @@ -3562,7 +3544,7 @@ def _initialize_impl(self) -> None: if not self._bindings: raise RuntimeError( f"OVPhysX could not create any articulation bindings for pattern {pattern!r}. " - f"Check that prim_path={prim_path!r} matches at least one " + f"Check that prim_path={self.cfg.prim_path!r} matches at least one " "UsdPhysics.ArticulationRootAPI prim." ) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py index 015c8f102f44..9ee266506480 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py @@ -18,9 +18,9 @@ from pxr import UsdPhysics -import isaaclab.sim as sim_utils from isaaclab.assets.rigid_object.base_rigid_object import BaseRigidObject from isaaclab.assets.rigid_object.rigid_object_cfg import RigidObjectCfg +from isaaclab.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source from isaaclab.utils.string import resolve_matching_names from isaaclab.utils.wrench_composer import WrenchComposer @@ -849,56 +849,32 @@ def _initialize_impl(self) -> None: # raise AttributeError (masked by hasattr) and fall back to "cuda:0" even when the # simulation is running on CPU, causing a device mismatch in binding.read(). self._device = OvPhysxManager.get_device() - # Convert IsaacLab prim-path notation to the glob patterns ovphysx expects. - # IsaacLab uses two conventions: - # /World/envs/env_.*/object -- regex dot-star for "any env index" - # /World/envs/{ENV_REGEX_NS}/object -- explicit placeholder - # ovphysx ``create_tensor_binding`` uses fnmatch-style globs, so both map to ``*``. - pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", self.cfg.prim_path) - pattern = re.sub(r"\.\*", "*", pattern) - self._binding_pattern = pattern - # Validate the prim tree before creating tensor bindings -- the wheel silently - # produces a 0-prim binding when the pattern matches nothing, which surfaces as an - # obscure ``TypeError`` deep in property accessors. - # 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, + # Resolve the rigid body root expression. + def has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) + + matches = resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString + root_prims = get_all_matching_child_prims( + walk_root, predicate=has_rigid_body_api, traverse_instance_prims=False ) - if len(root_prims) == 0: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] raise RuntimeError( - f"Failed to find a rigid body when resolving '{self.cfg.prim_path}'." - " Please ensure that the prim has 'USD RigidBodyAPI' applied." + f"Expected exactly one RigidBodyAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - 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" - f" objects. These are located at: '{articulation_prims}' under" - f" '{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." - ) + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] + + # IsaacLab paths may use ``.*`` regex or ``{ENV_REGEX_NS}`` placeholder; ovphysx + # ``create_tensor_binding`` expects fnmatch globs. + pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", root_prim_path_expr) + pattern = re.sub(r"\.\*", "*", pattern) + self._binding_pattern = pattern # Eagerly create every binding the data container reads at init, so failures # surface here with a helpful message rather than as a raw wheel exception diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py index 7a1d4e21e535..d712861ec34c 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py @@ -1040,60 +1040,30 @@ def _initialize_impl(self) -> None: self._prim_paths: list[str] = [] self._body_names_list: list[str] = [] - for name, obj_cfg in self.cfg.rigid_objects.items(): - # Convert IsaacLab prim-path notation to the fnmatch-style glob that - # OVPhysX create_tensor_binding expects. Two conventions are in use: - # /World/envs/env_.*/object -- regex dot-star for any env index - # /World/envs/{ENV_REGEX_NS}/object -- explicit placeholder - pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", obj_cfg.prim_path) - pattern = re.sub(r"\.\*", "*", pattern) - - # Validate the prim tree before creating tensor bindings. - # OVPhysX silently returns a zero-count binding when the pattern - # matches nothing; fail fast here with a clear message instead. - template_prim = sim_utils.find_first_matching_prim(obj_cfg.prim_path) - if template_prim is None: - raise RuntimeError(f"Failed to find prim for expression: '{obj_cfg.prim_path}' (body '{name}').") - template_prim_path = template_prim.GetPath().pathString + def has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) + for name, obj_cfg in self.cfg.rigid_objects.items(): + # Resolve the rigid body root expression. + matches = sim_utils.resolve_matching_prims_from_source(obj_cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{obj_cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString root_prims = sim_utils.get_all_matching_child_prims( - template_prim_path, - predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI), - traverse_instance_prims=False, + walk_root, predicate=has_rigid_body_api, traverse_instance_prims=False ) - if len(root_prims) == 0: - raise RuntimeError( - f"Failed to find a rigid body when resolving '{obj_cfg.prim_path}' (body '{name}')." - " Please ensure that the prim has 'USD RigidBodyAPI' applied." - ) - if len(root_prims) > 1: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] raise RuntimeError( - f"Failed to find a single rigid body when resolving '{obj_cfg.prim_path}' (body '{name}')." - f" Found multiple '{root_prims}' under '{template_prim_path}'." - " Please ensure that there is only one rigid body in the prim path tree." + f"Expected exactly one RigidBodyAPI prim under '{walk_root}'" + f" (resolved from '{obj_cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - - 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 '{obj_cfg.prim_path}' (body '{name}') in the" - f" rigid object collection. These are located at: '{articulation_prims}' under" - f" '{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 the regex expression - root_prim_path = root_prims[0].GetPath().pathString - suffix = root_prim_path[len(template_prim_path) :] - if suffix: - pattern = pattern + suffix - + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] + # IsaacLab paths may use ``.*`` regex or ``{ENV_REGEX_NS}`` placeholder; ovphysx + # ``create_tensor_binding`` expects fnmatch globs. + pattern = re.sub(r"\{ENV_REGEX_NS\}", "*", root_prim_path_expr) + pattern = re.sub(r"\.\*", "*", pattern) self._prim_paths.append(pattern) self._body_names_list.append(name) diff --git a/source/isaaclab_ovphysx/isaaclab_ovphysx/sensors/contact_sensor/contact_sensor.py b/source/isaaclab_ovphysx/isaaclab_ovphysx/sensors/contact_sensor/contact_sensor.py index 898fbbeb4de6..fb33c184472c 100644 --- a/source/isaaclab_ovphysx/isaaclab_ovphysx/sensors/contact_sensor/contact_sensor.py +++ b/source/isaaclab_ovphysx/isaaclab_ovphysx/sensors/contact_sensor/contact_sensor.py @@ -16,8 +16,8 @@ import warp as wp -import isaaclab.sim as sim_utils from isaaclab.sensors.contact_sensor import BaseContactSensor +from isaaclab.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source from isaaclab.utils.warp import ProxyArray import isaaclab_ovphysx.tensor_types as TT @@ -179,23 +179,32 @@ def _initialize_impl(self) -> None: raise RuntimeError("OvPhysxManager has not been initialized yet.") self._physx_instance = physx_instance - # Discover sensor bodies. Mirror the PhysX discovery path but use - # ``GetPrimTypeInfo().GetAppliedAPISchemas()`` (raw apiSchemas listOp) - # rather than ``GetAppliedSchemas()`` (filtered by USD's plugin - # registry). Under the kitless ovphysx flow the ``PhysxSchema`` USD - # plugin is registered by :meth:`OvPhysxManager.initialize` so the - # wheel-side schema check passes, but the Python-side filtered API - # still hides ``PhysxContactReportAPI`` because the schema TYPE - # registration only happens when the C++ plugin library is loaded by - # ``omni.physx``. The unfiltered API matches what the underlying - # USD apiSchemas listOp actually carries (verified against + # Discover sensor bodies. We use ``GetPrimTypeInfo().GetAppliedAPISchemas()`` + # (raw apiSchemas listOp) instead of ``GetAppliedSchemas()`` so that codeless + # USDs without ``omni.physx``'s plugin loaded still report + # ``PhysxContactReportAPI``. Under the kitless ovphysx flow the + # ``PhysxSchema`` USD plugin is registered by + # :meth:`OvPhysxManager.initialize` so the wheel-side schema check passes, + # but the Python-side filtered API still hides ``PhysxContactReportAPI`` + # because the schema TYPE registration only happens when the C++ plugin + # library is loaded by ``omni.physx``. The unfiltered API matches what + # the underlying USD apiSchemas listOp actually carries (verified against # :class:`pxr.Sdf.PrimSpec.GetInfo("apiSchemas")`). - leaf_pattern = self.cfg.prim_path.rsplit("/", 1)[-1] - template_prim_path = self._parent_prims[0].GetPath().pathString - body_names: list[str] = [] - for prim in sim_utils.find_matching_prims(template_prim_path + "/" + leaf_pattern): - if "PhysxContactReportAPI" in prim.GetPrimTypeInfo().GetAppliedAPISchemas(): - body_names.append(prim.GetPath().pathString.rsplit("/", 1)[-1]) + parent_expr, leaf_pattern = self.cfg.prim_path.rsplit("/", 1) + name_pattern = re.compile(leaf_pattern) + + def has_contact_report(prim) -> bool: + return bool(name_pattern.fullmatch(prim.GetName())) and ( + "PhysxContactReportAPI" in prim.GetPrimTypeInfo().GetAppliedAPISchemas() + ) + + matches = resolve_matching_prims_from_source(parent_expr) + if not matches: + raise RuntimeError(f"No prim found at '{parent_expr}'.") + asset_prim, body_parent = matches[0] + walk_root = asset_prim.GetPath().pathString + prims = get_all_matching_child_prims(walk_root, predicate=has_contact_report, traverse_instance_prims=False) + body_names = [prim.GetPath().pathString.rsplit("/", 1)[-1] for prim in prims] if not body_names: raise RuntimeError( f"Sensor at path '{self.cfg.prim_path}' could not find any bodies with contact reporter API." @@ -206,8 +215,7 @@ def _initialize_impl(self) -> None: # Build glob patterns: one per (env, sensor body). # IsaacLab path forms map to ovphysx fnmatch globs the same way Articulation does. - base_glob = self.cfg.prim_path.rsplit("/", 1)[0] - base_glob = re.sub(r"\{ENV_REGEX_NS\}", "*", base_glob) + base_glob = re.sub(r"\{ENV_REGEX_NS\}", "*", body_parent) base_glob = re.sub(r"\.\*", "*", base_glob) sensor_patterns = [f"{base_glob}/{name}" for name in body_names] diff --git a/source/isaaclab_physx/changelog.d/newton-clone-plan.skip b/source/isaaclab_physx/changelog.d/newton-clone-plan.skip new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 78fdb387aa04..eac74b91dd2c 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py @@ -23,16 +23,19 @@ from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator from isaaclab.assets.articulation.base_articulation import BaseArticulation - -_HAS_NEWTON_ACTUATORS = importlib.util.find_spec("isaaclab_newton.actuators") is not None - - -from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims +from isaaclab.sim.utils.queries import ( + find_first_matching_prim, + get_all_matching_child_prims, + resolve_matching_prims_from_source, +) 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 +_HAS_NEWTON_ACTUATORS = importlib.util.find_spec("isaaclab_newton.actuators") is not None + + from isaaclab_physx.assets import kernels as shared_kernels from isaaclab_physx.assets.articulation import kernels as articulation_kernels from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -3770,42 +3773,27 @@ def _initialize_impl(self): 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) - and prim.GetAttribute("physxArticulation:articulationEnabled").Get() is not False, - traverse_instance_prims=False, + + def has_articulation_root_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.ArticulationRootAPI)) + + matches = resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString + root_prims = get_all_matching_child_prims( + walk_root, predicate=has_articulation_root_api, 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: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] 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." + f"Expected exactly one ArticulationRootAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - - # 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 - + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] # -- articulation self._root_view = self._physics_sim_view.create_articulation_view(root_prim_path_expr.replace(".*", "*")) 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 759e31d9fb87..87706377173a 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 @@ -579,30 +579,24 @@ def _resolve_env_ids(self, env_ids): 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 deformable root prims + def has_deformable_body_api(prim) -> bool: + return "OmniPhysicsDeformableBodyAPI" in prim.GetAppliedSchemas() + + matches = sim_utils.resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString root_prims = sim_utils.get_all_matching_child_prims( - template_prim_path, - predicate=lambda prim: "OmniPhysicsDeformableBodyAPI" in prim.GetAppliedSchemas(), - traverse_instance_prims=False, + walk_root, predicate=has_deformable_body_api, traverse_instance_prims=False ) - if len(root_prims) == 0: - raise RuntimeError( - f"Failed to find a deformable body when resolving '{self.cfg.prim_path}'." - " Please ensure that the prim has 'OmniPhysicsDeformableBodyAPI' applied." - ) - if len(root_prims) > 1: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] raise RuntimeError( - f"Failed to find a single deformable body when resolving '{self.cfg.prim_path}'." - f" Found multiple '{root_prims}' under '{template_prim_path}'." - " Please ensure that there is only one deformable body in the prim path tree." + f"Expected exactly one OmniPhysicsDeformableBodyAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - # we only need the first one from the list root_prim = root_prims[0] # find deformable material prims @@ -655,10 +649,8 @@ def _initialize_impl(self): if has_mesh: self._deformable_type = "surface" - # resolve root path back into regex expression - # -- root prim expression - root_prim_path = root_prim.GetPath().pathString - root_prim_path_expr = self.cfg.prim_path + root_prim_path[len(template_prim_path) :] + # resolve root path back into the destination glob expression + root_prim_path_expr = root_expr + root_prim.GetPath().pathString[len(walk_root) :] # -- object view if self._deformable_type == "surface": # surface deformable @@ -690,8 +682,8 @@ def _initialize_impl(self): material_prim_path = material_prim.GetPath().pathString # check if the material prim is under the template prim # if not then we are assuming that the single material prim is used for all the deformable bodies - if template_prim_path in material_prim_path: - material_prim_path_expr = self.cfg.prim_path + material_prim_path[len(template_prim_path) :] + if walk_root in material_prim_path: + material_prim_path_expr = root_expr + material_prim_path[len(walk_root) :] else: material_prim_path_expr = material_prim_path # -- material view 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 930b8836859a..c2e0955c438a 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 @@ -15,9 +15,9 @@ 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.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source from isaaclab.utils.wrench_composer import WrenchComposer from isaaclab_physx.assets import kernels as shared_kernels @@ -898,47 +898,25 @@ def set_inertias_mask( 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, + + def has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) + + matches = resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString + root_prims = get_all_matching_child_prims( + walk_root, predicate=has_rigid_body_api, 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: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] 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." + f"Expected exactly one RigidBodyAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - - 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) :] + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] # -- object view self._root_view = self._physics_sim_view.create_rigid_body_view(root_prim_path_expr.replace(".*", "*")) 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 2031ded53b2f..b20bbe0d5f5f 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 @@ -1216,52 +1216,28 @@ def _initialize_impl(self): self._body_names_list.clear() # obtain global simulation view self._physics_sim_view = SimulationManager.get_physics_sim_view() + + def has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) + root_prim_path_exprs = [] for name, rigid_body_cfg in self.cfg.rigid_objects.items(): - # 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(rigid_body_cfg.prim_path) - if template_prim is None: - raise RuntimeError(f"Failed to find prim for expression: '{rigid_body_cfg.prim_path}'.") - template_prim_path = template_prim.GetPath().pathString - - # find rigid root prims + matches = sim_utils.resolve_matching_prims_from_source(rigid_body_cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{rigid_body_cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString root_prims = sim_utils.get_all_matching_child_prims( - template_prim_path, - predicate=lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI), - traverse_instance_prims=False, + walk_root, predicate=has_rigid_body_api, traverse_instance_prims=False ) - if len(root_prims) == 0: - raise RuntimeError( - f"Failed to find a rigid body when resolving '{rigid_body_cfg.prim_path}'." - " Please ensure that the prim has 'USD RigidBodyAPI' applied." - ) - if len(root_prims) > 1: + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] raise RuntimeError( - f"Failed to find a single rigid body when resolving '{rigid_body_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." + f"Expected exactly one RigidBodyAPI prim under '{walk_root}'" + f" (resolved from '{rigid_body_cfg.prim_path}'), found {len(root_prims)}: {matched}." ) - - # check that no rigid object has an articulation root API, which decreases simulation performance - 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 '{rigid_body_cfg.prim_path}' in the rigid object" - f" collection. These are 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 = rigid_body_cfg.prim_path + root_prim_path[len(template_prim_path) :] + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] root_prim_path_exprs.append(root_prim_path_expr.replace(".*", "*")) - self._body_names_list.append(name) # -- object view 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 6662582dac7c..df38e0e6d40b 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 @@ -453,36 +453,25 @@ def _initialize_impl(self) -> None: enable_extension("isaacsim.robot.surface_gripper") from isaacsim.robot.surface_gripper import GripperView - # 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 surface gripper prims + def is_surface_gripper(prim) -> bool: + return prim.GetTypeName() == "IsaacSurfaceGripper" + + matches = sim_utils.resolve_matching_prims_from_source(self._cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self._cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString gripper_prims = sim_utils.get_all_matching_child_prims( - template_prim_path, - predicate=lambda prim: prim.GetTypeName() == "IsaacSurfaceGripper", - traverse_instance_prims=False, + walk_root, predicate=is_surface_gripper, traverse_instance_prims=False ) - if len(gripper_prims) == 0: - raise RuntimeError( - f"Failed to find a surface gripper when resolving '{self._cfg.prim_path}'." - " Please ensure that the prim has type 'IsaacSurfaceGripper'." - ) - if len(gripper_prims) > 1: + if len(gripper_prims) != 1: + matched = [p.GetPath().pathString for p in gripper_prims] raise RuntimeError( - f"Failed to find a single surface gripper when resolving '{self._cfg.prim_path}'." - f" Found multiple '{gripper_prims}' under '{template_prim_path}'." - " Please ensure that there is only one surface gripper in the prim path tree." + f"Expected exactly one IsaacSurfaceGripper prim under '{walk_root}'" + f" (resolved from '{self._cfg.prim_path}'), found {len(gripper_prims)}: {matched}." ) - - # resolve gripper prim back into regex expression - gripper_prim_path = gripper_prims[0].GetPath().pathString - gripper_prim_path_expr = self._cfg.prim_path + gripper_prim_path[len(template_prim_path) :] - - # Count number of environments - self._prim_expr = gripper_prim_path_expr + gripper_prim = gripper_prims[0] + self._prim_expr = root_expr + gripper_prim.GetPath().pathString[len(walk_root) :] env_prim_path_expr = self._prim_expr.rsplit("/", 1)[0] self._parent_prims = sim_utils.find_matching_prims(env_prim_path_expr) self._num_envs = len(self._parent_prims) @@ -505,7 +494,7 @@ def _initialize_impl(self) -> None: ) # log information about the surface gripper - logger.info(f"Surface gripper initialized at: {self._cfg.prim_path} with root '{gripper_prim_path_expr}'.") + logger.info(f"Surface gripper initialized at: {self._cfg.prim_path} with root '{self._prim_expr}'.") logger.info(f"Number of instances: {self._num_envs}") # Reset grippers 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 c38a0fffe281..5d1fa0c7b65e 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 @@ -8,6 +8,7 @@ from __future__ import annotations +import re from collections.abc import Sequence from typing import TYPE_CHECKING @@ -16,10 +17,10 @@ import omni.physics.tensors.api as physx -import isaaclab.sim as sim_utils from isaaclab.app.settings_manager import get_settings_manager from isaaclab.markers import VisualizationMarkers from isaaclab.sensors.contact_sensor import BaseContactSensor +from isaaclab.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source from isaaclab.utils.warp import ProxyArray from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -288,26 +289,33 @@ def _initialize_impl(self): super()._initialize_impl() # obtain global simulation view self._physics_sim_view = SimulationManager.get_physics_sim_view() - # check that only rigid bodies are selected - leaf_pattern = self.cfg.prim_path.rsplit("/", 1)[-1] - template_prim_path = self._parent_prims[0].GetPath().pathString - body_names = list() - for prim in sim_utils.find_matching_prims(template_prim_path + "/" + leaf_pattern): - # check if prim has contact reporter API - if "PhysxContactReportAPI" in prim.GetAppliedSchemas(): - prim_path = prim.GetPath().pathString - body_names.append(prim_path.rsplit("/", 1)[-1]) - # check that there is at least one body with contact reporter API + + # Split the configured prim path into a parent expression and a leaf-name regex. + parent_expr, leaf_pattern = self.cfg.prim_path.rsplit("/", 1) + name_pattern = re.compile(leaf_pattern) + + def has_contact_report(prim) -> bool: + return bool(name_pattern.fullmatch(prim.GetName())) and ( + "PhysxContactReportAPI" in prim.GetAppliedSchemas() + ) + + # Resolve the asset subtree (clone-plan aware) and collect contact-reporting descendants. + matches = resolve_matching_prims_from_source(parent_expr) + if not matches: + raise RuntimeError(f"No prim found at '{parent_expr}'.") + asset_prim, body_parent = matches[0] + walk_root = asset_prim.GetPath().pathString + prims = get_all_matching_child_prims(walk_root, predicate=has_contact_report, traverse_instance_prims=False) + body_names = [prim.GetPath().pathString.rsplit("/", 1)[-1] for prim in prims] if not body_names: raise RuntimeError( f"Sensor at path '{self.cfg.prim_path}' could not find any bodies with contact reporter API." "\nHINT: Make sure to enable 'activate_contact_sensors' in the corresponding asset spawn configuration." ) - # construct regex expression for the body names + # construct regex expression for the body names and convert to PhysX glob form body_names_regex = r"(" + "|".join(body_names) + r")" - body_names_regex = f"{self.cfg.prim_path.rsplit('/', 1)[0]}/{body_names_regex}" - # convert regex expressions to glob expressions for PhysX + body_names_regex = f"{body_parent}/{body_names_regex}" body_names_glob = body_names_regex.replace(".*", "*") filter_prim_paths_glob = [expr.replace(".*", "*") for expr in self.cfg.filter_prim_paths_expr] @@ -325,7 +333,7 @@ def _initialize_impl(self): raise RuntimeError( "Failed to initialize contact reporter for specified bodies." f"\n\tInput prim path : {self.cfg.prim_path}" - f"\n\tResolved prim paths: {body_names_regex}" + f"\n\tResolved prim paths: {body_names_glob}" ) # check if filter paths are valid diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer.py b/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer.py index b7215062ff3f..52ad5d4d459d 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/frame_transformer/frame_transformer.py @@ -16,9 +16,9 @@ from pxr import UsdPhysics -import isaaclab.sim as sim_utils from isaaclab.markers import VisualizationMarkers from isaaclab.sensors.frame_transformer import BaseFrameTransformer +from isaaclab.sim.utils.queries import resolve_matching_prims_from_source from isaaclab.utils.math import is_identity_pose, normalize, quat_from_angle_axis from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -182,23 +182,19 @@ def _initialize_impl(self): frame_offsets = [None] + [target_frame.offset for target_frame in self.cfg.target_frames] frame_types = ["source"] + ["target"] * len(self.cfg.target_frames) for frame, prim_path, offset, frame_type in zip(frames, frame_prim_paths, frame_offsets, frame_types): - # Find correct prim - matching_prims = sim_utils.find_matching_prims(prim_path) - if len(matching_prims) == 0: + # Resolve the source-side env prims (filtered to rigid bodies) and their destination + # expressions. Plan-aware: with an active ``ClonePlan``, only env-0 representatives + # are walked and dest expressions are rebuilt against the plan's destination glob. + def has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) + + matches = resolve_matching_prims_from_source(prim_path, predicate=has_rigid_body_api) + if not matches: raise ValueError( f"Failed to create frame transformer for frame '{frame}' with path '{prim_path}'." - " No matching prims were found." + " No matching rigid-body prims were found." ) - for prim in matching_prims: - # Get the prim path of the matching prim - matching_prim_path = prim.GetPath().pathString - # Check if it is a rigid prim - if not prim.HasAPI(UsdPhysics.RigidBodyAPI): - raise ValueError( - f"While resolving expression '{prim_path}' found a prim '{matching_prim_path}' which is not a" - " rigid body. The class only supports transformations between rigid bodies." - ) - + for prim, matching_prim_path in matches: # Get the name of the body: use relative prim path for unique identification body_name = self._get_relative_body_path(matching_prim_path) # Use leaf name of prim path if frame name isn't specified by user @@ -244,7 +240,12 @@ def _initialize_impl(self): tracked_prim_paths = [body_names_to_frames[body_name]["prim_path"] for body_name in body_names_to_frames.keys()] tracked_body_names = [body_name for body_name in body_names_to_frames.keys()] - body_names_regex = [tracked_prim_path.replace("env_0", "env_*") for tracked_prim_path in tracked_prim_paths] + # Convert each tracked prim path to PhysX glob form for ``create_rigid_body_view``. + # Plan-mode dest expressions use ``env_.*`` (regex), legacy mode produces concrete + # ``env_0`` paths; chain both substitutions so each mode normalises to ``env_*``. + body_names_regex = [ + tracked_prim_path.replace(".*", "*").replace("env_0", "env_*") for tracked_prim_path in tracked_prim_paths + ] # obtain global simulation view self._physics_sim_view = SimulationManager.get_physics_sim_view() diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/imu/imu.py b/source/isaaclab_physx/isaaclab_physx/sensors/imu/imu.py index 75e8f7409d7b..21820ecb5f0b 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/imu/imu.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/imu/imu.py @@ -11,9 +11,6 @@ import torch import warp as wp -from pxr import UsdPhysics - -import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils from isaaclab.sensors.imu import BaseImu @@ -125,24 +122,8 @@ def _initialize_impl(self): """ super()._initialize_impl() self._physics_sim_view = SimulationManager.get_physics_sim_view() - prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) - if prim is None: - raise RuntimeError(f"Failed to find a prim at path expression: {self.cfg.prim_path}") - - ancestor_prim = sim_utils.get_first_matching_ancestor_prim( - prim.GetPath(), predicate=lambda _prim: _prim.HasAPI(UsdPhysics.RigidBodyAPI) - ) - if ancestor_prim is None: - raise RuntimeError(f"Failed to find a rigid body ancestor prim at path expression: {self.cfg.prim_path}") - - if ancestor_prim == prim: - self._rigid_parent_expr = self.cfg.prim_path - fixed_pos_b, fixed_quat_b = None, None - else: - relative_path = prim.GetPath().MakeRelativePath(ancestor_prim.GetPath()).pathString - self._rigid_parent_expr = self.cfg.prim_path.replace("/" + relative_path, "") - fixed_pos_b, fixed_quat_b = sim_utils.resolve_prim_pose(prim, ancestor_prim) + self._rigid_parent_expr, fixed_pos_b, fixed_quat_b = self._resolve_rigid_body_ancestor_expr() self._view = self._physics_sim_view.create_rigid_body_view(self._rigid_parent_expr.replace(".*", "*")) # Query world gravity and compute accelerometer bias (real IMUs always measure gravity) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/joint_wrench/joint_wrench_sensor.py b/source/isaaclab_physx/isaaclab_physx/sensors/joint_wrench/joint_wrench_sensor.py index 955a448f0b79..14386cb245a6 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/joint_wrench/joint_wrench_sensor.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/joint_wrench/joint_wrench_sensor.py @@ -17,7 +17,7 @@ from pxr import Usd, UsdPhysics from isaaclab.sensors.joint_wrench import BaseJointWrenchSensor -from isaaclab.sim.utils.queries import find_first_matching_prim, get_all_matching_child_prims +from isaaclab.sim.utils.queries import get_all_matching_child_prims, resolve_matching_prims_from_source from isaaclab_physx.physics import PhysxManager as SimulationManager @@ -122,7 +122,25 @@ def _initialize_impl(self) -> None: super()._initialize_impl() self._physics_sim_view = SimulationManager.get_physics_sim_view() - root_prim_path_expr = self._resolve_articulation_root_prim_path() + + def has_articulation_root_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.ArticulationRootAPI)) + + matches = resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + asset_prim, root_expr = matches[0] + walk_root = asset_prim.GetPath().pathString + root_prims = get_all_matching_child_prims( + walk_root, predicate=has_articulation_root_api, traverse_instance_prims=False + ) + if len(root_prims) != 1: + matched = [p.GetPath().pathString for p in root_prims] + raise RuntimeError( + f"Expected exactly one ArticulationRootAPI prim under '{walk_root}'" + f" (resolved from '{self.cfg.prim_path}'), found {len(root_prims)}: {matched}." + ) + root_prim_path_expr = root_expr + root_prims[0].GetPath().pathString[len(walk_root) :] self._root_view = self._physics_sim_view.create_articulation_view(root_prim_path_expr.replace(".*", "*")) if self._root_view._backend is None: raise RuntimeError(f"Failed to create articulation view at: {root_prim_path_expr}. Check PhysX logs.") @@ -137,44 +155,16 @@ def _initialize_impl(self) -> None: logger.info(f"Joint wrench sensor initialized: {self._num_envs} envs, {self._num_bodies} bodies") - def _resolve_articulation_root_prim_path(self) -> str: - """Resolve the articulation root prim path expression from the configured asset prim path.""" - 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 - - first_env_root_prims = get_all_matching_child_prims( - first_env_matching_prim_path, - predicate=lambda prim: prim.HasAPI(UsdPhysics.ArticulationRootAPI) - and prim.GetAttribute("physxArticulation:articulationEnabled").Get() is not False, - 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." - ) - - 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) :] - return self.cfg.prim_path + root_prim_path_relative_to_prim_path - def _create_joint_frame_buffers(self) -> None: """Create child-side joint frame transforms indexed by PhysX link order.""" joint_pos_b = np.zeros((self._num_bodies, 3), dtype=np.float32) joint_quat_b = np.zeros((self._num_bodies, 4), dtype=np.float32) joint_quat_b[:, 3] = 1.0 - 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}'.") + matches = resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: + raise RuntimeError(f"No prim found at '{self.cfg.prim_path}'.") + first_env_matching_prim = matches[0][0] link_name_to_index = {name: index for index, name in enumerate(self._data._body_names)} for prim in Usd.PrimRange(first_env_matching_prim): diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py b/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py index 38e24414cfb7..23f840cd4e12 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/pva/pva.py @@ -11,9 +11,8 @@ import torch import warp as wp -from pxr import UsdGeom, UsdPhysics +from pxr import UsdGeom -import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils from isaaclab.markers import VisualizationMarkers from isaaclab.sensors.pva import BasePva @@ -150,29 +149,8 @@ def _initialize_impl(self): super()._initialize_impl() # obtain global simulation view self._physics_sim_view = SimulationManager.get_physics_sim_view() - # check if the prim at path is a rigid prim - prim = sim_utils.find_first_matching_prim(self.cfg.prim_path) - if prim is None: - raise RuntimeError(f"Failed to find a prim at path expression: {self.cfg.prim_path}") - - # Find the first matching ancestor prim that implements rigid body API - ancestor_prim = sim_utils.get_first_matching_ancestor_prim( - prim.GetPath(), predicate=lambda _prim: _prim.HasAPI(UsdPhysics.RigidBodyAPI) - ) - if ancestor_prim is None: - raise RuntimeError(f"Failed to find a rigid body ancestor prim at path expression: {self.cfg.prim_path}") - # Convert ancestor prim path to expression - if ancestor_prim == prim: - self._rigid_parent_expr = self.cfg.prim_path - fixed_pos_b, fixed_quat_b = None, None - else: - # Convert ancestor prim path to expression by stripping the relative - # suffix (including its leading '/') so no trailing '/' remains. - relative_path = prim.GetPath().MakeRelativePath(ancestor_prim.GetPath()).pathString - self._rigid_parent_expr = self.cfg.prim_path.replace("/" + relative_path, "") - # Resolve the relative pose between the target prim and the ancestor prim - fixed_pos_b, fixed_quat_b = sim_utils.resolve_prim_pose(prim, ancestor_prim) + self._rigid_parent_expr, fixed_pos_b, fixed_quat_b = self._resolve_rigid_body_ancestor_expr() # Create the rigid body view on the ancestor self._view = self._physics_sim_view.create_rigid_body_view(self._rigid_parent_expr.replace(".*", "*")) diff --git a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py index 9597abd785d6..d21aa9199b6e 100644 --- a/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py +++ b/source/isaaclab_physx/isaaclab_physx/sensors/ray_caster/ray_caster.py @@ -20,27 +20,8 @@ from isaaclab_physx.physics import PhysxManager -def _find_physics_ancestor(prim): - """Return the nearest rigid-body ancestor for a sensor or target prim.""" - ancestor = prim - while ancestor and ancestor.IsValid() and ancestor.GetPath().pathString != "/": - if ancestor.HasAPI(UsdPhysics.RigidBodyAPI): - return ancestor - ancestor = ancestor.GetParent() - return None - - -def _body_expr_from_sensor_expr(sensor_expr: str, first_sensor_prim, first_body_prim) -> str: - """Convert a sensor/target expression to the matching rigid-body expression.""" - sensor_path = first_sensor_prim.GetPath().pathString - body_path = first_body_prim.GetPath().pathString - if sensor_path == body_path: - return sensor_expr - # Example: ``.../Robot/base/sensor`` target -> ``.../Robot/base`` body view. - suffix = sensor_path[len(body_path) :] - if suffix and sensor_expr.endswith(suffix): - return sensor_expr[: -len(suffix)] - return body_path +def _has_rigid_body_api(prim) -> bool: + return bool(prim.HasAPI(UsdPhysics.RigidBodyAPI)) def _physx_body_glob(body_expr: str) -> str: @@ -62,40 +43,40 @@ def count(self: Any) -> int: return self._view_count def _initialize_pose_tracking(self: Any) -> None: - """Initialize direct PhysX body tracking or a cached static pose table.""" - prims = sim_utils.find_matching_prims(self.cfg.prim_path) - if len(prims) == 0: + """Track the sensor frame through its PhysX rigid-body ancestor, else cache static poses.""" + # One clone-plan-/stage-aware resolution yields the sensor frame(s) and their + # multi-instance destination expressions; the rigid-body view is the frame's ancestor. + matches = sim_utils.resolve_matching_prims_from_source(self.cfg.prim_path) + if not matches: raise RuntimeError(f"No sensor prims matched: {self.cfg.prim_path}") - - # The base classes still use ``self._view.count`` in a few generic - # places. Point it at the sensor instead of constructing an adapter. + # Base classes read ``self._view.count``; the sensor doubles as its own view. self._view = self - body = _find_physics_ancestor(prims[0]) + prims = [prim for prim, _ in matches] + sensor_prim, sensor_expr = matches[0] + body = sim_utils.get_first_matching_ancestor_prim(sensor_prim.GetPath(), predicate=_has_rigid_body_api) if body is None: - self._initialize_static_pose_tracking(prims) + # No rigid-body ancestor: nothing spans envs, so cache every concrete env frame. + self._initialize_static_pose_tracking(sim_utils.find_matching_prims(self.cfg.prim_path)) return - requested_prim_path = getattr(self, "_requested_prim_path", self.cfg.prim_path) - # When the public prim path pointed at a rigid body, BaseRayCaster - # spawned a child sensor prim and preserved the original body path. - body_expr = ( - requested_prim_path - if self.cfg.prim_path != requested_prim_path - else _body_expr_from_sensor_expr(self.cfg.prim_path, prims[0], body) - ) + # The body view is ``sensor_expr`` with the sensor-relative suffix trimmed off. + sensor_path, body_path = sensor_prim.GetPath(), body.GetPath() + relative = sensor_path.MakeRelativePath(body_path).pathString + body_expr = sensor_expr if sensor_path == body_path else sensor_expr[: -(len(relative) + 1)] + physics_sim_view = PhysxManager.get_physics_sim_view() if physics_sim_view is None: raise RuntimeError("PhysX simulation view is not initialized.") - self._physx_body_view = physics_sim_view.create_rigid_body_view(body_expr.replace(".*", "*")) + self._physx_body_view = physics_sim_view.create_rigid_body_view(_physx_body_glob(body_expr)) self._view_count = self._physx_body_view.count - offset_pos = [] - offset_quat = [] + # Sensor-to-body offset per resolved frame; a lone frame broadcasts across all envs. + offset_pos, offset_quat = [], [] for prim in prims: - body_prim = _find_physics_ancestor(prim) - p, q = sim_utils.resolve_prim_pose(prim, body_prim) - offset_pos.append(p) - offset_quat.append(q) + prim_body = sim_utils.get_first_matching_ancestor_prim(prim.GetPath(), predicate=_has_rigid_body_api) + pos, quat = sim_utils.resolve_prim_pose(prim, prim_body) + offset_pos.append(pos) + offset_quat.append(quat) if len(offset_pos) == 1 and self._view_count > 1: offset_pos = offset_pos * self._view_count offset_quat = offset_quat * self._view_count @@ -152,7 +133,7 @@ def _create_tracked_target_view(self: Any, target_prim_paths: str | list[str]): body_paths.append(target_prim_path) continue for prim in prims: - body = _find_physics_ancestor(prim) + body = sim_utils.get_first_matching_ancestor_prim(prim.GetPath(), predicate=_has_rigid_body_api) if body is None: raise RuntimeError( f"Cannot track non-physics ray-cast target '{target_prim_path}' with PhysX. " diff --git a/source/isaaclab_physx/test/sim/test_cloner.py b/source/isaaclab_physx/test/sim/test_cloner.py index a4f1c730fc2d..882963076bb1 100644 --- a/source/isaaclab_physx/test/sim/test_cloner.py +++ b/source/isaaclab_physx/test/sim/test_cloner.py @@ -271,22 +271,22 @@ def test_direct_clone_plan_multi_asset(sim): mass_props=sim_utils.MassPropertiesCfg(mass=1.0), collision_props=sim_utils.CollisionPropertiesCfg(), ) - plan = make_clone_plan( + sources, destinations, clone_mask = make_clone_plan( [[f"/World/envs/env_{i}/Object" for i in range(len(cfg.assets_cfg))]], ["/World/envs/env_{}/Object"], num_clones, sequential, sim.cfg.device, ) - spawn_paths: list[str | None] = list(plan.sources) + spawn_paths: list[str | None] = list(sources) cfg.spawn_paths = spawn_paths prim = cfg.func("/World/unused", cfg) assert prim.IsValid() stage = sim_utils.get_current_stage() env_ids = torch.arange(num_clones, dtype=torch.long, device=sim.cfg.device) - physx_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask, device=sim.cfg.device) - usd_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask) + physx_replicate(stage, sources, destinations, env_ids, clone_mask, device=sim.cfg.device) + usd_replicate(stage, sources, destinations, env_ids, clone_mask) primitive_prims = sim_utils.get_all_matching_child_prims( "/World/envs", predicate=lambda prim: prim.GetTypeName() in ["Cone", "Cube", "Sphere"] @@ -315,7 +315,7 @@ def _run_colocation_collision_filter(sim, asset_cfg, expected_types, assert_coun sim_utils.create_prim(f"/World/envs/env_{i}", "Xform", translation=(0, 0, 0)) num_variants = len(asset_cfg.assets_cfg) if isinstance(asset_cfg, sim_utils.MultiAssetSpawnerCfg) else 1 - plan = make_clone_plan( + sources, destinations, clone_mask = make_clone_plan( [[f"/World/envs/env_{i}/Object" for i in range(num_variants)]], ["/World/envs/env_{}/Object"], num_clones, @@ -323,17 +323,17 @@ def _run_colocation_collision_filter(sim, asset_cfg, expected_types, assert_coun sim.cfg.device, ) if isinstance(asset_cfg, sim_utils.MultiAssetSpawnerCfg): - spawn_paths: list[str | None] = list(plan.sources) + spawn_paths: list[str | None] = list(sources) asset_cfg.spawn_paths = spawn_paths prim = asset_cfg.func("/World/unused", asset_cfg) else: - prim = asset_cfg.func(plan.sources[0], asset_cfg) + prim = asset_cfg.func(sources[0], asset_cfg) assert prim.IsValid() stage = sim_utils.get_current_stage() env_ids = torch.arange(num_clones, dtype=torch.long, device=sim.cfg.device) - physx_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask, device=sim.cfg.device) - usd_replicate(stage, plan.sources, plan.destinations, env_ids, plan.clone_mask) + physx_replicate(stage, sources, destinations, env_ids, clone_mask, device=sim.cfg.device) + usd_replicate(stage, sources, destinations, env_ids, clone_mask) primitive_prims = sim_utils.get_all_matching_child_prims( "/World/envs", predicate=lambda prim: prim.GetTypeName() in expected_types diff --git a/source/isaaclab_tasks/changelog.d/newton-clone-plan.rst b/source/isaaclab_tasks/changelog.d/newton-clone-plan.rst new file mode 100644 index 000000000000..bed038524315 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/newton-clone-plan.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed DexSuite point-cloud sampling in Newton replicated scenes to read + object geometry from clone-plan source prims. 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 fe666a7ecffe..34bcb7929ae7 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 @@ -234,15 +234,15 @@ def _depth_norm(self, images: torch.Tensor) -> torch.Tensor: return images def show_collage(self, images: torch.Tensor, save_path: str = "collage.png"): + import matplotlib import numpy as np - from matplotlib import cm from PIL import Image a = images.detach().cpu().numpy() n, h, w, c = a.shape s = int(np.ceil(np.sqrt(n))) canvas = np.full((s * h, s * w, 3), 255, np.uint8) - turbo = cm.get_cmap("turbo") + turbo = matplotlib.colormaps["turbo"] for i in range(n): r, col = divmod(i, s) img = a[i] diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/utils.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/utils.py index 367155ae7ad5..c092f7471d0f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/utils.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/dexsuite/mdp/utils.py @@ -14,6 +14,7 @@ from pxr import UsdGeom import isaaclab.sim as sim_utils +from isaaclab.cloner.cloner_utils import iter_clone_plan_matches # ---- module-scope caches ---- _PRIM_SAMPLE_CACHE: dict[tuple[str, int], np.ndarray] = {} # (prim_hash, num_points) -> (N,3) in root frame @@ -43,10 +44,12 @@ def sample_object_point_cloud(num_envs: int, num_points: int, prim_path: str, de # Obtain stage handle stage = sim_utils.get_current_stage() - for i in range(num_envs): - # Resolve prim path - obj_path = prim_path.replace(".*", str(i)) + sample_targets: list[tuple[str, tuple[int, ...]]] = [] + clone_plan = sim_utils.SimulationContext.instance().get_clone_plan() + for _, _, source_path, env_ids in iter_clone_plan_matches(clone_plan, prim_path): + sample_targets.append((source_path, env_ids)) + for obj_path, env_ids in sample_targets: # Gather prims prims = sim_utils.get_all_matching_child_prims( obj_path, predicate=lambda p: p.GetTypeName() in ("Mesh", "Cube", "Sphere", "Cylinder", "Capsule", "Cone") @@ -108,7 +111,9 @@ def sample_object_point_cloud(num_envs: int, num_points: int, prim_path: str, de # load from env-level in-memory cache if env_hash in _FINAL_SAMPLE_CACHE: arr = _FINAL_SAMPLE_CACHE[env_hash] # (num_points,3) in root frame - points[i] = torch.from_numpy(arr).to(device) * base_scale.unsqueeze(0) + scaled_samples = torch.from_numpy(arr).to(device) * base_scale.unsqueeze(0) + for env_id in env_ids: + points[env_id] = scaled_samples continue # otherwise build per-prim samples (with per-prim cache) @@ -164,7 +169,9 @@ def sample_object_point_cloud(num_envs: int, num_points: int, prim_path: str, de _FINAL_SAMPLE_CACHE[env_hash] = samples_final.detach().cpu().numpy() # apply root scale and write out - points[i] = samples_final * base_scale.unsqueeze(0) + scaled_samples = samples_final * base_scale.unsqueeze(0) + for env_id in env_ids: + points[env_id] = scaled_samples return points